Blame for backend/python-service/tabular_agent.py

from __future__ import annotations

import json
import os
import re
from typing import Any, Mapping, Optional
from urllib.parse import urlparse

from fastapi import HTTPException
from pydantic import BaseModel, Field, ValidationError

from inference.client import build_payload, post, text_from_response
from inference.config import resolve_model
from logger import current_trace_run, setup_logger
from models import ReviewFinding, ReviewResponse
from prompts.chat_system import chat_system_prompt
from prompts.json_repair_system import JSON_REPAIR_SYSTEM
from prompts.review_system import review_system_prompt
from prompts.review_user import review_user_prompt

logger = setup_logger(__name__)

try:
    from langsmith.run_trees import RunTree
except Exception:  # pragma: no cover - local tracing is optional.
    RunTree = None


ROOT_SCHEMA_PATH = os.path.join(
    os.path.dirname(__file__),
    "schemas",
    "word_review_output_schema.json",
)
LANGSMITH_PROJECT = os.environ.get("LANGSMITH_PROJECT", "legal-mike-oss")


class RawFinding(BaseModel):
    id: Optional[str] = None
    risk_level: Optional[str] = None
    clause_type: Optional[str] = None
    issue: Optional[str] = None
    comment: Optional[str] = None
    review_reasoning: Optional[dict] = None
    reasoning: Optional[dict] = None
    rationale: Optional[str] = None
    target_text: Optional[str] = None
    exact_text: Optional[str] = None
    start_paragraph_id: Optional[str] = None
    end_paragraph_id: Optional[str] = None
    start_anchor: Optional[str] = None
    end_anchor: Optional[str] = None
    before_anchor: Optional[str] = None
    after_anchor: Optional[str] = None
    anchor_text: Optional[str] = None
    anchor_mode: Optional[str] = None
    confidence: Optional[float] = None


class RawReview(BaseModel):
    schema_version: Optional[str] = None
    workflow_id: Optional[str] = None
    findings: list[RawFinding] = Field(default_factory=list)


class ChatSource(BaseModel):
    title: str
    url: str


class ChatParagraph(BaseModel):
    text: str
    citation_urls: list[str] = Field(default_factory=list)


class RawChatResponse(BaseModel):
    answer_paragraphs: list[ChatParagraph] = Field(default_factory=list)
    sources: list[ChatSource] = Field(default_factory=list)


def model_name(
    model_id: str | None = None,
    *,
    feature_id: str = "python-review",
) -> str:
    return resolve_model(request_model_id=model_id, feature_id=feature_id)


def public_agent_name() -> str:
    return "TabularAgent"


def _review_schema(max_findings: int) -> dict[str, Any]:
    if os.path.exists(ROOT_SCHEMA_PATH):
        try:
            with open(ROOT_SCHEMA_PATH, encoding="utf-8") as handle:
                schema = json.loads(handle.read())
            schema["properties"]["findings"]["maxItems"] = max_findings
            return schema
        except Exception:
            pass
    return ReviewResponse.model_json_schema()


def _reasoning_from_response(data: dict[str, Any]) -> str:
    try:
        parts = data["candidates"][0]["content"]["parts"]
        reasoning = "".join(
            str(part.get("text") or "") for part in parts if part.get("thought")
        )
    except (KeyError, IndexError, TypeError):
        reasoning = ""
    return reasoning


def _grounding_sources_from_response(data: dict[str, Any]) -> list[dict[str, str]]:
    try:
        metadata = data["candidates"][0].get("groundingMetadata") or {}
        chunks = metadata.get("groundingChunks") or []
    except (KeyError, IndexError, TypeError):
        return []
    sources: list[dict[str, str]] = []
    seen: set[str] = set()
    for chunk in chunks:
        web = chunk.get("web") if isinstance(chunk, dict) else None
        if not isinstance(web, dict):
            continue
        uri = str(web.get("uri") or "").strip()
        title = str(web.get("title") or "").strip() or uri
        if not uri or uri in seen:
            continue
        seen.add(uri)
        sources.append({"title": title, "uri": uri})
    return sources[:8]


