RAG Placeholder demo test

app.py

@@ -172,13 +172,29 @@ def llmguide_generate_response(prompt, stream=False):
     # torch.cuda.empty_cache()
     # gc.collect()
 
+# def rag(query, stream=False):
+#     retrieved_docs = retrieve(query)
+#     context = " ".join(retrieved_docs)
+#     prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
+#     return llmguide_generate_response(prompt, stream)
+
 def rag(query, stream=False):
     retrieved_docs = retrieve(query)
     context = " ".join(retrieved_docs)
     prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
-    return llmguide_generate_response(prompt, stream)
-
-
+    
+    generator = llmguide_generate_response(prompt, stream)
+    
+    if stream:
+        def stream_output():
+            for generated_text, tokens_per_second, ram_usage in generator:
+                yield generated_text, tokens_per_second, ram_usage
+        return stream_output()
+    else:
+        # For non-streaming, we just need to get the final output
+        for generated_text, tokens_per_second, ram_usage in generator:
+            pass  # This will iterate to the last yield
+        return generated_text, tokens_per_second, ram_usage
 
 #--------------------------------------------------------------------------------------------------------------------------------