71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""LLM client wrapper (OpenAI-compatible)."""
|
|
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
|
|
from openai import AsyncOpenAI
|
|
|
|
from steward.config import Settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LLMClient:
|
|
"""Thin async wrapper around the OpenAI chat-completions API."""
|
|
|
|
def __init__(self, settings: Settings) -> None:
|
|
self._settings = settings
|
|
self._client = AsyncOpenAI(
|
|
api_key=settings.openai_api_key,
|
|
base_url=settings.openai_base_url,
|
|
)
|
|
|
|
async def chat(
|
|
self,
|
|
user_message: str,
|
|
*,
|
|
history: list[dict[str, str]] | None = None,
|
|
system_prompt: str | None = None,
|
|
) -> str:
|
|
"""Send a user message (with optional history) and return the assistant reply."""
|
|
system = system_prompt or self._settings.openai_system_prompt
|
|
messages: list[dict[str, str]] = [{"role": "system", "content": system}]
|
|
if history:
|
|
messages.extend(history)
|
|
messages.append({"role": "user", "content": user_message})
|
|
|
|
logger.debug(
|
|
"LLM request: model=%s messages=%d", self._settings.openai_model, len(messages)
|
|
)
|
|
response = await self._client.chat.completions.create(
|
|
model=self._settings.openai_model,
|
|
messages=messages, # type: ignore[arg-type]
|
|
)
|
|
reply = response.choices[0].message.content or ""
|
|
logger.debug("LLM reply: %d chars", len(reply))
|
|
return reply
|
|
|
|
async def stream(
|
|
self,
|
|
user_message: str,
|
|
*,
|
|
history: list[dict[str, str]] | None = None,
|
|
system_prompt: str | None = None,
|
|
) -> AsyncIterator[str]:
|
|
"""Stream the assistant reply token by token."""
|
|
system = system_prompt or self._settings.openai_system_prompt
|
|
messages: list[dict[str, str]] = [{"role": "system", "content": system}]
|
|
if history:
|
|
messages.extend(history)
|
|
messages.append({"role": "user", "content": user_message})
|
|
|
|
stream = await self._client.chat.completions.create(
|
|
model=self._settings.openai_model,
|
|
messages=messages, # type: ignore[arg-type]
|
|
stream=True,
|
|
)
|
|
async for chunk in stream:
|
|
delta = chunk.choices[0].delta.content
|
|
if delta:
|
|
yield delta
|