Blame for backend/python-service/main.py

from __future__ import annotations

import json
import os
import subprocess
from pathlib import Path
from typing import Any, Optional

from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, PlainTextResponse, StreamingResponse
from pydantic import BaseModel

from auth import UserSession, current_user
from docx_comments import apply_docx_comments
from docx_text import extract_docx_paragraphs, render_review_text
from tabular_agent import run_chat, run_review, _new_trace, _end_trace
from inference.config import manifest_sha256, resolve_manifest_path
from models import ChatConversationDetail, PlaybookDetail, ReviewFinding
from playbooks import chat_role_brief, get_playbook, list_playbooks, update_playbook
from logger import setup_logger, current_trace_run

logger = setup_logger(__name__)

from storage import (
    DATA_DIR,
    get_conversation,
    get_commented_docx_path,
    get_run,
    init_db,
    list_conversations,
    list_runs,
    new_id,
    reviewed_docx_name,
    run_to_markdown,
    save_chat_message,
    save_workflow_run,
)

ROOT_DIR = Path(__file__).resolve().parent
UPLOAD_DIR = DATA_DIR / "uploads"
ARTIFACT_DIR = DATA_DIR / "artifacts"
LANGSMITH_PROJECT = os.environ.get("LANGSMITH_PROJECT", "legal-mike-oss")
COMMENTER_PROJECT = (
    ROOT_DIR.parent / "tools" / "docx-commenter" / "docx-commenter.csproj"
)

app = FastAPI(title="Tabularrr Playbook Review Service")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class PlaybookUpdate(BaseModel):
    yaml_text: str


class ReviewRunRequest(BaseModel):
    document_bytes_b64: str
    document_name: str
    playbook_id: str = "buyer"
    max_findings: int = 30


class ApplyCommentsRequest(BaseModel):
    document_bytes_b64: str
    document_name: str
    findings: list[ReviewFinding]
    track_changes: bool = False


class ChatRequest(BaseModel):
    message: str
    conversation_id: Optional[str] = None
    playbook_id: str = "buyer"
    model_id: Optional[str] = None
    document_name: Optional[str] = None
    document_text: Optional[str] = None


@app.on_event("startup")
def startup() -> None:
    init_db()
    UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)


def _sse(event: str, data: Any) -> str:
    return f"data: {json.dumps({'event': event, 'data': data}, ensure_ascii=False)}\n\n"


def _safe_filename(name: str) -> str:
    clean = "".join(ch if ch.isalnum() or ch in "._- " else "_" for ch in name).strip()
    return clean or "document.docx"


@app.get("/api/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/api/inference-manifest")
def inference_manifest() -> dict[str, str]:
    path = resolve_manifest_path()
    if not path.exists():
        raise HTTPException(status_code=500, detail="Inference manifest is unavailable.")
    return {
        "path": str(path),
        "sha256": manifest_sha256(path),
    }


@app.get("/api/playbooks")
def api_list_playbooks(user: UserSession = Depends(current_user)) -> dict[str, Any]:
    return {"playbooks": [item.model_dump() for item in list_playbooks()]}


@app.get("/api/playbooks/{playbook_id}")
def api_get_playbook(
    playbook_id: str, user: UserSession = Depends(current_user)
) -> PlaybookDetail:
    return get_playbook(playbook_id)


@app.put("/api/playbooks/{playbook_id}")
def api_update_playbook(
    playbook_id: str,
    payload: PlaybookUpdate,
    user: UserSession = Depends(current_user),
) -> PlaybookDetail:
    return update_playbook(playbook_id, payload.yaml_text)


@app.get("/api/workflow-runs")
def api_list_runs(
    limit: int = 50,
    playbook_id: Optional[str] = None,
    user: UserSession = Depends(current_user),
) -> dict[str, Any]:
    return {
        "runs": [
            item.model_dump()
            for item in list_runs(user.username, limit=limit, playbook_id=playbook_id)
        ]
    }


