shira-hermes/api/routes/kb_public.py

"""Public KB endpoints — direct search without the LLM layer.

These endpoints serve the EspoCRM KnowledgeBase extension UI. Auth is the
same X-Api-Key gate used by SmartAssistant (so the EspoCRM backend proxies
calls with the shared key).
"""

from __future__ import annotations

import logging
import os
from typing import Literal

from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from api.services.kb import search as kb_search
from api.services.kb import ingest as kb_ingest
from api.services.kb import s3 as kb_s3
from api.services.kb import topics as kb_topics

logger = logging.getLogger("shira.kb.public")

router = APIRouter(prefix="/kb", tags=["kb"])


def _verify_auth(request: Request) -> None:
    expected = os.environ.get("API_KEY", "")
    if not expected:
        raise HTTPException(status_code=503, detail="API_KEY not configured")
    provided = request.headers.get("X-Api-Key", "")
    if provided != expected:
        raise HTTPException(status_code=401, detail="Unauthorized")


class SearchRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=500)
    kind: Literal["any", "law", "regulation", "circular", "caselaw", "tool"] = "any"
    top_k: int = Field(8, ge=1, le=15)
    # Multi-query expansion: ask the LLM for alternative phrasings, retrieve
    # for each, merge, then rerank against the original query. Recovers
    # documents that don't share the user's exact wording.
    expand: bool = True
    # Domain scope. None means "lowest active topic" (back-compat — equivalent
    # to topic_id=1 'national-insurance' until additional topics are added).
    topic_id: int | None = None
    # Optional sub-topic filter (Phase 6 / kb_label).
    label_id: int | None = None


@router.get("/topics")
async def list_topics(request: Request):
    """Return the active KB topics for the topic <select> in the UI."""
    _verify_auth(request)
    topics = await kb_topics.list_topics(active_only=True)
    return {"topics": topics, "count": len(topics)}


