aegra-api 0.1.0__py3-none-any.whl

aegra_api/core/sse.py

@@ -0,0 +1,178 @@
+"""Server-Sent Events utilities and formatting"""
+
+import json
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Any
+
+# Import our serializer for handling complex objects
+from aegra_api.core.serializers import GeneralSerializer
+
+# Global serializer instance
+_serializer = GeneralSerializer()
+
+
+def get_sse_headers() -> dict[str, str]:
+    """Get standard SSE headers"""
+    return {
+        "Cache-Control": "no-cache",
+        "Connection": "keep-alive",
+        "Content-Type": "text/event-stream",
+        "Access-Control-Allow-Origin": "*",
+        "Access-Control-Allow-Headers": "Last-Event-ID",
+    }
+
+
+def format_sse_message(
+    event: str,
+    data: Any,
+    event_id: str | None = None,
+    serializer: Callable[[Any], Any] | None = None,
+) -> str:
+    """Format a message as Server-Sent Event following SSE standard
+
+    Args:
+        event: SSE event type
+        data: Data to serialize and send
+        event_id: Optional event ID
+        serializer: Optional custom serializer function
+    """
+    lines = []
+
+    lines.append(f"event: {event}")
+
+    # Convert data to JSON string
+    if data is None:
+        data_str = ""
+    else:
+        # Use our general serializer by default to handle complex objects
+        default_serializer = serializer or _serializer.serialize
+        data_str = json.dumps(data, default=default_serializer, separators=(",", ":"))
+
+    lines.append(f"data: {data_str}")
+
+    if event_id:
+        lines.append(f"id: {event_id}")
+
+    lines.append("")  # Empty line to end the event
+
+    return "\n".join(lines) + "\n"
+
+
+def create_metadata_event(run_id: str, event_id: str | None = None, attempt: int = 1) -> str:
+    """Create metadata event for LangSmith Studio compatibility"""
+    data = {"run_id": run_id, "attempt": attempt}
+    return format_sse_message("metadata", data, event_id)
+
+
+def create_debug_event(debug_data: dict[str, Any], event_id: str | None = None) -> str:
+    """Create debug event with checkpoint fields for LangSmith Studio compatibility"""
+
+    # Add checkpoint and parent_checkpoint fields if not present
+    if "payload" in debug_data and isinstance(debug_data["payload"], dict):
+        payload = debug_data["payload"]
+
+        # Extract checkpoint from config.configurable
+        if "checkpoint" not in payload and "config" in payload:
+            config = payload.get("config", {})
+            if isinstance(config, dict) and "configurable" in config:
+                configurable = config["configurable"]
+                if isinstance(configurable, dict):
+                    payload["checkpoint"] = {
+                        "thread_id": configurable.get("thread_id"),
+                        "checkpoint_id": configurable.get("checkpoint_id"),
+                        "checkpoint_ns": configurable.get("checkpoint_ns", ""),
+                    }
+
+        # Extract parent_checkpoint from parent_config.configurable
+        if "parent_checkpoint" not in payload and "parent_config" in payload:
+            parent_config = payload.get("parent_config")
+            if isinstance(parent_config, dict) and "configurable" in parent_config:
+                configurable = parent_config["configurable"]
+                if isinstance(configurable, dict):
+                    payload["parent_checkpoint"] = {
+                        "thread_id": configurable.get("thread_id"),
+                        "checkpoint_id": configurable.get("checkpoint_id"),
+                        "checkpoint_ns": configurable.get("checkpoint_ns", ""),
+                    }
+            elif parent_config is None:
+                payload["parent_checkpoint"] = None
+
+    return format_sse_message("debug", debug_data, event_id)
+
+
+def create_end_event(event_id: str | None = None) -> str:
+    """Create end event - signals completion of stream
+
+    Uses standard status: "success" instead of "completed"
+    """
+    return format_sse_message("end", {"status": "success"}, event_id)
+
+
+def create_error_event(error: str | dict[str, Any], event_id: str | None = None) -> str:
+    """Create error event with structured error information.
+
+    Error format: {"error": str, "message": str}
+    This format ensures compatibility with standard SSE error event consumers.
+
+    Args:
+        error: Either a simple error string, or a dict with structured error info.
+               Dict format: {"error": "ErrorType", "message": "detailed message"}
+        event_id: Optional SSE event ID for reconnection support.
+
+    Returns:
+        SSE-formatted error event string with standard error format.
+    """
+    if isinstance(error, dict):
+        # Structured error format - standard format: {error: str, message: str}
+        data = {
+            "error": error.get("error", "Error"),
+            "message": error.get("message", str(error)),
+        }
+    else:
+        # Simple string format - wrap it to standard format
+        data = {
+            "error": "Error",
+            "message": str(error),
+        }
+    return format_sse_message("error", data, event_id)
+
+
+def create_messages_event(messages_data: Any, event_type: str = "messages", event_id: str | None = None) -> str:
+    """Create messages event (messages, messages/partial, messages/complete, messages/metadata)"""
+    # Handle tuple format for token streaming: (message_chunk, metadata)
+    if isinstance(messages_data, tuple) and len(messages_data) == 2:
+        message_chunk, metadata = messages_data
+        # Format as expected by LangGraph SDK client
+        data = [message_chunk, metadata]
+        return format_sse_message(event_type, data, event_id)
+    else:
+        # Handle list of messages format
+        return format_sse_message(event_type, messages_data, event_id)
+
+
+# Legacy compatibility - used by event_store.py
+@dataclass
+class SSEEvent:
+    """SSE Event data structure for event storage"""
+
+    id: str
+    event: str
+    data: dict[str, Any]
+    timestamp: datetime | None = None
+
+    def __post_init__(self) -> None:
+        if self.timestamp is None:
+            self.timestamp = datetime.now(UTC)
+
+    def format(self) -> str:
+        """Format as proper SSE event"""
+        json_data = json.dumps(self.data, default=str)
+        return f"id: {self.id}\nevent: {self.event}\ndata: {json_data}\n\n"
+
+
+def format_sse_event(id: str, event: str, data: dict[str, Any]) -> str:
+    """Format SSE event (used by event_store)"""
+    json_data = json.dumps(data, default=str)
+    return f"id: {id}\nevent: {event}\ndata: {json_data}\n\n"

