chaim
b127cafb01
Two minimal additions for the upcoming v0.1.11 client work:
- /kb/search hits now include chunk_index, so the client can call
/source/{id}/chunks?around=N&ctx=K to fetch the matched chunk plus K
chunks before/after as context. Required by the new search-preview
layout for text sources, which until now duplicated the matched chunk
in both panes.
- /source/{id}/chunks accepts around (chunk_index) + ctx (window radius
clamped to 0..10). When around is given, returns only the windowed
range; otherwise unchanged (returns all chunks of the source).
Refs Task Master #1
315 lines
12 KiB
Python
315 lines
12 KiB
Python
"""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)
|
||
# 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
|
||
|
||
|
||
@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,
|
||
expand=body.expand,
|
||
)
|
||
# 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,
|
||
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
|
||
|
||
|
||
@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,
|
||
}
|