Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| import os | |
| import sqlite3 | |
| import uuid | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any, Optional | |
| from fastapi import HTTPException | |
| from models import ( | |
| ApplyResult, | |
| ChatConversation, | |
| ChatConversationDetail, | |
| ChatConversationSummary, | |
| ChatMessage, | |
| ReviewResponse, | |
| RunDetail, | |
| RunSummary, | |
| ) | |
| PUBLIC_AGENT_NAME = "TabularAgent" | |
| ROOT_DIR = Path(__file__).resolve().parent | |
| DATA_DIR = Path(os.environ.get("MIKE_PLAYBOOK_DATA_DIR", ROOT_DIR / "data")) | |
| DB_PATH = DATA_DIR / "tabular_review.sqlite3" | |
| def utc_now() -> str: | |
| return datetime.now(timezone.utc).isoformat() | |
| def new_id(prefix: str) -> str: | |
| return f"{prefix}_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:10]}" | |
| def reviewed_docx_name(document_name: str) -> str: | |
| stem = Path(document_name or "document.docx").stem or "document" | |
| safe_stem = "".join( | |
| ch if ch.isalnum() or ch in "._- " else "_" for ch in stem | |
| ).strip("._ ") | |
| safe_stem = safe_stem.replace(" ", "_") | |
| while "__" in safe_stem: | |
| safe_stem = safe_stem.replace("__", "_") | |
| return f"{safe_stem or 'document'}_tabular_reviewed.docx" | |
| def short_timestamp(value: str) -> str: | |
| try: | |
| dt = datetime.fromisoformat(value) | |
| except ValueError: | |
| return value[:16] | |
| return dt.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M") | |
| def connection() -> sqlite3.Connection: | |
| DATA_DIR.mkdir(parents=True, exist_ok=True) | |
| conn = sqlite3.connect(DB_PATH) | |
| conn.row_factory = sqlite3.Row | |
| return conn | |
| def init_db() -> None: | |
| with connection() as conn: | |
| conn.executescript( | |
| """ | |
| create table if not exists workflow_runs ( | |
| run_id text primary key, | |
| owner_username text not null default 'admin', | |
| created_at text not null, | |
| document_name text not null, | |
| original_docx_path text not null, | |
| commented_docx_path text, | |
| document_text text not null, | |
| playbook_id text not null, | |
| playbook_name text not null, | |
| playbook_text text not null, | |
| max_findings integer not null, | |
| status text not null, | |
| result_json text not null, | |
| apply_report_json text not null, | |
| model text, | |
| usage_json text | |
| ); | |
| create table if not exists chat_conversations ( | |
| conversation_id text primary key, | |
| owner_username text not null default 'admin', | |
| title text not null, | |
| created_at text not null, | |
| updated_at text not null | |
| ); | |
| create table if not exists chat_messages ( | |
| message_id text primary key, | |
| conversation_id text not null, | |
| owner_username text not null default 'admin', | |
| created_at text not null, | |
| role text not null, | |
| content text not null, | |
| document_name text, | |
| foreign key(conversation_id) references chat_conversations(conversation_id) | |
| ); | |
| """ | |
| ) | |
| _ensure_column( | |
| conn, "workflow_runs", "owner_username", "text not null default 'admin'" | |
| ) | |
| _ensure_column( | |
| conn, | |
| "chat_conversations", | |
| "owner_username", | |
| "text not null default 'admin'", | |
| ) | |
| _ensure_column( | |
| conn, "chat_messages", "owner_username", "text not null default 'admin'" | |
| ) | |
| def _ensure_column( | |
| conn: sqlite3.Connection, table: str, column: str, definition: str | |
| ) -> None: | |
| existing = { | |
| row["name"] for row in conn.execute(f"pragma table_info({table})").fetchall() | |
| } | |
| if column not in existing: | |
| conn.execute(f"alter table {table} add column {column} {definition}") | |
| def save_workflow_run( | |
| *, | |
| owner_username: str, | |
| run_id: str, | |
| document_name: str, | |
| original_docx_path: Path, | |
| commented_docx_path: Optional[Path], | |
| document_text: str, | |
| playbook_id: str, | |
| playbook_name: str, | |
| playbook_text: str, | |
| max_findings: int, | |
| status: str, | |
| result: ReviewResponse, | |
| apply_report: list[ApplyResult], | |
| model: Optional[str], | |
| usage: Optional[dict[str, Any]], | |
| ) -> None: | |
| with connection() as conn: | |
| conn.execute( | |
| """ | |
| insert or replace into workflow_runs ( | |
| run_id, owner_username, created_at, document_name, original_docx_path, commented_docx_path, | |
| document_text, playbook_id, playbook_name, playbook_text, max_findings, | |
| status, result_json, apply_report_json, model, usage_json | |
| ) | |
| values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) | |
| """, | |
| ( | |
| run_id, | |
| owner_username, | |
| utc_now(), | |
| document_name, | |
| str(original_docx_path), | |
| str(commented_docx_path) if commented_docx_path else None, | |
| document_text, | |
| playbook_id, | |
| playbook_name, | |
| playbook_text, | |
| max_findings, | |
| status, | |
| result.model_dump_json(), | |
| json.dumps( | |
| [item.model_dump() for item in apply_report], ensure_ascii=False | |
| ), | |
| model, | |
| json.dumps(usage, ensure_ascii=False) if usage else None, | |
| ), | |
| ) | |
| def _summary_from_row(row: sqlite3.Row) -> RunSummary: | |
| result = json.loads(row["result_json"]) | |
| apply_report = json.loads(row["apply_report_json"]) | |
| findings = result.get("findings") if isinstance(result, dict) else [] | |
| applied_count = sum(1 for item in apply_report if item.get("status") == "applied") | |
| return RunSummary( | |
| run_id=row["run_id"], | |
| created_at=row["created_at"], | |
| created_short=short_timestamp(row["created_at"]), | |
| document_name=row["document_name"], | |
| commented_docx_name=reviewed_docx_name(row["document_name"]), | |
| playbook_id=row["playbook_id"], | |
| playbook_name=row["playbook_name"], | |
| status=row["status"], | |
| finding_count=len(findings) if isinstance(findings, list) else 0, | |
| applied_count=applied_count, | |
| model=PUBLIC_AGENT_NAME, | |
| ) | |
| def list_runs( | |
| owner_username: str, limit: int = 50, playbook_id: str | None = None | |
| ) -> list[RunSummary]: | |
| with connection() as conn: | |
| if playbook_id: | |
| rows = conn.execute( | |
| "select * from workflow_runs where owner_username = ? and playbook_id = ? order by created_at desc limit ?", | |
| (owner_username, playbook_id, limit), | |
| ).fetchall() | |
| else: | |
| rows = conn.execute( | |
| "select * from workflow_runs where owner_username = ? order by created_at desc limit ?", | |
| (owner_username, limit), | |
| ).fetchall() | |
| return [_summary_from_row(row) for row in rows] | |
| def get_run(run_id: str, owner_username: str) -> RunDetail: | |
| with connection() as conn: | |
| row = conn.execute( | |
| "select * from workflow_runs where run_id = ? and owner_username = ?", | |
| (run_id, owner_username), | |
| ).fetchone() | |
| if not row: | |
| raise HTTPException(status_code=404, detail=f"Workflow run not found: {run_id}") | |
| summary = _summary_from_row(row) | |
| usage = json.loads(row["usage_json"]) if row["usage_json"] else None | |
| commented_path = ( | |
| Path(row["commented_docx_path"]) if row["commented_docx_path"] else None | |
| ) | |
| return RunDetail( | |
| **summary.model_dump(), | |
| document_text=row["document_text"], | |
| playbook_text=row["playbook_text"], | |
| result=ReviewResponse.model_validate_json(row["result_json"]), | |
| apply_report=[ | |
| ApplyResult.model_validate(item) | |
| for item in json.loads(row["apply_report_json"]) | |
| ], | |
| usage=usage, | |
| commented_docx_available=bool(commented_path and commented_path.exists()), | |
| ) | |
| def get_commented_docx_path(run_id: str, owner_username: str) -> Path: | |
| with connection() as conn: | |
| row = conn.execute( | |
| "select commented_docx_path from workflow_runs where run_id = ? and owner_username = ?", | |
| (run_id, owner_username), | |
| ).fetchone() | |
| if not row or not row["commented_docx_path"]: | |
| raise HTTPException( | |
| status_code=404, detail="Commented DOCX is not available for this run." | |
| ) | |
| path = Path(row["commented_docx_path"]) | |
| if not path.exists(): | |
| raise HTTPException(status_code=404, detail="Commented DOCX file is missing.") | |
| return path | |
| def save_chat_message( | |
| *, | |
| owner_username: str, | |
| conversation_id: Optional[str], | |
| role: str, | |
| content: str, | |
| document_name: Optional[str] = None, | |
| ) -> str: | |
| now = utc_now() | |
| conv_id = conversation_id or new_id("conv") | |
| title = content.strip().splitlines()[0][:80] or "New conversation" | |
| message_id = new_id("msg") | |
| with connection() as conn: | |
| existing = conn.execute( | |
| "select conversation_id from chat_conversations where conversation_id = ? and owner_username = ?", | |
| (conv_id, owner_username), | |
| ).fetchone() | |
| if not existing: | |
| conn.execute( | |
| "insert into chat_conversations (conversation_id, owner_username, title, created_at, updated_at) values (?, ?, ?, ?, ?)", | |
| (conv_id, owner_username, title, now, now), | |
| ) | |
| else: | |
| conn.execute( | |
| "update chat_conversations set updated_at = ? where conversation_id = ? and owner_username = ?", | |
| (now, conv_id, owner_username), | |
| ) | |
| conn.execute( | |
| """ | |
| insert into chat_messages ( | |
| message_id, conversation_id, owner_username, created_at, role, content, document_name | |
| ) | |
| values (?, ?, ?, ?, ?, ?, ?) | |
| """, | |
| ( | |
| message_id, | |
| conv_id, | |
| owner_username, | |
| now, | |
| role, | |
| content, | |
| document_name, | |
| ), | |
| ) | |
| return conv_id | |
| def list_conversations( | |
| owner_username: str, limit: int = 20 | |
| ) -> list[ChatConversationSummary]: | |
| with connection() as conn: | |
| rows = conn.execute( | |
| """ | |
| select conversation_id, title, created_at, updated_at | |
| from chat_conversations | |
| where owner_username = ? | |
| order by updated_at desc | |
| limit ? | |
| """, | |
| (owner_username, limit), | |
| ).fetchall() | |
| return [ | |
| ChatConversationSummary( | |
| conversation_id=row["conversation_id"], | |
| title=row["title"], | |
| created_at=row["created_at"], | |
| updated_at=row["updated_at"], | |
| ) | |
| for row in rows | |
| ] | |
| def get_conversation( | |
| conversation_id: str, owner_username: str | |
| ) -> ChatConversationDetail: | |
| with connection() as conn: | |
| row = conn.execute( | |
| """ | |
| select conversation_id, owner_username, title, created_at, updated_at | |
| from chat_conversations | |
| where conversation_id = ? and owner_username = ? | |
| """, | |
| (conversation_id, owner_username), | |
| ).fetchone() | |
| if not row: | |
| raise HTTPException( | |
| status_code=404, detail=f"Conversation not found: {conversation_id}" | |
| ) | |
| messages = conn.execute( | |
| """ | |
| select message_id, created_at, role, content, document_name | |
| from chat_messages | |
| where conversation_id = ? and owner_username = ? | |
| order by created_at asc | |
| """, | |
| (conversation_id, owner_username), | |
| ).fetchall() | |
| return ChatConversationDetail( | |
| conversation_id=row["conversation_id"], | |
| owner_username=row["owner_username"], | |
| title=row["title"], | |
| created_at=row["created_at"], | |
| updated_at=row["updated_at"], | |
| messages=[ | |
| ChatMessage( | |
| id=message["message_id"], | |
| role=message["role"], | |
| content=message["content"], | |
| created_at=message["created_at"], | |
| document_name=message["document_name"], | |
| ) | |
| for message in messages | |
| ], | |
| ) | |
| def run_to_markdown(run_id: str, owner_username: str) -> str: | |
| run = get_run(run_id, owner_username) | |
| lines = [ | |
| f"# Review Run: {run.document_name}", | |
| "", | |
| "## Metadata", | |
| "", | |
| f"- Run ID: `{run.run_id}`", | |
| f"- Created: `{run.created_short}`", | |
| f"- Playbook: `{run.playbook_name}` (`{run.playbook_id}`)", | |
| f"- Status: `{run.status}`", | |
| f"- Agent: `{PUBLIC_AGENT_NAME}`", | |
| f"- Findings: `{run.finding_count}`", | |
| f"- Comments applied: `{run.applied_count}`", | |
| "", | |
| "## Findings", | |
| "", | |
| ] | |
| for index, finding in enumerate(run.result.findings, start=1): | |
| apply_result = next( | |
| (item for item in run.apply_report if item.id == finding.id), None | |
| ) | |
| lines.extend( | |
| [ | |
| f"### {index}. {finding.clause_type} - {finding.risk_level}", | |
| "", | |
| f"**Issue:** {finding.issue}", | |
| "", | |
| f"**Comment inserted:** {finding.comment}", | |
| "", | |
| f"**Internal reasoning:**\n- **Clause/Section:** {finding.review_reasoning.get('clause_section', 'N/A') if isinstance(finding.review_reasoning, dict) else finding.review_reasoning}\n- **Violated Standard:** {finding.review_reasoning.get('violated_standard', 'N/A') if isinstance(finding.review_reasoning, dict) else 'N/A'}\n- **Zamil Exposure:** {finding.review_reasoning.get('zamil_exposure', 'N/A') if isinstance(finding.review_reasoning, dict) else 'N/A'}", | |
| "", | |
| f"**Confidence:** {finding.confidence}", | |
| "", | |
| f"**Apply status:** {apply_result.status if apply_result else 'not_applied'}", | |
| "", | |
| "**Target text:**", | |
| "", | |
| "```text", | |
| finding.target_text, | |
| "```", | |
| "", | |
| ] | |
| ) | |
| if not run.result.findings: | |
| lines.append("No findings were returned.") | |
| lines.append("") | |
| lines.extend( | |
| [ | |
| "## Apply Report", | |
| "", | |
| "```json", | |
| json.dumps( | |
| [item.model_dump() for item in run.apply_report], | |
| indent=2, | |
| ensure_ascii=False, | |
| ), | |
| "```", | |
| "", | |
| "## Document Text", | |
| "", | |
| "```text", | |
| run.document_text, | |
| "```", | |
| "", | |
| "## Playbook Text", | |
| "", | |
| "```yaml", | |
| run.playbook_text, | |
| "```", | |
| "", | |
| "## Raw Result JSON", | |
| "", | |
| "```json", | |
| run.result.model_dump_json(indent=2), | |
| "```", | |
| "", | |
| ] | |
| ) | |
| return "\n".join(lines) | |