Update codebase with appropriate coding standards

nemoguardrails/actions/action_dispatcher.py

@@ -32,12 +32,13 @@ class ActionDispatcher:
     def __init__(
         self, load_all_actions: bool = True, config_path: Optional[str] = None
     ):
-        """Initializes an actions dispatcher.
-
-        :param load_all_actions: When set, it will load all actions in the `actions` folder
-          both in the current working directory and in the package.
-        :param config_path: The path from which the configuration was loaded. If there are
-          actions at the specified path, we load them as well.
+        """
+        Initializes an actions dispatcher.
+        Args:
+            load_all_actions (bool, optional): When set to True, it loads all actions in the
+                'actions' folder both in the current working directory and in the package.
+            config_path (str, optional): The path from which the configuration was loaded.
+                If there are actions at the specified path, it loads them as well.
         """
         log.info("Initializing action dispatcher")
 
@@ -72,13 +73,21 @@ def __init__(
 
     @property
     def registered_actions(self):
+        """
+        Gets the dictionary of registered actions.
+        Returns:
+            dict: A dictionary where keys are action names and values are callable action functions.
+        """
         return self._registered_actions
 
     def load_actions_from_path(self, path: str):
         """Loads all actions from the specified path.
 
-        It will load all actions in the `actions.py` file if it exists and all actions
-        inside the `actions` folder if it exists.
+        This method loads all actions from the `actions.py` file if it exists and
+        all actions inside the `actions` folder if it exists.
+
+        Args:
+            path (str): A string representing the path from which to load actions.
         """
         actions_path = os.path.join(path, "actions")
         if os.path.exists(actions_path):
@@ -95,9 +104,10 @@ def register_action(
     ):
         """Registers an action with the given name.
 
-        :param name: The name of the action.
-        :param action: The action function.
-        :param override: If an action already exists, whether it should be overriden or not.
+        Args:
+            action (callable): The action function.
+            name (Optional[str]): The name of the action. Defaults to None.
+            override (bool): If an action already exists, whether it should be overridden or not.
         """
         if name is None:
             action_meta = getattr(action, "action_meta", None)
@@ -110,7 +120,13 @@ def register_action(
         self._registered_actions[name] = action
 
     def register_actions(self, actions_obj: any, override: bool = True):
-        """Registers all the actions from the given object."""
+        """Registers all the actions from the given object.
+
+        Args:
+            actions_obj (any): The object containing actions.
+            override (bool): If an action already exists, whether it should be overridden or not.
+        """
+
         # Register the actions
         for attr in dir(actions_obj):
             val = getattr(actions_obj, attr)
@@ -119,13 +135,27 @@ def register_actions(self, actions_obj: any, override: bool = True):
                 self.register_action(val, override=override)
 
     def get_action(self, name: str) -> callable:
+        """Get the registered action by name.
+
+        Args:
+            name (str): The name of the action.
+
+        Returns:
+            callable: The registered action.
+        """
         return self._registered_actions.get(name)
 
     async def execute_action(
         self, action_name: str, params: Dict[str, Any]
     ) -> Tuple[Union[str, Dict[str, Any]], str]:
-        """Endpoint called from action server to execute an action.
-        This endpoint interacts with different supported actions
+        """Execute a registered action.
+
+        Args:
+            action_name (str): The name of the action to execute.
+            params (Dict[str, Any]): Parameters for the action.
+
+        Returns:
+            Tuple[Union[str, Dict[str, Any]], str]: A tuple containing the result and status.
         """
 
         if action_name in self._registered_actions:
