Update crew.py

crew.py

@@ -220,12 +220,12 @@ def run_crew(question, file_path):
             raise RuntimeError(f"Processing failed: {str(e)}")
 
     @tool("Code Generation Tool")
-    def code_generation_tool(question: str, file_data: str) -> str:
-        """Given a question and file data, generate and execute code to answer the question.
+    def code_generation_tool(question: str, json_data: str) -> str:
+        """Given a question and json data, generate and execute code to answer the question.
 
            Args:
                question (str): Question to answer
-                file_path (str): The file data
+                file_path (str): The JSON data
 
            Returns:
                str: Answer to the question
@@ -237,7 +237,7 @@ def run_crew(question, file_path):
                     
             response = client.models.generate_content(
                 model=CODE_GENERATION_MODEL,
-                contents=[f"{question}\n{file_data}"],
+                contents=[f"{question}\n{json_data}"],
                 config=types.GenerateContentConfig(
                     tools=[types.Tool(code_execution=types.ToolCodeExecution)]
                 ),
@@ -392,8 +392,8 @@ def run_crew(question, file_path):
                     "- Video Analysis Agent requires a question and **.mp4, .mpeg, .mov, .avi, .x-flv, .mpg, .webm, .wmv, or .3gpp video file**.\n"
                     "- Document Analysis Agent requires a question and **.docx, .pptx, .pdf, .txt, .html, css, .js, .md, .xml, or .rtf document file**.\n"
                     "- YouTube Analysis Agent requires a question and **YouTube URL**.\n"
-                    "- Code Generation Agent requires a question and **file data**.\n"
-                    "- Code Execution Agent requires a question and **.py Python file**.\n"
+                    "- Code Generation Agent requires a question and **JSON data**.\n"
+                    "- Code Execution Agent requires a question and **Python file**.\n"
                     "Question: {question}",
         expected_output="The answer to the question."
     )
@@ -416,18 +416,21 @@ def run_crew(question, file_path):
 
     # Process
 
+    final_question = question
+    
     if file_path:
         if is_ext(file_path, ".csv") or is_ext(file_path, ".xls") or is_ext(file_path, ".xlsx") or is_ext(file_path, ".json") or is_ext(file_path, ".jsonl"):
-            file_data = read_file_json(file_path)
-            question = f"{question} File data: {file_data}."
+            json_data = read_file_json(file_path)
+            final_question = f"{question}\nJSON data:\n{json_data}."
         else:
-            question = f"{question} File path: {file_path}."
-
-    answer = crew.kickoff(inputs={"question": question})
-    final_answer = get_final_answer(FINAL_ANSWER_MODEL, question, str(answer))
+            final_question = f"{question}\nFile path:\n{file_path}."
+    
+    initial_answer = crew.kickoff(inputs={"question": final_question})
+    final_answer = get_final_answer(FINAL_ANSWER_MODEL, initial_question, str(initial_answer))
 
-    print(f"=> Question: {question}")
-    print(f"=> Answer: {answer}")
+    print(f"=> Initial question: {question}")
+    print(f"=> Final question: {final_question}")
+    print(f"=> Initial answer: {initial_answer}")
     print(f"=> Final answer: {final_answer}")
     
     return final_answer