99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
import os
|
|
from typing import Dict, List
|
|
from dotenv import load_dotenv
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_tavily import TavilySearchTool
|
|
from langgraph.prebuilt import create_chat_agent
|
|
from langgraph import add_messages, StateGraph
|
|
from src.state import CompareState
|
|
|
|
load_dotenv()
|
|
|
|
# LLM and Tavily tool
|
|
llm = ChatOpenAI(temperature=0.7)
|
|
tavily = TavilySearchTool(api_key=os.getenv("TAVILY_API_KEY"))
|
|
|
|
# Helper to format findings for LLM
|
|
def format_findings(findings: Dict[str, List[str]]) -> str:
|
|
parts = []
|
|
for entity, notes in findings.items():
|
|
parts.append(f"**{entity}**:")
|
|
for note in notes:
|
|
parts.append(f"- {note}")
|
|
return "\n".join(parts)
|
|
|
|
# Node: Generate comparison criteria
|
|
def plan_criteria(state: CompareState) -> CompareState:
|
|
entities = state["entities"]
|
|
prompt = (
|
|
f"Generate 3-5 concise comparison criteria for the following entities: "
|
|
f"{', '.join(entities)}. Return a numbered list."
|
|
)
|
|
response = llm.invoke(prompt)
|
|
# Parse numbered list
|
|
criteria = []
|
|
for line in response.splitlines():
|
|
line = line.strip()
|
|
if line:
|
|
# Remove leading numbers
|
|
if line[0].isdigit() and (len(line) > 1 and line[1] in ". "):
|
|
line = line[2:].strip()
|
|
criteria.append(line)
|
|
state["criteria"] = criteria
|
|
# Initialize findings dict
|
|
state["findings"] = {entity: [] for entity in entities}
|
|
return state
|
|
|
|
# Node: Research one entity-criterion pair
|
|
def research_entity(state: CompareState) -> CompareState:
|
|
entities = state["entities"]
|
|
criteria = state["criteria"]
|
|
findings = state["findings"]
|
|
|
|
# Find next entity needing research
|
|
for entity in entities:
|
|
if len(findings[entity]) < len(criteria):
|
|
idx = len(findings[entity])
|
|
criterion = criteria[idx]
|
|
query = f"{entity} {criterion}"
|
|
# Tavily search
|
|
results = tavily.invoke({"query": query, "max_results": 3})
|
|
# Take first result snippet
|
|
if results and "results" in results and len(results["results"]) > 0:
|
|
snippet = results["results"][0]["snippet"]
|
|
source = results["results"][0]["url"]
|
|
note = f"{criterion}: {snippet} (Source: {source})"
|
|
else:
|
|
note = f"{criterion}: No recent information found."
|
|
findings[entity].append(note)
|
|
break
|
|
state["findings"] = findings
|
|
return state
|
|
|
|
# Node: Build cohesive research brief
|
|
def build_brief(state: CompareState) -> CompareState:
|
|
findings_text = format_findings(state["findings"])
|
|
prompt = (
|
|
f"Using the following findings, write a cohesive research brief that summarizes "
|
|
f"the strengths and weaknesses of each entity. The brief should be clear, "
|
|
f"structured, and suitable for a technical audience.\n\n"
|
|
f"Findings:\n{findings_text}"
|
|
)
|
|
brief = llm.invoke(prompt)
|
|
state["final_brief"] = brief
|
|
return state
|
|
|
|
# Node: Generate verdict/recommendation
|
|
def verdict(state: CompareState) -> CompareState:
|
|
brief = state["final_brief"]
|
|
criteria = state["criteria"]
|
|
prompt = (
|
|
f"Based on the research brief below and the comparison criteria, provide a "
|
|
f"clear recommendation on which entity is best suited for a typical use case. "
|
|
f"Explain your reasoning in 2-4 sentences.\n\n"
|
|
f"Research Brief:\n{brief}\n\n"
|
|
f"Criteria:\n- " + "\n- ".join(criteria)
|
|
)
|
|
recommendation = llm.invoke(prompt)
|
|
state["verdict"] = recommendation
|
|
return state |