From 146f48d1409e59cad2adbc44c7eded7dc10b2923 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?=
 =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= <mirzagitov03@mail.ru>
Date: Wed, 17 Jun 2026 19:48:57 +0000
Subject: [PATCH] =?UTF-8?q?=D0=9F=D1=83=D0=B1=D0=BB=D0=B8=D0=BA=D0=B0?=
 =?UTF-8?q?=D1=86=D0=B8=D1=8F=20=D1=80=D0=B5=D1=88=D0=B5=D0=BD=D0=B8=D1=8F?=
 =?UTF-8?q?:=20update=20src/main.py?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/main.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/main.py b/src/main.py
index 89bb871..5752bb9 100644
--- a/src/main.py
+++ b/src/main.py
@@ -30,11 +30,11 @@ async def plan_criteria(state: CompareState) -> CompareState:
         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."
     )
-    response = await llm.agenerate([{"role": "user", "content": prompt}])
+    response = await llm.ainvoke({"role": "user", "content": prompt})
     # parse JSON
     import json
     try:
-        criteria = json.loads(response.generations[0][0].text.strip())
+        criteria = json.loads(response["content"].strip())
     except Exception as e:
         raise ValueError(f"Failed to parse criteria: {e}")
     state["criteria"] = criteria
@@ -79,8 +79,8 @@ async def verdict(state: CompareState) -> CompareState:
         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."
     )
-    response = await llm.agenerate([{"role": "user", "content": prompt}])
-    state["verdict"] = response.generations[0][0].text.strip()
+    response = await llm.ainvoke({"role": "user", "content": prompt})
+    state["verdict"] = response["content"].strip()
     return state
 
 # ---------- Graph ----------
@@ -109,8 +109,10 @@ app = builder.compile()
 if __name__ == "__main__":
     import argparse
     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()
+    if not args.entities:
+        args.entities = ["Chroma", "FAISS", "Qdrant"]
     init_state: CompareState = {
         "entities": list(args.entities),
         "criteria": [],