Adding APIs to UpsertGraphTemplate (#132)

* added models to create graph_template

* fixed spell check

* fixed comments by @coderabbitai

* added graph create apis

* removed extra f

* added GrpahTemplate to main

Author: NiveditJain

state-manager/app/controller/upsert_graph_template.py

@@ -0,0 +1,58 @@
+from app.singletons.logs_manager import LogsManager
+from app.models.graph_models import UpsertGraphTemplateRequest, UpsertGraphTemplateResponse
+from app.models.db.graph_template_model import GraphTemplate
+from app.models.graph_template_validation_status import GraphTemplateValidationStatus
+from beanie.operators import Set
+
+logger = LogsManager().get_logger()
+
+async def upsert_graph_template(namespace_name: str, graph_name: str, body: UpsertGraphTemplateRequest, x_exosphere_request_id: str) -> UpsertGraphTemplateResponse:
+    try:
+        graph_template = await GraphTemplate.find_one(
+            GraphTemplate.name == graph_name,
+            GraphTemplate.namespace == namespace_name
+        )
+        if graph_template:
+            logger.info(
+                "Graph template already exists in namespace", graph_template=graph_template,
+                namespace_name=namespace_name,
+                x_exosphere_request_id=x_exosphere_request_id)
+            
+            await graph_template.update(
+                Set({
+                    GraphTemplate.nodes: body.nodes, # type: ignore
+                    GraphTemplate.validation_status: GraphTemplateValidationStatus.PENDING, # type: ignore
+                    GraphTemplate.validation_errors: [] # type: ignore
+                })
+            )
+            
+        else:
+            logger.info(
+                "Graph template does not exist in namespace",
+                namespace_name=namespace_name,
+                graph_name=graph_name,
+                x_exosphere_request_id=x_exosphere_request_id)
+            
+            graph_template = await GraphTemplate.insert(
+                GraphTemplate(
+                    name=graph_name,
+                    namespace=namespace_name,
+                    nodes=body.nodes,
+                    validation_status=GraphTemplateValidationStatus.PENDING,
+                    validation_errors=[]
+                )
+            )
+
+        return UpsertGraphTemplateResponse(
+            name=graph_template.name,
+            namespace=graph_template.namespace,
+            nodes=graph_template.nodes,
+            validation_status=graph_template.validation_status,
+            validation_errors=graph_template.validation_errors,
+            created_at=graph_template.created_at,
+            updated_at=graph_template.updated_at
+        )
+    
+    except Exception as e:
+        logger.error("Error upserting graph template", error=e, x_exosphere_request_id=x_exosphere_request_id)
+        raise e

state-manager/app/main.py

@@ -20,6 +20,7 @@
 # injecting models
 from .models.db.state import State
 from .models.db.namespace import Namespace
+from .models.db.graph_template_model import GraphTemplate
 
 # injecting routes
 from .routes import router
@@ -35,7 +36,7 @@ async def lifespan(app: FastAPI):
     # initializing beanie
     client = AsyncMongoClient(os.getenv("MONGO_URI"))
     db = client[os.getenv("MONGO_DATABASE_NAME", "exosphere-state-manager")]
-    await init_beanie(db, document_models=[State, Namespace])
+    await init_beanie(db, document_models=[State, Namespace, GraphTemplate])
     logger.info("beanie dbs initialized")
 
     # initialize secret

state-manager/app/models/db/graph_template_model.py

@@ -0,0 +1,23 @@
+from .base import BaseDatabaseModel
+from pydantic import Field
+from typing import Optional, List
+from ..graph_template_validation_status import GraphTemplateValidationStatus
+from ..node_template_model import NodeTemplate
+from pymongo import IndexModel
+
+
+class GraphTemplate(BaseDatabaseModel):
+    name: str = Field(..., description="Name of the graph")
+    namespace: str = Field(..., description="Namespace of the graph")
+    nodes: List[NodeTemplate] = Field(..., description="Nodes of the graph")
+    validation_status: GraphTemplateValidationStatus = Field(..., description="Validation status of the graph")
+    validation_errors: Optional[List[str]] = Field(None, description="Validation errors of the graph")
+
+    class Settings:
+        indexes = [
+            IndexModel(
+                keys=[("name", 1), ("namespace", 1)],
+                unique=True,
+                name="unique_name_namespace"
+            )
+        ]

state-manager/app/models/graph_models.py

@@ -0,0 +1,21 @@
+from .node_template_model import NodeTemplate
+from pydantic import BaseModel, Field
+from typing import List, Optional
+from datetime import datetime
+from .graph_template_validation_status import GraphTemplateValidationStatus
+
+
+class UpsertGraphTemplateRequest(BaseModel):
+    name: str = Field(..., description="The name of the graph template")
+    namespace: str = Field(..., description="The namespace where the graph template will be stored")
+    nodes: List[NodeTemplate] = Field(..., description="List of node templates that define the graph structure")
+
+
+class UpsertGraphTemplateResponse(BaseModel):
+    name: str = Field(..., description="The name of the graph template")
+    namespace: str = Field(..., description="The namespace where the graph template is stored")
+    nodes: List[NodeTemplate] = Field(..., description="List of node templates that define the graph structure")
+    created_at: datetime = Field(..., description="Timestamp when the graph template was created")
+    updated_at: datetime = Field(..., description="Timestamp when the graph template was last updated")
+    validation_status: GraphTemplateValidationStatus = Field(..., description="Current validation status of the graph template")
+    validation_errors: Optional[List[str]] = Field(None, description="List of validation errors if the graph template is invalid")

state-manager/app/models/graph_template_validation_status.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class GraphTemplateValidationStatus(str, Enum):
+    VALID = "VALID"
+    INVALID = "INVALID"
+    PENDING = "PENDING"
+    ONGOING = "ONGOING"

state-manager/app/models/node_template_model.py

@@ -0,0 +1,11 @@
+from pydantic import Field, BaseModel
+from typing import Any, Optional, List
+
+
+class NodeTemplate(BaseModel):
+    node_name: str = Field(..., description="Name of the node")
+    namespace: str = Field(..., description="Namespace of the node")
+    identifier: str = Field(..., description="Identifier of the node")
+    inputs: dict[str, Any] = Field(..., description="Inputs of the node")
+    store: dict[str, Any] = Field(..., description="Upsert data to store object for the node")
+    next_nodes: Optional[List[str]] = Field(None, description="Next nodes to execute")

state-manager/app/routes.py

@@ -18,15 +18,18 @@
 from .models.errored_models import ErroredRequestModel, ErroredResponseModel
 from .controller.errored_state import errored_state
 
+from .models.graph_models import UpsertGraphTemplateRequest, UpsertGraphTemplateResponse
+from .controller.upsert_graph_template import upsert_graph_template as upsert_graph_template_controller
+
 
 
 logger = LogsManager().get_logger()
 
-router = APIRouter(prefix="/v0/namespace/{namespace_name}/states", tags=["state"])
+router = APIRouter(prefix="/v0/namespace/{namespace_name}", tags=["state"])
 
 
 @router.post(
-    "/enqueue",
+    "/states/enqueue",
     response_model=EnqueueResponseModel,
     status_code=status.HTTP_200_OK,
     response_description="State enqueued on node queue successfully"
@@ -45,7 +48,7 @@ async def enqueue_state(namespace_name: str, body: EnqueueRequestModel, request:
 
 
 @router.post(
-    "/create",
+    "/states/create",
     response_model=CreateResponseModel,
     status_code=status.HTTP_200_OK,
     response_description="States created successfully"
@@ -64,7 +67,7 @@ async def create_state(namespace_name: str, body: CreateRequestModel, request: R
 
 
 @router.post(
-    "/{state_id}/executed",
+    "/states/{state_id}/executed",
     response_model=ExecutedResponseModel,
     status_code=status.HTTP_200_OK,
     response_description="State executed successfully"
@@ -83,7 +86,7 @@ async def executed_state_route(namespace_name: str, state_id: str, body: Execute
 
 
 @router.post(
-    "/{state_id}/errored",
+    "/states/{state_id}/errored",
     response_model=ErroredResponseModel,
     status_code=status.HTTP_200_OK,
     response_description="State errored successfully"
@@ -98,4 +101,22 @@ async def errored_state_route(namespace_name: str, state_id: str, body: ErroredR
         logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
 
-    return await errored_state(namespace_name, ObjectId(state_id), body, x_exosphere_request_id)
+    return await errored_state(namespace_name, ObjectId(state_id), body, x_exosphere_request_id)
+
+
+@router.put(
+    "/graph-templates/{graph_name}",
+    response_model=UpsertGraphTemplateResponse,
+    status_code=status.HTTP_200_OK,
+    response_description="Graph template upserted successfully"
+)   
+async def upsert_graph_template(namespace_name: str, graph_name: str, body: UpsertGraphTemplateRequest, request: Request, api_key: str = Depends(check_api_key)):
+    x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4()))
+
+    if api_key:
+        logger.info(f"API key is valid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+    else:
+        logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
+
+    return await upsert_graph_template_controller(namespace_name, graph_name, body, x_exosphere_request_id)