From 3da8622cb5544168347e5a6b910adf7afd733095 Mon Sep 17 00:00:00 2001
From: artitw <artitw@gmail.com>
Date: Sun, 8 Dec 2024 21:33:41 +0000
Subject: [PATCH] RAG assistant fixes

---
 setup.py                   |  2 +-
 text2text/rag_assistant.py | 12 ++++--------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/setup.py b/setup.py
index 7efc5aa..47dc2c2 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
   name="text2text",
-  version="1.8.4",
+  version="1.8.5",
   author="artitw",
   author_email="artitw@gmail.com",
   description="Text2Text Language Modeling Toolkit",
diff --git a/text2text/rag_assistant.py b/text2text/rag_assistant.py
index f9402ef..fc3f0e4 100644
--- a/text2text/rag_assistant.py
+++ b/text2text/rag_assistant.py
@@ -114,14 +114,10 @@ def __init__(self, **kwargs):
   def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, schema=None, **kwargs):
     k = kwargs.get("k", 3)
     query = messages[-1]["content"]
-    question_check = f"Respond YES if this is a question; otherwise respond NO: {query}"
-    question_check = [{"role": "user", "content": question_check}]
-    response = t2t.Assistant.chat_completion(self, question_check)["message"]["content"]
-    docs = []
-    if is_affirmative(response):
-      reword_prompt = f"Reword this question to be a demand: {query}"
-      reword_prompt = [{"role": "user", "content": reword_prompt}]
-      demand = t2t.Assistant.chat_completion(self, reword_prompt)["message"]["content"]
+    if len(query) > 100:
+      summarize_prompt = f'Summarize succinctly what the user wants in a sentence less than 10 words:\n\n"{query}"'
+      summarize_prompt = [{"role": "user", "content": summarize_prompt}]
+      demand = t2t.Assistant.chat_completion(self, summarize_prompt)["message"]["content"]
       docs = self.index.retrieve([demand], k=k)[0]
     else:
       docs = self.index.retrieve([query], k=k)[0]