refactor(llmrails): consolidate extracted module contracts (7/9)

nemoguardrails/rails/llm/checks/rails_check.py

@@ -18,30 +18,19 @@
 import logging
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any, List, Optional, Protocol
+from typing import List, Optional
 
 from nemoguardrails.rails.llm.options import (
     GenerationResponse,
     RailsResult,
     RailStatus,
     RailType,
 )
+from nemoguardrails.rails.llm.types import RailsCheckSurface
 
 log = logging.getLogger(__name__)
 
-__all__ = ["RailsCheckRuntime", "check_messages"]
-
-
-class RailsCheckRuntime(Protocol):
-    config: Any
-    runtime: Any
-
-    async def generate_async(
-        self,
-        *,
-        messages: List[dict],
-        options: dict,
-    ) -> object: ...
+__all__ = ["RailsCheckSurface", "check_messages"]
 
 
 @dataclass
@@ -52,7 +41,7 @@ class RailsCheckPlan:
 
 
 async def check_messages(
-    rails: RailsCheckRuntime,
+    rails: RailsCheckSurface,
     messages: List[dict],
     rail_types: Optional[List[RailType]] = None,
 ) -> RailsResult:

nemoguardrails/rails/llm/conversation/conversation_events.py

@@ -16,27 +16,20 @@
 """Conversation message conversion to Colang events."""
 
 from collections.abc import MutableMapping
-from typing import Any, List, Protocol
+from typing import Any, List
 
 from nemoguardrails.colang.v2_x.runtime.flows import Action
+from nemoguardrails.rails.llm.types import ConversationEventSurface
 from nemoguardrails.rails.llm.utils import get_history_cache_key
 from nemoguardrails.utils import new_event_dict, new_uuid
 
 __all__ = [
-    "ConversationEventRails",
+    "ConversationEventSurface",
     "events_for_messages",
     "events_history_cache_prefix",
 ]
 
 
