kailash 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

kailash/api/database.py

@@ -9,9 +9,9 @@ This module provides:
 
 import uuid
 from contextlib import contextmanager
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 from sqlalchemy import (
     JSON,
@@ -383,7 +383,7 @@ class WorkflowRepository:
         tenant_id: str,
         name: str,
         description: str,
-        definition: Dict[str, Any],
+        definition: dict[str, Any],
         created_by: str = None,
     ) -> Workflow:
         """Create a new workflow"""
@@ -403,7 +403,7 @@ class WorkflowRepository:
         return workflow
 
     def update(
-        self, workflow_id: str, updates: Dict[str, Any], updated_by: str = None
+        self, workflow_id: str, updates: dict[str, Any], updated_by: str = None
     ) -> Workflow:
         """Update a workflow"""
         workflow = self.session.query(Workflow).filter_by(id=workflow_id).first()
@@ -433,7 +433,7 @@ class WorkflowRepository:
         self,
         workflow_id: str,
         version: int,
-        definition: Dict[str, Any],
+        definition: dict[str, Any],
         change_message: str,
         created_by: str = None,
     ):
@@ -448,11 +448,11 @@ class WorkflowRepository:
         self.session.add(version_record)
         self.session.commit()
 
-    def get(self, workflow_id: str) -> Optional[Workflow]:
+    def get(self, workflow_id: str) -> Workflow | None:
         """Get a workflow by ID"""
         return self.session.query(Workflow).filter_by(id=workflow_id).first()
 
-    def list(self, tenant_id: str, limit: int = 100, offset: int = 0) -> List[Workflow]:
+    def list(self, tenant_id: str, limit: int = 100, offset: int = 0) -> list[Workflow]:
         """List workflows for a tenant"""
         return (
             self.session.query(Workflow)
@@ -477,14 +477,14 @@ class CustomNodeRepository:
     def __init__(self, session: Session):
         self.session = session
 
-    def create(self, tenant_id: str, node_data: Dict[str, Any]) -> CustomNode:
+    def create(self, tenant_id: str, node_data: dict[str, Any]) -> CustomNode:
         """Create a custom node"""
         node = CustomNode(tenant_id=tenant_id, **node_data)
         self.session.add(node)
         self.session.commit()
         return node
 
-    def update(self, node_id: str, updates: Dict[str, Any]) -> CustomNode:
+    def update(self, node_id: str, updates: dict[str, Any]) -> CustomNode:
         """Update a custom node"""
         node = self.session.query(CustomNode).filter_by(id=node_id).first()
         if not node:
@@ -497,7 +497,7 @@ class CustomNodeRepository:
         self.session.commit()
         return node
 
-    def list(self, tenant_id: str) -> List[CustomNode]:
+    def list(self, tenant_id: str) -> list[CustomNode]:
         """List custom nodes for a tenant"""
         return (
             self.session.query(CustomNode)
@@ -506,7 +506,7 @@ class CustomNodeRepository:
             .all()
         )
 
-    def get(self, node_id: str) -> Optional[CustomNode]:
+    def get(self, node_id: str) -> CustomNode | None:
         """Get a custom node by ID"""
         return self.session.query(CustomNode).filter_by(id=node_id).first()
 
@@ -525,7 +525,7 @@ class ExecutionRepository:
         self.session = session
 
     def create(
-        self, workflow_id: str, tenant_id: str, parameters: Dict[str, Any] = None
+        self, workflow_id: str, tenant_id: str, parameters: dict[str, Any] = None
     ) -> WorkflowExecution:
         """Create an execution record"""
         execution = WorkflowExecution(
@@ -533,7 +533,7 @@ class ExecutionRepository:
             tenant_id=tenant_id,
             status="pending",
             parameters=parameters,
-            started_at=datetime.now(timezone.utc),
+            started_at=datetime.now(UTC),
         )
         self.session.add(execution)
         self.session.commit()
@@ -543,7 +543,7 @@ class ExecutionRepository:
         self,
         execution_id: str,
         status: str,
-        result: Dict[str, Any] = None,
+        result: dict[str, Any] = None,
         error: str = None,
     ):
         """Update execution status"""
@@ -560,7 +560,7 @@ class ExecutionRepository:
             execution.error = error
 
         if status in ["completed", "failed"]:
