dee/core/agent.py

"""Turing agent orchestrator — session/phase state + OpenRouter tool-calling loop.

Native Python rebuild of a pattern a former collaborator prototyped in Rust
(never deployed; that codebase is now abandoned). The OpenRouter request
shape and phase-FSM concept are ported from there, but auth and rate
limiting are NOT reimplemented in this module — that prototype's real gap
was that its LLM-calling route had neither. Routes built on top of this
module must go through the engine's existing dee/auth.py gate and be added
to server.py's _RL_RULES; this module has no opinion on who's allowed to
call it.

Tool-calling: run_agent_step accepts an `auth_anonymous` flag (a plain
bool, not the AuthContext object — dee/core/ stays auth-agnostic by
convention) and threads it through to dee/core/agent_tools.execute_tool,
which enforces the SAME sign-in requirement each tool's REST route already
enforces. See dee/core/agent_tools.py for which tools are wired and why
the rest (async DE library design, round-2 proposal, genome off-target
lookups) aren't yet.
"""

from __future__ import annotations

import json
import logging
import os
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import requests

from dee.core import agent_tools as _tools

logger = logging.getLogger("dee.agent")

PHASES = ("design", "build", "edit", "learn")


class AgentError(Exception):
    """Any agent-step failure worth surfacing to the caller as a clean message.

    kind mirrors the "kind" field the rest of the REST API already sends on
    error responses (e.g. "signin_required", "rate_limited") — lets the
    route pass it straight through instead of collapsing every AgentError
    into one generic value, so the frontend can render distinct cases
    (like "the model is temporarily unavailable") with their own copy."""

    def __init__(self, message: str, kind: str = "agent_error") -> None:
        super().__init__(message)
        self.kind = kind


# ─── Session state ──────────────────────────────────────────────────────
# Same in-memory + lock + TTL-sweep-on-write pattern as _SESSIONS in
# server.py — this process runs a single gthread worker, so in-process
# state is the established convention here, not a shortcut.

@dataclass
class AgentSession:
    session_id: str
    phase: str = "design"
    seed: int = 42
    created_at: float = field(default_factory=time.time)
    # Chat history as OpenRouter message dicts, system prompt excluded (that's
    # rebuilt fresh per step so phase/tool-list changes take effect immediately).
    history: List[Dict[str, Any]] = field(default_factory=list)


_SESSIONS: Dict[str, AgentSession] = {}
_SESSIONS_LOCK = threading.Lock()
_SESSION_TTL_SECONDS = 6 * 60 * 60  # matches server.py's _SESSION_TTL_SECONDS


def create_session() -> AgentSession:
    session = AgentSession(session_id=uuid.uuid4().hex)
    with _SESSIONS_LOCK:
        _SESSIONS[session.session_id] = session
        cutoff = time.time() - _SESSION_TTL_SECONDS
        stale = [k for k, v in _SESSIONS.items() if v.created_at < cutoff]
        for k in stale:
            _SESSIONS.pop(k, None)
    return session


def get_session(session_id: str) -> Optional[AgentSession]:
    with _SESSIONS_LOCK:
        return _SESSIONS.get(session_id)


def _save_session(session: AgentSession) -> None:
    with _SESSIONS_LOCK:
        _SESSIONS[session.session_id] = session


# ─── OpenRouter config ──────────────────────────────────────────────────

# google/gemini-3.5-flash's context window — verified against
# openrouter.ai/google/gemini-3.5-flash (2026-07-08, same check as the
# model slug above). Used only to compute a "% of context used" figure for
# the UI; there's no live per-model lookup here, so if AGENT_MODEL is ever
# changed to a different model, update this constant too.
CONTEXT_WINDOW_TOKENS = 1_000_000


@dataclass
class OpenRouterConfig:
    api_key: str
    model: str
    max_steps: int
    max_cost_usd: float

    @classmethod
    def from_env(cls) -> Optional["OpenRouterConfig"]:
        # .strip() — a copy-pasted secret with a trailing newline produced an
        # invalid HTTP header ("Bearer <key>\n\n"), which the underlying HTTP
        # client rejected with an exception whose message embedded the raw
        # header value. That exception text then reached the client via the
        # /api/agent/step error response — a real key leak, live, confirmed
        # in production. Stripping here prevents the malformed-header case
        # entirely; run_agent_step below no longer echoes exception text to
        # the client either way, as defense in depth.
        api_key = (os.environ.get("OPENROUTER_API_KEY") or "").strip()
        if not api_key:
            return None
        return cls(
            api_key=api_key,
            # Verified against openrouter.ai/google/gemini-3.5-flash directly
            # (2026-07-08) — don't change this without re-checking the slug,
            # Google's Gemini naming moves fast and a stale slug just 400s.
            model=os.environ.get("AGENT_MODEL", "google/gemini-3.5-flash"),
            max_steps=int(os.environ.get("AGENT_MAX_STEPS", "8")),
            max_cost_usd=float(os.environ.get("AGENT_MAX_COST_USD", "0.5")),
        )


