Update

app.py

@@ -20,289 +20,289 @@ from opentools.models.memory import Memory
 from opentools.models.executor import Executor
 from opentools.models.utlis import make_json_serializable
 
-# solver = None
-
-# class ChatMessage:
-#     def __init__(self, role: str, content: str, metadata: dict = None):
-#         self.role = role
-#         self.content = content
-#         self.metadata = metadata or {}
-
-# class Solver:
-#     def __init__(
-#         self,
-#         planner,
-#         memory,
-#         executor,
-#         task: str,
-#         task_description: str,
-#         output_types: str = "base,final,direct",
-#         index: int = 0,
-#         verbose: bool = True,
-#         max_steps: int = 10,
-#         max_time: int = 60,
-#         output_json_dir: str = "results",
-#         root_cache_dir: str = "cache"
-#     ):
-#         self.planner = planner
-#         self.memory = memory
-#         self.executor = executor
-#         self.task = task
-#         self.task_description = task_description
-#         self.index = index
-#         self.verbose = verbose
-#         self.max_steps = max_steps
-#         self.max_time = max_time
-#         self.output_json_dir = output_json_dir
-#         self.root_cache_dir = root_cache_dir
-
-#         self.output_types = output_types.lower().split(',')
-#         assert all(output_type in ["base", "final", "direct"] for output_type in self.output_types), "Invalid output type. Supported types are 'base', 'final', 'direct'."
-
-#         # self.benchmark_data = self.load_benchmark_data()
-
-
-
-#     def stream_solve_user_problem(self, user_query: str, user_image: Image.Image, messages: List[ChatMessage]) -> Iterator[List[ChatMessage]]:
-#         """
-#         Streams intermediate thoughts and final responses for the problem-solving process based on user input.
+solver = None
+
+class ChatMessage:
+    def __init__(self, role: str, content: str, metadata: dict = None):
+        self.role = role
+        self.content = content
+        self.metadata = metadata or {}
+
+class Solver:
+    def __init__(
+        self,
+        planner,
+        memory,
+        executor,
+        task: str,
+        task_description: str,
+        output_types: str = "base,final,direct",
+        index: int = 0,
+        verbose: bool = True,
+        max_steps: int = 10,
+        max_time: int = 60,
+        output_json_dir: str = "results",
+        root_cache_dir: str = "cache"
+    ):
+        self.planner = planner
+        self.memory = memory
+        self.executor = executor
+        self.task = task
+        self.task_description = task_description
+        self.index = index
+        self.verbose = verbose
+        self.max_steps = max_steps
+        self.max_time = max_time
+        self.output_json_dir = output_json_dir
+        self.root_cache_dir = root_cache_dir
+
+        self.output_types = output_types.lower().split(',')
+        assert all(output_type in ["base", "final", "direct"] for output_type in self.output_types), "Invalid output type. Supported types are 'base', 'final', 'direct'."
+
+        # self.benchmark_data = self.load_benchmark_data()
+
+
+
+    def stream_solve_user_problem(self, user_query: str, user_image: Image.Image, messages: List[ChatMessage]) -> Iterator[List[ChatMessage]]:
+        """
+        Streams intermediate thoughts and final responses for the problem-solving process based on user input.
         
-#         Args:
-#             user_query (str): The text query input from the user.
-#             user_image (Image.Image): The uploaded image from the user (PIL Image object).
-#             messages (list): A list of ChatMessage objects to store the streamed responses.
-#         """
-
-#         if user_image:
-#             # # Convert PIL Image to bytes (for processing)
-#             # img_bytes_io = io.BytesIO()
-#             # user_image.save(img_bytes_io, format="PNG")  # Convert image to PNG bytes
-#             # img_bytes = img_bytes_io.getvalue()  # Get bytes
+        Args:
+            user_query (str): The text query input from the user.
+            user_image (Image.Image): The uploaded image from the user (PIL Image object).
+            messages (list): A list of ChatMessage objects to store the streamed responses.
+        """
+
+        if user_image:
+            # # Convert PIL Image to bytes (for processing)
+            # img_bytes_io = io.BytesIO()
+            # user_image.save(img_bytes_io, format="PNG")  # Convert image to PNG bytes
+            # img_bytes = img_bytes_io.getvalue()  # Get bytes
             
-#             # Use image paths instead of bytes,
-#             os.makedirs(os.path.join(self.root_cache_dir, 'images'), exist_ok=True)
-#             img_path = os.path.join(self.root_cache_dir, 'images', str(uuid.uuid4()) + '.jpg')
-#             user_image.save(img_path)
-#         else:
-#             img_path = None
-
-#         # Set query cache
-#         _cache_dir = os.path.join(self.root_cache_dir)
-#         self.executor.set_query_cache_dir(_cache_dir)
+            # Use image paths instead of bytes,
+            os.makedirs(os.path.join(self.root_cache_dir, 'images'), exist_ok=True)
+            img_path = os.path.join(self.root_cache_dir, 'images', str(uuid.uuid4()) + '.jpg')
+            user_image.save(img_path)
+        else:
+            img_path = None
+
+        # Set query cache
+        _cache_dir = os.path.join(self.root_cache_dir)
+        self.executor.set_query_cache_dir(_cache_dir)
         
-#         # Step 1: Display the received inputs
-#         if user_image:
-#             messages.append(ChatMessage(role="assistant", content=f"📝 Received Query: {user_query}\n🖼️ Image Uploaded"))
-#         else:
-#             messages.append(ChatMessage(role="assistant", content=f"📝 Received Query: {user_query}"))
-#         yield messages
-
-#         # Step 2: Add "thinking" status while processing
-#         messages.append(ChatMessage(
-#             role="assistant",
-#             content="",
-#             metadata={"title": "⏳ Thinking: Processing input..."}
-#         ))
-
-#         # Step 3: Initialize problem-solving state
-#         start_time = time.time()
-#         step_count = 0
-#         json_data = {"query": user_query, "image": "Image received as bytes"}
-
-#         # Step 4: Query Analysis
-#         import pdb; pdb.set_trace()
-#         query_analysis = self.planner.analyze_query(user_query, img_path)
-#         json_data["query_analysis"] = query_analysis
-#         messages.append(ChatMessage(role="assistant", content=f"🔍 Query Analysis:\n{query_analysis}"))
-#         yield messages
-
-#         # Step 5: Execution loop (similar to your step-by-step solver)
-#         while step_count < self.max_steps and (time.time() - start_time) < self.max_time:
-#             step_count += 1
-#             messages.append(ChatMessage(role="assistant", content=f"🔄 Step {step_count}: Generating next step..."))
-#             yield messages
-
-#             # Generate the next step
-#             next_step = self.planner.generate_next_step(
-#                 user_query, img_path, query_analysis, self.memory, step_count, self.max_steps
-#             )
-#             context, sub_goal, tool_name = self.planner.extract_context_subgoal_and_tool(next_step)
-
-#             # Display the step information
-#             messages.append(ChatMessage(
-#                 role="assistant",
-#                 content=f"📌 Step {step_count} Details:\n- Context: {context}\n- Sub-goal: {sub_goal}\n- Tool: {tool_name}"
-#             ))
-#             yield messages
-
-#             # Handle tool execution or errors
-#             if tool_name not in self.planner.available_tools:
-#                 messages.append(ChatMessage(role="assistant", content=f"⚠️ Error: Tool '{tool_name}' is not available."))
-#                 yield messages
-#                 continue
-
-#             # Execute the tool command
-#             tool_command = self.executor.generate_tool_command(
-#                 user_query, img_path, context, sub_goal, tool_name, self.planner.toolbox_metadata[tool_name]
-#             )
-#             explanation, command = self.executor.extract_explanation_and_command(tool_command)
-#             result = self.executor.execute_tool_command(tool_name, command)
-#             result = make_json_serializable(result)
-
-#             messages.append(ChatMessage(role="assistant", content=f"✅ Step {step_count} Result:\n{json.dumps(result, indent=4)}"))
-#             yield messages
-
-#             # Step 6: Memory update and stopping condition
-#             self.memory.add_action(step_count, tool_name, sub_goal, tool_command, result)
-#             stop_verification = self.planner.verificate_memory(user_query, img_path, query_analysis, self.memory)
-#             conclusion = self.planner.extract_conclusion(stop_verification)
-
-#             messages.append(ChatMessage(role="assistant", content=f"🛑 Step {step_count} Conclusion: {conclusion}"))
-#             yield messages
-
-#             if conclusion == 'STOP':
-#                 break
-
-#         # Step 7: Generate Final Output (if needed)
-#         if 'final' in self.output_types:
-#             final_output = self.planner.generate_final_output(user_query, img_path, self.memory)
-#             messages.append(ChatMessage(role="assistant", content=f"🎯 Final Output:\n{final_output}"))
-#             yield messages
-
-#         if 'direct' in self.output_types:
-#             direct_output = self.planner.generate_direct_output(user_query, img_path, self.memory)
-#             messages.append(ChatMessage(role="assistant", content=f"🔹 Direct Output:\n{direct_output}"))
-#             yield messages
-
-#         # Step 8: Completion Message
-#         messages.append(ChatMessage(role="assistant", content="✅ Problem-solving process complete."))
-#         yield messages
+        # Step 1: Display the received inputs
+        if user_image:
+            messages.append(ChatMessage(role="assistant", content=f"📝 Received Query: {user_query}\n🖼️ Image Uploaded"))
+        else:
+            messages.append(ChatMessage(role="assistant", content=f"📝 Received Query: {user_query}"))
+        yield messages
+
+        # Step 2: Add "thinking" status while processing
+        messages.append(ChatMessage(
+            role="assistant",
+            content="",
+            metadata={"title": "⏳ Thinking: Processing input..."}
+        ))
+
+        # Step 3: Initialize problem-solving state
+        start_time = time.time()
+        step_count = 0
+        json_data = {"query": user_query, "image": "Image received as bytes"}
+
+        # Step 4: Query Analysis
+        import pdb; pdb.set_trace()
+        query_analysis = self.planner.analyze_query(user_query, img_path)
+        json_data["query_analysis"] = query_analysis
+        messages.append(ChatMessage(role="assistant", content=f"🔍 Query Analysis:\n{query_analysis}"))
+        yield messages
+
+        # Step 5: Execution loop (similar to your step-by-step solver)
+        while step_count < self.max_steps and (time.time() - start_time) < self.max_time:
+            step_count += 1
+            messages.append(ChatMessage(role="assistant", content=f"🔄 Step {step_count}: Generating next step..."))
+            yield messages
+
+            # Generate the next step
+            next_step = self.planner.generate_next_step(
+                user_query, img_path, query_analysis, self.memory, step_count, self.max_steps
+            )
+            context, sub_goal, tool_name = self.planner.extract_context_subgoal_and_tool(next_step)
+
+            # Display the step information
+            messages.append(ChatMessage(
+                role="assistant",
+                content=f"📌 Step {step_count} Details:\n- Context: {context}\n- Sub-goal: {sub_goal}\n- Tool: {tool_name}"
+            ))
+            yield messages
+
+            # Handle tool execution or errors
+            if tool_name not in self.planner.available_tools:
+                messages.append(ChatMessage(role="assistant", content=f"⚠️ Error: Tool '{tool_name}' is not available."))
+                yield messages
+                continue
+
+            # Execute the tool command
+            tool_command = self.executor.generate_tool_command(
+                user_query, img_path, context, sub_goal, tool_name, self.planner.toolbox_metadata[tool_name]
+            )
+            explanation, command = self.executor.extract_explanation_and_command(tool_command)
+            result = self.executor.execute_tool_command(tool_name, command)
+            result = make_json_serializable(result)
+
+            messages.append(ChatMessage(role="assistant", content=f"✅ Step {step_count} Result:\n{json.dumps(result, indent=4)}"))
+            yield messages
+
+            # Step 6: Memory update and stopping condition
+            self.memory.add_action(step_count, tool_name, sub_goal, tool_command, result)
+            stop_verification = self.planner.verificate_memory(user_query, img_path, query_analysis, self.memory)
+            conclusion = self.planner.extract_conclusion(stop_verification)
+
+            messages.append(ChatMessage(role="assistant", content=f"🛑 Step {step_count} Conclusion: {conclusion}"))
+            yield messages
+
+            if conclusion == 'STOP':
+                break
+
+        # Step 7: Generate Final Output (if needed)
+        if 'final' in self.output_types:
+            final_output = self.planner.generate_final_output(user_query, img_path, self.memory)
+            messages.append(ChatMessage(role="assistant", content=f"🎯 Final Output:\n{final_output}"))
+            yield messages
+
+        if 'direct' in self.output_types:
+            direct_output = self.planner.generate_direct_output(user_query, img_path, self.memory)
+            messages.append(ChatMessage(role="assistant", content=f"🔹 Direct Output:\n{direct_output}"))
+            yield messages
+
+        # Step 8: Completion Message
+        messages.append(ChatMessage(role="assistant", content="✅ Problem-solving process complete."))
+        yield messages
             
-# def parse_arguments():
-#     parser = argparse.ArgumentParser(description="Run the OpenTools demo with specified parameters.")
-#     parser.add_argument("--llm_engine_name", default="gpt-4o", help="LLM engine name.")
-#     parser.add_argument("--max_tokens", type=int, default=2000, help="Maximum tokens for LLM generation.")
-#     parser.add_argument("--run_baseline_only", type=bool, default=False, help="Run only the baseline (no toolbox).")
-#     parser.add_argument("--task", default="minitoolbench", help="Task to run.")
-#     parser.add_argument("--task_description", default="", help="Task description.")
-#     parser.add_argument(
-#         "--output_types",
-#         default="base,final,direct",
-#         help="Comma-separated list of required outputs (base,final,direct)"
-#     )
-#     parser.add_argument("--enabled_tools", default="Generalist_Solution_Generator_Tool", help="List of enabled tools.")
-#     parser.add_argument("--root_cache_dir", default="demo_solver_cache", help="Path to solver cache directory.")
-#     parser.add_argument("--output_json_dir", default="demo_results", help="Path to output JSON directory.")
-#     parser.add_argument("--max_steps", type=int, default=10, help="Maximum number of steps to execute.")
-#     parser.add_argument("--max_time", type=int, default=60, help="Maximum time allowed in seconds.")
-#     parser.add_argument("--verbose", type=bool, default=True, help="Enable verbose output.")
-#     return parser.parse_args()
-
-
-# def solve_problem_gradio(user_query, user_image):
-#     """
-#     Wrapper function to connect the solver to Gradio.
-#     Streams responses from `solver.stream_solve_user_problem` for real-time UI updates.
-#     """
-#     global solver  # Ensure we're using the globally defined solver
-
-#     if solver is None:
-#         return [["assistant", "⚠️ Error: Solver is not initialized. Please restart the application."]]
-
-#     messages = []  # Initialize message list
-#     for message_batch in solver.stream_solve_user_problem(user_query, user_image, messages):
-#         yield [[msg.role, msg.content] for msg in message_batch]  # Ensure correct format for Gradio Chatbot
-
-
-
-# def main(args):
-#     global solver
-#     # Initialize Tools
-#     enabled_tools = args.enabled_tools.split(",") if args.enabled_tools else []
-
-
-#     # Instantiate Initializer
-#     initializer = Initializer(
-#         enabled_tools=enabled_tools,
-#         model_string=args.llm_engine_name
-#     )
-
-#     # Instantiate Planner
-#     planner = Planner(
-#         llm_engine_name=args.llm_engine_name,
-#         toolbox_metadata=initializer.toolbox_metadata,
-#         available_tools=initializer.available_tools
-#     )
-
-#     # Instantiate Memory
-#     memory = Memory()
-
-#     # Instantiate Executor
-#     executor = Executor(
-#         llm_engine_name=args.llm_engine_name,
-#         root_cache_dir=args.root_cache_dir,
-#         enable_signal=False
-#     )
-
-#     # Instantiate Solver
-#     solver = Solver(
-#         planner=planner,
-#         memory=memory,
-#         executor=executor,
-#         task=args.task,
-#         task_description=args.task_description,
-#         output_types=args.output_types,  # Add new parameter
-#         verbose=args.verbose,
-#         max_steps=args.max_steps,
-#         max_time=args.max_time,
-#         output_json_dir=args.output_json_dir,
-#         root_cache_dir=args.root_cache_dir
-#     )
-
-#     # Test Inputs
-#     # user_query = "How many balls are there in the image?"
-#     # user_image_path = "/home/sheng/toolbox-agent/mathvista_113.png"  # Replace with your actual image path
-
-#     # # Load the image as a PIL object
-#     # user_image = Image.open(user_image_path).convert("RGB")  # Ensure it's in RGB mode
-
-#     # print("\n=== Starting Problem Solving ===\n")
-#     # messages = []
-#     # for message_batch in solver.stream_solve_user_problem(user_query, user_image, messages):
-#     #     for message in message_batch:
-#     #         print(f"{message.role}: {message.content}") 
-
-#     # messages = []
-#     # solver.stream_solve_user_problem(user_query, user_image, messages)
-
-
-#     # def solve_problem_stream(user_query, user_image):
-#     #     messages = []  # Ensure it's a list of [role, content] pairs
-
-#     #     for message_batch in solver.stream_solve_user_problem(user_query, user_image, messages):
-#     #         yield message_batch  # Stream messages correctly in tuple format
-
-#     # solve_problem_stream(user_query, user_image)
-
-#     # ========== Gradio Interface ==========
-#     with gr.Blocks() as demo:
-#         gr.Markdown("# 🧠 OctoTools AI Solver")  # Title
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Run the OpenTools demo with specified parameters.")
+    parser.add_argument("--llm_engine_name", default="gpt-4o", help="LLM engine name.")
+    parser.add_argument("--max_tokens", type=int, default=2000, help="Maximum tokens for LLM generation.")
+    parser.add_argument("--run_baseline_only", type=bool, default=False, help="Run only the baseline (no toolbox).")
+    parser.add_argument("--task", default="minitoolbench", help="Task to run.")
+    parser.add_argument("--task_description", default="", help="Task description.")
+    parser.add_argument(
+        "--output_types",
+        default="base,final,direct",
+        help="Comma-separated list of required outputs (base,final,direct)"
+    )
+    parser.add_argument("--enabled_tools", default="Generalist_Solution_Generator_Tool", help="List of enabled tools.")
+    parser.add_argument("--root_cache_dir", default="demo_solver_cache", help="Path to solver cache directory.")
+    parser.add_argument("--output_json_dir", default="demo_results", help="Path to output JSON directory.")
+    parser.add_argument("--max_steps", type=int, default=10, help="Maximum number of steps to execute.")
+    parser.add_argument("--max_time", type=int, default=60, help="Maximum time allowed in seconds.")
+    parser.add_argument("--verbose", type=bool, default=True, help="Enable verbose output.")
+    return parser.parse_args()
+
+
+def solve_problem_gradio(user_query, user_image):
+    """
+    Wrapper function to connect the solver to Gradio.
+    Streams responses from `solver.stream_solve_user_problem` for real-time UI updates.
+    """
+    global solver  # Ensure we're using the globally defined solver
+
+    if solver is None:
+        return [["assistant", "⚠️ Error: Solver is not initialized. Please restart the application."]]
+
+    messages = []  # Initialize message list
+    for message_batch in solver.stream_solve_user_problem(user_query, user_image, messages):
+        yield [[msg.role, msg.content] for msg in message_batch]  # Ensure correct format for Gradio Chatbot
+
+
+
+def main(args):
+    global solver
+    # Initialize Tools
+    enabled_tools = args.enabled_tools.split(",") if args.enabled_tools else []
+
+
+    # Instantiate Initializer
+    initializer = Initializer(
+        enabled_tools=enabled_tools,
+        model_string=args.llm_engine_name
+    )
+
+    # Instantiate Planner
+    planner = Planner(
+        llm_engine_name=args.llm_engine_name,
+        toolbox_metadata=initializer.toolbox_metadata,
+        available_tools=initializer.available_tools
+    )
+
+    # Instantiate Memory
+    memory = Memory()
+
+    # Instantiate Executor
+    executor = Executor(
+        llm_engine_name=args.llm_engine_name,
+        root_cache_dir=args.root_cache_dir,
+        enable_signal=False
+    )
+
+    # Instantiate Solver
+    solver = Solver(
+        planner=planner,
+        memory=memory,
+        executor=executor,
+        task=args.task,
+        task_description=args.task_description,
+        output_types=args.output_types,  # Add new parameter
+        verbose=args.verbose,
+        max_steps=args.max_steps,
+        max_time=args.max_time,
+        output_json_dir=args.output_json_dir,
+        root_cache_dir=args.root_cache_dir
+    )
+
+    # Test Inputs
+    # user_query = "How many balls are there in the image?"
+    # user_image_path = "/home/sheng/toolbox-agent/mathvista_113.png"  # Replace with your actual image path
+
+    # # Load the image as a PIL object
+    # user_image = Image.open(user_image_path).convert("RGB")  # Ensure it's in RGB mode
+
+    # print("\n=== Starting Problem Solving ===\n")
+    # messages = []
+    # for message_batch in solver.stream_solve_user_problem(user_query, user_image, messages):
+    #     for message in message_batch:
+    #         print(f"{message.role}: {message.content}") 
+
+    # messages = []
+    # solver.stream_solve_user_problem(user_query, user_image, messages)
+
+
+    # def solve_problem_stream(user_query, user_image):
+    #     messages = []  # Ensure it's a list of [role, content] pairs
+
+    #     for message_batch in solver.stream_solve_user_problem(user_query, user_image, messages):
+    #         yield message_batch  # Stream messages correctly in tuple format
+
+    # solve_problem_stream(user_query, user_image)
+
+    # ========== Gradio Interface ==========
+    with gr.Blocks() as demo:
+        gr.Markdown("# 🧠 OctoTools AI Solver")  # Title
 
-#         with gr.Row():
-#             user_query = gr.Textbox(label="Enter your query", placeholder="Type your question here...")
-#             user_image = gr.Image(type="pil", label="Upload an image")  # Accepts multiple formats
+        with gr.Row():
+            user_query = gr.Textbox(label="Enter your query", placeholder="Type your question here...")
+            user_image = gr.Image(type="pil", label="Upload an image")  # Accepts multiple formats
 
-#         run_button = gr.Button("Run")  # Run button
-#         chatbot_output = gr.Chatbot(label="Problem-Solving Output")
+        run_button = gr.Button("Run")  # Run button
+        chatbot_output = gr.Chatbot(label="Problem-Solving Output")
 
-#         # Link button click to function
-#         run_button.click(fn=solve_problem_gradio, inputs=[user_query, user_image], outputs=chatbot_output)
+        # Link button click to function
+        run_button.click(fn=solve_problem_gradio, inputs=[user_query, user_image], outputs=chatbot_output)
 
-#     # Launch the Gradio app
-#     demo.launch()
+    # Launch the Gradio app
+    demo.launch()
 
 
 
-# if __name__ == "__main__":
-#     args = parse_arguments()
-#     main(args)
+if __name__ == "__main__":
+    args = parse_arguments()
+    main(args)