Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
35ffd590ba | ||
|
|
bef151c92c |
@@ -543,6 +543,51 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
|
||||
status_token = set_status_emitter(_send_status)
|
||||
try:
|
||||
await emit_status("queued")
|
||||
|
||||
# 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 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,
|
||||
)
|
||||
final_model = _run.model_orchestrator or resolved_model
|
||||
await emit_status("writing")
|
||||
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)
|
||||
@@ -550,9 +595,6 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
|
||||
response_generator_or_dict = final_state["response_generator"]
|
||||
print("Response Generator: ", response_generator_or_dict)
|
||||
|
||||
full_response = ""
|
||||
tokens_in = tokens_out = None
|
||||
|
||||
if isinstance(response_generator_or_dict, dict):
|
||||
content = response_generator_or_dict.get("content", "")
|
||||
await self.send_json_message(
|
||||
@@ -572,15 +614,18 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
|
||||
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)(
|
||||
|
||||
@@ -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/<id>/``
|
||||
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/<id>/`` 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 "")
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user