def _append_sources(text: str, sources: list[dict[str, str]]) -> str:
    if not sources or any(source["uri"] in text for source in sources):
        return text
    lines = ["", "Sources:"]
    for index, source in enumerate(sources, start=1):
        lines.append(f"{index}. [{source['title']}]({source['uri']})")
    return text.rstrip() + "\n".join(lines)


def _clean_url(value: str) -> str | None:
    url = (value or "").strip()
    if not url:
        return None
    parsed = urlparse(url)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        return None
    return url


def _domain_label(url: str) -> str:
    try:
        host = urlparse(url).netloc.lower()
    except Exception:
        return url
    if host.startswith("www."):
        host = host[4:]
    return host or url


def _render_source_link(url: str, title: str | None = None) -> str:
    label = (title or "").strip() or _domain_label(url)
    return f"[{label}]({url})"


def _chat_json_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "answer_paragraphs": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "text": {"type": "string"},
                        "citation_urls": {"type": "array", "items": {"type": "string"}},
                    },
                    "required": ["text", "citation_urls"],
                    "additionalProperties": False,
                },
            },
            "sources": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "title": {"type": "string"},
                        "url": {"type": "string"},
                    },
                    "required": ["title", "url"],
                    "additionalProperties": False,
                },
            },
        },
        "required": ["answer_paragraphs", "sources"],
        "additionalProperties": False,
    }


def _normalize_chat_json(raw_text: str, grounding_sources: list[dict[str, str]]) -> str:
    obj = _extract_json_object(raw_text)
    parsed = RawChatResponse.model_validate(obj)
    source_map: dict[str, str] = {}
    for item in parsed.sources:
        clean = _clean_url(item.url)
        if not clean:
            continue
        source_map[clean] = item.title.strip() or _domain_label(clean)
    for item in grounding_sources:
        clean = _clean_url(item.get("uri", ""))
        if not clean:
            continue
        if clean not in source_map:
            source_map[clean] = (item.get("title") or "").strip() or _domain_label(
                clean
            )

    rendered_paragraphs: list[str] = []
    used_urls: list[str] = []
    seen_used: set[str] = set()
    for paragraph in parsed.answer_paragraphs:
        body = paragraph.text.strip()
        if not body:
            continue
        urls: list[str] = []
        for raw_url in paragraph.citation_urls:
            clean = _clean_url(raw_url)
            if not clean:
                continue
            urls.append(clean)
            if clean not in seen_used:
                used_urls.append(clean)
                seen_used.add(clean)
        if urls:
            inline_refs = " ".join(
                _render_source_link(url, source_map.get(url)) for url in urls
            )
            body = f"{body} {inline_refs}".strip()
        rendered_paragraphs.append(body)

    if not rendered_paragraphs:
        raise ValueError("Chat JSON response did not include answer_paragraphs.")

    source_lines = ["Sources:"]
    for index, url in enumerate(used_urls, start=1):
        source_lines.append(f"{index}. {_render_source_link(url, source_map.get(url))}")
    if not used_urls:
        source_lines.append("1. No external source links were provided.")

    return (
        ("\n\n".join(rendered_paragraphs)).rstrip()
        + "\n\n\n"
        + "\n\n".join(source_lines)
    )


def _extract_json_object(text: str) -> dict[str, Any]:
    stripped = text.strip()
    if stripped.startswith("```"):
        stripped = re.sub(r"^```(?:json)?\s*", "", stripped)
        stripped = re.sub(r"\s*```$", "", stripped)
    try:
        parsed = json.loads(stripped)
        if isinstance(parsed, dict):
            return parsed
    except json.JSONDecodeError:
        pass
    start = stripped.find("{")
    if start == -1:
        raise ValueError("No JSON object found in model response.")
    depth = 0
    in_string = False
    escaped = False
    for index, char in enumerate(stripped[start:], start=start):
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            continue
        if char == '"':
            in_string = True
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                parsed = json.loads(stripped[start : index + 1])
                if not isinstance(parsed, dict):
                    raise ValueError("Model JSON root was not an object.")
                return parsed
    raise ValueError("Could not parse a complete JSON object from model response.")


