40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
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
|
|
)
|