Implement Flashcard creation tool and update project naming

- Added FlashcardTool in notebook_tutor/tools.py for generating flashcards in CSV format, suitable for Anki import.
- Introduced FlashcardInput model to define the structure of flashcards input.
- Created TutorState in notebook_tutor/states.py to manage application states including flashcards creation status.
- Updated README.md to reflect the project's new name, AI-Notebook-Tutor, and updated instructions for running the application.
- Removed obsolete main.py

README.md

@@ -1,8 +1,8 @@
-# AIMS-Tutor
+# AI-Notebook-Tutor
 
 # RAG Application for QA in Jupyter Notebook
 
-AIMS-Tutor is designed to provide question-answering capabilities in a Jupyter Notebook using the Retrieval Augmented Generation (RAG) model. It's built on top of the LangChain and Chainlit platforms, and it uses the OpenAI API for the chat model.
+AI-Notebook-Tutor is designed to provide question-answering capabilities in a Jupyter Notebook using the Retrieval Augmented Generation (RAG) model. It's built on top of the LangChain and Chainlit platforms, and it uses the OpenAI API for the chat model.
 
 ## Features
 
@@ -28,7 +28,7 @@ OPENAI_API_KEY=your-key-here
 4. Run the application using the following command:
 
 ```bash
-chainlit run aims_tutor/app.py
+chainlit run notebook_tutor/app.py
 ```
 
 ## Usage

flashcards/flashcards_03fb423e-2087-4598-9ebb-b99a20db0b93.csv

@@ -0,0 +1,11 @@
+Front,Back
+What is the first step in the retriever creation process?,The first step is loading a fine-tuned embedding model from the hub using the HuggingFaceEmbeddings class. The model used is named 'JulsdL/e2erag-arctic-m' and it is configured to run on a CUDA device.
+Which class is used to load the embedding model in the retriever creation process?,The HuggingFaceEmbeddings class is used to load the embedding model.
+What is the name of the embedding model used in the retriever creation process?,The embedding model used is named 'JulsdL/e2erag-arctic-m'.
+On which device is the embedding model 'JulsdL/e2erag-arctic-m' configured to run?,The embedding model is configured to run on a CUDA device.
+What is the purpose of setting up a VectorStore in the retriever creation process?,The VectorStore is set up to power the dense vector search and is populated with documents that have been split into chunks and embedded using the embedding model.
+Which tool is used to set up the VectorStore in the retriever creation process?,Meta's FAISS (Facebook AI Similarity Search) is used to set up the VectorStore.
+How are documents prepared for the VectorStore in the retriever creation process?,Documents are split into chunks and embedded using the previously loaded embedding model before being added to the VectorStore.
+What is the final step in the retriever creation process?,The final step is converting the VectorStore into a retriever that can fetch relevant documents or chunks based on query embeddings.
+"In the context of retriever creation, what is the purpose of converting the VectorStore to a retriever?",Converting the VectorStore to a retriever allows it to be used for fetching relevant documents or chunks based on the query embeddings.
+How does a retriever enhance the model's ability in a Retrieval-Augmented Generation (RAG) setup?,"A retriever enhances the model's ability by providing contextually relevant responses based on the input queries, improving the quality and relevance of the generated content."

flashcards/flashcards_30a9e627-3d47-45fc-81db-a991269edb24.csv

@@ -0,0 +1,11 @@
+Front,Back
+What command is used to clone a repository from GitHub?,!git clone https://github.com/arcee-ai/DALM
+How do you install a package using pip and upgrade it if necessary?,!pip install --upgrade -q -e .
+"Which command is used to install the latest versions of langchain, langchain-core, langchain-community, and sentence_transformers?",!pip install -qU langchain langchain-core langchain-community sentence_transformers
+How do you install the pymupdf and faiss-cpu libraries using pip?,!pip install -qU pymupdf faiss-cpu
+What is the import statement for pandas?,import pandas as pd
+How do you import HuggingFaceEmbeddings from the langchain_community library?,from langchain_community.embeddings import HuggingFaceEmbeddings
+What is the import statement for FAISS from the langchain_community library?,from langchain_community.vectorstores import FAISS
+Which import statement is used for SimpleDirectoryReader from llama_index.core?,from llama_index.core import SimpleDirectoryReader
+How do you import SimpleNodeParser from the llama_index.core.node_parser module?,from llama_index.core.node_parser import SimpleNodeParser
+What is the import statement for MetadataMode from llama_index.core.schema?,from llama_index.core.schema import MetadataMode

