api.py

import asyncio
import os
import re
import uuid
from contextlib import asynccontextmanager
from urllib.parse import urlparse

from dotenv import load_dotenv

load_dotenv()

import jsonref
import requests
import ruamel.yaml
from elasticsearch import Elasticsearch
from fastapi import BackgroundTasks, FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama, OllamaEmbeddings
from pydantic import BaseModel

from api_swagger_to_es import embed_and_store_apis, extract_api_docs, merge_allof, sanitize_text
from db import append_turn, get_messages, init_db, list_sessions
from db import remove_session as db_remove_session
from db import delete_setting, get_setting, set_setting

ES_URL = os.environ.get("ES_URL", "http://localhost:9200")
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
OLLAMA_CHAT_MODEL = os.environ.get("OLLAMA_CHAT_MODEL", "gemma4:latest")
OLLAMA_EMBED_MODEL = "nomic-embed-text"


@asynccontextmanager
async def lifespan(app: FastAPI):
    init_db()
    yield


app = FastAPI(title="API Documentation Assistant", lifespan=lifespan)
_CORS_ORIGINS = os.environ.get(
    "CORS_ORIGINS",
    "http://localhost:5173,http://localhost:3000"
).split(",")

app.add_middleware(
    CORSMiddleware,
    allow_origins=_CORS_ORIGINS,
    allow_origin_regex=r"https?://.*",  # allow any origin in production; restrict via env in prod
    allow_methods=["*"],
    allow_headers=["*"],
)

# In-memory ingestion status keyed by index name
ingest_status: dict[str, dict] = {}


# ── Models ────────────────────────────────────────────────────────────────────

class IngestRequest(BaseModel):
    url: str
    index_name: str | None = None


class QueryRequest(BaseModel):
    query: str
    index_name: str
    k: int = 3


class QueryResponse(BaseModel):
    query: str
    answer: str
    index_name: str


class ChatRequest(BaseModel):
    message: str
    index_name: str
    session_id: str | None = None
    k: int = 3


class ChatMessage(BaseModel):
    role: str  # "user" | "assistant"
    content: str


class ChatResponse(BaseModel):
    session_id: str
    answer: str
    history: list[ChatMessage]


class SessionInfo(BaseModel):
    id: str
    index_name: str
    title: str | None
    created_at: str


# ── Helpers ───────────────────────────────────────────────────────────────────

def url_to_index_name(url: str) -> str:
    parsed = urlparse(url)
    segments = [s for s in parsed.path.split("/") if s]
    raw = segments[-1] if segments else parsed.netloc
    raw = re.sub(r'\.[^.]+$', '', raw)  # strip extension
    raw = re.sub(r'[^a-zA-Z0-9_-]', '_', raw)  # sanitize
    return raw.strip("_")[:64] or "api_docs"


def _do_ingestion(url: str, index_name: str) -> int:
    ingest_status[index_name]["message"] = "Fetching spec..."
    resp = requests.get(url, timeout=60)
    resp.raise_for_status()

    ingest_status[index_name]["message"] = "Parsing OpenAPI spec..."
    yaml_parser = ruamel.yaml.YAML(typ="safe")
    swagger_json = yaml_parser.load(resp.text)
    swagger_json = jsonref.replace_refs(swagger_json)
    swagger_json = merge_allof(swagger_json)

    ingest_status[index_name]["message"] = "Extracting API endpoints..."
    apis = extract_api_docs(swagger_json)

    ingest_status[index_name]["message"] = f"Embedding and storing {len(apis)} endpoints..."
    embed_and_store_apis(apis, ES_URL, index_name)
    return len(apis)


async def run_ingestion(url: str, index_name: str) -> None:
    try:
        ingest_status[index_name] = {"status": "ingesting", "message": "Starting...", "doc_count": None}
        doc_count = await asyncio.to_thread(_do_ingestion, url, index_name)
        ingest_status[index_name] = {
            "status": "done",
            "message": f"Ingested {doc_count} endpoints.",
            "doc_count": doc_count,
        }
    except Exception as exc:
        ingest_status[index_name] = {"status": "error", "message": str(exc), "doc_count": None}


