aisou/ai_search_agent/llm_client.py

from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncGenerator

import httpx

from .config import Settings
from .models import AnswerPayload
from .parsers import extract_json_object


class LLMClient:
    def __init__(self, settings: Settings):
        self.settings = settings

    async def answer(self, system_prompt: str, user_prompt: str, allowed_urls: set[str]) -> AnswerPayload:
        if self.settings.llm_api_style == "openai_chat":
            text = await self._call_openai_chat(system_prompt, user_prompt)
        else:
            text = await self._call_xai_responses(system_prompt, user_prompt)

        parsed = self._parse_answer(text)
        filtered_citations = [item for item in parsed.citations if item.url in allowed_urls]
        return parsed.model_copy(update={"citations": filtered_citations})

    async def _call_xai_responses(self, system_prompt: str, user_prompt: str) -> str:
        return await asyncio.to_thread(self._call_xai_responses_sync, system_prompt, user_prompt)

    def _call_xai_responses_sync(self, system_prompt: str, user_prompt: str) -> str:
        payload = {
            "model": self.settings.llm_model,
            "input": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
        }
        with httpx.Client(timeout=self.settings.model_timeout) as client:
            response = client.post(
                self.settings.llm_endpoint,
                headers={
                    "Authorization": f"Bearer {self.settings.llm_api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
            )
            response.raise_for_status()

        data = response.json()
        output_text: list[str] = []
        for item in data.get("output", []):
            for content in item.get("content", []):
                if content.get("type") == "output_text":
                    output_text.append(content.get("text", ""))
        return "\n".join(part for part in output_text if part).strip()

    async def _call_openai_chat(self, system_prompt: str, user_prompt: str) -> str:
        return await asyncio.to_thread(self._call_openai_chat_sync, system_prompt, user_prompt)

    def _call_openai_chat_sync(self, system_prompt: str, user_prompt: str) -> str:
        payload = {
            "model": self.settings.llm_model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0.2,
        }
        with httpx.Client(timeout=self.settings.model_timeout) as client:
            response = client.post(
                self.settings.llm_endpoint,
                headers={
                    "Authorization": f"Bearer {self.settings.llm_api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
            )
            response.raise_for_status()

        data = response.json()
        return data["choices"][0]["message"]["content"].strip()

    def _parse_answer(self, text: str) -> AnswerPayload:
        raw_json = extract_json_object(text)
        if raw_json:
            try:
                data = json.loads(raw_json)
                return AnswerPayload.model_validate(data)
            except Exception:
                pass

        return AnswerPayload(summary=text.strip(), key_points=[], caveats=[], citations=[])

    # ------------------------------------------------------------------ #
    # Token-level streaming                                                 #
    # ------------------------------------------------------------------ #

    async def answer_stream(
        self,
        system_prompt: str,
        user_prompt: str,
        allowed_urls: set[str],
    ) -> AsyncGenerator[tuple[str, AnswerPayload | None], None]:
        """Yield (token, None) for each streaming token, then ("", AnswerPayload) when done."""
        if self.settings.llm_api_style == "openai_chat":
            full_text = ""
            async for chunk in self._stream_openai_chat(system_prompt, user_prompt):
                full_text += chunk
                yield chunk, None
            parsed = self._parse_answer(full_text)
        else:
            text = await self._call_xai_responses(system_prompt, user_prompt)
            yield text, None
            parsed = self._parse_answer(text)

        filtered = [c for c in parsed.citations if c.url in allowed_urls]
        yield "", parsed.model_copy(update={"citations": filtered})

    async def _stream_openai_chat(
        self, system_prompt: str, user_prompt: str
    ) -> AsyncGenerator[str, None]:
        payload = {
            "model": self.settings.llm_model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0.2,
            "stream": True,
        }
        async with httpx.AsyncClient(timeout=self.settings.model_timeout) as client:
            async with client.stream(
                "POST",
                self.settings.llm_endpoint,
                headers={
                    "Authorization": f"Bearer {self.settings.llm_api_key}",
                    "Content-Type": "application/json",
                    "Accept": "text/event-stream",
                },
                json=payload,
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    line = line.strip()
                    if not line.startswith("data:"):
                        continue
                    data_str = line[len("data:") :].strip()
                    if data_str == "[DONE]":
                        break
                    try:
                        chunk_data = json.loads(data_str)
                        delta = chunk_data.get("choices", [{}])[0].get("delta", {})
                        content = delta.get("content", "")
                        if isinstance(content, str) and content:
                            yield content
                            continue
                        if isinstance(content, list):
                            for item in content:
                                text = item.get("text", "") if isinstance(item, dict) else ""
                                if text:
                                    yield text
                    except Exception:
                        pass