Merge pull request #2 from JulsdL/quiz_functionnality

Implementation of Quiz Functionality with LangGraph Integration

CHANGELOG.md

@@ -1,3 +1,18 @@
+version 0.2.0 [2024-05-14]
+
+## Added
+
+- Introduced a comprehensive quiz functionality with LangGraph integration, enabling dynamic quiz generation and question answering based on Jupyter notebook content.
+- Added new Python dependencies (`langgraph==0.0.48`) to support the quiz functionality and improved interaction flow.
+- Implemented a new `graph.py` module to define the quiz and QA agents, along with the supervisor logic for managing conversation flow between agents.
+- Enhanced the `chainlit_frontend.py` to integrate the LangGraph chain, facilitating real-time interaction with the quiz and QA functionality.
+- Updated the `document_processing.py` and `retrieval.py` modules to support the new quiz functionality, including adjustments to the OpenAI model configuration and retrieval logic.
+
+## Modified
+
+- Updated the OpenAI model used in `document_processing.py` from "gpt-4-turbo" to "gpt-4o" to improve the quality of document processing and retrieval.
+- Refined the retrieval logic in `retrieval.py` to include a new method for initializing the RAG QA chain, enhancing the system's ability to provide accurate and contextually relevant answers.
+
 version 0.1.1 [2024-05-13]
 
 ## Modified

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()
@@ -28,6 +30,7 @@ async def start_chat():
         ).send()
 
     file = files[0]  # Get the first file
+
     if file:
         notebook_path = file.path
         doc_manager = DocumentManager(notebook_path)
@@ -36,16 +39,51 @@ 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)
+        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=[], quiz_created=False, question_answered=False)
+
+
+    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}")
+
+        # 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:
+            # Check if the quiz was created and send it to the frontend
+            if state["quiz_created"]:
+                quiz_message = state["messages"][-1].content
+                await cl.Message(content=quiz_message).send()
+            # Check if a question was answered and send the response to the frontend
+            if state["question_answered"]:
+                qa_message = state["messages"][-1].content
+                await cl.Message(content=qa_message).send()
+
+            print("Reached end state.")
 
-    msg = cl.Message(content=response)
-    await msg.send()
+            break

aims_tutor/document_processing.py

@@ -13,7 +13,7 @@ load_dotenv()
 
 # Configuration for OpenAI
 OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
-openai_chat_model = ChatOpenAI(model="gpt-4-turbo", temperature=0.1)
+openai_chat_model = ChatOpenAI(model="gpt-4o", temperature=0.1)
 
 class DocumentManager:
     """
@@ -82,7 +82,7 @@ class DocumentManager:
 
         qdrant_vectorstore = Qdrant.from_documents(split_chunks, embedding_model, location=":memory:", collection_name="Notebook")
 
-        qdrant_retriever = qdrant_vectorstore.as_retriever() # Set the Qdrant vector store as a retriever
+        qdrant_retriever = qdrant_vectorstore.as_retriever()
 
         multiquery_retriever = MultiQueryRetriever.from_llm(retriever=qdrant_retriever, llm=openai_chat_model, include_original=True) # Create a multi-query retriever on top of the Qdrant retriever
 

aims_tutor/graph.py

@@ -0,0 +1,175 @@
+from typing import Annotated, List, TypedDict
+from dotenv import load_dotenv
+from langchain_core.tools import tool
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.messages import AIMessage, BaseMessage
+from langchain.agents import AgentExecutor, create_openai_functions_agent
+from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain_openai import ChatOpenAI
+from langgraph.graph import END, StateGraph
+import functools
+
+# 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)
+
+# Function to create agents
+def create_agent(
+    llm: ChatOpenAI,
+    tools: list,
+    system_prompt: str,
+) -> AgentExecutor:
+    """Create a function-calling agent and add it to the graph."""
+    system_prompt += "\nWork autonomously according to your specialty, using the tools available to you."
+    " Do not ask for clarification."
+    " Your other team members (and other teams) will collaborate with you with their own specialties."
+    " You are chosen for a reason! You are one of the following team members: {team_members}."
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                system_prompt,
+            ),
+            MessagesPlaceholder(variable_name="messages"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ]
+    )
+    agent = create_openai_functions_agent(llm, tools, prompt)
+    executor = AgentExecutor(agent=agent, tools=tools, handle_parsing_errors=True)
+    return executor
+
+# Function to create agent nodes
+def agent_node(state, agent, name):
+    result = agent.invoke(state)
+    if 'messages' not in result: # Check if messages are present in the agent state
+        raise ValueError(f"No messages found in agent state: {result}")
+    new_state = {"messages": state["messages"] + [AIMessage(content=result["output"], name=name)]}
+    if "next" in result:
+        new_state["next"] = result["next"]
+    if name == "QuizAgent" and "quiz_created" in state and not state["quiz_created"]:
+        new_state["quiz_created"] = True
+        new_state["next"] = "FINISH" # Finish the conversation after the quiz is created and wait for a new user input
+    if name == "QAAgent":
+        new_state["question_answered"] = True
+        new_state["next"] = "question_answered"
+    return new_state
+
+
+# Function to create the supervisor
+def create_team_supervisor(llm: ChatOpenAI, system_prompt, members) -> AgentExecutor:
+    """An LLM-based router."""
+    options = ["WAIT", "FINISH"] + members
+    function_def = {
+        "name": "route",
+        "description": "Select the next role.",
+        "parameters": {
+            "title": "routeSchema",
+            "type": "object",
+            "properties": {
+                "next": {
+                    "title": "Next",
+                    "anyOf": [
+                        {"enum": options},
+                    ],
+                },
+            },
+            "required": ["next"],
+        },
+    }
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            MessagesPlaceholder(variable_name="messages"),
+            (
+                "system",
+                "Given the conversation above, who should act next?"
+                " Or should we WAIT for user input? Select one of: {options}",
+            ),
+        ]
+    ).partial(options=str(options), team_members=", ".join(members))
+    return (
+        prompt
+        | llm.bind_functions(functions=[function_def], function_call="route")
+        | JsonOutputFunctionsParser()
+    )
+
+# Define the state for the system
+class AIMSState(TypedDict):
+    messages: List[BaseMessage]
+    next: str
+    quiz: List[dict]
+    quiz_created: bool
+    question_answered: bool
+
+
+# Create the LangGraph chain
+def create_aims_chain(retrieval_chain):
+
+    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,
+        [retrieve_information_tool],
+        "You are a quiz creator that generates quizzes based on the provided notebook content."
+
+        """First, You MUST Use the retrieval_inforation_tool to gather context from the notebook to gather relevant and accurate information.
+
+        Next, create a 5-question quiz based on the information you have gathered. Include the answers at the end of the quiz.
+
+        Present the quiz to the user in a clear and concise manner."""
+    )
+
+    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: "FINISH" if x.get("quiz_created") else ("FINISH" if x.get("question_answered") else x["next"]),
+        {"QAAgent": "QAAgent", "QuizAgent": "QuizAgent", "WAIT": END, "FINISH": END, "question_answered": 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")}
+        )

requirements.txt

@@ -1,4 +1,5 @@
 langchain==0.1.20
+langgraph==0.0.48
 crewai==0.30.0
 qdrant-client==1.9.1
 python-dotenv==1.0.1