From faca3fbc7d43a36dff3e60142408de2eae9f7bfd Mon Sep 17 00:00:00 2001
From: Vladimir Lipkin <lipkin@yandex-team.ru>
Date: Fri, 13 Dec 2024 21:27:16 +0300
Subject: [PATCH] Imrpove example for RAG

---
 .../assistants/assistant_with_search_index.py | 29 +++++++++++++------
 .../assistants/assistant_with_search_index.py | 29 +++++++++++++------
 2 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/examples/async/assistants/assistant_with_search_index.py b/examples/async/assistants/assistant_with_search_index.py
index 922f8c0..73ad350 100755
--- a/examples/async/assistants/assistant_with_search_index.py
+++ b/examples/async/assistants/assistant_with_search_index.py
@@ -33,15 +33,26 @@ async def main() -> None:
     assistant = await sdk.assistants.create('yandexgpt', tools=[tool])
     thread = await sdk.threads.create()
 
-    for search_query in (
-        local_path('search_query.txt').read_text().splitlines()[0],
-        "Cколько пошлина в Анталье"
-    ):
-        await thread.write(search_query)
-        run = await assistant.run(thread)
-        result = await run
-        print('Question', search_query)
-        print('Answer:', result.text)
+    search_query = local_path('search_query.txt').read_text().splitlines()[0]
+    await thread.write(search_query)
+    run = await assistant.run(thread)
+
+    # poll_inteval is 0.5s by default, but you could lower it to optimize
+    # wait time
+    result = await run.wait(poll_interval=0.05)
+    print('Question:', search_query)
+    print('Answer:', result.text)
+
+    search_query = "Cколько пошлина в Анталье"
+    await thread.write(search_query)
+
+    # You could also use run_stream method to start gettig response parts
+    # as soon it will be generated
+    run = await assistant.run_stream(thread)
+    print('Question:', search_query)
+    async for event in run:
+        print("Answer part:", event.text)
+        print("Answer status:", event.status)
 
     await search_index.delete()
     await thread.delete()
diff --git a/examples/sync/assistants/assistant_with_search_index.py b/examples/sync/assistants/assistant_with_search_index.py
index 4431c37..b24e95d 100755
--- a/examples/sync/assistants/assistant_with_search_index.py
+++ b/examples/sync/assistants/assistant_with_search_index.py
@@ -31,15 +31,26 @@ def main() -> None:
     assistant = sdk.assistants.create('yandexgpt', tools=[tool])
     thread = sdk.threads.create()
 
-    for search_query in (
-        local_path('search_query.txt').read_text().splitlines()[0],
-        "Cколько пошлина в Анталье"
-    ):
-        thread.write(search_query)
-        run = assistant.run(thread)
-        result = run.wait()
-        print('Question', search_query)
-        print('Answer:', result.text)
+    search_query = local_path('search_query.txt').read_text().splitlines()[0]
+    thread.write(search_query)
+    run = assistant.run(thread)
+
+    # poll_inteval is 0.5s by default, but you could lower it to optimize
+    # wait time
+    result = run.wait(poll_interval=0.05)
+    print('Question:', search_query)
+    print('Answer:', result.text)
+
+    search_query = "Cколько пошлина в Анталье"
+    thread.write(search_query)
+
+    # You could also use run_stream method to start gettig response parts
+    # as soon it will be generated
+    run = assistant.run_stream(thread)
+    print('Question:', search_query)
+    for event in run:
+        print("Answer part:", event.text)
+        print("Answer status:", event.status)
 
     search_index.delete()
     thread.delete()