def _vector_search(query: str, index_name: str, k: int, weights: []) -> list[dict]:
    embedder = OllamaEmbeddings(
        model=OLLAMA_EMBED_MODEL,
        base_url=OLLAMA_BASE_URL,
    )
    should_clauses = []

    for keyword, weight in weights:
        should_clauses.append({
            "multi_match": {
                "query": keyword,
                "boost": weight,
                "fields": [
                    "summary^5",
                    "text^3"
                ],
            }
        })

    query_vector = embedder.embed_query(query)
    response = Elasticsearch(ES_URL).search(
        index=index_name,
        knn={
            "field": "vector",
            "query_vector": query_vector,
            "k": k,
            "num_candidates": max(k * 10, 50),
        },
        query={
            "bool": {
                "should": should_clauses
            }
        },
        size=k,
        source=["metadata"],
    )
    return [hit["_source"] for hit in response["hits"]["hits"]]


def _rag_query(query: str, index_name: str, k: int) -> str:
    results = _vector_search(query, index_name, k)

    if not results:
        return "No relevant API documentation found for your query."

    context = "\n\n".join([
        f"path: {doc['metadata']['path']}\n\n"
        f"requestBody: {doc['metadata']['requestBody']}\n\n"
        f"responses: {doc['metadata']['responses']}"
        for doc in results
    ])

    prompt_template = ChatPromptTemplate.from_messages([
        (
            "system",
            "You are an API assistant. Given the following API documentation context, "
            "answer the user's query as helpfully as possible.",
        ),
        (
            "user",
            "Relevant API Information for your context:\n{context}\n\n"
            "Query user has provided: {query}\n\nAnswer:",
        ),
        (
            "system",
            "Your response should be in markdown language and should include code snippets if relevant. "
            "Always refer to the provided context when answering. "
            "Include URL path, HTTP method, parameters, and response structure in your answer if relevant.",
        ),
    ])

    prompt = prompt_template.format(context=context, query=query)
    llm = ChatOllama(model=OLLAMA_CHAT_MODEL, temperature=0, base_url=OLLAMA_BASE_URL)
    output_parser = StrOutputParser()
    answer = output_parser.parse(llm.invoke(prompt))
    return answer.content


def _restructure_query(query: str, llm: ChatOllama) -> str:
    lc_messages = [
        SystemMessage(content=(
            "You are helpful search query writer. Given a user's question, restructure this query using words from original query. Also make sure you are returning only new restructered query I don't want anything other than query to explaination needed."
        )),
        HumanMessage(
            content=f"User's original query: {query}\n\nWrite a restructured query that is optimized for retrieving relevant API documentation from a vector search engine. Use terms from the original query but feel free to rephrase for clarity and relevance. Return only the restructured query without any additional explanation."),
    ]
    return llm.invoke(lc_messages).content


def _chat_rag(session_id: str, message: str, index_name: str, k: int) -> tuple[str, list[dict]]:
    history = get_messages(session_id)

    queries = ' '.join((turn["content"] for turn in history if turn["role"] == "user"))
    query, weights = sanitize_text(f'{queries}. {message}')

    results = _vector_search(query, index_name, k, weights)

    context = "\n\n".join([
        f"path: {doc['metadata']['path']}\n\n"
        f"requestBody: {doc['metadata']['requestBody']}\n\n"
        f"responses: {doc['metadata']['responses']}"
        for doc in results
    ]) if results else "No relevant API documentation found."

    context += f"API Information: {results[0]['metadata'].get('infoDescription')}\n" if results else ""

    lc_messages = [
        SystemMessage(content=(
            "You are an API assistant engaged in a multi-turn conversation. "
            "Use the provided API documentation context and the conversation history to give accurate and to the point, helpful answers. "
            "Format responses in markdown and include code snippets where relevant."
        ))
    ]
    llm = ChatOllama(model=OLLAMA_CHAT_MODEL, temperature=0, base_url=OLLAMA_BASE_URL)

    context_message = _restructure_query(f'{queries}. {message}', llm)
    lc_messages.append(HumanMessage(content=(
        f"Relevant API Documentation:\n{context}\n\n"
        f"Current Question: {message}"
        f"Previous Question: {context_message}"
    )))

    answer = llm.invoke(lc_messages).content

    # Only set title on the first turn
    title = message[:80] if not history else None
    append_turn(session_id, index_name, title, message, answer)

    updated = history + [{"role": "user", "content": message}, {"role": "assistant", "content": answer}]
    return answer, updated


