shira-hermes/mcp_server/tools/legal_kb_tools.py

"""Legal knowledge base tool — topic-scoped semantic + lexical search.

Renamed from search_insurance_kb to search_legal_kb in Phase 1 of the
multi-topic refactor (Task Master KnowledgeBase #12). The handler now
takes a topic_id at registration time so each conversation is scoped to
its caller-chosen domain.
"""

from __future__ import annotations

import logging

from api.services.kb import search as kb_search
from mcp_server.tools._helpers import fail

logger = logging.getLogger("shira.tools.legal_kb")

_KIND_HEBREW = {
    "law": "חוק",
    "regulation": "תקנה",
    "circular": "חוזר",
}


def _format_hits(hits: list[dict]) -> str:
    if not hits:
        return "לא נמצאו תוצאות רלוונטיות בבסיס הידע."

    lines = [f"נמצאו {len(hits)} קטעים רלוונטיים:"]
    for i, h in enumerate(hits, 1):
        kind_he = _KIND_HEBREW.get(h["kind"], h["kind"])
        ident = f" {h['identifier']}" if h.get("identifier") else ""
        header = f"{i}. [{kind_he}{ident}] {h['title']}"

        # Prefer the full hierarchical path when present; fall back to section_ref.
        path = h.get("heading_path") or h.get("section_ref") or ""
        if path:
            header += f" — {path}"

        # Date annotation for time-sensitive citations.
        pub = h.get("published_at")
        if pub:
            header += f" (פורסם {pub})"

        lines.append(header)
        content = (h.get("content") or "").strip()
        if len(content) > 900:
            content = content[:900] + "…"
        lines.append(content)
        if h.get("source_url"):
            lines.append(f"מקור: {h['source_url']}")
        lines.append("")
    return "\n".join(lines).rstrip()


def register_legal_kb_tools(
    tools: dict,
    sources_used: list[dict] | None = None,
    topic_id: int | None = None,
    topic_name: str | None = None,
) -> None:
    """Register the search_legal_kb tool.

    topic_id pins retrieval to a single domain (kb_source.topic_id). When
    None, the underlying search falls back to the lowest active topic —
    the route layer (kb_public.py) normally resolves a concrete id before
    this is called, so None should be a transitional case only.

    topic_name is interpolated into the tool description so the LLM sees
    a domain-specific description (e.g. "ביטוח לאומי" vs "דיני עבודה").

    If `sources_used` is provided, each hit's (source_id, page_number,
    section_ref, title, kind, original_path) is appended — deduplicated
    by (source_id, page_number) — so the caller can build a "view PDF"
    link list alongside the agent's text answer.
    """
    domain_text = topic_name or "המשפט"

    async def search_legal_kb(
        query: str,
        kind: str = "any",
        top_k: int = 8,
    ) -> str:
        try:
            if kind not in ("law", "regulation", "circular", "any"):
                kind = "any"
            top_k = max(1, min(int(top_k or 8), 15))
            hits = await kb_search.search(
                query=query,
                kind=kind,
                top_k=top_k,
                topic_id=topic_id,
                topic_name=topic_name,
            )
            if sources_used is not None:
                seen = {(s.get("source_id"), s.get("page_number")) for s in sources_used}
                for h in hits:
                    key = (h.get("source_id"), h.get("page_number"))
                    if h.get("source_id") and key not in seen:
                        sources_used.append({
                            "source_id": h["source_id"],
                            "title": h.get("title"),
                            "kind": h.get("kind"),
                            "identifier": h.get("identifier"),
                            "heading_path": h.get("heading_path"),
                            "section_ref": h.get("section_ref"),
                            "page_number": h.get("page_number"),
                            "original_path": h.get("original_path"),
                        })
                        seen.add(key)
            return _format_hits(hits)
        except Exception as e:
            logger.exception("[search_legal_kb] failed (topic_id=%s)", topic_id)
            return fail(f"שגיאה בחיפוש בבסיס הידע: {e}")

    tools["search_legal_kb"] = {
        "description": (
            f"חיפוש בבסיס הידע המשפטי בנושא {domain_text} — חוקים, תקנות וחוזרים. "
            "מחזיר קטעים רלוונטיים עם ציטוט מדויק (סעיף/תקנה/מס' חוזר). "
            "יש להשתמש בכלי לפני טענה משפטית מבוססת-מקור, ולצטט את ה-section_ref "
            "המדויק בתשובה. החיפוש משלב דמיון סמנטי וחיפוש מילולי (hybrid)."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "שאילתת חיפוש בעברית — ניסוח חופשי או מונחים מדויקים.",
                },
                "kind": {
                    "type": "string",
                    "enum": ["law", "regulation", "circular", "any"],
                    "description": "סינון לפי סוג מקור: law=חוק, regulation=תקנה, circular=חוזר, any=הכל (ברירת מחדל).",
                },
                "top_k": {
                    "type": "integer",
                    "description": "מספר קטעים להחזיר (1–15, ברירת מחדל 8).",
                },
            },
            "required": ["query"],
        },
        "handler": search_legal_kb,
    }