-            execution.completed_at = datetime.now(timezone.utc)
+            execution.completed_at = datetime.now(UTC)
             if execution.started_at:
                 execution.execution_time_ms = int(
                     (execution.completed_at - execution.started_at).total_seconds()
@@ -569,13 +569,13 @@ class ExecutionRepository:
 
         self.session.commit()
 
-    def get(self, execution_id: str) -> Optional[WorkflowExecution]:
+    def get(self, execution_id: str) -> WorkflowExecution | None:
         """Get execution by ID"""
         return self.session.query(WorkflowExecution).filter_by(id=execution_id).first()
 
     def list_for_workflow(
         self, workflow_id: str, limit: int = 50
-    ) -> List[WorkflowExecution]:
+    ) -> list[WorkflowExecution]:
         """List executions for a workflow"""
         return (
             self.session.query(WorkflowExecution)

kailash/api/gateway.py

@@ -50,7 +50,7 @@ Example:
 import logging
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 from fastapi import FastAPI, WebSocket
 from fastapi.middleware.cors import CORSMiddleware
@@ -69,12 +69,12 @@ class WorkflowRegistration(BaseModel):
 
     name: str
     type: str = Field(description="embedded or proxied")
-    workflow: Optional[Workflow] = None
-    proxy_url: Optional[str] = None
-    health_check: Optional[str] = None
-    description: Optional[str] = None
+    workflow: Workflow | None = None
+    proxy_url: str | None = None
+    health_check: str | None = None
+    description: str | None = None
     version: str = "1.0.0"
-    tags: List[str] = Field(default_factory=list)
+    tags: list[str] = Field(default_factory=list)
 
 
 class WorkflowAPIGateway:
@@ -101,7 +101,7 @@ class WorkflowAPIGateway:
         description: str = "Unified API for Kailash workflows",
         version: str = "1.0.0",
         max_workers: int = 10,
-        cors_origins: List[str] = None,
+        cors_origins: list[str] = None,
     ):
         """Initialize the API gateway.
 
@@ -112,8 +112,8 @@ class WorkflowAPIGateway:
             max_workers: Maximum thread pool workers
             cors_origins: Allowed CORS origins
         """
-        self.workflows: Dict[str, WorkflowRegistration] = {}
-        self.mcp_servers: Dict[str, Any] = {}
+        self.workflows: dict[str, WorkflowRegistration] = {}
+        self.mcp_servers: dict[str, Any] = {}
         self.executor = ThreadPoolExecutor(max_workers=max_workers)
 
         # Create FastAPI app with lifespan
@@ -213,9 +213,9 @@ class WorkflowAPIGateway:
         self,
         name: str,
         workflow: Workflow,
-        description: Optional[str] = None,
+        description: str | None = None,
         version: str = "1.0.0",
-        tags: List[str] = None,
+        tags: list[str] = None,
         **kwargs,
     ):
         """Register an embedded workflow.
@@ -259,9 +259,9 @@ class WorkflowAPIGateway:
         name: str,
         proxy_url: str,
         health_check: str = "/health",
-        description: Optional[str] = None,
+        description: str | None = None,
         version: str = "1.0.0",
-        tags: List[str] = None,
+        tags: list[str] = None,
     ):
         """Register a proxied workflow.
 
@@ -306,7 +306,7 @@ class WorkflowAPIGateway:
         # TODO: Integrate MCP tools with workflows
         logger.info(f"Registered MCP server: {name}")
 
-    def _get_workflow_endpoints(self, name: str) -> List[str]:
+    def _get_workflow_endpoints(self, name: str) -> list[str]:
         """Get endpoints for a workflow."""
         reg = self.workflows.get(name)
         if not reg:
@@ -353,10 +353,10 @@ class WorkflowOrchestrator:
     def __init__(self, gateway: WorkflowAPIGateway):
         """Initialize orchestrator with a gateway."""
         self.gateway = gateway
-        self.chains: Dict[str, List[str]] = {}
-        self.dependencies: Dict[str, List[str]] = {}
+        self.chains: dict[str, list[str]] = {}
+        self.dependencies: dict[str, list[str]] = {}
 
-    def create_chain(self, name: str, workflow_sequence: List[str]):
+    def create_chain(self, name: str, workflow_sequence: list[str]):
         """Create a workflow chain.
 
         Args:
@@ -371,8 +371,8 @@ class WorkflowOrchestrator:
         self.chains[name] = workflow_sequence
 
     async def execute_chain(
-        self, chain_name: str, initial_input: Dict[str, Any]
-    ) -> Dict[str, Any]:
+        self, chain_name: str, initial_input: dict[str, Any]
+    ) -> dict[str, Any]:
         """Execute a workflow chain.
 
         Args:

kailash/api/mcp_integration.py

@@ -28,7 +28,8 @@ Example:
 
 import asyncio
 import logging
-from typing import Any, Callable, Dict, List, Optional, Union
+from collections.abc import Callable
+from typing import Any
 
 from pydantic import BaseModel, Field
 
@@ -42,10 +43,10 @@ class MCPTool(BaseModel):
 
     name: str
     description: str
-    parameters: Dict[str, Any] = Field(default_factory=dict)
-    function: Optional[Callable] = None
-    async_function: Optional[Callable] = None
-    metadata: Dict[str, Any] = Field(default_factory=dict)
+    parameters: dict[str, Any] = Field(default_factory=dict)
+    function: Callable | None = None
+    async_function: Callable | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
 
 
 class MCPResource(BaseModel):
@@ -55,7 +56,7 @@ class MCPResource(BaseModel):
     uri: str
     description: str
     mime_type: str = "text/plain"
