chore(types): Type-clean rails (86 errors)

.pre-commit-config.yaml

@@ -23,7 +23,14 @@ repos:
         args:
           - --license-filepath
           - LICENSE.md
-
+  - repo: local
+    hooks:
+      - id: pyright
+        name: pyright
+        entry: poetry run pyright
+        language: system
+        types: [python]
+        pass_filenames: false
 # Deactivating this for now.
 #  - repo: https://github.com/pycqa/pylint
 #    rev: v2.17.0

nemoguardrails/actions/llm/generation.py

@@ -82,7 +82,7 @@ class LLMGenerationActions:
     def __init__(
         self,
         config: RailsConfig,
-        llm: Union[BaseLLM, BaseChatModel],
+        llm: Optional[Union[BaseLLM, BaseChatModel]],
         llm_task_manager: LLMTaskManager,
         get_embedding_search_provider_instance: Callable[
             [Optional[EmbeddingSearchProvider]], EmbeddingsIndex

nemoguardrails/context.py

@@ -14,25 +14,45 @@
 # limitations under the License.
 
 import contextvars
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-streaming_handler_var = contextvars.ContextVar("streaming_handler", default=None)
+from nemoguardrails.logging.explain import LLMCallInfo
+
+if TYPE_CHECKING:
+    from nemoguardrails.logging.explain import ExplainInfo
+    from nemoguardrails.logging.stats import LLMStats
+    from nemoguardrails.rails.llm.options import GenerationOptions
+    from nemoguardrails.streaming import StreamingHandler
+
+streaming_handler_var: contextvars.ContextVar[
+    Optional["StreamingHandler"]
+] = contextvars.ContextVar("streaming_handler", default=None)
 
 # The object that holds additional explanation information.
-explain_info_var = contextvars.ContextVar("explain_info", default=None)
+explain_info_var: contextvars.ContextVar[
+    Optional["ExplainInfo"]
+] = contextvars.ContextVar("explain_info", default=None)
 
 # The current LLM call.
-llm_call_info_var = contextvars.ContextVar("llm_call_info", default=None)
+llm_call_info_var: contextvars.ContextVar[
+    Optional[LLMCallInfo]
+] = contextvars.ContextVar("llm_call_info", default=None)
 
 # All the generation options applicable to the current context.
-generation_options_var = contextvars.ContextVar("generation_options", default=None)
+generation_options_var: contextvars.ContextVar[
+    Optional["GenerationOptions"]
+] = contextvars.ContextVar("generation_options", default=None)
 
 # The stats about the LLM calls.
-llm_stats_var = contextvars.ContextVar("llm_stats", default=None)
+llm_stats_var: contextvars.ContextVar[Optional["LLMStats"]] = contextvars.ContextVar(
+    "llm_stats", default=None
+)
 
 # The raw LLM request that comes from the user.
 # This is used in passthrough mode.
-raw_llm_request = contextvars.ContextVar("raw_llm_request", default=None)
+raw_llm_request: contextvars.ContextVar[
+    Optional[Union[str, List[Dict[str, Any]]]]
+] = contextvars.ContextVar("raw_llm_request", default=None)
 
 reasoning_trace_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
     "reasoning_trace", default=None

nemoguardrails/rails/llm/buffer.py

@@ -138,7 +138,8 @@ async def process_stream(
             ...     print(f"Processing: {context_formatted}")
             ...     print(f"User: {user_formatted}")
         """
-        ...
+        raise NotImplementedError
+        yield
 
     async def __call__(self, streaming_handler) -> AsyncGenerator[ChunkBatch, None]:
         """Callable interface that delegates to process_stream.

nemoguardrails/rails/llm/config.py

@@ -487,7 +487,7 @@ class OutputRails(BaseModel):
         description="The names of all the flows that implement output rails.",
     )
 
-    streaming: Optional[OutputRailsStreamingConfig] = Field(
+    streaming: OutputRailsStreamingConfig = Field(
         default_factory=OutputRailsStreamingConfig,
         description="Configuration for streaming output rails.",
     )
@@ -1128,7 +1128,9 @@ def _load_path(
 
     # the first .railsignore file found from cwd down to its subdirectories
     railsignore_path = utils.get_railsignore_path(config_path)
-    ignore_patterns = utils.get_railsignore_patterns(railsignore_path)
+    ignore_patterns = (
+        utils.get_railsignore_patterns(railsignore_path) if railsignore_path else set()
+    )
 
     if os.path.isdir(config_path):
         for root, _, files in os.walk(config_path, followlinks=True):
@@ -1245,8 +1247,8 @@ def _parse_colang_files_recursively(
         current_file, current_path = colang_files[len(parsed_colang_files)]
 
         with open(current_path, "r", encoding="utf-8") as f:
+            content = f.read()
             try:
-                content = f.read()
                 _parsed_config = parse_colang_file(
                     current_file, content=content, version=colang_version
                 )
@@ -1748,7 +1750,7 @@ def streaming_supported(self):
             # if we have output rails streaming enabled
             # we keep it in case it was needed when we have
             # support per rails
-            if self.rails.output.streaming.enabled:
+            if self.rails.output.streaming and self.rails.output.streaming.enabled:
                 return True
             return False