117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import os
|
||
from typing import Dict, List, Tuple
|
||
|
||
from langgraph.graph import StateGraph, END
|
||
from langchain_openai import ChatOpenAI
|
||
from tavily import TavilyClient
|
||
from dotenv import load_dotenv
|
||
|
||
from .state import CompareState
|
||
|
||
load_dotenv()
|
||
|
||
# Initialize LLM and Tavily client
|
||
llm = ChatOpenAI(
|
||
temperature=0.2,
|
||
model="gpt-4o-mini",
|
||
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
||
)
|
||
|
||
tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
|
||
|
||
def plan_criteria(state: CompareState) -> CompareState:
|
||
"""
|
||
Generate 3–5 comparison criteria for the given entities.
|
||
"""
|
||
entities = state.get("entities", [])
|
||
if not entities:
|
||
raise ValueError("No entities provided for criteria planning.")
|
||
|
||
prompt = (
|
||
f"Given the following entities: {', '.join(entities)}.\n"
|
||
"Suggest 3 to 5 key criteria to compare them. "
|
||
"Return the criteria as a numbered list, one per line."
|
||
)
|
||
response = llm.invoke(prompt)
|
||
criteria_text = response.content.strip()
|
||
# Parse numbered list
|
||
criteria = []
|
||
for line in criteria_text.splitlines():
|
||
line = line.strip()
|
||
if line:
|
||
# Remove leading numbers if present
|
||
if line[0].isdigit() and (len(line) > 1 and line[1] in ". "):
|
||
line = line[2:].strip()
|
||
criteria.append(line)
|
||
state["criteria"] = criteria
|
||
return state
|
||
|
||
def research_entity(state: CompareState) -> CompareState:
|
||
"""
|
||
For each entity–criterion pair, perform a Tavily web search
|
||
and store a short note in findings.
|
||
"""
|
||
entities = state.get("entities", [])
|
||
criteria = state.get("criteria", [])
|
||
findings: Dict[str, List[str]] = {entity: [] for entity in entities}
|
||
|
||
for entity in entities:
|
||
for criterion in criteria:
|
||
query = f"{entity} {criterion}"
|
||
try:
|
||
result = tavily.search(query=query, max_results=1)
|
||
if result and result["results"]:
|
||
snippet = result["results"][0]["content"][:200]
|
||
else:
|
||
snippet = "No relevant information found."
|
||
except Exception as e:
|
||
snippet = f"Error during search: {e}"
|
||
findings[entity].append(snippet)
|
||
|
||
state["findings"] = findings
|
||
return state
|
||
|
||
def build_table(state: CompareState) -> CompareState:
|
||
"""
|
||
Build a Markdown table from findings.
|
||
Rows: criteria, Columns: entities.
|
||
"""
|
||
entities = state.get("entities", [])
|
||
criteria = state.get("criteria", [])
|
||
findings = state.get("findings", {})
|
||
|
||
header = "| Criterion | " + " | ".join(entities) + " |\n"
|
||
separator = "|---" * (len(entities) + 1) + "|\n"
|
||
|
||
rows = ""
|
||
for idx, criterion in enumerate(criteria):
|
||
row = f"| {criterion} | "
|
||
for entity in entities:
|
||
notes = findings.get(entity, [])
|
||
note = notes[idx] if idx < len(notes) else ""
|
||
# Escape pipe characters
|
||
note = note.replace("|", "\\|")
|
||
row += f"{note} | "
|
||
rows += row + "\n"
|
||
|
||
table = header + separator + rows
|
||
state["final_table"] = table
|
||
return state
|
||
|
||
def verdict(state: CompareState) -> CompareState:
|
||
"""
|
||
Generate a verdict recommendation based on the table and criteria.
|
||
"""
|
||
table = state.get("final_table", "")
|
||
criteria = state.get("criteria", [])
|
||
entities = state.get("entities", [])
|
||
|
||
prompt = (
|
||
f"Here is a comparative table of the following entities: {', '.join(entities)}.\n\n"
|
||
f"{table}\n\n"
|
||
f"Based on the criteria: {', '.join(criteria)}.\n"
|
||
"Provide a concise recommendation (2–4 sentences) indicating which entity is best suited for which use case."
|
||
)
|
||
response = llm.invoke(prompt)
|
||
state["verdict"] = response.content.strip()
|
||
return state |