geng

app/src/doctor_workstation/services/repository.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import mimetypes
 import re
 import time
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from contextlib import suppress
 from datetime import date
 from io import BytesIO
@@ -13,11 +13,14 @@ from os import PathLike
 from pathlib import Path
 from typing import Any, Final, Literal, Protocol
 
-from doctor_workstation.core.errors import (
-    ApiProtocolError,
-    ApiTransportError,
-    AuthenticationExpiredError,
-)
+from doctor_workstation.core.errors import (
+    ApiBusinessError,
+    ApiHttpError,
+    ApiProtocolError,
+    ApiTimeoutError,
+    ApiTransportError,
+    AuthenticationExpiredError,
+)
 from doctor_workstation.core.models import (
     Appointment,
     CallTicket,
@@ -191,6 +194,16 @@ class DoctorRepository(Protocol):
     ) -> dict[str, Any]:
         """Ask the first-party diagnosis assistant; the server selects the model."""
 
+    def stream_diagnosis_ai(
+        self,
+        diagnosis_id: int,
+        prompt: str,
+        *,
+        task: str = "custom",
+        cancelled: Callable[[], bool] | None = None,
+    ) -> Iterator[dict[str, Any]]:
+        """Yield normalized ``start``/``delta``/``done`` assistant events."""
+
     def get_diagnosis_ai_analysis(
         self,
         diagnosis_id: int,
@@ -1328,25 +1341,7 @@ class RemoteDoctorRepository:
     ) -> dict[str, Any]:
         """Submit a diagnosis question to the first-party server assistant."""
 
-        if diagnosis_id <= 0:
-            raise ValueError("diagnosis_id must be positive")
-        clean_prompt = prompt.strip()
-        if not clean_prompt:
-            raise ValueError("prompt is required")
-        if len(clean_prompt) > 500:
-            raise ValueError("prompt must not exceed 500 characters")
-        clean_task = task.strip().lower() or "custom"
-        if clean_task not in {
-            "summary",
-            "tcm_pattern",
-            "prescription_review",
-            "medication_review",
-            "exam_review",
-            "complication_risk",
-            "guideline_review",
-            "custom",
-        }:
-            raise ValueError("task is not supported")
+        clean_prompt, clean_task = _diagnosis_ai_request(diagnosis_id, prompt, task)
         payload = _client_request(
             self.client,
             "post",
@@ -1359,6 +1354,63 @@ class RemoteDoctorRepository:
         )
         return dict(_require_mapping(payload, "tcm.diagnosis/aiAssistant"))
 
+    def stream_diagnosis_ai(
+        self,
+        diagnosis_id: int,
+        prompt: str,
+        *,
+        task: str = "custom",
+        cancelled: Callable[[], bool] | None = None,
+    ) -> Iterator[dict[str, Any]]:
+        """Stream a diagnosis answer, with one legacy fallback before first content."""
+
+        clean_prompt, clean_task = _diagnosis_ai_request(diagnosis_id, prompt, task)
+        body = {"id": diagnosis_id, "prompt": clean_prompt, "task": clean_task}
+        received_delta = False
+        received_done = False
+        try:
+            for raw_event in self.client.post_event_stream(
+                "tcm.diagnosis/aiAssistantStream",
+                body,
+                timeout=105.0,
+                cancelled=cancelled,
+            ):
+                if cancelled is not None and cancelled():
+                    return
+                event = _normalise_diagnosis_ai_event(raw_event)
+                if event is None:
+                    continue
+                kind = event["event"]
+                if kind == "delta":
+                    received_delta = True
+                elif kind == "done":
+                    received_done = True
+                yield event
+                if kind == "done":
+                    return
+            if cancelled is not None and cancelled():
+                return
+            if not received_done:
+                raise ApiProtocolError("AI assistant stream ended before done")
+        except (ApiHttpError, ApiProtocolError, ApiTimeoutError, ApiTransportError):
+            if received_delta or (cancelled is not None and cancelled()):
+                raise
+
+        # Older deployments do not expose the stream route. Submit exactly one
+        # request through the confirmed non-streaming endpoint in that case.
+        result = self.analyze_diagnosis_ai(
+            diagnosis_id,
+            clean_prompt,
+            task=clean_task,
+        )
+        if cancelled is not None and cancelled():
+            return
+        yield {"event": "start", "fallback": True}
+        answer = str(result.get("answer") or result.get("content") or "")
+        if answer:
+            yield {"event": "delta", "text": answer, "fallback": True}
+        yield {**result, "event": "done", "fallback": True}
+
     def get_diagnosis_ai_analysis(
         self,
         diagnosis_id: int,
@@ -2638,7 +2690,66 @@ class RemoteDoctorRepository:
         return self.update_prescription_template(template, changes, **fields)
 
 
-def _page_params(page_no: int, page_size: int, filters: Mapping[str, Any]) -> dict[str, Any]:
+def _diagnosis_ai_request(diagnosis_id: int, prompt: str, task: str) -> tuple[str, str]:
+    if diagnosis_id <= 0:
+        raise ValueError("diagnosis_id must be positive")
+    clean_prompt = prompt.strip()
+    if not clean_prompt:
+        raise ValueError("prompt is required")
+    if len(clean_prompt) > 500:
+        raise ValueError("prompt must not exceed 500 characters")
+    clean_task = task.strip().lower() or "custom"
+    if clean_task not in {
+        "summary",
+        "tcm_pattern",
+        "prescription_review",
+        "medication_review",
+        "exam_review",
+        "complication_risk",
+        "guideline_review",
+        "custom",
+    }:
+        raise ValueError("task is not supported")
+    return clean_prompt, clean_task
+
+
+def _normalise_diagnosis_ai_event(raw: Mapping[str, Any]) -> dict[str, Any] | None:
+    kind = str(raw.get("event") or "message").strip().lower()
+    data = raw.get("data")
+    payload = dict(data) if isinstance(data, Mapping) else {}
+    if kind == "message":
+        kind = str(payload.get("event") or payload.get("type") or "message").lower()
+    if kind == "start":
+        payload.pop("event", None)
+        payload.pop("type", None)
+        return {**payload, "event": "start"}
+    if kind == "delta":
+        text = (
+            data
+            if isinstance(data, str)
+            else payload.get("delta")
+            or payload.get("content")
+            or payload.get("text")
+            or ""
+        )
+        if not isinstance(text, str):
+            raise ApiProtocolError("AI assistant delta content must be text", data=data)
+        return {"event": "delta", "text": text}
+    if kind == "done":
+        payload.pop("event", None)
+        payload.pop("type", None)
+        return {**payload, "event": "done"}
+    if kind == "error":
+        message = (
+            data
+            if isinstance(data, str)
+            else payload.get("message") or payload.get("msg") or payload.get("error")
+        )
+        raise ApiBusinessError(str(message or "AI assistant stream failed"), data=data)
+    return None
+
+
+def _page_params(page_no: int, page_size: int, filters: Mapping[str, Any]) -> dict[str, Any]:
     if page_no < 1 or page_size < 1:
         raise ValueError("page_no and page_size must be positive")
     result = {