nodes/RCAKangle-ComfyUI_LLM_Embeder
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add local OpenAI-compatible server support

Browse Source
This commit is contained in:
Kangle
2026-08-28 12:55:15 +08:00
parent 3103f14720
commit 2e0750e465
5 changed files with 303 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+23 -12
View File
@@ -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)
__pycache__/__init__.cpython-312.pyc
BIN
View File
Binary file not shown.
__pycache__/chat_optimize_nodes.cpython-312.pyc
BIN
View File
Binary file not shown.
chat_optimize_nodes.py
+76 -45
View File
@@ -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",
}
tests/test_openai_compatible_provider.py
+204
View File
@@ -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()