aegra-api 0.1.0__py3-none-any.whl

aegra_api/core/health.py

@@ -0,0 +1,131 @@
+"""Health check endpoints"""
+
+import contextlib
+
+from fastapi import APIRouter, HTTPException, Request
+from pydantic import BaseModel
+from sqlalchemy import text
+
+from aegra_api.core.database import db_manager
+
+router = APIRouter()
+
+
+class HealthResponse(BaseModel):
+    """Health check response model"""
+
+    status: str
+    database: str
+    langgraph_checkpointer: str
+    langgraph_store: str
+
+
+class InfoResponse(BaseModel):
+    """Info endpoint response model"""
+
+    name: str
+    version: str
+    description: str
+    status: str
+    flags: dict
+
+
+@router.get("/info", response_model=InfoResponse)
+async def info(_request: Request) -> InfoResponse:
+    """Simple service information endpoint"""
+    return InfoResponse(
+        name="Aegra",
+        version="0.1.0",
+        description="Production-ready Agent Protocol server built on LangGraph",
+        status="running",
+        flags={"assistants": True, "crons": False},
+    )
+
+
+@router.get("/health", response_model=HealthResponse)
+async def health_check(_request: Request) -> dict[str, str]:
+    """Core health check handler logic"""
+    health_status = {
+        "status": "healthy",
+        "database": "unknown",
+        "langgraph_checkpointer": "unknown",
+        "langgraph_store": "unknown",
+    }
+
+    # Database connectivity
+    try:
+        if db_manager.engine:
+            async with db_manager.engine.begin() as conn:
+                await conn.execute(text("SELECT 1"))
+            health_status["database"] = "connected"
+        else:
+            health_status["database"] = "not_initialized"
+            health_status["status"] = "unhealthy"
+    except Exception as e:
+        health_status["database"] = f"error: {str(e)}"
+        health_status["status"] = "unhealthy"
+
+    # LangGraph checkpointer (lazy-init)
+    try:
+        checkpointer = db_manager.get_checkpointer()
+        # probe - will raise if connection is bad; tuple may not exist which is fine
+        with contextlib.suppress(Exception):
+            await checkpointer.aget_tuple({"configurable": {"thread_id": "health-check"}})
+        health_status["langgraph_checkpointer"] = "connected"
+    except Exception as e:
+        health_status["langgraph_checkpointer"] = f"error: {str(e)}"
+        health_status["status"] = "unhealthy"
+
+    # LangGraph store (lazy-init)
+    try:
+        store = db_manager.get_store()
+        with contextlib.suppress(Exception):
+            await store.aget(("health",), "check")
+        health_status["langgraph_store"] = "connected"
+    except Exception as e:
+        health_status["langgraph_store"] = f"error: {str(e)}"
+        health_status["status"] = "unhealthy"
+
+    if health_status["status"] == "unhealthy":
+        raise HTTPException(status_code=503, detail="Service unhealthy")
+
+    return health_status
+
+
+@router.get("/ready")
+async def readiness_check(_request: Request) -> dict[str, str]:
+    """Kubernetes readiness probe endpoint"""
+    # Engine must exist and respond to a trivial query
+    if not db_manager.engine:
+        raise HTTPException(
+            status_code=503,
+            detail="Service not ready - database engine not initialized",
+        )
+    try:
+        async with db_manager.engine.begin() as conn:
+            await conn.execute(text("SELECT 1"))
+    except Exception as e:
+        raise HTTPException(status_code=503, detail=f"Service not ready - database error: {str(e)}") from e
+
+    # Check that LangGraph components can be obtained (lazy init) and respond
+    try:
+        checkpointer = db_manager.get_checkpointer()
+        store = db_manager.get_store()
+        # lightweight probes
+        with contextlib.suppress(Exception):
+            await checkpointer.aget_tuple({"configurable": {"thread_id": "ready-check"}})
+        with contextlib.suppress(Exception):
+            await store.aget(("ready",), "check")
+    except Exception as e:
+        raise HTTPException(
+            status_code=503,
+            detail=f"Service not ready - components unavailable: {str(e)}",
+        ) from e
+
+    return {"status": "ready"}
+
+
+@router.get("/live")
+async def liveness_check(_request: Request) -> dict[str, str]:
+    """Kubernetes liveness probe endpoint"""
+    return {"status": "alive"}

