feat: multi-step ReAct diagnosis

pyproject.toml

@@ -11,6 +11,9 @@ dependencies = [
     "chonkie>=1.7.0",
     "fastapi>=0.139.2",
     "fpdf2>=2.8.7",
+    "langchain>=1.3.14",
+    "langchain-openai>=1.3.5",
+    "langgraph>=1.2.9",
     "mcp>=1.28.1",
     "openai>=2.46.0",
     "paho-mqtt>=2.1.0",

src/torq/agent/diagnose.py

@@ -1,7 +1,11 @@
-"""Diagnosis agent: retrieve manual + history context, call the LLM, return a Diagnosis.
+"""Diagnosis agent: retrieve manual + history context, reason, return a Diagnosis.
+
+The agent runs a multi-step reason/act loop: it searches manuals, then repair
+history, and may search again with a refined query before answering. If the
+agent loop fails for any reason it falls back to a single-shot diagnosis.
 
 Retrieval goes through the MCP knowledge server (proving plant data stays
-on-premise).  If the MCP server is unreachable the agent falls back to
+on-premise).  If the MCP server is unreachable retrieval falls back to
 calling ``ingest.search`` directly.
 """
 
@@ -11,7 +15,12 @@
 
 from openai import OpenAI
 
-from torq.agent.prompts import SYSTEM, build_user_prompt
+from torq.agent.prompts import (
+    REACT_SYSTEM,
+    SYSTEM,
+    build_react_task,
+    build_user_prompt,
+)
 from torq.agent.schemas import Diagnosis
 from torq.config import settings
 from torq.ingest import search  # direct fallback
@@ -131,37 +140,70 @@ def _chat(client: OpenAI, messages: list[dict], json_mode: bool = True):
     return client.chat.completions.create(model=settings.llm_model, messages=messages, stream=False)
 
 
-# ── diagnosis cache ──────────────────────────────────────────────────────────
-
-# Recent diagnoses keyed by (machine, fault_code, context). A repeat fault with
-# the same context within the TTL reuses the cached result and skips the
-# retrieval + LLM round-trip. Context is part of the key so a recurring fault
-# described differently (new symptom text) re-diagnoses instead of reusing a
-# stale answer. No lock: a rare race just costs a redundant diagnosis, it never
-# corrupts state. Copies go in and out so a caller mutating a Diagnosis cannot
-# poison the cache.
-_CACHE: dict[tuple[str, str, str], tuple[float, Diagnosis]] = {}
-
-
-# ── main entry point ─────────────────────────────────────────────────────────
-
+def _merge_sources(data: dict, extra: list[str]) -> None:
+    """Add retrieved source ids to data['sources'] without duplicates."""
+    data.setdefault("sources", [])
+    for s in extra:
+        if s and s not in data["sources"]:
+            data["sources"].append(s)
 
-def diagnose(fault_code: str, machine: str = "", context: str = "") -> Diagnosis:
-    key = (machine, fault_code, context)
-    ttl = settings.diagnose_cache_ttl
-    if ttl > 0:
-        hit = _CACHE.get(key)
-        if hit and time.monotonic() < hit[0]:
-            log.info("Diagnosis cache hit for %s %s (reused, no LLM call)", machine, fault_code)
-            return hit[1].model_copy(deep=True)
 
+# ── multi-step (ReAct) agent ─────────────────────────────────────────────────
+
+
+def _diagnose_react(fault_code: str, machine: str, context: str) -> Diagnosis:
+    """Reason/act loop: the LLM decides when to search manuals/history vs answer."""
+    # Imported lazily so a missing/broken install degrades to one-shot.
+    from langchain.agents import create_agent
+    from langchain_core.tools import tool
+    from langchain_openai import ChatOpenAI
+
+    collected: list[str] = []  # source ids surfaced by tool calls this run
+    steps: list[str] = []  # ordered record of what the agent searched
+
+    @tool
+    def search_manuals(query: str) -> str:
+        """Search the plant's OEM manuals for excerpts relevant to a fault or symptom."""
+        steps.append(f"Searched manuals: {query}")
+        txt, src = _join_manuals(_fetch_manuals(query))
+        collected.extend(src)
+        return txt or "no manual excerpts found"
+
+    @tool
+    def search_history(query: str) -> str:
+        """Search past repair records for fixes to similar faults."""
+        steps.append(f"Searched repair history: {query}")
+        txt, src = _join_history(_fetch_history(query, machine=machine))
+        collected.extend(src)
+        return txt or "no past repairs found"
+
+    model = ChatOpenAI(
+        model=settings.agent_model,
+        api_key=settings.llm_api_key,
+        base_url=settings.llm_base_url,
+        temperature=0,
+    )
+    agent = create_agent(model, [search_manuals, search_history], system_prompt=REACT_SYSTEM)
+    result = agent.invoke(
+        {"messages": [("user", build_react_task(fault_code, machine, context))]},
+        {"recursion_limit": settings.agent_max_steps * 2},
+    )
+    data = _parse_json(result["messages"][-1].content)
+    _merge_sources(data, collected)
+    data["investigation"] = steps
+    return Diagnosis(fault_code=fault_code, machine=machine, **data)
+
+
+# ── single-shot fallback ─────────────────────────────────────────────────────
+
+
+def _diagnose_oneshot(fault_code: str, machine: str, context: str) -> Diagnosis:
+    """One retrieval pass then one LLM call. Used when the agent loop fails."""
     query = f"{fault_code} {machine} {context}".strip()
 