@app.get("/api/workflow-runs/{run_id}")
def api_get_run(
    run_id: str, user: UserSession = Depends(current_user)
) -> dict[str, Any]:
    return get_run(run_id, user.username).model_dump()


@app.get("/api/workflow-runs/{run_id}/commented-docx")
def api_download_commented_docx(
    run_id: str, user: UserSession = Depends(current_user)
) -> FileResponse:
    detail = get_run(run_id, user.username)
    path = get_commented_docx_path(run_id, user.username)
    return FileResponse(
        path,
        media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        filename=detail.commented_docx_name or reviewed_docx_name(detail.document_name),
    )


@app.get("/api/workflow-runs/{run_id}/markdown")
def api_download_markdown(
    run_id: str, user: UserSession = Depends(current_user)
) -> PlainTextResponse:
    return PlainTextResponse(
        run_to_markdown(run_id, user.username),
        media_type="text/markdown; charset=utf-8",
        headers={"Content-Disposition": f'attachment; filename="{run_id}.md"'},
    )


@app.post("/api/workflow-runs")
async def api_create_workflow_run(
    request: Request,
    payload: ReviewRunRequest,
    user: UserSession = Depends(current_user),
) -> StreamingResponse:
    import base64

    run_id = new_id("run")
    run_dir = ARTIFACT_DIR / run_id
    run_dir.mkdir(parents=True, exist_ok=True)
    upload_dir = UPLOAD_DIR / run_id
    upload_dir.mkdir(parents=True, exist_ok=True)
    document_name = _safe_filename(payload.document_name)
    upload_path = upload_dir / document_name
    
    logger.info(f"Starting workflow run {run_id} for document '{document_name}' with playbook '{payload.playbook_id}'")

    docx_bytes = base64.b64decode(payload.document_bytes_b64)
    upload_path.write_bytes(docx_bytes)

    playbook_id = payload.playbook_id
    max_findings = max(1, min(100, int(payload.max_findings)))

    async def generate():
        trace = _new_trace(
            name="workflow_run",
            run_type="chain",
            inputs={"document_name": document_name, "playbook_id": playbook_id, "max_findings": max_findings},
            tags=["workflow", "api"],
            trace_headers=request.headers,
        )
        token = current_trace_run.set(trace)
        try:
            yield _sse("status", "Reading DOCX and extracting non-table contract text.")
            paragraphs = extract_docx_paragraphs(upload_path)
            if not paragraphs:
                raise HTTPException(
                    status_code=400,
                    detail="No non-table paragraph text was found in this DOCX.",
                )
            review_text = render_review_text(paragraphs)
            document_text = "\n\n".join(para.text for para in paragraphs)

            playbook = get_playbook(playbook_id)
            yield _sse("status", f"Loaded {playbook.name}.")
            yield _sse(
                "status",
                f"Sending {len(document_text):,} non-table characters to LLM for review.",
            )

            result_key = payload.playbook_id
            result, usage, _ = await run_review(
                review_text=review_text,
                document_text=document_text,
                document_name=document_name,
                playbook_id=playbook.id,
                playbook_name=playbook.name,
                playbook_text=playbook.yaml_text,
                max_findings=max_findings,
                trace_headers=request.headers,
            )
            yield _sse(
                "status", f"Validated {len(result.findings)} structured findings."
            )

            output_docx = run_dir / reviewed_docx_name(document_name)
            yield _sse(
                "status", "Inserting native Word comments into matched DOCX text."
            )
            apply_report = apply_docx_comments(
                input_docx=upload_path,
                output_docx=output_docx,
                findings=result.findings,
                work_dir=run_dir,
            )
            applied_count = sum(1 for item in apply_report if item.status == "applied")
            yield _sse("status", f"Applied {applied_count} native Word comments.")

            save_workflow_run(
                run_id=run_id,
                owner_username=user.username,
                document_name=document_name,
                original_docx_path=upload_path,
                commented_docx_path=output_docx if output_docx.exists() else None,
                document_text=document_text,
                playbook_id=playbook.id,
                playbook_name=playbook.name,
                playbook_text=playbook.yaml_text,
                max_findings=max_findings,
                status="completed",
                result=result,
                apply_report=apply_report,
                model=result.model,
                usage=usage,
            )
            logger.info(f"Workflow run {run_id} completed successfully.")
            _end_trace(trace, outputs={"run_id": run_id, "applied_count": applied_count})
            yield _sse("complete", get_run(run_id, user.username).model_dump())
        except HTTPException as exc:
            logger.warning(f"HTTPException in workflow run {run_id}: {exc.detail}")
            _end_trace(trace, error=str(exc.detail))
            yield _sse("error", exc.detail)
        except Exception as exc:
            logger.error(f"Exception in workflow run {run_id}: {exc}", exc_info=True)
            _end_trace(trace, error=str(exc))
            yield _sse("error", str(exc))
        finally:
            current_trace_run.reset(token)

    return StreamingResponse(generate(), media_type="text/event-stream")