@router.post("/search")
async def search(body: SearchRequest, request: Request):
    """Hybrid search + rerank. Returns ranked chunks with citations."""
    _verify_auth(request)
    try:
        topic_id = await kb_topics.resolve_topic_id(body.topic_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    topic = await kb_topics.get_topic(topic_id)
    hits = await kb_search.search(
        query=body.query,
        kind=body.kind,
        top_k=body.top_k,
        expand=body.expand,
        topic_id=topic_id,
        topic_name=(topic or {}).get("name"),
        label_id=body.label_id,
    )
    # Dates come back as datetime.date; serialize as ISO strings for the UI.
    serialized = []
    for h in hits:
        item = dict(h)
        for k in ("published_at", "effective_at"):
            if item.get(k) is not None:
                item[k] = item[k].isoformat()
        serialized.append(item)
    return {"query": body.query, "kind": body.kind, "count": len(serialized), "hits": serialized}


@router.get("/sources")
async def list_sources(request: Request, topic_id: int | None = None):
    """Browse mode: list all active sources with basic metadata."""
    _verify_auth(request)
    try:
        topic_id = await kb_topics.resolve_topic_id(topic_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    from api.services.kb.db import get_pool
    pool = await get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            """
            SELECT s.id, s.kind, s.title, s.identifier, s.published_at,
                   s.effective_at, s.source_url, s.topic_id,
                   (SELECT COUNT(*) FROM kb_chunk c WHERE c.source_id = s.id) AS chunk_count
            FROM kb_source s
            WHERE s.superseded_by IS NULL
              AND s.topic_id = $1
            ORDER BY s.kind, s.published_at DESC NULLS LAST, s.title
            """,
            topic_id,
        )
    sources = []
    for r in rows:
        d = dict(r)
        for k in ("published_at", "effective_at"):
            if d.get(k) is not None:
                d[k] = d[k].isoformat()
        sources.append(d)
    return {"sources": sources, "count": len(sources)}


@router.get("/source/{source_id}/chunks")
async def source_chunks(
    source_id: int,
    request: Request,
    around: int | None = None,
    ctx: int = 2,
):
    """Browse mode: return chunks of a single source, ordered.

    By default returns all chunks. With `around=<chunk_index>&ctx=<k>` it
    returns only chunks with chunk_index in [around-k, around+k] inclusive,
    so the search UI can show a matched chunk together with its surrounding
    sections for context (instead of duplicating the same chunk in both
    panes when the source has no PDF).
    """
    _verify_auth(request)
    from api.services.kb.db import get_pool
    pool = await get_pool()
    async with pool.acquire() as conn:
        source = await conn.fetchrow(
            "SELECT id, kind, title, identifier, published_at, source_url, original_path "
            "FROM kb_source WHERE id = $1",
            source_id,
        )
        if not source:
            raise HTTPException(status_code=404, detail="source not found")
        if around is not None:
            ctx = max(0, min(ctx, 10))
            chunks = await conn.fetch(
                """
                SELECT chunk_index, heading_path, section_ref, content, token_count
                FROM kb_chunk
                WHERE source_id = $1
                  AND chunk_index BETWEEN $2 AND $3
                ORDER BY chunk_index
                """,
                source_id, around - ctx, around + ctx,
            )
        else:
            chunks = await conn.fetch(
                """
                SELECT chunk_index, heading_path, section_ref, content, token_count
                FROM kb_chunk WHERE source_id = $1 ORDER BY chunk_index
                """,
                source_id,
            )
    src = dict(source)
    if src.get("published_at") is not None:
        src["published_at"] = src["published_at"].isoformat()
    return {"source": src, "chunks": [dict(c) for c in chunks]}


def _parse_s3_uri(uri: str) -> tuple[str, str] | None:
    """Split 's3://<bucket>/<key...>' into (bucket, key). None if not s3."""
    if not uri or not uri.startswith("s3://"):
        return None
    rest = uri[len("s3://"):]
    if "/" not in rest:
        return None
    bucket, _, key = rest.partition("/")
    return bucket, key


@router.get("/source/{source_id}/pdf")
async def source_pdf(source_id: int, request: Request):
    """Stream the original PDF binary of a source straight from MinIO.

    Used by the EspoCRM KnowledgeBase browse tab to render the PDF inline
    without exposing shira-hermes or MinIO to the end-user's browser.
    The EspoCRM EntryPoint is the only caller and proxies this response.

    Returns 404 if the source doesn't exist. Returns 409 if the source
    has no stored PDF (e.g. the law, which was ingested from plain text
    and has an `s3://.../file.txt` original_path).
    """
    _verify_auth(request)
    from api.services.kb.db import get_pool

    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT title, original_path FROM kb_source WHERE id = $1",
            source_id,
        )
    if not row:
        raise HTTPException(status_code=404, detail="source not found")

    original_path = (row["original_path"] or "").strip()
    if not original_path.lower().endswith(".pdf"):
        raise HTTPException(
            status_code=409,
            detail="source has no PDF original (stored as non-PDF text)",
        )

    parsed = _parse_s3_uri(original_path)
    if not parsed:
        raise HTTPException(status_code=500, detail="original_path is not an s3 URI")
    bucket, key = parsed

    # The KB bucket is the only one we serve from here; fetching from an
    # unexpected bucket would be a bug or a tampered DB row.
    if bucket != kb_s3._bucket():
        raise HTTPException(status_code=500, detail="source bucket mismatch")

    try:
        data = kb_s3.fetch(key)
    except Exception as e:
        logger.exception("[kb.pdf] fetch failed for source %s key %s", source_id, key)
        raise HTTPException(status_code=502, detail=f"storage fetch failed: {e}")

    # Chunk the response so we don't keep the full 2 MB buffer per request
    # beyond this function's scope.
    def _chunks(blob: bytes, size: int = 64 * 1024):
        for i in range(0, len(blob), size):
            yield blob[i : i + size]

    filename = (row["title"] or f"source-{source_id}").strip() + ".pdf"
    # Content-Disposition filename must be ASCII-safe on the fallback param,
    # plus RFC 5987 UTF-8 for the real name.
    from urllib.parse import quote as _q
    ascii_fallback = f"source-{source_id}.pdf"
    disposition = (
        f"inline; filename=\"{ascii_fallback}\"; filename*=UTF-8''{_q(filename)}"
    )
    return StreamingResponse(
        _chunks(data),
        media_type="application/pdf",
        headers={
            "Content-Disposition": disposition,
            "Content-Length": str(len(data)),
            "X-Content-Type-Options": "nosniff",
            "Cache-Control": "private, max-age=300",
        },
    )


class AskRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=2000)
    conversation_id: str | None = None
    conversation_history: list[dict] | None = None
    # Domain scope. None means "lowest active topic" (back-compat).
    topic_id: int | None = None


async def _build_ask_runner_context(topic_id: int | None) -> tuple:
    """Common construction shared by /kb/ask and /kb/ask/stream.

    Resolves the topic and pulls its system_prompt_addendum from the DB,
    so adding a new domain (Phase 4) automatically updates Shira's
    instructions without a code change. Returns the resolved topic too
    so the caller can pin tools (search_legal_kb) to it.
    """
    from api.services.agent_runner import AgentRunner
    from api.services.prompt_builder import build_office_prompt
    from api.services.skills import list_skills

    resolved_id = await kb_topics.resolve_topic_id(topic_id)
    topic = await kb_topics.get_topic(resolved_id) or {}

    runner = AgentRunner(
        ai_gateway_url=os.environ.get("AI_GATEWAY_URL", "http://localhost:3000"),
        ai_gateway_api_key=os.environ.get("AI_GATEWAY_API_KEY", ""),
        model=os.environ.get("CLAUDE_MODEL", "sonnet"),
        max_tokens=int(os.environ.get("MAX_OUTPUT_TOKENS", "4096")),
        max_iterations=int(os.environ.get("MAX_AGENT_ITERATIONS", "10")),
    )
    context = {"currentUser": {"id": "kb-ui", "name": "KB UI"}}
    system_prompt = build_office_prompt(
        context,
        user_profile="",
        available_skills=list_skills(),
    )
    addendum = (topic.get("system_prompt_addendum") or "").strip()
    if addendum:
        system_prompt += "\n\n" + addendum
    return runner, context, system_prompt, topic


