SupplyAgent/supply_agent/llm/client.py

from __future__ import annotations

from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any

from openai import AsyncOpenAI, OpenAI

from supply_agent.config import Settings
from supply_agent.types import Message, Role, ToolCall, ToolDefinition

if TYPE_CHECKING:
    from supply_agent.logging.logger import AgentLogger


class LLMClient:
    """OpenRouter LLM client using the OpenAI-compatible API."""

    def __init__(
        self,
        settings: Settings,
        logger: AgentLogger | None = None,
        *,
        reasoning_effort: str | None = None,
    ) -> None:
        self.settings = settings
        self.model = settings.openrouter_model
        self.logger = logger
        # OpenRouter reasoning effort: low / medium / high / max / xhigh.
        # None or "" disables the parameter.
        self.reasoning_effort = reasoning_effort

        extra_headers: dict[str, str] = {}
        if settings.openrouter_site_url:
            extra_headers["HTTP-Referer"] = settings.openrouter_site_url
        if settings.openrouter_site_name:
            extra_headers["X-Title"] = settings.openrouter_site_name

        self._client = OpenAI(
            api_key=settings.openrouter_api_key,
            base_url=settings.openrouter_base_url,
            default_headers=extra_headers or None,
        )
        self._async_client = AsyncOpenAI(
            api_key=settings.openrouter_api_key,
            base_url=settings.openrouter_base_url,
            default_headers=extra_headers or None,
        )

    def _extra_body(self) -> dict[str, Any] | None:
        """OpenRouter-only params (e.g. Claude reasoning effort)."""
        effort = (self.reasoning_effort or "").strip()
        if not effort:
            return None
        return {"reasoning": {"effort": effort}}

    def chat(
        self,
        messages: list[Message],
        tools: list[ToolDefinition] | None = None,
        temperature: float | None = None,
        *,
        iteration: int = 0,
    ) -> Message:
        """Send a chat completion request and return the assistant message."""
        temp = temperature if temperature is not None else self.settings.agent_temperature

        if self.logger:
            self.logger.log_llm_input(iteration, self.model, messages, tools, temp)

        kwargs: dict[str, Any] = {
            "model": self.model,
            "messages": [m.to_api_dict() for m in messages],
            "tools": [t.to_api_dict() for t in tools] if tools else None,
            "temperature": temp,
        }
        extra_body = self._extra_body()
        if extra_body:
            kwargs["extra_body"] = extra_body

        response = self._client.chat.completions.create(**kwargs)
        raw_message = response.choices[0].message
        result = self._parse_response(raw_message)

        if self.logger:
            self.logger.log_llm_output(iteration, result, raw_response=response)

        return result

    async def achat(
        self,
        messages: list[Message],
        tools: list[ToolDefinition] | None = None,
        temperature: float | None = None,
        *,
        iteration: int = 0,
    ) -> Message:
        """Async chat completion."""
        temp = temperature if temperature is not None else self.settings.agent_temperature

        if self.logger:
            self.logger.log_llm_input(iteration, self.model, messages, tools, temp)

        kwargs: dict[str, Any] = {
            "model": self.model,
            "messages": [m.to_api_dict() for m in messages],
            "tools": [t.to_api_dict() for t in tools] if tools else None,
            "temperature": temp,
        }
        extra_body = self._extra_body()
        if extra_body:
            kwargs["extra_body"] = extra_body

        response = await self._async_client.chat.completions.create(**kwargs)
        raw_message = response.choices[0].message
        result = self._parse_response(raw_message)

        if self.logger:
            self.logger.log_llm_output(iteration, result, raw_response=response)

        return result

    def stream(
        self,
        messages: list[Message],
        tools: list[ToolDefinition] | None = None,
        temperature: float | None = None,
        *,
        iteration: int = 0,
    ) -> Iterator[str]:
        """Stream text content chunks from the model."""
        temp = temperature if temperature is not None else self.settings.agent_temperature

        if self.logger:
            self.logger.log_llm_input(iteration, self.model, messages, tools, temp)

        kwargs: dict[str, Any] = {
            "model": self.model,
            "messages": [m.to_api_dict() for m in messages],
            "tools": [t.to_api_dict() for t in tools] if tools else None,
            "temperature": temp,
            "stream": True,
        }
        extra_body = self._extra_body()
        if extra_body:
            kwargs["extra_body"] = extra_body

        stream = self._client.chat.completions.create(**kwargs)
        chunks: list[str] = []
        for chunk in stream:
            delta = chunk.choices[0].delta
            if delta.content:
                chunks.append(delta.content)
                yield delta.content

        if self.logger:
            full_content = "".join(chunks)
            self.logger.log_llm_output(
                iteration,
                Message(role=Role.ASSISTANT, content=full_content),
            )

    async def astream(
        self,
        messages: list[Message],
        tools: list[ToolDefinition] | None = None,
        temperature: float | None = None,
        *,
        iteration: int = 0,
    ) -> AsyncIterator[str]:
        """Async stream text content chunks."""
        temp = temperature if temperature is not None else self.settings.agent_temperature

        if self.logger:
            self.logger.log_llm_input(iteration, self.model, messages, tools, temp)

        kwargs: dict[str, Any] = {
            "model": self.model,
            "messages": [m.to_api_dict() for m in messages],
            "tools": [t.to_api_dict() for t in tools] if tools else None,
            "temperature": temp,
            "stream": True,
        }
        extra_body = self._extra_body()
        if extra_body:
            kwargs["extra_body"] = extra_body

        stream = await self._async_client.chat.completions.create(**kwargs)
        chunks: list[str] = []
        async for chunk in stream:
            delta = chunk.choices[0].delta
            if delta.content:
                chunks.append(delta.content)
                yield delta.content

        if self.logger:
            full_content = "".join(chunks)
            self.logger.log_llm_output(
                iteration,
                Message(role=Role.ASSISTANT, content=full_content),
            )

    def _parse_response(self, choice_message: Any) -> Message:
        tool_calls = None
        if choice_message.tool_calls:
            tool_calls = [
                ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=tc.function.arguments,
                )
                for tc in choice_message.tool_calls
            ]
        reasoning = getattr(choice_message, "reasoning", None)
        if reasoning is not None and not isinstance(reasoning, str):
            reasoning = str(reasoning)
        return Message(
            role=Role.ASSISTANT,
            content=choice_message.content,
            tool_calls=tool_calls,
            reasoning=reasoning,
        )

    def set_model(self, model: str) -> None:
        """Switch to a different model at runtime."""
        self.model = model