app/agents/base.py

"""Abstract base for LexGuard agents.

Concrete agents define:
  - `name`: matches `AgentName` Literal in schemas
  - `prompt_file`: filename (without .md) under `app/prompts/`
  - `run(...)`: the agent-specific call producing structured output

The base class owns:
  - prompt loading + system-block assembly with caching
  - LLM invocation
  - event-bus emission of `AgentMessage` for streaming
"""

from __future__ import annotations

import logging
import time
from abc import ABC
from typing import Any, ClassVar

from app.prompts import SHARED_PREAMBLE, load_prompt
from app.schemas import AgentMessage, AgentName
from app.services import llm

logger = logging.getLogger(__name__)


class BaseAgent(ABC):
    name: ClassVar[AgentName]
    prompt_file: ClassVar[str]

    def __init__(self, *, document: str | None = None, emit=None) -> None:
        self._document = document
        self._emit = emit

    @property
    def role_instructions(self) -> str:
        return load_prompt(self.prompt_file)

    def system_blocks(self) -> list[dict[str, Any]]:
        return llm.build_system_blocks(
            shared_preamble=SHARED_PREAMBLE,
            document=self._document,
            role_instructions=self.role_instructions,
        )

    async def _call(
        self,
        user_prompt: str,
        *,
        max_tokens: int = 2048,
        temperature: float = 0.2,
        clause_id: str | None = None,
        log_label: str = "",
    ) -> llm.LLMResponse:
        started = time.perf_counter()
        response = await llm.complete(
            system_blocks=self.system_blocks(),
            user_prompt=user_prompt,
            max_tokens=max_tokens,
            temperature=temperature,
        )
        elapsed = int((time.perf_counter() - started) * 1000)
        logger.info(
            "agent=%s %s clause_id=%s elapsed_ms=%d cache_read=%d",
            self.name,
            log_label,
            clause_id,
            elapsed,
            response.usage.cache_read_input_tokens,
        )
        return response

    def _emit_message(self, content: str, clause_id: str | None = None) -> None:
        if self._emit is None:
            return
        self._emit(
            AgentMessage(
                agent=self.name,
                clause_id=clause_id,
                content=content,
                timestamp=time.time(),
            )
        )