agent_workflow.py

# ============================================================================
# agent_workflow.py — Workflow backend (fixed 2-step prompt chain)
# ============================================================================
"""Deterministic two-step prompt chain backend.

Purpose
-------
The simplest possible backend in the workbench: a hand-written, two-step
prompt chain (clarify -> answer). No tools, no agent loop, no framework.
The developer (not the model) decides what runs and in what order, so this
file is the baseline students can compare every other backend against.

Inputs
------
    api_key       : str  — provider API key (may be empty when env var is set).
    user_message  : str  — the researcher's chat message.
    provider      : str  — one of providers.LLM_PROVIDERS keys. Default Mistral.

Outputs
-------
    run(...) returns {"reply": str, "steps": list[dict], "extracted": dict}
    matching the workbench's universal backend contract.

Side effects
------------
    Two LLM completion calls per run() invocation. No disk I/O, no network
    state mutation beyond the LLM provider's own request log.

Contract
--------
    BACKEND_NAME, get_client, run, build_code_snippets.
"""

from __future__ import annotations

from typing import Any

from parameters import MAX_TOKENS, TEMPERATURE
from prompts import WORKFLOW_STEP1_CLARIFY, WORKFLOW_STEP2_ANSWER
import providers


BACKEND_NAME: str = "Workflow"


# ----------------------------------------------------------------
# PUBLIC: client factory
# ----------------------------------------------------------------
def get_client(api_key: str, provider: str = "Mistral") -> Any:
    """Build a provider-agnostic LLM client.

    All adapter logic lives in providers.py; this is a one-line forwarder so
    app.py never imports a specific SDK directly.

    Args:
        api_key:  Provider API key. Empty string means "fall back to env var".
        provider: Name from providers.LLM_PROVIDERS.

    Returns:
        An object exposing ``client.chat.complete(model, messages, ...)``
        whose response shape matches Mistral 1.x SDK.

    Raises:
        ValueError: when ``provider`` is not registered.

    Example:
        >>> client = get_client("", provider="Mistral")
    """
    return providers.get_llm_client(provider, api_key)


# ----------------------------------------------------------------
# INTERNAL: single LLM call, provider-agnostic
# ----------------------------------------------------------------
def _call_llm(client: Any, system_prompt: str, user_text: str, provider: str) -> str:
    """One round-trip to the LLM with a system+user pair.

    Args:
        client:        Object returned from get_client().
        system_prompt: System role message (instruction).
        user_text:     User role message (content to act on).
        provider:      Provider name; used to pick the default model.

    Returns:
        The assistant's text content; empty string if the model produced none.

    Raises:
        RuntimeError: if the provider client raises during completion. The
            original exception's message is included so the caller can log it.
    """
    try:
        response = client.chat.complete(
            model=providers.get_llm_model(provider),
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_text},
            ],
        )
    except Exception as exc:  # narrow: re-raise with workbench framing
        raise RuntimeError(
            f"Workflow LLM call failed (provider={provider}): {exc}"
        ) from exc

    return response.choices[0].message.content or ""


# ----------------------------------------------------------------
# Step pipeline — declared as data, executed as data.
# Each entry is (step_no, label, system_prompt). The body of run() is
# a list comprehension over this table — no if/elif, no per-step glue.
# ----------------------------------------------------------------
_STEPS: list[tuple[int, str, str]] = [
    (1, "clarify", WORKFLOW_STEP1_CLARIFY),
    (2, "answer", WORKFLOW_STEP2_ANSWER),
]


# ----------------------------------------------------------------
# PUBLIC: run
# ----------------------------------------------------------------
def run(client: Any, user_message: str, provider: str = "Mistral") -> dict[str, Any]:
    """Execute the fixed clarify -> answer chain.

    Each step's output becomes the next step's user input. The final step's
    output is the user-facing reply. Both steps are recorded in the step log
    so the UI can render them in the Process tab.

    Args:
        client:       From get_client().
        user_message: Researcher's chat message.
        provider:     Provider name; passed through to _call_llm.

    Returns:
        Dict with keys ``reply`` (str), ``steps`` (list), ``extracted`` (dict).

    Raises:
        RuntimeError: bubbled up from _call_llm if the LLM call fails.

    Example:
        >>> result = run(client, "what is grounded theory?")
        >>> result["reply"]
        'Grounded theory is...'
    """
    steps: list[dict[str, Any]] = []
    current_input: str = user_message

    # Declarative loop: walk the table, no per-step special-casing.
    for step_no, label, system_prompt in _STEPS:
        output = _call_llm(client, system_prompt, current_input, provider)
        steps.append({
            "step": step_no,
            "type": "llm_call",
            "tool": label,
            "args": current_input,
            "result": output,
        })
        current_input = output

    return {
        "reply": current_input,
        "steps": steps,
        "extracted": {"clarified_question": steps[0]["result"] if steps else ""},
    }


# ----------------------------------------------------------------
# PUBLIC: build_code_snippets
# ----------------------------------------------------------------
def build_code_snippets(user_message: str, steps: list[dict[str, Any]]) -> str:
    """Render an annotated code snippet of what this backend just did.

    Args:
        user_message: The original researcher message.
        steps:        The step log from run().

    Returns:
        A multi-line string the UI's Code tab displays verbatim.
    """
    header = [
        "# Backend: Workflow",
        "# Fixed 2-step prompt chain (clarify -> answer). No tools, no framework.",
        f"# User message: {user_message}",
        "",
        "step1 = client.chat.complete(",
        "    model=MODEL,",
        "    messages=[",
        "        {'role': 'system', 'content': WORKFLOW_STEP1_CLARIFY},",
        f"        {{'role': 'user', 'content': {user_message!r}}},",
        "    ],",
        ").choices[0].message",
        "clarified = step1.content",
        "",
        "step2 = client.chat.complete(",
        "    model=MODEL,",
        "    messages=[",
        "        {'role': 'system', 'content': WORKFLOW_STEP2_ANSWER},",
        "        {'role': 'user', 'content': clarified},",
        "    ],",
        ").choices[0].message",
        "answer = step2.content",
        "",
        "# ---------- actual step log ----------",
    ]
    log_lines = [
        line
        for s in steps
        for line in (
            f"# Step {s['step']} [{s['type']}] {s['tool']}",
            f"#   input:  {s['args']!r}",
            f"#   output: {s['result']!r}",
        )
    ]
    return "\n".join(header + log_lines)