@@ -176,12 +206,23 @@ async def execute_action(
         return None, "failed"
 
     def get_registered_actions(self) -> List[str]:
-        """Endpoint called from action server to get the list of available actions"""
+        """Get the list of available actions.
+
+        Returns:
+            List[str]: List of available actions.
+        """
         return list(self._registered_actions.keys())
 
     @staticmethod
     def _load_actions_from_module(filepath: str):
-        """Loads the actions from the specified python module."""
+        """Loads the actions from the specified python module.
+
+        Args:
+            filepath (str): The path of the Python module.
+
+        Returns:
+            Dict: Dictionary of loaded actions.
+        """
         action_objects = {}
         filename = os.path.basename(filepath)
 
@@ -218,7 +259,13 @@ def _load_actions_from_module(filepath: str):
     @staticmethod
     def _find_actions(directory) -> Dict:
         """Loop through all the subdirectories and check for the class with @action
-        decorator and add in action_classes dict
+        decorator and add in action_classes dict.
+
+        Args:
+            directory: The directory to search for actions.
+
+        Returns:
+            Dict: Dictionary of found actions.
         """
         action_objects = {}
 

nemoguardrails/actions/actions.py

@@ -19,7 +19,22 @@
 
 # A decorator that sets a property on the function to indicate if it's a system action or not.
 def action(is_system_action: bool = False, name: Optional[str] = None):
+    """Decorator to mark a function or class as an action.
+
+    Args:
+        is_system_action (bool): Flag indicating if the action is a system action.
+        name (Optional[str]): The name to associate with the action.
+
+    Returns:
+        callable: The decorated function or class.
+    """
+
     def decorator(fn_or_cls):
+        """Inner decorator function to add metadata to the action.
+
+        Args:
+            fn_or_cls: The function or class being decorated.
+        """
         fn_or_cls.action_meta = {
             "name": name or fn_or_cls.__name__,
             "is_system_action": is_system_action,
@@ -31,6 +46,13 @@ def decorator(fn_or_cls):
 
 @dataclass
 class ActionResult:
+    """Data class representing the result of an action.
+
+    Attributes:
+        return_value (Optional[Any]): The value returned by the action.
+        events (Optional[List[dict]]): The events to be added to the stream.
+        context_updates (Optional[dict]): Updates made to the context by this action.
+    """
     # The value returned by the action
     return_value: Optional[Any] = None
 

nemoguardrails/actions/core.py

@@ -27,7 +27,15 @@ async def create_event(
     event: dict,
     context: Optional[dict] = None,
 ):
-    """Checks the facts for the bot response."""
+    """Creates an event for the bot based on the provided data.
+
+    Args:
+        event (dict): The input event data.
+        context (Optional[dict]): The context for the action. Defaults to None.
+
+    Returns:
+        ActionResult: An action result containing the created event.
+    """
 
     event_dict = new_event_dict(
         event["_type"], **{k: v for k, v in event.items() if k != "_type"}

nemoguardrails/actions/math.py

@@ -34,10 +34,17 @@
 async def wolfram_alpha_request(
     query: Optional[str] = None, context: Optional[dict] = None
 ):
-    """Makes a request to the Wolfram Alpha API
+    """Makes a request to the Wolfram Alpha API.
 
-    :param context: The context for the execution of the action.
-    :param query: The query for Wolfram.
+    Args:
+        query (Optional[str]): The query for Wolfram Alpha. Defaults to None.
+        context (Optional[dict]): The context for the execution of the action. Defaults to None.
+
+    Returns:
+        ActionResult or str: The result of the Wolfram Alpha request.
+
+    Raises:
+        Exception: If no query is provided to Wolfram Alpha.
     """
     # If we don't have an explicit query, we take the last user message
     if query is None and context is not None:

nemoguardrails/actions/retrieve_relevant_chunks.py

@@ -27,7 +27,26 @@ async def retrieve_relevant_chunks(
     context: Optional[dict] = None,
     kb: Optional[KnowledgeBase] = None,
 ):
-    """Retrieve relevant chunks from the knowledge base and add them to the context."""
+    """Retrieve relevant knowledge chunks and update the context.
+
+    Args:
+        context (Optional[dict]): The context for the execution of the action. Defaults to None.
+        kb (Optional[KnowledgeBase]): The KnowledgeBase to search for relevant chunks. Defaults to None.
+
+    Returns:
+        ActionResult: An action result containing the retrieved relevant chunks.
+
+    Note:
+        This action retrieves relevant chunks from the KnowledgeBase based on the user's last message
+        and updates the context with the information.
+
+    Example:
+        ```
+        result = await retrieve_relevant_chunks(context=my_context, kb=my_knowledge_base)
+        print(result.return_value)  # Relevant chunks as a string
+        print(result.context_updates)  # Updated context with relevant chunks
+        ```
+    """
     user_message = context.get("last_user_message")
     context_updates = {}
 

nemoguardrails/actions/summarize_document.py

@@ -22,9 +22,20 @@
 
 @action(name="summarize_document")
 class SummarizeDocument:
-    """Sample implementation of a document summarization action.
+    """Action for summarizing a document.
 
-    The implementation uses the summarization chain from LangChain.
+    This class provides a sample implementation of document summarization using LangChain's summarization chain.
+
+    Args:
+        document_path (str): The path to the document to be summarized.
+        llm (BaseLLM): The Language Model for the summarization process.
+
+    Example:
+        ```python
+        summarizer = SummarizeDocument(document_path="path/to/document.txt", llm=my_language_model)
+        result = summarizer.run()
+        print(result)  # The summarized document
+        ```
     """
 
     def __init__(self, document_path: str, llm: BaseLLM):

nemoguardrails/actions_server/actions_server.py

@@ -38,13 +38,15 @@
 
 
 class RequestBody(BaseModel):
+    """Request body for executing an action."""
     action_name: str = ""
     action_parameters: Dict = Field(
         default={}, description="The list of action parameters."
     )
 
 
 class ResponseBody(BaseModel):
+    """Response body for action execution."""
     status: str = "success"  # success / failed
     result: Optional[str]
 
@@ -55,7 +57,14 @@ class ResponseBody(BaseModel):
     response_model=ResponseBody,
 )
 async def run_action(body: RequestBody):
-    """Execute action_name with action_parameters and return result."""
+    """Execute the specified action and return the result.
+
+    Args:
+        body (RequestBody): The request body containing action_name and action_parameters.
+
+    Returns:
+        dict: The response containing the execution status and result.
+    """
 
     log.info(f"Request body: {body}")
     result, status = await app.action_dispatcher.execute_action(

nemoguardrails/cli/chat.py

@@ -26,6 +26,14 @@
 
 
 async def input_async(prompt_message: str = "") -> str:
+    """Asynchronously read user input with a prompt.
+
+    Args:
+        prompt_message (str): The message to display as a prompt. Defaults to an empty string.
+
+    Returns:
+        str: The user's input.
+    """
     loop = asyncio.get_event_loop()
     return await loop.run_in_executor(None, input, prompt_message)
 
@@ -37,6 +45,15 @@ async def run_chat_async(
     server_url: Optional[str] = None,
     config_id: Optional[str] = None,
 ):
+    """Asynchronously run a chat session in the terminal.
+
+    Args:
+        config_path (Optional[str]): The path to the configuration file. Defaults to None.
+        verbose (bool): Whether to run in verbose mode. Defaults to False.
+        streaming (bool): Whether to enable streaming mode. Defaults to False.
+        server_url (Optional[str]): The URL of the chat server. Defaults to None.
+        config_id (Optional[str]): The configuration ID. Defaults to None.
+    """
     if config_path is None and server_url is None:
         raise RuntimeError(
             "At least one of `config_path` or `server-url` must be provided."
@@ -114,7 +131,15 @@ def run_chat(
     server_url: Optional[str] = None,
     config_id: Optional[str] = None,
 ):
-    """Runs a chat session in the terminal."""
+    """Run a chat session in the terminal.
+
+    Args:
+        config_path (Optional[str]): The path to the configuration file. Defaults to None.
+        verbose (bool): Whether to run in verbose mode. Defaults to False.
+        streaming (bool): Whether to enable streaming mode. Defaults to False.
+        server_url (Optional[str]): The URL of the chat server. Defaults to None.
+        config_id (Optional[str]): The configuration ID. Defaults to None.
+    """
     asyncio.run(
         run_chat_async(
             config_path=config_path,