AIyaodianzhushou/backend-python/app/core/llm_client.py

from typing import AsyncIterator, Optional

from openai import AsyncOpenAI

from app.core.config import get_settings


class LLMClient:
    def __init__(self):
        self._client = None

    def _ensure_client(self):
        if self._client is not None:
            return
        settings = get_settings()
        is_local = settings.qwen_base_url.startswith("http://localhost")
        api_key = "local" if is_local else settings.qwen_api_key
        if not api_key and not is_local:
            raise ValueError("QWEN_API_KEY 未设置，请在 .env 文件中填入 API Key")
        self._client = AsyncOpenAI(
            api_key=api_key,
            base_url=settings.qwen_base_url,
        )
        self.model = settings.qwen_model

    async def chat(
        self,
        messages: list[dict],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ) -> str:
        self._ensure_client()
        settings = get_settings()
        response = await self._client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature or settings.qwen_temperature,
            max_tokens=max_tokens or settings.qwen_max_tokens,
        )
        return response.choices[0].message.content or ""

    async def chat_stream(
        self,
        messages: list[dict],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ) -> AsyncIterator[str]:
        self._ensure_client()
        settings = get_settings()
        stream = await self._client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature or settings.qwen_temperature,
            max_tokens=max_tokens or settings.qwen_max_tokens,
            stream=True,
        )
        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content


llm_client = LLMClient()