diff --git a/llm_be/chat_backend/consumers_graph.py b/llm_be/chat_backend/consumers_graph.py index 7bcf195..50746bf 100644 --- a/llm_be/chat_backend/consumers_graph.py +++ b/llm_be/chat_backend/consumers_graph.py @@ -532,44 +532,89 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): status_token = set_status_emitter(_send_status) try: await emit_status("queued") - # Run Graph (moderation emits moderating; grounding emits evaluating/…) - final_state = await app.ainvoke(initial_state) - print("Final State: ", final_state) - response_generator_or_dict = final_state["response_generator"] - print("Response Generator: ", response_generator_or_dict) + # Agentic path (#63) — same gate as consumers.py; skip the + # fixed LangGraph when the heuristic says multi-step. + from chat_backend.services.agent import ( + run_agentic_turn, + should_use_agent, + ) full_response = "" tokens_in = tokens_out = None + citations = [] + final_model = resolved_model - if isinstance(response_generator_or_dict, dict): - content = response_generator_or_dict.get("content", "") - await self.send_json_message( - json.dumps(response_generator_or_dict) + if should_use_agent(message): + try: + await enforce_feature_gate(chat_user, "agentic_tasks") + except FeatureNotAllowed as exc: + await self.send_json_message( + json.dumps( + { + "type": "error", + "code": exc.code, + "content": exc.message, + } + ) + ) + await self.send("END_OF_THE_STREAM_ENDER_GAME_42") + return + + async def _ws_send(raw: str): + await self.send_json_message(raw) + + _run, answer = await run_agentic_turn( + user=chat_user, + scope=tenant_scope, + conversation_id=conversation_id, + goal=message, + prompt=prompt_instance, + ws_send=_ws_send, ) - full_response = content - tokens_in, tokens_out = extract_token_usage( - response_generator_or_dict - ) - else: + final_model = _run.model_orchestrator or resolved_model await emit_status("writing") - usage = TokenUsageCollector() - async for chunk in aiter_text_chunks( - response_generator_or_dict, usage - ): - full_response += chunk - await self.send_json_message(chunk) - tokens_in, tokens_out = usage.pair + await self.send_json_message(answer) + full_response = answer + else: + # Run Graph (moderation emits moderating; grounding emits evaluating/…) + final_state = await app.ainvoke(initial_state) + print("Final State: ", final_state) + + response_generator_or_dict = final_state["response_generator"] + print("Response Generator: ", response_generator_or_dict) + + if isinstance(response_generator_or_dict, dict): + content = response_generator_or_dict.get("content", "") + await self.send_json_message( + json.dumps(response_generator_or_dict) + ) + full_response = content + tokens_in, tokens_out = extract_token_usage( + response_generator_or_dict + ) + else: + await emit_status("writing") + usage = TokenUsageCollector() + async for chunk in aiter_text_chunks( + response_generator_or_dict, usage + ): + full_response += chunk + await self.send_json_message(chunk) + tokens_in, tokens_out = usage.pair + + citations = final_state.get("citations") or [] + final_model = ( + final_state.get("resolved_model") or resolved_model + ) await self.send("END_OF_THE_STREAM_ENDER_GAME_42") - citations = final_state.get("citations") or [] if citations: await self.send_json_message( json.dumps(citations_frame(citations)) ) - final_model = final_state.get("resolved_model") or resolved_model if final_model and final_model != prompt_metric.model_name: prompt_metric.model_name = final_model await database_sync_to_async(prompt_metric.save)( diff --git a/llm_be/chat_backend/services/agent/runner.py b/llm_be/chat_backend/services/agent/runner.py index 098a95f..f325ec5 100644 --- a/llm_be/chat_backend/services/agent/runner.py +++ b/llm_be/chat_backend/services/agent/runner.py @@ -11,6 +11,7 @@ full history for clients that missed frames. from __future__ import annotations +import json import logging from asgiref.sync import sync_to_async @@ -103,11 +104,19 @@ def _upsert_step(run_id: int, event_type: str, data: dict): return step -async def _broadcast(run, event_type: str, data: dict) -> None: +async def _broadcast(run, event_type: str, data: dict, *, ws_send=None) -> None: + """Emit one agent frame both on the originating WS (if bound) and the + run's channel-layer group (so a reconnect / ``GET /api/agent_runs//`` + still sees it).""" + frame = agent_frame(event_type, {"run_id": str(run.pk), **data}) + if ws_send is not None: + try: + await ws_send(json.dumps(frame)) + except Exception: # pragma: no cover - WS send must never break a run + logger.exception("Failed to send agent frame on WS run=%s type=%s", run.pk, event_type) layer = get_channel_layer() if layer is None: return - frame = agent_frame(event_type, {"run_id": str(run.pk), **data}) try: await layer.group_send( run.channel_group_name(), @@ -131,8 +140,15 @@ def _build_llms(tools: list): return planner_llm, subagent_llm_factory -async def execute_agent_run(run_id: int) -> None: - """Load, execute, and persist the outcome of one :class:`AgentRun`.""" +async def execute_agent_run(run_id: int, *, ws_send=None) -> None: + """Load, execute, and persist the outcome of one :class:`AgentRun`. + + ``ws_send`` is optional — set by :func:`run_agentic_turn` when a run is + started inline on a live WS connection so frames land there immediately, + in addition to the channel-layer group broadcast every run always gets + (for reconnects / ``GET /api/agent_runs//`` backfill). Background + dispatch via :mod:`chat_backend.services.agent.tasks` leaves it unset. + """ run = await _load_run(run_id) if run is None: logger.error("execute_agent_run: AgentRun %s not found", run_id) @@ -146,7 +162,12 @@ async def execute_agent_run(run_id: int) -> None: run.status = AgentRunModel.Status.PLANNING run.started_at = timezone.now() await _save_run(run, ["status", "started_at"]) - await _broadcast(run, "run_started", {"title": run.title or run.goal[:80], "status": run.status}) + await _broadcast( + run, + "run_started", + {"title": run.title or run.goal[:80], "status": run.status}, + ws_send=ws_send, + ) scope, workspace = await _resolve_scope_and_workspace( run.user, run.conversation_id @@ -166,7 +187,7 @@ async def execute_agent_run(run_id: int) -> None: async def on_event(event_type: str, data: dict) -> None: if event_type in ("step_started", "step_completed", "step_failed"): await _upsert_step(run.pk, event_type, data) - await _broadcast(run, event_type, data) + await _broadcast(run, event_type, data, ws_send=ws_send) async def is_cancelled() -> bool: return await _refresh_cancel_flag(run.pk) @@ -210,6 +231,7 @@ async def execute_agent_run(run_id: int) -> None: run, "run_completed", {"status": run.status, "result": final_answer, "title": run.title}, + ws_send=ws_send, ) except AgentCancelled: run.status = AgentRunModel.Status.CANCELLED @@ -219,7 +241,9 @@ async def execute_agent_run(run_id: int) -> None: await _save_run( run, ["status", "completed_at", "tool_call_count", "iteration_count"] ) - await _broadcast(run, "run_completed", {"status": run.status, "title": run.title}) + await _broadcast( + run, "run_completed", {"status": run.status, "title": run.title}, ws_send=ws_send + ) except AgentRunLimitExceeded as exc: run.status = AgentRunModel.Status.FAILED run.error = str(exc) @@ -231,7 +255,10 @@ async def execute_agent_run(run_id: int) -> None: ["status", "error", "completed_at", "tool_call_count", "iteration_count"], ) await _broadcast( - run, "run_completed", {"status": run.status, "error": str(exc), "title": run.title} + run, + "run_completed", + {"status": run.status, "error": str(exc), "title": run.title}, + ws_send=ws_send, ) except Exception as exc: # pragma: no cover - defensive top-level guard logger.exception("Agent run %s failed", run_id) @@ -245,7 +272,10 @@ async def execute_agent_run(run_id: int) -> None: ["status", "error", "completed_at", "tool_call_count", "iteration_count"], ) await _broadcast( - run, "run_completed", {"status": run.status, "error": str(exc), "title": run.title} + run, + "run_completed", + {"status": run.status, "error": str(exc), "title": run.title}, + ws_send=ws_send, ) @@ -260,8 +290,12 @@ async def run_agentic_turn( ): """Create an AgentRun and execute it inline for the originating WS turn. - ``ws_send`` / ``scope`` accepted for API compatibility with consumers. - Progress broadcasts via the channel-layer group. Returns ``(run, answer)``. + ``scope`` is accepted for API compatibility with consumers (tenant scope + was already validated there) but re-derived from ``user``/ + ``conversation_id`` inside :func:`execute_agent_run` to keep a single + resolution path. ``ws_send`` — when provided — receives every + ``agent_*`` frame live, in addition to the channel-layer group broadcast + every run always gets. Returns ``(run, answer)``. """ from django.conf import settings from chat_backend.models import AgentRun @@ -273,7 +307,6 @@ async def run_agentic_turn( from chat_backend.services.status_context import emit_status del scope - del ws_send if not getattr(settings, "ALLOW_AGENTIC_TASKS", False): raise RuntimeError("ALLOW_AGENTIC_TASKS is disabled") @@ -295,7 +328,7 @@ async def run_agentic_turn( ), ) await emit_status("evaluating", "Planning multi-step task") - await execute_agent_run(run.pk) + await execute_agent_run(run.pk, ws_send=ws_send) finished = await _load_run(run.pk) answer = ( (finished.result if finished else "") diff --git a/llm_be/chat_backend/tests/test_agent_orchestrator.py b/llm_be/chat_backend/tests/test_agent_orchestrator.py new file mode 100644 index 0000000..f105016 --- /dev/null +++ b/llm_be/chat_backend/tests/test_agent_orchestrator.py @@ -0,0 +1,343 @@ +"""Offline unit tests for the agent orchestrator + runner (#63). + +No live Ollama/Redis/network — planner/subagent LLMs and tools are fakes, +following the ``async_to_sync`` convention already used in +``test_agent_tools.py`` for exercising async code from Django's sync +``TestCase``. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest import mock + +from asgiref.sync import async_to_sync +from django.test import SimpleTestCase, TestCase, override_settings + +from chat_backend.services.agent.orchestrator import ( + AgentCancelled, + AgentOrchestrator, + AgentRunLimitExceeded, + RunLimits, + fallback_plan, + parse_plan, +) +from chat_backend.tests.factories import make_conversation, make_prompt, make_user + + +class FakeMessage: + def __init__(self, content="", tool_calls=None): + self.content = content + self.tool_calls = tool_calls or [] + + +class FakePlannerLLM: + """Returns a fixed JSON plan, then a fixed synthesis string on the 2nd call.""" + + def __init__(self, plan_json: str, synthesis_text: str = "Final answer."): + self.plan_json = plan_json + self.synthesis_text = synthesis_text + self.calls = 0 + + async def ainvoke(self, messages): + self.calls += 1 + if self.calls == 1: + return FakeMessage(content=self.plan_json) + return FakeMessage(content=self.synthesis_text) + + +class FailingLLM: + async def ainvoke(self, messages): + raise RuntimeError("ollama down") + + +def make_fake_tool(name: str, output: str = "tool result"): + async def _ainvoke(_input): + return output + + return SimpleNamespace(name=name, ainvoke=_ainvoke) + + +class PlanParsingTestCase(SimpleTestCase): + def test_parse_plan_valid_json(self): + raw = ( + '{"title": "T", "steps": [' + '{"step_id": "s1", "title": "Search", "tool": "web_search"},' + '{"step_id": "s2", "title": "Write", "tool": null}]}' + ) + plan = parse_plan(raw, goal="goal", max_steps=8) + self.assertEqual(plan.title, "T") + self.assertEqual(len(plan.steps), 2) + self.assertEqual(plan.steps[0].tool, "web_search") + self.assertIsNone(plan.steps[1].tool) + + def test_parse_plan_strips_code_fences(self): + raw = '```json\n{"title": "T", "steps": [{"step_id": "s1", "title": "X"}]}\n```' + plan = parse_plan(raw, goal="goal", max_steps=8) + self.assertEqual(len(plan.steps), 1) + + def test_parse_plan_falls_back_on_garbage(self): + plan = parse_plan("not json at all", goal="latest news today", max_steps=8) + self.assertTrue(plan.steps) + self.assertEqual(plan.steps[0].tool, "web_search") + + def test_fallback_plan_no_tool_hints(self): + plan = fallback_plan("write a haiku about rain", max_steps=8) + self.assertEqual(len(plan.steps), 1) + self.assertIsNone(plan.steps[0].tool) + + def test_parse_plan_caps_at_max_steps(self): + steps = ",".join( + f'{{"step_id": "s{i}", "title": "Step {i}", "tool": null}}' for i in range(10) + ) + raw = f'{{"title": "T", "steps": [{steps}]}}' + plan = parse_plan(raw, goal="goal", max_steps=3) + self.assertEqual(len(plan.steps), 3) + + +class AgentOrchestratorRunTestCase(SimpleTestCase): + def _run(self, orchestrator): + return async_to_sync(orchestrator.run)() + + def test_full_run_with_tool_and_synthesis(self): + plan_json = ( + '{"title": "Research", "steps": [' + '{"step_id": "s1", "title": "Search web", "tool": "web_search", ' + '"tool_input": {"query": "x"}}]}' + ) + planner_llm = FakePlannerLLM(plan_json, synthesis_text="Here is the synthesis.") + tool = make_fake_tool("web_search", output="search results about x") + events: list[tuple[str, dict]] = [] + + async def on_event(event_type, data): + events.append((event_type, data)) + + orchestrator = AgentOrchestrator( + goal="Research x", + history_text="", + planner_llm=planner_llm, + subagent_llm_factory=lambda: FailingLLM(), + tools=[tool], + limits=RunLimits(), + on_event=on_event, + ) + plan, results, answer = self._run(orchestrator) + + self.assertEqual(plan.title, "Research") + self.assertEqual(results["s1"], "search results about x") + self.assertEqual(answer, "Here is the synthesis.") + event_types = [e[0] for e in events] + self.assertEqual( + event_types, ["plan_ready", "step_started", "step_completed"] + ) + + def test_planner_failure_uses_fallback_plan(self): + # plan() itself is exercised directly here (rather than the full + # run(), which also calls the same LLM for synthesis) so a fully + # unavailable planner model is isolated to the planning step. + orchestrator = AgentOrchestrator( + goal="latest news today", + history_text="", + planner_llm=FailingLLM(), + subagent_llm_factory=lambda: FailingLLM(), + tools=[make_fake_tool("web_search")], + limits=RunLimits(), + ) + plan = async_to_sync(orchestrator.plan)() + self.assertTrue(plan.steps) + self.assertEqual(plan.steps[0].tool, "web_search") + + def test_step_failure_is_recorded_but_run_continues(self): + plan_json = ( + '{"title": "T", "steps": [' + '{"step_id": "s1", "title": "Broken", "tool": "broken_tool"}]}' + ) + planner_llm = FakePlannerLLM(plan_json, synthesis_text="Done despite failure.") + orchestrator = AgentOrchestrator( + goal="do a thing", + history_text="", + planner_llm=planner_llm, + subagent_llm_factory=lambda: FailingLLM(), + tools=[], # "broken_tool" is not registered + limits=RunLimits(), + ) + _plan, results, answer = self._run(orchestrator) + self.assertIn("Unknown tool", results["s1"]) + self.assertEqual(answer, "Done despite failure.") + + def test_cancellation_raises_agent_cancelled(self): + plan_json = ( + '{"title": "T", "steps": [' + '{"step_id": "s1", "title": "Search", "tool": "web_search"}]}' + ) + planner_llm = FakePlannerLLM(plan_json) + tool = make_fake_tool("web_search") + + async def is_cancelled(): + return True + + orchestrator = AgentOrchestrator( + goal="goal", + history_text="", + planner_llm=planner_llm, + subagent_llm_factory=lambda: FailingLLM(), + tools=[tool], + limits=RunLimits(), + is_cancelled=is_cancelled, + ) + with self.assertRaises(AgentCancelled): + self._run(orchestrator) + + def test_wall_clock_exceeded_fails_step_but_run_still_completes(self): + # A blown wall-clock budget fails the in-flight step (caught by the + # generic per-step exception handler in run()) rather than aborting + # the whole run — synthesis still runs over whatever is available. + plan_json = ( + '{"title": "T", "steps": [' + '{"step_id": "s1", "title": "Search", "tool": "web_search"}]}' + ) + planner_llm = FakePlannerLLM(plan_json, synthesis_text="Best effort answer.") + tool = make_fake_tool("web_search") + orchestrator = AgentOrchestrator( + goal="goal", + history_text="", + planner_llm=planner_llm, + subagent_llm_factory=lambda: FailingLLM(), + tools=[tool], + limits=RunLimits(wall_clock_seconds=1), + ) + # Force the deadline into the past without sleeping in the test. + orchestrator._deadline = 0.0 + _plan, results, answer = self._run(orchestrator) + self.assertIn("Wall-clock budget exceeded", results["s1"]) + self.assertEqual(answer, "Best effort answer.") + + def test_execute_step_raises_when_deadline_passed(self): + orchestrator = AgentOrchestrator( + goal="goal", + history_text="", + planner_llm=FailingLLM(), + subagent_llm_factory=lambda: FailingLLM(), + tools=[], + limits=RunLimits(), + ) + orchestrator._deadline = 0.0 + from chat_backend.services.agent.orchestrator import PlanStep + + with self.assertRaises(AgentRunLimitExceeded): + async_to_sync(orchestrator.execute_step)( + PlanStep(step_id="s1", title="x", tool="web_search"), {} + ) + + +@override_settings(ALLOW_AGENTIC_TASKS=True) +class RunAgenticTurnTestCase(TestCase): + def setUp(self): + self.user = make_user(email="agent-runner@example.com") + self.conversation = make_conversation(user=self.user) + self.prompt = make_prompt(self.conversation, message="Research the top 5 things") + + def _patch_llms(self, plan_json="", synthesis_text="All done."): + planner_llm = FakePlannerLLM( + plan_json + or '{"title": "T", "steps": [{"step_id": "s1", "title": "Answer", "tool": null}]}', + synthesis_text=synthesis_text, + ) + return mock.patch( + "chat_backend.services.agent.runner._build_llms", + return_value=(planner_llm, lambda: FailingLLM()), + ) + + def test_run_agentic_turn_completes_and_persists(self): + from chat_backend.models import AgentRun + + sent_frames = [] + + async def ws_send(raw): + sent_frames.append(raw) + + with self._patch_llms(synthesis_text="The answer is 42."): + run, answer = async_to_sync(self._call_run_agentic_turn)(ws_send) + + run.refresh_from_db() + self.assertEqual(run.status, AgentRun.Status.COMPLETED) + self.assertEqual(run.result, "The answer is 42.") + self.assertEqual(answer, "The answer is 42.") + self.assertTrue(sent_frames) # progress frames reached the WS callback + + async def _call_run_agentic_turn(self, ws_send): + from chat_backend.services.agent.runner import run_agentic_turn + from chat_backend.services.chat_tenant_scope import resolve_chat_company_scope + from asgiref.sync import sync_to_async + + scope = await sync_to_async(resolve_chat_company_scope)( + self.user, self.conversation.id + ) + return await run_agentic_turn( + user=self.user, + scope=scope, + conversation_id=self.conversation.id, + goal=self.prompt.message, + prompt=self.prompt, + ws_send=ws_send, + ) + + def test_disabled_flag_short_circuits(self): + from chat_backend.services.chat_tenant_scope import resolve_chat_company_scope + + with override_settings(ALLOW_AGENTIC_TASKS=False): + with self.assertRaises(RuntimeError): + async_to_sync(self._call_run_agentic_turn)(None) + + +@override_settings(ALLOW_AGENTIC_TASKS=True) +class ExecuteAgentRunTestCase(TestCase): + def setUp(self): + self.user = make_user(email="agent-bg@example.com") + self.conversation = make_conversation(user=self.user) + + def test_execute_agent_run_marks_cancelled(self): + from chat_backend.models import AgentRun, AgentStep + + run = AgentRun.objects.create( + user=self.user, + conversation=self.conversation, + goal="Research the top 5 things and compare", + status=AgentRun.Status.PENDING, + ) + run.cancel_requested = True + run.save(update_fields=["cancel_requested"]) + + planner_llm = FakePlannerLLM( + '{"title": "T", "steps": [{"step_id": "s1", "title": "Search", ' + '"tool": "web_search"}]}' + ) + with mock.patch( + "chat_backend.services.agent.runner._build_llms", + return_value=(planner_llm, lambda: FailingLLM()), + ), mock.patch( + "chat_backend.services.agent.runner.build_agent_tools", + return_value=[make_fake_tool("web_search")], + ): + from chat_backend.services.agent.runner import execute_agent_run + + async_to_sync(execute_agent_run)(run.pk) + + run.refresh_from_db() + self.assertEqual(run.status, AgentRun.Status.CANCELLED) + self.assertEqual(AgentStep.objects.filter(run=run).count(), 0) + + def test_execute_agent_run_skips_non_pending(self): + from chat_backend.models import AgentRun + from chat_backend.services.agent.runner import execute_agent_run + + run = AgentRun.objects.create( + user=self.user, + conversation=self.conversation, + goal="already running", + status=AgentRun.Status.RUNNING, + ) + # Should return immediately without touching the run. + async_to_sync(execute_agent_run)(run.pk) + run.refresh_from_db() + self.assertEqual(run.status, AgentRun.Status.RUNNING)