aegra_api/core/orm.py

@@ -0,0 +1,165 @@
+"""SQLAlchemy ORM setup for persistent assistant/thread/run records.
+
+This module creates:
+• `Base` – the declarative base used by our models.
+• `Assistant`, `Thread`, `Run` – ORM models mirroring the bootstrap tables
+  already created in ``DatabaseManager._create_metadata_tables``.
+• `async_session_maker` – a factory that hands out `AsyncSession` objects
+  bound to the shared engine managed by `db_manager`.
+• `get_session` – FastAPI dependency helper for routers.
+
+Nothing is auto-imported by FastAPI yet; routers will `from ...core.db import get_session`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from datetime import datetime
+
+from sqlalchemy import (
+    TIMESTAMP,
+    ForeignKey,
+    Index,
+    Integer,
+    Text,
+    text,
+)
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from sqlalchemy.orm import Mapped, declarative_base, mapped_column
+
+Base = declarative_base()
+
+
+class Assistant(Base):
+    __tablename__ = "assistant"
+
+    # TEXT PK with DB-side generation using uuid_generate_v4()::text
+    assistant_id: Mapped[str] = mapped_column(
+        Text, primary_key=True, server_default=text("public.uuid_generate_v4()::text")
+    )
+    name: Mapped[str] = mapped_column(Text, nullable=False)
+    description: Mapped[str | None] = mapped_column(Text)
+    graph_id: Mapped[str] = mapped_column(Text, nullable=False)
+    config: Mapped[dict] = mapped_column(JSONB, server_default=text("'{}'::jsonb"))
+    context: Mapped[dict] = mapped_column(JSONB, server_default=text("'{}'::jsonb"))
+    user_id: Mapped[str] = mapped_column(Text, nullable=False)
+    version: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("1"))
+    metadata_dict: Mapped[dict] = mapped_column(JSONB, server_default=text("'{}'::jsonb"), name="metadata")
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+    updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+
+    # Indexes for performance
+    __table_args__ = (
+        Index("idx_assistant_user", "user_id"),
+        Index("idx_assistant_user_assistant", "user_id", "assistant_id", unique=True),
+        Index(
+            "idx_assistant_user_graph_config",
+            "user_id",
+            "graph_id",
+            "config",
+            unique=True,
+        ),
+    )
+
+
+class AssistantVersion(Base):
+    __tablename__ = "assistant_versions"
+
+    assistant_id: Mapped[str] = mapped_column(
+        Text, ForeignKey("assistant.assistant_id", ondelete="CASCADE"), primary_key=True
+    )
+    version: Mapped[int] = mapped_column(Integer, primary_key=True)
+    graph_id: Mapped[str] = mapped_column(Text, nullable=False)
+    config: Mapped[dict | None] = mapped_column(JSONB)
+    context: Mapped[dict | None] = mapped_column(JSONB)
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+    metadata_dict: Mapped[dict] = mapped_column(JSONB, server_default=text("'{}'::jsonb"), name="metadata")
+    name: Mapped[str | None] = mapped_column(Text)
+    description: Mapped[str | None] = mapped_column(Text)
+
+
+class Thread(Base):
+    __tablename__ = "thread"
+
+    thread_id: Mapped[str] = mapped_column(Text, primary_key=True)
+    status: Mapped[str] = mapped_column(Text, server_default=text("'idle'"))
+    # Database column is 'metadata_json' (per database.py). ORM attribute 'metadata_json' must map to that column.
+    metadata_json: Mapped[dict] = mapped_column("metadata_json", JSONB, server_default=text("'{}'::jsonb"))
+    user_id: Mapped[str] = mapped_column(Text, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+    updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+
+    # Indexes for performance
+    __table_args__ = (Index("idx_thread_user", "user_id"),)
+
+
+class Run(Base):
+    __tablename__ = "runs"
+
+    # TEXT PK with DB-side generation using uuid_generate_v4()::text
+    run_id: Mapped[str] = mapped_column(Text, primary_key=True, server_default=text("public.uuid_generate_v4()::text"))
+    thread_id: Mapped[str] = mapped_column(Text, ForeignKey("thread.thread_id", ondelete="CASCADE"), nullable=False)
+    assistant_id: Mapped[str | None] = mapped_column(Text, ForeignKey("assistant.assistant_id", ondelete="CASCADE"))
+    status: Mapped[str] = mapped_column(Text, server_default=text("'pending'"))
+    input: Mapped[dict | None] = mapped_column(JSONB, server_default=text("'{}'::jsonb"))
+    # Some environments may not yet have a 'config' column; make it nullable without default to match existing DB.
+    # If migrations add this column later, it's already represented here.
+    config: Mapped[dict | None] = mapped_column(JSONB, nullable=True)
+    context: Mapped[dict | None] = mapped_column(JSONB, nullable=True)
+    output: Mapped[dict | None] = mapped_column(JSONB)
+    error_message: Mapped[str | None] = mapped_column(Text)
+    user_id: Mapped[str] = mapped_column(Text, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+    updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+
+    # Indexes for performance
+    __table_args__ = (
+        Index("idx_runs_thread_id", "thread_id"),
+        Index("idx_runs_user", "user_id"),
+        Index("idx_runs_status", "status"),
+        Index("idx_runs_assistant_id", "assistant_id"),
+        Index("idx_runs_created_at", "created_at"),
+    )
+
+
+class RunEvent(Base):
+    __tablename__ = "run_events"
+
+    id: Mapped[str] = mapped_column(Text, primary_key=True)
+    run_id: Mapped[str] = mapped_column(Text, nullable=False)
+    seq: Mapped[int] = mapped_column(Integer, nullable=False)
+    event: Mapped[str] = mapped_column(Text, nullable=False)
+    data: Mapped[dict | None] = mapped_column(JSONB, nullable=True)
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=text("now()"))
+
+    # Indexes for performance
+    __table_args__ = (
+        Index("idx_run_events_run_id", "run_id"),
+        Index("idx_run_events_seq", "run_id", "seq"),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Session factory
+# ---------------------------------------------------------------------------
+
+async_session_maker: async_sessionmaker[AsyncSession] | None = None
+
+
+def _get_session_maker() -> async_sessionmaker[AsyncSession]:
+    """Return a cached async_sessionmaker bound to db_manager.engine."""
+    global async_session_maker
+    if async_session_maker is None:
+        from aegra_api.core.database import db_manager
+
+        engine = db_manager.get_engine()
+        async_session_maker = async_sessionmaker(engine, expire_on_commit=False)
+    return async_session_maker
+
+
+async def get_session() -> AsyncIterator[AsyncSession]:
+    """FastAPI dependency that yields an AsyncSession."""
+    maker = _get_session_maker()
+    async with maker() as session:
+        yield session

aegra_api/core/route_merger.py

@@ -0,0 +1,69 @@
+"""Route merging utilities for combining custom apps with Aegra core routes"""
+
+from collections.abc import Callable
+from contextlib import asynccontextmanager
+
+import structlog
+from fastapi import FastAPI
+
+logger = structlog.get_logger(__name__)
+
+
+def merge_lifespans(user_app: FastAPI, core_lifespan: Callable) -> FastAPI:
+    """Merge user lifespan with Aegra's core lifespan.
+
+    Both lifespans will run, with core lifespan wrapping user lifespan.
+    This ensures Aegra's initialization (database, services) happens before
+    user initialization, and cleanup happens in reverse order.
+
+    Args:
+        user_app: User's FastAPI/Starlette application
+        core_lifespan: Aegra's core lifespan context manager
+
+    Returns:
+        Modified user_app with merged lifespan
+    """
+    user_lifespan = user_app.router.lifespan_context
+
+    # Check for deprecated on_startup/on_shutdown handlers
+    if user_app.router.on_startup or user_app.router.on_shutdown:
+        raise ValueError(
+            f"Cannot merge lifespans with on_startup or on_shutdown handlers. "
+            f"Please use lifespan context manager instead. "
+            f"Found: on_startup={user_app.router.on_startup}, "
+            f"on_shutdown={user_app.router.on_shutdown}"
+        )
+
+    @asynccontextmanager
+    async def combined_lifespan(app):
+        async with core_lifespan(app):
+            if user_lifespan:
+                async with user_lifespan(app):
+                    yield
+            else:
+                yield
+
+    user_app.router.lifespan_context = combined_lifespan
+    return user_app
+
+
+def merge_exception_handlers(user_app: FastAPI, core_exception_handlers: dict[type, Callable]) -> FastAPI:
+    """Merge core exception handlers with user exception handlers.
+
+    Core handlers are added only if user hasn't defined a handler for that exception type.
+    User handlers take precedence.
+
+    Args:
+        user_app: User's FastAPI/Starlette application
+        core_exception_handlers: Aegra's core exception handlers
+
+    Returns:
+        Modified user_app with merged exception handlers
+    """
+    for exc_type, handler in core_exception_handlers.items():
+        if exc_type not in user_app.exception_handlers:
+            user_app.exception_handlers[exc_type] = handler
+        else:
+            logger.debug(f"User app overrides exception handler for {exc_type}")
+
+    return user_app

aegra_api/core/serializers/__init__.py

@@ -0,0 +1,7 @@
+"""Serialization layer for LangGraph and general objects"""
+
+from aegra_api.core.serializers.base import Serializer
+from aegra_api.core.serializers.general import GeneralSerializer
+from aegra_api.core.serializers.langgraph import LangGraphSerializer
+
+__all__ = ["Serializer", "GeneralSerializer", "LangGraphSerializer"]

aegra_api/core/serializers/base.py

@@ -0,0 +1,22 @@
+"""Base serialization interface"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Serializer(ABC):
+    """Abstract base class for object serialization"""
+
+    @abstractmethod
+    def serialize(self, obj: Any) -> Any:
+        """Serialize an object to a JSON-compatible format"""
+        pass
+
+
+class SerializationError(Exception):
+    """Raised when serialization fails"""
+
+    def __init__(self, message: str, obj_type: str, original_error: Exception | None = None):
+        super().__init__(message)
+        self.obj_type = obj_type
+        self.original_error = original_error