-class ConversationEventRails(Protocol):
-    @property
-    def config(self) -> Any: ...
-
-    @property
-    def events_history_cache(self) -> MutableMapping[str, list[dict]]: ...
-
-
 def events_history_cache_prefix(
     events_history_cache: MutableMapping[str, list[dict]],
     messages: List[dict],
@@ -58,7 +51,7 @@ def events_history_cache_prefix(
     return 0, []
 
 
-def events_for_messages(rails: ConversationEventRails, messages: List[dict], state: Any) -> List[dict]:
+def events_for_messages(rails: ConversationEventSurface, messages: List[dict], state: Any) -> List[dict]:
     """Return Colang events corresponding to OpenAI-style messages."""
     events = []
 

nemoguardrails/rails/llm/generation/generation_workflow.py

@@ -15,7 +15,7 @@
 
 import logging
 import time
-from typing import Any, Optional, Union
+from typing import Optional, Union
 
 from nemoguardrails.actions.llm.utils import get_colang_history
 from nemoguardrails.colang.v2_x.runtime.flows import State
@@ -36,15 +36,16 @@
 )
 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
 
 __all__ = ["generate_standard_async"]
 
-log = logging.getLogger("nemoguardrails.rails.llm.llmrails")
+log = logging.getLogger(__name__)
 
 
 async def generate_standard_async(
-    rails: Any,
+    rails: StandardGenerationSurface,
     *,
     prompt: Optional[str],
     messages: list[dict],

nemoguardrails/rails/llm/llmrails.py

@@ -110,8 +110,10 @@ class LLMRails(BaseGuardrails):
     embedding_search: EmbeddingSearchState
     _explain_info: Optional[ExplainInfo]
     _kb: Any
+    _log_adapters: Any
     _llm_generation_actions: Any
     _verbose: bool
+    events_history_cache: dict[str, list[dict]]
     llm: Optional[LLMModel]
     runtime: Runtime
 

nemoguardrails/rails/llm/runtime/colang_turns.py

@@ -19,13 +19,14 @@
 import json
 import logging
 import time
-from typing import Any, List, Optional, Protocol, Tuple, Union, cast
+from typing import Any, List, Optional, Tuple, Union, cast
 
 from nemoguardrails.actions.llm.utils import get_colang_history
 from nemoguardrails.colang.v2_x.runtime.flows import State
 from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x
 from nemoguardrails.context import llm_stats_var, streaming_handler_var
 from nemoguardrails.logging.stats import LLMStats
+from nemoguardrails.rails.llm.types import ColangTurnSurface
 from nemoguardrails.streaming import END_OF_STREAM
 from nemoguardrails.utils import extract_error_json
 
@@ -34,27 +35,16 @@
 process_events_semaphore = asyncio.Semaphore(1)
 
 __all__ = [
-    "ColangTurnRails",
+    "ColangTurnSurface",
     "generate_colang_events",
     "process_colang_events",
     "process_events_semaphore",
     "run_colang_turn",
 ]
 
 
-class ColangTurnRails(Protocol):
-    @property
-    def config(self) -> Any: ...
-
-    @property
-    def runtime(self) -> Any: ...
-
-    @property
-    def verbose(self) -> bool: ...
-
-
 async def run_colang_turn(
-    rails: ColangTurnRails,
+    rails: ColangTurnSurface,
     events: List[dict],
     state: Any,
     processing_log: List[dict],
@@ -66,7 +56,7 @@ async def run_colang_turn(
 
 
 async def _run_colang_1_turn(
-    rails: ColangTurnRails,
+    rails: ColangTurnSurface,
     events: List[dict],
     state: Any,
     processing_log: List[dict],
@@ -94,7 +84,7 @@ async def _run_colang_1_turn(
 
 
 async def _run_colang_2_turn(
-    rails: ColangTurnRails,
+    rails: ColangTurnSurface,
     events: List[dict],
     state: Any,
 ) -> List[dict]:
@@ -113,7 +103,7 @@ async def _run_colang_2_turn(
     return new_events
 
 
-async def generate_colang_events(rails: ColangTurnRails, events: List[dict]) -> List[dict]:
+async def generate_colang_events(rails: ColangTurnSurface, events: List[dict]) -> List[dict]:
     """Generate the next Colang 1.0 events for an event history."""
     t0 = time.time()
 
@@ -138,7 +128,7 @@ async def generate_colang_events(rails: ColangTurnRails, events: List[dict]) ->
 
 
 async def process_colang_events(
-    rails: ColangTurnRails,
+    rails: ColangTurnSurface,
     events: List[dict],
     state: Union[Optional[dict], State] = None,
     blocking: bool = False,

nemoguardrails/rails/llm/startup/colang_flows.py

@@ -29,7 +29,7 @@
 from nemoguardrails.exceptions import InvalidRailsConfigurationError
 from nemoguardrails.rails.llm.config import RailsConfig
 
-log = logging.getLogger("nemoguardrails.rails.llm.llmrails")
+log = logging.getLogger(__name__)
 
 _NEMOGUARDRAILS_PACKAGE = "nemoguardrails"
 

nemoguardrails/rails/llm/startup/generation_actions.py

@@ -15,49 +15,27 @@
 
 """LLM generation action registration."""
 
-from typing import Any, Protocol, Type
+from typing import Any, Type
 
 from nemoguardrails.actions.llm.generation import LLMGenerationActions
 from nemoguardrails.actions.v2_x.generation import LLMGenerationActionsV2dotx
+from nemoguardrails.rails.llm.types import GenerationActionsSurface
 
 __all__ = [
-    "GenerationActionRails",
-    "GenerationActionRuntime",
+    "GenerationActionsSurface",
     "LLMGenerationActions",
     "LLMGenerationActionsV2dotx",
     "generation_actions_class_for_colang_version",
     "register_llm_generation_actions",
 ]
 
 
-class GenerationActionRuntime(Protocol):
-    @property
-    def llm_task_manager(self) -> Any: ...
-
-    def register_actions(self, actions_obj: Any, /, override: bool = True) -> None: ...
-
-
-class GenerationActionRails(Protocol):
-    _llm_generation_actions: Any
-
-    @property
-    def config(self) -> Any: ...
-
-    @property
-    def llm(self) -> Any: ...
-
-    @property
-    def runtime(self) -> GenerationActionRuntime: ...
-
-    def _get_embeddings_search_provider_instance(self, esp_config: Any = None) -> Any: ...
-
-
 def generation_actions_class_for_colang_version(colang_version: str) -> Type[Any]:
     """Return the LLM generation actions class selected by the Colang version."""
     return LLMGenerationActions if colang_version == "1.0" else LLMGenerationActionsV2dotx
 
 
-def register_llm_generation_actions(rails: GenerationActionRails, verbose: bool) -> None:
+def register_llm_generation_actions(rails: GenerationActionsSurface, verbose: bool) -> None:
     """Create and register the LLM generation actions for an LLMRails instance."""
     llm_generation_actions_class = generation_actions_class_for_colang_version(rails.config.colang_version)
     rails._llm_generation_actions = llm_generation_actions_class(

nemoguardrails/rails/llm/startup/knowledge_base.py

@@ -23,6 +23,7 @@
 from nemoguardrails.kb.kb import KnowledgeBase
 from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop
 from nemoguardrails.rails.llm.config import EmbeddingSearchProvider, RailsConfig
+from nemoguardrails.rails.llm.types import KnowledgeBaseSurface
 from nemoguardrails.utils import get_or_create_event_loop
 
 __all__ = ["build_knowledge_base_for_docs", "init_knowledge_base"]
@@ -47,7 +48,7 @@ async def build_knowledge_base_for_docs(
     return kb
 
 
-def init_knowledge_base(rails) -> None:
+def init_knowledge_base(rails: KnowledgeBaseSurface) -> None:
     """Initialize and register the LLMRails knowledge base."""
     rails._kb = None
 

nemoguardrails/rails/llm/startup/llm_action_caches.py

@@ -16,12 +16,13 @@
 """LLM action model caches."""
 
 import logging
-from typing import Dict, Protocol
+from typing import Dict
 
 from nemoguardrails.llm.cache import CacheInterface, LFUCache
 from nemoguardrails.rails.llm.config import Model, RailsConfig
+from nemoguardrails.rails.llm.types import LLMActionCacheSurface
 
-log = logging.getLogger("nemoguardrails.rails.llm.llmrails")
+log = logging.getLogger(__name__)
 
 __all__ = [
     "build_llm_action_cache",
@@ -30,18 +31,6 @@
 ]
 
 
-class LLMActionCacheRuntime(Protocol):
-    def register_action_param(self, name: str, value: object) -> None: ...
-
-
-class LLMActionCacheRails(Protocol):
-    @property
-    def config(self) -> RailsConfig: ...
-
-    @property
-    def runtime(self) -> LLMActionCacheRuntime: ...
-
-
 def build_llm_action_cache(model: Model) -> LFUCache:
     """Build the cache configured for one LLM action model."""
     if model.cache is None:
@@ -87,7 +76,7 @@ def build_llm_action_caches(config: RailsConfig) -> Dict[str, CacheInterface]:
     return model_caches
 
 
-def initialize_llm_action_caches(rails: LLMActionCacheRails) -> None:
+def initialize_llm_action_caches(rails: LLMActionCacheSurface) -> None:
     """Initialize configured action model caches for an LLMRails instance."""
     model_caches = build_llm_action_caches(rails.config)
     if model_caches:

nemoguardrails/rails/llm/startup/llm_action_models.py

@@ -17,41 +17,26 @@
 
 import logging
 import os
-from typing import Any, Callable, Dict, Protocol
+from typing import Any, Callable, Dict
 
 from nemoguardrails.exceptions import InvalidModelConfigurationError
 from nemoguardrails.llm.models.initializer import ModelInitializationError
+from nemoguardrails.rails.llm.types import LLMActionModelsSurface
 from nemoguardrails.types import LLMModel
 
-log = logging.getLogger("nemoguardrails.rails.llm.llmrails")
+log = logging.getLogger(__name__)
 
 InitLLM = Callable[..., LLMModel]
 
 __all__ = [
     "InitLLM",
-    "LLMActionRails",
-    "LLMActionRuntime",
+    "LLMActionModelsSurface",
     "load_llm_action_models",
     "model_kwargs_from_config",
     "sync_update_llm_bindings",
 ]
 
 
-class LLMActionRuntime(Protocol):
-    def register_action_param(self, name: str, value: Any) -> None: ...
-
-
-class LLMActionRails(Protocol):
-    llm: Any
-    _llm_generation_actions: Any
-
-    @property
-    def config(self) -> Any: ...
-
-    @property
-    def runtime(self) -> LLMActionRuntime: ...
-
-
 def model_kwargs_from_config(model_config: Any) -> Dict[str, Any]:
     """Prepare model kwargs, including optional API key environment variables."""
     kwargs = dict(model_config.parameters or {})
@@ -64,14 +49,14 @@ def model_kwargs_from_config(model_config: Any) -> Dict[str, Any]:
     return kwargs
 
 
-def sync_update_llm_bindings(rails: LLMActionRails, llm: LLMModel) -> None:
+def sync_update_llm_bindings(rails: LLMActionModelsSurface, llm: LLMModel) -> None:
     """Synchronize the main LLM bindings after a public update."""
     rails.llm = llm
     rails._llm_generation_actions.llm = llm
     rails.runtime.register_action_param("llm", llm)
 
 
-def load_llm_action_models(rails: LLMActionRails, init_llm: InitLLM) -> None:
+def load_llm_action_models(rails: LLMActionModelsSurface, init_llm: InitLLM) -> None:
     """Load the main and action LLMs configured for an LLMRails instance."""
     from nemoguardrails._compat.langchain_kwargs import check_langchain_kwargs
     from nemoguardrails.llm.frameworks import get_default_framework