.

runners/service.py

@@ -1,7 +1,7 @@
 # service.py
 import asyncio
 from dataclasses import dataclass
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Callable, Awaitable
 
 from config import settings
 from models import LLMServiceObj, ResultObj
@@ -20,72 +20,76 @@ class LLMService:
     Python/Gradio equivalent of your .NET LLMService.
     Keeps identical field names and queue semantics when talking to RabbitMQ.
     """
-    def __init__(self, publisher: RabbitRepo, runner_factory):
-        self._pub = publisher
-        self._runner_factory = runner_factory  # async factory: dict|LLMServiceObj -> ILLMRunner
+    def __init__(
+        self,
+        publisher: RabbitRepo,
+        runner_factory: Callable[[Dict[str, Any]], Awaitable[ILLMRunner]],
+    ):
+        self._pub: RabbitRepo = publisher
+        self._runner_factory = runner_factory  # async factory: dict -> ILLMRunner
         self._sessions: Dict[str, _Session] = {}
         self._ready = asyncio.Event()
-        # If you need async load (history, etc.), call self._ready.clear() and later set
-        self._ready.set()
+        self._ready.set()  # call clear()/set() if you preload history
+        self._service_id_lc = settings.SERVICE_ID.lower()
 
-    async def init(self):
-        """
-        Hook to preload history/sessions if needed; call self._ready.set() when finished.
-        """
-        # Example:
-        # self._ready.clear()
-        # await load_history()
-        # self._ready.set()
+    async def init(self) -> None:
+        """Hook to preload history/sessions; call self._ready.set() when finished."""
         pass
 
     # ---------------------------- helpers ----------------------------
 
     def _to_model(self, data: Any) -> LLMServiceObj:
-        """
-        Accepts LLMServiceObj or dict and returns a validated LLMServiceObj.
-        """
+        """Accepts LLMServiceObj or dict and returns a validated LLMServiceObj."""
         if isinstance(data, LLMServiceObj):
             return data
         if isinstance(data, dict):
             return LLMServiceObj(**data)
-        # If your pipeline ever sends compressed strings here, decompress+parse first.
         raise TypeError("LLMService expects an object payload (dict/LLMServiceObj).")
 
     async def _emit_result(
         self,
-        obj: LLMServiceObj,
+        obj: LLMServiceObj | Dict[str, Any],
         message: str,
         success: bool,
         queue: str,
         *,
         check_system: bool = False,
-        include_llm_message: bool = True
-    ):
+        include_llm_message: bool = True,
+    ) -> None:
         """
         Build a ResultObj-style message on the wire, mirroring your .NET usage.
         check_system=True -> don't publish if obj.IsSystemLlm is True (matches your rule).
         """
-        obj.ResultMessage = message
-        obj.ResultSuccess = success
+        llm = obj if isinstance(obj, LLMServiceObj) else LLMServiceObj(**obj)
+
+        llm.ResultMessage = message
+        llm.ResultSuccess = success
         if include_llm_message:
-            obj.LlmMessage = f"<Success>{message}</Success>" if success else f"<Error>{message}</Error>"
+            llm.LlmMessage = f"<Success>{message}</Success>" if success else f"<Error>{message}</Error>"
 
-        if check_system and obj.IsSystemLlm:
+        if check_system and llm.IsSystemLlm:
             return
 
         # You publish LLMServiceObj on "llmServiceMessage"/"llmSessionMessage" in .NET
-        # That lets the coordinator show the assistant text and statuses.
-        await self._pub.publish(queue, obj)
+        await self._pub.publish(queue, llm)
 
     def _session_for(self, session_id: str) -> Optional[_Session]:
         return self._sessions.get(session_id)
 
     # ---------------------------- API methods ----------------------------
 
-    async def StartProcess(self, payload: Any):
+    async def StartProcess(self, payload: Any) -> None:
         llm = self._to_model(payload)
 
-        # Construct Python-side session id like C#: RequestSessionId + "_" + LLMRunnerType
+        # Validate critical fields
+        if not llm.RequestSessionId:
+            await self._emit_result(llm, "Error: RequestSessionId is required.", False, "llmServiceMessage")
+            return
+        if not llm.LLMRunnerType:
+            await self._emit_result(llm, "Error: LLMRunnerType is required.", False, "llmServiceMessage")
+            return
+
+        # Construct session id like C#: RequestSessionId + "_" + LLMRunnerType
         session_id = f"{llm.RequestSessionId}_{llm.LLMRunnerType}"
         llm.SessionId = session_id
 
@@ -99,19 +103,19 @@ class LLMService:
             return
 
         sess = self._session_for(session_id)
-        is_runner_null = (sess is None) or (sess.Runner is None)
-        create_new = is_runner_null or (sess and sess.Runner and sess.Runner.IsStateFailed)
+        runner = sess.Runner if sess else None
+        create_new = (runner is None) or getattr(runner, "IsStateFailed", False)
 
         if create_new:
             # Remove previous runner if exists
-            if sess and sess.Runner:
+            if runner:
                 try:
-                    await sess.Runner.RemoveProcess(session_id)
+                    await runner.RemoveProcess(session_id)
                 except Exception:
                     pass
 
-            # Create runner from factory
-            runner: ILLMRunner = await self._runner_factory(llm.model_dump())
+            # Create runner from factory (pass a plain dict for decoupling)
+            runner = await self._runner_factory(llm.model_dump())
             if not runner.IsEnabled:
                 await self._emit_result(
                     llm,
@@ -129,8 +133,8 @@ class LLMService:
 
             self._sessions[session_id] = _Session(Runner=runner, FullSessionId=session_id)
 
-            # Mirror your friendly greeting, gated by service id (you renamed yours to gradllm)
-            if settings.SERVICE_ID.lower() in {"monitor", "gradllm"}:
+            # Friendly greeting for your renamed service
+            if self._service_id_lc in {"monitor", "gradllm"}:
                 await self._emit_result(
                     llm,
                     f"Hi i'm {runner.Type} your {settings.SERVICE_ID} Assistant. How can I help you.",
@@ -139,16 +143,20 @@ class LLMService:
                     check_system=True,
                 )
 
-        # Notify "started"
+        # Notify "started" (full LLMServiceObj)
         await self._pub.publish("llmServiceStarted", llm)
 
-    async def RemoveSession(self, payload: Any):
+    async def RemoveSession(self, payload: Any) -> None:
         llm = self._to_model(payload)
         base = (llm.SessionId or "").split("_")[0]
-        targets = [k for k in list(self._sessions.keys()) if k.startswith(base + "_")]
+        if not base:
+            await self._emit_result(llm, "Error: SessionId is required to remove sessions.", False, "llmServiceMessage")
+            return
 
-        msgs = []
+        targets = [k for k in list(self._sessions.keys()) if k.startswith(base + "_")]
+        msgs: list[str] = []
         ok = True
+
         for sid in targets:
             s = self._sessions.get(sid)
             if not s or not s.Runner:
@@ -172,7 +180,7 @@ class LLMService:
         else:
             await self._emit_result(llm, " ".join(msgs), False, "llmServiceMessage")
 
-    async def StopRequest(self, payload: Any):
+    async def StopRequest(self, payload: Any) -> None:
         llm = self._to_model(payload)
         sid = llm.SessionId or ""
         s = self._session_for(sid)
@@ -189,7 +197,7 @@ class LLMService:
             check_system=True,
         )
 
-    async def UserInput(self, payload: Any):
+    async def UserInput(self, payload: Any) -> None:
         llm = self._to_model(payload)
         sid = llm.SessionId or ""
         s = self._session_for(sid)
@@ -198,17 +206,17 @@ class LLMService:
             return
 
         r: ILLMRunner = s.Runner
-        if r.IsStateStarting:
+        if getattr(r, "IsStateStarting", False):
             await self._emit_result(llm, "Please wait, the assistant is starting...", False, "llmServiceMessage")
             return
-        if r.IsStateFailed:
+        if getattr(r, "IsStateFailed", False):
             await self._emit_result(llm, "The Assistant is stopped. Try reloading.", False, "llmServiceMessage")
             return
 
         # Let runner push partials itself if desired; we still return a small ack
         await r.SendInputAndGetResponse(llm.model_dump())
 
-    async def QueryIndexResult(self, payload: Any):
+    async def QueryIndexResult(self, payload: Any) -> None:
         """
         Your .NET listener concatenates RAG outputs, sets ResultObj, and notifies the coordinator.
         Here, we forward a service message containing the same info so the UI can reflect completion.
@@ -219,22 +227,21 @@ class LLMService:
             outputs = data.get("QueryResults") or []
             rag_data = "\n".join([x.get("Output", "") for x in outputs if isinstance(x, dict)])
 
-            # Shape compatible with your coordinator expectations
             await self._pub.publish(
                 "llmServiceMessage",
                 ResultObj(
                     Message=data.get("Message", ""),
                     Success=bool(data.get("Success", False)),
                     Data=rag_data,
-                )
+                ),
             )
         except Exception as e:
             await self._pub.publish(
                 "llmServiceMessage",
-                ResultObj(Message=str(e), Success=False)
+                ResultObj(Message=str(e), Success=False),
             )
 
-    async def GetFunctionRegistry(self, filtered: bool = False):
+    async def GetFunctionRegistry(self, filtered: bool = False) -> None:
         """
         Wire up to your real registry when ready.
         For now, mimic your success message payload.
@@ -243,5 +250,5 @@ class LLMService:
         msg = f"Success : Got GetFunctionCatalogJson : {catalog}"
         await self._pub.publish(
             "llmServiceMessage",
-            ResultObj(Message=msg, Success=True)
+            ResultObj(Message=msg, Success=True),
         )