Add main.py

main.py

@@ -1,8 +1,11 @@
-import os
-import asyncio
+import asyncio, os
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
 from langchain.tools import tool
+from deepagents import create_deep_agent
+from deepagents.backends import CompositeBackend, LocalShellBackend
+from pydantic import BaseModel, Field
+from langchain_core.output_parsers import PydanticOutputParser
 
 # LLM configuration – always OpenRouter
 llm = ChatOpenAI(
@@ -12,28 +15,67 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# Tool that performs a simple comparison of three entities
+# Backend – local shell in virtual mode
+backend = CompositeBackend(
+    default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
+    routes={},
+)
+
+# Pydantic model for structured output
+class ComparisonResult(BaseModel):
+    entity: str = Field(description="Name of the entity")
+    strengths: str = Field(description="Key strengths")
+    weaknesses: str = Field(description="Key weaknesses")
+    overall_score: int = Field(description="Overall score out of 10")
+
+parser = PydanticOutputParser(pydantic_object=ComparisonResult)
+
+# Tool that returns a mock comparison for three entities
 @tool
-def compare_entities(entity1: str, entity2: str, entity3: str) -> str:
-    """Return a concise comparison of three entities."""
-    comparison = (
-        f"Comparison of {entity1}, {entity2}, and {entity3}:\n"
-        f"1. {entity1}: Feature A, Feature B, Feature C.\n"
-        f"2. {entity2}: Feature D, Feature E, Feature F.\n"
-        f"3. {entity3}: Feature G, Feature H, Feature I.\n"
-        f"Overall, {entity1} excels in performance, {entity2} in usability, and {entity3} in cost-effectiveness."
-    )
-    return comparison
+def compare_entities(entities: str) -> str:
+    """Return a JSON array with comparison of three entities.
+    Expected input: comma separated list of three entity names.
+    """
+    names = [e.strip() for e in entities.split(',')]
+    if len(names) != 3:
+        return "Error: please provide exactly three entities separated by commas."
+    results = []
+    for name in names:
+        results.append({
+            "entity": name,
+            "strengths": f"Strengths of {name}",
+            "weaknesses": f"Weaknesses of {name}",
+            "overall_score": 7,  # placeholder score
+        })
+    import json
+    return json.dumps(results)
+
+# Create the deep agent
+agent = create_deep_agent(
+    model=llm,
+    tools=[compare_entities],
+    backend=backend,
+    system_prompt="You are a research assistant. Use the compare_entities tool to compare three entities and return the result in the specified JSON format.",
+)
 
 async def main():
-    user_query = "Compare Tavily, Google, and Bing for web search capabilities."
-    parts = [p.strip() for p in user_query.replace("?", "").split(" ") if p.strip()]
-    if len(parts) >= 3:
-        e1, e2, e3 = parts[-3], parts[-2], parts[-1]
-    else:
-        e1, e2, e3 = "Entity1", "Entity2", "Entity3"
-    result = await compare_entities(e1, e2, e3)
-    print(result)
+    # Ask the agent to compare three entities
+    user_query = "Compare Tavily, Google, and Bing."
+    result = await agent.ainvoke(
+        {"messages": [HumanMessage(content=user_query)]},
+        {"configurable": {"thread_id": "session-compare-1"}},
+    )
+    # The agent will return the tool output; parse it into structured data
+    tool_output = result["messages"][-1].content
+    try:
+        import json
+        parsed = json.loads(tool_output)
+        structured = [ComparisonResult(**item) for item in parsed]
+        for item in structured:
+            print(item.json(indent=2))
+    except Exception as e:
+        print("Failed to parse output:", e)
+        print(tool_output)
 
 if __name__ == "__main__":
     asyncio.run(main())