Integrate LangGraph chain for AIMS quiz functionality and extend retrieval.py with RAG QA chain

- Added AIMessage and HumanMessage imports to support message processing in chainlit_frontend.py.
- Implemented AIMSState and create_aims_chain function to initialize LangGraph chain with retrieval capabilities.
- Modified main message handler in chainlit_frontend.py to utilize the LangGraph chain for processing user messages and generating responses.
- Extended retrieval.py with get_RAG_QA_chain method to facilitate retrieval-augmented question answering within the AIMS quiz functionality.

aims_tutor/chainlit_frontend.py

@@ -2,6 +2,8 @@ import chainlit as cl
 from dotenv import load_dotenv
 from document_processing import DocumentManager
 from retrieval import RetrievalManager
+from langchain_core.messages import AIMessage, HumanMessage
+from graph import create_aims_chain, AIMSState
 
 # Load environment variables
 load_dotenv()
@@ -36,16 +38,40 @@ async def start_chat():
         cl.user_session.set("docs", doc_manager.get_documents())
         cl.user_session.set("retrieval_manager", RetrievalManager(doc_manager.get_retriever()))
 
+        # Initialize LangGraph chain with the retrieval chain
+        retrieval_chain = cl.user_session.get("retrieval_manager").get_RAG_QA_chain()
+        cl.user_session.set("retrieval_chain", retrieval_chain)  # Store the retrieval chain in the session
+        aims_chain = create_aims_chain(retrieval_chain)
+        cl.user_session.set("aims_chain", aims_chain)
+
 @cl.on_message
 async def main(message: cl.Message):
-    # Retrieve the multi-query retriever from session
-    retrieval_manager = cl.user_session.get("retrieval_manager")
-    if not retrieval_manager:
+    # Retrieve the LangGraph chain from the session
+    aims_chain = cl.user_session.get("aims_chain")
+
+    if not aims_chain:
         await cl.Message(content="No document processing setup found. Please upload a Jupyter notebook first.").send()
         return
 
-    question = message.content
-    response = retrieval_manager.notebook_QA(question)  # Process the question
+    # Create the initial state with the user message
+    user_message = message.content
+    state = AIMSState(messages=[HumanMessage(content=user_message)], next="supervisor", quiz=[])
+
+    print(f"Initial state: {state}")
+
+    # Process the message through the LangGraph chain
+    for s in aims_chain.stream(state, {"recursion_limit": 10}):
+        print(f"State after processing: {s}")
 
-    msg = cl.Message(content=response)
-    await msg.send()
+        # Extract messages from the state
+        if "__end__" not in s:
+            agent_state = next(iter(s.values()))
+            if "messages" in agent_state:
+                response = agent_state["messages"][-1].content
+                print(f"Response: {response}")
+                await cl.Message(content=response).send()
+            else:
+                print("Error: No messages found in agent state.")
+        else:
+            print("Reached end state.")
+            break

aims_tutor/{test.py → graph.py}

@@ -9,10 +9,31 @@ from langchain_core.runnables import RunnablePassthrough
 from langchain_openai import ChatOpenAI
 from langgraph.graph import END, StateGraph
 import functools
+from retrieval import RetrievalManager
 
 # Load environment variables
 load_dotenv()
 
