添加在输出模型output后输出citation的功能(beta)

app.py

@@ -172,6 +172,7 @@ def chat_llama3_8b(message: str,
              ) -> str:
     """
     Generate a streaming response using the llama3-8b model.
+    Will display citations after the response if citations are available.
     """
     # Get citations from vector store
     citation = query_vector_store(vector_store, message, 4, 0.7)
@@ -214,7 +215,18 @@ def chat_llama3_8b(message: str,
     outputs = []
     for text in streamer:
         outputs.append(text)
-        yield "".join(outputs)
+        current_output = "".join(outputs)
+        
+        # If we have citations, append them at the end
+        if citation and text == streamer[-1]:  # On the last chunk
+            citation_display = "\n\nReferences:\n" + "\n".join(
+                f"[{i+1}] {cite.strip()}" 
+                for i, cite in enumerate(citation.split('\n')) 
+                if cite.strip()
+            )
+            current_output += citation_display
+            
+        yield current_output
         
 
 # Gradio block