flashcards/flashcards_37136597-b97c-4889-9dc1-c3b5e10084c1.csv

@@ -0,0 +1,11 @@
+Front,Back
+What is the first step in the retriever creation process?,The first step in the retriever creation process is loading the embedding model.
+Which model is used as the embedding model in the retriever creation process?,The embedding model used is 'JulsdL/e2erag-arctic-m'.
+Which class and module are used to load the embedding model?,The embedding model is loaded using the HuggingFaceEmbeddings class from the langchain_community.embeddings module.
+On which device is the embedding model configured to run?,The embedding model is configured to run on a CUDA device.
+What is the purpose of the VectorStore in the retriever creation process?,The VectorStore is created to power dense vector searches.
+Which technology is used to set up the VectorStore?,Meta's FAISS (Facebook AI Similarity Search) is used to set up the VectorStore.
+How are the documents prepared for creating the VectorStore?,The documents are split into chunks before creating the VectorStore.
+Which method is called to create the VectorStore from the documents and embedding model?,The method FAISS.from_documents is called to create the VectorStore.
+How is the VectorStore converted into a retriever?,The VectorStore is converted into a retriever by invoking the as_retriever() method on the vector_store object.
+What are the benefits of the retriever created through this process?,"The retriever leverages the power of dense embeddings and efficient search capabilities provided by FAISS, making it effective for retrieval-augmented tasks."

flashcards/flashcards_4dbaa0d1-6363-4584-9691-824540220571.csv

@@ -0,0 +1,11 @@
+Front,Back
+What is the first step in creating a retriever for a document?,The first step is loading the documents using a loader like `PyMuPDFLoader` to load documents from a PDF file.
+Which function is used to load documents from a PDF file?,The function `PyMuPDFLoader` is used to load documents from a PDF file.
+How are documents split into chunks for processing?,"Documents are split into chunks using the `RecursiveCharacterTextSplitter`, which divides the documents based on predefined rules like token limits and split characters."
+What is the purpose of chunking documents?,Chunking documents into smaller parts helps in managing and processing them more efficiently for embedding and retrieval.
+Which library is used for loading a pre-trained embedding model?,The `HuggingFaceEmbeddings` library is used for loading a pre-trained embedding model.
+How do you specify the device for running the embedding model?,"You specify the device (e.g., 'cuda') in the `model_kwargs` parameter when loading the embedding model."
+What is Meta's FAISS used for in the retriever creation process?,Meta's FAISS is used to power dense vector search by creating a vector store from the document chunks and embedding model.
+How do you convert a vector store into a retriever?,You convert a vector store into a retriever using the `as_retriever` method.
+What is the role of the retriever in the retrieval-augmented generation (RAG) system?,"The retriever fetches relevant document chunks based on query vectors, providing contextually relevant responses."
+Which embedding model is used in the example provided for retriever creation?,The embedding model 'JulsdL/e2erag-arctic-m' from Hugging Face is used in the example.

flashcards/flashcards_52137f3d-8a8d-4891-9216-ab3bbc6ee66a.csv

@@ -0,0 +1,11 @@
+Front,Back
+What library is used to load the fine-tuned embedding model for retriever creation?,The HuggingFaceEmbeddings class is used to load the fine-tuned embedding model.
+What model is loaded for embedding in the retriever creation process?,The model 'JulsdL/e2erag-arctic-m' is loaded for embedding.
+On which device is the embedding model set to run?,The embedding model is set to run on a CUDA device.
+Which class is used to set up the vector store in retriever creation?,The FAISS class from Meta is used to set up the vector store.
+What is the purpose of the vector store in the retriever creation process?,The vector store powers the dense vector search by storing documents that have been split into chunks and embedded.
+How is the vector store created in the retriever creation process?,The vector store is created from documents that have been split into chunks and embedded using the loaded embedding model.
+How is the vector store converted into a retriever?,The vector store is converted into a retriever using the 'as_retriever()' method.
+What is a retriever used for in the context of retriever creation?,A retriever is used to retrieve context based on a query for a language model.
+What is the benefit of creating a retriever in a RAG setup?,Creating a retriever enhances the capabilities of language models by providing them with relevant context for generating responses.
+What does 'RAG' stand for in the context of retriever creation?,'RAG' stands for Retrieval-Augmented Generation.

flashcards/flashcards_545f9f54-f8b0-4549-b27d-66b7303a017e.csv

