diff --git a/nodes/image_to_prompt.py b/nodes/image_to_prompt.py index ddf16f4..93d60bb 100644 --- a/nodes/image_to_prompt.py +++ b/nodes/image_to_prompt.py @@ -45,6 +45,14 @@ class MoltbotImageToPrompt: def __init__(self): self.llm_client = LLMClient() + def _get_request_llm_client(self): + # CRITICAL: this node instance can persist across multiple UI runs. + # Refresh the default LLMClient per call so provider/key changes from Settings/UI + # apply without restarting ComfyUI. Preserve injected mocks/fakes for tests. + if isinstance(self.llm_client, LLMClient): + self.llm_client = LLMClient() + return self.llm_client + @classmethod def INPUT_TYPES(cls): return { @@ -147,9 +155,11 @@ Do not use markdown blocks. try: # 4. Call Vision LLM logger.info("Sending vision request to LLM...") + # IMPORTANT: resolve client at call time to avoid stale provider/key state. + llm_client = self._get_request_llm_client() # Using updated client signature - response = self.llm_client.complete( + response = llm_client.complete( system=system_prompt, user_message=user_message, image_base64=image_b64 ) diff --git a/services/automation_composer.py b/services/automation_composer.py index 42b32ee..e0df465 100644 --- a/services/automation_composer.py +++ b/services/automation_composer.py @@ -40,6 +40,14 @@ class AutomationComposerService: def __init__(self): self.llm_client = LLMClient() + def _get_request_llm_client(self): + # CRITICAL: Assist compose service is held by long-lived route handlers. + # Refresh the default LLMClient per request so UI-saved provider/key updates + # take effect without backend restart. Keep injected fakes intact for tests. + if isinstance(self.llm_client, LLMClient): + self.llm_client = LLMClient() + return self.llm_client + def compose_payload( self, *, @@ -184,8 +192,10 @@ class AutomationComposerService: system = self._build_system_prompt(kind, template_id, profile_id) user = self._build_user_prompt(intent, inputs_hint) + # IMPORTANT: resolve client at request time to avoid stale provider/key. + llm_client = self._get_request_llm_client() - response = self.llm_client.complete( + response = llm_client.complete( system=system, user_message=user, tools=tools, diff --git a/services/planner.py b/services/planner.py index 442f913..01e2ea2 100644 --- a/services/planner.py +++ b/services/planner.py @@ -53,6 +53,15 @@ class PlannerService: def __init__(self): self.llm_client = LLMClient() + def _get_request_llm_client(self): + # CRITICAL: refresh the default LLMClient per request. + # Assist handlers are long-lived singletons, so caching the startup client here + # causes stale provider/key state after UI Save (requires backend restart). + # Keep custom test fakes/injected clients intact by only rotating real LLMClient. + if isinstance(self.llm_client, LLMClient): + self.llm_client = LLMClient() + return self.llm_client + def plan_generation( self, profile_id: str, @@ -110,6 +119,9 @@ Style: {style_directives} """ try: + # IMPORTANT: resolve client at request time (not service init time). + # This keeps Planner aligned with the latest runtime config + server-side key store. + llm_client = self._get_request_llm_client() # F25: Check if tool calling is enabled use_tool_calling = ( TOOL_CALLING_AVAILABLE @@ -131,7 +143,7 @@ Style: {style_directives} tools = [PLANNER_TOOL_SCHEMA] # Call LLM with tool - response = self.llm_client.complete( + response = llm_client.complete( system_prompt, user_message, tools=tools, tool_choice="auto" ) @@ -169,7 +181,7 @@ Style: {style_directives} else: # Traditional mode: Call LLM normally logger.info(f"Sending request to LLM for profile {profile_id}...") - response = self.llm_client.complete( + response = llm_client.complete( system_prompt, user_message, streaming=on_text_delta is not None, diff --git a/services/plugins/builtin/audit_log.py b/services/plugins/builtin/audit_log.py index a9b3de6..8efa8ed 100644 --- a/services/plugins/builtin/audit_log.py +++ b/services/plugins/builtin/audit_log.py @@ -27,7 +27,14 @@ class AuditLogPlugin: """Hook: llm.audit_request (PARALLEL).""" try: # R28: Build structured audit event - from services.audit_events import build_audit_event, emit_audit_event + # CRITICAL: keep package-relative import first for ComfyUI custom-node loaders. + # Some runtime contexts do not expose a top-level `services` package, which + # produces noisy non-fatal audit errors (`No module named 'services.audit_events'`). + # Fallback to `services.*` only for test/direct-import contexts. + try: + from ...audit_events import build_audit_event, emit_audit_event + except ImportError: + from services.audit_events import build_audit_event, emit_audit_event event = build_audit_event( event_type="llm.request", diff --git a/services/refiner.py b/services/refiner.py index f12a02b..0bc81ef 100644 --- a/services/refiner.py +++ b/services/refiner.py @@ -44,6 +44,15 @@ class RefinerService: def __init__(self): self.llm_client = LLMClient() + def _get_request_llm_client(self): + # CRITICAL: refresh the default LLMClient per request. + # Refiner shares the same long-lived assist handler lifecycle as Planner; keeping + # the startup client causes stale provider/key state after UI Save. + # Preserve injected fakes by only rotating real LLMClient instances. + if isinstance(self.llm_client, LLMClient): + self.llm_client = LLMClient() + return self.llm_client + def refine_prompt( self, image_b64: str, @@ -107,6 +116,9 @@ Issue: {issue} """ try: + # IMPORTANT: resolve client at request time so UI-saved provider/key changes + # apply without restarting ComfyUI. + llm_client = self._get_request_llm_client() # F25: Optional tool calling (OpenAI-compat only; fallback to JSON parsing) use_tool_calling = ( TOOL_CALLING_AVAILABLE @@ -124,7 +136,7 @@ Issue: {issue} except ImportError: tools = [REFINER_TOOL_SCHEMA] - response = self.llm_client.complete( + response = llm_client.complete( system=system_prompt, user_message=user_message, image_base64=image_b64, @@ -175,7 +187,7 @@ Issue: {issue} data = extract_json_object(content) else: # 5. Call Vision LLM (traditional JSON) - response = self.llm_client.complete( + response = llm_client.complete( system=system_prompt, user_message=user_message, image_base64=image_b64, diff --git a/tests/test_assist_llm_client_hot_reload.py b/tests/test_assist_llm_client_hot_reload.py new file mode 100644 index 0000000..cc4e206 --- /dev/null +++ b/tests/test_assist_llm_client_hot_reload.py @@ -0,0 +1,138 @@ +import json +import unittest +from unittest.mock import patch + + +class _PlannerDynamicFakeLLMClient: + _next_id = 0 + + def __init__(self): + type(self)._next_id += 1 + self.instance_id = type(self)._next_id + self.calls = [] + + def complete(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return { + "text": json.dumps( + { + "positive_prompt": f"p-{self.instance_id}", + "negative_prompt": "", + "params": { + "width": 1024, + "height": 1024, + "steps": 20, + "cfg": 7.0, + "sampler_name": "euler", + "scheduler": "normal", + }, + } + ), + "raw": {}, + } + + +class _RefinerDynamicFakeLLMClient: + _next_id = 0 + + def __init__(self): + type(self)._next_id += 1 + self.instance_id = type(self)._next_id + self.calls = [] + + def complete(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return { + "text": json.dumps( + { + "refined_positive": f"rp-{self.instance_id}", + "refined_negative": "", + "param_patch": {"steps": 25}, + "rationale": "ok", + } + ), + "raw": {}, + } + + +class _InjectedFakeLLMClient: + def __init__(self): + self.calls = [] + + def complete(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return { + "text": json.dumps( + { + "positive_prompt": "custom", + "negative_prompt": "", + "params": {"width": 1024, "height": 1024}, + } + ), + "raw": {}, + } + + +class TestAssistLLMClientHotReload(unittest.TestCase): + def test_planner_refreshes_default_llm_client_per_request(self): + import services.planner as planner_mod + + _PlannerDynamicFakeLLMClient._next_id = 0 + with patch.object(planner_mod, "LLMClient", _PlannerDynamicFakeLLMClient): + planner = planner_mod.PlannerService() + + pos1, _, _ = planner.plan_generation("SDXL-v1", "req", "style", seed=1) + first_client = planner.llm_client + pos2, _, _ = planner.plan_generation("SDXL-v1", "req", "style", seed=2) + second_client = planner.llm_client + + self.assertNotEqual(pos1, pos2) + self.assertIsNot(first_client, second_client) + self.assertEqual(first_client.instance_id, 2) + self.assertEqual(second_client.instance_id, 3) + + def test_refiner_refreshes_default_llm_client_per_request(self): + import services.refiner as refiner_mod + + _RefinerDynamicFakeLLMClient._next_id = 0 + with patch.object(refiner_mod, "LLMClient", _RefinerDynamicFakeLLMClient): + refiner = refiner_mod.RefinerService() + + rp1, _, _, _ = refiner.refine_prompt( + image_b64="dummy", + orig_positive="op", + orig_negative="on", + issue="fix", + params_json=json.dumps({"width": 1024, "height": 1024}), + ) + first_client = refiner.llm_client + rp2, _, _, _ = refiner.refine_prompt( + image_b64="dummy", + orig_positive="op", + orig_negative="on", + issue="fix", + params_json=json.dumps({"width": 1024, "height": 1024}), + ) + second_client = refiner.llm_client + + self.assertNotEqual(rp1, rp2) + self.assertIsNot(first_client, second_client) + self.assertEqual(first_client.instance_id, 2) + self.assertEqual(second_client.instance_id, 3) + + def test_planner_keeps_injected_custom_llm_client(self): + from services.planner import PlannerService + + planner = PlannerService() + fake = _InjectedFakeLLMClient() + planner.llm_client = fake + + pos, _, _ = planner.plan_generation("SDXL-v1", "req", "style", seed=1) + + self.assertEqual(pos, "custom") + self.assertIs(planner.llm_client, fake) + self.assertEqual(len(fake.calls), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_llm_client_hot_reload_non_assist.py b/tests/test_llm_client_hot_reload_non_assist.py new file mode 100644 index 0000000..b01d5fd --- /dev/null +++ b/tests/test_llm_client_hot_reload_non_assist.py @@ -0,0 +1,114 @@ +import json +import os +import unittest +from unittest.mock import patch + + +class _DynamicComposerLLMClient: + _next_id = 0 + + def __init__(self): + type(self)._next_id += 1 + self.instance_id = type(self)._next_id + + def complete(self, *args, **kwargs): + # No tool call -> compose path falls back to deterministic payload. + return { + "text": "", + "raw": {"choices": [{"message": {"content": f"plain-{self.instance_id}"}}]}, + } + + +class _DynamicVisionLLMClient: + _next_id = 0 + + def __init__(self): + type(self)._next_id += 1 + self.instance_id = type(self)._next_id + + def complete(self, *args, **kwargs): + return { + "text": json.dumps( + { + "caption": f"caption-{self.instance_id}", + "tags": ["tag1", "tag2"], + "prompt_suggestion": f"prompt-{self.instance_id}", + } + ), + "raw": {}, + } + + +class TestLLMClientHotReloadNonAssist(unittest.TestCase): + def test_automation_composer_refreshes_default_llm_client_per_request(self): + import services.automation_composer as composer_mod + + _DynamicComposerLLMClient._next_id = 0 + with ( + patch.object(composer_mod, "LLMClient", _DynamicComposerLLMClient), + patch.object(composer_mod, "TOOL_CALLING_AVAILABLE", True), + patch.object(composer_mod, "is_template_allowed", return_value=True), + patch.dict(os.environ, {"OPENCLAW_ENABLE_TOOL_CALLING": "1"}), + ): + svc = composer_mod.AutomationComposerService() + first_init_client = svc.llm_client + + res1 = svc.compose_payload( + kind="trigger", + template_id="tmpl", + intent="compose 1", + inputs_hint={"requirements": "a"}, + ) + first_request_client = svc.llm_client + res2 = svc.compose_payload( + kind="trigger", + template_id="tmpl", + intent="compose 2", + inputs_hint={"requirements": "b"}, + ) + second_request_client = svc.llm_client + + self.assertIsNot(first_init_client, first_request_client) + self.assertIsNot(first_request_client, second_request_client) + self.assertEqual(first_request_client.instance_id, 2) + self.assertEqual(second_request_client.instance_id, 3) + self.assertFalse(res1["used_tool_calling"]) + self.assertFalse(res2["used_tool_calling"]) + self.assertTrue(any("tool_call_fallback" in w for w in res1["warnings"])) + + def test_image_to_prompt_refreshes_default_llm_client_per_request(self): + import nodes.image_to_prompt as vision_mod + + _DynamicVisionLLMClient._next_id = 0 + with patch.object(vision_mod, "LLMClient", _DynamicVisionLLMClient): + node = vision_mod.MoltbotImageToPrompt() + init_client = node.llm_client + + with patch.object(node, "_tensor_to_base64_png", return_value="ZmFrZQ=="): + cap1, tags1, prompt1 = node.generate_prompt( + image=object(), + goal="goal", + detail_level="medium", + max_image_side=512, + ) + first_request_client = node.llm_client + cap2, tags2, prompt2 = node.generate_prompt( + image=object(), + goal="goal", + detail_level="medium", + max_image_side=512, + ) + second_request_client = node.llm_client + + self.assertIsNot(init_client, first_request_client) + self.assertIsNot(first_request_client, second_request_client) + self.assertEqual(first_request_client.instance_id, 2) + self.assertEqual(second_request_client.instance_id, 3) + self.assertNotEqual(cap1, cap2) + self.assertEqual(tags1, "tag1, tag2") + self.assertNotEqual(prompt1, prompt2) + self.assertEqual(tags2, "tag1, tag2") + + +if __name__ == "__main__": + unittest.main()