switched to chat engine

backend.py

@@ -55,62 +55,27 @@ def build_index():
 
 @spaces.GPU(duration=20)
 def handle_query(query_str, chathistory):
-    
     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(
+    
+    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)
         for token in response.response_gen:
-            yield token
-          
-        
+            if token.strip():  # Only yield non-empty tokens
+                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

@@ -68,9 +68,8 @@ class GemmaLLMInterface(CustomLLM):
     @llm_completion_callback()
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
         streamer, generate_kwargs = self._prepare_generation(prompt)
-        
-        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
-        t.start()
-        
+
+        self.model.generate(**generate_kwargs)  # Run synchronously.
+
         for new_token in streamer:
             yield CompletionResponse(text=new_token)