From 5b85073c364ff9421646e91ef622f914c9ed6176 Mon Sep 17 00:00:00 2001
From: Matt Bertrand <mrbertrand@gmail.com>
Date: Fri, 20 Dec 2024 16:21:54 -0500
Subject: [PATCH] Track chatbot activity in posthog (#1923)

---
 ai_chat/agents.py      | 23 ++++++++++++++++++++---
 ai_chat/agents_test.py | 15 +++++++++++++--
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/ai_chat/agents.py b/ai_chat/agents.py
index a9dc69963b..7a63406be3 100644
--- a/ai_chat/agents.py
+++ b/ai_chat/agents.py
@@ -5,6 +5,7 @@
 from abc import ABC, abstractmethod
 from typing import Optional
 
+import posthog
 import pydantic
 import requests
 from django.conf import settings
@@ -146,6 +147,7 @@ def get_completion(self, message: str, *, debug: bool = settings.AI_DEBUG) -> st
 
         Append the response with debugging metadata and/or errors.
         """
+        full_response = ""
         if not self.agent:
             error = "Create agent before running"
             raise ValueError(error)
@@ -154,7 +156,9 @@ def get_completion(self, message: str, *, debug: bool = settings.AI_DEBUG) -> st
                 message,
             )
             response_gen = response.response_gen
-            yield from response_gen
+            for response in response_gen:
+                full_response += response
+                yield response
         except BadRequestError as error:
             # Format and yield an error message inside a hidden comment
             if hasattr(error, "response"):
@@ -178,13 +182,26 @@ def get_completion(self, message: str, *, debug: bool = settings.AI_DEBUG) -> st
             self.save_chat_history()
         if debug:
             yield f"\n\n<!-- {self.get_comment_metadata()} -->\n\n".encode()
+        hog_client = posthog.Posthog(
+            settings.POSTHOG_PROJECT_API_KEY, host=settings.POSTHOG_API_HOST
+        )
+        hog_client.capture(
+            self.user_id,
+            event=self.JOB_ID,
+            properties={
+                "question": message,
+                "answer": full_response,
+                "metadata": self.get_comment_metadata(),
+                "user": self.user_id,
+            },
+        )
 
 
 class SearchAgent(BaseChatAgent):
     """Service class for the AI search function agent"""
 
-    JOB_ID = "SEARCH_JOB"
-    TASK_NAME = "SEARCH_TASK"
+    JOB_ID = "recommendation_agent"
+    TASK_NAME = "recommendation_task"
 
     INSTRUCTIONS = f"""You are an assistant helping users find courses from a catalog
 of learning resources. Users can ask about specific topics, levels, or recommendations
diff --git a/ai_chat/agents_test.py b/ai_chat/agents_test.py
index f5fd5e6029..309f03183c 100644
--- a/ai_chat/agents_test.py
+++ b/ai_chat/agents_test.py
@@ -233,12 +233,13 @@ def test_get_completion(mocker):
             "system_prompt": SearchAgent.INSTRUCTIONS,
         }
     }
-    expected_return_value = [b"Here ", b"are ", b"some ", b"results"]
+    expected_return_value = ["Here ", "are ", "some ", "results"]
     mocker.patch(
         "ai_chat.agents.OpenAIAgent.stream_chat",
         return_value=mocker.Mock(response_gen=iter(expected_return_value)),
     )
-    search_agent = SearchAgent("test agent")
+    mock_hog = mocker.patch("ai_chat.agents.posthog.Posthog")
+    search_agent = SearchAgent("test agent", user_id="testuser@email.edu")
     search_agent.search_parameters = metadata["metadata"]["search_parameters"]
     search_agent.search_results = metadata["metadata"]["search_results"]
     search_agent.instructions = metadata["metadata"]["system_prompt"]
@@ -256,4 +257,14 @@ def test_get_completion(mocker):
         ]
     )
     search_agent.agent.stream_chat.assert_called_once_with("I want to learn physics")
+    mock_hog.return_value.capture.assert_called_once_with(
+        "testuser@email.edu",
+        event=search_agent.JOB_ID,
+        properties={
+            "question": "I want to learn physics",
+            "answer": "Here are some results",
+            "metadata": search_agent.get_comment_metadata(),
+            "user": "testuser@email.edu",
+        },
+    )
     assert "".join([str(value) for value in expected_return_value]) in results