src/matrix_context/adapters/agent_generator.py

"""agent-generator integration — emit a Matrix Context memory layer.

When `agent-generator` is invoked with `--context-provider matrix-context`, it
calls :func:`emit_template` to obtain the files that wire Matrix Context into the
generated project. Two variants are supported, selected by the ``mcp`` flag:

* **in-process** (default): a local :class:`~matrix_context.ContextManager` backed
  by a default SQLite store, with ``build_pack`` called *before* each model call
  and ``remember`` called *after* each turn, so the generated agent actually
  accumulates and uses memory.
* **MCP** (``mcp=True``): instead of an in-process client, emit an MCP server
  launch configuration pointing at ``matrix-context serve --transport stdio`` so
  the agent reaches the same engine over the protocol.

The emitter is framework-aware for the ``crewai``, ``langgraph`` and ``react``
targets — the core ``ContextManager`` wiring is identical, only the call-site
shape (memory hook, graph node, plain tool) differs.

This is the proof that the engine is usable from generated code: the emitted
client must import :class:`ContextManager`, reference a SQLite path, and call
``build_pack``. The unit test under ``tests/unit/test_agent_generator.py``
asserts exactly that.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import Dict, Optional

FRAMEWORKS = ("react", "crewai", "langgraph")

# How the engine should be reached from generated code.
IN_PROCESS = "in_process"
MCP = "mcp"


@dataclass
class EmittedTemplate:
    """The result of :func:`emit_template`.

    ``files`` maps a relative path to file content; ``entrypoint`` names the
    primary client module so the caller can wire imports. ``code`` / ``config``
    are convenience views over the primary client and the MCP/launch config.
    """

    framework: str
    variant: str  # IN_PROCESS | MCP
    slug: str
    scopes: Dict[str, str]
    files: Dict[str, str] = field(default_factory=dict)
    entrypoint: str = "matrix_memory.py"

    @property
    def code(self) -> str:
        return self.files.get(self.entrypoint, "")

    @property
    def config(self) -> str:
        return self.files.get("mcp.json", "")


def _slug(text: str, fallback: str = "agent") -> str:
    s = re.sub(r"[^a-z0-9]+", "-", (text or "").lower()).strip("-")
    s = "-".join(s.split("-")[:4])  # keep it short
    return s or fallback


def _default_scopes(slug: str, purpose: str) -> Dict[str, str]:
    """Example scopes appropriate to the agent's purpose.

    Profile is always-injectable identity; semantic/episodic are recalled by
    routing. A policy scope is added when the purpose hints at governance.
    """
    base = f"/{slug}"
    scopes = {
        "profile": f"{base}/profile",
        "semantic": f"{base}/knowledge",
        "episodic": f"{base}/history",
    }
    if re.search(r"govern|policy|complian|audit|secure", purpose or "", re.I):
        scopes["policy"] = f"{base}/policy"
    return scopes


# --------------------------------------------------------------------------- #
# In-process client (shared core + framework-specific call sites)
# --------------------------------------------------------------------------- #
def _client_module(slug: str, purpose: str, scopes: Dict[str, str],
                   store_path: str, max_tokens: int, framework: str) -> str:
    seed = []
    for expert, scope in scopes.items():
        if expert == "profile":
            seed.append(
                f'    ctx.remember("purpose: {purpose}", expert="profile", '
                f'scope=SCOPES["profile"], importance=0.9)')
    seed_block = "\n".join(seed) or "    pass"

    framework_note = {
        "crewai": "Wire `build_context` into a CrewAI Task's context and call "
                  "`record_turn` from a step/`task_callback`.",
        "langgraph": "Use `build_context` inside a node before the model call "
                     "and `record_turn` in the node that closes the turn.",
        "react": "Expose `build_context` and `record_turn` as plain tools the "
                 "ReAct loop can call.",
    }[framework]

    scopes_literal = json.dumps(scopes, indent=4).replace("null", "None")
    return f'''"""Matrix Context memory layer for `{slug}` ({framework}).

Generated by agent-generator (--context-provider matrix-context). {framework_note}

The two calls that matter:
  * `build_context(query)` -> run BEFORE each model call (routes + budgets memory)
  * `record_turn(user, agent)` -> run AFTER each turn (remember what happened)
"""
from __future__ import annotations

from matrix_context import ContextManager

# Default local SQLite store — the source of truth, vectors are an accelerator.
STORE_PATH = "{store_path}"
MAX_TOKENS = {max_tokens}

# Example scopes appropriate to this agent's purpose.
SCOPES = {scopes_literal}

ctx = ContextManager.create("{slug}", path=STORE_PATH)


def bootstrap() -> None:
    """Seed always-injectable profile facts (idempotent enough for a demo)."""
{seed_block}


def build_context(query: str, max_tokens: int = MAX_TOKENS) -> str:
    """Route + retrieve + budget memory into a compact prompt block.

    Call this BEFORE every model call and prepend the result to the prompt.
    """
    pack = ctx.build_pack(query, max_tokens=max_tokens)
    return pack.to_prompt()


def record_turn(user_message: str, agent_message: str,
                importance: float = 0.5) -> None:
    """Remember the turn AFTER it happens, so memory accumulates across calls."""
    ctx.remember(user_message, expert="episodic",
                 scope=SCOPES["episodic"], importance=importance)
    ctx.remember(agent_message, expert="semantic",
                 scope=SCOPES["semantic"], importance=importance)


def explain(query: str) -> str:
    """Inspect why the engine selected what it did (every choice is explainable)."""
    return ctx.inspect(query, max_tokens=MAX_TOKENS)


if __name__ == "__main__":
    bootstrap()
    record_turn("I prefer concise answers.", "Understood — I will be concise.")
    print(build_context("what does the user prefer?"))
'''


# --------------------------------------------------------------------------- #
# MCP variant (launch config + thin client)
# --------------------------------------------------------------------------- #
def _mcp_config(slug: str, store_path: str) -> str:
    cfg = {
        "mcpServers": {
            "matrix-context": {
                "command": "matrix-context",
                "args": ["serve", "--transport", "stdio"],
                "env": {
                    "MATRIX_CONTEXT_NAME": slug,
                    "MATRIX_CONTEXT_PATH": store_path,
                },
            }
        }
    }
    return json.dumps(cfg, indent=2)


def _mcp_client_module(slug: str, purpose: str, scopes: Dict[str, str],
                       store_path: str, max_tokens: int, framework: str) -> str:
    scopes_literal = json.dumps(scopes, indent=4).replace("null", "None")
    return f'''"""Matrix Context (MCP) memory layer for `{slug}` ({framework}).

Generated by agent-generator (--context-provider matrix-context --mcp).

This variant does NOT embed the engine in-process. It launches the standards
compliant server via `matrix-context serve --transport stdio` (see mcp.json) and
talks to it over MCP. The two tools that matter are `build_pack` (before a model
call) and `remember` (after a turn). A local fallback `ContextManager` keeps the
generated project runnable offline before the MCP host is attached.
"""
from __future__ import annotations

from matrix_context import ContextManager

# The MCP server is configured in mcp.json -> `matrix-context serve --transport stdio`.
STORE_PATH = "{store_path}"
MAX_TOKENS = {max_tokens}
SCOPES = {scopes_literal}

# Offline fallback so the project runs before an MCP host wires the server in.
_local = ContextManager.create("{slug}", path=STORE_PATH)


def build_context(query: str, max_tokens: int = MAX_TOKENS) -> str:
    """Before each model call: ask the MCP `build_pack` tool (local fallback here)."""
    return _local.build_pack(query, max_tokens=max_tokens).to_prompt()


def record_turn(user_message: str, agent_message: str) -> None:
    """After each turn: call the MCP `remember` tool (local fallback here)."""
    _local.remember(user_message, expert="episodic", scope=SCOPES["episodic"])
    _local.remember(agent_message, expert="semantic", scope=SCOPES["semantic"])
'''


def _readme(slug: str, variant: str, framework: str) -> str:
    how = ("Launch the MCP server with `matrix-context serve --transport stdio` "
           "(configured in `mcp.json`)."
           if variant == MCP else
           "The memory layer runs in-process against a local SQLite store.")
    return (f"# {slug} — Matrix Context memory\n\n"
            f"Framework: **{framework}**  ·  Variant: **{variant}**\n\n{how}\n\n"
            "- `build_context(query)` before every model call\n"
            "- `record_turn(user, agent)` after every turn\n")


def emit_template(purpose: str = "", framework: str = "react", *,
                  mcp: bool = False, scopes: Optional[Dict[str, str]] = None,
                  store_path: Optional[str] = None, name: Optional[str] = None,
                  max_tokens: int = 256) -> EmittedTemplate:
    """Emit the Matrix Context client code + config for a generated project.

    Parameters
    ----------
    purpose:    natural-language description of the agent (drives example scopes).
    framework:  one of ``react`` | ``crewai`` | ``langgraph``.
    mcp:        emit the MCP server launch config instead of an in-process client.
    scopes:     override the example scopes (expert -> scope path).
    store_path: SQLite path for the default local store.
    name:       project/agent name (defaults to a slug of ``purpose``).
    max_tokens: per-turn pack budget (compact-injection discipline).
    """
    framework = (framework or "react").lower()
    if framework not in FRAMEWORKS:
        raise ValueError(f"unknown framework: {framework!r}; "
                         f"expected one of {FRAMEWORKS}")
    slug = _slug(name or purpose)
    scopes = scopes or _default_scopes(slug, purpose)
    store_path = store_path or f"./{slug}.matrix-context.db"
    variant = MCP if mcp else IN_PROCESS

    files: Dict[str, str] = {}
    if mcp:
        files["matrix_memory.py"] = _mcp_client_module(
            slug, purpose, scopes, store_path, max_tokens, framework)
        files["mcp.json"] = _mcp_config(slug, store_path)
    else:
        files["matrix_memory.py"] = _client_module(
            slug, purpose, scopes, store_path, max_tokens, framework)
    files["MATRIX_CONTEXT.md"] = _readme(slug, variant, framework)

    return EmittedTemplate(framework=framework, variant=variant, slug=slug,
                           scopes=scopes, files=files)