update

Author: kausmeows

libs/agno/tests/integration/agent/test_agent_with_storage_and_memory.py

@@ -27,91 +27,91 @@ def memory_agent(agent_storage, memory):
     )
 
 
-# def test_agent_session_state(chat_agent, agent_storage):
-#     session_id = "session_1"
-
-#     chat_agent.session_id = session_id
-#     chat_agent.session_name = "my_test_session"
-#     chat_agent.session_state = {"test_key": "test_value"}
-#     chat_agent.team_session_state = {"team_test_key": "team_test_value"}
-
-#     response = chat_agent.run("Hello, how are you?")
-#     assert response.run_id is not None
-#     assert chat_agent.session_id == session_id
-#     assert chat_agent.session_name == "my_test_session"
-#     assert chat_agent.session_state == {"current_session_id": session_id, "test_key": "test_value"}
-#     assert chat_agent.team_session_state == {"team_test_key": "team_test_value"}
-#     session_from_storage = agent_storage.read(session_id=session_id)
-#     assert session_from_storage is not None
-#     assert session_from_storage.session_id == session_id
-#     assert session_from_storage.session_data["session_name"] == "my_test_session"
-#     assert session_from_storage.session_data["session_state"] == {
-#         "current_session_id": session_id,
-#         "test_key": "test_value",
-#     }
-
-#     # Run again with the same session ID
-#     response = chat_agent.run("What can you do?", session_id=session_id)
-#     assert response.run_id is not None
-#     assert chat_agent.session_id == session_id
-#     assert chat_agent.session_name == "my_test_session"
-#     assert chat_agent.session_state == {"current_session_id": session_id, "test_key": "test_value"}
-
-#     # Run with a different session ID
-#     response = chat_agent.run("What can you do?", session_id="session_2")
-#     assert response.run_id is not None
-#     assert chat_agent.session_id == "session_2"
-#     assert chat_agent.session_name is None
-#     assert chat_agent.session_state == {"current_session_id": "session_2"}
-#     assert chat_agent.team_session_state is None
-
-#     # Run again with original session ID
-#     response = chat_agent.run("What name should I call you?", session_id=session_id)
-#     assert response.run_id is not None
-#     assert chat_agent.session_id == session_id
-#     assert chat_agent.session_name == "my_test_session"
-#     assert chat_agent.session_state == {"current_session_id": session_id, "test_key": "test_value"}
-
-
-# def test_agent_runs_in_memory(chat_agent):
-#     session_id = "test_session"
-#     response = chat_agent.run("Hello, how are you?", session_id=session_id)
-#     assert response is not None
-#     assert response.content is not None
-#     assert response.run_id is not None
-
-#     assert len(chat_agent.memory.runs[session_id]) == 1
-#     stored_run_response = chat_agent.memory.runs[session_id][0]
-#     assert stored_run_response.run_id == response.run_id
-#     assert len(stored_run_response.messages) == 2
-
-#     # Check that the run is also stored in the agent session
-#     assert len(chat_agent.agent_session.memory["runs"]) == 1
-
-
-# def test_agent_runs_in_memory_legacy(chat_agent):
-#     chat_agent.memory = AgentMemory()
-#     session_id = "test_session"
-#     response = chat_agent.run(
-#         "What can you do?",
-#         messages=[
-#             Message(role="user", content="Hello, how are you?"),
-#             Message(role="assistant", content="I'm good, thank you!"),
-#         ],
-#         session_id=session_id,
-#     )
-#     assert response is not None
-#     assert response.content is not None
-#     assert response.run_id is not None
-
-#     assert len(chat_agent.memory.runs) == 1
-#     stored_agent_run = chat_agent.memory.runs[0]
-#     assert stored_agent_run.response.run_id == response.run_id
-#     assert len(stored_agent_run.response.messages) == 4
-#     assert len(stored_agent_run.messages) == 2
-
-#     # Check that the run is also stored in the agent session
-#     assert len(chat_agent.agent_session.memory["runs"]) == 1
+def test_agent_session_state(chat_agent, agent_storage):
+    session_id = "session_1"
+
+    chat_agent.session_id = session_id
+    chat_agent.session_name = "my_test_session"
+    chat_agent.session_state = {"test_key": "test_value"}
+    chat_agent.team_session_state = {"team_test_key": "team_test_value"}
+
+    response = chat_agent.run("Hello, how are you?")
+    assert response.run_id is not None
+    assert chat_agent.session_id == session_id
+    assert chat_agent.session_name == "my_test_session"
+    assert chat_agent.session_state == {"current_session_id": session_id, "test_key": "test_value"}
+    assert chat_agent.team_session_state == {"team_test_key": "team_test_value"}
+    session_from_storage = agent_storage.read(session_id=session_id)
+    assert session_from_storage is not None
+    assert session_from_storage.session_id == session_id
+    assert session_from_storage.session_data["session_name"] == "my_test_session"
+    assert session_from_storage.session_data["session_state"] == {
+        "current_session_id": session_id,
+        "test_key": "test_value",
+    }
+
+    # Run again with the same session ID
+    response = chat_agent.run("What can you do?", session_id=session_id)
+    assert response.run_id is not None
+    assert chat_agent.session_id == session_id
+    assert chat_agent.session_name == "my_test_session"
+    assert chat_agent.session_state == {"current_session_id": session_id, "test_key": "test_value"}
+
+    # Run with a different session ID
+    response = chat_agent.run("What can you do?", session_id="session_2")
+    assert response.run_id is not None
+    assert chat_agent.session_id == "session_2"
+    assert chat_agent.session_name is None
+    assert chat_agent.session_state == {"current_session_id": "session_2"}
+    assert chat_agent.team_session_state is None
+
+    # Run again with original session ID
+    response = chat_agent.run("What name should I call you?", session_id=session_id)
+    assert response.run_id is not None
+    assert chat_agent.session_id == session_id
+    assert chat_agent.session_name == "my_test_session"
+    assert chat_agent.session_state == {"current_session_id": session_id, "test_key": "test_value"}
+
+
+def test_agent_runs_in_memory(chat_agent):
+    session_id = "test_session"
+    response = chat_agent.run("Hello, how are you?", session_id=session_id)
+    assert response is not None
+    assert response.content is not None
+    assert response.run_id is not None
+
+    assert len(chat_agent.memory.runs[session_id]) == 1
+    stored_run_response = chat_agent.memory.runs[session_id][0]
+    assert stored_run_response.run_id == response.run_id
+    assert len(stored_run_response.messages) == 2
+
+    # Check that the run is also stored in the agent session
+    assert len(chat_agent.agent_session.memory["runs"]) == 1
+
+
+def test_agent_runs_in_memory_legacy(chat_agent):
+    chat_agent.memory = AgentMemory()
+    session_id = "test_session"
+    response = chat_agent.run(
+        "What can you do?",
+        messages=[
+            Message(role="user", content="Hello, how are you?"),
+            Message(role="assistant", content="I'm good, thank you!"),
+        ],
+        session_id=session_id,
+    )
+    assert response is not None
+    assert response.content is not None
+    assert response.run_id is not None
+
+    assert len(chat_agent.memory.runs) == 1
+    stored_agent_run = chat_agent.memory.runs[0]
+    assert stored_agent_run.response.run_id == response.run_id
+    assert len(stored_agent_run.response.messages) == 4
+    assert len(stored_agent_run.messages) == 2
+
+    # Check that the run is also stored in the agent session
+    assert len(chat_agent.agent_session.memory["runs"]) == 1
 
 
 @pytest.mark.asyncio