refactor(llmrails): fix package ownership boundaries (8/9)

nemoguardrails/rails/llm/startup/tracing.py → nemoguardrails/rails/llm/embedding/__init__.py

@@ -13,16 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nemoguardrails.rails.llm.config import RailsConfig
-from nemoguardrails.tracing.adapters.base import InteractionLogAdapter
-
-__all__ = ["create_startup_tracing_adapters"]
-
-
-def create_startup_tracing_adapters(config: RailsConfig) -> list[InteractionLogAdapter] | None:
-    if not config.tracing:
-        return None
-
-    from nemoguardrails.tracing import create_log_adapters
-
-    return create_log_adapters(config.tracing)
+__all__ = []

nemoguardrails/rails/llm/startup/embedding_search.py → nemoguardrails/rails/llm/embedding/embedding_search.py

@@ -16,10 +16,10 @@
 """Embedding search provider setup."""
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Type
+from typing import Any, Dict, Optional, Type
 
 from nemoguardrails.embeddings.index import EmbeddingsIndex
-from nemoguardrails.rails.llm.config import EmbeddingSearchProvider, RailsConfig
+from nemoguardrails.rails.llm.config import EmbeddingSearchProvider
 
 DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2"
 DEFAULT_EMBEDDING_ENGINE = "FastEmbed"
@@ -28,7 +28,6 @@
     "DEFAULT_EMBEDDING_ENGINE",
     "DEFAULT_EMBEDDING_MODEL",
     "EmbeddingSearchState",
-    "apply_embedding_model_config",
     "get_embedding_search_provider_instance",
 ]
 
@@ -76,37 +75,6 @@ def get_provider_instance(
         )
 
 
-def apply_embedding_model_config(
-    config: RailsConfig,
-    default_embedding_model: Optional[str],
-    default_embedding_engine: Optional[str],
-    default_embedding_params: Dict[str, Any],
-) -> Tuple[Optional[str], Optional[str], Dict[str, Any]]:
-    """Apply an embeddings model config to the default embedding search settings."""
-    for model in config.models:
-        if model.type != "embeddings":
-            continue
-
-        default_embedding_model = model.model
-        default_embedding_engine = model.engine
-        default_embedding_params = model.parameters or {}
-
-        for esp in [
-            config.core.embedding_search_provider,
-            config.knowledge_base.embedding_search_provider,
-        ]:
-            if esp.name != "default":
-                continue
-            if "embedding_model" not in esp.parameters and model.model is not None:
-                esp.parameters["embedding_model"] = model.model
-            if "embedding_engine" not in esp.parameters and model.engine is not None:
-                esp.parameters["embedding_engine"] = model.engine
-
-        break
-
-    return default_embedding_model, default_embedding_engine, default_embedding_params
-
-
 def get_embedding_search_provider_instance(
     embedding_search_providers: Dict[str, Type[EmbeddingsIndex]],
     default_embedding_model: Optional[str],

nemoguardrails/rails/llm/generation/generation_request.py

@@ -28,6 +28,7 @@
     "PreparedGenerationRequest",
     "normalize_generation_request",
     "prepare_generation_request_for_runtime",
+    "validate_prompt_or_messages",
 ]
 
 
@@ -49,6 +50,18 @@ class PreparedGenerationRequest:
     needs_llm: bool
 
 
