kailash 0.4.2__py3-none-any.whl → 0.6.0__py3-none-any.whl

kailash/gateway/api.py

@@ -0,0 +1,289 @@
+"""FastAPI integration for Enhanced Gateway.
+
+This module provides REST API endpoints for the enhanced gateway
+with resource management and async workflow support.
+"""
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union
+
+from fastapi import APIRouter, BackgroundTasks, Depends, FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, Field
+
+from ..resources.registry import ResourceRegistry
+from .enhanced_gateway import (
+    EnhancedDurableAPIGateway,
+    ResourceReference,
+    WorkflowRequest,
+    WorkflowResponse,
+)
+from .security import SecretManager
+
+logger = logging.getLogger(__name__)
+
+
+# Pydantic models for API
+class ResourceReferenceModel(BaseModel):
+    """Model for resource reference in API."""
+
+    type: str = Field(
+        ..., description="Resource type (database, http_client, cache, etc.)"
+    )
+    config: Dict[str, Any] = Field(..., description="Resource configuration")
+    credentials_ref: Optional[str] = Field(
+        None, description="Reference to credentials secret"
+    )
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "type": "database",
+                "config": {"host": "localhost", "port": 5432, "database": "myapp"},
+                "credentials_ref": "db_credentials",
+            }
+        }
+
+
+class WorkflowRequestModel(BaseModel):
+    """API model for workflow requests."""
+
+    inputs: Dict[str, Any] = Field(..., description="Workflow input parameters")
+    resources: Optional[Dict[str, Union[str, ResourceReferenceModel]]] = Field(
+        None,
+        description="Resource references (@name for registered, or inline definition)",
+    )
+    context: Optional[Dict[str, Any]] = Field(
+        None, description="Additional context variables"
+    )
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "inputs": {"user_id": 123, "action": "process"},
+                "resources": {
+                    "db": "@main_database",
+                    "api": {
+                        "type": "http_client",
+                        "config": {"base_url": "https://api.example.com"},
+                        "credentials_ref": "api_key_secret",
+                    },
+                },
+                "context": {"environment": "production", "trace_id": "abc123"},
+            }
+        }
+
+
+class WorkflowResponseModel(BaseModel):
+    """API model for workflow responses."""
+
+    request_id: str
+    workflow_id: str
+    status: str
+    result: Optional[Any] = None
+    error: Optional[str] = None
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    execution_time: Optional[float] = None
+
+
+class WorkflowRegistrationModel(BaseModel):
+    """Model for workflow registration."""
+
+    workflow_definition: Dict[str, Any] = Field(..., description="Workflow definition")
+    required_resources: Optional[List[str]] = Field(
+        None, description="Required resource names"
+    )
+    description: Optional[str] = Field(None, description="Workflow description")
+
+
+# Create router
+router = APIRouter(prefix="/api/v1", tags=["workflows"])
+
+# Dependency to get gateway instance
+_gateway_instance: Optional[EnhancedDurableAPIGateway] = None
+
+
+async def get_gateway() -> EnhancedDurableAPIGateway:
+    """Get or create gateway instance."""
+    global _gateway_instance
+    if not _gateway_instance:
+        _gateway_instance = EnhancedDurableAPIGateway()
+    return _gateway_instance
+
+
+# API Endpoints
+@router.post("/workflows/{workflow_id}/execute", response_model=WorkflowResponseModel)
+async def execute_workflow(
+    workflow_id: str,
+    request: WorkflowRequestModel,
+    background_tasks: BackgroundTasks,
+    gateway: EnhancedDurableAPIGateway = Depends(get_gateway),
+):
+    """Execute a workflow with resource support."""
+    # Convert API model to internal request
+    resources = {}
+    if request.resources:
+        for name, ref in request.resources.items():
+            if isinstance(ref, str):
+                resources[name] = ref
+            elif isinstance(ref, ResourceReferenceModel):
+                resources[name] = ResourceReference(
+                    type=ref.type,
+                    config=ref.config,
+                    credentials_ref=ref.credentials_ref,
+                )
+            elif isinstance(ref, dict):
+                resources[name] = ResourceReference(**ref)
+
+    workflow_request = WorkflowRequest(
+        inputs=request.inputs, resources=resources, context=request.context or {}
+    )
+
+    # Execute workflow
+    try:
+        response = await gateway.execute_workflow(workflow_id, workflow_request)
+        return WorkflowResponseModel(**response.to_dict())
+    except Exception as e:
+        logger.error(f"Workflow execution failed: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get(
+    "/workflows/{workflow_id}/status/{request_id}", response_model=WorkflowResponseModel
+)
+async def get_workflow_status(
+    workflow_id: str,
+    request_id: str,
+    gateway: EnhancedDurableAPIGateway = Depends(get_gateway),
+):
+    """Get status of a workflow execution."""
+    try:
+        response = await gateway.get_workflow_status(request_id)
+        if response.workflow_id != workflow_id:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Request {request_id} is for workflow {response.workflow_id}, not {workflow_id}",
+            )
+        return WorkflowResponseModel(**response.to_dict())
+    except ValueError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+
+
+@router.get("/workflows")
+async def list_workflows(gateway: EnhancedDurableAPIGateway = Depends(get_gateway)):
+    """List all registered workflows."""
+    return gateway.list_workflows()
+
+
+@router.get("/workflows/{workflow_id}")
+async def get_workflow_details(
+    workflow_id: str, gateway: EnhancedDurableAPIGateway = Depends(get_gateway)
+):
+    """Get details of a specific workflow."""
+    workflows = gateway.list_workflows()
+    if workflow_id not in workflows:
+        raise HTTPException(status_code=404, detail=f"Workflow {workflow_id} not found")
+    return workflows[workflow_id]
+
+
+@router.post("/workflows/{workflow_id}/register")
+async def register_workflow(
+    workflow_id: str,
+    registration: WorkflowRegistrationModel,
+    gateway: EnhancedDurableAPIGateway = Depends(get_gateway),
+):
+    """Register a new workflow."""
+    try:
+        # Build workflow from definition
+        from ..workflow import WorkflowBuilder
+
+        builder = WorkflowBuilder()
+
+        # Parse workflow definition
+        for node in registration.workflow_definition.get("nodes", []):
+            builder.add_node(node["type"], node["id"], node.get("config", {}))
+
+        for conn in registration.workflow_definition.get("connections", []):
+            builder.add_connection(
+                conn["from_node"],
+                conn.get("from_output"),
+                conn["to_node"],
+                conn.get("to_input"),
+            )
+
+        workflow = builder.build()
+        workflow.name = workflow_id
+
+        # Register with gateway
+        gateway.register_workflow(
+            workflow_id,
+            workflow,
+            required_resources=registration.required_resources,
+            description=registration.description,
+        )
+
+        return {
+            "status": "registered",
+            "workflow_id": workflow_id,
+            "message": f"Workflow {workflow_id} registered successfully",
+        }
+    except Exception as e:
+        logger.error(f"Workflow registration failed: {e}", exc_info=True)
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.get("/health")
+async def health_check(gateway: EnhancedDurableAPIGateway = Depends(get_gateway)):
+    """Check gateway and resource health."""
+    return await gateway.health_check()
+
+
+@router.get("/resources")
+async def list_resources(gateway: EnhancedDurableAPIGateway = Depends(get_gateway)):
+    """List available resources."""
+    return gateway.resource_registry.list_resources()
+
+
+# Create FastAPI app
+def create_gateway_app(
+    resource_registry: Optional[ResourceRegistry] = None,
+    secret_manager: Optional[SecretManager] = None,
+    title: str = "Kailash Enhanced Gateway",
+    description: str = "API Gateway for async workflows with resource management",
+    version: str = "1.0.0",
+) -> FastAPI:
+    """Create FastAPI app for gateway."""
+    app = FastAPI(title=title, description=description, version=version)
+
+    # Set up gateway instance
+    global _gateway_instance
+    _gateway_instance = EnhancedDurableAPIGateway(
+        resource_registry=resource_registry,
+        secret_manager=secret_manager,
+        title=title,
+        description=description,
+        version=version,
+    )
+
+    # Include router
+    app.include_router(router)
+
+    # Startup event
+    @app.on_event("startup")
+    async def startup_event():
+        """Initialize gateway on startup."""
+        logger.info("Enhanced Gateway starting up...")
+        # Could load workflows from storage here
+
+    # Shutdown event
+    @app.on_event("shutdown")
+    async def shutdown_event():
+        """Cleanup on shutdown."""
+        logger.info("Enhanced Gateway shutting down...")
+        if _gateway_instance and _gateway_instance.resource_registry:
+            await _gateway_instance.resource_registry.cleanup()
+
+    return app

