524 lines
18 KiB
Python
524 lines
18 KiB
Python
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"}
|
|
|
|
def _clean_llm_output(text: str) -> str:
|
|
"""
|
|
Strictly clean the LLM output to ensure only the raw prompt remains.
|
|
Removes Markdown code blocks, JSON wrapping if present, and extra whitespace.
|
|
"""
|
|
text = text.strip()
|
|
if text.startswith("```") and text.endswith("```"):
|
|
lines = text.splitlines()
|
|
if lines[0].startswith("```"):
|
|
lines = lines[1:]
|
|
if lines and lines[-1].strip() == "```":
|
|
lines = lines[:-1]
|
|
text = "\n".join(lines).strip()
|
|
return text
|
|
|
|
def _call_ollama_chat(
|
|
base_url: str,
|
|
model: str,
|
|
messages: List[Dict[str, str]],
|
|
seed: Optional[int] = None,
|
|
other_options: Optional[Dict] = None,
|
|
timeout: int = REQUEST_TIMEOUT,
|
|
) -> str:
|
|
url = base_url.rstrip("/") + "/api/chat"
|
|
payload = {"model": model, "messages": messages, "stream": False}
|
|
|
|
options = {}
|
|
if seed is not None:
|
|
options["seed"] = seed
|
|
if other_options:
|
|
options.update(other_options)
|
|
|
|
if options:
|
|
payload["options"] = options
|
|
|
|
resp = requests.post(url, headers=_build_headers(), data=json.dumps(payload), timeout=timeout)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
return data.get("message", {}).get("content", "")
|
|
|
|
def _build_hf_prompt(messages: List[Dict[str, str]]) -> str:
|
|
lines: List[str] = []
|
|
for msg in messages:
|
|
role = msg.get("role", "user")
|
|
content = msg.get("content", "")
|
|
if role == "system":
|
|
lines.append(f"System: {content}")
|
|
elif role == "user":
|
|
lines.append(f"User: {content}")
|
|
elif role == "assistant":
|
|
lines.append(f"Assistant: {content}")
|
|
else:
|
|
lines.append(f"{role.capitalize()}: {content}")
|
|
lines.append("Assistant:")
|
|
return "\n".join(lines).strip()
|
|
|
|
def _call_hf_inference(
|
|
model: str,
|
|
token: str,
|
|
messages: List[Dict[str, str]],
|
|
temperature: Optional[float] = None,
|
|
top_p: Optional[float] = None,
|
|
max_new_tokens: Optional[int] = None,
|
|
api_url: Optional[str] = None,
|
|
timeout: int = REQUEST_TIMEOUT,
|
|
) -> str:
|
|
url = api_url or f"https://api-inference.huggingface.co/models/{model}"
|
|
headers = _build_headers()
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
|
|
prompt = _build_hf_prompt(messages)
|
|
parameters: Dict[str, object] = {}
|
|
if temperature is not None:
|
|
parameters["temperature"] = temperature
|
|
if top_p is not None:
|
|
parameters["top_p"] = top_p
|
|
if max_new_tokens is not None:
|
|
parameters["max_new_tokens"] = int(max_new_tokens)
|
|
|
|
payload: Dict[str, object] = {"inputs": prompt}
|
|
if parameters:
|
|
payload["parameters"] = parameters
|
|
|
|
resp = requests.post(url, headers=headers, data=json.dumps(payload), timeout=timeout)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
if isinstance(data, dict) and data.get("error"):
|
|
raise RuntimeError(data.get("error"))
|
|
|
|
item = data
|
|
if isinstance(data, list) and data:
|
|
item = data[0]
|
|
|
|
generated_text = ""
|
|
if isinstance(item, dict):
|
|
generated_text = item.get("generated_text", "")
|
|
elif isinstance(item, str):
|
|
generated_text = item
|
|
|
|
if generated_text.startswith(prompt):
|
|
generated_text = generated_text[len(prompt):]
|
|
|
|
return generated_text.strip()
|
|
|
|
def _build_openai_compatible_url(base_url: str) -> str:
|
|
base = (base_url or "").rstrip("/")
|
|
if base.endswith("/chat/completions"):
|
|
return base
|
|
return base + "/chat/completions"
|
|
|
|
def _call_openai_compatible_chat(
|
|
base_url: str,
|
|
model: str,
|
|
api_key: str,
|
|
messages: List[Dict[str, str]],
|
|
temperature: Optional[float] = None,
|
|
top_p: Optional[float] = None,
|
|
max_new_tokens: Optional[int] = None,
|
|
timeout: int = REQUEST_TIMEOUT,
|
|
) -> str:
|
|
url = _build_openai_compatible_url(base_url)
|
|
headers = _build_headers()
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
payload: Dict[str, object] = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"stream": False,
|
|
}
|
|
if temperature is not None:
|
|
payload["temperature"] = temperature
|
|
if top_p is not None:
|
|
payload["top_p"] = top_p
|
|
if max_new_tokens is not None:
|
|
payload["max_tokens"] = int(max_new_tokens)
|
|
|
|
resp = requests.post(url, headers=headers, data=json.dumps(payload), timeout=timeout)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
if isinstance(data, dict) and data.get("error"):
|
|
raise RuntimeError(data.get("error"))
|
|
|
|
choices = data.get("choices") if isinstance(data, dict) else None
|
|
if isinstance(choices, list) and choices:
|
|
message = choices[0].get("message", {})
|
|
return message.get("content", "").strip()
|
|
|
|
return ""
|
|
|
|
def _split_system_messages(messages: List[Dict[str, str]]):
|
|
system_parts: List[str] = []
|
|
non_system: List[Dict[str, str]] = []
|
|
for msg in messages:
|
|
if msg.get("role") == "system":
|
|
system_parts.append(msg.get("content", ""))
|
|
else:
|
|
non_system.append(msg)
|
|
return ("\n\n".join([p for p in system_parts if p]), non_system)
|
|
|
|
def _call_anthropic_messages(
|
|
base_url: str,
|
|
model: str,
|
|
api_key: str,
|
|
messages: List[Dict[str, str]],
|
|
temperature: Optional[float] = None,
|
|
top_p: Optional[float] = None,
|
|
max_new_tokens: Optional[int] = None,
|
|
timeout: int = REQUEST_TIMEOUT,
|
|
) -> str:
|
|
base = (base_url or "").rstrip("/")
|
|
if base.endswith("/v1/messages"):
|
|
url = base
|
|
elif base.endswith("/v1"):
|
|
url = base + "/messages"
|
|
else:
|
|
url = base + "/v1/messages"
|
|
headers = _build_headers()
|
|
if api_key:
|
|
headers["x-api-key"] = api_key
|
|
headers["anthropic-version"] = "2023-06-01"
|
|
|
|
system_prompt, non_system = _split_system_messages(messages)
|
|
formatted_messages: List[Dict[str, object]] = []
|
|
for msg in non_system:
|
|
role = msg.get("role", "user")
|
|
content = msg.get("content", "")
|
|
formatted_messages.append(
|
|
{"role": role, "content": [{"type": "text", "text": content}]}
|
|
)
|
|
|
|
payload: Dict[str, object] = {
|
|
"model": model,
|
|
"messages": formatted_messages,
|
|
"max_tokens": int(max_new_tokens or 256),
|
|
}
|
|
if system_prompt:
|
|
payload["system"] = system_prompt
|
|
if temperature is not None:
|
|
payload["temperature"] = temperature
|
|
if top_p is not None:
|
|
payload["top_p"] = top_p
|
|
|
|
resp = requests.post(url, headers=headers, data=json.dumps(payload), timeout=timeout)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
if isinstance(data, dict) and data.get("error"):
|
|
raise RuntimeError(data.get("error"))
|
|
|
|
content = data.get("content") if isinstance(data, dict) else None
|
|
if isinstance(content, list) and content:
|
|
first = content[0]
|
|
if isinstance(first, dict):
|
|
return first.get("text", "").strip()
|
|
|
|
return ""
|
|
|
|
def _reset_history(session_id: str, system_prompt: str) -> List[Dict[str, str]]:
|
|
history: List[Dict[str, str]] = []
|
|
if system_prompt:
|
|
history.append({"role": "system", "content": system_prompt})
|
|
_chat_sessions[session_id] = history
|
|
return history
|
|
|
|
class ChatNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": ("STRING", {"default": "llama3"}),
|
|
"base_url": ("STRING", {"default": "http://127.0.0.1:11434"}),
|
|
"user_message": ("STRING", {"multiline": True, "default": ""}),
|
|
"action": (["send", "regenerate", "clear", "deliver_to_optimizer"], {"default": "send"}),
|
|
},
|
|
"optional": {
|
|
"session_id": ("STRING", {"default": "default"}),
|
|
"system_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"refresh_session": ("BOOLEAN", {"default": False}),
|
|
"auto_clear_input": ("BOOLEAN", {"default": True}),
|
|
"llm_config": ("LLM_CONFIG",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("assistant_response", "readable_history")
|
|
FUNCTION = "chat"
|
|
CATEGORY = "ChatOptimize"
|
|
OUTPUT_NODE = False
|
|
|
|
def chat(
|
|
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,
|
|
):
|
|
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 = ""
|
|
hf_api_url = ""
|
|
api_key = ""
|
|
temperature = None
|
|
top_p = None
|
|
max_new_tokens = None
|
|
if llm_config:
|
|
provider = llm_config.get("provider", provider)
|
|
model_name = llm_config.get("model_name", model_name)
|
|
base_url = llm_config.get("base_url", base_url)
|
|
hf_token = llm_config.get("hf_token", "")
|
|
hf_api_url = llm_config.get("hf_api_url", "")
|
|
api_key = llm_config.get("api_key", "")
|
|
if "temperature" in llm_config:
|
|
temperature = llm_config["temperature"]
|
|
if "top_p" in llm_config:
|
|
top_p = llm_config["top_p"]
|
|
if "max_new_tokens" in llm_config:
|
|
max_new_tokens = llm_config["max_new_tokens"]
|
|
|
|
if refresh_session or session_id not in _chat_sessions:
|
|
history = _reset_history(session_id, system_prompt)
|
|
else:
|
|
history = _chat_sessions.get(session_id, [])
|
|
|
|
if action == "clear":
|
|
history = _reset_history(session_id, system_prompt)
|
|
return ("", "")
|
|
|
|
if system_prompt:
|
|
if not history or history[0].get("role") != "system":
|
|
history = _reset_history(session_id, system_prompt)
|
|
|
|
messages = list(history)
|
|
|
|
if action == "send":
|
|
if user_message.strip():
|
|
messages.append({"role": "user", "content": user_message})
|
|
elif action == "regenerate":
|
|
for idx in range(len(messages) - 1, -1, -1):
|
|
if messages[idx].get("role") == "assistant":
|
|
messages = messages[:idx]
|
|
break
|
|
elif action == "deliver_to_optimizer":
|
|
latest_response = ""
|
|
for msg in reversed(messages):
|
|
if msg.get("role") == "assistant":
|
|
latest_response = msg.get("content", "")
|
|
break
|
|
readable_history = ""
|
|
for msg in messages:
|
|
role = msg.get("role", "unknown").capitalize()
|
|
content = msg.get("content", "")
|
|
readable_history += f"{role}: {content}\n\n"
|
|
return (latest_response, readable_history)
|
|
|
|
response_text = ""
|
|
if messages and messages[-1].get("role") == "user":
|
|
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
|
|
|
|
readable_history = ""
|
|
for msg in messages:
|
|
role = msg.get("role", "unknown").capitalize()
|
|
content = msg.get("content", "")
|
|
readable_history += f"{role}: {content}\n\n"
|
|
|
|
return ("", readable_history)
|
|
|
|
class LLMConfigNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"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}),
|
|
"top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"max_new_tokens": ("INT", {"default": 256, "min": 1, "max": 8192, "step": 1}),
|
|
"hf_token": ("STRING", {"default": ""}),
|
|
"hf_api_url": ("STRING", {"default": ""}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLM_CONFIG",)
|
|
RETURN_NAMES = ("llm_config",)
|
|
FUNCTION = "config"
|
|
CATEGORY = "ChatOptimize"
|
|
|
|
def config(self, provider, base_url, model_name, temperature, top_p, max_new_tokens, hf_token, hf_api_url, api_key):
|
|
return (
|
|
{
|
|
"provider": provider,
|
|
"base_url": base_url,
|
|
"model_name": model_name,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_new_tokens": max_new_tokens,
|
|
"hf_token": hf_token,
|
|
"hf_api_url": hf_api_url,
|
|
"api_key": api_key,
|
|
},
|
|
)
|
|
|
|
class ChatHistoryViewer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"history": ("STRING", {"default": "", "multiline": True}),
|
|
}
|
|
}
|
|
|
|
INPUT_IS_LIST = False
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("history",)
|
|
FUNCTION = "show"
|
|
CATEGORY = "ChatOptimize"
|
|
OUTPUT_NODE = True
|
|
|
|
def show(self, history=""):
|
|
return {"ui": {"text": [history]}, "result": (history,)}
|
|
|
|
def _chat_action(data: Dict[str, object]):
|
|
node = ChatNode()
|
|
return node.chat(
|
|
model_name=data.get("model_name", "llama3"),
|
|
base_url=data.get("base_url", "http://127.0.0.1:11434"),
|
|
user_message=data.get("user_message", ""),
|
|
action=data.get("action", "send"),
|
|
session_id=data.get("session_id", "default"),
|
|
system_prompt=data.get("system_prompt", ""),
|
|
refresh_session=bool(data.get("refresh_session", False)),
|
|
llm_config=data.get("llm_config"),
|
|
)
|
|
|
|
@PromptServer.instance.routes.post("/chat_optimize/chat")
|
|
async def chat_endpoint(request):
|
|
payload = await request.json()
|
|
loop = asyncio.get_event_loop()
|
|
try:
|
|
assistant_response, readable_history = await loop.run_in_executor(None, _chat_action, payload)
|
|
return web.json_response(
|
|
{
|
|
"assistant_response": assistant_response,
|
|
"readable_history": readable_history,
|
|
}
|
|
)
|
|
except Exception as exc:
|
|
return web.json_response({"error": str(exc)}, status=500)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ChatNode": ChatNode,
|
|
"LLMConfigNode": LLMConfigNode,
|
|
"ChatHistoryViewer": ChatHistoryViewer,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ChatNode": "Chat (LLM)",
|
|
"LLMConfigNode": "LLM Config",
|
|
"ChatHistoryViewer": "Chat History Viewer",
|
|
}
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|