# ── Endpoints ─────────────────────────────────────────────────────────────────

@app.post("/ingest", status_code=202)
async def ingest_swagger(request: IngestRequest, background_tasks: BackgroundTasks):
    index_name = request.index_name or url_to_index_name(request.url)
    if ingest_status.get(index_name, {}).get("status") == "ingesting":
        raise HTTPException(status_code=409, detail=f"Ingestion already running for '{index_name}'")
    ingest_status[index_name] = {"status": "ingesting", "message": "Queued...", "doc_count": None}
    background_tasks.add_task(run_ingestion, request.url, index_name)
    return {"index_name": index_name, "status": "ingesting"}


@app.get("/ingest/status/{index_name}")
async def get_ingest_status(index_name: str):
    record = ingest_status.get(index_name)
    if record is None:
        raise HTTPException(status_code=404, detail="No ingestion record for this index")
    return {"index_name": index_name, **record}


@app.get("/indexes")
async def list_indexes():
    client = Elasticsearch(ES_URL)
    try:
        raw = client.cat.indices(format="json")
        names = [idx["index"] for idx in raw if not idx["index"].startswith(".")]
        default_index = await asyncio.to_thread(get_setting, "default_index")
        return [
            {
                "uid": name,
                "status": ingest_status.get(name, {}).get("status", "ready"),
                "doc_count": ingest_status.get(name, {}).get("doc_count"),
                "is_default": name == default_index,
            }
            for name in names
        ]
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.put("/indexes/{index_name}/default", status_code=204)
async def set_default_index(index_name: str):
    await asyncio.to_thread(set_setting, "default_index", index_name)


@app.delete("/indexes/{index_name}/default", status_code=204)
async def unset_default_index(index_name: str):
    default = await asyncio.to_thread(get_setting, "default_index")
    if default == index_name:
        await asyncio.to_thread(delete_setting, "default_index")


@app.delete("/indexes/{index_name}", status_code=204)
async def delete_index(index_name: str):
    try:
        client = Elasticsearch(ES_URL)
        client.indices.delete(index=index_name)
        ingest_status.pop(index_name, None)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.post("/query", response_model=QueryResponse)
async def query_api_docs(request: QueryRequest):
    try:
        answer = await asyncio.to_thread(_rag_query, request.query, request.index_name, request.k)
        return QueryResponse(query=request.query, answer=answer, index_name=request.index_name)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.get("/chat/sessions", response_model=list[SessionInfo])
async def get_sessions():
    return await asyncio.to_thread(list_sessions)


@app.get("/chat/{session_id}", response_model=ChatResponse)
async def get_session(session_id: str):
    history = await asyncio.to_thread(get_messages, session_id)
    if not history:
        raise HTTPException(status_code=404, detail="Session not found")
    return ChatResponse(
        session_id=session_id,
        answer=history[-1]["content"] if history else "",
        history=[ChatMessage(role=t["role"], content=t["content"]) for t in history],
    )


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    session_id = request.session_id or str(uuid.uuid4())
    try:
        answer, history = await asyncio.to_thread(
            _chat_rag, session_id, request.message, request.index_name, request.k
        )
        return ChatResponse(
            session_id=session_id,
            answer=answer,
            history=[ChatMessage(role=t["role"], content=t["content"]) for t in history],
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.delete("/chat/{session_id}", status_code=204)
async def delete_session(session_id: str):
    await asyncio.to_thread(db_remove_session, session_id)


@app.get("/health")
async def health():
    return {"status": "ok"}