kailash/gateway/enhanced_gateway.py

@@ -0,0 +1,357 @@
+"""Enhanced Gateway with resource management and async workflow support.
+
+This module provides an enhanced version of the DurableAPIGateway that adds:
+- Resource reference resolution for non-serializable objects
+- Integration with ResourceRegistry for shared resources
+- Secret management for credentials
+- Async workflow execution support
+- Health checks for resources
+"""
+
+import asyncio
+import json
+import logging
+import uuid
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from typing import Any, Dict, List, Optional, Set, Union
+
+from ..middleware.gateway.durable_gateway import DurableAPIGateway
+from ..resources.registry import ResourceRegistry
+from ..runtime.async_local import AsyncLocalRuntime, ExecutionContext
+from ..workflow import Workflow
+from .resource_resolver import ResourceReference, ResourceResolver
+from .security import SecretManager
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowNotFoundError(Exception):
+    """Raised when workflow is not found."""
+
+    pass
+
+
+@dataclass
+class WorkflowRequest:
+    """Enhanced workflow request with resource support."""
+
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    inputs: Dict[str, Any] = field(default_factory=dict)
+    resources: Dict[str, Union[str, ResourceReference]] = field(default_factory=dict)
+    context: Dict[str, Any] = field(default_factory=dict)
+    timestamp: datetime = field(default_factory=datetime.utcnow)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to JSON-serializable dict."""
+        return {
+            "request_id": self.request_id,
+            "inputs": self.inputs,
+            "resources": {
+                k: v if isinstance(v, str) else v.to_dict()
+                for k, v in self.resources.items()
+            },
+            "context": self.context,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+
+@dataclass
+class WorkflowResponse:
+    """Response from workflow execution."""
+
+    request_id: str
+    workflow_id: str
+    status: str  # pending, running, completed, failed
+    result: Optional[Any] = None
+    error: Optional[str] = None
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    execution_time: Optional[float] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to JSON-serializable dict."""
+        return {
+            "request_id": self.request_id,
+            "workflow_id": self.workflow_id,
+            "status": self.status,
+            "result": self.result,
+            "error": self.error,
+            "started_at": self.started_at.isoformat() if self.started_at else None,
+            "completed_at": (
+                self.completed_at.isoformat() if self.completed_at else None
+            ),
+            "execution_time": self.execution_time,
+        }
+
+
+class EnhancedDurableAPIGateway(DurableAPIGateway):
+    """Gateway with resource management and async workflow support."""
+
+    def __init__(
+        self,
+        resource_registry: ResourceRegistry = None,
+        secret_manager: SecretManager = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.resource_registry = resource_registry or ResourceRegistry()
+        self.secret_manager = secret_manager or SecretManager()
+        self._workflow_resources: Dict[str, Set[str]] = {}
+        self._resource_resolver = ResourceResolver(
+            self.resource_registry, self.secret_manager
+        )
+        self._runtime = AsyncLocalRuntime(resource_registry=self.resource_registry)
+        self._active_requests: Dict[str, WorkflowResponse] = {}
+        self._cleanup_tasks: List[asyncio.Task] = []
+
+    def register_workflow(
+        self,
+        workflow_id: str,
+        workflow: Workflow,
+        required_resources: List[str] = None,
+        description: str = None,
+    ):
+        """Register workflow with resource requirements."""
+        # Use parent's register_workflow method
+        super().register_workflow(workflow_id, workflow)
+
+        # Track resource requirements
+        if required_resources:
+            self._workflow_resources[workflow_id] = set(required_resources)
+
+        # Extract requirements from workflow metadata
+        if hasattr(workflow, "metadata"):
+            declared_resources = workflow.metadata.get("required_resources", [])
+            if workflow_id not in self._workflow_resources:
+                self._workflow_resources[workflow_id] = set()
+            self._workflow_resources[workflow_id].update(declared_resources)
+
+        # Store workflow description
+        if description and hasattr(workflow, "metadata"):
+            workflow.metadata["description"] = description
+
+    async def execute_workflow(
+        self, workflow_id: str, request: WorkflowRequest
+    ) -> WorkflowResponse:
+        """Execute workflow with resource injection."""
+        # Create response object
+        response = WorkflowResponse(
+            request_id=request.request_id,
+            workflow_id=workflow_id,
+            status="pending",
+            started_at=datetime.now(UTC),
+        )
+
+        # Store active request
+        self._active_requests[request.request_id] = response
+
+        try:
+            # Validate workflow exists
+            if workflow_id not in self.workflows:
+                raise WorkflowNotFoundError(f"Workflow {workflow_id} not found")
+
+            workflow_reg = self.workflows[workflow_id]
+            workflow = workflow_reg.workflow
+
+            # Update status
+            response.status = "running"
+
+            # Prepare execution context
+            context = await self._prepare_execution_context(workflow_id, request)
+
+            # Execute with resource injection
+            result = await self._execute_with_resources(
+                workflow, request.inputs, context
+            )
+
+            # Update response
+            response.status = "completed"
+            response.result = (
+                result.get("results", result) if isinstance(result, dict) else result
+            )
+            response.completed_at = datetime.now(UTC)
+            response.execution_time = (
+                response.completed_at - response.started_at
+            ).total_seconds()
+
+        except Exception as e:
+            # Handle error
+            response.status = "failed"
+            response.error = str(e)
+            response.completed_at = datetime.now(UTC)
+            response.execution_time = (
+                response.completed_at - response.started_at
+            ).total_seconds()
+
+            # Log error
+            logger.error(f"Workflow {workflow_id} failed: {e}", exc_info=True)
+
+        finally:
+            # Clean up active request after a delay
+            cleanup_task = asyncio.create_task(
+                self._cleanup_request(request.request_id)
+            )
+            self._cleanup_tasks.append(cleanup_task)
+
+        return response
+
+    async def _prepare_execution_context(
+        self, workflow_id: str, request: WorkflowRequest
+    ) -> ExecutionContext:
+        """Prepare execution context with resources."""
+        context = ExecutionContext()
+        context.resource_registry = self.resource_registry
+
+        # Add request context variables
+        for key, value in request.context.items():
+            context.set_variable(key, value)
+
+        # Handle resource references in request
+        if request.resources:
+            for name, ref in request.resources.items():
+                if isinstance(ref, ResourceReference):
+                    # Resolve resource reference
+                    resource = await self._resource_resolver.resolve(ref)
+
+                    # Register the resource under the expected name
+                    # Create a wrapper factory that returns the already-created resource
+                    class ExistingResourceFactory:
+                        def __init__(self, resource):
+                            self._resource = resource
+
+                        async def create(self):
+                            return self._resource
+
+                    self.resource_registry.register_factory(
+                        name, ExistingResourceFactory(resource)
+                    )
+
+                elif isinstance(ref, str) and ref.startswith("@"):
+                    # Reference to registered resource
+                    resource_name = ref[1:]  # Remove @ prefix
+                    # Ensure resource exists
+                    if not self.resource_registry.has_factory(resource_name):
+                        raise ValueError(f"Resource '{resource_name}' not registered")
+                    # Resource will be fetched on demand
+
+                elif isinstance(ref, dict) and "type" in ref:
+                    # Inline resource reference
+                    resource_ref = ResourceReference(**ref)
+                    resource = await self._resource_resolver.resolve(resource_ref)
+
+        # Add required resources to context
+        required = self._workflow_resources.get(workflow_id, set())
+        for resource_name in required:
+            if not self.resource_registry.has_factory(resource_name):
+                raise ValueError(f"Required resource '{resource_name}' not available")
+
+        return context
+
+    async def _execute_with_resources(
+        self, workflow: Workflow, inputs: Dict[str, Any], context: ExecutionContext
+    ) -> Any:
+        """Execute workflow with resources."""
+        # Use async runtime for execution
+        result = await self._runtime.execute_workflow_async(workflow, inputs, context)
+        return result
+
+    async def _cleanup_request(self, request_id: str, delay: int = 3600):
+        """Clean up request after delay."""
+        try:
+            await asyncio.sleep(delay)
+            if request_id in self._active_requests:
+                del self._active_requests[request_id]
+        except asyncio.CancelledError:
+            # Task was cancelled during shutdown
+            pass
+
+    async def get_workflow_status(self, request_id: str) -> WorkflowResponse:
+        """Get status of workflow execution."""
+        if request_id in self._active_requests:
+            return self._active_requests[request_id]
+
+        # Could check persistent storage here
+        raise ValueError(f"Request {request_id} not found")
+
+    def list_workflows(self) -> Dict[str, Dict[str, Any]]:
+        """List all registered workflows with metadata."""
+        workflows = {}
+
+        for workflow_id, workflow_reg in self.workflows.items():
+            workflow = workflow_reg.workflow
+            metadata = getattr(workflow, "metadata", {})
+            workflows[workflow_id] = {
+                "name": workflow.name,
+                "description": metadata.get(
+                    "description", workflow_reg.description or ""
+                ),
+                "required_resources": list(
+                    self._workflow_resources.get(workflow_id, [])
+                ),
+                "async_workflow": metadata.get("async_workflow", False),
+                "node_count": len(workflow.nodes),
+                "type": workflow_reg.type,
+                "version": workflow_reg.version,
+                "tags": workflow_reg.tags,
+            }
+
+        return workflows
+
+    async def shutdown(self):
+        """Shutdown the gateway and cleanup resources."""
+        # Cancel all cleanup tasks
+        for task in self._cleanup_tasks:
+            if not task.done():
+                task.cancel()
+
+        # Wait for all tasks to complete
+        if self._cleanup_tasks:
+            await asyncio.gather(*self._cleanup_tasks, return_exceptions=True)
+
+        # Clear the task list
+        self._cleanup_tasks.clear()
+
+        # Clear active requests
+        self._active_requests.clear()
+
+        # Cleanup runtime if it has a cleanup method
+        if hasattr(self._runtime, "cleanup"):
+            await self._runtime.cleanup()
+
+        # Call parent's close method to cleanup middleware components
+        if hasattr(super(), "close"):
+            await super().close()
+
+    async def health_check(self) -> Dict[str, Any]:
+        """Perform health check on gateway and resources."""
+        health = {
+            "status": "healthy",
+            "timestamp": datetime.now(UTC).isoformat(),
+            "workflows": len(self.workflows),
+            "active_requests": len(self._active_requests),
+            "resources": {},
+        }
+
+        # Check resource health
+        for resource_name in self.resource_registry.list_resources():
+            try:
+                resource = await self.resource_registry.get_resource(resource_name)
+                # Try to get health check
+                factory = self.resource_registry._factories.get(resource_name)
+                if (
+                    factory
+                    and hasattr(factory, "health_check")
+                    and factory.health_check
+                ):
+                    is_healthy = await factory.health_check(resource)
+                    health["resources"][resource_name] = (
+                        "healthy" if is_healthy else "unhealthy"
+                    )
+                else:
+                    health["resources"][resource_name] = "healthy"
+            except Exception as e:
+                health["resources"][resource_name] = f"unhealthy: {str(e)}"
+                health["status"] = "degraded"
+
+        return health