diff --git a/README.md b/README.md index a798454..6fcdabf 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ Local LLM chat nodes for ComfyUI, with a clean handoff path to downstream prompt ## Features - Multi-turn chat with Ollama `/api/chat` - Optional Hugging Face Inference API support via LLM Config -- OpenAI-compatible chat support (OpenAI / DeepSeek / Qwen) via custom `base_url` +- OpenAI-compatible chat support (local servers / OpenAI / DeepSeek / Qwen) via custom `base_url` - Anthropic Claude support via custom `base_url` - Session memory by `session_id` (in-memory only) - Optional `system_prompt` @@ -15,10 +15,10 @@ Local LLM chat nodes for ComfyUI, with a clean handoff path to downstream prompt ## Nodes -### Chat (Ollama) +### Chat (LLM) Inputs: -- `model_name`: model id (Ollama or HF, depending on provider) -- `base_url`: Ollama base URL (ignored when provider is `huggingface`) +- `model_name`: model name or model id for the selected provider +- `base_url`: provider base URL (ignored when provider is `huggingface`) - `user_message`: user input text - `action`: `send` / `regenerate` / `clear` / `deliver_to_optimizer` - `session_id`: conversation id @@ -41,7 +41,7 @@ Behavior: Outputs a `LLM_CONFIG` struct that can be shared across Chat nodes. Fields: -- `provider`: `ollama`, `huggingface`, `openai`, `deepseek`, `qwen`, `claude` +- `provider`: `ollama`, `openai_compatible`, `huggingface`, `openai`, `deepseek`, `qwen`, `claude` - `base_url`: provider base URL (see sections below) - `model_name`: model name or model id (provider specific) - `temperature` @@ -62,15 +62,26 @@ Notes: - Many public models still require an HF token. - Chat history is converted into a single prompt before the API call. -## OpenAI-Compatible API (OpenAI / DeepSeek / Qwen) -When `provider=openai|deepseek|qwen`, the Chat node will call: +## OpenAI-Compatible API +When `provider=openai_compatible|openai|deepseek|qwen`, the Chat node will call: `{base_url}/chat/completions` Notes: -- Use an OpenAI-compatible base URL for your provider. -- Set `api_key` in LLM Config. +- Use an OpenAI-compatible base URL that includes `/v1`, or enter the full `/v1/chat/completions` URL. +- Set `api_key` when required by the server; it can remain empty for local servers without authentication. - `model_name` should be the provider's model id. +### Local OpenAI-Compatible Servers +Select `provider=openai_compatible` in LLM Config and use the base URL exposed by the local server: + +| Server | Example `base_url` | +| --- | --- | +| LM Studio | `http://127.0.0.1:1234/v1` | +| llama.cpp (`llama-server`) | `http://127.0.0.1:8080/v1` | +| LMDeploy | `http://127.0.0.1:23333/v1` | + +Use the exact model id reported by the server as `model_name`. Ports may differ if the server was started with custom settings. + ## Anthropic Claude API When `provider=claude`, the Chat node will call: `{base_url}/v1/messages` @@ -83,12 +94,12 @@ Notes: 1) Place this folder under `ComfyUI/custom_nodes/` 2) Restart ComfyUI 3) Add nodes: - - `Chat (Ollama)` + - `Chat (LLM)` - `LLM Config` - `Chat History Viewer` ## Usage -1) Create `Chat (Ollama)` and set `base_url` / `model_name` +1) Create `Chat (LLM)` and set `base_url` / `model_name` 2) (Optional) Add `LLM Config` to select provider and set `api_key` 3) Set `action=send`, click Execute to chat 4) When satisfied, set `action=deliver_to_optimizer` and Execute @@ -97,7 +108,7 @@ Notes: ## Screenshot ![ComfyUI node graph example](LLM_embeder_nodes.png) -Example graph: `LLM Config` drives `Chat (Ollama)` via `llm_config`, and outputs route to a `text` node plus `Chat History Viewer` for full history review. +Example graph: `LLM Config` drives `Chat (LLM)` via `llm_config`, and outputs route to a `text` node plus `Chat History Viewer` for full history review. ## Example Workflow ![Example ComfyUI workflow](Example%20workflow.png) diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..a39b7f4 Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ diff --git a/__pycache__/chat_optimize_nodes.cpython-312.pyc b/__pycache__/chat_optimize_nodes.cpython-312.pyc new file mode 100644 index 0000000..a537b50 Binary files /dev/null and b/__pycache__/chat_optimize_nodes.cpython-312.pyc differ diff --git a/chat_optimize_nodes.py b/chat_optimize_nodes.py index 594b858..4a70467 100644 --- a/chat_optimize_nodes.py +++ b/chat_optimize_nodes.py @@ -1,15 +1,24 @@ import asyncio import json import requests +import threading from typing import Dict, List, Optional from aiohttp import web from server import PromptServer # Simple in-memory stores for chat state. _chat_sessions: Dict[str, List[Dict[str, str]]] = {} +_chat_session_locks = {} +_chat_session_locks_guard = threading.Lock() REQUEST_TIMEOUT = 60 +def _get_chat_session_lock(session_id: str): + with _chat_session_locks_guard: + if session_id not in _chat_session_locks: + _chat_session_locks[session_id] = threading.Lock() + return _chat_session_locks[session_id] + def _build_headers() -> Dict[str, str]: return {"Content-Type": "application/json"} @@ -277,6 +286,31 @@ class ChatNode: refresh_session: bool = False, auto_clear_input: bool = True, llm_config: Optional[Dict] = None, + ): + with _get_chat_session_lock(session_id): + return self._chat_locked( + model_name=model_name, + base_url=base_url, + user_message=user_message, + action=action, + session_id=session_id, + system_prompt=system_prompt, + refresh_session=refresh_session, + auto_clear_input=auto_clear_input, + llm_config=llm_config, + ) + + def _chat_locked( + self, + model_name: str, + base_url: str, + user_message: str, + action: str = "send", + session_id: str = "default", + system_prompt: str = "", + refresh_session: bool = False, + auto_clear_input: bool = True, + llm_config: Optional[Dict] = None, ): provider = "ollama" hf_token = "" @@ -337,49 +371,46 @@ class ChatNode: response_text = "" if messages and messages[-1].get("role") == "user": - try: - if provider == "huggingface": - response_text = _call_hf_inference( - model_name, - hf_token, - messages, - temperature=temperature, - top_p=top_p, - max_new_tokens=max_new_tokens, - api_url=hf_api_url or None, - ) - elif provider in ("openai", "deepseek", "qwen"): - response_text = _call_openai_compatible_chat( - base_url, - model_name, - api_key, - messages, - temperature=temperature, - top_p=top_p, - max_new_tokens=max_new_tokens, - ) - elif provider == "claude": - response_text = _call_anthropic_messages( - base_url, - model_name, - api_key, - messages, - temperature=temperature, - top_p=top_p, - max_new_tokens=max_new_tokens, - ) - else: - options = {} - if temperature is not None: - options["temperature"] = temperature - if top_p is not None: - options["top_p"] = top_p - if max_new_tokens is not None: - options["num_predict"] = max_new_tokens - response_text = _call_ollama_chat(base_url, model_name, messages, other_options=options) - messages.append({"role": "assistant", "content": response_text}) - except Exception as exc: - response_text = f"[chat error] {exc}" + if provider == "huggingface": + response_text = _call_hf_inference( + model_name, + hf_token, + messages, + temperature=temperature, + top_p=top_p, + max_new_tokens=max_new_tokens, + api_url=hf_api_url or None, + ) + elif provider in ("openai_compatible", "openai", "deepseek", "qwen"): + response_text = _call_openai_compatible_chat( + base_url, + model_name, + api_key, + messages, + temperature=temperature, + top_p=top_p, + max_new_tokens=max_new_tokens, + ) + elif provider == "claude": + response_text = _call_anthropic_messages( + base_url, + model_name, + api_key, + messages, + temperature=temperature, + top_p=top_p, + max_new_tokens=max_new_tokens, + ) + else: + options = {} + if temperature is not None: + options["temperature"] = temperature + if top_p is not None: + options["top_p"] = top_p + if max_new_tokens is not None: + options["num_predict"] = max_new_tokens + response_text = _call_ollama_chat(base_url, model_name, messages, other_options=options) + messages.append({"role": "assistant", "content": response_text}) _chat_sessions[session_id] = messages @@ -396,7 +427,7 @@ class LLMConfigNode: def INPUT_TYPES(cls): return { "required": { - "provider": (["ollama", "huggingface", "openai", "deepseek", "qwen", "claude"], {"default": "ollama"}), + "provider": (["ollama", "openai_compatible", "huggingface", "openai", "deepseek", "qwen", "claude"], {"default": "ollama"}), "base_url": ("STRING", {"default": "http://127.0.0.1:11434"}), "model_name": ("STRING", {"default": "llama3"}), "temperature": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), @@ -482,7 +513,7 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "ChatNode": "Chat (Ollama)", + "ChatNode": "Chat (LLM)", "LLMConfigNode": "LLM Config", "ChatHistoryViewer": "Chat History Viewer", } diff --git a/tests/test_openai_compatible_provider.py b/tests/test_openai_compatible_provider.py new file mode 100644 index 0000000..e49a370 --- /dev/null +++ b/tests/test_openai_compatible_provider.py @@ -0,0 +1,204 @@ +import importlib.util +import sys +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +class _DummyRoutes: + def post(self, _path): + return lambda handler: handler + + +def _load_chat_nodes(): + fake_aiohttp = types.ModuleType("aiohttp") + fake_aiohttp.web = types.SimpleNamespace( + json_response=lambda payload, status=200: (payload, status) + ) + + fake_server = types.ModuleType("server") + fake_server.PromptServer = types.SimpleNamespace( + instance=types.SimpleNamespace(routes=_DummyRoutes()) + ) + + module_path = PROJECT_ROOT / "chat_optimize_nodes.py" + spec = importlib.util.spec_from_file_location( + "chat_optimize_nodes_under_test", module_path + ) + module = importlib.util.module_from_spec(spec) + with patch.dict(sys.modules, {"aiohttp": fake_aiohttp, "server": fake_server}): + spec.loader.exec_module(module) + return module + + +chat_nodes = _load_chat_nodes() + + +class OpenAICompatibleProviderTests(unittest.TestCase): + def setUp(self): + chat_nodes._chat_sessions.clear() + chat_nodes._chat_session_locks.clear() + + def test_module_loader_restores_dependency_modules(self): + original_modules = { + name: sys.modules.get(name) for name in ("aiohttp", "server") + } + + _load_chat_nodes() + + for name, original_module in original_modules.items(): + with self.subTest(name=name): + self.assertIs(sys.modules.get(name), original_module) + + def test_llm_config_exposes_openai_compatible_provider(self): + providers = chat_nodes.LLMConfigNode.INPUT_TYPES()["required"]["provider"][0] + + self.assertIn("openai_compatible", providers) + + def test_openai_compatible_provider_reuses_existing_chat_client(self): + config = { + "provider": "openai_compatible", + "base_url": "http://127.0.0.1:3762/v1", + "model_name": "local-model", + "api_key": "", + } + + with patch.object( + chat_nodes, "_call_openai_compatible_chat", return_value="local response" + ) as compatible_call, patch.object(chat_nodes, "_call_ollama_chat") as ollama_call: + _, history = chat_nodes.ChatNode().chat( + model_name="unused", + base_url="http://127.0.0.1:11434", + user_message="hello", + session_id="openai-compatible-test", + refresh_session=True, + llm_config=config, + ) + + compatible_call.assert_called_once() + ollama_call.assert_not_called() + self.assertEqual(compatible_call.call_args.args[:3], ( + "http://127.0.0.1:3762/v1", + "local-model", + "", + )) + self.assertIn("Assistant: local response", history) + + def test_provider_error_propagates_without_storing_failed_turn(self): + session_id = "provider-error-test" + config = { + "provider": "openai_compatible", + "base_url": "http://127.0.0.1:3762/v1", + "model_name": "local-model", + "api_key": "", + } + + with patch.object( + chat_nodes, + "_call_openai_compatible_chat", + side_effect=RuntimeError("connection refused"), + ): + with self.assertRaisesRegex(RuntimeError, "connection refused"): + chat_nodes.ChatNode().chat( + model_name="unused", + base_url="http://127.0.0.1:11434", + user_message="hello", + session_id=session_id, + refresh_session=True, + llm_config=config, + ) + + self.assertEqual(chat_nodes._chat_sessions[session_id], []) + + def test_same_session_requests_are_serialized(self): + session_id = "concurrent-session-test" + config = { + "provider": "openai_compatible", + "base_url": "http://127.0.0.1:3762/v1", + "model_name": "local-model", + "api_key": "", + } + first_started = threading.Event() + second_started = threading.Event() + release_first = threading.Event() + call_count = 0 + call_count_lock = threading.Lock() + errors = [] + + def fake_chat(_base_url, _model_name, _api_key, messages, **_kwargs): + nonlocal call_count + with call_count_lock: + call_count += 1 + current_call = call_count + if current_call == 1: + first_started.set() + if not release_first.wait(timeout=2): + raise TimeoutError("first request was not released") + else: + second_started.set() + return f"{messages[-1]['content']}-reply" + + def run_chat(message): + try: + chat_nodes.ChatNode().chat( + model_name="unused", + base_url="http://127.0.0.1:11434", + user_message=message, + session_id=session_id, + llm_config=config, + ) + except Exception as exc: + errors.append(exc) + + with patch.object(chat_nodes, "_call_openai_compatible_chat", side_effect=fake_chat): + first_thread = threading.Thread(target=run_chat, args=("first",)) + second_thread = threading.Thread(target=run_chat, args=("second",)) + first_thread.start() + self.assertTrue(first_started.wait(timeout=1)) + second_thread.start() + second_started_before_release = second_started.wait(timeout=0.2) + release_first.set() + first_thread.join(timeout=2) + second_thread.join(timeout=2) + + self.assertFalse(first_thread.is_alive()) + self.assertFalse(second_thread.is_alive()) + self.assertEqual(errors, []) + self.assertFalse(second_started_before_release) + self.assertEqual( + chat_nodes._chat_sessions[session_id], + [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "first-reply"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "second-reply"}, + ], + ) + + def test_chat_node_has_provider_neutral_display_name(self): + self.assertEqual(chat_nodes.NODE_DISPLAY_NAME_MAPPINGS["ChatNode"], "Chat (LLM)") + + def test_readme_documents_supported_local_server_examples(self): + readme = (PROJECT_ROOT / "README.md").read_text(encoding="utf-8") + + expected_text = ( + "`provider=openai_compatible`", + "LM Studio", + "http://127.0.0.1:1234/v1", + "llama.cpp", + "http://127.0.0.1:8080/v1", + "LMDeploy", + "http://127.0.0.1:23333/v1", + ) + for text in expected_text: + with self.subTest(text=text): + self.assertIn(text, readme) + + +if __name__ == "__main__": + unittest.main()