-    manuals_hits = _fetch_manuals(query)
-    manuals_txt, m_src = _join_manuals(manuals_hits)
-
-    history_hits = _fetch_history(query, machine=machine)
-    history_txt, h_src = _join_history(history_hits)
+    manuals_txt, m_src = _join_manuals(_fetch_manuals(query))
+    history_txt, h_src = _join_history(_fetch_history(query, machine=machine))
+    steps = [f"Searched manuals: {query}", f"Searched repair history: {query}"]
 
     client = OpenAI(api_key=settings.llm_api_key, base_url=settings.llm_base_url)
     messages = [
@@ -177,13 +219,41 @@ def diagnose(fault_code: str, machine: str = "", context: str = "") -> Diagnosis
         resp = _chat(client, messages)
         data = _parse_json(resp.choices[0].message.content)
 
-    # Prefer retrieved sources; let the model add any it names.
-    data.setdefault("sources", [])
-    for s in m_src + h_src:
-        if s not in data["sources"]:
-            data["sources"].append(s)
+    _merge_sources(data, m_src + h_src)
+    data["investigation"] = steps
+    return Diagnosis(fault_code=fault_code, machine=machine, **data)
+
+
+# ── diagnosis cache ──────────────────────────────────────────────────────────
+
+# Recent diagnoses keyed by (machine, fault_code, context). A repeat fault with
+# the same context within the TTL reuses the cached result and skips the whole
+# reason/act loop. Context is part of the key so a recurring fault described
+# differently (new symptom text) re-diagnoses instead of reusing a stale answer.
+# No lock: a rare race just costs a redundant diagnosis, it never corrupts state.
+# Copies go in and out so a caller mutating a Diagnosis cannot poison the cache.
+_CACHE: dict[tuple[str, str, str], tuple[float, Diagnosis]] = {}
+
+
+# ── main entry point ─────────────────────────────────────────────────────────
+
+
+def diagnose(fault_code: str, machine: str = "", context: str = "") -> Diagnosis:
+    """Cached multi-step (ReAct) diagnosis, falling back to single-shot on failure."""
+    key = (machine, fault_code, context)
+    ttl = settings.diagnose_cache_ttl
+    if ttl > 0:
+        hit = _CACHE.get(key)
+        if hit and time.monotonic() < hit[0]:
+            log.info("Diagnosis cache hit for %s %s (reused, no LLM call)", machine, fault_code)
+            return hit[1].model_copy(deep=True)
+
+    try:
+        diag = _diagnose_react(fault_code, machine, context)
+    except Exception:  # noqa: BLE001 - any agent failure degrades to one-shot
+        log.warning("ReAct diagnosis failed, falling back to single-shot", exc_info=True)
+        diag = _diagnose_oneshot(fault_code, machine, context)
 
-    diag = Diagnosis(fault_code=fault_code, machine=machine, **data)
     if ttl > 0:
         _CACHE[key] = (time.monotonic() + ttl, diag.model_copy(deep=True))
     return diag

src/torq/agent/prompts.py

@@ -11,6 +11,30 @@
 )
 
 