@@ -0,0 +1,11 @@
+Front,Back
+What tool is used to load documents from a PDF file in the retriever creation process?,The PyMuPDFLoader is used to load documents from a PDF file.
+What is the purpose of the RecursiveCharacterTextSplitter in the retriever creation process?,The RecursiveCharacterTextSplitter is used to divide the documents into manageable chunks based on specific rules such as token limits and preferred split characters.
+Which library is used to create a dense vector store in the retriever creation process?,Meta's FAISS library is used to create a dense vector store from the document chunks and the embedding model.
+How do you load a pre-trained embedding model from Hugging Face to run on a CUDA device?,Use the HuggingFaceEmbeddings class with the model name and specify the device as CUDA in the model_kwargs parameter.
+What is the role of the embedding model in the retriever creation process?,The embedding model is used to convert text into vectors that can be stored in the vector store for efficient retrieval.
+How do you convert a vector store into a retriever?,Use the as_retriever method on the vector store to convert it into a retriever.
+What is the first step in the retriever creation process?,The first step is loading the documents using PyMuPDFLoader.
+Which model name is used in the provided example for the embedding model?,"The model name used is ""JulsdL/e2erag-arctic-m""."
+What is the purpose of the vector store in the retriever creation process?,"The vector store holds the dense vector representations of document chunks, allowing for efficient retrieval based on query embeddings."
+Which Python import is necessary to use FAISS for creating a vector store?,You need to import FAISS from langchain_community.vectorstores.

flashcards/flashcards_af3650ed-34b5-4040-b672-036e1cf3b8e3.csv

@@ -0,0 +1,11 @@
+Front,Back
+What is the first step in creating a retriever?,The first step in creating a retriever is loading a fine-tuned embedding model from the hub using the HuggingFaceEmbeddings class.
+Which class is used to load the embedding model?,The HuggingFaceEmbeddings class is used to load the embedding model.
+How is the model 'JulsdL/e2erag-arctic-m' loaded onto the GPU?,"The model 'JulsdL/e2erag-arctic-m' is loaded onto the GPU by setting the model_kwargs parameter to {""device"": ""cuda""} in the HuggingFaceEmbeddings class."
+What library is used to set up the vector store?,Meta's FAISS library is used to set up the vector store.
+What does the vector store manage?,The vector store manages the dense vectors generated by the embedding model.
+How is the vector store converted into a retriever?,The vector store is converted into a retriever using the as_retriever() method.
+What is the purpose of the retriever in this context?,The purpose of the retriever is to efficiently fetch relevant document vectors based on the query vectors.
+What is the role of the embedding model in the retriever creation process?,"The embedding model generates dense vectors for the documents, which are then managed by the vector store and used by the retriever to find relevant information."
+Which parameter in the HuggingFaceEmbeddings class specifies the device to be used?,"The model_kwargs parameter specifies the device to be used, such as ""cuda"" for GPU."
+What does the FAISS library stand for?,FAISS stands for Facebook AI Similarity Search.

main.py

@@ -1,118 +0,0 @@
-import os
-from operator import itemgetter
-
-import chainlit as cl
-import tiktoken
-from dotenv import load_dotenv
-
-
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.retrievers import MultiQueryRetriever
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnablePassthrough
-from langchain_community.document_loaders import PyMuPDFLoader, PythonLoader, NotebookLoader
-from langchain_community.vectorstores import Qdrant
-from langchain_openai import ChatOpenAI
-from langchain_openai.embeddings import OpenAIEmbeddings
-
-# Load environment variables
-load_dotenv()
-
-# Configuration for OpenAI
-OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
-openai_chat_model = ChatOpenAI(model="gpt-4-turbo", temperature=0.1)
-
-# Define the RAG prompt
-RAG_PROMPT = """
-CONTEXT:
-{context}
-
-QUERY:
-{question}
-
-Answer the query in a pretty format if the context is related to it; otherwise, answer: 'Sorry, I can't answer.'
-"""
-rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
-
-
-# ChainLit setup for chat interaction
-@cl.on_chat_start
-async def start_chat():
-    settings = {
-        "model": "gpt-3.5-turbo",
-        "temperature": 0,
-        "top_p": 1,
-        "frequency_penalty": 0,
-        "presence_penalty": 0,
-    }
-    cl.user_session.set("settings", settings)
-
-    # Display a welcoming message with instructions
-    welcome_message = "Welcome to the AIMS-Tutor! Please upload a Jupyter notebook (.ipynb and max. 5mb) to start."
-    await cl.Message(content=welcome_message).send()
-
-    # Wait for the user to upload a file
-    files = None
-    while files is None:
-        files = await cl.AskFileMessage(
-            content="Please upload a Jupyter notebook (.ipynb, max. 5mb):",
-            accept={"application/x-ipynb+json": [".ipynb"]},
-            max_size_mb=5
-        ).send()
-
-    file = files[0] # Get the first file
-
-    if file:
-        # Load the Jupyter notebook
-        notebook_path = file.path # Extract the path from the AskFileResponse object
-
-        loader = NotebookLoader(
-            notebook_path,
-            include_outputs=False,
-            max_output_length=20,
-            remove_newline=True,
-            traceback=False
-        )
-        docs = loader.load()
-        cl.user_session.set("docs", docs) # Store the docs in the user session
-
-        # Initialize the retriever components after loading document
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50, length_function=tiktoken_len) # Initialize the text splitter
-        split_chunks = text_splitter.split_documents(docs) # Split the documents into chunks
-        embedding_model = OpenAIEmbeddings(model="text-embedding-3-small") # Initialize the embedding model
-        qdrant_vectorstore = Qdrant.from_documents(split_chunks, embedding_model, location=":memory:", collection_name="Notebook") # Create a Qdrant vector store
-        qdrant_retriever = qdrant_vectorstore.as_retriever() # Set the Qdrant vector store as a 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
-
-        # Store the multiquery_retriever in the user session
-        cl.user_session.set("multiquery_retriever", multiquery_retriever)
-
-
-@cl.on_message
-async def main(message: cl.Message):
-    # Retrieve the multi-query retriever from session
-    multiquery_retriever = cl.user_session.get("multiquery_retriever")
-    if not multiquery_retriever:
-        await cl.Message(content="No document processing setup found. Please upload a Jupyter notebook first.").send()
-        return
-
-    question = message.content
-    response = handle_query(question, multiquery_retriever)  # Process the question
-
-    msg = cl.Message(content=response)
-    await msg.send()
-
-def handle_query(question, retriever):
-    # Define the retrieval augmented query-answering chain
-    retrieval_augmented_qa_chain = (
-        {"context": itemgetter("question") | retriever, "question": itemgetter("question")}
-        | RunnablePassthrough.assign(context=itemgetter("context"))
-        | {"response": rag_prompt | openai_chat_model, "context": itemgetter("context")}
-    )
-    response = retrieval_augmented_qa_chain.invoke({"question": question})
-    return response["response"].content
-
-# Tokenization function
-def tiktoken_len(text):
-    tokens = tiktoken.encoding_for_model("gpt-3.5-turbo").encode(text)
-    return len(tokens)

{aims_tutor → notebook_tutor}/app.py

@@ -1,6 +1,6 @@
 import os
 from dotenv import load_dotenv
-import aims_tutor.chainlit_frontend as cl_frontend
+import notebook_tutor.chainlit_frontend as cl_frontend
 
 # Load environment variables
 load_dotenv()

{aims_tutor → notebook_tutor}/chainlit_frontend.py

@@ -1,13 +1,20 @@
+import os
+import logging
 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
+from graph import create_tutor_chain, TutorState
 
 # Load environment variables
 load_dotenv()
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+
+logger = logging.getLogger(__name__)
+
 @cl.on_chat_start
 async def start_chat():
     settings = {
@@ -18,7 +25,7 @@ async def start_chat():
         "presence_penalty": 0,
     }
     cl.user_session.set("settings", settings)
-    welcome_message = "Welcome to the AIMS-Tutor! Please upload a Jupyter notebook (.ipynb and max. 5mb) to start."
+    welcome_message = "Welcome to the Notebook-Tutor! Please upload a Jupyter notebook (.ipynb and max. 5mb) to start."
     await cl.Message(content=welcome_message).send()
 
     files = None
@@ -42,27 +49,36 @@ async def start_chat():
         # 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)
+        tutor_chain = create_tutor_chain(retrieval_chain)
+        cl.user_session.set("tutor_chain", tutor_chain)
+
+        logger.info("Chat started and notebook uploaded successfully.")
 
 @cl.on_message
 async def main(message: cl.Message):
     # Retrieve the LangGraph chain from the session
-    aims_chain = cl.user_session.get("aims_chain")
+    tutor_chain = cl.user_session.get("tutor_chain")
 
