add: main.py
This commit is contained in:
@@ -0,0 +1,207 @@
|
||||
import os
|
||||
import argparse
|
||||
import asyncio
|
||||
from typing import TypedDict, List, Dict, Optional
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain.tools import tool
|
||||
from deepagents import create_deep_agent
|
||||
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
|
||||
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from tavily import TavilyClient
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
|
||||
|
||||
# LLM configuration
|
||||
llm = ChatOpenAI(
|
||||
model="openai/gpt-oss-20b:free",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_key=OPENAI_API_KEY,
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
# Tavily client
|
||||
tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
|
||||
|
||||
# State definition
|
||||
class CompareState(TypedDict):
|
||||
entities: List[str]
|
||||
criteria: List[str]
|
||||
findings: Dict[str, List[str]]
|
||||
final_table: Optional[str]
|
||||
verdict: Optional[str]
|
||||
|
||||
# Pydantic models for parsing
|
||||
class CriteriaOutput(BaseModel):
|
||||
criteria: List[str] = Field(description="List of comparison criteria")
|
||||
|
||||
class VerdictOutput(BaseModel):
|
||||
verdict: str = Field(description="Verdict text")
|
||||
|
||||
criteria_parser = PydanticOutputParser(pydantic_object=CriteriaOutput)
|
||||
verdict_parser = PydanticOutputParser(pydantic_object=VerdictOutput)
|
||||
|
||||
# Node: plan_criteria
|
||||
def plan_criteria(state: CompareState) -> CompareState:
|
||||
prompt = (
|
||||
f"Given the following entities: {', '.join(state['entities'])}. "
|
||||
"Generate 3 to 5 distinct criteria for comparing these entities. "
|
||||
"Return a JSON object with a field 'criteria' that is a list of strings."
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
parsed = criteria_parser.parse(response.content)
|
||||
state["criteria"] = parsed.criteria
|
||||
# Initialize findings dict
|
||||
state["findings"] = {entity: [] for entity in state["entities"]}
|
||||
return state
|
||||
|
||||
# Node: research_entity
|
||||
def research_entity(state: CompareState) -> CompareState:
|
||||
# Find first entity with missing notes
|
||||
for entity in state["entities"]:
|
||||
if len(state["findings"][entity]) < len(state["criteria"]):
|
||||
idx = len(state["findings"][entity])
|
||||
criterion = state["criteria"][idx]
|
||||
query = f"{entity} {criterion}"
|
||||
# Perform Tavily search
|
||||
results = tavily_client.search(query)
|
||||
# Take first result snippet
|
||||
if results and results[0].snippet:
|
||||
note = results[0].snippet.strip()
|
||||
else:
|
||||
note = f"No relevant info found for {entity} on {criterion}."
|
||||
state["findings"][entity].append(note)
|
||||
break
|
||||
return state
|
||||
|
||||
# Node: check_done
|
||||
def check_done(state: CompareState) -> str:
|
||||
for entity in state["entities"]:
|
||||
if len(state["findings"][entity]) < len(state["criteria"]):
|
||||
return "continue"
|
||||
return "done"
|
||||
|
||||
# Node: build_table
|
||||
def build_table(state: CompareState) -> CompareState:
|
||||
header = ["Criterion"] + state["entities"]
|
||||
rows = []
|
||||
for criterion in state["criteria"]:
|
||||
row = [criterion]
|
||||
for entity in state["entities"]:
|
||||
notes = state["findings"][entity]
|
||||
idx = state["criteria"].index(criterion)
|
||||
note = notes[idx] if idx < len(notes) else ""
|
||||
row.append(note)
|
||||
rows.append(row)
|
||||
# Build markdown table
|
||||
table_lines = ["| " + " | ".join(header) + " |"]
|
||||
table_lines.append("|" + "|".join(["---"] * len(header)) + "|")
|
||||
for row in rows:
|
||||
table_lines.append("| " + " | ".join(row) + " |")
|
||||
table_md = "\n".join(table_lines)
|
||||
state["final_table"] = table_md
|
||||
return state
|
||||
|
||||
# Node: verdict
|
||||
def verdict(state: CompareState) -> CompareState:
|
||||
prompt = (
|
||||
f"Here is a markdown table comparing the entities:\n\n{state['final_table']}\n\n"
|
||||
"Based on this table, provide a concise verdict recommending which entity "
|
||||
"is best suited for a typical vector database use case. "
|
||||
"Return a JSON object with a field 'verdict' that is a string."
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
parsed = verdict_parser.parse(response.content)
|
||||
state["verdict"] = parsed.verdict
|
||||
return state
|
||||
|
||||
# Build LangGraph
|
||||
graph = StateGraph(CompareState)
|
||||
graph.add_node("plan_criteria", plan_criteria)
|
||||
graph.add_node("research_entity", research_entity)
|
||||
graph.add_node("check_done", check_done)
|
||||
graph.add_node("build_table", build_table)
|
||||
graph.add_node("verdict", verdict)
|
||||
|
||||
graph.add_edge(START, "plan_criteria")
|
||||
graph.add_edge("plan_criteria", "research_entity")
|
||||
graph.add_edge("research_entity", "check_done")
|
||||
graph.add_conditional_edges(
|
||||
"check_done",
|
||||
lambda x: x,
|
||||
{
|
||||
"continue": "research_entity",
|
||||
"done": "build_table",
|
||||
},
|
||||
)
|
||||
graph.add_edge("build_table", "verdict")
|
||||
graph.add_edge("verdict", END)
|
||||
|
||||
compiled_graph = graph.compile()
|
||||
|
||||
# Tool: compare_entities
|
||||
@tool
|
||||
def compare_entities(query: str) -> str:
|
||||
"""
|
||||
Compare three entities based on user query.
|
||||
Expected format: "Compare 3 vector DBs: Chroma, FAISS, Qdrant"
|
||||
"""
|
||||
# Extract entities after colon
|
||||
if ":" in query:
|
||||
parts = query.split(":", 1)
|
||||
entities_part = parts[1]
|
||||
else:
|
||||
entities_part = query
|
||||
entities = [e.strip() for e in entities_part.split(",") if e.strip()]
|
||||
if len(entities) != 3:
|
||||
return "Please provide exactly three entities separated by commas."
|
||||
initial_state: CompareState = {
|
||||
"entities": entities,
|
||||
"criteria": [],
|
||||
"findings": {},
|
||||
"final_table": None,
|
||||
"verdict": None,
|
||||
}
|
||||
final_state = compiled_graph.invoke(initial_state)
|
||||
table = final_state["final_table"] or ""
|
||||
verdict_text = final_state["verdict"] or ""
|
||||
return f"{table}\n\nVerdict:\n{verdict_text}"
|
||||
|
||||
# DeepAgent setup
|
||||
backend = CompositeBackend([LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend()])
|
||||
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[compare_entities],
|
||||
backend=backend,
|
||||
system_prompt="You are a helpful assistant that can compare three entities.",
|
||||
)
|
||||
|
||||
# CLI
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Compare three entities.")
|
||||
parser.add_argument(
|
||||
"--query",
|
||||
type=str,
|
||||
default="Compare 3 vector DBs: Chroma, FAISS, Qdrant",
|
||||
help="Comparison query in the format 'Compare 3 vector DBs: A, B, C'",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
async def run():
|
||||
result = await agent.ainvoke(
|
||||
{"messages": [{"role": "user", "content": args.query}]},
|
||||
{"configurable": {"thread_id": "session-1"}},
|
||||
)
|
||||
print(result["messages"][-1]["content"])
|
||||
asyncio.run(run())
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user