topoteretes--cognee/cognee/api/client.py

"""FastAPI server for the Cognee API."""

import os

import uvicorn
from traceback import format_exc
from contextlib import asynccontextmanager
from fastapi import Request
from fastapi import FastAPI, status
from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import RequestValidationError
from fastapi.openapi.utils import get_openapi

from cognee.exceptions import CogneeApiError
from cognee.shared.logging_utils import get_logger, setup_logging
from cognee.api.v1.cloud.routers import get_checks_router
from cognee.api.v1.permissions.routers import get_permissions_router
from cognee.api.v1.settings.routers import get_settings_router
from cognee.api.v1.datasets.routers import get_datasets_router
from cognee.api.v1.cognify.routers import get_cognify_router
from cognee.api.v1.search.routers import get_search_router
from cognee.api.v1.ontologies.routers.get_ontology_router import get_ontology_router
from cognee.api.v1.memify.routers import get_memify_router
from cognee.api.v1.add.routers import get_add_router
from cognee.api.v1.delete.routers import get_delete_router
from cognee.api.v1.remember.routers import get_remember_router
from cognee.api.v1.recall.routers import get_recall_router
from cognee.api.v1.improve.routers import get_improve_router
from cognee.api.v1.forget.routers import get_forget_router
from cognee.api.v1.responses.routers import get_responses_router
from cognee.api.v1.llm.routers import get_llm_router
from cognee.api.v1.sync.routers import get_sync_router
from cognee.api.v1.health.routers import get_health_router
from cognee.api.v1.update.routers import get_update_router
from cognee.api.v1.users.routers import (
    get_auth_router,
    get_register_router,
    get_reset_password_router,
    get_verify_router,
    get_users_router,
    get_visualize_router,
    get_configuration_router,
    get_user_id_by_email_router,
)
from cognee.api.v1.api_keys.routers import get_api_key_management_router
from cognee.api.v1.agents.routers import get_agents_router
from cognee.api.v1.visualize.routers import get_schema_router
from cognee.api.v1.skills.routers import get_skills_router
from cognee.api.v1.proposals.routers import get_proposals_router
from cognee.api.v1.activity.routers import get_activity_router
from cognee.api.v1.sessions import get_sessions_router
from cognee.modules.users.methods.get_authenticated_user import REQUIRE_AUTHENTICATION

# Ensure application logging is configured for container stdout/stderr
setup_logging()
logger = get_logger()

app_environment = os.getenv("ENV", "prod")


@asynccontextmanager
async def lifespan(app: FastAPI):
    # from cognee.modules.data.deletion import prune_system, prune_data
    # await prune_data()
    # await prune_system(metadata = True)
    # if app_environment == "local" or app_environment == "dev":
    from cognee.infrastructure.databases.relational import get_relational_engine
    from cognee.run_migrations import run_migrations

    try:
        await run_migrations()
    except Exception:
        db_engine = get_relational_engine()
        await db_engine.create_database()

        await run_migrations()

    from cognee.modules.users.methods import get_default_user

    await get_default_user()
    from cognee.modules.cognify.recovery import recover_stale_cognify_runs_on_startup

    await recover_stale_cognify_runs_on_startup()

    # Emit a clear startup message for docker logs
    logger.info("Backend server has started")

    yield

    # Flush and close all cached database adapters so Ladybug can
    # CHECKPOINT its WAL before the process exits.  Without this,
    # a SIGTERM during an active WAL write leaves a half-written
    # record on disk → "Corrupted wal file" on next startup.
    logger.info("Shutting down: closing cached database engines")
    from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
    from cognee.infrastructure.databases.vector.create_vector_engine import (
        _create_vector_engine,
    )

    _create_graph_engine.cache_clear()
    _create_vector_engine.cache_clear()


app = FastAPI(debug=app_environment != "prod", lifespan=lifespan)


# Read allowed origins from environment variable (comma-separated)
CORS_ALLOWED_ORIGINS = os.getenv("CORS_ALLOWED_ORIGINS")
if CORS_ALLOWED_ORIGINS:
    allowed_origins = [
        origin.strip() for origin in CORS_ALLOWED_ORIGINS.split(",") if origin.strip()
    ]
