diff --git a/__pycache__/openai_compatible_loader.cpython-312.pyc b/__pycache__/openai_compatible_loader.cpython-312.pyc index 337741d..ca2b497 100644 Binary files a/__pycache__/openai_compatible_loader.cpython-312.pyc and b/__pycache__/openai_compatible_loader.cpython-312.pyc differ diff --git a/openai_compatible_loader.py b/openai_compatible_loader.py index 365ab27..7965e9c 100644 --- a/openai_compatible_loader.py +++ b/openai_compatible_loader.py @@ -17,8 +17,9 @@ class OpenAICompatibleLoader: def INPUT_TYPES(cls): return { "required": { - "model": ("STRING", {"default": "gpt-3.5-turbo"}), - "api_base": ("STRING", {"default": "https://api.openai.com"}), + "auto_detect": ("BOOLEAN", {"default": True, "label": "自动检测模型"}), + "model": ("STRING", {"default": "", "label": "手动指定模型(当自动检测失败时)"}), + "base_url": ("STRING", {"default": "https://api.openai.com"}), "api_key": ("STRING", {"default": ""}), "prompt": ("STRING", {"multiline": True}), "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0}), @@ -34,10 +35,25 @@ class OpenAICompatibleLoader: FUNCTION = "generate" CATEGORY = "DeepSeek_Toolkit" - async def async_generate(self, session: ClientSession, payload: dict, api_base: str, api_key: str): + async def fetch_models(self, session: ClientSession, base_url: str, api_key: str): + try: + async with session.get( + f"{base_url}/models", + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}" + } + ) as response: + response.raise_for_status() + data = await response.json() + return [model['id'] for model in data.get('data', [])] + except ClientError as e: + raise Exception(f"Failed to fetch models: {str(e)}") + + async def async_generate(self, session: ClientSession, payload: dict, base_url: str, api_key: str): try: async with session.post( - f"{api_base}/chat/completions", + f"{base_url}/chat/completions", headers={ "Content-Type": "application/json", "Authorization": f"Bearer {api_key}" @@ -50,21 +66,57 @@ class OpenAICompatibleLoader: except ClientError as e: raise Exception(f"API request failed: {str(e)}") - def generate(self, api_base: str, api_key: str, prompt: str, model: str, + def generate(self, base_url: str, api_key: str, prompt: str, model: str, temperature: float, max_tokens: int, context: Optional[str] = None): - messages = [{ - "role": "user", - "content": prompt - }] + # 模型检测逻辑 + async def detect_model(): + models = await self.fetch_models(session, base_url, api_key) + if not models: + raise ValueError("无法从API获取模型列表") + return models + # 构建消息体 + messages = [] if context: - messages.insert(0, { + messages.append({ "role": "system", "content": context }) - + messages.append({ + "role": "user", + "content": prompt + }) + + # 确定最终使用的模型 + available_models = [] + selected_model = model + async def run_model_detection(): + if auto_detect: + try: + models = await self.fetch_models(session, base_url, api_key) + if not models: + raise ValueError("无法从API获取模型列表") + return models + except Exception as e: + if not model: + raise ValueError(f"自动检测模型失败且未提供手动模型: {str(e)}") + return [] + + # 模型检测逻辑 + available_models = await run_model_detection() + if model and model not in available_models: + print(f"警告: 手动指定的模型 '{model}' 不在可用模型列表中") + elif not model: + selected_model = available_models[0] if available_models else "" + except Exception as e: + if not model: + raise ValueError(f"自动检测模型失败且未提供手动模型: {str(e)}") + + if not selected_model: + raise ValueError("无法确定使用的模型,请检查base url或手动指定有效模型") + payload = { - "model": model, + "model": selected_model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens @@ -74,7 +126,10 @@ class OpenAICompatibleLoader: # 使用异步请求(需要ComfyUI运行在支持async的环境中) async def run_async(): async with ClientSession() as session: - return await self.async_generate(session, payload, api_base, api_key) + models = await self.fetch_models(session, base_url, api_key) + if model not in models: + raise ValueError(f"Model '{model}' not found in available models.") + return await self.async_generate(session, payload, base_url, api_key) import asyncio response = asyncio.run(run_async()) @@ -86,7 +141,7 @@ class OpenAICompatibleLoader: try: import requests response = requests.post( - f"{api_base}/v1/chat/completions", + f"{base_url}/v1/chat/completions", headers={ "Content-Type": "application/json", "Authorization": f"Bearer {api_key}"