agent/agent_client.py

import logging
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable

from agent.agent import Agent, AgentResponse
from agent.agent_exceptions import AgentException
from agent.conversation_history import ConversationHistory
from agent.documents import build_documents_block
from helpers.impacts_tracker_helper import get_champ_impacts
from helpers.language import (
    TRANSLATION_CORRECTION_PROMPT,
    detect_language,
    find_leakage,
)


logger = logging.getLogger(__name__)


@dataclass
class EnvImpact:
    gwp_kgcoeq: float = 0.0
    water_L: float = 0.0
    electricity_kWh: float = 0.0


@dataclass
class ChatOutcome:
    reply: str
    reply_id: str
    success: bool
    env_impact: EnvImpact = field(default_factory=EnvImpact)
    n_tokens: int = 0
    context: list = field(default_factory=list)
    triage_meta: dict = field(default_factory=dict)
    error: str | None = None
    error_type: str | None = None
    inference_impacts: Any = None  # raw EcoLogits Impacts for env logging
    # Compact per-turn observability trace: tool calls with reasoning,
    # sub-agent pipeline steps (recalled pages, drafts, judge/attribution
    # verdicts), translation-correction passes. Logged to DDB; never user-facing.
    trace: list = field(default_factory=list)


LoggerFn = Callable[["ChatOutcome"], None]

# Long free-text fields in the trace are clipped to keep DDB items well under
# the 400 KB item limit. Full recalled pages still land in `context`.
_TRACE_CLIP = 1500


def _clip(text: str | None, limit: int = _TRACE_CLIP) -> str | None:
    if text is None or len(text) <= limit:
        return text
    return text[:limit] + f"…[+{len(text) - limit} chars]"


def _avg(v) -> float:
    """RangeValue (.min/.max) → midpoint; plain number → float."""
    if hasattr(v, "min") and hasattr(v, "max"):
        return (v.min + v.max) / 2
    return float(v)


def _to_env_impact(impacts: Any) -> EnvImpact:
    if impacts is None:
        return EnvImpact()
    return EnvImpact(
        gwp_kgcoeq=_avg(impacts.usage.gwp.value),
        water_L=_avg(impacts.usage.wcf.value),
        electricity_kWh=_avg(impacts.usage.energy.value),
    )


