from dataclasses import dataclass import json from urllib.error import URLError from urllib.request import Request, urlopen class LLMConfigurationError(ValueError): pass class EmbeddingConfigurationError(ValueError): pass @dataclass class LLMResponse: content: str = "" model_name: str = "" success: bool = True error: Exception | None = None class MockLLMProvider: def __init__(self, model_name: str = "mock-model"): self.model_name = model_name or "mock-model" def generate(self, messages: list[dict], response_format: dict | None = None) -> LLMResponse: user_content = "" for message in reversed(messages): if message.get("role") == "user": user_content = message.get("content", "") break return LLMResponse( content=f"模拟模型回答:{user_content}", model_name=self.model_name, success=True, ) class OpenAICompatibleProvider: def __init__(self, api_key: str, base_url: str, model_name: str): self.api_key = api_key self.base_url = base_url self.model_name = model_name def generate(self, messages: list[dict], response_format: dict | None = None) -> LLMResponse: if not self.api_key: return LLMResponse( model_name=self.model_name, success=False, error=LLMConfigurationError("LLM_API_KEY 未配置,无法调用 OpenAI 兼容模型接口"), ) payload = { "model": self.model_name, "messages": messages, } if response_format: payload["response_format"] = response_format try: data = _post_json( base_url=self.base_url, endpoint="chat/completions", api_key=self.api_key, payload=payload, ) choice = data.get("choices", [{}])[0] content = choice.get("message", {}).get("content", "") return LLMResponse( content=content, model_name=data.get("model", self.model_name), success=True, ) except Exception as exc: return LLMResponse(model_name=self.model_name, success=False, error=exc) class OpenAICompatibleEmbeddingProvider: def __init__(self, api_key: str, base_url: str, model_name: str): self.api_key = api_key self.base_url = base_url self.model_name = model_name def embed_texts(self, texts: list[str]) -> list[list[float]]: if not self.api_key: raise EmbeddingConfigurationError("EMBEDDING_API_KEY 未配置,无法调用 OpenAI 兼容 Embedding 接口") data = _post_json( base_url=self.base_url, endpoint="embeddings", api_key=self.api_key, payload={"model": self.model_name, "input": texts}, ) return [item.get("embedding", []) for item in data.get("data", [])] def _post_json(base_url: str, endpoint: str, api_key: str, payload: dict) -> dict: url = f"{base_url.rstrip('/')}/{endpoint}" request = Request( url, data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", }, method="POST", ) try: with urlopen(request, timeout=60) as response: return json.loads(response.read().decode("utf-8")) except URLError as exc: raise RuntimeError(f"OpenAI 兼容接口调用失败:{exc}") from exc def create_llm_provider(config: dict | None = None): config = config or {} provider_name = config.get("LLM_PROVIDER", "mock") model_name = config.get("LLM_MODEL", "mock-model") if provider_name == "mock": return MockLLMProvider(model_name=model_name) return OpenAICompatibleProvider( api_key=config.get("LLM_API_KEY", ""), base_url=config.get("LLM_BASE_URL", "https://api.openai.com/v1"), model_name=model_name, ) def create_embedding_provider(config: dict | None = None): config = config or {} return OpenAICompatibleEmbeddingProvider( api_key=config.get("EMBEDDING_API_KEY", config.get("LLM_API_KEY", "")), base_url=config.get("EMBEDDING_BASE_URL", config.get("LLM_BASE_URL", "https://api.openai.com/v1")), model_name=config.get("EMBEDDING_MODEL", "text-embedding-3-small"), )