_SYSTEM_PROMPT_TEMPLATE = (
    "You are Turing, the conversational orchestrator for TuringDNA, a "
    "directed-evolution workbench (design variant libraries with ESM-2, "
    "build/map plasmids, edit with CRISPR, learn from bench results). "
    "Current loop phase: {phase}. You have four tools available right now: "
    "fetch_sequence, design_crispr_guides, design_primers, and "
    "design_variant_library. If a request could be answered by calling one "
    "of these, call it — do not describe what you would do, and do not "
    "answer a design/analysis question from your own knowledge instead of "
    "running the real tool. If the user names a gene, protein, or "
    "accession instead of pasting a sequence (e.g. \"human GFP\", "
    "\"NM_001301717\"), call fetch_sequence first to resolve it, then feed "
    "the sequence it returns into whichever design tool the request "
    "actually needs — in the same reply, without asking the user to go "
    "paste it themselves. Only ask the user to paste a sequence directly "
    "if fetch_sequence fails or the input doesn't look like a resolvable "
    "name/accession. This also applies across turns, not just within one: "
    "if you already fetched or were given a sequence earlier in THIS "
    "conversation and the user refers back to it (\"it\", \"that gene\", "
    "\"the sequence you found\", or just a follow-up request with no new "
    "sequence attached), reuse the one already in the conversation instead "
    "of asking for it again — check your own conversation history before "
    "concluding you don't have something. design_variant_library accepts "
    "either a protein or a DNA coding sequence (it auto-translates), so "
    "the output of fetch_sequence can be passed straight through with no "
    "separate translation step. design_variant_library is a chat-sized version of the Directed "
    "Evolution tool (capped protein length and search budget so it fits "
    "in one reply); if a request clearly needs the full tool — a longer "
    "protein, or a deeper search — say so and point the user to the "
    "Directed Evolution tool in the sidebar instead of running a "
    "truncated version silently. Round-2 proposals from bench results "
    "aren't wired up as a tool yet; if asked for that, say so plainly "
    "instead of fabricating a result, and point the user to the Directed "
    "Evolution tool.\n\n"
    "This is a scientific tool feeding real lab decisions — a wrong "
    "number or an invented fact has real consequences, so accuracy "
    "matters more than sounding confident or complete:\n"
    "- Never invent sequences, scores, or variant/guide/primer results a "
    "tool didn't actually return. Only state numbers/values that appear "
    "verbatim in a tool's result.\n"
    "- For molecular-biology facts NOT covered by a tool result (e.g. "
    "explaining what a CFD score means, why ΔLL is used, general CRISPR "
    "mechanism) — if you are not confident a specific claim is correct, "
    "say so explicitly rather than guessing. A brief, honest 'I'm not "
    "certain about X' beats a fluent wrong answer here.\n"
    "- Don't pad a short factual answer with unrequested elaboration — "
    "every additional claim is another chance to be wrong.\n\n"
    "Formatting: this chat renders plain text with light markdown only — "
    "paragraphs, `inline code`, **bold**, '# ' headings, and '* ' bullet "
    "lists all display correctly. It does NOT render LaTeX: never use $...$ "
    "math delimiters or LaTeX macros like \\rightarrow — write arrows as "
    "the plain character → directly in the text. Put each DNA/protein "
    "sequence on its own line (plain, or inside a ```text fenced block) "
    "rather than folding it into a sentence — the interface collapses a "
    "long sequence into a clean expandable block automatically, but only "
    "when it isn't tangled up with surrounding prose.\n\n"
    "Company facts (the only ones you're grounded on — do not add to this "
    "list from your own training data, it has no reliable information "
    "about a company this small and specific): TuringDNA was created by "
    "Tengo Gzirishvili. For anything about the company, team, funding, or "
    "history beyond that one fact — including any other name — say you "
    "don't have verified information about it rather than guessing a "
    "name or detail. A wrong guess here (e.g. inventing a founder's name) "
    "is a real credibility failure, not a harmless one."
)


# ─── OpenRouter chat-completions ────────────────────────────────────────

