Публикация решения: update src/main.py
This commit is contained in:
+7
-5
@@ -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": [],
|
||||||
|
|||||||
Reference in New Issue
Block a user