aegra_api/core/serializers/general.py

@@ -0,0 +1,54 @@
+"""General-purpose object serialization for complex objects"""
+
+from typing import Any
+
+from aegra_api.core.serializers.base import SerializationError, Serializer
+
+
+class GeneralSerializer(Serializer):
+    """Simple object serializer for complex Python objects"""
+
+    def serialize(self, obj: Any) -> Any:
+        """Serialize any object to JSON-compatible format"""
+        try:
+            return self._serialize_object(obj)
+        except Exception as e:
+            raise SerializationError(f"Failed to serialize object: {str(e)}", obj.__class__.__name__, e) from e
+
+    def _serialize_object(self, obj: Any) -> Any:
+        """Core serialization logic for Python objects"""
+        # Handle Pydantic v2 models (model_dump method)
+        if hasattr(obj, "model_dump") and callable(obj.model_dump):
+            return obj.model_dump()
+
+        # Handle LangChain objects and Pydantic v1 models (dict method)
+        elif hasattr(obj, "dict") and callable(obj.dict):
+            return obj.dict()
+
+        # Handle LangGraph Interrupt objects (they don't have .dict() method)
+        elif obj.__class__.__name__ == "Interrupt" and hasattr(obj, "value") and hasattr(obj, "id"):
+            return {"value": self._serialize_object(obj.value), "id": obj.id}
+
+        # Handle NamedTuples (like PregelTask) - they have _asdict() method
+        elif hasattr(obj, "_asdict") and callable(obj._asdict):
+            return {k: self._serialize_object(v) for k, v in obj._asdict().items()}
+
+        # Handle sets and frozensets
+        elif isinstance(obj, (set, frozenset)):
+            return list(obj)
+
+        # Handle tuples and lists recursively
+        elif isinstance(obj, (tuple, list)):
+            return [self._serialize_object(item) for item in obj]
+
+        # Handle dictionaries recursively
+        elif isinstance(obj, dict):
+            return {k: self._serialize_object(v) for k, v in obj.items()}
+
+        # Handle basic JSON-serializable types
+        elif isinstance(obj, (str, int, float, bool, type(None))):
+            return obj
+
+        # Fallback to string representation for unknown types
+        else:
+            return str(obj)