+REACT_SYSTEM = (
+    "You are a maintenance diagnosis engine for industrial machines. Work in steps.\n"
+    "1. Call search_manuals ONCE with your best focused query about the fault.\n"
+    "2. Call search_history ONCE to find how similar faults were fixed before.\n"
+    "3. Reason about the root cause. Only search again if a result was empty or "
+    "clearly contradicts the fault, and only with a meaningfully different query. "
+    "Do not repeat similar searches.\n"
+    "Ground every claim in the retrieved context; do not invent part numbers.\n"
+    "When confident, STOP calling tools and reply with ONE JSON object only, no prose, "
+    "no markdown fences, matching:\n"
+    '{"root_cause": str, "confidence": float 0-1, "repair_steps": [str], '
+    '"parts": [str], "tools": [str], "safety_warnings": [str], "sources": [str]}'
+)
+
+
+def build_react_task(fault_code: str, machine: str, context: str) -> str:
+    return (
+        f"FAULT CODE: {fault_code}\n"
+        f"MACHINE: {machine or 'unknown'}\n"
+        f"SITUATION: {context or 'n/a'}\n\n"
+        "Diagnose this fault. Retrieve manuals and repair history first, then answer."
+    )
+
+
 def build_user_prompt(
     fault_code: str, machine: str, context: str, manuals: str, history: str
 ) -> str:

src/torq/agent/schemas.py

@@ -19,6 +19,7 @@ class Diagnosis(BaseModel):
     tools: list[str] = []
     safety_warnings: list[str] = []
     sources: list[str] = []
+    investigation: list[str] = []  # ordered retrieval steps the agent took
 
 
 class WorkOrder(BaseModel):
@@ -32,6 +33,7 @@ class WorkOrder(BaseModel):
     safety_warnings: list[str] = []
     required_skill: str = "general"
     sources: list[str] = []  # manual sections + past repairs the diagnosis cited
+    investigation: list[str] = []  # ordered retrieval steps the agent took
     content: dict[str, str] = {}  # language code (fr/ar/en) -> formatted text
     status: str = "pending"  # pending | approved | rejected | dispatched | resolved | failed
     assigned_to: str | None = None

src/torq/config.py

@@ -17,6 +17,10 @@ class Settings(BaseSettings):
     llm_api_key: str = ""
     llm_base_url: str = "https://api.deepseek.com"
     llm_model: str = "deepseek-reasoner"
+    # Model for the multi-step (ReAct) agent. Needs tool-calling support, which
+    # the reasoner model lacks, so default to the chat model.
+    agent_model: str = "deepseek-chat"
+    agent_max_steps: int = 8  # reason/act ceiling; generous so the agent answers vs erroring into fallback
 
     # Vector DB (Qdrant)
     qdrant_url: str = ""

src/torq/workorder/generate.py

