feat: solution for 'Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily)'
This commit is contained in:
+21
-20
@@ -1,44 +1,45 @@
|
||||
import argparse
|
||||
from typing import List
|
||||
import os
|
||||
from src.graph import build_graph
|
||||
from src.state import CompareState
|
||||
|
||||
from .state import CompareState
|
||||
from .graph import create_graph
|
||||
|
||||
def parse_entities(arg: str) -> List[str]:
|
||||
def parse_entities(arg: str) -> list[str]:
|
||||
return [e.strip() for e in arg.split(",") if e.strip()]
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="LangGraph Comparative Review Agent")
|
||||
parser = argparse.ArgumentParser(description="Research Brief Generator")
|
||||
parser.add_argument(
|
||||
"-e",
|
||||
"--entities",
|
||||
type=str,
|
||||
help="Comma-separated list of three entities to compare. "
|
||||
"If omitted, defaults to Chroma, FAISS, Qdrant.",
|
||||
help="Comma-separated list of 3 entities to compare",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.entities:
|
||||
entities = parse_entities(args.entities)
|
||||
if len(entities) != 3:
|
||||
raise ValueError("Please provide exactly three entities.")
|
||||
print("Please provide exactly 3 entities.")
|
||||
return
|
||||
else:
|
||||
# Default entities
|
||||
entities = ["Chroma", "FAISS", "Qdrant"]
|
||||
|
||||
initial_state: CompareState = {
|
||||
# Initial state
|
||||
state: CompareState = {
|
||||
"entities": entities,
|
||||
"criteria": [],
|
||||
"findings": {},
|
||||
"final_brief": None,
|
||||
"verdict": None,
|
||||
}
|
||||
|
||||
graph = create_graph()
|
||||
final_state = graph.invoke(initial_state)
|
||||
graph = build_graph()
|
||||
final_state = graph.invoke(state)
|
||||
|
||||
print("\n=== Comparison Criteria ===")
|
||||
for idx, crit in enumerate(final_state["criteria"], 1):
|
||||
print(f"{idx}. {crit}")
|
||||
|
||||
print("\n=== Findings Table ===")
|
||||
print(final_state["final_table"])
|
||||
|
||||
print("\n=== Verdict ===")
|
||||
print("\n=== Research Brief ===\n")
|
||||
print(final_state["final_brief"])
|
||||
print("\n=== Verdict ===\n")
|
||||
print(final_state["verdict"])
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+26
-8
@@ -1,21 +1,39 @@
|
||||
from langgraph.graph import StateGraph
|
||||
from .state import CompareState
|
||||
from .nodes import plan_criteria, research_entity, build_table, verdict
|
||||
from langgraph import StateGraph
|
||||
from src.state import CompareState
|
||||
from src.nodes import (
|
||||
plan_criteria,
|
||||
research_entity,
|
||||
build_brief,
|
||||
verdict,
|
||||
)
|
||||
|
||||
def create_graph() -> StateGraph:
|
||||
def build_graph() -> StateGraph:
|
||||
graph = StateGraph(CompareState)
|
||||
|
||||
# Add nodes
|
||||
graph.add_node("plan_criteria", plan_criteria)
|
||||
graph.add_node("research_entity", research_entity)
|
||||
graph.add_node("build_table", build_table)
|
||||
graph.add_node("build_brief", build_brief)
|
||||
graph.add_node("verdict", verdict)
|
||||
|
||||
# Define edges
|
||||
graph.set_entry_point("plan_criteria")
|
||||
graph.add_edge("plan_criteria", "research_entity")
|
||||
graph.add_edge("research_entity", "build_table")
|
||||
graph.add_edge("build_table", "verdict")
|
||||
graph.add_edge("verdict", END)
|
||||
|
||||
# Conditional loop: if research still needed, stay in research_entity
|
||||
def research_done(state: CompareState) -> str:
|
||||
entities = state["entities"]
|
||||
criteria = state["criteria"]
|
||||
findings = state["findings"]
|
||||
# Check if all entities have enough findings
|
||||
for entity in entities:
|
||||
if len(findings[entity]) < len(criteria):
|
||||
return "research_entity"
|
||||
return "build_brief"
|
||||
|
||||
graph.add_conditional_edges("research_entity", research_done)
|
||||
|
||||
graph.add_edge("build_brief", "verdict")
|
||||
graph.add_edge("verdict", "__end__")
|
||||
|
||||
return graph
|
||||
+3
-1
@@ -1,4 +1,6 @@
|
||||
from .cli import main
|
||||
# Entry point for the package
|
||||
# This file simply calls the CLI main function
|
||||
from src.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+65
-83
@@ -1,117 +1,99 @@
|
||||
import os
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from langgraph.graph import StateGraph, END
|
||||
from langchain_openai import ChatOpenAI
|
||||
from tavily import TavilyClient
|
||||
from typing import Dict, List
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from .state import CompareState
|
||||
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()
|
||||
|
||||
# Initialize LLM and Tavily client
|
||||
llm = ChatOpenAI(
|
||||
temperature=0.2,
|
||||
model="gpt-4o-mini",
|
||||
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
||||
)
|
||||
# LLM and Tavily tool
|
||||
llm = ChatOpenAI(temperature=0.7)
|
||||
tavily = TavilySearchTool(api_key=os.getenv("TAVILY_API_KEY"))
|
||||
|
||||
tavily = TavilyClient(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:
|
||||
"""
|
||||
Generate 3–5 comparison criteria for the given entities.
|
||||
"""
|
||||
entities = state.get("entities", [])
|
||||
if not entities:
|
||||
raise ValueError("No entities provided for criteria planning.")
|
||||
|
||||
entities = state["entities"]
|
||||
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."
|
||||
f"Generate 3-5 concise comparison criteria for the following entities: "
|
||||
f"{', '.join(entities)}. Return a numbered list."
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
criteria_text = response.content.strip()
|
||||
# Parse numbered list
|
||||
criteria = []
|
||||
for line in criteria_text.splitlines():
|
||||
for line in response.splitlines():
|
||||
line = line.strip()
|
||||
if line:
|
||||
# Remove leading numbers if present
|
||||
# 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:
|
||||
"""
|
||||
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}
|
||||
entities = state["entities"]
|
||||
criteria = state["criteria"]
|
||||
findings = state["findings"]
|
||||
|
||||
# Find next entity needing research
|
||||
for entity in entities:
|
||||
for criterion in criteria:
|
||||
if len(findings[entity]) < len(criteria):
|
||||
idx = len(findings[entity])
|
||||
criterion = criteria[idx]
|
||||
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)
|
||||
|
||||
# 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
|
||||
|
||||
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
|
||||
# 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:
|
||||
"""
|
||||
Generate a verdict recommendation based on the table and criteria.
|
||||
"""
|
||||
table = state.get("final_table", "")
|
||||
criteria = state.get("criteria", [])
|
||||
entities = state.get("entities", [])
|
||||
|
||||
brief = state["final_brief"]
|
||||
criteria = state["criteria"]
|
||||
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."
|
||||
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)
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
state["verdict"] = response.content.strip()
|
||||
recommendation = llm.invoke(prompt)
|
||||
state["verdict"] = recommendation
|
||||
return state
|
||||
+3
-3
@@ -1,8 +1,8 @@
|
||||
from typing import TypedDict, List, Dict, Optional
|
||||
|
||||
class CompareState(TypedDict, total=False):
|
||||
class CompareState(TypedDict):
|
||||
entities: List[str] # 3 names to compare
|
||||
criteria: List[str] # 3–5 comparison criteria
|
||||
findings: Dict[str, List[str]] # entity -> list of notes per criterion
|
||||
final_table: Optional[str]
|
||||
verdict: Optional[str]
|
||||
final_brief: Optional[str] # cohesive research brief
|
||||
verdict: Optional[str] # recommendation
|
||||
Reference in New Issue
Block a user