Fixed agent output parsing and async execution

app.py

@@ -13,7 +13,9 @@ import tempfile
 import shutil
 import gradio as gr
 import networkx as nx
-from langchain.prompts import PromptTemplate
+from langchain.prompts import PromptTemplate, MessagesPlaceholder
+from langchain.memory import ConversationBufferMemory
+from langchain.agents import Tool, initialize_agent, AgentType
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain.llms import HuggingFacePipeline
@@ -70,27 +72,42 @@ def get_agent(agent_type):
     """Get or initialize an agent with the specified type."""
     try:
         llm = get_llm()
+        
+        tools = [
+            Tool(
+                name="Code Formatter",
+                func=lambda x: subprocess.run(["black", "-"], input=x.encode(), capture_output=True).stdout.decode(),
+                description="Formats code using Black.",
+            ),
+            Tool(
+                name="API Generator",
+                func=lambda x: json.dumps({"endpoints": {"example": "POST - Example endpoint."}}),
+                description="Generates API details from code.",
+            ),
+            Tool(
+                name="Task Decomposer",
+                func=lambda x: json.dumps({"tasks": ["Design UI", "Develop Backend", "Test App", "Deploy App"]}),
+                description="Breaks down app requirements into smaller tasks.",
+            ),
+        ]
+        
+        memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        )
+        
+        agent_kwargs = {
+            "extra_prompt_messages": [MessagesPlaceholder(variable_name="chat_history")],
+        }
+        
         return initialize_agent(
-            tools=[
-                Tool(
-                    name="Code Formatter",
-                    func=lambda x: subprocess.run(["black", "-"], input=x.encode(), capture_output=True).stdout.decode(),
-                    description="Formats code using Black.",
-                ),
-                Tool(
-                    name="API Generator",
-                    func=lambda x: json.dumps({"endpoints": {"example": "POST - Example endpoint."}}),
-                    description="Generates API details from code.",
-                ),
-                Tool(
-                    name="Task Decomposer",
-                    func=lambda x: json.dumps({"tasks": ["Design UI", "Develop Backend", "Test App", "Deploy App"]}),
-                    description="Breaks down app requirements into smaller tasks.",
-                ),
-            ],
+            tools=tools,
             llm=llm,
-            agent="zero-shot-react-description",
+            agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
             verbose=True,
+            memory=memory,
+            agent_kwargs=agent_kwargs,
+            handle_parsing_errors=True,
         )
     except Exception as e:
         logger.error(f"Failed to get agent: {str(e)}")
@@ -468,7 +485,7 @@ class AIFlow:
                 raise ValueError(f"Unknown agent role: {role}")
             
             # Execute the agent's task
-            result = agent.run(prompt)
+            result = await asyncio.to_thread(agent.run, prompt)
             
             # Log the execution
             logger.info(f"Agent {role.value} completed task")
@@ -1017,7 +1034,7 @@ For any issues or questions, please refer to the documentation or create an issu
             agent = get_agent(role)
             
             # Execute the agent's task
-            result = await agent.arun(prompt)
+            result = await asyncio.to_thread(agent.run, prompt)
             
             # Process and return the result
             return {"output": result}