aegra_api/main.py

@@ -0,0 +1,303 @@
+"""FastAPI application for Aegra (Agent Protocol Server)"""
+
+import asyncio
+import sys
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Add graphs directory to Python path so react_agent can be imported
+# This MUST happen before importing any modules that depend on graphs/
+current_dir = Path(__file__).parent.parent.parent  # Go up to aegra root
+graphs_dir = current_dir / "graphs"
+if str(graphs_dir) not in sys.path:
+    sys.path.insert(0, str(graphs_dir))
+
+# ruff: noqa: E402 - imports below require sys.path modification above
+import structlog
+from asgi_correlation_id import CorrelationIdMiddleware
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+from aegra_api.api.assistants import router as assistants_router
+from aegra_api.api.runs import router as runs_router
+from aegra_api.api.store import router as store_router
+from aegra_api.api.threads import router as threads_router
+from aegra_api.config import HttpConfig, get_config_dir, load_http_config
+from aegra_api.core.app_loader import load_custom_app
+from aegra_api.core.auth_deps import auth_dependency
+from aegra_api.core.database import db_manager
+from aegra_api.core.health import router as health_router
+from aegra_api.core.route_merger import (
+    merge_exception_handlers,
+    merge_lifespans,
+)
+from aegra_api.middleware import DoubleEncodedJSONMiddleware, StructLogMiddleware
+from aegra_api.models.errors import AgentProtocolError, get_error_type
+from aegra_api.observability.setup import setup_observability
+from aegra_api.services.event_store import event_store
+from aegra_api.services.langgraph_service import get_langgraph_service
+from aegra_api.settings import settings
+from aegra_api.utils.setup_logging import setup_logging
+
+# Task management for run cancellation
+active_runs: dict[str, asyncio.Task] = {}
+
+setup_logging()
+logger = structlog.getLogger(__name__)
+
+# Default CORS headers required for LangGraph SDK stream reconnection
+DEFAULT_EXPOSE_HEADERS = ["Content-Location", "Location"]
+
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
+    """FastAPI lifespan context manager for startup/shutdown"""
+    # Startup: Initialize database and LangGraph components
+    await db_manager.initialize()
+
+    # Observability
+    setup_observability()
+
+    # Initialize LangGraph service
+    langgraph_service = get_langgraph_service()
+    await langgraph_service.initialize()
+
+    # Initialize event store cleanup task
+    await event_store.start_cleanup_task()
+
+    yield
+
+    # Shutdown: Clean up connections and cancel active runs
+    for task in active_runs.values():
+        if not task.done():
+            task.cancel()
+
+    # Stop event store cleanup task
+    await event_store.stop_cleanup_task()
+
+    await db_manager.close()
+
+
+# Define core exception handlers
+async def agent_protocol_exception_handler(_request: Request, exc: HTTPException) -> JSONResponse:
+    """Convert HTTP exceptions to Agent Protocol error format"""
+    return JSONResponse(
+        status_code=exc.status_code,
+        content=AgentProtocolError(
+            error=get_error_type(exc.status_code),
+            message=exc.detail,
+            details=getattr(exc, "details", None),
+        ).model_dump(),
+    )
+
+
+async def general_exception_handler(_request: Request, exc: Exception) -> JSONResponse:
+    """Handle unexpected exceptions"""
+    return JSONResponse(
+        status_code=500,
+        content=AgentProtocolError(
+            error="internal_error",
+            message="An unexpected error occurred",
+            details={"exception": str(exc)},
+        ).model_dump(),
+    )
+
+
+exception_handlers = {
+    HTTPException: agent_protocol_exception_handler,
+    Exception: general_exception_handler,
+}
+
+
+# Define root endpoint handler
+async def root_handler() -> dict[str, str]:
+    """Root endpoint"""
+    return {
+        "message": settings.app.PROJECT_NAME,
+        "version": settings.app.VERSION,
+        "status": "running",
+    }
+
+
+def _apply_auth_to_routes(app: FastAPI, auth_deps: list[Any]) -> None:
+    """Apply auth dependency to all existing routes in the FastAPI app.
+
+    This function recursively processes all routes including nested routers,
+    adding the auth dependency to each route that doesn't already have it.
+    Auth dependencies are prepended to ensure they run first (fail-fast).
+
+    Args:
+        app: FastAPI application instance
+        auth_deps: List of dependencies to apply (e.g., [Depends(require_auth)])
+    """
+    from fastapi.routing import APIRoute, APIRouter
+
+    def process_routes(routes: list) -> None:
+        """Recursively process routes and nested routers."""
+        for route in routes:
+            if isinstance(route, APIRoute):
+                # Add auth dependency if not already present
+                existing_deps = list(route.dependencies or [])
+                # Check if auth dependency is already present
+                auth_dep_ids = {id(dep) for dep in auth_deps}
+                existing_dep_ids = {id(dep) for dep in existing_deps}
+                if not auth_dep_ids.intersection(existing_dep_ids):
+                    # Prepend auth deps so they run first (fail-fast)
+                    route.dependencies = auth_deps + existing_deps
+            elif isinstance(route, APIRouter):
+                # Process nested router
+                process_routes(route.routes)
+            elif hasattr(route, "routes"):
+                # Handle other route types that have nested routes
+                process_routes(route.routes)
+
+    process_routes(app.routes)
+    logger.info("Applied authentication dependency to custom routes")
+
+
+def _add_cors_middleware(app: FastAPI, cors_config: dict[str, Any] | None) -> None:
+    """Add CORS middleware with config or defaults.
+
+    Args:
+        app: FastAPI application instance
+        cors_config: CORS configuration dict or None for defaults
+    """
+    if cors_config:
+        app.add_middleware(
+            CORSMiddleware,
+            allow_origins=cors_config.get("allow_origins", ["*"]),
+            allow_credentials=cors_config.get("allow_credentials", True),
+            allow_methods=cors_config.get("allow_methods", ["*"]),
+            allow_headers=cors_config.get("allow_headers", ["*"]),
+            expose_headers=cors_config.get("expose_headers", DEFAULT_EXPOSE_HEADERS),
+            max_age=cors_config.get("max_age", 600),
+        )
+    else:
+        app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+            expose_headers=DEFAULT_EXPOSE_HEADERS,
+        )
+
+
+def _add_common_middleware(app: FastAPI, cors_config: dict[str, Any] | None) -> None:
+    """Add common middleware stack in correct order.
+
+    Middleware runs in reverse registration order, so we register:
+    1. DoubleEncodedJSONMiddleware (outermost - runs first)
+    2. CORSMiddleware (handles preflight early)
+    3. CorrelationIdMiddleware (adds request ID)
+    4. StructLogMiddleware (innermost - logs with correlation ID)
+
+    Args:
+        app: FastAPI application instance
+        cors_config: CORS configuration dict or None for defaults
+    """
+    app.add_middleware(StructLogMiddleware)
+    app.add_middleware(CorrelationIdMiddleware)
+    _add_cors_middleware(app, cors_config)
+    app.add_middleware(DoubleEncodedJSONMiddleware)
+
+
+def _include_core_routers(app: FastAPI) -> None:
+    """Include all core API routers with auth dependency.
+
+    Routers are included in consistent order:
+    1. Health (no auth)
+    2. Assistants (with auth)
+    3. Threads (with auth)
+    4. Runs (with auth)
+    5. Store (with auth)
+
+    Args:
+        app: FastAPI application instance
+    """
+    app.include_router(health_router, prefix="", tags=["Health"])
+    app.include_router(assistants_router, dependencies=auth_dependency, prefix="", tags=["Assistants"])
+    app.include_router(threads_router, dependencies=auth_dependency, prefix="", tags=["Threads"])
+    app.include_router(runs_router, dependencies=auth_dependency, prefix="", tags=["Runs"])
+    app.include_router(store_router, dependencies=auth_dependency, prefix="", tags=["Store"])
+
+
+def create_app() -> FastAPI:
+    """Create and configure the FastAPI application.
+
+    Returns:
+        Configured FastAPI application instance
+    """
+    http_config: HttpConfig | None = load_http_config()
+    cors_config = http_config.get("cors") if http_config else None
+
+    # Try to load custom app if configured
+    user_app = None
+    if http_config and http_config.get("app"):
+        try:
+            config_dir = get_config_dir()
+            user_app = load_custom_app(http_config["app"], base_dir=config_dir)
+            logger.info("Custom app loaded successfully")
+        except Exception as e:
+            logger.error(f"Failed to load custom app: {e}", exc_info=True)
+            raise
+
+    if user_app:
+        if not isinstance(user_app, FastAPI):
+            raise TypeError(
+                "Custom apps must be FastAPI applications. Use: from fastapi import FastAPI; app = FastAPI()"
+            )
+
+        application = user_app
+        _include_core_routers(application)
+
+        # Add root endpoint if not already defined
+        if not any(route.path == "/" for route in application.routes if hasattr(route, "path")):
+            application.get("/")(root_handler)
+
+        application = merge_lifespans(application, lifespan)
+        application = merge_exception_handlers(application, exception_handlers)
+        _add_common_middleware(application, cors_config)
+
+        # Apply auth to custom routes if enabled
+        if http_config and http_config.get("enable_custom_route_auth", False):
+            _apply_auth_to_routes(application, auth_dependency)
+    else:
+        application = FastAPI(
+            title=settings.app.PROJECT_NAME,
+            description="Production-ready Agent Protocol server",
+            version=settings.app.VERSION,
+            debug=settings.app.DEBUG,
+            docs_url="/docs",
+            redoc_url="/redoc",
+            lifespan=lifespan,
+        )
+
+        _add_common_middleware(application, cors_config)
+        _include_core_routers(application)
+
+        for exc_type, handler in exception_handlers.items():
+            application.exception_handler(exc_type)(handler)
+
+        application.get("/")(root_handler)
+
+    return application
+
+
+# Create application instance
+app = create_app()
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    port = int(settings.app.PORT)
+    uvicorn.run(app, host=settings.app.HOST, port=port)  # nosec B104 - binding to all interfaces is intentional

