agent.py

#!/usr/bin/env python3
import re
import argparse
import os
import json
import sqlite3
import sys
import time
from pathlib import Path
from typing import Any, Optional, List
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIModel
from langchain_core.tools import tool as lc_tool
from deepagents import create_deep_agent
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command
import packager
from openai import OpenAI as _OpenAI

TRANSLATE_MODEL = "gemma4:26b"
OLLAMA_BASE_URL = "http://localhost:11434/v1"
_TRANSLATE_CLIENT = _OpenAI(base_url=OLLAMA_BASE_URL, api_key="ollama")

LANGUAGES: dict[str, str] = {
    "zh": "Simplified Chinese (Mandarin)",
    "ar": "Modern Standard Arabic",
}

TRANSLATE_PROMPT = (
    "You are a professional translator specialising in plain-language safety guidance. "
    "Translate the following English text into {lang_name}. "
    "Preserve numbered lists, bullet structure, and proper nouns. "
    "Return ONLY the translated text — no explanation, no preamble, no quotation marks.\n\n"
    "TEXT:\n{text}"
)

# ── Pydantic Models for Tool Returns ──────────────────────────────────

class ChunkInfo(BaseModel):
    text: str
    section: str
    language: str
    source_file: str
    page_num: int
    token_estimate: int

class ChunkResult(BaseModel):
    row_count: int
    name: str
    tier: str

class TranslatedChunk(BaseModel):
    text: str
    language: str
    original_text: str

class ProvenanceDict(BaseModel):
    publisher: str
    license: str
    source_url: str
    reviewed_at: str
    cultural_sensitivity: Optional[str] = None
    expires_at: Optional[str] = None
    language: Optional[str] = None

class GeoPoint(BaseModel):
    name: str
    lat: float
    lon: float
    category: str

class ValidationResult(BaseModel):
    valid: bool
    checks: dict[str, Any] = Field(default_factory=dict)
    error: Optional[str] = None

# ── Helpers ───────────────────────────────────────────────────────────

def _validate_path(path: str) -> str:
    """Validate that the path is within the current working directory.
    
    This addresses CR-04 (Path Traversal).
    """
    root = Path.cwd().resolve()
    candidate = Path(path).resolve()
    try:
        candidate.relative_to(root)
    except ValueError as exc:
        raise ValueError(
            f"Path traversal detected: {path} is outside allowed directory."
        ) from exc
    return str(candidate)

def _ensure_layer(pack_path: str, layer_name: str, tier: str = "static_reference"):
    """Ensure a layer exists in the database."""
    conn = sqlite3.connect(pack_path)
    try:
        exists = conn.execute("SELECT 1 FROM layers WHERE name = ?", (layer_name,)).fetchone()
        if not exists:
            with conn:
                conn.execute(
                    "INSERT INTO layers (name, tier, description, added_at) VALUES (?, ?, ?, ?)",
                    (layer_name, tier, f"Auto-created layer for {layer_name}", packager._now())
                )
    finally:
        conn.close()

def _read_markdown_inputs(input_path: str) -> list[tuple[str, str]]:
    """Read markdown inputs from a file or directory before invoking the compiler model."""
    input_p = Path(input_path)
    if input_p.is_dir():
        files = sorted(list(input_p.glob("*.md")))
    else:
        files = [input_p]
    return [(f_path.name, f_path.read_text(encoding="utf-8")) for f_path in files]


def _batch_markdown(content: str, max_chars: int = 60_000) -> list[str]:
    """Split markdown into section-aligned batches each under max_chars.

    Splits on H1/H2 boundaries so sections are never torn mid-paragraph.
    Falls back to the whole content as one batch if no headings exist.
    """
    # Split keeping the heading delimiter with each section
    sections = re.split(r'(?=^#{1,2} )', content, flags=re.MULTILINE)
    sections = [s for s in sections if s.strip()]

    batches: list[str] = []
    current_parts: list[str] = []
    current_len = 0

    for section in sections:
        if current_len + len(section) > max_chars and current_parts:
            batches.append("".join(current_parts))
            current_parts = [section]
            current_len = len(section)
        else:
            current_parts.append(section)
            current_len += len(section)

    if current_parts:
        batches.append("".join(current_parts))

    return batches or [content]


# ── PydanticAI Compiler Agent ─────────────────────────────────────────

# Default model as per D-15
DEFAULT_MODEL = "offlineaid-compiler"
DEFAULT_BASE_URL = "http://localhost:11434/v1"

# Initialize with a dummy key for Ollama compatibility and to avoid early exit if env is missing
def get_model(model_name: str = DEFAULT_MODEL, base_url: str = DEFAULT_BASE_URL):
    return OpenAIModel(model_name, base_url=base_url, api_key="ollama")

compiler_agent = Agent(
    model=get_model(),
    result_type=str,
    system_prompt=(
        "You are the OfflineAid compiler agent. Transform raw markdown into a structured knowledge pack. "
        "After calling extract_chunks for any content, call translate_chunk(chunk, 'zh', pack_path) "
        "and translate_chunk(chunk, 'ar', pack_path) for each extracted chunk to add Mandarin and Arabic versions. "
        "Always pass pack_path explicitly on every tool call."
    )
)

@compiler_agent.tool_plain
def extract_chunks(text: str, pack_path: str, source_file: str, language: str) -> dict:
    """Split inspected markdown into chunks and save them to the pack.

    Args:
        text: The content to chunk.
        pack_path: Path to the SQLite .db pack.
        source_file: Filename for attribution.
        language: BCP-47 language tag.
    """
    print(f"[agent] tool=extract_chunks  source={source_file}  lang={language}  chars={len(text)}", flush=True)
    _validate_path(pack_path)
    chunks = []
    current_section = "Introduction"
    
    # Simple split by Markdown H2 headings
    sections = re.split(r'^(##\s+.*)$', text, flags=re.MULTILINE)
    
    # Initial content before any H2
    if sections and not sections[0].startswith("## "):
        intro_text = sections[0].strip()
        if intro_text:
            chunks.append(ChunkInfo(
                text=intro_text,
                section=current_section,
                language=language,
                source_file=source_file,
                page_num=1,
                token_estimate=packager._estimate_tokens(intro_text)
            ).model_dump())
        sections = sections[1:]

    # Subsequent sections
    for i in range(0, len(sections), 2):
        header = sections[i].strip("# ").strip()
        content = sections[i+1].strip() if i+1 < len(sections) else ""
        if content:
            chunks.append(ChunkInfo(
                text=content,
                section=header,
                language=language,
                source_file=source_file,
                page_num=1,
                token_estimate=packager._estimate_tokens(content)
            ).model_dump())
            
    print(f"[agent]   -> {len(chunks)} chunks  pack={pack_path}", flush=True)
    rows = [{"text": c["text"], "section": c.get("section", "Introduction")} for c in chunks]
    res = packager.add_layer(
        pack_path=pack_path,
        name=source_file,
        tier="static_reference",
        rows=rows,
        source_info=json.dumps({"file": source_file}),
        provenance={"language": language}
    )
    return ChunkResult(**res).model_dump()

@compiler_agent.tool_plain
def translate_chunk(chunk: dict, target_lang: str, pack_path: str) -> dict:
    """Translate a chunk dict and persist as a language-namespaced layer.

    Args:
        chunk: The chunk dict to translate. Expected keys: text, language, source_file, section.
        target_lang: Target language tag ('zh' or 'ar' triggers Ollama call; unknown tags are no-ops per D-04 fallback).
        pack_path: Path to the SQLite .db pack.
    """
    print(f"[agent] tool=translate_chunk  target={target_lang}  chars={len(chunk.get('text', ''))}", flush=True)
    _validate_path(pack_path)
    original_text = chunk["text"]
    source_lang = chunk.get("language")

    if target_lang in LANGUAGES:
        # D-04/D-05: real Ollama call to gemma4:26b stock
        lang_name = LANGUAGES[target_lang]
        response = _TRANSLATE_CLIENT.chat.completions.create(
            model=TRANSLATE_MODEL,
            messages=[{
                "role": "user",
                "content": TRANSLATE_PROMPT.format(lang_name=lang_name, text=original_text),
            }],
            temperature=0.3,
        )
        translated_text = response.choices[0].message.content.strip()

        # Pitfall 2: language-namespaced layer name avoids INSERT OR REPLACE collision
        source_file = chunk.get("source_file") or chunk.get("section") or "unknown"
        layer_name = f"{source_file}.{target_lang}"
        rows = [{"text": translated_text, "section": chunk.get("section", "")}]
        packager.add_layer(
            pack_path=pack_path,
            name=layer_name,
            tier="static_reference",
            rows=rows,
            source_info=json.dumps({"file": source_file}),
            provenance={"language": target_lang},
        )
        print(f"[agent]   -> translated  layer={layer_name}  chars={len(translated_text)}", flush=True)
    else:
        # Unknown lang fallback (D-04): preserves test_translate_chunk_uses_pack_scoped_contract contract
        translated_text = original_text

    res = {
        "text": translated_text,
        "language": target_lang if target_lang != source_lang else source_lang or target_lang,
        "original_text": original_text,
    }
    return TranslatedChunk(**res).model_dump()

@compiler_agent.tool_plain
def derive_provenance(source_url: str, pack_path: str) -> dict:
    """Derive provenance metadata for a source.
    
    Args:
        source_url: Canonical URL of the source.
        pack_path: Path to the SQLite .db pack.
    """
    print(f"[agent] tool=derive_provenance  url={source_url}", flush=True)
    _validate_path(pack_path)
    res = {
        "publisher": "Auto-derived",
        "license": "CC-BY-4.0",
        "source_url": source_url,
        "reviewed_at": packager._now()
    }
    return ProvenanceDict(**res).model_dump()

@compiler_agent.tool_plain
def propose_geo_points(text: str, region: str, pack_path: str) -> List[dict]:
    """Extract named lat/lon candidates from text using regex. (Addresses CR-03)
    
    Args:
        text: Text to scan for coordinates (e.g. 'Hospital: 22.3, 114.1').
        region: Geographic region for context.
        pack_path: Path to the SQLite .db pack.
    """
    print(f"[agent] tool=propose_geo_points  region={region}  chars={len(text)}", flush=True)
    _validate_path(pack_path)
    points: List[dict] = []
    
    # Matches: "Hospital: 22.3, 114.1" or "lat: 22.3, lon: 114.1"
    # Improved regex: look for a word just before the colon
    pattern = r'(\b\w+):\s*(?:lat:\s*)?(-?\d+\.\d+),\s*(?:lon:\s*)?(-?\d+\.\d+)'
    matches = re.finditer(pattern, text)
    
    for match in matches:
        name = match.group(1).strip()
        lat, lon = float(match.group(2)), float(match.group(3))
        
        # Filter out "lat" being captured as name if pattern was "lat: X, lon: Y"
        if name.lower() == "lat":
            name = "Unnamed Point"
        
        points.append(GeoPoint(
            name=name,
            lat=lat,
            lon=lon,
            category="candidate"
        ).model_dump())
            
    if points:
        _ensure_layer(pack_path, "geo_candidates")
        packager.add_geo_points(pack_path, layer="geo_candidates", points=points)
    print(f"[agent]   -> {len(points)} geo point(s) found", flush=True)
    return points

@compiler_agent.tool_plain
def validate_pack(pack_path: str) -> dict:
    """Validate the pack archive and structure.
    
    Args:
        pack_path: Path to the .db file OR the .oapack.zip archive.
    """
    print(f"[agent] tool=validate_pack  path={pack_path}", flush=True)
    _validate_path(pack_path)
    try:
        if pack_path.endswith(".db"):
            # Per D-10, wrap archive validation
            archive_path = packager.archive_pack(pack_path, force=True)
            checks = packager.verify_pack_archive(archive_path)
            return ValidationResult(valid=True, checks=checks).model_dump()
        else:
            checks = packager.verify_pack_archive(pack_path)
            return ValidationResult(valid=True, checks=checks).model_dump()
    except Exception as e:
        return ValidationResult(valid=False, error=str(e)).model_dump()

# ── Pipeline logic ────────────────────────────────────────────────────

def run_pipeline(
    input_path: str,
    output_path: str,
    model: str = DEFAULT_MODEL,
    base_url: str = DEFAULT_BASE_URL,
    deterministic: bool = False
) -> dict:
    """Main pipeline for building a pack from input markdown."""
    print(f"[agent] pipeline start  input={input_path}  output={output_path}  model={model}", flush=True)
    # Ensure output directory exists
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)

    # 1. Initialize pack DB (D-12)
    db_path = output_path.replace(".oapack.zip", ".db")
    if db_path == output_path:
        db_path = output_path + ".db"
        
    pack_name = os.path.basename(output_path)
    if pack_name.endswith(".oapack.zip"):
        pack_name = pack_name[:-11]
    elif pack_name.endswith(".zip"):
        pack_name = pack_name[:-4]
    
    print(f"[agent] init pack  db={db_path}  name={pack_name}", flush=True)
    packager.init_pack(db_path, name=pack_name, force=True)

    if deterministic:
        for source_file, content in _read_markdown_inputs(input_path):
            result = extract_chunks(content, db_path, source_file, "en")
            if result["row_count"] > 0:
                provenance = {
                    "publisher": "Deterministic test fixture",
                    "license": "CC-BY-4.0",
                    "source_url": str(Path(input_path) / source_file),
                    "reviewed_at": packager._now(),
                    "language": "en",
                }
                derive_provenance(provenance["source_url"], db_path)
                conn = sqlite3.connect(db_path)
                try:
                    with conn:
                        conn.execute(
                            """UPDATE layers
                                  SET publisher = ?, license = ?, source_url = ?,
                                      reviewed_at = ?, language = ?
                                WHERE name = ?""",
                            (
                                provenance["publisher"],
                                provenance["license"],
                                provenance["source_url"],
                                provenance["reviewed_at"],
                                provenance["language"],
                                source_file,
                            ),
                        )
                finally:
                    conn.close()
                geo_points = []
                for match in re.finditer(
                    r'(\b\w+):\s*(?:lat:\s*)?(-?\d+\.\d+),\s*(?:lon:\s*)?(-?\d+\.\d+)',
                    content,
                ):
                    name = match.group(1).strip()
                    if name.lower() == "lat":
                        name = "Unnamed Point"
                    geo_points.append(
                        {
                            "name": name,
                            "category": "candidate",
                            "lat": float(match.group(2)),
                            "lon": float(match.group(3)),
                        }
                    )
                if geo_points:
                    packager.add_geo_points(db_path, layer=source_file, points=geo_points)
    else:
        # Step 2: Run PydanticAI Agent loop
        ollama_model = get_model(model_name=model, base_url=base_url)
        markdown_inputs = _read_markdown_inputs(input_path)
        print(f"[agent] loaded {len(markdown_inputs)} markdown file(s): {[f for f, _ in markdown_inputs]}", flush=True)
        for source_file, content in markdown_inputs:
            batches = _batch_markdown(content)
            n = len(batches)
            print(f"[agent] {source_file}  {len(content)} chars  {n} batch(es)", flush=True)
            # Provenance is set once per file, deterministically, using the filename as URL
            derive_provenance(f"file://{source_file}", db_path)

            for idx, batch in enumerate(batches):
                section_line = f"Section: {idx + 1} of {n}\n" if n > 1 else ""
                prompt = (
                    f"Build a knowledge pack layer from the extracted markdown below.\n"
                    f"Source file: '{source_file}'\n"
                    f"{section_line}"
                    f"Target database: '{db_path}'.\n"
                    "Call extract_chunks(text, pack_path, source_file, language) with the full content. "
                    "Call propose_geo_points if coordinates appear in the text. "
                    "Every tool call must pass pack_path explicitly.\n\n"
                    f"{batch}"
                )
                part_label = f" part {idx + 1}/{n}" if n > 1 else ""
                print(f"[agent] invoking compiler  {source_file}{part_label}  prompt_chars={len(prompt)}", flush=True)
                t0 = time.monotonic()
                compiler_agent.run_sync(prompt, model=ollama_model)
                print(f"[agent] done  {source_file}{part_label}  {time.monotonic() - t0:.1f}s", flush=True)

    # 3. Build index
    print(f"[agent] building FTS index  db={db_path}", flush=True)
    packager.build_index(db_path)

    # 4. Archive pack
    print(f"[agent] archiving  -> {output_path}", flush=True)
    archive_path = packager.archive_pack(db_path, output=output_path, force=True)

    # 5. Validate
    if deterministic:
        try:
            val = {"valid": True, "checks": packager.verify_pack_archive(archive_path)}
        except Exception as exc:
            val = {"valid": False, "error": str(exc)}
    else:
        val = validate_pack(archive_path)
    if not val["valid"]:
        print(f"[agent] validation FAILED: {val.get('error')}", flush=True)
        return {"status": "error", "error": val.get("error"), "pack_path": db_path}

    print(f"[agent] pipeline complete  archive={archive_path}", flush=True)
    return {
        "status": "success",
        "pack_path": db_path,
        "archive_path": archive_path,
        "validation": val
    }

