All but one error left to clean

Author: tgasser-nv

nemoguardrails/actions/llm/generation.py

@@ -972,7 +972,9 @@ async def generate_bot_message(
 
                     # We use the potentially updated $user_message. This means that even
                     # in passthrough mode, input rails can still alter the input.
-                    prompt: Optional[str] = context.get(
+                    prompt: Optional[
+                        str
+                    ] = context.get(  # pyright: ignore (TODO Refactor these branches into separate methods)
                         "user_message"
                     )  # pyright: ignore (TODO - refactor nested `prompt` definitions)
                     if not prompt:

nemoguardrails/actions/v2_x/generation.py

@@ -23,6 +23,7 @@
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.language_models.llms import BaseLLM
+from langchain_text_splitters import ElementType
 from pytest_asyncio.plugin import event_loop
 from rich.text import Text
 
@@ -215,7 +216,8 @@ async def _collect_user_intent_and_examples(
         # We add all currently active user intents (heads on match statements)
         heads = find_all_active_event_matchers(state)
         for head in heads:
-            element = get_element_from_head(state, head)
+            el = get_element_from_head(state, head)
+            element = el if type(el) == SpecOp else SpecOp(**cast(Dict, el))
             flow_state = state.flow_states[head.flow_state_uid]
             event = get_event_from_element(state, flow_state, element)
             if (
@@ -235,9 +237,11 @@ async def _collect_user_intent_and_examples(
                     ):
                         if flow_config.elements[1]["_type"] == "doc_string_stmt":
                             examples += "user action: <" + (
-                                flow_config.elements[1]["elements"][0]["elements"][
-                                    0
+                                flow_config.elements[1]["elements"][0][
+                                    "elements"
                                 ][  # pyright: ignore (TODO - Don't know where to even start with this line of code)
+                                    0
+                                ][
                                     "elements"
                                 ][
                                     0
@@ -279,7 +283,7 @@ async def generate_user_intent(  # pyright: ignore (TODO - Signature completely
         """Generate the canonical form for what the user said i.e. user intent."""
 
         # Use action specific llm if registered else fallback to main llm
-        llm = llm or self.llm
+        generation_llm: Union[BaseLLM, BaseChatModel] = llm if llm else self.llm
 
         log.info("Phase 1 :: Generating user intent")
         (
@@ -311,8 +315,8 @@ async def generate_user_intent(  # pyright: ignore (TODO - Signature completely
         )
 
         # We make this call with lowest temperature to have it as deterministic as possible.
-        with llm_params(llm, temperature=self.config.lowest_temperature):
-            result = await llm_call(llm, prompt, stop=stop)
+        with llm_params(generation_llm, temperature=self.config.lowest_temperature):
+            result = await llm_call(generation_llm, prompt, stop=stop)
 
         # Parse the output using the associated parser
         result = self.llm_task_manager.parse_task_output(
@@ -356,7 +360,7 @@ async def generate_user_intent_and_bot_action(
         """Generate the canonical form for what the user said i.e. user intent and a suitable bot action."""
 
         # Use action specific llm if registered else fallback to main llm
-        llm = llm or self.llm
+        generation_llm: Union[BaseLLM, BaseChatModel] = llm if llm else self.llm
 
         log.info("Phase 1 :: Generating user intent and bot action")
 
@@ -389,8 +393,8 @@ async def generate_user_intent_and_bot_action(
         )
 
         # We make this call with lowest temperature to have it as deterministic as possible.
-        with llm_params(llm, temperature=self.config.lowest_temperature):
-            result = await llm_call(llm, prompt, stop=stop)
+        with llm_params(generation_llm, temperature=self.config.lowest_temperature):
+            result = await llm_call(generation_llm, prompt, stop=stop)
 
         # Parse the output using the associated parser
         result = self.llm_task_manager.parse_task_output(
@@ -439,7 +443,12 @@ async def passthrough_llm_action(
         events: List[dict],
         llm: Optional[BaseLLM] = None,
     ):
+        if not llm:
+            raise Exception("No LLM provided to passthrough LLM Action")
+
         event = get_last_user_utterance_event_v2_x(events)
+        if not event:
+            raise Exception("Passthrough LLM Action couldn't find last user utterance")
 
         # We check if we have a raw request. If the guardrails API is using
         # the `generate_events` API, this will not be set.
@@ -465,7 +474,10 @@ async def passthrough_llm_action(
         # Initialize the LLMCallInfo object
         llm_call_info_var.set(LLMCallInfo(task=Task.GENERAL.value))
 
-        generation_options: GenerationOptions = generation_options_var.get()
+        generation_options: Optional[GenerationOptions] = generation_options_var.get()
+
+        streaming_handler: Optional[StreamingHandler] = streaming_handler_var.get()
+        custom_callback_handlers = [streaming_handler] if streaming_handler else None
 
         with llm_params(
             llm,
@@ -474,7 +486,7 @@ async def passthrough_llm_action(
             text = await llm_call(
                 llm,
                 user_message,
-                custom_callback_handlers=[streaming_handler_var.get()],
+                custom_callback_handlers=custom_callback_handlers,
             )
 
             text = self.llm_task_manager.parse_task_output(Task.GENERAL, output=text)
@@ -526,12 +538,12 @@ async def generate_flow_from_instructions(
             raise RuntimeError("No instruction flows index has been created.")
 
         # Use action specific llm if registered else fallback to main llm
-        llm = llm or self.llm
+        generation_llm: Union[BaseLLM, BaseChatModel] = llm if llm else self.llm
 
         log.info("Generating flow for instructions: %s", instructions)
 
         results = await self.instruction_flows_index.search(
-            text=instructions, max_results=5
+            text=instructions, max_results=5, threshold=None
         )
 
         examples = ""
@@ -557,8 +569,8 @@ async def generate_flow_from_instructions(
         )
 
         # We make this call with temperature 0 to have it as deterministic as possible.
-        with llm_params(llm, temperature=self.config.lowest_temperature):
-            result = await llm_call(llm, prompt)
+        with llm_params(generation_llm, temperature=self.config.lowest_temperature):
+            result = await llm_call(generation_llm, prompt)
 
         result = self.llm_task_manager.parse_task_output(
             task=Task.GENERATE_FLOW_FROM_INSTRUCTIONS, output=result
@@ -604,12 +616,15 @@ async def generate_flow_from_name(
             raise RuntimeError("No flows index has been created.")
 
         # Use action specific llm if registered else fallback to main llm
-        llm = llm or self.llm
+        generation_llm: Union[BaseLLM, BaseChatModel] = llm if llm else self.llm
 
         log.info("Generating flow for name: {name}")
 
+        if not self.instruction_flows_index:
+            raise Exception("No instruction flows index has been created.")
+
         results = await self.instruction_flows_index.search(
-            text=f"flow {name}", max_results=5
+            text=f"flow {name}", max_results=5, threshold=None
         )
 
         examples = ""
@@ -631,8 +646,8 @@ async def generate_flow_from_name(
         stop = self.llm_task_manager.get_stop_tokens(Task.GENERATE_FLOW_FROM_NAME)
 
         # We make this call with temperature 0 to have it as deterministic as possible.
-        with llm_params(llm, temperature=self.config.lowest_temperature):
-            result = await llm_call(llm, prompt, stop)
+        with llm_params(generation_llm, temperature=self.config.lowest_temperature):
+            result = await llm_call(generation_llm, prompt, stop=stop)
 
         result = self.llm_task_manager.parse_task_output(
             task=Task.GENERATE_FLOW_FROM_NAME, output=result
@@ -666,7 +681,7 @@ async def generate_flow_continuation(
             raise RuntimeError("No instruction flows index has been created.")
 
         # Use action specific llm if registered else fallback to main llm
-        llm = llm or self.llm
+        generation_llm: Union[BaseLLM, BaseChatModel] = llm if llm else self.llm
 
         log.info("Generating flow continuation.")
 
@@ -675,7 +690,11 @@ async def generate_flow_continuation(
         # We use the last line from the history to search for relevant flows
         search_text = colang_history.split("\n")[-1]
 
-        results = await self.flows_index.search(text=search_text, max_results=10)
+        if self.flows_index is None:
+            raise RuntimeError("No flows index has been created.")
+        results = await self.flows_index.search(
+            text=search_text, max_results=10, threshold=None
+        )
 
         examples = ""
         for result in reversed(results):
@@ -697,8 +716,8 @@ async def generate_flow_continuation(
         )
 
         # We make this call with temperature 0 to have it as deterministic as possible.
-        with llm_params(llm, temperature=temperature):
-            result = await llm_call(llm, prompt)
+        with llm_params(generation_llm, temperature=temperature):
+            result = await llm_call(generation_llm, prompt)
 
         # TODO: Currently, we only support generating a bot action as continuation. This could be generalized
         # Colang statements.
@@ -775,7 +794,7 @@ async def create_flow(
         }
 
     @action(name="GenerateValueAction", is_system_action=True, execute_async=True)
-    async def generate_value(
+    async def generate_value(  # pyright: ignore (TODO - different arguments to base-class)
         self,
         state: State,
         instructions: str,
@@ -791,15 +810,21 @@ async def generate_value(
         :param llm: Custom llm model to generate_value
         """
         # Use action specific llm if registered else fallback to main llm
-        llm = llm or self.llm
+        generation_llm: Union[BaseLLM, BaseChatModel] = llm if llm else self.llm
 
         # We search for the most relevant flows.
         examples = ""
         if self.flows_index:
-            if var_name:
-                results = await self.flows_index.search(
-                    text=f"${var_name} = ", max_results=5
+            results = (
+                await self.flows_index.search(
+                    text=f"${var_name} = ", max_results=5, threshold=None
                 )
+                if var_name
+                else None
+            )
+
+            if not results:
+                raise Exception("No results found while generating value")
 
             # We add these in reverse order so the most relevant is towards the end.
             for result in reversed(results):
@@ -827,8 +852,8 @@ async def generate_value(
             Task.GENERATE_USER_INTENT_FROM_USER_ACTION
         )
 
-        with llm_params(llm, temperature=0.1):
-            result = await llm_call(llm, prompt, stop)
+        with llm_params(generation_llm, temperature=0.1):
+            result = await llm_call(generation_llm, prompt, stop=stop)
 
         # Parse the output using the associated parser
         result = self.llm_task_manager.parse_task_output(
@@ -871,11 +896,17 @@ async def generate_flow(
     ) -> dict:
         """Generate the body for a flow."""
         # Use action specific llm if registered else fallback to main llm
-        llm = llm or self.llm
+        generation_llm: Union[BaseLLM, BaseChatModel] = llm if llm else self.llm
 
         triggering_flow_id = flow_id
+        if not triggering_flow_id:
+            raise Exception(
+                f"No flow_id provided to generate flow."
+            )  # TODO! Should flow_id be mandatory?
 
         flow_config = state.flow_configs[triggering_flow_id]
+        if not flow_config.source_code:
+            raise Exception(f"No source_code in flow_config {flow_config}")
         docstrings = re.findall(r'"""(.*?)"""', flow_config.source_code, re.DOTALL)
 
         if len(docstrings) > 0:
@@ -897,6 +928,10 @@ async def generate_flow(
         for flow_config in state.flow_configs.values():
             if flow_config.decorators.get("meta", {}).get("tool") is True:
                 # We get rid of the first line, which is the decorator
+
+                if not flow_config.source_code:
+                    raise Exception(f"No source_code in flow_config {flow_config}")
+
                 body = flow_config.source_code.split("\n", maxsplit=1)[1]
 
                 # We only need the part up to the docstring
@@ -936,8 +971,8 @@ async def generate_flow(
             Task.GENERATE_FLOW_CONTINUATION_FROM_NLD
         )
 
-        with llm_params(llm, temperature=self.config.lowest_temperature):
-            result = await llm_call(llm, prompt, stop)
+        with llm_params(generation_llm, temperature=self.config.lowest_temperature):
+            result = await llm_call(generation_llm, prompt, stop=stop)
 
         # Parse the output using the associated parser
         result = self.llm_task_manager.parse_task_output(

nemoguardrails/actions/validation/base.py

@@ -14,15 +14,15 @@
 # limitations under the License.
 import json
 import re
-from typing import List
+from typing import List, Sequence
 from urllib.parse import quote
 
 from nemoguardrails.actions.validation.filter_secrets import contains_secrets
 
 MAX_LEN = 50
 
 
-def validate_input(attribute: str, validators: List[str] = (), **validation_args):
+def validate_input(attribute: str, validators: Sequence[str] = (), **validation_args):
     """A generic decorator that can be used by any action (class method or function) for input validation.
 
     Supported validation choices are: length and quote.

nemoguardrails/actions/validation/filter_secrets.py

@@ -22,7 +22,7 @@ def contains_secrets(resp):
     ArtifactoryDetector    : False
     """
     try:
-        import detect_secrets
+        import detect_secrets  # pyright: ignore (Assume user installs detect_secrets with instructions below)
     except ModuleNotFoundError:
         raise ValueError(
             "Could not import detect_secrets. Please install using `pip install detect-secrets`"