Refactor chainlit_frontend.py to streamline message processing and integrate new agent states

- Reduce recursion limit in tutor_chain.stream to improve performance
- Simplify message processing logic by removing conditional checks for end state and directly processing agent states
- Implement direct handling for QAAgent, QuizAgent, and FlashcardsAgent with specific actions for each agent type
- Update flashcards creation flow to use a dynamic path stored in agent state, enhancing flexibility
- Replace langchain.pydantic_v1 import with pydantic in tools.py for BaseModel and Field, aligning with updated dependencies
- Add explicit directory creation for flashcard files to ensure path validity
- Introduce debug prints across chainlit_frontend.py, tools.py, and graph.py for better traceability of the processing flow
- Adjust states.py to replace flashcard_filename with flashcard_path, reflecting the new method of handling flashcard file paths

notebook_tutor/chainlit_frontend.py

@@ -78,52 +78,43 @@ async def main(message: cl.Message):
     print("\033[93m" + f"Initial state: {state}" + "\033[0m")
 
     # Process the message through the LangGraph chain
-    for s in tutor_chain.stream(state, {"recursion_limit": 10}):
+    for s in tutor_chain.stream(state, {"recursion_limit": 3}):
         print("\033[93m" + f"State after processing: {s}" + "\033[0m")
 
-        # 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("\033[93m" + f"Response: {response}" + "\033[0m")
-                await cl.Message(content=response).send()
-            else:
-                print("Error: No messages found in agent state.")
-        else:
-            # Extract the final state
-            final_state = next(iter(s.values()))
-            print("\033[93m" + f"Final state: {final_state}" + "\033[0m")
-
-            # Check if the quiz was created and send it to the frontend
-            if final_state.get("quiz_created"):
-                quiz_message = final_state["messages"][-1].content
-                await cl.Message(content=quiz_message).send()
+        agent_state = next(iter(s.values()))
+        print("\033[93m" + f"Agent state: {agent_state}" + "\033[0m")
 
-            # Check if a question was answered and send the response to the frontend
-            if final_state.get("question_answered"):
-                qa_message = final_state["messages"][-1].content
+        if "QAAgent" in s:
+            if s['QAAgent']['question_answered']:
+                print("\033[93m" + "************************Question answered**********************." + "\033[0m")
+                qa_message = agent_state["messages"][-1].content
                 await cl.Message(content=qa_message).send()
 
-            # 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
+        if "QuizAgent" in s:
+            if s['QuizAgent']['quiz_created']:
+                print("\033[93m" + "************************Quiz created**********************." + "\033[0m")
+                quiz_message = agent_state["messages"][-1].content
+                await cl.Message(content=quiz_message).send()
+
+        if "FlashcardsAgent" in s:
+            if s['FlashcardsAgent']['flashcards_created']:
+                print("\033[93m" + "************************Flashcards created**********************." + "\033[0m")
+                flashcards_message = agent_state["messages"][-1].content
                 await cl.Message(content=flashcards_message).send()
 
-                # Create a relative path to the file
-                flashcard_filename = final_state["flashcard_filename"]
-                print("\033[93m" + f"Flashcard filename: {flashcard_filename}" + "\033[0m")
-                flashcard_path = os.path.join(".files", flashcard_filename)
+                flashcard_path = agent_state["flashcard_path"]
                 print("\033[93m" + f"Flashcard path: {flashcard_path}" + "\033[0m")
 
+
                 # Use the File class to send the file
-                file_element = cl.File(name=os.path.basename(flashcard_path), path=flashcard_path)
+                file_element = cl.File(name="Flashcards", path=flashcard_path)
                 print("\033[93m" + f"Sending flashcards file: {file_element}" + "\033[0m")
                 await cl.Message(
                     content="Here are your flashcards:",
                     elements=[file_element]
                 ).send()
 
-            print("\033[93m" + "Reached END state." + "\033[0m")
+        final_state = s  # Save the final state after processing
+        print("\033[93m" + f"Final state: {final_state}" + "\033[0m")
 
-            break
+    print("\033[93m" + "Reached END state." + "\033[0m")

notebook_tutor/graph.py

@@ -69,6 +69,8 @@ def agent_node(state, agent, name):
         raise ValueError(f"No messages found in agent state: {result}")
     new_state = {"messages": state["messages"] + [AIMessage(content=result["output"], name=name)]}
 
+    print("\033[93m" + f"agent_node function state {state}" + "\033[0m")
+
     # Set the appropriate flags and next state
     if name == "QuizAgent":
         new_state["quiz_created"] = True
@@ -76,7 +78,8 @@ def agent_node(state, agent, name):
         new_state["question_answered"] = True
     elif name == "FlashcardsAgent":
         new_state["flashcards_created"] = True
-        new_state["flashcard_filename"] = result["output"].split('(')[-1].strip(')')
+        print("\033[93m" + f"agent_node function result_output {result}" + "\033[0m")
+        # new_state["flashcard_path"] = result["output"]
 
     new_state["next"] = "FINISH"
     return new_state

notebook_tutor/states.py

@@ -9,4 +9,4 @@ class TutorState(TypedDict):
     quiz_created: bool
     question_answered: bool
     flashcards_created: bool
-    flashcard_filename: str
+    flashcard_path: str

notebook_tutor/tools.py

@@ -1,5 +1,5 @@
 from typing import Optional, Type
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 from langchain.callbacks.manager import (
     AsyncCallbackManagerForToolRun,
@@ -23,7 +23,9 @@ class FlashcardTool(BaseTool):
         """Use the tool to create flashcards."""
         filename = f"flashcards_{uuid.uuid4()}.csv"
         save_path = os.path.join('.files', filename)
-        # os.makedirs(os.path.dirname(save_path), exist_ok=True) # Create directory if it doesn't exist
+        print("\033[91m" + f"Saving flashcards to {save_path}" + "\033[0m")
+
+        os.makedirs(os.path.dirname(save_path), exist_ok=True) # Create directory if it doesn't exist
         with open(save_path, 'w', newline='') as csvfile:
             fieldnames = ['Front', 'Back']
             writer = csv.DictWriter(csvfile, fieldnames=fieldnames)