From f65a7723e053df027fa0dbe0e31739a723a2f673 Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 30 Jul 2026 12:20:40 +0200 Subject: [PATCH 1/3] Wire DocumentAnalyzer to NATSLLMChatAdapter for centralized LLM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace placeholder backend-handles-llm keys with a real NATS chat adapter when USE_BACKEND and USE_CENTRALIZED_LLM are set. Adapter currently lives in report_analyst/core/ — folder placement to review later. --- report_analyst/core/analyzer.py | 30 ++- report_analyst/core/llm_providers.py | 14 ++ report_analyst_enterprise/nats_llm_adapter.py | 174 ++++++++++++++++++ tests/test_centralized_llm.py | 80 ++++++++ 4 files changed, 288 insertions(+), 10 deletions(-) create mode 100644 report_analyst_enterprise/nats_llm_adapter.py create mode 100644 tests/test_centralized_llm.py diff --git a/report_analyst/core/analyzer.py b/report_analyst/core/analyzer.py index 92c9cb45..8576eaf0 100644 --- a/report_analyst/core/analyzer.py +++ b/report_analyst/core/analyzer.py @@ -48,14 +48,12 @@ use_full_backend_analysis, ) -# If using backend for LLM, don't require local API keys -if use_backend and (use_centralized_llm or use_full_backend_analysis): - logger.info("Using backend for LLM functionality - local API keys not required") - # Set placeholder values for compatibility - if not openai_key: - openai_key = "backend-handles-llm" - if not gemini_key: - gemini_key = "backend-handles-llm" +# Centralized LLM: platform handles inference via NATS — never use placeholder OpenAI keys +_centralized_llm_active = use_backend and (use_centralized_llm or use_full_backend_analysis) +if _centralized_llm_active: + logger.info( + "Centralized LLM mode — inference via platform/NATS (no local OpenAI/Gemini)" + ) else: # Only check for API keys if not using backend LLM # Check if we need to force the default model based on available keys @@ -144,8 +142,20 @@ def __init__(self): self.use_backend_llm = use_backend and (use_centralized_llm or use_full_backend_analysis) if self.use_backend_llm: - # Set minimal placeholders for compatibility - self.llm = None + log_analysis_step( + "Initializing NATS-backed LLM (platform centralized)", + "info", + ) + try: + from report_analyst_enterprise.nats_llm_adapter import NATSLLMChatAdapter + except ImportError as exc: + raise ImportError( + "Centralized LLM requires the enterprise package " + "(report_analyst_enterprise.nats_llm_adapter). " + "Install enterprise extras or disable USE_CENTRALIZED_LLM." + ) from exc + + self.llm = NATSLLMChatAdapter(model=os.getenv("PLATFORM_CHAT_MODEL", "gemma3-4b")) self.embeddings = None else: self._initialize_llm_clients() diff --git a/report_analyst/core/llm_providers.py b/report_analyst/core/llm_providers.py index 0fd9808c..b2b12acc 100644 --- a/report_analyst/core/llm_providers.py +++ b/report_analyst/core/llm_providers.py @@ -14,6 +14,13 @@ logger = logging.getLogger(__name__) +def centralized_llm_requested() -> bool: + use_backend = os.getenv("USE_BACKEND", "false").lower() == "true" + use_centralized = os.getenv("USE_CENTRALIZED_LLM", "false").lower() == "true" + use_full = os.getenv("USE_FULL_BACKEND_ANALYSIS", "false").lower() == "true" + return use_backend and (use_centralized or use_full) + + def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: """ Factory function to get LLM implementations based on model name. @@ -29,7 +36,14 @@ def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: Raises: ValueError: If the API key for the selected model is not available ValueError: If the model type is not supported + RuntimeError: If centralized LLM is enabled (must use NATSLLMChatAdapter) """ + if centralized_llm_requested(): + raise RuntimeError( + "Centralized LLM mode is active (USE_BACKEND + USE_CENTRALIZED_LLM). " + "Local OpenAI/Gemini clients are disabled; use DocumentAnalyzer with NATS LLM." + ) + # OpenAI models if model_name.startswith("gpt-"): api_key = os.getenv("OPENAI_API_KEY") diff --git a/report_analyst_enterprise/nats_llm_adapter.py b/report_analyst_enterprise/nats_llm_adapter.py new file mode 100644 index 00000000..19579052 --- /dev/null +++ b/report_analyst_enterprise/nats_llm_adapter.py @@ -0,0 +1,174 @@ +"""LlamaIndex-compatible LLM adapter routing .achat() to platform via NATS. + +Enterprise module: platform NATS chat for USE_CENTRALIZED_LLM deployments. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import uuid +from dataclasses import dataclass +from typing import Any, List, Optional, Union + +import nats +from llama_index.core.llms import ChatMessage, MessageRole + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = os.getenv("PLATFORM_CHAT_MODEL", "gemma3-4b") +REQUEST_TIMEOUT = float(os.getenv("NATS_LLM_TIMEOUT", "120")) + + +@dataclass +class _ChatResponseMessage: + content: str + + +@dataclass +class _ChatResponseWrapper: + message: _ChatResponseMessage + + +class NATSLLMChatAdapter: + """ + Implements the subset of LlamaIndex LLM interface used by DocumentAnalyzer: + await self.llm.achat(prompt=...) and await self.llm.achat(messages). + """ + + def __init__( + self, + model: Optional[str] = None, + nats_url: Optional[str] = None, + nats_token: Optional[str] = None, + ): + self.model = model or DEFAULT_MODEL + self._nats_url = nats_url or os.getenv("NATS_URL", "nats://localhost:4222") + self._nats_token = nats_token or os.getenv("NATS_TOKEN") + self._nc: Optional[nats.NATS] = None + self._lock = asyncio.Lock() + self._allowed_models: Optional[List[str]] = None + + def _connection_url(self) -> str: + url = self._nats_url + if self._nats_token and "@" not in url: + protocol, rest = url.split("://", 1) + url = f"{protocol}://{self._nats_token}@{rest}" + return url + + async def _ensure_connected(self): + if self._nc and self._nc.is_connected: + return + async with self._lock: + if self._nc and self._nc.is_connected: + return + self._nc = await nats.connect(self._connection_url(), connect_timeout=15) + + async def close(self): + if self._nc: + await self._nc.close() + self._nc = None + + async def list_allowed_models(self) -> List[str]: + if self._allowed_models is not None: + return self._allowed_models + text = await self._request_llm( + prompt="", + system_prompt=None, + request_type="list_models", + ) + try: + parsed = json.loads(text) + self._allowed_models = parsed.get("models", [DEFAULT_MODEL]) + except json.JSONDecodeError: + self._allowed_models = [DEFAULT_MODEL] + return self._allowed_models + + def _messages_to_prompt(self, messages: Union[str, List[ChatMessage], List[Any]]) -> tuple[str, Optional[str]]: + if isinstance(messages, str): + return messages, None + system_parts = [] + user_parts = [] + for msg in messages: + if isinstance(msg, ChatMessage): + role, content = msg.role, msg.content + elif isinstance(msg, dict): + role = msg.get("role", MessageRole.USER) + content = msg.get("content", "") + else: + role = getattr(msg, "role", MessageRole.USER) + content = getattr(msg, "content", str(msg)) + if role == MessageRole.SYSTEM or str(role).lower() == "system": + system_parts.append(content) + else: + user_parts.append(content) + system_prompt = "\n".join(system_parts) if system_parts else None + prompt = "\n".join(user_parts) if user_parts else "" + return prompt, system_prompt + + async def achat( + self, + messages: Union[str, List[ChatMessage], List[Any]] = None, + prompt: Optional[str] = None, + **kwargs: Any, + ) -> _ChatResponseWrapper: + if prompt is not None and messages is None: + user_prompt, system_prompt = prompt, kwargs.get("system_prompt") + else: + user_prompt, system_prompt = self._messages_to_prompt(messages or prompt or "") + model = kwargs.get("model", self.model) + text = await self._request_llm( + prompt=user_prompt, + system_prompt=system_prompt, + model=model, + ) + return _ChatResponseWrapper(message=_ChatResponseMessage(content=text)) + + async def _request_llm( + self, + prompt: str, + system_prompt: Optional[str], + model: Optional[str] = None, + request_type: str = "custom", + ) -> str: + await self._ensure_connected() + request_id = str(uuid.uuid4()) + reply_subject = f"llm.response.{request_id}" + payload = { + "request_id": request_id, + "request_type": request_type, + "prompt": prompt, + "system_prompt": system_prompt, + "model": model or self.model, + "reply_subject": reply_subject, + "metadata": {"source": "report_analyst"}, + } + + future: asyncio.Future = asyncio.get_event_loop().create_future() + + async def on_reply(msg): + try: + data = json.loads(msg.data.decode()) + if data.get("request_id") != request_id: + return + if data.get("error"): + if not future.done(): + future.set_exception(RuntimeError(data["error"])) + elif not future.done(): + future.set_result(data.get("response", "")) + except Exception as exc: + if not future.done(): + future.set_exception(exc) + + sub = await self._nc.subscribe(reply_subject, cb=on_reply) + try: + js = self._nc.jetstream() + try: + await js.publish("llm.request", json.dumps(payload).encode()) + except Exception: + await self._nc.publish("llm.request", json.dumps(payload).encode()) + return await asyncio.wait_for(future, timeout=REQUEST_TIMEOUT) + finally: + await sub.unsubscribe() diff --git a/tests/test_centralized_llm.py b/tests/test_centralized_llm.py new file mode 100644 index 00000000..372dfae8 --- /dev/null +++ b/tests/test_centralized_llm.py @@ -0,0 +1,80 @@ +"""Centralized LLM mode: fail-closed local OpenAI, NATS adapter for .achat().""" + +from __future__ import annotations + +import asyncio +import json +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +# Minimal stubs so we can test llm_providers / adapter without full DocumentAnalyzer import chain +if "nats" not in sys.modules: + sys.modules["nats"] = MagicMock() + + +@pytest.fixture +def centralized_env(monkeypatch): + monkeypatch.setenv("USE_BACKEND", "true") + monkeypatch.setenv("USE_CENTRALIZED_LLM", "true") + monkeypatch.setenv("USE_FULL_BACKEND_ANALYSIS", "false") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + +def test_get_llm_fails_closed_when_centralized(centralized_env): + """Hypothesis: get_llm() must not build OpenAI client in centralized mode.""" + from report_analyst.core.llm_providers import get_llm + + with pytest.raises(RuntimeError, match="Centralized LLM"): + get_llm("gpt-4o-mini") + + +@pytest.mark.asyncio +async def test_nats_llm_adapter_achat_uses_per_request_reply_subject(): + """Hypothesis: adapter waits on llm.response.{request_id}, not broadcast llm.response.""" + pytest.importorskip("llama_index.core") + from report_analyst_enterprise.nats_llm_adapter import NATSLLMChatAdapter + + request_id_holder = [] + + mock_nc = MagicMock() + mock_js = MagicMock() + mock_nc.jetstream.return_value = mock_js + mock_nc.is_connected = True + mock_nc.drain = AsyncMock() + mock_nc.close = AsyncMock() + + async def fake_subscribe(subject, cb=None): + async def deliver(): + if cb and subject.startswith("llm.response."): + rid = subject.split(".", 2)[-1] + request_id_holder.append(rid) + msg = MagicMock() + msg.data = json.dumps( + { + "request_id": rid, + "response": "platform says hi", + "error": None, + } + ).encode() + await cb(msg) + + asyncio.get_event_loop().call_soon(lambda: asyncio.create_task(deliver())) + sub_mock = MagicMock() + sub_mock.unsubscribe = AsyncMock() + return sub_mock + + mock_nc.subscribe = AsyncMock(side_effect=fake_subscribe) + mock_js.publish = AsyncMock() + + adapter = NATSLLMChatAdapter(model="gemma3-4b", nats_url="nats://localhost:4222") + adapter._nc = mock_nc + + response = await adapter.achat(prompt="Hello") + assert response.message.content == "platform says hi" + assert request_id_holder + publish_args = mock_js.publish.call_args + assert publish_args[0][0] == "llm.request" + payload = json.loads(publish_args[0][1].decode()) + assert payload["reply_subject"] == f"llm.response.{payload['request_id']}" From 1d59f64422cdc758ca4f3bf1dfb094352e583969 Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 30 Jul 2026 12:55:02 +0200 Subject: [PATCH 2/3] Place NATS LLM adapter under report_analyst_enterprise Centralized LLM chat adapter is an enterprise dependency; core DocumentAnalyzer imports it only when USE_CENTRALIZED_LLM is on. --- report_analyst/core/llm_providers.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/report_analyst/core/llm_providers.py b/report_analyst/core/llm_providers.py index b2b12acc..5ed8d319 100644 --- a/report_analyst/core/llm_providers.py +++ b/report_analyst/core/llm_providers.py @@ -36,12 +36,13 @@ def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: Raises: ValueError: If the API key for the selected model is not available ValueError: If the model type is not supported - RuntimeError: If centralized LLM is enabled (must use NATSLLMChatAdapter) + RuntimeError: If centralized LLM is enabled (must use enterprise NATSLLMChatAdapter) """ if centralized_llm_requested(): raise RuntimeError( "Centralized LLM mode is active (USE_BACKEND + USE_CENTRALIZED_LLM). " - "Local OpenAI/Gemini clients are disabled; use DocumentAnalyzer with NATS LLM." + "Local OpenAI/Gemini clients are disabled; use DocumentAnalyzer with " + "report_analyst_enterprise.nats_llm_adapter.NATSLLMChatAdapter." ) # OpenAI models From e417dcd8c78f19c1d3618207efbc55d33dba9d27 Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 20 Aug 2026 22:17:39 +0200 Subject: [PATCH 3/3] Fix NATS LLM lint: black, unused imports, exception types. --- report_analyst/core/analyzer.py | 4 +--- report_analyst/core/llm_providers.py | 3 +-- report_analyst_enterprise/nats_llm_adapter.py | 7 ++++--- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/report_analyst/core/analyzer.py b/report_analyst/core/analyzer.py index 8576eaf0..c5e38645 100644 --- a/report_analyst/core/analyzer.py +++ b/report_analyst/core/analyzer.py @@ -51,9 +51,7 @@ # Centralized LLM: platform handles inference via NATS — never use placeholder OpenAI keys _centralized_llm_active = use_backend and (use_centralized_llm or use_full_backend_analysis) if _centralized_llm_active: - logger.info( - "Centralized LLM mode — inference via platform/NATS (no local OpenAI/Gemini)" - ) + logger.info("Centralized LLM mode — inference via platform/NATS (no local OpenAI/Gemini)") else: # Only check for API keys if not using backend LLM # Check if we need to force the default model based on available keys diff --git a/report_analyst/core/llm_providers.py b/report_analyst/core/llm_providers.py index 5ed8d319..4e06f47e 100644 --- a/report_analyst/core/llm_providers.py +++ b/report_analyst/core/llm_providers.py @@ -2,8 +2,7 @@ import logging import os -from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Optional from llama_index.llms.gemini import Gemini diff --git a/report_analyst_enterprise/nats_llm_adapter.py b/report_analyst_enterprise/nats_llm_adapter.py index 19579052..ea1974e1 100644 --- a/report_analyst_enterprise/nats_llm_adapter.py +++ b/report_analyst_enterprise/nats_llm_adapter.py @@ -15,6 +15,7 @@ import nats from llama_index.core.llms import ChatMessage, MessageRole +from nats.errors import Error as NatsError logger = logging.getLogger(__name__) @@ -110,7 +111,7 @@ def _messages_to_prompt(self, messages: Union[str, List[ChatMessage], List[Any]] async def achat( self, - messages: Union[str, List[ChatMessage], List[Any]] = None, + messages: Optional[Union[str, List[ChatMessage], List[Any]]] = None, prompt: Optional[str] = None, **kwargs: Any, ) -> _ChatResponseWrapper: @@ -158,7 +159,7 @@ async def on_reply(msg): future.set_exception(RuntimeError(data["error"])) elif not future.done(): future.set_result(data.get("response", "")) - except Exception as exc: + except (json.JSONDecodeError, UnicodeDecodeError, TypeError, ValueError) as exc: if not future.done(): future.set_exception(exc) @@ -167,7 +168,7 @@ async def on_reply(msg): js = self._nc.jetstream() try: await js.publish("llm.request", json.dumps(payload).encode()) - except Exception: + except NatsError: await self._nc.publish("llm.request", json.dumps(payload).encode()) return await asyncio.wait_for(future, timeout=REQUEST_TIMEOUT) finally: