init commit

app.py

@@ -0,0 +1,92 @@
+import asyncio
+import gradio as gr
+from fastapi import FastAPI
+
+from app.config import settings
+from app.listener import RabbitListenerBase
+from app.rabbit_repo import RabbitRepo
+from app.service import LLMService
+from app.runners.base import ILLMRunner
+
+# --- Runner factory (stub) ---
+class EchoRunner(ILLMRunner):
+    Type = "EchoRunner"
+    async def StartProcess(self, llmServiceObj: dict): pass
+    async def RemoveProcess(self, sessionId: str): pass
+    async def StopRequest(self, sessionId: str): pass
+    async def SendInputAndGetResponse(self, llmServiceObj: dict):
+        # Emits a message back (you can choose queue names per your topology)
+        pass
+
+async def runner_factory(llmServiceObj: dict) -> ILLMRunner:
+    # Use llmServiceObj["LLMRunnerType"] to instantiate different runners
+    return EchoRunner()
+
+# --- Publisher and Service ---
+publisher = RabbitRepo(external_source="https://space.external")  # put your ExternalUrl if you have one
+service = LLMService(publisher, runner_factory)
+
+# --- Handlers mapping .NET FuncName -> service method ---
+async def h_start(data): await service.StartProcess(data or {})
+async def h_user(data):  await service.UserInput(data or {})
+async def h_remove(data): await service.RemoveSession(data or {})
+async def h_stop(data): await service.StopRequest(data or {})
+async def h_qir(data): await service.QueryIndexResult(data or {})
+async def h_getreg(data): await service.GetFunctionRegistry(False)
+async def h_getreg_f(data): await service.GetFunctionRegistry(True)
+
+handlers = {
+    "llmStartSession": h_start,
+    "llmUserInput": h_user,
+    "llmRemoveSession": h_remove,
+    "llmStopRequest": h_stop,
+    "queryIndexResult": h_qir,
+    "getFunctionRegistry": h_getreg,
+    "getFunctionRegistryFiltered": h_getreg_f,
+}
+
+listener = RabbitListenerBase(service_id=settings.SERVICE_ID, handlers=handlers)
+
+# Declarations mirror your C# InitRabbitMQObjs()
+DECLS = [
+    {"ExchangeName":"llmStartSession"+settings.SERVICE_ID, "FuncName":"llmStartSession", "MessageTimeout":600000, "RoutingKeys":[settings.RABBIT_ROUTING_KEY]},
+    {"ExchangeName":"llmUserInput"+settings.SERVICE_ID, "FuncName":"llmUserInput", "MessageTimeout":600000, "RoutingKeys":[settings.RABBIT_ROUTING_KEY]},
+    {"ExchangeName":"llmRemoveSession"+settings.SERVICE_ID, "FuncName":"llmRemoveSession", "MessageTimeout":60000, "RoutingKeys":[settings.RABBIT_ROUTING_KEY]},
+    {"ExchangeName":"llmStopRequest"+settings.SERVICE_ID, "FuncName":"llmStopRequest", "MessageTimeout":60000, "RoutingKeys":[settings.RABBIT_ROUTING_KEY]},
+    {"ExchangeName":"queryIndexResult"+settings.SERVICE_ID, "FuncName":"queryIndexResult", "MessageTimeout":60000, "RoutingKeys":[settings.RABBIT_ROUTING_KEY]},
+    {"ExchangeName":"getFunctionRegistry"+settings.SERVICE_ID, "FuncName":"getFunctionRegistry", "MessageTimeout":60000, "RoutingKeys":[settings.RABBIT_ROUTING_KEY]},
+    {"ExchangeName":"getFunctionRegistryFiltered"+settings.SERVICE_ID, "FuncName":"getFunctionRegistryFiltered", "MessageTimeout":60000, "RoutingKeys":[settings.RABBIT_ROUTING_KEY]},
+]
+
+# --- Gradio UI (for smoke test) ---
+async def ping():
+    return "ok"
+
+with gr.Blocks() as demo:
+    gr.Markdown("### LLM Runner (Python) listening on RabbitMQ")
+    btn = gr.Button("Ping")
+    out = gr.Textbox()
+    btn.click(ping, inputs=None, outputs=out)
+
+# --- FastAPI mount + lifecycle ---
+app = FastAPI()
+app = gr.mount_gradio_app(app, demo, path="/")
+
+@app.get("/health")
+async def health():
+    return {"status":"ok"}
+
+@app.on_event("startup")
+async def on_start():
+    await publisher.connect()
+    await service.init()
+    await listener.start(DECLS)
+
+@app.on_event("shutdown")
+async def on_stop():
+    # aio-pika RobustConnection closes on GC; optionally add explicit closes if you add references
+    pass
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