def _call_openrouter(config: OpenRouterConfig, body: Dict[str, Any]) -> Tuple[Dict[str, Any], float, int]:
    """POST one chat-completions request. Returns (message_dict, cost_usd, prompt_tokens).

    Every exception below is logged in FULL server-side and surfaced to the
    caller as a fixed, generic AgentError message only — never str(exc).
    A real incident (2026-07-08) is why: a malformed API key (trailing
    newline from a copy-paste) made the underlying HTTP client raise an
    "invalid header value" error whose message embeds the offending header
    VERBATIM — i.e. the bearer token. That exception's str() reached the
    client through this route's error response, live, in production. Never
    again: config.api_key itself must also never appear in a raised
    AgentError message, even indirectly."""
    try:
        resp = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {config.api_key}",
                "Content-Type": "application/json",
            },
            json=body,
            timeout=30,
        )
    except Exception as exc:  # noqa: BLE001 — deliberately broad, see comment above
        logger.exception("OpenRouter request failed")
        raise AgentError("could not reach OpenRouter") from exc

    if resp.status_code == 402:
        # OpenRouter's status specifically for "account/key is out of
        # credits" (distinct from 401 auth / 403 permissions / 429 rate
        # limit) — https://openrouter.ai/docs, confirmed 2026-07-10. Never
        # surface billing language to end users: "insufficient credits"
        # reads as "this product ran out of money," which is alarming and
        # none of their business. "Still in development" is equally true
        # (this genuinely is pre-launch) without the financial detail.
        logger.warning("OpenRouter out of credits (402): %s", resp.text[:500])
        raise AgentError(
            "Turing is still in active development and is temporarily "
            "unavailable — check back shortly.",
            kind="agent_unavailable",
        )

    if resp.status_code != 200:
        # OpenRouter's error body is small and safe to log; never echo it
        # verbatim to the client (could contain the key's own account info).
        logger.warning("OpenRouter error %s: %s", resp.status_code, resp.text[:500])
        raise AgentError(f"OpenRouter error (status {resp.status_code})")

    try:
        parsed = resp.json()
        choices = parsed.get("choices") or []
        if not choices:
            raise AgentError("empty response from OpenRouter")
        message = choices[0].get("message") or {}
    except AgentError:
        raise
    except Exception as exc:  # noqa: BLE001 — malformed upstream JSON, etc.
        logger.exception("failed to parse OpenRouter response")
        raise AgentError("could not parse OpenRouter's response") from exc

    usage = parsed.get("usage") or {}
    cost = usage.get("total_cost") or usage.get("cost") or 0.0
    # prompt_tokens on a chat-completions response is the size of the WHOLE
    # messages array just sent — already includes session.history, so this
    # naturally grows call over call as a conversation gets longer. Used to
    # show users how much of the model's context window they've used.
    prompt_tokens = int(usage.get("prompt_tokens") or 0)
    return message, cost, prompt_tokens


def run_agent_step(
    config: OpenRouterConfig,
    session: AgentSession,
    user_message: str,
    auth_anonymous: bool,
) -> Dict[str, Any]:
    """Run one user turn, including any tool calls the model makes along the
    way. Bounded to config.max_steps OpenRouter round-trips (a tool call +
    its follow-up reply is 2 — plenty of headroom for the two tools wired up
    today; see agent_tools.py). auth_anonymous is checked per-tool-call, not
    once up front, so a signed-out user can still chat freely and only hits
    the sign-in wall if they actually ask for a gated tool."""
    system = _SYSTEM_PROMPT_TEMPLATE.format(phase=session.phase)
    turn_messages: List[Dict[str, Any]] = [{"role": "user", "content": user_message}]
    total_cost = 0.0
    tool_result: Optional[Dict[str, Any]] = None
    tool_name: Optional[str] = None
    content = ""
    # The LAST call's prompt_tokens is what matters for "how full is the
    # context window right now" — it already reflects everything sent this
    # turn (history + any tool-call/tool-result messages appended so far).
    prompt_tokens = 0

    for _ in range(config.max_steps):
        body = {
            "model": config.model,
            "messages": [{"role": "system", "content": system}] + session.history + turn_messages,
            "tools": _tools.TOOL_SPECS,
            "temperature": 0,
            "top_p": 1,
            "parallel_tool_calls": False,
        }
        message, cost, prompt_tokens = _call_openrouter(config, body)
        total_cost += cost
        if total_cost > config.max_cost_usd:
            # Already spent — this caps the SESSION from continuing, not the
            # money already charged for calls made so far. Real hard-stop is
            # the OpenRouter account credit limit (set at the provider).
            logger.warning("agent session %s exceeded max_cost_usd (%.4f > %.2f)",
                            session.session_id, total_cost, config.max_cost_usd)
            raise AgentError("cost cap exceeded for this session")

        tool_calls = message.get("tool_calls") or []
        if not tool_calls:
            content = message.get("content") or ""
            turn_messages.append({"role": "assistant", "content": content})
            break

        turn_messages.append({
            "role": "assistant",
            "content": message.get("content"),
            "tool_calls": tool_calls,
        })
        # parallel_tool_calls=False above asks the model for one at a time;
        # only act on the first entry either way.
        tc = tool_calls[0]
        fn = tc.get("function") or {}
        name = fn.get("name") or ""
        try:
            args = json.loads(fn.get("arguments") or "{}")
        except (TypeError, ValueError):
            args = {}
        result = _tools.execute_tool(name, args, auth_anonymous)
        tool_result = result
        tool_name = name
        turn_messages.append({
            "role": "tool",
            "tool_call_id": tc.get("id") or "",
            # Bounded so one oversized tool result can't blow up the cost or
            # context of the follow-up call — mirrors the 4000-char cap
            # server.py already puts on the incoming user message.
            "content": json.dumps(result)[:8000],
        })
    else:
        content = "I ran out of steps working through that — try asking in smaller pieces."

    session.history.extend(turn_messages)
    _save_session(session)

    return {
        "ok": True,
        "phase": session.phase,
        "assistant_message": content,
        "tool_result": tool_result,
        "tool_name": tool_name,
        "cost_usd": round(total_cost, 6),
        "context_tokens_used": prompt_tokens,
        "context_tokens_limit": CONTEXT_WINDOW_TOKENS,
    }