@@ -79,6 +79,7 @@ def build_work_order(
         safety_warnings=diag.safety_warnings,
         required_skill=_required_skill(diag.fault_code),
         sources=diag.sources,
+        investigation=diag.investigation,
         content=content,
         confidence=diag.confidence,
         fault_arrived_at=fault_arrived_at,

tests/test_diagnose_cache.py

@@ -1,60 +1,55 @@
-"""Tests for the diagnosis cache: repeat faults skip the LLM within the TTL."""
+"""Tests for the diagnosis cache: repeat faults skip recomputation within the TTL."""
 
 import time
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 from torq.agent import diagnose as diag_mod
+from torq.agent.schemas import Diagnosis
 from torq.config import settings
 
 
-def _resp(content: str) -> MagicMock:
-    return MagicMock(choices=[MagicMock(message=MagicMock(content=content))])
-
-
-CANNED = '{"root_cause": "seized bearing", "repair_steps": ["swap bearing"]}'
+def _diag() -> Diagnosis:
+    # Fresh object per computation so cache copy-isolation is exercised.
+    return Diagnosis(fault_code="E-201", root_cause="seized bearing", repair_steps=["swap bearing"])
 
 
 class DiagnoseCacheTests(unittest.TestCase):
     def setUp(self) -> None:
         diag_mod._CACHE.clear()
-        # Patch retrieval + LLM so no vector DB or network is touched.
-        self._patchers = {
-            "man": patch("torq.agent.diagnose._fetch_manuals", return_value=[]),
-            "hist": patch("torq.agent.diagnose._fetch_history", return_value=[]),
-            "oai": patch("torq.agent.diagnose.OpenAI", return_value=MagicMock()),
-            "chat": patch("torq.agent.diagnose._chat", return_value=_resp(CANNED)),
-        }
-        started = {name: p.start() for name, p in self._patchers.items()}
-        self.mock_chat = started["chat"]
+        # Mock the whole diagnosis computation (ReAct path) so the cache is tested
+        # in isolation with no vector DB or network. call_count == cache misses.
+        self._patcher = patch(
+            "torq.agent.diagnose._diagnose_react", side_effect=lambda *a, **k: _diag()
+        )
+        self.mock_diag = self._patcher.start()
 
     def tearDown(self) -> None:
-        for p in self._patchers.values():
-            p.stop()
+        self._patcher.stop()
         diag_mod._CACHE.clear()
 
     def test_repeat_fault_served_from_cache(self) -> None:
         d1 = diag_mod.diagnose("E-201", "CM-350")
         d2 = diag_mod.diagnose("E-201", "CM-350")
-        self.assertEqual(self.mock_chat.call_count, 1)  # second call reused
+        self.assertEqual(self.mock_diag.call_count, 1)  # second call reused
         self.assertEqual(d1.root_cause, d2.root_cause)
 
     def test_distinct_keys_not_shared(self) -> None:
         diag_mod.diagnose("E-201", "CM-350")
         diag_mod.diagnose("J-108", "CM-350")  # different fault_code
         diag_mod.diagnose("E-201", "PK-9")  # different machine
-        self.assertEqual(self.mock_chat.call_count, 3)
+        self.assertEqual(self.mock_diag.call_count, 3)
 
     def test_ttl_zero_disables_cache(self) -> None:
         with patch.object(settings, "diagnose_cache_ttl", 0):
             diag_mod.diagnose("E-201", "CM-350")
             diag_mod.diagnose("E-201", "CM-350")
-        self.assertEqual(self.mock_chat.call_count, 2)
+        self.assertEqual(self.mock_diag.call_count, 2)
 
     def test_different_context_not_shared(self) -> None:
         diag_mod.diagnose("E-471", "CM-350", "overtemp after 6h runtime")
         diag_mod.diagnose("E-471", "CM-350", "recurring overtemp, lint buildup")
-        self.assertEqual(self.mock_chat.call_count, 2)  # context busts the cache
+        self.assertEqual(self.mock_diag.call_count, 2)  # context busts the cache
 
     def test_expired_entry_recomputes(self) -> None:
         diag_mod.diagnose("E-201", "CM-350")
@@ -63,7 +58,7 @@ def test_expired_entry_recomputes(self) -> None:
         _exp, diag = diag_mod._CACHE[key]
         diag_mod._CACHE[key] = (time.monotonic() - 1, diag)
         diag_mod.diagnose("E-201", "CM-350")
-        self.assertEqual(self.mock_chat.call_count, 2)
+        self.assertEqual(self.mock_diag.call_count, 2)
 
     def test_cache_copy_isolation(self) -> None:
         d1 = diag_mod.diagnose("E-201", "CM-350")