cloud_event.py

@@ -0,0 +1,23 @@
+import json
+from dataclasses import dataclass, asdict
+from datetime import datetime, timezone
+from typing import Any
+
+@dataclass
+class CloudEvent:
+    id: str
+    type: str
+    source: str
+    time: str
+    data: Any
+
+    @staticmethod
+    def wrap(obj: Any, *, event_type: str, source: str, id: str) -> bytes:
+        evt = CloudEvent(
+            id=id,
+            type=event_type or (obj.__class__.__name__ if obj is not None else "NullOrEmpty"),
+            source=source,
+            time=datetime.now(timezone.utc).isoformat(),
+            data=obj,
+        )
+        return json.dumps(asdict(evt), ensure_ascii=False).encode("utf-8")

config.py

@@ -0,0 +1,22 @@
+from pydantic import BaseSettings, AnyUrl
+from typing import Optional, Dict
+
+class Settings(BaseSettings):
+    # AMQP
+    AMQP_URL: AnyUrl  # e.g. amqps://user:pass@host:5671/%2F?heartbeat=30
+    RABBIT_INSTANCE_NAME: str = "prod"
+    RABBIT_EXCHANGE_TYPE: str = "topic"  # match your .NET Type
+    RABBIT_ROUTING_KEY: str = ""         # match your .NET RoutingKeys ("" ok)
+    RABBIT_PREFETCH: int = 1
+
+    # Service identity
+    SERVICE_ID: str = "monitor"          # "monitor"|"nmap"|...
+    USE_TLS: bool = True
+
+    # Optional exchange type overrides by prefix, like .NET ExchangeTypes
+    EXCHANGE_TYPES: Dict[str, str] = {}
+
+    class Config:
+        case_sensitive = True
+
+settings = Settings()  # env-driven

listener.py

@@ -0,0 +1,47 @@
+import json
+from typing import Callable, Dict, List, Optional
+import aio_pika
+from .rabbit_base import RabbitBase
+from .config import settings
+
+# Maps FuncName -> handler coroutine
+Handler = Callable[[dict], "awaitable[None]"]
+
+class RabbitListenerBase(RabbitBase):
+    def __init__(self, service_id: str, handlers: Dict[str, Handler]):
+        super().__init__()
+        self._service_id = service_id
+        self._handlers = handlers
+        self._consumers: List[aio_pika.abc.AbstractRobustQueue] = []
+
+    def _qname(self, exchange: str, routing_keys: List[str]) -> str:
+        rk_part = "-".join(sorted([rk for rk in (routing_keys or [""]) if rk != ""])) or ""
+        suffix = f"-{rk_part}" if rk_part else ""
+        return f"{settings.RABBIT_INSTANCE_NAME}-{exchange}{suffix}"
+
+    async def start(self, declarations: List[dict]):
+        """
+        declarations: list of {ExchangeName, FuncName, MessageTimeout, Type?, RoutingKeys?}
+        """
+        for d in declarations:
+            exch = d["ExchangeName"]
+            rks = d.get("RoutingKeys") or [settings.RABBIT_ROUTING_KEY]
+            ttl = d.get("MessageTimeout") or None
+            q = await self.declare_queue_bind(exchange=exch, queue_name=self._qname(exch, rks), routing_keys=rks, ttl_ms=ttl)
+            await q.consume(self._make_consumer(d["FuncName"]))
+            self._consumers.append(q)
+
+    def _make_consumer(self, func_name: str):
+        handler = self._handlers.get(func_name)
+        async def _on_msg(msg: aio_pika.IncomingMessage):
+            async with msg.process():
+                try:
+                    # Expect CloudEvent JSON
+                    envelope = json.loads(msg.body.decode("utf-8"))
+                    data = envelope.get("data")
+                    if handler:
+                        await handler(data)
+                except Exception as e:
+                    # swallow to avoid nack loops; your logger can capture details
+                    pass
+        return _on_msg

