Update app.py

app.py

@@ -9,9 +9,11 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub import HfApi, login
 from datetime import datetime
 from langchain_openai import ChatOpenAI
-from langchain.chains import create_retrieval_chain
-from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain_core.prompts import ChatPromptTemplate
+
+
+from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
 
 from manabUtils import retrieve_chunks
 
@@ -63,32 +65,27 @@ PQC_rules="""
 #===========================
 def create_qa_chain():
     retriever = retrieve_chunks(repo_id)
-    system_prompt = (
-        "You are a helpful assistant. Use the following context to answer the question. "
-        "If you don't know the answer, say so.\n\n{context}"
+    prompt = ChatPromptTemplate.from_template(
+        "Use context to answer: {context}\n\nQ: {input}"
+    )
+    chain = (
+        {"context": retriever | (lambda docs: "\n\n".join(doc.page_content for doc in docs)),
+         "input": RunnablePassthrough()}
+        | prompt
+        | llm
+        | StrOutputParser()
     )
-    prompt = ChatPromptTemplate.from_messages([
-        ("system", system_prompt),
-        ("human", "{input}")
-    ])
-    doc_chain = create_stuff_documents_chain(llm, prompt)
-    return create_retrieval_chain(retriever, doc_chain)
+    return chain
 
 qa_chain = create_qa_chain()
 #=======================
 def chat(message, history):
-    response = qa_chain.invoke({"input": message})
-    answer = response["answer"]
-    
-    # Extract references with pages
-    docs = response["context"]
-    refs = []
-    for i, doc in enumerate(docs):
-        page = doc.metadata.get("page", "N/A")
-        refs.append(f"Ref {i+1}: Page {page}")
-    
-    full_response = f"{answer}\n\n**References:**\n" + "\n".join(refs)
-    history.append([message, full_response])
+    answer = qa_chain.invoke(message)
+    # Get docs for refs
+    docs = retriever.invoke(message)
+    refs = [f"Page {d.metadata.get('page', 'N/A')}" for d in docs]
+    full = f"{answer}\n\nRefs: {' | '.join(refs)}"
+    history.append([message, full])
     return history, ""
 #============starting extract_docx_text