Wrap-up retrieval chain into LangGraph

bin/chat-chainlit.py

@@ -7,6 +7,7 @@
 import chainlit as cl
 from dotenv import load_dotenv
 
+from conversational_chain.graph import RAGGraphWithMemory
 from retreival_chain import initialize_retrieval_chain
 from util.embedding_environment import EmbeddingEnvironment
 
@@ -57,14 +58,14 @@ async def start() -> None:
     chat_profile = cl.user_session.get("chat_profile")
 
     embeddings_directory = EmbeddingEnvironment.get_dir(env)
-    llm_chain = initialize_retrieval_chain(
+    llm_graph = initialize_retrieval_chain(
         env,
         embeddings_directory,
         False,
         False,
         hf_model=EmbeddingEnvironment.get_model(env),
     )
-    cl.user_session.set("llm_chain", llm_chain)
+    cl.user_session.set("llm_graph", llm_graph)
 
     initial_message: str = f"""Welcome to {chat_profile} your interactive chatbot for exploring Reactome!
         Ask me about biological pathways and processes"""
@@ -73,13 +74,15 @@ async def start() -> None:
 
 @cl.on_message
 async def main(message: cl.Message) -> None:
-    llm_chain = cl.user_session.get("llm_chain")
+    llm_graph: RAGGraphWithMemory = cl.user_session.get("llm_graph")
     cb = cl.AsyncLangchainCallbackHandler(
         stream_final_answer=True, answer_prefix_tokens=["FINAL", "ANSWER"]
     )
-    cb.answer_reached = True
-
-    res = await llm_chain.ainvoke(message.content, callbacks=[cb])
+    res = await llm_graph.ainvoke(
+        message.content,
+        callbacks = [cb],
+        configurable = {"thread_id": "0"}  # single thread
+    )
     if cb.has_streamed_final_answer and cb.final_stream is not None:
         await cb.final_stream.update()
     else:

pyproject.toml

@@ -35,6 +35,7 @@ torch = {version = "^2.4.0+cpu", source = "pytorch_cpu"}
 langchain-chroma = "^0.1.4"
 langchain-ollama = "^0.2.0"
 lark = "^1.2.2"
+langgraph = "^0.2.39"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.7.1"

src/conversational_chain/chain.py

@@ -1,12 +1,17 @@
-from langchain.chains import (create_history_aware_retriever,
-                              create_retrieval_chain)
 from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains.history_aware_retriever import \
+    create_history_aware_retriever
+from langchain.chains.retrieval import create_retrieval_chain
+from langchain_core.language_models import LanguageModelLike
+from langchain_core.retrievers import RetrieverLike
 
 from system_prompt.reactome_prompt import contextualize_q_prompt, qa_prompt
 
 
 class RAGChainWithMemory:
-    def __init__(self, memory, retriever, llm):
+    def __init__(
+        self, memory, retriever: RetrieverLike, llm: LanguageModelLike
+    ):
         """
         Initializes the Retrieval-Augmented Generation (RAG) chain with memory.
         """