+# Instantiate the language model
+llm = ChatOpenAI(model="gpt-4o")
+
+class RetrievalChainWrapper:
+    def __init__(self, retrieval_chain):
+        self.retrieval_chain = retrieval_chain
+
+    def retrieve_information(
+        self,
+        query: Annotated[str, "query to ask the RAG tool"]
+    ):
+        """Use this tool to retrieve information about the provided notebook."""
+        response = self.retrieval_chain.invoke({"question": query})
+        return response["response"].content
+
+# Create an instance of the wrapper
+def get_retrieve_information_tool(retrieval_chain):
+    wrapper_instance = RetrievalChainWrapper(retrieval_chain)
+    return tool(wrapper_instance.retrieve_information)
+
 @tool
 def generate_quiz(
     documents: Annotated[List[str], "List of documents to generate quiz from"],
@@ -24,14 +45,6 @@ def generate_quiz(
     questions = [{"question": f"Question {i+1}", "options": ["Option 1", "Option 2", "Option 3"], "answer": "Option 1"} for i in range(num_questions)]
     return questions
 
-@tool
-def retrieve_information(
-    query: Annotated[str, "query to ask the retrieve information tool"]
-    ):
-    """Use Retrieval Augmented Generation to retrieve information about the provided content."""
-    return {"response": "This is a placeholder response for retrieval information."}
-
-
 # Define a function to create agents
 def create_agent(
     llm: ChatOpenAI,
@@ -107,63 +120,48 @@ class AIMSState(TypedDict):
     quiz: List[dict]
 
 
-# Instantiate the language model
-llm = ChatOpenAI(model="gpt-4o")
+# Create the LangGraph chain
+def create_aims_chain(retrieval_chain):
 
-# Create QA Agent
-qa_agent = create_agent(
-    llm,
-    [retrieve_information],  # Existing QA tool
-    "You are a QA assistant who answers questions about the provided notebook content.",
-)
-qa_node = functools.partial(agent_node, agent=qa_agent, name="QAAgent")
-
-# Create Quiz Agent
-quiz_agent = create_agent(
-    llm,
-    [generate_quiz],
-    "You are a quiz creator that generates quizzes based on the provided notebook content.",
-)
-quiz_node = functools.partial(agent_node, agent=quiz_agent, name="QuizAgent")
-
-# Create Supervisor Agent
-supervisor_agent = create_team_supervisor(
-    llm,
-    "You are a supervisor tasked with managing a conversation between the following agents: QAAgent, QuizAgent. Given the user request, decide which agent should act next.",
-    ["QAAgent", "QuizAgent"],
-)
-
-# Build the LangGraph
-aims_graph = StateGraph(AIMSState)
-aims_graph.add_node("QAAgent", qa_node)
-aims_graph.add_node("QuizAgent", quiz_node)
-aims_graph.add_node("supervisor", supervisor_agent)
-
-aims_graph.add_edge("QAAgent", "supervisor")
-aims_graph.add_edge("QuizAgent", "supervisor")
-aims_graph.add_conditional_edges(
-    "supervisor",
-    lambda x: x["next"],
-    {"QAAgent": "QAAgent", "QuizAgent": "QuizAgent", "WAIT": END, "FINISH": END},
-)
-
-aims_graph.set_entry_point("supervisor")
-chain = aims_graph.compile()
-
-if __name__ == "__main__":
-
-    # Define the function to enter the chain
-    def enter_chain(message: str):
-        results = {
-            "messages": [HumanMessage(content="I'd like to take a quiz based on the uploaded notebook.")],
-        }
-        return results
-
-    aims_chain = enter_chain | chain
-
-    for s in aims_chain.stream(
-        "I'd like to take a quiz based on the uploaded notebook.", {"recursion_limit": 15}
-    ):
-        if "__end__" not in s:
-            print(s)
-            print("---")
+    retrieve_information_tool = get_retrieve_information_tool(retrieval_chain)
+
+    # Create QA Agent
+    qa_agent = create_agent(
+        llm,
+        [retrieve_information_tool],
+        "You are a QA assistant who answers questions about the provided notebook content.",
+    )
+
+    qa_node = functools.partial(agent_node, agent=qa_agent, name="QAAgent")
+
+    # Create Quiz Agent
+    quiz_agent = create_agent(
+        llm,
+        [generate_quiz, retrieve_information_tool],
+        "You are a quiz creator that generates quizzes based on the provided notebook content.",
+    )
+    quiz_node = functools.partial(agent_node, agent=quiz_agent, name="QuizAgent")
+
+    # Create Supervisor Agent
+    supervisor_agent = create_team_supervisor(
+        llm,
+        "You are a supervisor tasked with managing a conversation between the following agents: QAAgent, QuizAgent. Given the user request, decide which agent should act next.",
+        ["QAAgent", "QuizAgent"],
+    )
+
+    # Build the LangGraph
+    aims_graph = StateGraph(AIMSState)
+    aims_graph.add_node("QAAgent", qa_node)
+    aims_graph.add_node("QuizAgent", quiz_node)
+    aims_graph.add_node("supervisor", supervisor_agent)
+
+    aims_graph.add_edge("QAAgent", "supervisor")
+    aims_graph.add_edge("QuizAgent", "supervisor")
+    aims_graph.add_conditional_edges(
+        "supervisor",
+        lambda x: x["next"],
+        {"QAAgent": "QAAgent", "QuizAgent": "QuizAgent", "WAIT": END, "FINISH": END},
+    )
+
+    aims_graph.set_entry_point("supervisor")
+    return aims_graph.compile()

aims_tutor/retrieval.py

@@ -42,3 +42,10 @@ class RetrievalManager:
         response = retrieval_augmented_qa_chain.invoke({"question": question})
 
         return response["response"].content
+
+    def get_RAG_QA_chain(self):
+        return (
+            {"context": itemgetter("question") | self.retriever, "question": itemgetter("question")}
+            | RunnablePassthrough.assign(context=itemgetter("context"))
+            | {"response": self.prompts.get_rag_qa_prompt() | self.chat_model, "context": itemgetter("context")}
+        )