Add limit in ui

1 files changed, 1 insertions, 1 deletions

diff --git a/rag/ui.py b/rag/ui.py
index 253ad87..23c44f9 100644
--- a/rag/ui.py
+++ b/rag/ui.py
@@ -102,7 +102,7 @@ def generate_chat(query: str):
     retriever = ss.retriever
     generator = ss.generator
 
-    documents = retriever.retrieve(query)
+    documents = retriever.retrieve(query, limit=5)
     prompt = Prompt(query, documents)
 
     with st.chat_message(ss.bot):