Update deprecated chains to LCEL (#198)

* update chains
* Format

Author: tomasonjo

api.py

@@ -128,10 +128,7 @@ def qstream(question: Question = Depends()):
     q = Queue()
 
     def cb():
-        output_function(
-            {"question": question.text, "chat_history": []},
-            callbacks=[QueueCallback(q)],
-        )
+        output_function.invoke(question.text, config={"callbacks": [QueueCallback(q)]})
 
     def generate():
         yield json.dumps({"init": True, "model": llm_name})
@@ -146,9 +143,7 @@ async def ask(question: Question = Depends()):
     output_function = llm_chain
     if question.rag:
         output_function = rag_chain
-    result = output_function(
-        {"question": question.text, "chat_history": []}, callbacks=[]
-    )
+    result = output_function.invoke(question.text)
 
     return {"result": result["answer"], "model": llm_name}
 

bot.py

@@ -92,10 +92,10 @@ def chat_input():
         with st.chat_message("assistant"):
             st.caption(f"RAG: {name}")
             stream_handler = StreamHandler(st.empty())
-            result = output_function(
-                {"question": user_input, "chat_history": []}, callbacks=[stream_handler]
-            )["answer"]
-            output = result
+            output = output_function.invoke(
+                user_input, config={"callbacks": [stream_handler]}
+            )
+
             st.session_state[f"user_input"].append(user_input)
             st.session_state[f"generated"].append(output)
             st.session_state[f"rag_mode"].append(name)

chains.py

@@ -1,4 +1,3 @@
-
 from langchain_openai import OpenAIEmbeddings
 from langchain_ollama import OllamaEmbeddings
 from langchain_aws import BedrockEmbeddings
@@ -10,17 +9,17 @@
 
 from langchain_neo4j import Neo4jVector
 
-from langchain.chains import RetrievalQAWithSourcesChain
-from langchain.chains.qa_with_sources import load_qa_with_sources_chain
+from langchain_core.runnables import RunnableParallel, RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
 
 from langchain.prompts import (
     ChatPromptTemplate,
     HumanMessagePromptTemplate,
-    SystemMessagePromptTemplate
+    SystemMessagePromptTemplate,
 )
 
 from typing import List, Any
-from utils import BaseLogger, extract_title_and_question
+from utils import BaseLogger, extract_title_and_question, format_docs
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
 
 AWS_MODELS = (
@@ -32,6 +31,7 @@
     "mistral.mi",
 )
 
+
 def load_embedding_model(embedding_model_name: str, logger=BaseLogger(), config={}):
     if embedding_model_name == "ollama":
         embeddings = OllamaEmbeddings(
@@ -47,10 +47,8 @@ def load_embedding_model(embedding_model_name: str, logger=BaseLogger(), config=
         embeddings = BedrockEmbeddings()
         dimension = 1536
         logger.info("Embedding: Using AWS")
-    elif embedding_model_name == "google-genai-embedding-001":        
-        embeddings = GoogleGenerativeAIEmbeddings(
-            model="models/embedding-001"
-        )
+    elif embedding_model_name == "google-genai-embedding-001":
+        embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
         dimension = 768
         logger.info("Embedding: Using Google Generative AI Embeddings")
     else:
@@ -112,17 +110,8 @@ def configure_llm_only_chain(llm):
     chat_prompt = ChatPromptTemplate.from_messages(
         [system_message_prompt, human_message_prompt]
     )
-
-    def generate_llm_output(
-        user_input: str, callbacks: List[Any], prompt=chat_prompt
-    ) -> str:
-        chain = prompt | llm
-        answer = chain.invoke(
-            {"question": user_input}, config={"callbacks": callbacks}
-        ).content
-        return {"answer": answer}
-
-    return generate_llm_output
+    chain = chat_prompt | llm | StrOutputParser()
+    return chain
 
 
 def configure_qa_rag_chain(llm, embeddings, embeddings_store_url, username, password):
@@ -152,12 +141,6 @@ def configure_qa_rag_chain(llm, embeddings, embeddings_store_url, username, pass
     ]
     qa_prompt = ChatPromptTemplate.from_messages(messages)
 
-    qa_chain = load_qa_with_sources_chain(
-        llm,
-        chain_type="stuff",
-        prompt=qa_prompt,
-    )
-
     # Vector + Knowledge Graph response
     kg = Neo4jVector.from_existing_index(
         embedding=embeddings,
@@ -183,12 +166,16 @@ def configure_qa_rag_chain(llm, embeddings, embeddings_store_url, username, pass
     ORDER BY similarity ASC // so that best answers are the last
     """,
     )
-
-    kg_qa = RetrievalQAWithSourcesChain(
-        combine_documents_chain=qa_chain,
-        retriever=kg.as_retriever(search_kwargs={"k": 2}),
-        reduce_k_below_max_tokens=False,
-        max_tokens_limit=3375,
+    kg_qa = (
+        RunnableParallel(
+            {
+                "summaries": kg.as_retriever(search_kwargs={"k": 2}) | format_docs,
+                "question": RunnablePassthrough(),
+            }
+        )
+        | qa_prompt
+        | llm
+        | StrOutputParser()
     )
     return kg_qa
 

pdf_bot.py

@@ -1,16 +1,19 @@
 import os
 
 import streamlit as st
-from langchain.chains import RetrievalQA
 from PyPDF2 import PdfReader
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.prompts import ChatPromptTemplate
 from langchain_neo4j import Neo4jVector
 from streamlit.logger import get_logger
 from chains import (
     load_embedding_model,
     load_llm,
 )
+from langchain_core.runnables import RunnableParallel, RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+from utils import format_docs
 
 # load api key lib
 from dotenv import load_dotenv
@@ -67,6 +70,14 @@ def main():
         )
 
         chunks = text_splitter.split_text(text=text)
+        qa_prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "human",
+                    "Based on the provided summary: {summaries} \n Answer the following question:{question}",
+                )
+            ]
+        )
 
         # Store the chunks part in db (vector)
         vectorstore = Neo4jVector.from_texts(
@@ -79,16 +90,25 @@ def main():
             node_label="PdfBotChunk",
             pre_delete_collection=True,  # Delete existing PDF data
         )
-        qa = RetrievalQA.from_chain_type(
-            llm=llm, chain_type="stuff", retriever=vectorstore.as_retriever()
+        qa = (
+            RunnableParallel(
+                {
+                    "summaries": vectorstore.as_retriever(search_kwargs={"k": 2})
+                    | format_docs,
+                    "question": RunnablePassthrough(),
+                }
+            )
+            | qa_prompt
+            | llm
+            | StrOutputParser()
         )
 
         # Accept user questions/query
         query = st.text_input("Ask questions about your PDF file")
 
         if query:
             stream_handler = StreamHandler(st.empty())
-            qa.run(query, callbacks=[stream_handler])
+            qa.invoke(query, {"callbacks": [stream_handler]})
 
 
 if __name__ == "__main__":

utils.py

@@ -32,7 +32,9 @@ def create_vector_index(driver) -> None:
         driver.query(index_query)
     except:  # Already exists
         pass
-    index_query = "CREATE VECTOR INDEX top_answers IF NOT EXISTS FOR (m:Answer) ON m.embedding"
+    index_query = (
+        "CREATE VECTOR INDEX top_answers IF NOT EXISTS FOR (m:Answer) ON m.embedding"
+    )
     try:
         driver.query(index_query)
     except:  # Already exists
@@ -52,3 +54,7 @@ def create_constraints(driver):
     driver.query(
         "CREATE CONSTRAINT tag_name IF NOT EXISTS FOR (t:Tag) REQUIRE (t.name) IS UNIQUE"
     )
+
+
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)