rabbit_base.py

@@ -0,0 +1,40 @@
+import asyncio, json, uuid
+import aio_pika
+from typing import Callable, Dict, List, Optional
+from .config import settings
+
+ExchangeResolver = Callable[[str], str]  # exchangeName -> exchangeType
+
+class RabbitBase:
+    def __init__(self, exchange_type_resolver: Optional[ExchangeResolver] = None):
+        self._conn: Optional[aio_pika.RobustConnection] = None
+        self._chan: Optional[aio_pika.RobustChannel] = None
+        self._exchanges: Dict[str, aio_pika.Exchange] = {}
+        self._exchange_type_resolver = exchange_type_resolver or (lambda _: settings.RABBIT_EXCHANGE_TYPE)
+
+    async def connect(self):
+        if self._conn and not self._conn.is_closed:
+            return
+        self._conn = await aio_pika.connect_robust(str(settings.AMQP_URL))
+        self._chan = await self._conn.channel()
+        await self._chan.set_qos(prefetch_count=settings.RABBIT_PREFETCH)
+
+    async def ensure_exchange(self, name: str) -> aio_pika.Exchange:
+        await self.connect()
+        if name in self._exchanges:
+            return self._exchanges[name]
+        ex_type = self._exchange_type_resolver(name)
+        ex = await self._chan.declare_exchange(name, getattr(aio_pika.ExchangeType, ex_type), durable=True)
+        self._exchanges[name] = ex
+        return ex
+
+    async def declare_queue_bind(self, exchange: str, queue_name: str, routing_keys: List[str], ttl_ms: Optional[int]):
+        await self.connect()
+        ex = await self.ensure_exchange(exchange)
+        args = {}
+        if ttl_ms:
+            args["x-message-ttl"] = ttl_ms
+        q = await self._chan.declare_queue(queue_name, durable=True, exclusive=False, auto_delete=True, arguments=args)
+        for rk in routing_keys or [""]:
+            await q.bind(ex, rk)
+        return q

rabbit_repo.py

@@ -0,0 +1,34 @@
+import uuid
+from typing import Any
+from .rabbit_base import RabbitBase
+from .cloud_event import CloudEvent
+from .config import settings
+from .utils import to_json, json_compress_str
+
+class RabbitRepo(RabbitBase):
+    def __init__(self, external_source: str):
+        super().__init__(exchange_type_resolver=self._resolve_type)
+        self._external_source = external_source  # like SystemUrl.ExternalUrl
+
+    def _resolve_type(self, exch: str) -> str:
+        # longest prefix wins (like your .NET mapping)
+        matches = [k for k in settings.EXCHANGE_TYPES.keys() if exch.lower().startswith(k.lower())]
+        if matches:
+            return settings.EXCHANGE_TYPES[max(matches, key=len)]
+        return settings.RABBIT_EXCHANGE_TYPE
+
+    async def publish(self, exchange: str, obj: Any, routing_key: str = ""):
+        ex = await self.ensure_exchange(exchange)
+        payload = CloudEvent.wrap(obj, event_type=(obj.__class__.__name__ if obj is not None else "NullOrEmpty"),
+                                  source=self._external_source, id=str(uuid.uuid4()))
+        await ex.publish(aio_pika.Message(body=payload), routing_key=routing_key)
+
+    async def publish_jsonz(self, exchange: str, obj: Any, routing_key: str = "", with_id: str | None = None):
+        ex = await self.ensure_exchange(exchange)
+        json_str = to_json(obj)
+        datajsonZ = json_compress_str(json_str)
+        to_send = (datajsonZ, with_id) if with_id else datajsonZ
+        payload = CloudEvent.wrap(to_send, event_type=(obj.__class__.__name__ if obj is not None else "NullOrEmpty"),
+                                  source=self._external_source, id=str(uuid.uuid4()))
+        await ex.publish(aio_pika.Message(body=payload), routing_key=routing_key)
+        return datajsonZ