else:
    allowed_origins = [
        os.getenv("UI_APP_URL", "http://localhost:3000"),
    ]  # Block all except explicitly set origins

app.add_middleware(
    CORSMiddleware,
    allow_origins=allowed_origins,  # Now controlled by env var
    allow_credentials=True,
    allow_methods=["OPTIONS", "GET", "PUT", "POST", "DELETE"],
    allow_headers=["*"],
)
# To allow origins, set CORS_ALLOWED_ORIGINS env variable to a comma-separated list, e.g.:
# CORS_ALLOWED_ORIGINS="https://yourdomain.com,https://another.com"


def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema

    openapi_schema = get_openapi(
        title="Cognee API",
        version="1.0.0",
        description="Cognee API with Bearer token and Cookie auth",
        routes=app.routes,
    )

    # Define security schemes with valid identifiers (no hyphens)
    openapi_schema["components"]["securitySchemes"] = {
        "ApiKeyAuth": {"type": "apiKey", "in": "header", "name": "X-Api-Key"},
        "BearerAuth": {"type": "http", "scheme": "bearer"},
    }

    our_security = [{"BearerAuth": []}, {"ApiKeyAuth": []}]

    if REQUIRE_AUTHENTICATION:
        # Set global security fallback
        openapi_schema["security"] = our_security

        # Replace per-operation security refs auto-generated by fastapi-users
        # (they reference internal names like OAuth2PasswordBearer, APIKeyHeader, APIKeyCookie
        # which no longer exist after we replaced securitySchemes above)
        for path_data in openapi_schema.get("paths", {}).values():
            for operation in path_data.values():
                if isinstance(operation, dict) and "security" in operation:
                    operation["security"] = our_security

    app.openapi_schema = openapi_schema

    return app.openapi_schema


app.openapi = custom_openapi


@app.exception_handler(RequestValidationError)
async def request_validation_exception_handler(request: Request, exc: RequestValidationError):
    if request.url.path == "/api/v1/auth/login":
        return JSONResponse(
            status_code=400,
            content={"detail": "LOGIN_BAD_CREDENTIALS"},
        )

    return JSONResponse(
        status_code=400,
        content=jsonable_encoder({"detail": exc.errors(), "body": exc.body}),
    )


@app.exception_handler(CogneeApiError)
async def exception_handler(_: Request, exc: CogneeApiError) -> JSONResponse:
    detail = {}

    if exc.name and exc.message and exc.status_code:
        status_code = exc.status_code
        detail["message"] = f"{exc.message} [{exc.name}]"
    else:
        # Log an error indicating the exception is improperly defined
        logger.error("Improperly defined exception: %s", exc)
        # Provide a default error response
        detail["message"] = "An unexpected error occurred."
        status_code = status.HTTP_418_IM_A_TEAPOT

    # log the stack trace for easier serverside debugging
    logger.error(format_exc())
    return JSONResponse(status_code=status_code, content={"detail": detail["message"]})


app.include_router(get_auth_router(), prefix="/api/v1/auth", tags=["auth"])

app.include_router(
    get_register_router(),
    prefix="/api/v1/auth",
    tags=["auth"],
)

app.include_router(
    get_reset_password_router(),
    prefix="/api/v1/auth",
    tags=["auth"],
)

app.include_router(
    get_verify_router(),
    prefix="/api/v1/auth",
    tags=["auth"],
)

app.include_router(get_api_key_management_router(), prefix="/api/v1/auth", tags=["auth"])

app.include_router(get_add_router(), prefix="/api/v1/add", tags=["add"])

app.include_router(get_cognify_router(), prefix="/api/v1/cognify", tags=["cognify"])

app.include_router(get_memify_router(), prefix="/api/v1/memify", tags=["memify"])

app.include_router(get_search_router(), prefix="/api/v1/search", tags=["search"])

app.include_router(
    get_permissions_router(),
    prefix="/api/v1/permissions",
    tags=["permissions"],
)

