Added a call method to the unfinished adalflow.components.memory.memory.Memory class. This method returns a formatted conversation history, handles empty conversations, and formats user and assistant messages.Additionally, added add_dialog_turn() method to handle new back-and-forth messages. This adds an id to each dialog turn for tracking and appends it to the current_conversation list. Finally, added 3 new tests for memory -empty memory, single dialog turn, and multiple dialog turn tests. Also updated the chatbot usage notebook to replace the deprecated class with the updated one.

lu-ny · lu-ny · commit 5b699b5a460f · 2024-11-28T13:59:48.000-05:00
diff --git a/adalflow/adalflow/components/memory/memory.py b/adalflow/adalflow/components/memory/memory.py
@@ -1,23 +1,77 @@
-"""Memory for user-assistant conversations. [Not completed]
+"""Memory component for user-assistant conversations.
 
 Memory can include data modeling, in-memory data storage, local file data storage, cloud data persistence, data pipeline, data retriever.
 It is itself an LLM application and different use cases can do it differently.
 
+This component handles the storage and retrieval of conversation history between users
+and assistants. It provides local memory experience with the ability to format and
+return conversation history.
 
-This implementation covers the minimal and local memory experience for the user-assistant conversation.
+Attributes:
+    current_conversation (Conversation): Stores the current active conversation.
+    turn_db (LocalDB): Database for storing all conversation turns.
+    conver_db (LocalDB): Database for storing complete conversations.
 """
 
+from uuid import uuid4
+from adalflow.core.component import Component
+from adalflow.core.db import LocalDB
 from adalflow.core.types import (
     Conversation,
+    DialogTurn,
+    UserQuery,
+    AssistantResponse,
 )
 
-from adalflow.core.db import LocalDB
-from adalflow.core.component import Component
-
 
 class Memory(Component):
     def __init__(self, turn_db: LocalDB = None):
+        """Initialize the Memory component.
+
+        Args:
+            turn_db (LocalDB, optional): Database for storing conversation turns.
+                Defaults to None, in which case a new LocalDB is created.
+        """
         super().__init__()
-        self.current_convesation = Conversation()
+        self.current_conversation = Conversation()
         self.turn_db = turn_db or LocalDB()  # all turns
         self.conver_db = LocalDB()  # a list of conversations
+
+    def call(self) -> str:
+        """Returns the current conversation history as a formatted string.
+
+        Returns:
+            str: Formatted conversation history with alternating user and assistant messages.
+                Returns empty string if no conversation history exists.
+        """
+        if not self.current_conversation.dialog_turns:
+            return ""
+
+        formatted_history = []
+        for turn in self.current_conversation.dialog_turns.values():
+            formatted_history.extend(
+                [
+                    f"User: {turn.user_query.query_str}",
+                    f"Assistant: {turn.assistant_response.response_str}",
+                ]
+            )
+        return "\n".join(formatted_history)
+
+    def add_dialog_turn(self, user_query: str, assistant_response: str):
+        """Add a new dialog turn to the current conversation.
+
+        Args:
+            user_query (str): The user's input message.
+            assistant_response (str): The assistant's response message.
+        """
+        dialog_turn = DialogTurn(
+            id=str(uuid4()),
+            user_query=UserQuery(query_str=user_query),
+            assistant_response=AssistantResponse(response_str=assistant_response),
+        )
+
+        self.current_conversation.append_dialog_turn(dialog_turn)
+
+        self.turn_db.add(
+            {"user_query": user_query, "assistant_response": assistant_response}
+        )
diff --git a/adalflow/tests/test_memory.py b/adalflow/tests/test_memory.py
@@ -0,0 +1,23 @@
+from adalflow.components.memory.memory import Memory
+
+
+def test_empty_memory():
+    memory = Memory()
+    assert memory() == ""
+
+
+def test_add_dialog_turn():
+    memory = Memory()
+    memory.add_dialog_turn("Hello", "Hi! How can I help you?")
+    expected = "User: Hello\nAssistant: Hi! How can I help you?"
+    assert memory() == expected
+
+
+def test_multiple_turns():
+    memory = Memory()
+    memory.add_dialog_turn("Hello", "Hi!")
+    memory.add_dialog_turn("How are you?", "I'm good!")
+    expected = (
+        "User: Hello\n" "Assistant: Hi!\n" "User: How are you?\n" "Assistant: I'm good!"
+    )
+    assert memory() == expected
diff --git a/use_cases/question_answering/chatbot.ipynb b/use_cases/question_answering/chatbot.ipynb
@@ -16,19 +16,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Import needed modules from LightRAG\n",
-    "from adalflow.core.component import Component\n",
+    "# Import needed modules from LightRAGfrom adalflow.core.component import Component\n",
     "from adalflow.core.generator import Generator\n",
-    "from adalflow.core.memory import Memory"
+    "from adalflow.components.memory.memory import Memory"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -39,78 +38,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "ChatBot(\n",
-      "  (generator): Generator(\n",
-      "    model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
-      "    (system_prompt): Prompt(\n",
-      "      template: {# task desc #}\n",
-      "      {% if task_desc_str %}\n",
-      "      {{task_desc_str}}\n",
-      "      {% else %}\n",
-      "      Answer user query.\n",
-      "      {% endif %}\n",
-      "      {# output format #}\n",
-      "      {% if output_format_str %}\n",
-      "      <OUTPUT_FORMAT>\n",
-      "      {{output_format_str}}\n",
-      "      </OUTPUT_FORMAT>\n",
-      "      {% endif %}\n",
-      "      {# tools #}\n",
-      "      {% if tools_str %}\n",
-      "      <TOOLS>\n",
-      "      {{tools_str}}\n",
-      "      </TOOLS>\n",
-      "      {% endif %}\n",
-      "      {# example #}\n",
-      "      {% if examples_str %}\n",
-      "      <EXAMPLES>\n",
-      "      {{examples_str}}\n",
-      "      </EXAMPLES>\n",
-      "      {% endif %}\n",
-      "      {# chat history #}\n",
-      "      {% if chat_history_str %}\n",
-      "      <CHAT_HISTORY>\n",
-      "      {{chat_history_str}}\n",
-      "      </CHAT_HISTORY>\n",
-      "      {% endif %}\n",
-      "      {#contex#}\n",
-      "      {% if context_str %}\n",
-      "      <CONTEXT>\n",
-      "      {{context_str}}\n",
-      "      </CONTEXT>\n",
-      "      {% endif %}\n",
-      "      {# steps #}\n",
-      "      {% if steps_str %}\n",
-      "      <STEPS>\n",
-      "      {{steps_str}}\n",
-      "      </STEPS>\n",
-      "      {% endif %}\n",
-      "      {% if input_str %}\n",
-      "      <Inputs>\n",
-      "      {{input_str}}\n",
-      "      </Inputs>\n",
-      "      {% endif %}\n",
-      "      {% if output_str %}\n",
-      "      <Outputs>\n",
-      "      {{output_str}}\n",
-      "      </Outputs>\n",
-      "      {% endif %}\n",
-      "      , prompt_variables: ['context_str', 'task_desc_str', 'tools_str', 'chat_history_str', 'input_str', 'output_str', 'output_format_str', 'steps_str', 'examples_str']\n",
-      "    )\n",
-      "    (model_client): OpenAIClient()\n",
-      "  )\n",
-      "  (chat_history): Memory()\n",
-      ")\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Build the ChatBot pipeline\n",
     "class ChatBot(Component):\n",
@@ -149,30 +79,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Welcome to the ChatBot. Type anything to chat. Type 'exit' to end.\n",
-      "ChatBot: GeneratorOutput(data=\"Learning to drive can be an exciting and rewarding experience. Here are some general steps to help you get started on your journey to becoming a safe and confident driver:\\n\\n1. Get a learner's permit: In most places, you will need to obtain a learner's permit before you can start learning how to drive. Check with your local department of motor vehicles for the specific requirements in your area.\\n\\n2. Take a driver's education course: Consider enrolling in a driver's education course to learn the rules of the road and get some hands-on practice with a qualified instructor.\\n\\n3. Practice with a licensed driver: Before you can get your driver's license, you will need to log a certain number of supervised driving hours with a licensed adult. This is a great opportunity to get comfortable behind the wheel and practice your skills.\\n\\n4. Study the driver's manual: Make sure to familiarize yourself with the driver's manual for your state or country. It contains important information about traffic laws, road signs, and safe driving practices.\\n\\n5. Practice, practice, practice: The more time you spend behind the wheel, the more confident you will become as a driver. Practice in a variety of different conditions and situations to hone your skills.\\n\\n6. Take a driving test: Once you feel ready, schedule a driving test with your local department of motor vehicles. If you pass the test, you will receive your driver's license and be able to drive independently.\\n\\nRemember, learning to drive takes time and practice, so be patient with yourself and don't be afraid to ask for help if you need it. Good luck on your journey to becoming a licensed driver!\", error=None, raw_response=\"Learning to drive can be an exciting and rewarding experience. Here are some general steps to help you get started on your journey to becoming a safe and confident driver:\\n\\n1. Get a learner's permit: In most places, you will need to obtain a learner's permit before you can start learning how to drive. Check with your local department of motor vehicles for the specific requirements in your area.\\n\\n2. Take a driver's education course: Consider enrolling in a driver's education course to learn the rules of the road and get some hands-on practice with a qualified instructor.\\n\\n3. Practice with a licensed driver: Before you can get your driver's license, you will need to log a certain number of supervised driving hours with a licensed adult. This is a great opportunity to get comfortable behind the wheel and practice your skills.\\n\\n4. Study the driver's manual: Make sure to familiarize yourself with the driver's manual for your state or country. It contains important information about traffic laws, road signs, and safe driving practices.\\n\\n5. Practice, practice, practice: The more time you spend behind the wheel, the more confident you will become as a driver. Practice in a variety of different conditions and situations to hone your skills.\\n\\n6. Take a driving test: Once you feel ready, schedule a driving test with your local department of motor vehicles. If you pass the test, you will receive your driver's license and be able to drive independently.\\n\\nRemember, learning to drive takes time and practice, so be patient with yourself and don't be afraid to ask for help if you need it. Good luck on your journey to becoming a licensed driver!\")\n",
-      "ChatBot: GeneratorOutput(data=\"To get a driver's license in California, you can follow these general steps:\\n\\n1. Obtain a learner's permit: Applicants must be at least 15 and a half years old to apply for a learner's permit in California. You will need to pass a written knowledge test and a vision test to obtain your permit.\\n\\n2. Complete driver's education: If you are under 17 and a half years old, you must complete a driver's education course before applying for a provisional permit.\\n\\n3. Practice driving: With your learner's permit, you can start practicing driving with a licensed adult who is at least 25 years old.\\n\\n4. Apply for a provisional license: After holding your learner's permit for at least 6 months and completing at least 50 hours of practice (including 10 hours at night), you can apply for a provisional license.\\n\\n5. Pass the driving test: Schedule and pass a driving test at a local DMV office. Make sure to bring all required documents and fees.\\n\\n6. Receive your driver's license: If you pass the driving test, you will receive your provisional driver's license. With this license, you will have certain restrictions, such as driving with no passengers under 20 years old for the first year.\\n\\nRemember to check with the California Department of Motor Vehicles (DMV) for the most up-to-date and specific requirements for obtaining a driver's license in the state. Good luck with your journey to becoming a licensed driver in California!\", error=None, raw_response=\"To get a driver's license in California, you can follow these general steps:\\n\\n1. Obtain a learner's permit: Applicants must be at least 15 and a half years old to apply for a learner's permit in California. You will need to pass a written knowledge test and a vision test to obtain your permit.\\n\\n2. Complete driver's education: If you are under 17 and a half years old, you must complete a driver's education course before applying for a provisional permit.\\n\\n3. Practice driving: With your learner's permit, you can start practicing driving with a licensed adult who is at least 25 years old.\\n\\n4. Apply for a provisional license: After holding your learner's permit for at least 6 months and completing at least 50 hours of practice (including 10 hours at night), you can apply for a provisional license.\\n\\n5. Pass the driving test: Schedule and pass a driving test at a local DMV office. Make sure to bring all required documents and fees.\\n\\n6. Receive your driver's license: If you pass the driving test, you will receive your provisional driver's license. With this license, you will have certain restrictions, such as driving with no passengers under 20 years old for the first year.\\n\\nRemember to check with the California Department of Motor Vehicles (DMV) for the most up-to-date and specific requirements for obtaining a driver's license in the state. Good luck with your journey to becoming a licensed driver in California!\")\n",
-      "Goodbye!\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "chatbot.call()"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "lightrag-project",
+   "display_name": ".venv",
    "language": "python",
-   "name": "light-rag-project"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -184,7 +103,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.3"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,