runners/base.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+class ILLMRunner(ABC):
+    Type: str = "BaseLLM"
+    IsEnabled: bool = True
+    IsStateStarting: bool = False
+    IsStateFailed: bool = False
+
+    @abstractmethod
+    async def StartProcess(self, llmServiceObj: dict) -> None: ...
+    @abstractmethod
+    async def RemoveProcess(self, sessionId: str) -> None: ...
+    @abstractmethod
+    async def StopRequest(self, sessionId: str) -> None: ...
+    @abstractmethod
+    async def SendInputAndGetResponse(self, llmServiceObj: dict) -> None: ...

runners/service.py

@@ -0,0 +1,127 @@
+import asyncio
+from typing import Dict, Optional
+from collections import defaultdict
+
+from .rabbit_repo import RabbitRepo
+from .config import settings
+from .runners.base import ILLMRunner
+
+class LLMService:
+    def __init__(self, publisher: RabbitRepo, runner_factory):
+        self._pub = publisher
+        self._runner_factory = runner_factory
+        self._sessions: Dict[str, dict] = {}  # sessionId -> {"Runner": ILLMRunner, "FullSessionId": str}
+        self._ready = asyncio.Event()
+        self._ready.set()  # if you have async load, clear and set after
+
+    async def init(self):
+        # If you have history to load, do here then self._ready.set()
+        pass
+
+    async def _set_result(self, obj: dict, message: str, success: bool, queue: str, check_system: bool=False):
+        obj["ResultMessage"] = message
+        obj["ResultSuccess"] = success
+        obj["LlmMessage"] = (f"<Success>{message}</Success>" if success else f"<Error>{message}</Error>")
+        # mirror your .NET rule (don’t publish for system llm if check_system is True)
+        if not (check_system and obj.get("IsSystemLlm")):
+            await self._pub.publish(queue, obj)
+
+    async def StartProcess(self, llmServiceObj: dict):
+        session_id = f"{llmServiceObj['RequestSessionId']}_{llmServiceObj['LLMRunnerType']}"
+        llmServiceObj["SessionId"] = session_id
+
+        # wait ready (max ~120s like .NET)
+        try:
+            await asyncio.wait_for(self._ready.wait(), timeout=120)
+        except asyncio.TimeoutError:
+            await self._set_result(llmServiceObj, "Timed out waiting for initialization.", False, "llmServiceMessage", True)
+            return
+
+        sess = self._sessions.get(session_id)
+        is_runner_null = not sess or not sess.get("Runner")
+
+        create_new = is_runner_null or sess["Runner"].IsStateFailed
+        if create_new:
+            if sess and sess.get("Runner"):
+                try:
+                    await sess["Runner"].RemoveProcess(session_id)
+                except: pass
+
+            runner: ILLMRunner = await self._runner_factory(llmServiceObj)
+            if not runner.IsEnabled:
+                await self._set_result(llmServiceObj, f"{llmServiceObj['LLMRunnerType']} {settings.SERVICE_ID} not started as it is disabled.", True, "llmServiceMessage")
+                return
+
+            await self._set_result(llmServiceObj, f"Starting {runner.Type} {settings.SERVICE_ID} Expert", True, "llmServiceMessage", True)
+            await runner.StartProcess(llmServiceObj)
+
+            self._sessions[session_id] = {"Runner": runner, "FullSessionId": session_id}
+            if settings.SERVICE_ID == "monitor":
+                await self._set_result(llmServiceObj, f"Hi i'm {runner.Type} your Network Monitor Assistant. How can I help you.", True, "llmServiceMessage", True)
+
+        await self._pub.publish("llmServiceStarted", llmServiceObj)
+
+    async def RemoveSession(self, llmServiceObj: dict):
+        # Behaves like your RemoveAllSessionIdProcesses (prefix match)
+        base = llmServiceObj.get("SessionId","").split("_")[0]
+        targets = [k for k in self._sessions.keys() if k.startswith(base + "_")]
+        msgs = []
+        ok = True
+        for sid in targets:
+            s = self._sessions.get(sid)
+            if s and s.get("Runner"):
+                try:
+                    await s["Runner"].RemoveProcess(sid)
+                    s["Runner"] = None
+                    msgs.append(sid)
+                except Exception as e:
+                    ok = False
+                    msgs.append(f"Error {sid}: {e}")
+        if ok:
+            await self._set_result(llmServiceObj, f"Success: Removed sessions for {' '.join(msgs)}", True, "llmSessionMessage", True)
+        else:
+            await self._set_result(llmServiceObj, " ".join(msgs), False, "llmServiceMessage")
+
+    async def StopRequest(self, llmServiceObj: dict):
+        sid = llmServiceObj.get("SessionId","")
+        s = self._sessions.get(sid)
+        if not s or not s.get("Runner"):
+            await self._set_result(llmServiceObj, f"Error: Runner missing for session {sid}.", False, "llmServiceMessage")
+            return
+        await s["Runner"].StopRequest(sid)
+        await self._set_result(llmServiceObj, f"Success {s['Runner'].Type} {settings.SERVICE_ID} Assistant output has been halted", True, "llmServiceMessage", True)
+
+    async def UserInput(self, llmServiceObj: dict):
+        sid = llmServiceObj.get("SessionId","")
+        s = self._sessions.get(sid)
+        if not s or not s.get("Runner"):
+            await self._set_result(llmServiceObj, f"Error: SessionId {sid} has no running process.", False, "llmServiceMessage")
+            return
+        r: ILLMRunner = s["Runner"]
+        if r.IsStateStarting:
+            await self._set_result(llmServiceObj, "Please wait, the assistant is starting...", False, "llmServiceMessage")
+            return
+        if r.IsStateFailed:
+            await self._set_result(llmServiceObj, "The Assistant is stopped. Try reloading.", False, "llmServiceMessage")
+            return
+        await r.SendInputAndGetResponse(llmServiceObj)
+        # emitter side can push partials directly to queues if desired
+
+    async def QueryIndexResult(self, queryIndexRequest: dict):
+        # Adapted to your behavior: concatenate outputs, publish completion via internal coordinator if needed
+        try:
+            rag_data = "\n".join([qr.get("Output","") for qr in (queryIndexRequest.get("QueryResults") or [])])
+            # You signal _queryCoordinator.CompleteQuery in .NET; here you may forward/publish result…
+            # Example: include rag data in a service message to the session
+            await self._pub.publish("llmServiceMessage", {
+                "ResultSuccess": queryIndexRequest.get("Success", False),
+                "ResultMessage": queryIndexRequest.get("Message",""),
+                "Data": rag_data,
+            })
+        except Exception as e:
+            await self._pub.publish("llmServiceMessage", {"ResultSuccess": False, "ResultMessage": str(e)})
+
+    async def GetFunctionRegistry(self, filtered: bool = False):
+        # Plug in your registry
+        data = {"FunctionCatalogJson": "{}", "Filtered": filtered}
+        await self._pub.publish("llmServiceMessage", {"ResultSuccess": True, "ResultMessage": f"Success : Got GetFunctionCatalogJson : {data}"})

utils.py

@@ -0,0 +1,11 @@
+import json, zlib, base64
+from typing import Any
+
+def to_json(obj: Any) -> str:
+    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"))
+
+def json_compress_str(s: str) -> str:
+    return base64.b64encode(zlib.compress(s.encode("utf-8"), level=6)).decode("ascii")
+
+def json_decompress_str(s: str) -> str:
+    return zlib.decompress(base64.b64decode(s)).decode("utf-8")