-    metadata: Dict[str, Any] = Field(default_factory=dict)
+    metadata: dict[str, Any] = Field(default_factory=dict)
 
 
 class MCPIntegration:
@@ -74,7 +75,7 @@ class MCPIntegration:
     """
 
     def __init__(
-        self, name: str, description: str = "", capabilities: List[str] = None
+        self, name: str, description: str = "", capabilities: list[str] = None
     ):
         """Initialize MCP integration.
 
@@ -86,16 +87,16 @@ class MCPIntegration:
         self.name = name
         self.description = description
         self.capabilities = capabilities or ["tools", "resources"]
-        self.tools: Dict[str, MCPTool] = {}
-        self.resources: Dict[str, MCPResource] = {}
-        self._context: Dict[str, Any] = {}
+        self.tools: dict[str, MCPTool] = {}
+        self.resources: dict[str, MCPResource] = {}
+        self._context: dict[str, Any] = {}
 
     def add_tool(
         self,
         name: str,
-        function: Union[Callable, Callable[..., asyncio.Future]],
+        function: Callable | Callable[..., asyncio.Future],
         description: str = "",
-        parameters: Dict[str, Any] = None,
+        parameters: dict[str, Any] = None,
     ):
         """Add a tool to the MCP server.
 
@@ -134,7 +135,7 @@ class MCPIntegration:
         self.resources[name] = resource
         logger.info(f"Added resource '{name}' to MCP server '{self.name}'")
 
-    async def execute_tool(self, tool_name: str, parameters: Dict[str, Any]) -> Any:
+    async def execute_tool(self, tool_name: str, parameters: dict[str, Any]) -> Any:
         """Execute a tool asynchronously.
 
         Args:
@@ -173,7 +174,7 @@ class MCPIntegration:
         else:
             raise ValueError(f"Tool '{tool_name}' has no implementation")
 
-    def execute_tool_sync(self, tool_name: str, parameters: Dict[str, Any]) -> Any:
+    def execute_tool_sync(self, tool_name: str, parameters: dict[str, Any]) -> Any:
         """Execute a tool synchronously.
 
         Args:
@@ -200,7 +201,7 @@ class MCPIntegration:
         else:
             raise ValueError(f"Tool '{tool_name}' requires async execution")
 
-    def get_resource(self, resource_name: str) -> Optional[MCPResource]:
+    def get_resource(self, resource_name: str) -> MCPResource | None:
         """Get a resource by name.
 
         Args:
@@ -231,7 +232,7 @@ class MCPIntegration:
         """
         return self._context.get(key)
 
-    def list_tools(self) -> List[Dict[str, Any]]:
+    def list_tools(self) -> list[dict[str, Any]]:
         """List all available tools.
 
         Returns:
@@ -246,7 +247,7 @@ class MCPIntegration:
             for tool in self.tools.values()
         ]
 
-    def list_resources(self) -> List[Dict[str, Any]]:
+    def list_resources(self) -> list[dict[str, Any]]:
         """List all available resources.
 
         Returns:
@@ -262,7 +263,7 @@ class MCPIntegration:
             for resource in self.resources.values()
         ]
 
-    def to_mcp_protocol(self) -> Dict[str, Any]:
+    def to_mcp_protocol(self) -> dict[str, Any]:
         """Convert to MCP protocol format.
 
         Returns:
@@ -298,7 +299,7 @@ class MCPToolNode(AsyncNode):
     """
 
     def __init__(
-        self, mcp_server: str, tool_name: str, parameter_mapping: Dict[str, str] = None
+        self, mcp_server: str, tool_name: str, parameter_mapping: dict[str, str] = None
     ):
         """Initialize MCP tool node.
 
@@ -311,13 +312,13 @@ class MCPToolNode(AsyncNode):
         self.mcp_server = mcp_server
         self.tool_name = tool_name
         self.parameter_mapping = parameter_mapping or {}
-        self._mcp_integration: Optional[MCPIntegration] = None
+        self._mcp_integration: MCPIntegration | None = None
 
     def set_mcp_integration(self, mcp: MCPIntegration):
         """Set the MCP integration instance."""
         self._mcp_integration = mcp
 
-    def get_parameters(self) -> Dict[str, Any]:
+    def get_parameters(self) -> dict[str, Any]:
         """Get node parameters.
 
         Returns:
@@ -326,7 +327,7 @@ class MCPToolNode(AsyncNode):
         # For MCP tools, parameters are dynamic based on the tool
         return {}
 
-    def validate_inputs(self, **kwargs) -> Dict[str, Any]:
+    def validate_inputs(self, **kwargs) -> dict[str, Any]:
         """Validate runtime inputs.
 
         For MCPToolNode, we accept any inputs since the parameters
@@ -373,7 +374,7 @@ class MCPToolNode(AsyncNode):
             return {"result": result}
         return result
 
-    async def async_run(self, **kwargs) -> Dict[str, Any]:
+    async def async_run(self, **kwargs) -> dict[str, Any]:
         """Run the node asynchronously.
 
         Args: