_proposer.py

"""Schema proposal & refinement engine.

Two responsibilities:

1. **propose**: given a few sample inputs + (optional) an LLM, suggest an
   initial Schema. Three flavours:
    - `propose_from_rule(meta)`: rule extractor → hard-coded preset
    - `propose_generic()`:       no extractor / no LLM → minimal Entity/Relation
    - `propose_with_llm(...)`:   LLM looks at samples, returns Schema

2. **refine**: turn a free-form instruction into a list of structured
   `SchemaEdit` ops, applied to a live Schema by `apply_edit()`.
   Two paths: regex first (fast, deterministic), LLM fallback (free-form).

This module never touches Session state directly — Session calls the
functions and consumes their return values.
"""
from __future__ import annotations

import json
import re
from typing import Any, Callable

from .schema import EntityType, RelationType, Schema


# ---------------------------------------------------------------------------
# Propose
# ---------------------------------------------------------------------------


def propose_generic() -> Schema:
    """Bare-bones starter Schema for unfamiliar artifacts.

    Two entity types and two relations — enough to bootstrap; the user
    refines from here.
    """
    s = Schema(name="generic", description="Generic Entity/Relation starter.")
    s.add_entity("Entity", description="Anything worth a node.")
    s.add_entity("Concept", description="An abstract idea referenced by Entities.")
    s.add_relation("relates_to", description="Generic association.",
                   domain=["Entity", "Concept"], range=["Entity", "Concept"])
    s.add_relation("mentions", description="Source mentions a target.",
                   domain=["Entity"], range=["Concept"])
    return s


def propose_from_rule(meta: Any) -> Schema:
    """Pick a built-in preset that matches a rule extractor."""
    from .schema import load_preset

    name_to_preset = {
        "ast": "codebase",
        "kg_code": "papers",
        "markdown": "obsidian-vault",
        "openapi": "openapi",
        "fstree": "fstree",
        "chatlog": "chat-log",
    }
    preset_name = name_to_preset.get(meta.name)
    if preset_name:
        try:
            return load_preset(preset_name)
        except FileNotFoundError:
            pass
    return propose_generic()


_META_PROMPT = """You are a knowledge-graph schema designer.

I will paste short excerpts from {n} sample artifacts. Suggest an initial
schema with 3-8 entity types and 3-10 relation types that covers most of
what you see. Prefer specific types over generic ones.

Output strict JSON, no prose, no fences:

{{
  "name": "<kebab-case-name>",
  "description": "<one sentence>",
  "entities": [
    {{"name": "Foo", "description": "...", "attrs": ["bar", "baz"]}}
  ],
  "relations": [
    {{"name": "uses", "description": "...", "domain": ["Foo"], "range": ["Bar"]}}
  ]
}}

SAMPLES:
{samples}
"""


def propose_with_llm(
    sample_paths: list[str],
    *,
    llm: Any,
    on_chunk: Callable[[dict], None] | None = None,
    max_chars_per_sample: int = 4_000,
) -> Schema:
    """Ask an LLM to propose a Schema from sample artifacts."""
    from pathlib import Path

    excerpts = []
    for p in sample_paths:
        path = Path(p)
        if not path.exists() or not path.is_file():
            continue
        try:
            text = path.read_text(errors="replace")[:max_chars_per_sample]
        except Exception:
            continue
        excerpts.append(f"--- {path.name} ---\n{text}\n")

    if not excerpts:
        return propose_generic()

    prompt = _META_PROMPT.format(n=len(excerpts), samples="\n".join(excerpts))

    raw = _llm_call(llm, prompt, on_chunk=on_chunk)
    try:
        data = _coerce_json(raw)
    except Exception:
        return propose_generic()
    if not isinstance(data, dict):
        return propose_generic()
    return Schema.from_dict(data)


# ---------------------------------------------------------------------------
# Refine
# ---------------------------------------------------------------------------


# An "edit" is a structured instruction applicable to a Schema.
SchemaEdit = dict[str, Any]


_RE_ADD_ENTITY = re.compile(
    r"^\s*add\s+"
    r"(?:entity\s+([A-Za-z_][A-Za-z0-9_]*)|"           # add entity Foo
    r"([A-Za-z_][A-Za-z0-9_]*)\s+entity|"              # add Foo entity
    r"([A-Za-z_][A-Za-z0-9_]*))"                       # add Foo
    r"\s*(?::\s*(.+?))?\s*$",
    re.IGNORECASE,
)
_RE_RM_ENTITY = re.compile(
    r"^\s*(?:remove|drop|delete)\s+(?:entity\s+)?([A-Za-z_][A-Za-z0-9_]*)\s*$",
    re.IGNORECASE,
)
_RE_ADD_RELATION = re.compile(
    r"^\s*add\s+relation\s+([A-Za-z_][A-Za-z0-9_]*)"
    r"(?:\s+from\s+([A-Za-z_][A-Za-z0-9_, ]*?))?"
    r"(?:\s+to\s+([A-Za-z_][A-Za-z0-9_, ]*?))?"
    r"\s*(?::\s*(.+?))?\s*$",
    re.IGNORECASE,
)
_RE_RM_RELATION = re.compile(
    r"^\s*(?:remove|drop|delete)\s+relation\s+([A-Za-z_][A-Za-z0-9_]*)\s*$",
    re.IGNORECASE,
)
_RE_RENAME_ENTITY = re.compile(
    r"^\s*rename\s+entity\s+([A-Za-z_][A-Za-z0-9_]*)\s+(?:to|→)\s+([A-Za-z_][A-Za-z0-9_]*)\s*$",
    re.IGNORECASE,
)


def parse_refine_instruction(instruction: str) -> list[SchemaEdit]:
    """Regex pass over deterministic forms. Returns [] if nothing matched."""
    edits: list[SchemaEdit] = []

    for line in instruction.splitlines():
        line = line.strip().rstrip(".")
        if not line:
            continue

        m = _RE_ADD_ENTITY.match(line)
        if m and not _RE_ADD_RELATION.match(line):
            # _RE_ADD_ENTITY has 3 alternative name capture groups (1, 2, 3);
            # exactly one will be non-None.
            name = m.group(1) or m.group(2) or m.group(3)
            edits.append({
                "op": "add_entity",
                "name": name,
                "description": (m.group(4) or "").strip(),
            })
            continue

        m = _RE_RM_ENTITY.match(line)
        if m:
            edits.append({"op": "remove_entity", "name": m.group(1)})
            continue

        m = _RE_ADD_RELATION.match(line)
        if m:
            edits.append({
                "op": "add_relation",
                "name": m.group(1),
                "domain": _split(m.group(2)),
                "range": _split(m.group(3)),
                "description": (m.group(4) or "").strip(),
            })
            continue

        m = _RE_RM_RELATION.match(line)
        if m:
            edits.append({"op": "remove_relation", "name": m.group(1)})
            continue

        m = _RE_RENAME_ENTITY.match(line)
        if m:
            edits.append({
                "op": "rename_entity",
                "old": m.group(1),
                "new": m.group(2),
            })
            continue

    return edits


def _split(s: str | None) -> list[str]:
    if not s:
        return []
    return [t.strip() for t in s.split(",") if t.strip()]


_REFINE_PROMPT = """You translate one freeform schema instruction into a
list of structured edits to apply.

CURRENT SCHEMA:
{schema_summary}

INSTRUCTION:
{instruction}

OUTPUT (strict JSON list, no prose, no fences). Each edit is one of:
  {{"op": "add_entity",     "name": "...", "description": "...", "attrs": []}}
  {{"op": "remove_entity",  "name": "..."}}
  {{"op": "rename_entity",  "old": "...", "new": "..."}}
  {{"op": "add_relation",   "name": "...", "domain": [...], "range": [...]}}
  {{"op": "remove_relation","name": "..."}}

Only emit edits that the instruction *clearly* implies. If unsure, omit.
"""


def parse_refine_with_llm(
    instruction: str,
    schema: Schema,
    *,
    llm: Any,
    on_chunk: Callable[[dict], None] | None = None,
) -> list[SchemaEdit]:
    """LLM fallback when regex didn't match."""
    schema_summary = json.dumps({
        "entities": [e.name for e in schema.entities],
        "relations": [{"name": r.name, "domain": r.domain, "range": r.range} for r in schema.relations],
    }, ensure_ascii=False)

    prompt = _REFINE_PROMPT.format(schema_summary=schema_summary, instruction=instruction)
    raw = _llm_call(llm, prompt, on_chunk=on_chunk)
    try:
        data = _coerce_json(raw)
    except Exception:
        return []

    if isinstance(data, dict) and "edits" in data:
        data = data["edits"]
    if not isinstance(data, list):
        return []
    return [e for e in data if isinstance(e, dict) and "op" in e]


def apply_edit(schema: Schema, edit: SchemaEdit) -> bool:
    """Apply one edit to a live Schema. Returns True if anything changed."""
    op = edit.get("op")
    if op == "add_entity":
        before = len(schema.entities)
        schema.add_entity(
            edit["name"],
            description=edit.get("description", "") or "",
            attrs=edit.get("attrs") or [],
        )
        return len(schema.entities) > before
    if op == "remove_entity":
        return schema.remove_entity(edit["name"])
    if op == "rename_entity":
        old = edit["old"]
        new = edit["new"]
        for e in schema.entities:
            if e.name == old:
                e.name = new
                schema.version += 1
                return True
        return False
    if op == "add_relation":
        before = len(schema.relations)
        schema.add_relation(
            edit["name"],
            description=edit.get("description", "") or "",
            domain=edit.get("domain") or [],
            range=edit.get("range") or [],
        )
        return len(schema.relations) > before
    if op == "remove_relation":
        return schema.remove_relation(edit["name"])
    return False


# ---------------------------------------------------------------------------
# LLM call helper (mirrors extractors/llm_entity._call_llm logic)
# ---------------------------------------------------------------------------


def _llm_call(llm: Any, prompt: str, *, on_chunk) -> str:
    """Lightweight LLM-call adapter, same shape as the extractor's."""
    messages = [{"role": "user", "content": prompt}]
    if hasattr(llm, "chat_json"):
        for arg in (prompt, messages):
            try:
                data = llm.chat_json(arg)
            except TypeError:
                continue
            except Exception:
                break
            _record_chunk(on_chunk, data if isinstance(data, dict) else {"content": data}, prompt)
            return json.dumps(data) if isinstance(data, (dict, list)) else str(data)
    if hasattr(llm, "chat"):
        for arg in (prompt, messages):
            try:
                out = llm.chat(arg)
            except TypeError:
                continue
            text = out.get("content") if isinstance(out, dict) else out
            _record_chunk(on_chunk, out if isinstance(out, dict) else {"content": text}, prompt)
            return text or ""
    raise TypeError("LLM client must expose .chat or .chat_json")


def _record_chunk(on_chunk, payload, prompt):
    if on_chunk is None:
        return
    if not isinstance(payload, dict):
        on_chunk({"prompt_chars": len(prompt)})
        return
    usage = payload.get("usage") or {}
    on_chunk({
        "tokens_in": int(usage.get("prompt_tokens") or 0),
        "tokens_out": int(usage.get("completion_tokens") or 0),
        "dollars": float(usage.get("estimated_cost") or 0.0),
        "prompt_chars": len(prompt),
    })


def _coerce_json(text: str) -> Any:
    text = (text or "").strip()
    if text.startswith("```"):
        text = re.sub(r"^```(?:json)?\s*", "", text, count=1)
        text = re.sub(r"\s*```$", "", text, count=1).strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        for opener, closer in (("{", "}"), ("[", "]")):
            i, j = text.find(opener), text.rfind(closer)
            if 0 <= i < j:
                try:
                    return json.loads(text[i : j + 1])
                except json.JSONDecodeError:
                    continue
        raise