# ── Deep Agents / LangChain Wrapper ───────────────────────────────────

@lc_tool
def run_compiler_pipeline(input_markdown_path: str, pack_path: str) -> str:
    """Runs the full compiler pipeline to transform markdown into a .db pack."""
    # Note: pack_path here is likely the target .oapack.zip or final destination
    res = run_pipeline(input_markdown_path, pack_path)
    return json.dumps(res)

@lc_tool
def emit_pack(pack_path: str) -> str:
    """Final tool to emit the completed knowledge pack archive."""
    # In this mock/agent flow, we just verify it exists
    if os.path.exists(pack_path):
        return f"Pack emitted successfully at {pack_path}"
    return f"Error: Pack not found at {pack_path}"

def create_pack_builder_agent():
    """Creates the outer Deep Agent wrapper."""
    # Load SKILL.md for instructions if available
    instructions = "You are a Pack Builder orchestrator. Use run_compiler_pipeline to build packs and emit_pack to finalize."
    skill_path = os.path.join(os.path.dirname(__file__), "SKILL.md")
    if os.path.exists(skill_path):
        with open(skill_path, "r") as f:
            instructions += "\n\n" + f.read()

    return create_deep_agent(
        tools=[run_compiler_pipeline, emit_pack],
        checkpointer=MemorySaver(),
        instructions=instructions,
        interrupt_config={"emit_pack": True}
    )

def start_agent(input_path: str, output_path: str):
    """Start a fresh Deep Agent run."""
    agent = create_pack_builder_agent()
    thread_id = f"fresh-{packager._now()}"
    config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 25}
    prompt = f"Run the compiler pipeline for {input_path} and emit the pack to {output_path}."
    res = agent.invoke({"messages": [("user", prompt)]}, config=config)
    return {"thread_id": thread_id, "result": str(res)}

def resume_agent(thread_id: str):
    """Resume an agent after HITL approval."""
    agent = create_pack_builder_agent()
    # D-10: LangGraph HITL resume shape
    command = Command(resume={"decisions": [{"type": "approve"}]})
    config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 25}
    return agent.invoke(command, config=config)

# ── CLI ───────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="OfflineAid Pack Builder Agent CLI")
    parser.add_argument("--input", "-i", required=True, help="Input markdown path")
    parser.add_argument("--output", "-o", required=True, help="Output .oapack.zip path")
    parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help=f"Compiler model (default: {DEFAULT_MODEL})")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL, help=f"Ollama base URL (default: {DEFAULT_BASE_URL})")
    parser.add_argument("--deterministic", action="store_true", help="Run without LLM for testing")
    parser.add_argument("--agent", action="store_true", help="Use Deep Agent wrapper instead of direct pipeline")
    parser.add_argument("--thread-id", help="Thread ID for resuming an existing run")
    parser.add_argument("--resume", action="store_true", help="Resume an existing run")

    args = parser.parse_args()

    if args.resume:
        if not args.thread_id:
            print("Error: --thread-id required for resume")
            sys.exit(1)
        res = resume_agent(args.thread_id)
        print(json.dumps(res, indent=2, default=str))
    elif args.agent:
        res = start_agent(args.input, args.output)
        print(json.dumps(res, indent=2))
    else:
        result = run_pipeline(
            input_path=args.input,
            output_path=args.output,
            model=args.model,
            base_url=args.base_url,
            deterministic=args.deterministic
        )
        print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()