aegra_api/core/serializers/langgraph.py

@@ -0,0 +1,102 @@
+"""LangGraph-specific serialization"""
+
+import json
+from typing import Any
+
+import structlog
+
+from aegra_api.core.serializers.base import SerializationError, Serializer
+from aegra_api.core.serializers.general import GeneralSerializer
+
+logger = structlog.getLogger(__name__)
+
+
+class LangGraphSerializer(Serializer):
+    """Handles serialization of LangGraph objects (tasks, interrupts, snapshots)"""
+
+    def __init__(self):
+        self.general_serializer = GeneralSerializer()
+
+    def serialize(self, obj: Any) -> Any:
+        """Main serialization entry point"""
+        return json.loads(json.dumps(obj, default=self.general_serializer.serialize))
+
+    def serialize_task(self, task: Any) -> dict[str, Any]:
+        """Serialize a LangGraph task to ThreadTask format"""
+        try:
+            if hasattr(task, "id") and hasattr(task, "name"):
+                # Proper task object
+                task_dict = {
+                    "id": getattr(task, "id", ""),
+                    "name": getattr(task, "name", ""),
+                    "error": getattr(task, "error", None),
+                    "interrupts": [],
+                    "checkpoint": None,
+                    "state": getattr(task, "state", None),
+                    "result": getattr(task, "result", None),
+                }
+
+                # Handle task interrupts
+                if hasattr(task, "interrupts") and task.interrupts:
+                    task_dict["interrupts"] = self.serialize(task.interrupts)
+
+                return task_dict
+            else:
+                # Raw task data - serialize as-is but safely
+                serialized_task = self.serialize(task)
+                if isinstance(serialized_task, dict):
+                    return serialized_task
+                else:
+                    raise SerializationError(
+                        f"Task serialization resulted in non-dict: {type(serialized_task)}",
+                        task.__class__.__name__,
+                    )
+        except Exception as e:
+            if isinstance(e, SerializationError):
+                raise
+            raise SerializationError(f"Failed to serialize task: {str(e)}", task.__class__.__name__, e) from e
+
+    def serialize_interrupt(self, interrupt: Any) -> dict[str, Any]:
+        """Serialize a LangGraph interrupt"""
+        try:
+            return self.serialize(interrupt)
+        except Exception as e:
+            raise SerializationError(
+                f"Failed to serialize interrupt: {str(e)}",
+                interrupt.__class__.__name__,
+                e,
+            ) from e
+
+    def extract_tasks_from_snapshot(self, snapshot: Any) -> list[dict[str, Any]]:
+        """Extract and serialize tasks from a snapshot"""
+        tasks = []
+
+        if not (hasattr(snapshot, "tasks") and snapshot.tasks):
+            return tasks
+
+        for task in snapshot.tasks:
+            try:
+                serialized_task = self.serialize_task(task)
+                tasks.append(serialized_task)
+            except SerializationError as e:
+                logger.warning(
+                    f"Task serialization failed, skipping task: {e} "
+                    f"(task_type={type(task).__name__}, task_id={getattr(task, 'id', 'unknown')})"
+                )
+                continue
+
+        return tasks
+
+    def extract_interrupts_from_snapshot(self, snapshot: Any) -> list[dict[str, Any]]:
+        """Extract and serialize interrupts from a snapshot"""
+        interrupts = []
+        if hasattr(snapshot, "interrupts") and snapshot.interrupts:
+            try:
+                interrupts = self.serialize(snapshot.interrupts)
+                if interrupts:
+                    return interrupts
+            except Exception as e:
+                logger.warning(
+                    f"Snapshot interrupt serialization failed: {e} (snapshot_type={type(snapshot).__name__})"
+                )
+        return interrupts if isinstance(interrupts, list) else []