different system message for reasoning vs tool use chains

agent.py

@@ -18,12 +18,36 @@ names_to_functions = {func: functools.partial(
     getattr(functions, func)) for func in all_functions}
 
 
-def create_message(prompt):
-    logging.debug(f"Creating message with prompt: {prompt}")
+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 = (
+            "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 = (
+            "You are a movie search assistant bot who uses TMDB to help users "
+            "find movies. Think step by step and identify the sequence of "
+            "function calls that will help to answer the user's query. Use the "
+            "available functions to gather the necessary data. "
+            "Do not call multiple functions when they need to be executed in sequence. "
+            "Only call multiple functions when they can be executed in parallel. "
+            "Stop with a discover_movie function call that returns a list of movie ids. "
+            "Ensure the discover_movie function call includes all the necessary parameters to filter the movies accurately."
+        )
+    else:
+        raise ValueError(
+            "Invalid message type. Expected 'reasoning_chain' or 'function_call'")
+
     return [
         {
             "role": "system",
-            "content": "You are a movie search assistant bot who uses TMDB to help users find movies. Think step by step and identify the sequence of function calls that will help to answer.",
+            "content": system_message,
         },
         {
             "role": "user",
@@ -35,7 +59,7 @@ def create_message(prompt):
 def get_response(client, model, messages, tool_choice="auto"):
     logging.debug(
         f"Getting response with model: {model}, messages: {messages}, tool_choice: {tool_choice}")
-    return client.chat.completions.create(
+    response = client.chat.completions.create(
         model=model,
         messages=messages,
         tools=tools,
@@ -43,31 +67,38 @@ def get_response(client, model, messages, tool_choice="auto"):
         temperature=0,
         max_tokens=4096,
     )
+    logging.debug(f"Response: {response}")
+    return response
 
 
-def generate_reasoning_chain(messages):
+def generate_reasoning_chain(user_prompt):
+    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":
-        messages.append({
-            'role': cot_response.choices[0].message.role,
-            'content': cot_response.choices[0].message.content})
-        return messages
+        return cot_response.choices[0]
     else:
-        return None
+        logging.error(
+            "Failed to generate reasoning chain. Got response: ", cot_response)
+        raise Exception("Failed to generate reasoning chain")
 
 
 def gather_movie_data(messages, iteration=0, max_iterations=2):
     logging.debug(
         f"Gathering movie data with messages: {messages}, iteration: {iteration}")
     response = get_response(client, MODEL, messages, tool_choice="required")
+    logging.info(
+        f"Gathering movie data response: {response}")
     if response.choices[0].finish_reason == "tool_calls":
         tool_calls = response.choices[0].message.tool_calls
         updated_messages = messages.copy()
         for tool_call in tool_calls:
+            logging.info(
+                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: {tool_call.function.name}, output: {tool_output}")
+            logging.info(
+                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
             else:
@@ -81,6 +112,9 @@ def gather_movie_data(messages, iteration=0, max_iterations=2):
                 )
         if iteration < max_iterations:
             return gather_movie_data(updated_messages, iteration + 1)
+    else:
+        raise Exception(
+            "Failed to gather movie data. Got response: ", response)
 
 
 def execute_tool(tool_call):
@@ -91,8 +125,11 @@ def execute_tool(tool_call):
 
 
 def chatbot(user_prompt):
-    messages = create_message(user_prompt)
-    cot = generate_reasoning_chain(messages)
+    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
 

functions.py

@@ -7,7 +7,7 @@ BASE_URL = "https://api.themoviedb.org/3"
 API_KEY = os.environ['TMDB_API_KEY']
 
 # Set up logging
-logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO)
 
 
 def query_tmdb(endpoint, params={}):