Spaces:

bardd
/

mike-backend

Sleeping

mike-backend / backend /python-service /tabular_agent.py

Deploy Mike backend-only Space

06671aa verified 3 months ago

22.7 kB

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

	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