Update app/rag_pipeline.py

app/rag_pipeline.py

@@ -9,6 +9,7 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.embeddings import SentenceTransformerEmbeddings
 
 CHROMA_DIR = "data/chroma_db"
+CHROMA_IMG_DIR = "data/image_db"
 
 load_dotenv()
 HUGGINGFACE_API_KEY = os.getenv("HF_API_KEY")  # set this in .env or directly
@@ -83,3 +84,32 @@ def generate_answer(query, top_k=3):
 
     answer = call_huggingface_mistral(prompt)
     return {"answer": answer}
+
+
+def generate_answer_am(query, top_k=3):
+
+    CHROMA_DIR = "data/chroma_db"
+    embeddings = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
+
+    db = Chroma(persist_directory=CHROMA_IMG_DIR,
+                embedding_function=embeddings)
+    retriever = db.as_retriever(
+        search_type='similarity',
+        search_kwargs={'k': 4}
+    )
+    relevant_document_chunks = retriever.get_relevant_documents(
+        query=query, k=top_k)
+    context_list = [d.page_content for d in relevant_document_chunks]
+    print(f'context_list: {context_list}')
+
+    # Combine document chunks into a single context
+    context = ". ".join(context_list)
+    user_message = qna_user_message_template.replace(
+        '{context}', context)
+    user_message = user_message.replace('{question}', query)
+
+    prompt = qna_system_message + '\n' + user_message
+    print(f'Prompt: {prompt}')
+
+    answer = call_huggingface_mistral(prompt)
+    return {"answer": answer}