Публикация решения: update src/main.py

This commit is contained in:
2026-06-17 19:48:57 +00:00
parent 1ce37efc52
commit 146f48d140
+7 -5
View File
@@ -30,11 +30,11 @@ async def plan_criteria(state: CompareState) -> CompareState:
f"You are a helpful assistant. Given the following entities: {entities_str}. " f"You are a helpful assistant. Given the following entities: {entities_str}. "
"Generate 3-5 concise criteria for comparing them. Return a JSON array of strings." "Generate 3-5 concise criteria for comparing them. Return a JSON array of strings."
) )
response = await llm.agenerate([{"role": "user", "content": prompt}]) response = await llm.ainvoke({"role": "user", "content": prompt})
# parse JSON # parse JSON
import json import json
try: try:
criteria = json.loads(response.generations[0][0].text.strip()) criteria = json.loads(response["content"].strip())
except Exception as e: except Exception as e:
raise ValueError(f"Failed to parse criteria: {e}") raise ValueError(f"Failed to parse criteria: {e}")
state["criteria"] = criteria state["criteria"] = criteria
@@ -79,8 +79,8 @@ async def verdict(state: CompareState) -> CompareState:
f"Based on the following table:\n{state['final_table']}\n" f"Based on the following table:\n{state['final_table']}\n"
"Provide a short recommendation (2-4 sentences) on which entity is best for each use case." "Provide a short recommendation (2-4 sentences) on which entity is best for each use case."
) )
response = await llm.agenerate([{"role": "user", "content": prompt}]) response = await llm.ainvoke({"role": "user", "content": prompt})
state["verdict"] = response.generations[0][0].text.strip() state["verdict"] = response["content"].strip()
return state return state
# ---------- Graph ---------- # ---------- Graph ----------
@@ -109,8 +109,10 @@ app = builder.compile()
if __name__ == "__main__": if __name__ == "__main__":
import argparse import argparse
parser = argparse.ArgumentParser(description="Compare three entities.") parser = argparse.ArgumentParser(description="Compare three entities.")
parser.add_argument("entities", nargs=3, help="Three entities to compare") parser.add_argument("entities", nargs='*', help="Three entities to compare (default: Chroma, FAISS, Qdrant)")
args = parser.parse_args() args = parser.parse_args()
if not args.entities:
args.entities = ["Chroma", "FAISS", "Qdrant"]
init_state: CompareState = { init_state: CompareState = {
"entities": list(args.entities), "entities": list(args.entities),
"criteria": [], "criteria": [],