def _anchor(text: str, *, start: bool) -> str:
    value = " ".join(text.split())
    if len(value) <= 220:
        return value
    return value[:220] if start else value[-220:]


def _clean_id(value: str | None, fallback: str) -> str:
    if value and re.fullmatch(r"(P_[0-9]{4,}|BMK_[PT]_[a-z0-9]+)", value):
        return value
    return fallback


def _normalize_review(
    obj: dict[str, Any], *, workflow_id: str, max_findings: int
) -> ReviewResponse:
    raw = RawReview.model_validate(obj)
    findings: list[ReviewFinding] = []
    for index, item in enumerate(raw.findings[:max_findings], start=1):
        target = (item.target_text or item.exact_text or item.anchor_text or "").strip()
        issue = (item.issue or "").strip()
        comment = (item.comment or "").strip()
        if not issue or not comment:
            continue
        risk = (item.risk_level or "medium").lower()
        if risk not in {"critical", "high", "medium", "low", "info"}:
            risk = "medium"
        paragraph_id = _clean_id(item.start_paragraph_id, "P_0001")
        end_paragraph_id = _clean_id(item.end_paragraph_id, paragraph_id)
        confidence = item.confidence if item.confidence is not None else 0.75
        confidence = max(0, min(1, float(confidence)))
        findings.append(
            ReviewFinding(
                id=item.id or f"finding_{index:03d}",
                risk_level=risk,  # type: ignore[arg-type]
                clause_type=(item.clause_type or "General").strip() or "General",
                issue=issue,
                comment=comment,
                review_reasoning=(
                    item.review_reasoning
                    or item.reasoning
                    or {
                        "clause_section": "N/A",
                        "violated_standard": item.rationale or "N/A",
                        "zamil_exposure": "Model provided no structured rationale.",
                    }
                ),
                target_text=target,
                anchor_mode="exact"
                if (item.anchor_mode or "").lower() == "exact"
                else "paragraph",
                start_paragraph_id=paragraph_id,
                end_paragraph_id=end_paragraph_id,
                start_anchor=(item.start_anchor or _anchor(target, start=True)).strip()[
                    :254
                ],
                end_anchor=(item.end_anchor or _anchor(target, start=False)).strip()[
                    :254
                ],
                before_anchor=(
                    item.before_anchor or _anchor(target, start=True)
                ).strip()[:254],
                after_anchor=(
                    item.after_anchor or _anchor(target, start=False)
                ).strip()[:254],
                confidence=confidence,
            )
        )
    return ReviewResponse(
        workflow_id=workflow_id,
        findings=findings,
        model=public_agent_name(),
    )


async def _post(
    payload: dict[str, Any],
    *,
    trace_run: Any = None,
    model_id: str | None = None,
    feature_id: str = "python-review",
) -> tuple[str, dict[str, Any] | None, dict[str, Any]]:
    model = model_name(model_id, feature_id=feature_id)
    try:
        text, usage, data = await post(
            payload,
            model=model,
            feature_id=feature_id,
        )
    except HTTPException as exc:
        detail = str(exc.detail)
        if trace_run is not None:
            _end_trace(trace_run, error=detail)
        raise
    reasoning = _reasoning_from_response(data)
    if trace_run is not None:
        _end_trace(
            trace_run,
            outputs={
                "text": text,
                "thought_process": reasoning if reasoning else "[No thoughts captured]",
                "usage": usage,
            },
        )
    return text, usage, data


