cursor commit

Author: NiveditJain

state-manager/app/controller/register_nodes.py

@@ -0,0 +1,67 @@
+from ..models.register_nodes_request import RegisterNodesRequestModel
+from ..models.register_nodes_response import RegisterNodesResponseModel, RegisteredNodeModel
+from ..models.db.registered_node import RegisteredNode
+
+from app.singletons.logs_manager import LogsManager
+from beanie.operators import Set
+
+logger = LogsManager().get_logger()
+
+
+async def register_nodes(namespace_name: str, body: RegisterNodesRequestModel, x_exosphere_request_id: str) -> RegisterNodesResponseModel:
+    
+    try:
+        logger.info(f"Registering nodes for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        # Check if nodes already exist and update them, or create new ones
+        registered_nodes = []
+        
+        for node_data in body.nodes:
+            # Check if node already exists
+            existing_node = await RegisteredNode.find_one(
+                RegisteredNode.name == node_data.name,
+                RegisteredNode.namespace == node_data.namespace
+            )
+            
+            if existing_node:
+                # Update existing node
+                await existing_node.update(
+                    Set({
+                        RegisteredNode.runtime_name: body.runtime_name,
+                        RegisteredNode.runtime_namespace: body.runtime_namespace,
+                        RegisteredNode.inputs_schema: node_data.inputs_schema, # type: ignore
+                        RegisteredNode.outputs_schema: node_data.outputs_schema # type: ignore
+                }))
+                logger.info(f"Updated existing node {node_data.name} in namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+            else:
+                # Create new node
+                new_node = RegisteredNode(
+                    name=node_data.name,
+                    namespace=node_data.namespace,
+                    runtime_name=body.runtime_name,
+                    runtime_namespace=body.runtime_namespace,
+                    inputs_schema=node_data.inputs_schema,
+                    outputs_schema=node_data.outputs_schema
+                )
+                inserted_node = await new_node.insert()
+                logger.info(f"Created new node {node_data.name} in namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+            
+            registered_nodes.append(
+                RegisteredNodeModel(
+                    name=node_data.name,
+                    namespace=node_data.namespace
+                )
+            )
+
+        response = RegisterNodesResponseModel(
+            runtime_name=body.runtime_name,
+            runtime_namespace=body.runtime_namespace,
+            registered_nodes=registered_nodes,
+        )
+        
+        logger.info(f"Successfully registered {len(registered_nodes)} nodes for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        return response
+    
+    except Exception as e:
+        logger.error(f"Error registering nodes for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id, error=e)
+        raise e 

state-manager/app/main.py

@@ -21,6 +21,7 @@
 from .models.db.state import State
 from .models.db.namespace import Namespace
 from .models.db.graph_template_model import GraphTemplate
+from .models.db.registered_node import RegisteredNode
 
 # injecting routes
 from .routes import router
@@ -36,7 +37,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, GraphTemplate])
+    await init_beanie(db, document_models=[State, Namespace, GraphTemplate, RegisteredNode])
     logger.info("beanie dbs initialized")
 
     # initialize secret

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

@@ -0,0 +1,12 @@
+from .base import BaseDatabaseModel
+from pydantic import Field
+from typing import Any
+
+
+class RegisteredNode(BaseDatabaseModel):
+    name: str = Field(..., description="Unique name of the registered node")
+    namespace: str = Field(..., description="Namespace of the registered node")
+    runtime_name: str = Field(..., description="Name of the runtime that registered this node")
+    runtime_namespace: str = Field(..., description="Namespace of the runtime that registered this node")
+    inputs_schema: dict[str, Any] = Field(..., description="JSON schema for node inputs")
+    outputs_schema: dict[str, Any] = Field(..., description="JSON schema for node outputs") 

state-manager/app/models/register_nodes_request.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, Field
+from typing import Any, List
+
+
+class NodeRegistrationModel(BaseModel):
+    name: str = Field(..., description="Unique name of the node")
+    namespace: str = Field(..., description="Namespace of the node")
+    inputs_schema: dict[str, Any] = Field(..., description="JSON schema for node inputs")
+    outputs_schema: dict[str, Any] = Field(..., description="JSON schema for node outputs")
+
+
+class RegisterNodesRequestModel(BaseModel):
+    runtime_name: str = Field(..., description="Name of the runtime registering the nodes")
+    runtime_namespace: str = Field(..., description="Namespace of the runtime")
+    nodes: List[NodeRegistrationModel] = Field(..., description="List of nodes to register") 

state-manager/app/models/register_nodes_response.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel, Field
+from typing import List
+
+
+class RegisteredNodeModel(BaseModel):
+    name: str = Field(..., description="Name of the registered node")
+    namespace: str = Field(..., description="Namespace of the registered node")
+
+
+class RegisterNodesResponseModel(BaseModel):
+    runtime_name: str = Field(..., description="Name of the runtime that registered the nodes")
+    runtime_namespace: str = Field(..., description="Namespace of the runtime")
+    registered_nodes: List[RegisteredNodeModel] = Field(..., description="List of successfully registered nodes")

state-manager/app/routes.py

@@ -21,6 +21,10 @@
 from .models.graph_models import UpsertGraphTemplateRequest, UpsertGraphTemplateResponse
 from .controller.upsert_graph_template import upsert_graph_template as upsert_graph_template_controller
 
+from .models.register_nodes_request import RegisterNodesRequestModel
+from .models.register_nodes_response import RegisterNodesResponseModel
+from .controller.register_nodes import register_nodes
+
 
 
 logger = LogsManager().get_logger()
@@ -124,4 +128,23 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
         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)
+    return await upsert_graph_template_controller(namespace_name, graph_name, body, x_exosphere_request_id)
+
+
+@router.put(
+    "/nodes/",
+    response_model=RegisterNodesResponseModel,
+    status_code=status.HTTP_200_OK,
+    response_description="Nodes registered successfully",
+    tags=["nodes"]
+)
+async def register_nodes_route(namespace_name: str, body: RegisterNodesRequestModel, 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 register_nodes(namespace_name, body, x_exosphere_request_id)