app.include_router(get_datasets_router(), prefix="/api/v1/datasets", tags=["datasets"])

app.include_router(get_ontology_router(), prefix="/api/v1/ontologies", tags=["ontologies"])

app.include_router(get_settings_router(), prefix="/api/v1/settings", tags=["settings"])

app.include_router(get_visualize_router(), prefix="/api/v1/visualize", tags=["visualize"])

app.include_router(get_schema_router(), prefix="/api/v1/schema", tags=["schema"])

app.include_router(get_skills_router(), prefix="/api/v1/skills", tags=["skills"])
app.include_router(get_proposals_router(), prefix="/api/v1/proposals", tags=["skills"])

app.include_router(
    get_configuration_router(),
    prefix="/api/v1/configuration",
    tags=["configuration"],
)

app.include_router(get_delete_router(), prefix="/api/v1/delete", tags=["delete"])

app.include_router(get_update_router(), prefix="/api/v1/update", tags=["update"])

app.include_router(get_responses_router(), prefix="/api/v1/responses", tags=["responses"])

app.include_router(get_llm_router(), prefix="/api/v1/llm", tags=["llm"])

app.include_router(get_sync_router(), prefix="/api/v1/sync", tags=["sync"])

app.include_router(
    get_users_router(),
    prefix="/api/v1/users",
    tags=["users"],
)

app.include_router(
    get_user_id_by_email_router(),
    prefix="/api/v1/users",
    tags=["users"],
)

app.include_router(
    get_checks_router(),
    prefix="/api/v1/checks",
    tags=["checks"],
)

app.include_router(
    get_health_router(),
    prefix="/health",
    tags=["health"],
)

app.include_router(get_agents_router(), prefix="/api/v1/agents")

# Activity / observability
app.include_router(
    get_activity_router(),
    prefix="/api/v1/activity",
    tags=["activity"],
)

# Sessions lifecycle + dashboard aggregates
app.include_router(
    get_sessions_router(),
    prefix="/api/v1/sessions",
    tags=["sessions"],
)

app.include_router(get_remember_router(), prefix="/api/v1/remember", tags=["remember"])
app.include_router(get_recall_router(), prefix="/api/v1/recall", tags=["recall"])
app.include_router(get_improve_router(), prefix="/api/v1/improve", tags=["improve"])
app.include_router(get_forget_router(), prefix="/api/v1/forget", tags=["forget"])


@app.get("/")
async def root():
    """
    Root endpoint that returns a welcome message.
    """
    return {"message": "Hello, World, I am alive!"}


def start_api_server(host: str = "0.0.0.0", port: int = 8000):
    """
    Start the API server using uvicorn.
    Parameters:
    host (str): The host for the server.
    port (int): The port for the server.
    """
    import socket

    try:
        logger.info("Starting server at %s:%s", host, port)

        # Bind before serving so the port is held during startup (including the
        # lifespan migration): uvicorn runs the lifespan before it accepts on this
        # socket, so while migrating, a second server on the same host:port fails
        # fast with EADDRINUSE and no endpoint is served yet (requests queue).
        # reuse_port=False keeps that guard — SO_REUSEPORT would let a second
        # process share the port.
        sock = socket.create_server((host, port), reuse_port=False)

        config = uvicorn.Config(app, host=host, port=port)
        uvicorn.Server(config).run(sockets=[sock])
    except Exception as e:
        logger.exception(f"Failed to start server: {e}")
        # Here you could add any cleanup code or error recovery code.
        raise e


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Cognee API server")
    parser.add_argument(
        "--agent-mode",
        action="store_true",
        default=None,
        help="Enable agent mode (overrides COGNEE_AGENT_MODE env var)",
    )
    args = parser.parse_args()

    from cognee.modules.agents.agent_mode import is_agent_mode_enabled, set_agent_mode

    if args.agent_mode:
        set_agent_mode(True)

    default_port = 8011 if is_agent_mode_enabled() else 8000

    logger = setup_logging()

    start_api_server(
        host=os.getenv("HTTP_API_HOST", "0.0.0.0"),
        port=int(os.getenv("HTTP_API_PORT", default_port)),
    )