@app.post("/api/apply-comments")
async def api_apply_comments(
    request: Request,
    payload: ApplyCommentsRequest,
    user: UserSession = Depends(current_user),
) -> dict[str, Any]:
    import base64

    run_id = new_id("comment")
    run_dir = ARTIFACT_DIR / run_id
    run_dir.mkdir(parents=True, exist_ok=True)
    
    document_name = _safe_filename(payload.document_name)
    logger.info(f"Applying comments to '{document_name}' in run {run_id}. Number of findings: {len(payload.findings)}")

    trace = _new_trace(
        name="apply_comments",
        run_type="chain",
        inputs={"document_name": document_name, "findings_count": len(payload.findings)},
        tags=["apply", "api"],
        trace_headers=request.headers,
    )
    token = current_trace_run.set(trace)
    
    input_path = run_dir / f"input_{document_name}"
    output_docx = run_dir / reviewed_docx_name(document_name)

    docx_bytes = base64.b64decode(payload.document_bytes_b64)
    input_path.write_bytes(docx_bytes)

    try:
        # Build sanitized inputs for LangSmith — no binary bytes, no credentials
        safe_findings = []
        for f in payload.findings:
            d = f.model_dump() if hasattr(f, "model_dump") else dict(f)
            safe = {
                "id": d.get("id"),
                "comment_key": d.get("comment_key"),
                "model_block_id": d.get("model_block_id"),
                "risk_level": d.get("risk_level"),
                "clause_type": d.get("clause_type"),
                "comment": d.get("comment"),
                "target_text": d.get("target_text"),
                "anchor_mode": d.get("anchor_mode"),
                "start_paragraph_id": d.get("start_paragraph_id"),
                "end_paragraph_id": d.get("end_paragraph_id"),
                "stable_id": d.get("stable_id"),
                "native_para_id": d.get("native_para_id"),
                "text_hash": d.get("text_hash"),
                "xml_path": d.get("xml_path"),
                "op_type": d.get("op_type"),
                "bookmark_name": d.get("bookmark_name"),
                "section_heading": d.get("section_heading"),
            }
            safe_findings.append(safe)

        apply_report = apply_docx_comments(
            input_docx=input_path,
            output_docx=output_docx,
            findings=payload.findings,
            work_dir=run_dir,
            track_changes=payload.track_changes,
        )
        
        # Return the output bytes as B64 so the TS side can handle storage
        output_bytes = output_docx.read_bytes() if output_docx.exists() else b""
        out_b64 = base64.b64encode(output_bytes).decode("ascii")

        applied = sum(1 for r in apply_report if r.status == "applied")
        not_applied = sum(1 for r in apply_report if r.status != "applied")
        failure_counts = {}
        for r in apply_report:
            if r.status != "applied" and r.failure_code:
                failure_counts[r.failure_code] = failure_counts.get(r.failure_code, 0) + 1
        
        finding_lookup = {d.get("id"): d for d in safe_findings}
        enriched_apply_report = []
        for item in apply_report:
            entry = item.model_dump()
            src = finding_lookup.get(item.id, {})
            entry["op_type"] = entry.get("op_type") or src.get("op_type")
            entry["bookmark_name"] = entry.get("bookmark_name") or src.get("bookmark_name")
            entry["section_heading"] = entry.get("section_heading") or src.get("section_heading")
            enriched_apply_report.append(entry)

        _end_trace(trace, outputs={
            "document_name": document_name,
            "findings_count": len(payload.findings),
            "applied": applied,
            "not_applied": not_applied,
            "failure_code_counts": failure_counts,
            "apply_report": enriched_apply_report,
            "sanitized_findings": safe_findings,
        })
        return {
            "document_name": document_name,
            "document_bytes_b64": out_b64,
            "commented_bytes_b64": out_b64,
            "apply_report": [item.model_dump() for item in apply_report],
        }
    except  Exception as exc:
        logger.error(f"Error applying comments: {exc}", exc_info=True)
        _end_trace(trace, error=str(exc))
        raise HTTPException(status_code=500, detail=str(exc))
    finally:
        current_trace_run.reset(token)


