main.py

import json
import logging
import time
import uuid
from contextlib import asynccontextmanager
from os import getenv
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.responses import (
    FileResponse,
    HTMLResponse,
    JSONResponse,
    StreamingResponse,
)
from pydantic import BaseModel, Field
from websearch_agent.agent import get_workflow_closure
from websearch_agent.tracing import enable_tracing
from websearch_agent.workflow import InputEvent, ToolCallEvent

logger = logging.getLogger(__name__)


# OpenAI-compatible request/response models
class ChatMessage(BaseModel):
    """A message in the conversation."""

    role: str = Field(
        ...,
        description="The role of the message author.",
        examples=["user", "assistant", "system", "tool"],
    )
    content: str = Field(
        ...,
        description="The contents of the message.",
        examples=["Which company is consider the best?"],
    )


class ChatCompletionRequest(BaseModel):
    """Creates a model response for the given chat conversation.

    [See OpenAI docs](https://platform.openai.com/docs/api-reference/chat/create)
    """

    messages: list[ChatMessage] = Field(
        ..., description="A list of messages comprising the conversation so far."
    )
    model: str | None = Field(
        None,
        description="ID of the model to use. Defaults to the server's configured MODEL_ID.",
    )
    stream: bool = Field(
        False,
        description="If true, partial message deltas will be sent as SSE `data: {json}\\n\\n` events, terminated by `data: [DONE]\\n\\n`.",
    )


class ChoiceMessage(BaseModel):
    """A chat completion message generated by the model."""

    role: str = Field(
        "assistant", description="The role of the author of this message."
    )
    content: str = Field(..., description="The contents of the message.")


class Choice(BaseModel):
    index: int = Field(
        ..., description="The index of the choice in the list of choices."
    )
    message: ChoiceMessage
    finish_reason: str = Field(
        ...,
        description="The reason the model stopped generating tokens.",
        examples=["stop", "tool_calls"],
    )


class ChatCompletionResponse(BaseModel):
    """Represents a chat completion response returned by the model, based on the provided input."""

    id: str = Field(
        ...,
        description="A unique identifier for the chat completion.",
        examples=["chatcmpl-abc123def456"],
    )
    object: str = Field(
        "chat.completion",
        description="The object type, which is always `chat.completion`.",
    )
    created: int = Field(
        ...,
        description="The Unix timestamp (in seconds) of when the chat completion was created.",
    )
    model: str = Field(..., description="The model used for the chat completion.")
    choices: list[Choice] = Field(..., description="A list of chat completion choices.")

    usage: dict | None = Field(
        None, description="Usage statistics for the completion request."
    )


class HealthResponse(BaseModel):
    """Service health status."""

    status: str = Field(
        ..., description="Current service status.", examples=["healthy"]
    )
    agent_initialized: bool = Field(
        ...,
        description="Whether the agent has been initialized and is ready to serve requests.",
    )


