diff --git a/nemoguardrails/rails/llm/runtime/__init__.py b/nemoguardrails/rails/llm/colang_turns/__init__.py similarity index 100% rename from nemoguardrails/rails/llm/runtime/__init__.py rename to nemoguardrails/rails/llm/colang_turns/__init__.py diff --git a/nemoguardrails/rails/llm/runtime/colang_turns.py b/nemoguardrails/rails/llm/colang_turns/colang_turns.py similarity index 100% rename from nemoguardrails/rails/llm/runtime/colang_turns.py rename to nemoguardrails/rails/llm/colang_turns/colang_turns.py diff --git a/nemoguardrails/rails/llm/startup/tracing.py b/nemoguardrails/rails/llm/embedding/__init__.py similarity index 61% rename from nemoguardrails/rails/llm/startup/tracing.py rename to nemoguardrails/rails/llm/embedding/__init__.py index b5ca19987a..dc33922794 100644 --- a/nemoguardrails/rails/llm/startup/tracing.py +++ b/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__ = [] diff --git a/nemoguardrails/rails/llm/startup/embedding_search.py b/nemoguardrails/rails/llm/embedding/embedding_search.py similarity index 75% rename from nemoguardrails/rails/llm/startup/embedding_search.py rename to nemoguardrails/rails/llm/embedding/embedding_search.py index 01a63bcc71..9da833e2e4 100644 --- a/nemoguardrails/rails/llm/startup/embedding_search.py +++ b/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], diff --git a/nemoguardrails/rails/llm/generation/generation_request.py b/nemoguardrails/rails/llm/generation/generation_request.py index a984792a0c..74956e3897 100644 --- a/nemoguardrails/rails/llm/generation/generation_request.py +++ b/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}] diff --git a/nemoguardrails/rails/llm/generation/generation_workflow.py b/nemoguardrails/rails/llm/generation/generation_workflow.py index 9bc2200e7e..f95f321478 100644 --- a/nemoguardrails/rails/llm/generation/generation_workflow.py +++ b/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 diff --git a/nemoguardrails/rails/llm/generation/generation_tracing.py b/nemoguardrails/rails/llm/generation/tracing.py similarity index 87% rename from nemoguardrails/rails/llm/generation/generation_tracing.py rename to nemoguardrails/rails/llm/generation/tracing.py index b2a23db330..3d9d72a251 100644 --- a/nemoguardrails/rails/llm/generation/generation_tracing.py +++ b/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] diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index eb4abe8ec4..3e24078c4e 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -43,7 +43,13 @@ 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, @@ -51,20 +57,17 @@ ) 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, diff --git a/nemoguardrails/rails/llm/runtime/colang_runtime.py b/nemoguardrails/rails/llm/startup/colang_runtime.py similarity index 100% rename from nemoguardrails/rails/llm/runtime/colang_runtime.py rename to nemoguardrails/rails/llm/startup/colang_runtime.py diff --git a/nemoguardrails/rails/llm/startup/embedding_config.py b/nemoguardrails/rails/llm/startup/embedding_config.py new file mode 100644 index 0000000000..a4cd733df7 --- /dev/null +++ b/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 diff --git a/nemoguardrails/rails/llm/types.py b/nemoguardrails/rails/llm/types.py index 5bc56bada4..5eb44bcafa 100644 --- a/nemoguardrails/rails/llm/types.py +++ b/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): diff --git a/tests/guardrails/test_public_api_deprecations.py b/tests/guardrails/test_public_api_deprecations.py index 0495fdb3d6..3a5a584c3d 100644 --- a/tests/guardrails/test_public_api_deprecations.py +++ b/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 diff --git a/tests/rails/llm/test_colang_runtime.py b/tests/rails/llm/test_colang_runtime.py index 9f6f0077de..eab3387e34 100644 --- a/tests/rails/llm/test_colang_runtime.py +++ b/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(): diff --git a/tests/rails/llm/test_colang_turns.py b/tests/rails/llm/test_colang_turns.py index 75f0c8453c..fa3790434a 100644 --- a/tests/rails/llm/test_colang_turns.py +++ b/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 diff --git a/tests/rails/llm/test_embedding_search.py b/tests/rails/llm/test_embedding_search.py index 4676b12651..b60f44520a 100644 --- a/tests/rails/llm/test_embedding_search.py +++ b/tests/rails/llm/test_embedding_search.py @@ -18,22 +18,21 @@ from nemoguardrails.embeddings.basic import BasicEmbeddingsIndex from nemoguardrails.embeddings.index import EmbeddingsIndex from nemoguardrails.rails.llm.config import EmbeddingSearchProvider, Model, RailsConfig -from nemoguardrails.rails.llm.startup import embedding_search as startup_embedding_search -from nemoguardrails.rails.llm.startup.embedding_search import ( +from nemoguardrails.rails.llm.embedding import embedding_search as embedding_search_module +from nemoguardrails.rails.llm.embedding.embedding_search import ( DEFAULT_EMBEDDING_ENGINE, DEFAULT_EMBEDDING_MODEL, EmbeddingSearchState, - apply_embedding_model_config, get_embedding_search_provider_instance, ) +from nemoguardrails.rails.llm.startup.embedding_config import apply_embedding_model_config -def test_embedding_search_startup_module_exports_public_helpers(): - assert startup_embedding_search.__all__ == [ +def test_embedding_search_module_exports_public_helpers(): + assert embedding_search_module.__all__ == [ "DEFAULT_EMBEDDING_ENGINE", "DEFAULT_EMBEDDING_MODEL", "EmbeddingSearchState", - "apply_embedding_model_config", "get_embedding_search_provider_instance", ] diff --git a/tests/rails/llm/test_generation_request.py b/tests/rails/llm/test_generation_request.py index b0852fa6e1..243f1fa384 100644 --- a/tests/rails/llm/test_generation_request.py +++ b/tests/rails/llm/test_generation_request.py @@ -39,6 +39,7 @@ def test_generation_request_exports(): "PreparedGenerationRequest", "normalize_generation_request", "prepare_generation_request_for_runtime", + "validate_prompt_or_messages", ] diff --git a/tests/rails/llm/test_generation_tracing.py b/tests/rails/llm/test_generation_tracing.py index b69f201ce2..01d450fd9f 100644 --- a/tests/rails/llm/test_generation_tracing.py +++ b/tests/rails/llm/test_generation_tracing.py @@ -17,8 +17,8 @@ import pytest -from nemoguardrails.rails.llm.generation import generation_tracing -from nemoguardrails.rails.llm.generation.generation_tracing import ( +from nemoguardrails.rails.llm.generation import tracing as generation_tracing +from nemoguardrails.rails.llm.generation.tracing import ( export_generation_trace, prepare_generation_tracing, restore_generation_trace_log, @@ -35,6 +35,7 @@ def test_generation_tracing_exports(): assert generation_tracing.__all__ == [ "GenerationTracingState", + "create_startup_tracing_adapters", "export_generation_trace", "prepare_generation_tracing", "restore_generation_trace_log", diff --git a/tests/rails/llm/test_llmrails_startup_boundary.py b/tests/rails/llm/test_llmrails_startup_boundary.py index 2e0649abcc..d5a17e98b0 100644 --- a/tests/rails/llm/test_llmrails_startup_boundary.py +++ b/tests/rails/llm/test_llmrails_startup_boundary.py @@ -18,8 +18,8 @@ from nemoguardrails.embeddings.index import EmbeddingsIndex from nemoguardrails.rails.llm import llmrails from nemoguardrails.rails.llm.config import RailsConfig, TracingConfig -from nemoguardrails.rails.llm.startup.embedding_search import DEFAULT_EMBEDDING_ENGINE, DEFAULT_EMBEDDING_MODEL -from nemoguardrails.rails.llm.startup.tracing import create_startup_tracing_adapters +from nemoguardrails.rails.llm.embedding.embedding_search import DEFAULT_EMBEDDING_ENGINE, DEFAULT_EMBEDDING_MODEL +from nemoguardrails.rails.llm.generation.tracing import create_startup_tracing_adapters class RuntimeWithRegistries: