Update rag_pipeline.py

rag_pipeline.py

@@ -5,13 +5,18 @@ from embedding_functions import embed_text
 
 llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct")
 
-def retrieve_multimodal(query, vector_store, k=5):
-    query_embedding = embed_text(query)
-    results = vector_store.similarity_search_by_vector(
-        embedding = query_embedding,
-        k = k
-    )
-    return results
+def retrieve_multimodal(query, vector_store):
+    # 1. Cast a much wider net so images don't get pushed out by text scores
+    all_results = vector_store.similarity_search(query=query, k=20)
+    
+    # 2. Separate the results manually
+    text_docs = [doc for doc in all_results if doc.metadata.get("type") == "text"]
+    image_docs = [doc for doc in all_results if doc.metadata.get("type") == "image"]
+    
+    # 3. Guarantee that we return the top 5 texts AND the top 2 images
+    final_docs = text_docs[:5] + image_docs[:2]
+    
+    return final_docs
 
 def create_multimodal_message(query, retrieved_docs, image_data_store):
     content = []