# Global variable for workflow closure (get_agent callable)
get_agent = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize the LlamaIndex workflow closure on startup and clear it on shutdown."""
    global get_agent

    enable_tracing()

    base_url = getenv("BASE_URL")
    model_id = getenv("MODEL_ID")

    if base_url and not base_url.endswith("/v1"):
        base_url = base_url.rstrip("/") + "/v1"

    get_agent = get_workflow_closure(model_id=model_id, base_url=base_url)

    yield

    get_agent = None


# Create FastAPI app
app = FastAPI(
    title="LlamaIndex Websearch Agent API",
    description="FastAPI service for LlamaIndex Websearch Agent with OpenAI-compatible chat completions API.",
    lifespan=lifespan,
    openapi_tags=[
        {"name": "Chat", "description": "Chat completion operations"},
        {"name": "Health", "description": "Service health monitoring"},
    ],
)


def _get_message_content(msg) -> str:
    """Extract text content from a LlamaIndex ChatMessage."""
    if hasattr(msg, "blocks") and msg.blocks:
        # Find the first block with text content (skip ToolCallBlock)
        for block in msg.blocks:
            if hasattr(block, "text"):
                return block.text or ""
        return ""
    if hasattr(msg, "content"):
        if isinstance(msg.content, str):
            return msg.content
        if isinstance(msg.content, list) and msg.content:
            first = msg.content[0]
            if isinstance(first, dict) and "text" in first:
                return first["text"] or ""
    return ""


def _message_to_response_dict(msg):
    """Map a LlamaIndex ChatMessage to OpenAI-compatible format."""
    role = getattr(msg, "role", "user")
    content = _get_message_content(msg)

    if role == "user":
        return {"role": "user", "content": content}

    if role == "assistant":
        msg_data = {"role": "assistant", "content": content or ""}
        tool_calls = getattr(msg, "tool_calls", None)
        if not tool_calls and getattr(msg, "additional_kwargs", None):
            tool_calls = msg.additional_kwargs.get("tool_calls")
        if tool_calls:
            if hasattr(tool_calls[0], "tool_id"):  # ToolSelection-like
                msg_data["tool_calls"] = [
                    {
                        "id": tc.tool_id,
                        "type": "function",
                        "function": {
                            "name": tc.tool_name,
                            "arguments": json.dumps(tc.tool_kwargs),
                        },
                    }
                    for tc in tool_calls
                ]
            elif hasattr(tool_calls[0], "id") and hasattr(tool_calls[0], "function"):
                # ChatCompletionMessageFunctionToolCall object
                msg_data["tool_calls"] = []
                for tc in tool_calls:
                    fn = tc.function
                    args = fn.arguments if hasattr(fn, "arguments") else ""
                    if isinstance(args, dict):
                        args = json.dumps(args)
                    msg_data["tool_calls"].append(
                        {
                            "id": tc.id,
                            "type": getattr(tc, "type", "function"),
                            "function": {
                                "name": fn.name if hasattr(fn, "name") else "",
                                "arguments": args,
                            },
                        }
                    )
            else:  # dict format (e.g. from additional_kwargs)
                msg_data["tool_calls"] = []
                for tc in tool_calls:
                    fn = tc.get("function", {}) or {}
                    args = fn.get("arguments", "")
                    if isinstance(args, dict):
                        args = json.dumps(args)
                    msg_data["tool_calls"].append(
                        {
                            "id": tc.get("id", ""),
                            "type": "function",
                            "function": {"name": fn.get("name", ""), "arguments": args},
                        }
                    )
        return msg_data

    if role == "tool":
        additional = getattr(msg, "additional_kwargs", {}) or {}
        return {
            "role": "tool",
            "tool_call_id": additional.get("tool_call_id", ""),
            "name": additional.get("name", ""),
            "content": content,
        }

    return None  # skip system or unknown


def _build_user_message(messages: list[ChatMessage]) -> str:
    """Extract the last user message from the OpenAI-format messages list."""
    for msg in reversed(messages):
        if msg.role == "user":
            return msg.content
    raise ValueError("No user message found in messages list")


def _make_completion_id() -> str:
    return f"chatcmpl-{uuid.uuid4().hex[:12]}"


@app.post(
    "/chat/completions",
    response_model=ChatCompletionResponse,
    summary="Create chat completion",
    description="Creates a model response for the given chat conversation. When `stream=false`, returns a complete `chat.completion` JSON object. When `stream=true`, returns Server-Sent Events with `chat.completion.chunk` deltas.",
    tags=["Chat"],
)
async def chat_completions(request: ChatCompletionRequest):
    global get_agent

    if get_agent is None:
        raise HTTPException(status_code=503, detail="Agent not initialized")

    user_message = _build_user_message(request.messages)
    model_id = request.model or getenv("MODEL_ID", "model")

    if request.stream:
        return await _handle_stream(user_message, model_id)
    else:
        return await _handle_chat(user_message, model_id)


async def _handle_chat(user_message: str, model_id: str):
    """Handle non-streaming chat completion."""
    global get_agent

    try:
        agent = get_agent()
        messages = [{"role": "user", "content": user_message}]

        result = await agent.run(input=messages)

        # Extract the final assistant message content
        assistant_content = ""
        context_messages = []

        if result and "messages" in result and len(result["messages"]) > 0:
            for message in result["messages"]:
                if getattr(message, "role", None) == "system":
                    continue
                item = _message_to_response_dict(message)
                if item is not None:
                    context_messages.append(item)

            # Final assistant content is the last assistant message with content
            for item in reversed(context_messages):
                if item["role"] == "assistant" and item.get("content"):
                    assistant_content = item["content"]
                    break

        return {
            "id": _make_completion_id(),
            "object": "chat.completion",
            "created": int(time.time()),
            "model": model_id,
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": assistant_content,
                    },
                    "finish_reason": "stop",
                }
            ],
            "context": context_messages,
            "usage": None,
        }

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Error processing request: {str(e)}"
        )


async def _handle_stream(user_message: str, model_id: str):
    """Handle streaming chat completion with OpenAI-compatible SSE chunks."""
    global get_agent

    completion_id = _make_completion_id()
    created = int(time.time())

    async def event_generator():
        try:
            agent = get_agent()
            messages = [{"role": "user", "content": user_message}]

            handler = agent.run(input=messages)

            async for event in handler.stream_events():
                if isinstance(event, ToolCallEvent):
                    tool_calls_delta = [
                        {
                            "index": i,
                            "id": getattr(tc, "tool_id", ""),
                            "type": "function",
                            "function": {
                                "name": tc.tool_name,
                                "arguments": json.dumps(tc.tool_kwargs),
                            },
                        }
                        for i, tc in enumerate(event.tool_calls)
                    ]
                    data = {
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model_id,
                        "choices": [
                            {
                                "index": 0,
                                "delta": {
                                    "role": "assistant",
                                    "tool_calls": tool_calls_delta,
                                },
                                "finish_reason": None,
                            }
                        ],
                    }
                    yield f"data: {json.dumps(data)}\n\n"

                elif isinstance(event, InputEvent):
                    if event.input:
                        last_msg = event.input[-1]
                        if getattr(last_msg, "role", None) == "tool":
                            additional = (
                                getattr(last_msg, "additional_kwargs", {}) or {}
                            )
                            data = {
                                "id": completion_id,
                                "object": "chat.completion.chunk",
                                "created": created,
                                "model": model_id,
                                "choices": [
                                    {
                                        "index": 0,
                                        "delta": {
                                            "role": "tool",
                                            "content": _get_message_content(last_msg),
                                            "name": additional.get("name", ""),
                                        },
                                        "finish_reason": None,
                                    }
                                ],
                            }
                            yield f"data: {json.dumps(data)}\n\n"

            result = await handler
            # Extract final answer from the result
            if result and "response" in result:
                content = _get_message_content(result["response"].message)
                if content:
                    data = {
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model_id,
                        "choices": [
                            {
                                "index": 0,
                                "delta": {"content": content},
                                "finish_reason": None,
                            }
                        ],
                    }
                    yield f"data: {json.dumps(data)}\n\n"

            # Send final chunk with finish_reason
            final_data = {
                "id": completion_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model_id,
                "choices": [
                    {
                        "index": 0,
                        "delta": {},
                        "finish_reason": "stop",
                    }
                ],
            }
            yield f"data: {json.dumps(final_data)}\n\n"
            yield "data: [DONE]\n\n"

        except Exception:
            logger.exception("Error in stream event_generator")
            error_data = {
                "error": {
                    "message": "Internal server error",
                    "type": "server_error",
                }
            }
            yield f"data: {json.dumps(error_data)}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.get(
    "/health", response_model=HealthResponse, summary="Health check", tags=["Health"]
)
async def health():
    initialized = get_agent is not None
    body = {
        "status": "healthy" if initialized else "not_ready",
        "agent_initialized": initialized,
    }
    if not initialized:
        return JSONResponse(status_code=503, content=body)
    return body


# ── Playground API aliases (so the same index.html works in both modes) ───────
@app.get("/api/health", response_model=HealthResponse, include_in_schema=False)
async def api_health():
    return await health()


@app.post("/api/chat", include_in_schema=False)
async def api_chat(request: ChatCompletionRequest):
    return await chat_completions(request)


# ── Playground UI ────────────────────────────────────────────────────────────
_BASE_DIR = Path(__file__).resolve().parent
_PLAYGROUND_HTML = _BASE_DIR / "playground" / "templates" / "index.html"
# In Docker the images are copied to /opt/app-root/src/images; locally they live at the repo root
_IMAGES_DIR = _BASE_DIR / "images"
if not _IMAGES_DIR.is_dir():
    _IMAGES_DIR = _BASE_DIR.parent.parent.parent / "images"


@app.get("/", response_class=HTMLResponse, include_in_schema=False)
async def playground():
    """Serve the playground chat UI."""
    return FileResponse(_PLAYGROUND_HTML)


@app.get("/images/{filename:path}", include_in_schema=False)
async def serve_image(filename: str):
    """Serve images from the project-level images directory."""
    base = _IMAGES_DIR.resolve()
    file_path = (base / filename).resolve()
    try:
        file_path.relative_to(base)
    except ValueError:
        raise HTTPException(status_code=404, detail="Image not found")
    if not file_path.is_file():
        raise HTTPException(status_code=404, detail="Image not found")
    return FileResponse(file_path)


if __name__ == "__main__":
    import uvicorn

    port = int(getenv("PORT", 8000))
    uvicorn.run(app, host="0.0.0.0", port=port)