switched to chat engine

backend.py

@@ -55,27 +55,63 @@ def build_index():
 
 @spaces.GPU(duration=20)
 def handle_query(query_str, chathistory):
-    index = build_index()
     
-    memory = ChatMemoryBuffer.from_defaults(token_limit=1500)
-    chat_engine = index.as_chat_engine(
+    index = build_index()
+  
+    qa_prompt_str = (
+        "Context information is below.\n"
+        "---------------------\n"
+        "{context_str}\n"
+        "---------------------\n"
+        "Given the context information and not prior knowledge, "
+        "answer the question: {query_str}\n"
+    )
+
+    # Text QA Prompt
+    chat_text_qa_msgs = [
+        (
+            "system",
+            "Sei un assistente italiano di nome Ossy che risponde solo alle domande o richieste pertinenti. ",
+        ),
+        ("user", qa_prompt_str),
+    ]
+    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
+
+    try:
+        # Create a streaming query engine
+        """query_engine = index.as_query_engine(text_qa_template=text_qa_template, streaming=False, similarity_top_k=1)
+        
+        # Execute the query
+        streaming_response = query_engine.query(query_str)
+        
+        r = streaming_response.response
+        cleaned_result = r.replace("<end_of_turn>", "").strip()
+        yield cleaned_result"""
+      
+        # Stream the response
+        """outputs = []
+        for text in streaming_response.response_gen:
+            
+            outputs.append(str(text))
+            yield "".join(outputs)"""
+          
+        memory = ChatMemoryBuffer.from_defaults(token_limit=1500)
+        chat_engine = index.as_chat_engine(
         chat_mode="context",
         memory=memory,
         system_prompt=(
             "Sei un assistente italiano di nome Ossy che risponde solo alle domande o richieste pertinenti. "
-        ),
-    )
-    
-    try:
+          ),
+        )
+        
         response = chat_engine.stream_chat(query_str)
+        #response = chat_engine.chat(query_str)
         for token in response.response_gen:
-            if token.strip():  # Only yield non-empty tokens
-                yield token
+            yield token
+          
+        
     except Exception as e:
         yield f"Error processing query: {str(e)}"
-    finally:
-        # You might want to add any cleanup code here
-        pass
 
 
 

interface.py

@@ -69,7 +69,11 @@ class GemmaLLMInterface(CustomLLM):
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
         streamer, generate_kwargs = self._prepare_generation(prompt)
 
-        self.model.generate(**generate_kwargs)  # Run synchronously.
+        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
+        t.start()
 
-        for new_token in streamer:
-            yield CompletionResponse(text=new_token)
+        try:
+            for new_token in streamer:
+                yield CompletionResponse(text=new_token)
+        except StopIteration:
+            return