-    if not aims_chain:
+    if not tutor_chain:
         await cl.Message(content="No document processing setup found. Please upload a Jupyter notebook first.").send()
         return
 
     # 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)
-
+    state = TutorState(
+        messages=[HumanMessage(content=user_message)],
+        next="supervisor",
+        quiz=[],
+        quiz_created=False,
+        question_answered=False,
+        flashcards_created=False,
+        flashcard_filename="",
+    )
 
     print(f"Initial state: {state}")
 
     # Process the message through the LangGraph chain
-    for s in aims_chain.stream(state, {"recursion_limit": 10}):
+    for s in tutor_chain.stream(state, {"recursion_limit": 10}):
         print(f"State after processing: {s}")
 
         # Extract messages from the state
@@ -75,15 +91,38 @@ async def main(message: cl.Message):
             else:
                 print("Error: No messages found in agent state.")
         else:
+            # Extract the final state
+            final_state = next(iter(s.values()))
+
             # Check if the quiz was created and send it to the frontend
-            if state["quiz_created"]:
-                quiz_message = state["messages"][-1].content
+            if final_state.get("quiz_created"):
+                quiz_message = final_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
+            if final_state.get("question_answered"):
+                qa_message = final_state["messages"][-1].content
                 await cl.Message(content=qa_message).send()
 
-            print("Reached end state.")
+            # Check if flashcards are ready and send the file to the frontend
+            if final_state.get("flashcards_created"):
+                flashcards_message = final_state["messages"][-1].content
+                await cl.Message(content=flashcards_message).send()
+
+                # Create a full path to the file
+                flashcard_filename = final_state["flashcard_filename"]
+                print(f"Flashcard filename: {flashcard_filename}")
+                flashcard_path = os.path.abspath(flashcard_filename)
+                print(f"Flashcard path: {flashcard_path}")
+
+                # Use the File class to send the file
+                file_element = cl.File(name=os.path.basename(flashcard_filename), path=flashcard_path)
+                print(f"Sending flashcards file: {file_element}")
+                await cl.Message(
+                    content="Here are your flashcards:",
+                    elements=[file_element]
+                ).send()
+
+            print("Reached END state.")
 
             break

{aims_tutor → notebook_tutor}/document_processing.py

@@ -6,7 +6,7 @@ from langchain.retrievers import MultiQueryRetriever
 from langchain_openai.embeddings import OpenAIEmbeddings
 from langchain_openai import ChatOpenAI
 from dotenv import load_dotenv
-from aims_tutor.utils import tiktoken_len
+from notebook_tutor.utils import tiktoken_len
 
 # Load environment variables
 load_dotenv()

{aims_tutor → notebook_tutor}/graph.py

@@ -1,12 +1,14 @@
-from typing import Annotated, List, TypedDict
+from typing import Annotated
 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_core.messages import AIMessage
 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
+from tools import create_flashcards_tool
+from states import TutorState
 import functools
 
 # Load environment variables
@@ -32,6 +34,9 @@ def get_retrieve_information_tool(retrieval_chain):
     wrapper_instance = RetrievalChainWrapper(retrieval_chain)
     return tool(wrapper_instance.retrieve_information)
 
