This repository has been archived on 2026-07-19. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
shira-hermes/api/routes/kb_public.py
T
chaim 4a3a8945af feat(kb): track page_number per chunk + expose sources_used from /kb/ask
To let the browse/ask UIs jump straight to the right PDF page:

- ingest.py: embed \x00KB_PAGE:<N>\x00 sentinels at the start of each
  PDF page during _parse_pdf.
- chunker.py: _strip_page_markers removes the sentinels and records the
  first page seen per chunk as page_number. Chunk dataclass + _as_dicts
  propagate page_number through both chunk_statute and chunk_circular.
- ingest.py INSERT: new kb_chunk.page_number column (DB migration
  applied manually).
- search.py SQL: SELECT c.page_number + s.id + s.original_path so the
  rerank results carry them forward.
- legal_kb_tools.py: register_legal_kb_tools(sources_used=list) — the
  tool appends {source_id, title, kind, identifier, heading_path,
  section_ref, page_number, original_path} for each hit, deduped by
  (source_id, page_number).
- agent_runner.py: new kb_sources_used kwarg threads the list through
  to the tool registration.
- kb_public.py /kb/ask: collects sources_used during the run, filters
  to PDF-only sources, and returns them alongside text.

Chunks stored before this change have page_number = NULL; PDFs need to
be re-ingested to populate them. The TXT law source is unaffected.

Refs Task Master #2

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-24 12:20:52 +00:00

285 lines
10 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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
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"] = "any"
top_k: int = Field(8, ge=1, le=15)
@router.post("/search")
async def search(body: SearchRequest, request: Request):
"""Hybrid search + rerank. Returns ranked chunks with citations."""
_verify_auth(request)
hits = await kb_search.search(
query=body.query,
kind=body.kind,
top_k=body.top_k,
)
# 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):
"""Browse mode: list all active sources with basic metadata."""
_verify_auth(request)
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,
(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
ORDER BY s.kind, s.published_at DESC NULLS LAST, s.title
"""
)
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):
"""Browse mode: return all chunks of a single source, ordered."""
_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")
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
@router.post("/ask")
async def ask(body: AskRequest, request: Request):
"""Full agent loop: the user's question goes to shira, she calls
search_insurance_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)
from api.services.agent_runner import AgentRunner
from api.services.prompt_builder import build_office_prompt
from api.services.skills import list_skills
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(),
)
# Pin the question to the insurance KB context so generic-law answers
# (e.g. "תקנה 37 לתקנות סדר הדין האזרחי") don't hijack the response.
system_prompt += (
"\n\nContext: The user is asking from the Knowledge Base UI of the "
"Israeli National Insurance firm. Every question should be answered "
"from the KB of חוק/תקנות/חוזרי הביטוח הלאומי. Before answering, "
"call search_insurance_kb to retrieve the relevant sections, and "
"cite them explicitly with section_ref (ס' N, תקנה N, or חוזר N/YYYY). "
"Never answer from generic legal training if the KB has a matching "
"section — if no KB section matches, say so explicitly."
)
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})
# Collect the sources Shira actually searched through during this turn.
# The legal_kb tool appends to this list on every search_insurance_kb
# call; the UI uses it to render a side PDF viewer next to the answer.
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,
)
except Exception as e:
logger.exception("[kb.ask] error")
raise HTTPException(status_code=500, detail=str(e))
# Keep only the ones that actually have a PDF attached.
pdf_sources = [
s for s in sources_used
if (s.get("original_path") or "").lower().endswith(".pdf")
]
return {
"text": text,
"conversationId": body.conversation_id,
"sources": pdf_sources,
}