aegra_api/middleware/__init__.py

@@ -0,0 +1,4 @@
+from aegra_api.middleware.double_encoded_json import DoubleEncodedJSONMiddleware
+from aegra_api.middleware.logger_middleware import StructLogMiddleware
+
+__all__ = ["DoubleEncodedJSONMiddleware", "StructLogMiddleware"]

aegra_api/middleware/double_encoded_json.py

@@ -0,0 +1,74 @@
+import json
+
+import structlog
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+logger = structlog.getLogger(__name__)
+
+
+class DoubleEncodedJSONMiddleware:
+    """Middleware to handle double-encoded JSON payloads from frontend.
+
+    Some frontend clients may send JSON that's been stringified twice,
+    resulting in payloads like '"{\"key\":\"value\"}"' instead of '{"key":"value"}'.
+    This middleware detects and corrects such cases.
+    """
+
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        method = scope["method"]
+        headers = dict(scope.get("headers", []))
+        content_type = headers.get(b"content-type", b"").decode("latin1")
+
+        if method in ["POST", "PUT", "PATCH"] and content_type:
+            body_parts = []
+
+            async def receive_wrapper() -> dict:
+                message = await receive()
+                if message["type"] == "http.request":
+                    body_parts.append(message.get("body", b""))
+
+                    if not message.get("more_body", False):
+                        body = b"".join(body_parts)
+
+                        if body:
+                            try:
+                                decoded = body.decode("utf-8")
+                                parsed = json.loads(decoded)
+
+                                if isinstance(parsed, str):
+                                    parsed = json.loads(parsed)
+
+                                new_body = json.dumps(parsed).encode("utf-8")
+
+                                if b"content-type" in headers and content_type != "application/json":
+                                    new_headers = []
+                                    for name, value in scope.get("headers", []):
+                                        if name != b"content-type":
+                                            new_headers.append((name, value))
+                                    new_headers.append((b"content-type", b"application/json"))
+                                    scope["headers"] = new_headers
+
+                                return {
+                                    "type": "http.request",
+                                    "body": new_body,
+                                    "more_body": False,
+                                }
+                            except (
+                                json.JSONDecodeError,
+                                ValueError,
+                                UnicodeDecodeError,
+                            ):
+                                pass
+
+                return message
+
+            await self.app(scope, receive_wrapper, send)
+        else:
+            await self.app(scope, receive, send)