@app.post("/api/chat")
async def api_chat(
    request: Request,
    payload: ChatRequest,
    user: UserSession = Depends(current_user),
) -> StreamingResponse:
    document_name = payload.document_name
    document_text = payload.document_text
    playbook = get_playbook(payload.playbook_id)
    
    logger.info(f"Starting chat. Conversation ID: {payload.conversation_id}, Playbook: {payload.playbook_id}, Document: {document_name}")

    conv_id = save_chat_message(
        owner_username=user.username,
        conversation_id=payload.conversation_id,
        role="user",
        content=payload.message,
        document_name=document_name,
    )

    async def generate():
        trace = _new_trace(
            name="chat_api",
            run_type="chain",
            inputs={"message": payload.message, "document_name": document_name, "playbook_id": payload.playbook_id},
            tags=["chat", "api"],
            trace_headers=request.headers,
        )
        token = current_trace_run.set(trace)
        try:
            if document_name:
                yield _sse("status", f"Attached {document_name} to this chat.")
            answer, _usage, _ = await run_chat(
                message=payload.message,
                document_text=document_text,
                document_name=document_name,
                playbook_id=playbook.id,
                playbook_name=playbook.name,
                role_brief=chat_role_brief(playbook.id),
                model_id=payload.model_id,
                trace_headers=request.headers,
            )
            words = answer.split(" ")
            for index in range(0, len(words), 18):
                yield _sse(
                    "delta",
                    " ".join(words[index : index + 18])
                    + (" " if index + 18 < len(words) else ""),
                )
            save_chat_message(
                owner_username=user.username,
                conversation_id=conv_id,
                role="assistant",
                content=answer,
                document_name=None,
            )
            logger.info("Chat completed successfully.")
            _end_trace(trace, outputs={"answer_length": len(answer)})
            yield _sse("complete", answer)
        except Exception as exc:
            logger.error(f"Chat error: {exc}", exc_info=True)
            _end_trace(trace, error=str(exc))
            yield _sse("error", str(exc))
        finally:
            current_trace_run.reset(token)

    return StreamingResponse(generate(), media_type="text/event-stream")


@app.get("/api/chat/conversations")
def api_list_conversations(
    limit: int = 20, user: UserSession = Depends(current_user)
) -> dict[str, Any]:
    return {
        "conversations": [
            item.model_dump() for item in list_conversations(user.username, limit=limit)
        ]
    }


@app.get("/api/chat/conversations/{conversation_id}")
def api_get_conversation(
    conversation_id: str,
    user: UserSession = Depends(current_user),
) -> ChatConversationDetail:
    return get_conversation(conversation_id, user.username).model_dump()