fix clear_history option and test

mitodl · Dec 19, 2024 · a7cb6ad · a7cb6ad
1 parent 9c6202e
commit a7cb6ad
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/ai_agents/consumers.py b/ai_agents/consumers.py
@@ -27,7 +27,7 @@ def process_message(message_json, agent) -> str:
     model = serializer.validated_data.pop("model", None)
 
     if clear_history:
-        agent.agent.clear_chat_history()
+        agent.clear_chat_history()
     if model:
         agent.agent.agent_worker._llm.model = model  # noqa: SLF001
     if temperature:

diff --git a/ai_agents/consumers_test.py b/ai_agents/consumers_test.py
@@ -118,6 +118,9 @@ async def test_clear_history(mocker, clear_history, recommendation_consumer):
     mock_clear = mocker.patch(
         "ai_agents.consumers.RecommendationAgent.clear_chat_history"
     )
+    mocker.patch(
+        "ai_agents.agents.RecommendationAgent.get_completion",
+    )
     await recommendation_consumer.connect()
     await recommendation_consumer.receive(
         json.dumps({"clear_history": clear_history, "message": "hello"})