aegra_api/middleware/logger_middleware.py

@@ -0,0 +1,95 @@
+import time
+from typing import TypedDict
+
+import structlog
+from asgi_correlation_id import correlation_id
+from starlette.types import ASGIApp, Receive, Scope, Send
+from uvicorn.protocols.utils import get_path_with_query_string
+
+from aegra_api.settings import settings
+
+app_logger = structlog.stdlib.get_logger("app.app_logs")
+access_logger = structlog.stdlib.get_logger("app.access_logs")
+
+
+class AccessInfo(TypedDict, total=False):
+    status_code: int
+    start_time: float
+
+
+class StructLogMiddleware:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        # If the request is not an HTTP request, we don't need to do anything special
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        structlog.contextvars.clear_contextvars()
+        structlog.contextvars.bind_contextvars(request_id=correlation_id.get())
+
+        info = AccessInfo()
+
+        # Inner send function
+        async def inner_send(message):
+            if message.get("type") == "http.response.start":
+                info["status_code"] = message.get("status", 500)
+            await send(message)
+
+        try:
+            info["start_time"] = time.perf_counter_ns()
+            await self.app(scope, receive, inner_send)
+        except Exception as e:
+            # Log the exception here, but re-raise so application-level
+            # exception handlers (e.g. Agent Protocol formatter) can run.
+            app_logger.exception(
+                "An unhandled exception was caught by middleware; re-raising to allow app handlers to format response",
+                exception_class=e.__class__.__name__,
+                exc_info=e,
+                stack_info=True,
+            )
+            raise
+        finally:
+            process_time = time.perf_counter_ns() - info["start_time"]
+            client_host, client_port = scope["client"]
+            http_method = scope["method"]
+            http_version = scope["http_version"]
+            url = get_path_with_query_string(scope)
+
+            # Recreate the Uvicorn access log format, but add all parameters as structured information
+            log_data = {
+                "url": str(url),
+                "status_code": info.get("status_code", 500),
+                "method": http_method,
+                "version": http_version,
+            }
+            if settings.app.LOG_VERBOSITY == "verbose":
+                log_data["request_id"] = correlation_id.get()
+
+            status_code = info.get("status_code", 500)
+            if 400 <= status_code < 500:
+                # Log as warning for client errors (4xx)
+                access_logger.warning(
+                    f"""{client_host}:{client_port} - "{http_method} {scope["path"]} HTTP/{http_version}" {status_code}""",
+                    http=log_data,
+                    network={"client": {"ip": client_host, "port": client_port}},
+                    duration=process_time,
+                )
+            elif 500 <= status_code < 600:
+                # Log as error for server errors (5xx)
+                access_logger.error(
+                    f"""{client_host}:{client_port} - "{http_method} {scope["path"]} HTTP/{http_version}" {status_code}""",
+                    http=log_data,
+                    network={"client": {"ip": client_host, "port": client_port}},
+                    duration=process_time,
+                )
+            else:
+                # Normal log for successful responses (2xx, 3xx)
+                access_logger.info(
+                    f"""{client_host}:{client_port} - "{http_method} {scope["path"]} HTTP/{http_version}" {status_code}""",
+                    http=log_data,
+                    network={"client": {"ip": client_host, "port": client_port}},
+                    duration=process_time,
+                )