async def _repair_json(
    raw_text: str,
    error: Exception,
    *,
    workflow_id: str,
    max_findings: int,
    trace_run: Any = None,
    model_id: str | None = None,
) -> ReviewResponse:
    payload = build_payload(
        system_prompt=JSON_REPAIR_SYSTEM,
        user_prompt="\n\n".join(
            [
                "VALIDATION ERROR:",
                str(error)[:3000],
                "---",
                "SCHEMA:",
                json.dumps(_review_schema(max_findings), indent=2),
                "---",
                "RAW MODEL RESPONSE:",
                raw_text[:50000],
            ]
        ),
        max_tokens=int(os.environ.get("TABULAR_AGENT_REPAIR_MAX_TOKENS", "8192")),
        temperature=0,
        json_mode=True,
        model_id=model_id,
    )
    text, _usage, _data = await _post(
        payload,
        trace_run=trace_run,
        model_id=model_id,
        feature_id="python-review",
    )
    return _normalize_review(
        _extract_json_object(text), workflow_id=workflow_id, max_findings=max_findings
    )


def _new_trace(
    name: str,
    run_type: str,
    inputs: dict[str, Any],
    *,
    tags: list[str],
    trace_headers: Mapping[str, Any] | None = None,
) -> Any:
    if RunTree is None or not os.environ.get("LANGSMITH_API_KEY"):
        return None
    try:
        if trace_headers:
            run = RunTree.from_headers(
                trace_headers,
                name=name,
                run_type=run_type,
                inputs=inputs,
                project_name=LANGSMITH_PROJECT,
                tags=tags,
            )
            if run is not None:
                run.post()
                return run
        run = RunTree(
            name=name,
            run_type=run_type,
            inputs=inputs,
            project_name=LANGSMITH_PROJECT,
            tags=tags,
        )
        run.post()
        return run
    except Exception:
        return None


def _child_trace(
    parent: Any, name: str, run_type: str, inputs: dict[str, Any], *, tags: list[str]
) -> Any:
    if parent is None:
        return None
    try:
        child = parent.create_child(
            name=name, run_type=run_type, inputs=inputs, tags=tags
        )
        child.post()
        return child
    except Exception:
        return None


def _end_trace(
    run: Any, *, outputs: dict[str, Any] | None = None, error: str | None = None
) -> None:
    if run is None:
        return
    try:
        if error:
            run.end(error=error)
        else:
            run.end(outputs=outputs or {})
        run.patch()
    except Exception:
        pass


async def run_review(
    *,
    review_text: str,
    document_text: str,
    document_name: str,
    playbook_id: str,
    playbook_name: str,
    playbook_text: str,
    max_findings: int,
    model_id: str | None = None,
    trace_headers: Mapping[str, Any] | None = None,
) -> tuple[ReviewResponse, dict[str, Any] | None, str | None]:
    trace = _new_trace(
        "workflow.review",
        "chain",
        {
            "document": document_name,
            "playbook": playbook_name,
            "max_findings": max_findings,
        },
        tags=["workflow", "tabular-agent", "docx-comments"],
        trace_headers=trace_headers,
    )
    token = current_trace_run.set(trace)
    logger.info(
        f"Starting review for document '{document_name}' against playbook '{playbook_name}' with max {max_findings} findings."
    )
    try:
        payload = build_payload(
            system_prompt=review_system_prompt(max_findings),
            user_prompt=review_user_prompt(
                review_text=review_text,
                document_name=document_name,
                playbook_name=playbook_name,
                playbook_text=playbook_text,
            ),
            max_tokens=int(os.environ.get("TABULAR_AGENT_REVIEW_MAX_TOKENS", "16384")),
            temperature=0.1,
            json_mode=True,
            google_search=True,
            model_id=model_id,
        )
        llm_trace = _child_trace(
            trace,
            "tabular_agent.workflow.review",
            "llm",
            {"request": payload},
            tags=["workflow", "llm"],
        )
        logger.info(f"Sending LLM request for review of {document_name}")
        raw_text, usage, _data = await _post(
            payload,
            trace_run=llm_trace,
            model_id=model_id,
            feature_id="python-review",
        )
        logger.info(f"Received LLM response for review of {document_name}")
        parser_trace = _child_trace(
            trace,
            "validate.review_json",
            "parser",
            {"raw": raw_text},
            tags=["workflow", "parser"],
        )
        try:
            result = _normalize_review(
                _extract_json_object(raw_text),
                workflow_id=playbook_id,
                max_findings=max_findings,
            )
            result.usage = usage
            _end_trace(parser_trace, outputs={"finding_count": len(result.findings)})
            logger.info(f"Successfully parsed {len(result.findings)} findings.")
        except (ValueError, ValidationError, json.JSONDecodeError) as exc:
            logger.warning(
                f"Failed to parse initial LLM JSON response: {exc}. Attempting repair."
            )
            _end_trace(parser_trace, error=str(exc))
            repair_trace = _child_trace(
                trace,
                "tabular_agent.workflow.json_repair",
                "llm",
                {"raw": raw_text},
                tags=["workflow", "repair"],
            )
            result = await _repair_json(
                raw_text,
                exc,
                workflow_id=playbook_id,
                max_findings=max_findings,
                trace_run=repair_trace,
                model_id=model_id,
            )
            result.usage = usage
            logger.info(
                f"Successfully repaired JSON and parsed {len(result.findings)} findings."
            )
        result.model = public_agent_name()
        _end_trace(
            trace,
            outputs={"finding_count": len(result.findings), "workflow_id": playbook_id},
        )
        logger.info(f"Review complete for {document_name}.")
        return result, usage, getattr(trace, "id", None) if trace is not None else None
    except Exception as e:
        logger.error(f"Error during review of {document_name}: {e}", exc_info=True)
        raise
    finally:
        current_trace_run.reset(token)


