main.py

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

from crewai import LLM
from crewai_web_search.crew import AssistanceAgents
from crewai_web_search.tracing import enable_tracing
from fastapi import FastAPI, HTTPException
from fastapi.responses import (
    FileResponse,
    HTMLResponse,
    JSONResponse,
    StreamingResponse,
)
from pydantic import BaseModel, Field

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=["What is the best cluster hosting service?"],
    )


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 LLM instance
llm = None

# Patterns that indicate CrewAI internal scaffolding in the output
_REACT_NOISE = re.compile(
    r"(^|\n)\s*(Thought:\s*|Action:\s*|Action Input:\s*|Observation:\s*|Final Answer:\s*).*",
    re.DOTALL,
)
_CREWAI_PROMPT_MARKER = "\n\n\nYou ONLY have access to"


def _clean_content(text: str) -> str:
    """Strip CrewAI internal ReAct scaffolding and prompt noise from output."""
    # Strip appended retry instructions
    idx = text.find(_CREWAI_PROMPT_MARKER)
    if idx != -1:
        text = text[:idx]

    # Strip ReAct format artifacts (Thought:/Action:/Final Answer: prefixes)
    text = _REACT_NOISE.sub("", text)

    return text.strip()


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]}"


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize the CrewAI LLM on startup."""
    global llm
    enable_tracing()

    base_url = getenv("BASE_URL")
    model_id = getenv("MODEL_ID")
    api_key = getenv("API_KEY", "no-key")

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

    llm = LLM(
        model=f"openai/{model_id}",
        base_url=base_url,
        api_key=api_key,
        temperature=0.7,
    )

    yield

    llm = None


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


@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 llm

    if llm 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 llm

    try:
        inputs = {
            "user_prompt": user_message,
            "custom_instruction": "",
        }

        crew = AssistanceAgents(llm=llm).crew()
        result = await asyncio.to_thread(crew.kickoff, inputs=inputs)

        assistant_content = _clean_content(str(result))

        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",
                }
            ],
            "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 llm

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

    async def event_generator():
        try:
            inputs = {
                "user_prompt": user_message,
                "custom_instruction": "",
            }

            crew = AssistanceAgents(llm=llm, stream=True).crew()
            streaming_output = await crew.kickoff_async(inputs=inputs)

            # Buffer tokens until we see "Final Answer:" — everything before
            # that is internal ReAct reasoning (Thought/Action/Observation).
            buffer = ""
            emitting = False

            async for chunk in streaming_output:
                if chunk.chunk_type.value != "text" or not chunk.content:
                    continue

                if emitting:
                    data = {
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model_id,
                        "choices": [
                            {
                                "index": 0,
                                "delta": {"content": chunk.content},
                                "finish_reason": None,
                            }
                        ],
                    }
                    yield f"data: {json.dumps(data)}\n\n"
                else:
                    buffer += chunk.content
                    marker = "Final Answer:"
                    idx = buffer.find(marker)
                    if idx != -1:
                        emitting = True
                        remainder = buffer[idx + len(marker) :]
                        if remainder.strip():
                            data = {
                                "id": completion_id,
                                "object": "chat.completion.chunk",
                                "created": created,
                                "model": model_id,
                                "choices": [
                                    {
                                        "index": 0,
                                        "delta": {"content": remainder.lstrip()},
                                        "finish_reason": None,
                                    }
                                ],
                            }
                            yield f"data: {json.dumps(data)}\n\n"

            # If no "Final Answer:" was found, send the cleaned full buffer
            if not emitting and buffer.strip():
                cleaned = _clean_content(buffer)
                if cleaned:
                    data = {
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model_id,
                        "choices": [
                            {
                                "index": 0,
                                "delta": {"content": cleaned},
                                "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 = llm 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)