+# Instantiate the tools
+flashcard_tool = create_flashcards_tool
+
 # Function to create agents
 def create_agent(
     llm: ChatOpenAI,
@@ -60,20 +65,24 @@ def create_agent(
 # 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
+    if 'messages' not in result:
         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"]:
+
+    # Set the appropriate flags and next state
+    if name == "QuizAgent":
         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":
+    elif name == "QAAgent":
         new_state["question_answered"] = True
-        new_state["next"] = "question_answered"
+    elif name == "FlashcardsAgent":
+        new_state["flashcards_created"] = True
+        new_state["flashcard_filename"] = result["output"].split('(')[-1].strip(')')
+
+    new_state["next"] = "FINISH"
     return new_state
 
 
+
 # Function to create the supervisor
 def create_team_supervisor(llm: ChatOpenAI, system_prompt, members) -> AgentExecutor:
     """An LLM-based router."""
@@ -112,17 +121,9 @@ def create_team_supervisor(llm: ChatOpenAI, system_prompt, members) -> AgentExec
         | 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):
+def create_tutor_chain(retrieval_chain):
 
     retrieve_information_tool = get_retrieve_information_tool(retrieval_chain)
 
@@ -139,37 +140,55 @@ def create_aims_chain(retrieval_chain):
     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.
-
+        """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 Flashcards Agent
+    flashcards_agent = create_agent(
+        llm,
+        [retrieve_information_tool, flashcard_tool],
+        """
+        You are the Flashcard creator. Your mission is to create effective and concise flashcards based on the user's query and the content of the provided notebook. Your role involves the following tasks:
+        1. Analyze User Query: Understand the user's request and determine the key concepts and information they need to learn.
+        2. Search Notebook Content: Use the notebook content to gather relevant information and generate accurate and informative flashcards.
+        3. Generate Flashcards: Create a series of flashcards content with clear questions on the front and detailed answers on the back. Ensure that the flashcards cover the essential points and concepts requested by the user.
+        4. Export Flashcards: Use the flashcard_tool to create and export the flashcards in a format that can be easily imported into a flashcard management system, such as Anki.
+
+        Remember, your goal is to help the user learn efficiently and effectively by breaking down the notebook content into manageable, repeatable flashcards."""
+    )
+
+    flashcards_node = functools.partial(agent_node, agent=flashcards_agent, name="FlashcardsAgent")
+
     # 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"],
+        "You are a supervisor tasked with managing a conversation between the following agents: QAAgent, QuizAgent, FlashcardsAgent. Given the user request, decide which agent should act next.",
+        ["QAAgent", "QuizAgent", "FlashcardsAgent"],
     )
 
     # 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(
+    tutor_graph = StateGraph(TutorState)
+    tutor_graph.add_node("QAAgent", qa_node)
+    tutor_graph.add_node("QuizAgent", quiz_node)
+    tutor_graph.add_node("FlashcardsAgent", flashcards_node)
+    tutor_graph.add_node("supervisor", supervisor_agent)
+
+    tutor_graph.add_edge("QAAgent", "supervisor")
+    tutor_graph.add_edge("QuizAgent", "supervisor")
+    tutor_graph.add_edge("FlashcardsAgent", "supervisor")
+    tutor_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},
+        lambda x: "FINISH" if x.get("quiz_created") or x.get("question_answered") or x.get("flashcards_created") else x["next"],
+        {"QAAgent": "QAAgent",
+        "QuizAgent": "QuizAgent",
+        "FlashcardsAgent": "FlashcardsAgent",
+        "FINISH": END},
     )
 
-    aims_graph.set_entry_point("supervisor")
-    return aims_graph.compile()
+    tutor_graph.set_entry_point("supervisor")
+    return tutor_graph.compile()

notebook_tutor/states.py

@@ -0,0 +1,12 @@
+from typing import List, TypedDict
+from langchain_core.messages import BaseMessage
+
+# Define the state for the system
+class TutorState(TypedDict):
+    messages: List[BaseMessage]
+    next: str
+    quiz: List[dict]
+    quiz_created: bool
+    question_answered: bool
+    flashcards_created: bool
+    flashcard_filename: str

notebook_tutor/tools.py

@@ -0,0 +1,43 @@
+from typing import Optional, Type
+from langchain.pydantic_v1 import BaseModel, Field
+from langchain.tools import BaseTool
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForToolRun,
+    CallbackManagerForToolRun,
+)
+import csv
+import uuid
+import os
+
+class FlashcardInput(BaseModel):
+    flashcards: list = Field(description="A list of flashcards. Each flashcard should be a dictionary with 'question' and 'answer' keys.")
+
+class FlashcardTool(BaseTool):
+    name = "create_flashcards"
+    description = "Create flashcards in a .csv format suitable for import into Anki"
+    args_schema: Type[BaseModel] = FlashcardInput
+
+    def _run(
+        self, flashcards: list, run_manager: Optional[CallbackManagerForToolRun] = None
+    ) -> str:
+        """Use the tool to create flashcards."""
+        filename = f"flashcards_{uuid.uuid4()}.csv"
+        save_path = os.path.join('flashcards', filename)  # Save in 'flashcards' directory
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        with open(save_path, 'w', newline='') as csvfile:
+            fieldnames = ['Front', 'Back']
+            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+
+            writer.writeheader()
+            for card in flashcards:
+                writer.writerow({'Front': card['question'], 'Back': card['answer']})
+        return save_path
+
+    async def _arun(
+        self, flashcards: list, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
+    ) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError("create_flashcards does not support async")
+
+# Instantiate the tool
+create_flashcards_tool = FlashcardTool()