import os

from openai import AsyncOpenAI
from openai.types.chat import (
    ChatCompletion,
    ChatCompletionMessageParam,
    ChatCompletionToolChoiceOptionParam,
    ChatCompletionToolParam,
)


class AsyncLLMClient:
    def __init__(
        self,
        model: str | None = None,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: int | None = None,
    ) -> None:
        self.model = model or os.getenv("LLM_MODEL_ID")
        self.base_url = base_url or os.getenv("LLM_BASE_URL")
        self.api_key = api_key or os.getenv("LLM_API_KEY")
        self.timeout = timeout if timeout is not None else int(os.getenv("LLM_TIMEOUT", "60"))

        if not all([self.model, self.base_url, self.api_key]):
            raise ValueError("模型ID、API密钥和服务地址必须被提供或在.env文件中定义。")
        self.client = AsyncOpenAI(base_url=self.base_url, api_key=self.api_key, timeout=self.timeout)

    async def think(
        self,
        messages: list[ChatCompletionMessageParam],
        tools: list[ChatCompletionToolParam],
        tool_choice: ChatCompletionToolChoiceOptionParam = "auto",
        **kwargs,
    ) -> ChatCompletion:
        assert self.model is not None
        return await self.client.chat.completions.create(
            model=self.model, messages=messages, tools=tools, tool_choice=tool_choice, **kwargs
        )