cortex_ai/api/server.py

"""OpenAI-compatible HTTP API for CORTEX AI.

Any client that speaks the OpenAI chat-completions protocol -- the official
Python SDK, LangChain, LlamaIndex, a curl script -- can talk to CORTEX AI by
changing only the base URL.

Endpoints:
    GET  /health
    GET  /v1/models
    POST /v1/chat/completions
"""

from __future__ import annotations

import time
import uuid
from typing import Any

from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel, Field

from ..adapters.base import ModelAdapter
from ..config import CortexConfig
from ..engine.agent import CortexAgent
from ..identity import identity_language, identity_response, is_identity_question
from ..tools.registry import registry_from_names


class ChatMessage(BaseModel):
    role: str
    content: str = ""
    name: str | None = None


class ChatCompletionRequest(BaseModel):
    model: str | None = None
    messages: list[ChatMessage]
    temperature: float | None = None
    max_tokens: int | None = None
    stream: bool = False
    tools: list[dict[str, Any]] | None = None


class Usage(BaseModel):
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChatCompletionChoice(BaseModel):
    index: int = 0
    message: ChatMessage
    finish_reason: str = "stop"


class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: list[ChatCompletionChoice]
    usage: Usage
    # CORTEX extension: the reasoning trace, when thinking mode is on.
    reasoning_content: str = ""
    tool_calls: list[dict[str, Any]] = Field(default_factory=list)


def create_app(adapter: ModelAdapter, config: CortexConfig | None = None) -> FastAPI:
    """Build the FastAPI application around a model adapter."""
    cfg = config or CortexConfig()
    tools = registry_from_names(cfg.enabled_tools)
    agent = CortexAgent(
        adapter,
        tools,
        cfg.engine,
        system_prompt=cfg.system_prompt,
    )

    app = FastAPI(
        title="CORTEX AI API",
        version="1.0.0",
        description="API compatible OpenAI pour CORTEX AI, un projet de Frankenstein-Labs.",
    )
    app.state.cortex_config = cfg
    app.state.cortex_agent = agent
    app.state.identity_interception = True

    def _check_auth(authorization: str | None) -> None:
        if not cfg.server.requires_auth:
            return
        expected = f"Bearer {cfg.server.api_key}"
        if authorization != expected:
            raise HTTPException(status_code=401, detail="invalid API key")

    @app.get("/health")
    def health() -> dict[str, Any]:
        return {
            "status": "ok",
            "model": cfg.model_id,
            "tools": tools.names(),
            "thinking_mode": cfg.engine.thinking_mode,
            "identity_interception": "deterministic",
        }

    @app.get("/v1/models")
    def list_models(authorization: str | None = Header(default=None)) -> dict[str, Any]:
        _check_auth(authorization)
        return {
            "object": "list",
            "data": [
                {
                    "id": cfg.model_id,
                    "object": "model",
                    "created": int(time.time()),
                    "owned_by": "Frankenstein-Labs",
                }
            ],
        }

    @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
    def chat_completions(
        request: ChatCompletionRequest,
        authorization: str | None = Header(default=None),
    ) -> ChatCompletionResponse:
        _check_auth(authorization)

        if request.stream:
            raise HTTPException(
                status_code=400,
                detail="stream=true is not supported yet; use stream=false",
            )
        if not request.messages:
            raise HTTPException(status_code=400, detail="messages must not be empty")

        # Deterministic identity boundary: answer before system prompts, tools,
        # or model inference can alter the canonical creator attribution.
        last_user_message = next(
            (m.content for m in reversed(request.messages) if m.role == "user"),
            None,
        )
        if last_user_message is not None and is_identity_question(last_user_message):
            content = identity_response(identity_language(last_user_message))
            prompt_tokens = adapter.count_tokens(last_user_message)
            completion_tokens = adapter.count_tokens(content)
            return ChatCompletionResponse(
                id=f"chatcmpl-{uuid.uuid4().hex[:24]}",
                created=int(time.time()),
                model=request.model or cfg.model_id,
                choices=[
                    ChatCompletionChoice(
                        index=0,
                        message=ChatMessage(role="assistant", content=content),
                        finish_reason="stop",
                    )
                ],
                usage=Usage(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=prompt_tokens + completion_tokens,
                ),
            )

        system_override: list[dict[str, Any]] = []
        turns: list[dict[str, Any]] = []
        for m in request.messages:
            if m.role == "system":
                system_override.append({"role": "system", "content": m.content})
            else:
                turns.append({"role": m.role, "content": m.content})

        if system_override:
            agent.system_prompt = system_override[-1]["content"]

        result = agent.run(turns)
        prompt_tokens = sum(adapter.count_tokens(m["content"]) for m in turns)
        completion_tokens = adapter.count_tokens(result.content)

        return ChatCompletionResponse(
            id=f"chatcmpl-{uuid.uuid4().hex[:24]}",
            created=int(time.time()),
            model=request.model or cfg.model_id,
            choices=[
                ChatCompletionChoice(
                    index=0,
                    message=ChatMessage(role="assistant", content=result.content),
                    finish_reason="stop",
                )
            ],
            usage=Usage(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
            ),
            reasoning_content=result.reasoning,
            tool_calls=[
                {"name": c.name, "arguments": c.arguments, "result": c.result, "ok": c.ok}
                for c in result.tool_calls
            ],
        )

    return app