From 4d4e4e3790534192073591b3801261a62dfef34f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= <v.mayoralv@gmail.com>
Date: Fri, 10 Jan 2025 08:03:11 +0000
Subject: [PATCH] Various improvements in docs and debug prints
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
---
 cai/core.py | 51 +++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 47 insertions(+), 4 deletions(-)

diff --git a/cai/core.py b/cai/core.py
index 93cbddc4..50f7c1c6 100644
--- a/cai/core.py
+++ b/cai/core.py
@@ -36,7 +36,7 @@ __CTX_VARS_NAME__ = "context_variables"
 
 class CAI:
     """
-    Main class for the CAI library.
+    Cybersecurity AI (CAI) object
     """
 
     def __init__(self,
@@ -95,7 +95,13 @@ class CAI:
 
     def handle_function_result(self, result, debug) -> Result:
         """
-        Handle the result of a function call.
+        Handle the result of a function call by
+        converting it into a standardized Result type.
+
+        The Result type encapsulates the possible
+        return values (Result, Agent, or context variables)
+        that functions can produce into a consistent
+        format for the framework to process.
         """
         match result:
             case Result() as result:
@@ -122,7 +128,37 @@ class CAI:
         debug: bool,
     ) -> Response:
         """
-        Handle the tool calls for the given agent.
+        Execute and handle tool calls made by the AI agent.
+
+        Processes a list of tool calls by:
+        1. Looking up each function in the provided function map
+        2. Handling missing tools gracefully by skipping them
+        3. Parsing and validating function arguments
+        4. Executing functions with provided arguments and context
+        5. Processing results into standardized Response format
+        6. Accumulating results from multiple tool calls
+
+        Args:
+            tool_calls (List[ChatCompletionMessageToolCall]): Tool
+                calls requested by AI agent
+            functions (List[AgentFunction]): Available functions
+                that can be called
+            context_variables (dict): Context variables to pass
+                to functions
+            debug (bool): Flag to enable debug logging
+
+        Returns:
+            Response: Object containing:
+                messages (List): Tool call results
+                agent (Optional[Agent]): Updated agent
+                    if returned by a function
+                context_variables (dict): Updated context variables
+
+        Note:
+            Results from multiple tool calls are accumulated
+                into a single Response.
+            Context variables are updated iteratively as
+                functions are called.
         """
         function_map = {f.__name__: f for f in functions}
         partial_response = Response(
@@ -144,7 +180,11 @@ class CAI:
                 continue
             args = json.loads(tool_call.function.arguments)
             debug_print(
-                debug, f"Processing tool call: {name} with arguments {args}")
+                debug,
+                "Processing tool call",
+                name,
+                "with arguments",
+                args)
 
             func = function_map[name]
             # pass context_variables to agent functions
@@ -183,6 +223,9 @@ class CAI:
     ):
         """
         Run the cai and stream the results.
+
+        The key difference from run() is that this streams results
+        incrementally, while run() returns everything at once.
         """
         active_agent = agent
         context_variables = copy.deepcopy(context_variables)