RAG responses

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
   name="text2text",
-  version="1.7.4",
+  version="1.7.5",
   author="artitw",
   author_email="artitw@gmail.com",
   description="Text2Text: Crosslingual NLP/G toolkit",
@@ -18,7 +18,7 @@
     "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
   ],
-  keywords='multilingual crosslingual gpt chatgpt bert natural language processing nlp nlg text generation gpt question answer answering information retrieval tfidf tf-idf bm25 search index summary summarizer summarization tokenizer tokenization translation backtranslation data augmentation science machine learning colab embedding levenshtein sub-word edit distance conversational dialog chatbot llama',
+  keywords='multilingual crosslingual gpt chatgpt bert natural language processing nlp nlg text generation gpt question answer answering information retrieval tfidf tf-idf bm25 search index summary summarizer summarization tokenizer tokenization translation backtranslation data augmentation science machine learning colab embedding levenshtein sub-word edit distance conversational dialog chatbot llama rag',
   install_requires=[
     'faiss-cpu',
     'flask',

text2text/rag_assistant.py

@@ -26,6 +26,21 @@ def is_valid_url(url):
   except Exception:
     return False
 
+def is_affirmative(response):
+    affirmative_keywords = [
+        "yes", "yeah", "yep", "sure", "absolutely", "definitely", 
+        "certainly", "of course", "indeed", "affirmative", "correct", 
+        "right", "exactly", "true", "positive"
+    ]
+
+    response_lower = response.lower()
+
+    for keyword in affirmative_keywords:
+        if keyword in response_lower:
+            return True
+
+    return False
+
 class RagAssistant(t2t.Assistant):
   def __init__(self, **kwargs):
     super().__init__(**kwargs)
@@ -57,7 +72,17 @@ 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"]
-    docs = self.index.retrieve([query], k=k)[0]
-    grounding_information = "Ground your response based on the following information:\n\n" + "\n- ".join(docs)
-    messages[-1] = {"role": "user", "content": query + "\n\n" + grounding_information}
+    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"]
+      docs = self.index.retrieve([demand], k=k)[0]
+    else:
+      docs = self.index.retrieve([query], k=k)[0]
+    grounding_prompt = "Base your response on the following information:\n\n" + "\n- ".join(docs)
+    messages[-1] = {"role": "user", "content": query + "\n\n" + grounding_prompt}
     return t2t.Assistant.chat_completion(self, messages=messages, stream=stream, schema=schema, **kwargs)