async def run_chat(
    *,
    message: str,
    document_text: str | None,
    document_name: str | None,
    playbook_id: str,
    playbook_name: str,
    role_brief: str,
    model_id: str | None = None,
    trace_headers: Mapping[str, Any] | None = None,
) -> tuple[str, dict[str, Any] | None, str | None]:
    trace = _new_trace(
        "chat.answer",
        "chain",
        {
            "message": message,
            "document": document_name,
            "playbook": playbook_name,
        },
        tags=["chat", "tabular-agent", "google-search"],
        trace_headers=trace_headers,
    )
    token = current_trace_run.set(trace)
    logger.info(
        f"Starting chat for document '{document_name}' against playbook '{playbook_name}'."
    )
    try:
        system = chat_system_prompt(playbook_name=playbook_name, role_brief=role_brief)
        user = "\n\n".join(
            part
            for part in [
                f"DOCUMENT NAME: {document_name}" if document_name else "",
                "ATTACHED DOCUMENT TEXT:\n" + document_text if document_text else "",
                "USER MESSAGE:\n" + message,
            ]
            if part
        )
        payload = build_payload(
            system_prompt=system,
            user_prompt=user,
            max_tokens=int(os.environ.get("TABULAR_AGENT_CHAT_MAX_TOKENS", "8192")),
            temperature=0.2,
            json_schema=_chat_json_schema(),
            google_search=True,
            model_id=model_id,
        )
        llm_trace = _child_trace(
            trace, "tabular_agent.chat", "llm", {"request": payload}, tags=["chat", "llm"]
        )
        logger.info("Sending LLM chat request")
        raw_text, usage, data = await _post(
            payload,
            trace_run=llm_trace,
            model_id=model_id,
            feature_id="python-chat",
        )
        logger.info("Received LLM chat response")
        grounding_sources = _grounding_sources_from_response(data)
        try:
            text = _normalize_chat_json(raw_text, grounding_sources)
            logger.info("Successfully normalized chat response.")
        except Exception as exc:
            logger.warning(
                f"Failed to normalize chat JSON response: {exc}. Appending sources to raw text."
            )
            text = _append_sources(raw_text, grounding_sources)
        _end_trace(trace, outputs={"answer": text, "usage": usage})
        logger.info("Chat response complete.")
        return text, usage, getattr(trace, "id", None) if trace is not None else None
    except Exception as e:
        logger.error(f"Error during chat: {e}", exc_info=True)
        raise
    finally:
        current_trace_run.reset(token)


# Backward-compatible exports used by tests and callers.
_payload = build_payload