class AgentClient:
    """One conversation, one model. Owns ConversationHistory and the error boundary."""

    def __init__(
        self, agent: Agent | None = None, logger_fn: LoggerFn | None = None
    ) -> None:
        self.agent = agent
        self.conversation = ConversationHistory()
        self.logger_fn = logger_fn
        # Index into the event log where the current turn started; lets
        # _extract_trace / _turn_events report per-turn data instead of
        # re-logging the whole conversation every message.
        self._turn_start = 0
        # Session documents ({file_name: text}) for the CURRENT call — set per
        # call, not per client, because uploads/deletes can happen mid-session.
        self._documents: dict[str, str] = {}

    def call(
        self,
        query: str,
        lang: str | None = None,
        documents: dict[str, str] | None = None,
    ) -> ChatOutcome:
        reply_id = str(uuid.uuid4())
        self._turn_start = len(self.conversation.ordered_transcript())
        self._documents = documents or {}
        try:
            response = self._invoke(query, lang=lang)
        except AgentException as e:
            outcome = ChatOutcome(
                reply="",
                reply_id=reply_id,
                success=False,
                error=str(e),
                error_type=type(e).__name__,
            )
            logger.warning("agent call failed: %s: %s", type(e).__name__, e)
        else:
            impacts = self._inference_impacts(response.n_tokens)
            outcome = ChatOutcome(
                reply=response.content,
                reply_id=reply_id,
                success=True,
                n_tokens=response.n_tokens,
                env_impact=_to_env_impact(impacts),
                context=self._extract_context(),
                triage_meta=self._triage_meta(),
                inference_impacts=impacts,
            )
        # Trace is attached on failure too — that's when it matters most.
        outcome.trace = self._extract_trace()
        if self.logger_fn:
            self.logger_fn(outcome)
        return outcome

    def _invoke(self, query: str, lang: str | None = None) -> AgentResponse:
        assert self.agent is not None, "base _invoke requires self.agent"
        return self.agent.chat(query, self.conversation)

    def _inference_impacts(self, n_tokens: int) -> Any:
        """Return raw EcoLogits Impacts for this call, or None if unavailable."""
        return None

    def _documents_system_block(self) -> str | None:
        """Uploaded session documents rendered for the system prompt, or None."""
        return build_documents_block(self._documents)

    def _extract_context(self) -> list:
        return []

    def _triage_meta(self) -> dict:
        return {}

    def _turn_events(self) -> list[dict]:
        """Events recorded during the current call() only."""
        return self.conversation.ordered_transcript()[self._turn_start :]

    def _extract_trace(self) -> list:
        """Compact, DDB-safe rendering of everything this turn did internally.

        Covers the outer agent loop (tool calls + model reasoning), the wiki
        sub-agent's pipeline_step events (recalled pages, drafts, judge and
        attribution verdicts, improve cycles), internal follow-up prompts
        (e.g. translation correction), and errors. The user's own message and
        the final reply are already logged as human_message/reply, so they are
        skipped/clipped here.
        """
        trace: list[dict] = []
        first_user_skipped = False
        last_tool_call: str | None = None
        for e in self._turn_events():
            etype = e.get("type")
            if etype == "user":
                # First user event is the human message itself; later ones are
                # internal prompts (translation correction, retry nudges).
                if not first_user_skipped:
                    first_user_skipped = True
                    continue
                trace.append(
                    {"type": "internal_prompt", "content": _clip(e.get("content"), 300)}
                )
            elif etype == "assistant":
                trace.append(
                    {
                        "type": "assistant",
                        "content": _clip(e.get("content"), 300),
                        "reasoning": _clip(e.get("reasoning")),
                    }
                )
            elif etype == "tool_call":
                last_tool_call = e.get("function_name")
                trace.append(
                    {
                        "type": "tool_call",
                        "function": last_tool_call,
                        "arguments": e.get("arguments"),
                        "reasoning": _clip(e.get("reasoning")),
                    }
                )
            elif etype == "tool_result":
                # activate_skill returns the full SKILL.md — static content we
                # already have in the repo; the skill name (in the matching
                # tool_call's arguments) is what matters.
                content = (
                    "<skill instructions elided>"
                    if last_tool_call == "activate_skill"
                    else _clip(e.get("content"), 300)
                )
                trace.append(
                    {"type": "tool_result", "for": last_tool_call, "content": content}
                )
            elif etype == "pipeline_step":
                step = e.get("step") or ""
                row: dict = {"type": "pipeline_step", "step": step}
                if "iteration" in e:
                    row["iteration"] = e["iteration"]
                if step == "subagent/tool_call":
                    row["tool_calls"] = e.get("tool_calls")
                    row["reasoning"] = _clip(e.get("reasoning"))
                elif step == "subagent/tool_result":
                    # Full recalled pages go to `context`; keep a stub here.
                    row["content"] = _clip(e.get("content"), 200)
                elif step in ("subagent/draft_answer", "subagent/answer_before_read"):
                    row["content"] = _clip(e.get("content"))
                    row["reasoning"] = _clip(e.get("reasoning"))
                elif step in ("subagent/judge", "subagent/attribution"):
                    row["verdict"] = e.get("verdict")
                    row["reasoning"] = _clip(e.get("reasoning"))
                elif step == "subagent/improve":
                    row["judge_reasoning"] = _clip(e.get("judge_reasoning"))
                else:  # start / end / future steps: keep scalar fields
                    row.update(
                        {
                            k: v
                            for k, v in e.items()
                            if k not in ("type", "step")
                            and isinstance(v, (str, int, float, bool))
                        }
                    )
                trace.append(row)
            elif etype == "error":
                trace.append({"type": "error", "content": _clip(e.get("content"))})
        return trace

    def get_ordered_transcript(self) -> list[dict]:
        return self.conversation.ordered_transcript()

    def get_conversation_transcript(self) -> list[dict[str, str]]:
        return self.conversation.user_facing_transcript()

    def clear_chat_history(self) -> None:
        self.conversation.reset()

    def upload_document(self, document: str) -> None:
        raise NotImplementedError()


class SkillsAgentClient(AgentClient):
    """gpt-oss-style client: post-call language-leak detection + one correction pass."""

    def _invoke(self, query: str, lang: str | None = None) -> AgentResponse:
        lang = lang or detect_language(query) or "en"
        docs = self._documents or None
        response = self.agent.chat(query, self.conversation, documents=docs)
        leaks = find_leakage(response.content, target_lang=lang)
        if not leaks:
            return response
        prompt = TRANSLATION_CORRECTION_PROMPT.format(
            language=lang, mistranslated_words=leaks
        )
        return self.agent.chat(prompt, self.conversation, documents=docs)

    def _extract_context(self) -> list:
        return [
            e["content"]
            for e in self._turn_events()
            if e.get("type") == "pipeline_step"
            and e.get("step") == "subagent/tool_result"
            and e.get("content")
        ]

    def _inference_impacts(self, n_tokens: int) -> Any:
        # Skills runs on gpt-oss-20b — same coefficients (OSS_AVG_*) the
        # champ path uses.
        return get_champ_impacts(n_tokens) if n_tokens > 0 else None