添加基础的 LLM 支持

konabot/common/llm/__init__.py

@@ -0,0 +1,64 @@
+from typing import Any
+import openai
+
+from loguru import logger
+from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageParam
+from pydantic import BaseModel, Field
+
+from konabot.common.path import CONFIG_PATH
+
+LLM_CONFIG_PATH = CONFIG_PATH / 'llm.json'
+
+if not LLM_CONFIG_PATH.exists():
+    LLM_CONFIG_PATH.write_text("{}")
+
+
+class LLMInfo(BaseModel):
+    base_url: str
+    api_key: str
+    model_name: str
+
+    def get_openai_client(self):
+        return openai.AsyncClient(
+            api_key=self.api_key,
+            base_url=self.base_url,
+        )
+
+    async def chat(
+        self,
+        messages: list[ChatCompletionMessageParam],
+        timeout: float | None = 30.0,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> ChatCompletionMessage:
+        logger.info(f"调用 LLM: BASE_URL={self.base_url} MODEL_NAME={self.model_name}")
+        completion: ChatCompletion = await self.get_openai_client().chat.completions.create(
+            messages=messages,
+            model=self.model_name,
+            max_tokens=max_tokens,
+            timeout=timeout,
+            stream=False,
+            **kwargs,
+        )
+        choice = completion.choices[0]
+        logger.info(
+            f"调用 LLM 完成: BASE_URL={self.base_url} MODEL_NAME={self.model_name} REASON={choice.finish_reason}"
+        )
+        return choice.message
+
+
+class LLMConfig(BaseModel):
+    llms: dict[str, LLMInfo] = Field(default_factory=dict)
+    default_llm: str = "Qwen2.5-7B-Instruct"
+
+
+llm_config = LLMConfig.model_validate_json(LLM_CONFIG_PATH.read_text())
+
+
+def get_llm(llm_model: str | None = None):
+    if llm_model is None:
+        llm_model = llm_config.default_llm
+    if llm_model not in llm_config.llms:
+        raise NotImplementedError("LLM 未配置，该功能无法使用")
+    return llm_config.llms[llm_model]
+

konabot/common/path.py

@@ -6,6 +6,7 @@ FONTS_PATH = ASSETS_PATH / "fonts"
 SRC_PATH = Path(__file__).resolve().parent.parent
 DATA_PATH = SRC_PATH.parent / "data"
 LOG_PATH = DATA_PATH / "logs"
+CONFIG_PATH = DATA_PATH / "config"
 
 DOCS_PATH = SRC_PATH / "docs"
 DOCS_PATH_MAN1 = DOCS_PATH / "user"
@@ -19,3 +20,5 @@ if not DATA_PATH.exists():
 if not LOG_PATH.exists():
     LOG_PATH.mkdir()
 
+CONFIG_PATH.mkdir(exist_ok=True)
+