+def validate_prompt_or_messages(
+    prompt: Optional[str],
+    messages: Optional[List[dict]],
+) -> None:
+    """Validate that exactly one of prompt or messages is provided."""
+    if prompt is None and messages is None:
+        raise ValueError("Either prompt or messages must be provided.")
+
+    if prompt is not None and messages is not None:
+        raise ValueError("Only one of prompt or messages can be provided.")
+
+
 def normalize_generation_request(
     *,
     prompt: Optional[str],
@@ -57,11 +70,7 @@ def normalize_generation_request(
     state: Optional[Union[dict, State]],
 ) -> GenerationRequest:
     """Normalize public generate inputs into the internal request shape."""
-    if prompt is None and messages is None:
-        raise ValueError("Either prompt or messages must be provided.")
-
-    if prompt is not None and messages is not None:
-        raise ValueError("Only one of prompt or messages can be provided.")
+    validate_prompt_or_messages(prompt, messages)
 
     if prompt is not None:
         messages = [{"role": "user", "content": prompt}]

nemoguardrails/rails/llm/generation/generation_workflow.py

@@ -19,6 +19,7 @@
 
 from nemoguardrails.actions.llm.utils import get_colang_history
 from nemoguardrails.colang.v2_x.runtime.flows import State
+from nemoguardrails.rails.llm.colang_turns.colang_turns import run_colang_turn
 from nemoguardrails.rails.llm.conversation.conversation_events import events_for_messages
 from nemoguardrails.rails.llm.generation.bot_messages import bot_message_from_colang_events
 from nemoguardrails.rails.llm.generation.generation_context import (
@@ -29,13 +30,12 @@
     generation_event_metadata,
     generation_response_from_colang_turn,
 )
-from nemoguardrails.rails.llm.generation.generation_tracing import (
+from nemoguardrails.rails.llm.generation.tracing import (
     export_generation_trace,
     prepare_generation_tracing,
     restore_generation_trace_log,
 )
 from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse
-from nemoguardrails.rails.llm.runtime.colang_turns import run_colang_turn
 from nemoguardrails.rails.llm.types import StandardGenerationSurface
 from nemoguardrails.rails.llm.utils import get_history_cache_key
 

nemoguardrails/rails/llm/generation/generation_tracing.py → nemoguardrails/rails/llm/generation/tracing.py

@@ -18,20 +18,33 @@
 from dataclasses import dataclass
 from typing import Any, List, Optional
 
+from nemoguardrails.rails.llm.config import RailsConfig
 from nemoguardrails.rails.llm.options import (
     GenerationLogOptions,
     GenerationOptions,
     GenerationResponse,
 )
+from nemoguardrails.tracing.adapters.base import InteractionLogAdapter
 
 __all__ = [
     "GenerationTracingState",
+    "create_startup_tracing_adapters",
     "export_generation_trace",
     "prepare_generation_tracing",
     "restore_generation_trace_log",
 ]
 
 
+def create_startup_tracing_adapters(config: RailsConfig) -> list[InteractionLogAdapter] | None:
+    """Build the tracing log adapters configured at startup (None when tracing is off)."""
+    if not config.tracing:
+        return None
+
+    from nemoguardrails.tracing import create_log_adapters
+
+    return create_log_adapters(config.tracing)
+
+
 @dataclass
 class GenerationTracingState:
     gen_options: Optional[GenerationOptions]

nemoguardrails/rails/llm/llmrails.py

@@ -43,28 +43,31 @@
 from nemoguardrails.logging.verbose import set_verbose
 from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop
 from nemoguardrails.rails.llm.checks import rails_check
+from nemoguardrails.rails.llm.colang_turns.colang_turns import (
+    generate_colang_events,
+    process_colang_events,
+    process_events_semaphore,
+)
 from nemoguardrails.rails.llm.config import OutputRailsStreamingConfig, RailsConfig
+from nemoguardrails.rails.llm.embedding.embedding_search import EmbeddingSearchState
 from nemoguardrails.rails.llm.generation.generation_context import (
     ensure_explain_info,
     explain_info_for_current_context,
     start_generation_request_context,
 )
 from nemoguardrails.rails.llm.generation.generation_request import (
     prepare_generation_request_for_runtime,
+    validate_prompt_or_messages,
     validate_public_state,
 )
 from nemoguardrails.rails.llm.generation.generation_workflow import generate_standard_async
+from nemoguardrails.rails.llm.generation.tracing import create_startup_tracing_adapters
 from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse, RailsResult, RailType
-from nemoguardrails.rails.llm.runtime.colang_runtime import runtime_for_colang_version
-from nemoguardrails.rails.llm.runtime.colang_turns import (
-    generate_colang_events,
-    process_colang_events,
-    process_events_semaphore,
-)
+from nemoguardrails.rails.llm.startup.colang_runtime import runtime_for_colang_version
 from nemoguardrails.rails.llm.startup.config_preparation import prepare_llmrails_config
 from nemoguardrails.rails.llm.startup.config_py import load_config_py_modules, run_config_py_init_hooks
 from nemoguardrails.rails.llm.startup.config_validation import validate_llmrails_config
-from nemoguardrails.rails.llm.startup.embedding_search import EmbeddingSearchState, apply_embedding_model_config
+from nemoguardrails.rails.llm.startup.embedding_config import apply_embedding_model_config
 from nemoguardrails.rails.llm.startup.generation_actions import register_llm_generation_actions
 from nemoguardrails.rails.llm.startup.knowledge_base import init_knowledge_base
 from nemoguardrails.rails.llm.startup.llm_action_caches import initialize_llm_action_caches
@@ -73,7 +76,6 @@
     model_kwargs_from_config,
     sync_update_llm_bindings,
 )
-from nemoguardrails.rails.llm.startup.tracing import create_startup_tracing_adapters
 from nemoguardrails.rails.llm.streaming.generation_stream import (
     generation_token_stream,
     validate_streaming_with_output_rails,
@@ -112,7 +114,7 @@ class LLMRails(BaseGuardrails):
     _kb: Any
     _log_adapters: Any
     _llm_generation_actions: Any
-    _verbose: bool
+    verbose: bool
     events_history_cache: dict[str, list[dict]]
     llm: Optional[LLMModel]
     runtime: Runtime
@@ -230,14 +232,6 @@ def passthrough_fn(self):
     def passthrough_fn(self, fn):
         self._llm_generation_actions._passthrough_fn = fn
 
-    @property
-    def verbose(self) -> bool:
-        return self._verbose
-
-    @verbose.setter
-    def verbose(self, verbose: bool) -> None:
-        self._verbose = verbose
-
     def __init__(
         self,
         config: RailsConfig,
@@ -397,12 +391,7 @@ async def generate_async(
             The completion (when a prompt is provided) or the next message.
 
         System messages are not yet supported."""
-        if prompt is None and messages is None:
-            raise ValueError("Either prompt or messages must be provided.")
-
-        if prompt is not None and messages is not None:
-            raise ValueError("Only one of prompt or messages can be provided.")
-
+        validate_prompt_or_messages(prompt, messages)
         validate_public_state(self.config, state)
         prepared_request = prepare_generation_request_for_runtime(
             prompt=prompt,

nemoguardrails/rails/llm/startup/embedding_config.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Startup-phase derivation of the default embedding-search settings from config."""
+
+from typing import Any, Dict, Optional, Tuple
+
+from nemoguardrails.rails.llm.config import RailsConfig
+
+__all__ = ["apply_embedding_model_config"]
+
+
+def apply_embedding_model_config(
+    config: RailsConfig,
+    default_embedding_model: Optional[str],
+    default_embedding_engine: Optional[str],
+    default_embedding_params: Dict[str, Any],
+) -> Tuple[Optional[str], Optional[str], Dict[str, Any]]:
+    """Apply an embeddings model config to the default embedding search settings."""
+    for model in config.models:
+        if model.type != "embeddings":
+            continue
+
+        default_embedding_model = model.model
+        default_embedding_engine = model.engine
+        default_embedding_params = model.parameters or {}
+
+        for esp in [
+            config.core.embedding_search_provider,
+            config.knowledge_base.embedding_search_provider,
+        ]:
+            if esp.name != "default":
+                continue
+            if "embedding_model" not in esp.parameters and model.model is not None:
+                esp.parameters["embedding_model"] = model.model
+            if "embedding_engine" not in esp.parameters and model.engine is not None:
+                esp.parameters["embedding_engine"] = model.engine
+
+        break
+
+    return default_embedding_model, default_embedding_engine, default_embedding_params

nemoguardrails/rails/llm/types.py

@@ -131,7 +131,7 @@ def events_history_cache(self) -> Dict[str, List[dict]]: ...
 
 
 class HasVerbose(Protocol):
-    """Reads the verbose flag (a ``property`` on the concrete class)."""
+    """Reads the verbose flag."""
 
     @property
     def verbose(self) -> bool: ...
@@ -177,7 +177,7 @@ class ConversationEventSurface(HasConfig, HasEventsHistoryCache, Protocol):
 
 
 class ColangTurnSurface(HasConfig, HasRuntime, HasVerbose, Protocol):
-    """Surface used by ``runtime.colang_turns`` to run a Colang turn."""
+    """Surface used by ``colang_turns.colang_turns`` to run a Colang turn."""
 
 
 class StandardGenerationSurface(ConversationEventSurface, ColangTurnSurface, HasExplainInfo, HasLogAdapters, Protocol):

tests/guardrails/test_public_api_deprecations.py

@@ -35,8 +35,8 @@
 from nemoguardrails import Guardrails
 from nemoguardrails.guardrails.iorails import IORails
 from nemoguardrails.rails.llm.config import RailsConfig
+from nemoguardrails.rails.llm.embedding.embedding_search import EmbeddingSearchState
 from nemoguardrails.rails.llm.llmrails import LLMRails
-from nemoguardrails.rails.llm.startup.embedding_search import EmbeddingSearchState
 from tests.guardrails.test_data import CONTENT_SAFETY_CONFIG
 
 

tests/rails/llm/test_colang_runtime.py

@@ -18,15 +18,15 @@
 from nemoguardrails.colang.v1_0.runtime.runtime import RuntimeV1_0
 from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x
 from nemoguardrails.exceptions import InvalidRailsConfigurationError
-from nemoguardrails.rails.llm import runtime as runtime_package
+from nemoguardrails.rails.llm import colang_turns as colang_turns_package
 from nemoguardrails.rails.llm.config import RailsConfig
 from nemoguardrails.rails.llm.llmrails import LLMRails
-from nemoguardrails.rails.llm.runtime import colang_runtime
-from nemoguardrails.rails.llm.runtime.colang_runtime import runtime_for_colang_version
+from nemoguardrails.rails.llm.startup import colang_runtime
+from nemoguardrails.rails.llm.startup.colang_runtime import runtime_for_colang_version
 
 
-def test_runtime_package_has_no_star_exports():
-    assert runtime_package.__all__ == []
+def test_colang_turns_package_has_no_star_exports():
+    assert colang_turns_package.__all__ == []
 
 
 def test_runtime_for_colang_version_uses_v1_runtime():

tests/rails/llm/test_colang_turns.py

@@ -23,12 +23,12 @@
 from nemoguardrails.context import llm_stats_var, streaming_handler_var
 from nemoguardrails.logging.stats import LLMStats
 from nemoguardrails.rails.llm import llmrails as llmrails_module
-from nemoguardrails.rails.llm.llmrails import LLMRails
-from nemoguardrails.rails.llm.runtime.colang_turns import (
+from nemoguardrails.rails.llm.colang_turns.colang_turns import (
     generate_colang_events,
     process_colang_events,
     run_colang_turn,
 )
+from nemoguardrails.rails.llm.llmrails import LLMRails
 from nemoguardrails.streaming import END_OF_STREAM