def _build_ask_messages(body: AskRequest) -> list[dict]:
    messages: list[dict] = []
    for m in (body.conversation_history or []):
        role = m.get("role", "user")
        messages.append({
            "role": "assistant" if role == "assistant" else "user",
            "content": m.get("content", ""),
        })
    messages.append({"role": "user", "content": body.message})
    return messages


def _filter_pdf_sources(sources_used: list[dict]) -> list[dict]:
    return [
        s for s in sources_used
        if (s.get("original_path") or "").lower().endswith(".pdf")
    ]


@router.post("/ask")
async def ask(body: AskRequest, request: Request):
    """Full agent loop: the user's question goes to shira, she calls
    search_legal_kb as needed and returns a summarized answer.

    Thin wrapper around the existing smart-assistant chat endpoint so the
    KnowledgeBase UI can expose an ask mode without rebuilding prompt
    plumbing.
    """
    _verify_auth(request)
    try:
        runner, context, system_prompt, topic = await _build_ask_runner_context(body.topic_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    messages = _build_ask_messages(body)
    sources_used: list[dict] = []
    try:
        text = await runner.run(
            system_prompt=system_prompt,
            messages=messages,
            context=context,
            espocrm_url=os.environ.get("ESPOCRM_URL", ""),
            espocrm_api_key=os.environ.get("ESPOCRM_API_KEY", ""),
            allowed_toolsets=["legal"],
            kb_sources_used=sources_used,
            kb_topic_id=topic.get("id"),
            kb_topic_name=topic.get("name"),
        )
    except Exception as e:
        logger.exception("[kb.ask] error")
        raise HTTPException(status_code=500, detail=str(e))
    return {
        "text": text,
        "conversationId": body.conversation_id,
        "sources": _filter_pdf_sources(sources_used),
    }


@router.post("/ask/stream")
async def ask_stream(body: AskRequest, request: Request):
    """Streaming variant of /kb/ask. Returns text/event-stream with one
    SSE event per agent step (thinking, tool_start, tool_done) and a
    final {type:"answer",text,sources} event. The KnowledgeBase UI
    consumes this via EventSource (proxied through an EspoCRM EntryPoint
    so the auth token stays on the server).
    """
    import asyncio
    import json as _json
    from fastapi.responses import StreamingResponse

    _verify_auth(request)
    try:
        runner, context, system_prompt, topic = await _build_ask_runner_context(body.topic_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    messages = _build_ask_messages(body)
    sources_used: list[dict] = []

    queue: asyncio.Queue = asyncio.Queue()

    def _on_event(ev: dict) -> None:
        # AgentRunner calls this synchronously between awaits; an unbounded
        # queue lets us put_nowait without blocking the loop.
        try:
            queue.put_nowait(ev)
        except Exception as e:
            logger.warning("[kb.ask/stream] queue put failed: %s", e)

    async def _runner_task():
        try:
            text = await runner.run(
                system_prompt=system_prompt,
                messages=messages,
                context=context,
                espocrm_url=os.environ.get("ESPOCRM_URL", ""),
                espocrm_api_key=os.environ.get("ESPOCRM_API_KEY", ""),
                allowed_toolsets=["legal"],
                kb_sources_used=sources_used,
                kb_topic_id=topic.get("id"),
                kb_topic_name=topic.get("name"),
                on_event=_on_event,
            )
            await queue.put({
                "type": "answer",
                "text": text,
                "sources": _filter_pdf_sources(sources_used),
                "conversationId": body.conversation_id,
            })
        except Exception as e:
            logger.exception("[kb.ask/stream] runner error")
            await queue.put({"type": "error", "message": str(e)})
        finally:
            await queue.put(None)  # sentinel

    asyncio.create_task(_runner_task())

    async def _gen():
        # Heartbeat every 15s keeps the connection open through any
        # timeout-happy proxies; comment lines (`: ...`) are ignored by
        # EventSource but keep TCP/HTTP alive.
        while True:
            try:
                ev = await asyncio.wait_for(queue.get(), timeout=15.0)
            except asyncio.TimeoutError:
                yield ": keep-alive\n\n"
                continue
            if ev is None:
                break
            yield "data: " + _json.dumps(ev, ensure_ascii=False) + "\n\n"

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