refactor

agent.py

@@ -2,6 +2,7 @@ import functools
 import json
 import os
 import logging
+from typing import List, Dict, Any, Optional
 from groq import Groq
 import functions
 from utils import python_type, raise_error
@@ -13,25 +14,19 @@ logging.basicConfig(level=logging.DEBUG)
 client = Groq(api_key=os.environ["GROQ_API_KEY"])
 
 MODEL = "llama3-groq-70b-8192-tool-use-preview"
-all_functions = [func for func in dir(functions) if callable(
-    getattr(functions, func)) and not func.startswith("__")]
-names_to_functions = {func: functools.partial(
-    getattr(functions, func)) for func in all_functions}
-
-
-def create_message(prompt, message_type):
-    logging.debug(
-        f"Creating message with prompt: {prompt} and message type: {message_type}")
-    system_message = ""
-
-    if message_type == "reasoning_chain":
-        system_message = (
+ALL_FUNCTIONS = [func for func in dir(functions) if callable(getattr(functions, func)) and not func.startswith("__")]
+NAMES_TO_FUNCTIONS = {func: functools.partial(getattr(functions, func)) for func in ALL_FUNCTIONS}
+
+def create_message(prompt: str, message_type: str) -> List[Dict[str, str]]:
+    logging.debug(f"Creating message with prompt: {prompt} and message type: {message_type}")
+    
+    system_messages = {
+        "reasoning_chain": (
             "You are a movie search assistant bot who uses TMDB to help users "
             "find movies. Think step by step and identify the sequence of "
             "reasoning steps that will help to answer the user's query."
-        )
-    elif message_type == "function_call":
-        system_message = (
+        ),
+        "function_call": (
             "You are a movie search assistant bot that utilizes TMDB to help users find movies. "
             "Approach each query step by step, determining the sequence of function calls needed to gather the necessary information. "
             "Execute functions sequentially, using the output from one function to inform the next function call when required. "
@@ -40,180 +35,144 @@ def create_message(prompt, message_type):
             "finalize your process with a discover_movie function call that returns a list of movie IDs. "
             "Ensure that this call includes all necessary parameters to accurately filter the movies."
         )
-    else:
-        raise ValueError(
-            "Invalid message type. Expected 'reasoning_chain' or 'function_call'")
-
+    }
+    
+    if message_type not in system_messages:
+        raise ValueError("Invalid message type. Expected 'reasoning_chain' or 'function_call'")
+    
     return [
-        {
-            "role": "system",
-            "content": system_message,
-        },
-        {
-            "role": "user",
-            "content": prompt,
-        },
+        {"role": "system", "content": system_messages[message_type]},
+        {"role": "user", "content": prompt},
     ]
 
+def get_groq_response(messages: List[Dict[str, str]], tool_choice: str = "auto") -> Any:
+    logging.debug(f"Getting response with model: {MODEL}, \nmessages: {json.dumps(messages, indent=2)}, \ntool_choice: {tool_choice}")
+    try:
+        response = client.chat.completions.create(
+            model=MODEL,
+            messages=messages,
+            tools=tools,
+            tool_choice=tool_choice,
+            temperature=0,
+            max_tokens=4096,
+        )
+        logging.debug(f"Response: {response}")
+        return response
+    except Exception as e:
+        logging.error(f"Error getting response from Groq: {str(e)}")
+        raise
 
-def get_response(client, model, messages, tool_choice="auto"):
-    logging.debug(
-        f"Getting response with model: {model}, \nmessages: {json.dumps(messages, indent=2)}, \ntool_choice: {tool_choice}")
-    response = client.chat.completions.create(
-        model=model,
-        messages=messages,
-        tools=tools,
-        tool_choice=tool_choice,
-        temperature=0,
-        max_tokens=4096,
-    )
-    logging.debug(f"Response: {response}")
-    return response
-
-
-def generate_reasoning_chain(user_prompt):
+def generate_reasoning_chain(user_prompt: str) -> Any:
     messages = create_message(user_prompt, "reasoning_chain")
     logging.debug(f"Generating reasoning chain with messages: {messages}")
-    cot_response = get_response(client, MODEL, messages, tool_choice="none")
-    logging.info(f"COT response: {cot_response.choices[0].message.content}")
-    if cot_response.choices[0].finish_reason == "stop":
-        return cot_response.choices[0]
-    else:
-        raise_error("Failed to generate reasoning chain. Got response: " +
-                    str(cot_response), Exception)
-
-
-def validate_params(tool_params, param_name, param_value):
-    """
-    Checks if the parameter value matches with the one defined in tools.py
-    """
-    logging.debug(
-        f"Validating parameter: {param_name} with value: {param_value}")
-    param_def = tool_params.get(param_name, None)
+    try:
+        cot_response = get_groq_response(messages, tool_choice="none")
+        logging.info(f"COT response: {cot_response.choices[0].message.content}")
+        if cot_response.choices[0].finish_reason == "stop":
+            return cot_response.choices[0]
+        else:
+            raise_error("Failed to generate reasoning chain. Got response: " + str(cot_response), Exception)
+    except Exception as e:
+        logging.error(f"Error generating reasoning chain: {str(e)}")
+        raise
+
+def validate_parameter(param_name: str, param_value: Any, tool_params: Dict[str, Any]) -> bool:
+    logging.debug(f"Validating parameter: {param_name} with value: {param_value}")
+    param_def = tool_params.get(param_name)
     if param_def is None:
-        logging.error(
-            f"Parameter {param_name} not found in tools. Dropping this tool call.")
+        logging.error(f"Parameter {param_name} not found in tools. Dropping this tool call.")
         return False
     try:
-        param_value = python_type(param_def["type"])(param_value)
+        python_type(param_def["type"])(param_value)
+        return True
     except ValueError:
-        logging.error(
-            f"Parameter {param_name} value cannot be cast to {param_def['type']}. Dropping this tool call.")
+        logging.error(f"Parameter {param_name} value cannot be cast to {param_def['type']}. Dropping this tool call.")
         return False
-    return True
-
 
-def is_tool_valid(tool_name):
-    """Check if the tool name is valid and return its definition."""
+def is_tool_valid(tool_name: str) -> Optional[Dict[str, Any]]:
     return next((tool for tool in tools if tool["function"]["name"] == tool_name), None)
 
-
-def validate_tool_parameters(tool_def, tool_args):
-    """Validate the parameters of the tool against its definition."""
+def validate_tool_parameters(tool_def: Dict[str, Any], tool_args: Dict[str, Any]) -> bool:
     tool_params = tool_def.get("function", {}).get("parameters", {}).get("properties", {})
     if not tool_params:
-        # No parameters defined for this tool, so all arguments are valid
         return True
-    for param_name, param_value in tool_args.items():
-        if not validate_params(tool_params, param_name, param_value):
-            logging.error(
-                f"Invalid parameter {param_name} for tool {tool_def['function']['name']}. Dropping this tool call.")
-            return False
-    return True
-
-
-def are_arguments_valid(tool_args, user_query_values, previous_values):
-    """Check if all argument values are valid."""
-    arg_values = tool_args.values()
-    return all(str(value) in user_query_values or value in previous_values for value in arg_values)
-
-
-def verify_tool_calls(tool_calls):
-    """
-    Verify tool calls based on user query and previous tool outputs.
-
-    :param tool_calls: List of tool calls with arguments.
-    :return: List of valid tool calls.
-    """
+    return all(validate_parameter(param_name, param_value, tool_params) for param_name, param_value in tool_args.items())
 
+def verify_tool_calls(tool_calls: List[Any]) -> List[Any]:
     valid_tool_calls = []
-
     for tool_call in tool_calls:
         tool_name = tool_call.function.name
         tool_args = json.loads(tool_call.function.arguments)
-
         tool_def = is_tool_valid(tool_name)
-        if tool_def:
-            if validate_tool_parameters(tool_def, tool_args):
-                valid_tool_calls.append(tool_call)
+        if tool_def and validate_tool_parameters(tool_def, tool_args):
+            valid_tool_calls.append(tool_call)
         else:
-            logging.error(
-                f"Tool {tool_name} not found in tools. Dropping this tool call.")
-
-    tool_calls_str = json.dumps(
-        [tool_call.__dict__ for tool_call in valid_tool_calls], default=str, indent=2)
-    logging.info(
-        'Tool calls validated successfully. Valid tool calls are: %s', tool_calls_str)
+            logging.error(f"Invalid tool call: {tool_name}. Dropping this tool call.")
+    
+    tool_calls_str = json.dumps([tool_call.__dict__ for tool_call in valid_tool_calls], default=str, indent=2)
+    logging.info('Tool calls validated successfully. Valid tool calls are: %s', tool_calls_str)
     return valid_tool_calls
 
+def execute_tool(tool_call: Any) -> Dict[str, Any]:
+    logging.info(f"Executing tool: \n Name: {tool_call.function.name}\n Parameters: {tool_call.function.arguments}")
+    function_to_call = NAMES_TO_FUNCTIONS[tool_call.function.name]
+    function_args = json.loads(tool_call.function.arguments)
+    return function_to_call(**function_args)
 
-def gather_movie_data(messages):
+def gather_movie_data(messages: List[Dict[str, str]], max_tool_calls: int = 3) -> Optional[List[Dict[str, Any]]]:
     logging.debug(f"Gathering movie data with messages: {messages}")
-    response = get_response(client, MODEL, messages, tool_choice="required")
-    logging.debug(f"Calling tools based on the response: {response}")
-    if response.choices[0].finish_reason == "tool_calls":
+    try:
+        response = get_groq_response(messages, tool_choice="required")
+        logging.debug(f"Calling tools based on the response: {response}")
+        
+        if response.choices[0].finish_reason != "tool_calls":
+            raise Exception("Failed to gather movie data. Got response: " + str(response))
+        
         tool_calls = response.choices[0].message.tool_calls
-        # validate tool calls
         valid_tool_calls = verify_tool_calls(tool_calls)
-        # valid_tool_calls = tool_calls
-        updated_messages = messages.copy()
-        tool_messages_count = len(
-            [msg for msg in messages if msg["role"] == "tool"])
-        if tool_messages_count <= 3 and valid_tool_calls:
-            tool_call = valid_tool_calls[0]  # Run one tool call at a time
-            logging.debug(
-                f"Tool call: {tool_call.function.name}, Tool call parameters: {tool_call.function.arguments}")
-            tool_output = execute_tool(tool_call)
-            logging.debug(
-                f"Tool call output: {json.dumps(tool_output, indent=2)}")
-            if tool_call.function.name == "discover_movie" or tool_messages_count > 3:
-                return tool_output["results"]  # A list of movies
-            else:
-                updated_messages.append(
-                    {
-                        "tool_call_id": tool_call.id,
-                        "role": "tool",
-                        "name": tool_call.function.name,
-                        "content": str(tool_output),
-                    }
-                )
-            return gather_movie_data(updated_messages)
-        else:
-            return "No results found"
-    else:
-        raise Exception(
-            "Failed to gather movie data. Got response: ", response)
-
-
-def execute_tool(tool_call):
-    logging.info(
-        f"Executing tool: \n Name: {tool_call.function.name}\n Parameters: {tool_call.function.arguments}")
-
-    function_to_call = names_to_functions[tool_call.function.name]
-    function_args = json.loads(tool_call.function.arguments)
-    return function_to_call(**function_args)
-
-
-def chatbot(user_prompt):
-    cot_response_choice = generate_reasoning_chain(user_prompt)
-    cot = create_message(user_prompt, "function_call")
-    cot.append({
-        'role': cot_response_choice.message.role,
-        'content': cot_response_choice.message.content})
-    movie_list = gather_movie_data(cot)
-    return movie_list
-
+        
+        tool_messages_count = len([msg for msg in messages if msg["role"] == "tool"])
+        
+        if tool_messages_count >= max_tool_calls or not valid_tool_calls:
+            return None  # No results found or max tool calls reached
+        
+        tool_call = valid_tool_calls[0]  # Run one tool call at a time
+        logging.debug(f"Tool call: {tool_call.function.name}, Tool call parameters: {tool_call.function.arguments}")
+        
+        tool_output = execute_tool(tool_call)
+        logging.debug(f"Tool call output: {json.dumps(tool_output, indent=2)}")
+        
+        if tool_call.function.name == "discover_movie":
+            return tool_output["results"]  # A list of movies
+        
+        updated_messages = messages + [
+            {
+                "tool_call_id": tool_call.id,
+                "role": "tool",
+                "name": tool_call.function.name,
+                "content": str(tool_output),
+            }
+        ]
+        return gather_movie_data(updated_messages, max_tool_calls)
+    
+    except Exception as e:
+        logging.error(f"Error gathering movie data: {str(e)}")
+        return None
+
+def chatbot(user_prompt: str) -> Optional[List[Dict[str, Any]]]:
+    try:
+        cot_response_choice = generate_reasoning_chain(user_prompt)
+        cot = create_message(user_prompt, "function_call")
+        cot.append({
+            'role': cot_response_choice.message.role,
+            'content': cot_response_choice.message.content
+        })
+        movie_list = gather_movie_data(cot)
+        return movie_list
+    except Exception as e:
+        logging.error(f"Error in chatbot: {str(e)}")
+        return None
 
 if __name__ == "__main__":
-    # print(json.dumps(chatbot("Movies of Tom Hanks that were released in 1993"), indent=2))
-    print(json.dumps(chatbot("List some movies of Tom Cruise"), indent=2))
+    result = chatbot("List some movies of Tom Cruise")
+    print(json.dumps(result, indent=2) if result else "No results found")