From 3541d6a770c919e1f2e55a1ae53c4fc3abe31aa7 Mon Sep 17 00:00:00 2001
From: Nate Sesti <sestinj@gmail.com>
Date: Sun, 13 Aug 2023 22:00:02 -0700
Subject: feat: :white_check_mark: update test and add model telemetry

---
 continuedev/src/continuedev/libs/util/telemetry.py     | 5 +++++
 continuedev/src/continuedev/plugins/steps/chat.py      | 7 ++++++-
 continuedev/src/continuedev/plugins/steps/core/core.py | 5 +++++
 docs/docs/telemetry.md                                 | 2 ++
 extension/src/test-suite/environmentSetup.test.ts      | 4 ++--
 5 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/continuedev/src/continuedev/libs/util/telemetry.py b/continuedev/src/continuedev/libs/util/telemetry.py
index 0f66ad8d..2c76f4a3 100644
--- a/continuedev/src/continuedev/libs/util/telemetry.py
+++ b/continuedev/src/continuedev/libs/util/telemetry.py
@@ -26,6 +26,11 @@ class PostHogLogger:
         self.unique_id = unique_id or "NO_UNIQUE_ID"
         self.allow_anonymous_telemetry = allow_anonymous_telemetry or True
 
+        # Capture initial event
+        self.capture_event("session_start", {
+            "os": os.name
+        })
+
     def capture_event(self, event_name: str, event_properties: Any):
         # logger.debug(
         #     f"Logging to PostHog: {event_name} ({self.unique_id}, {self.allow_anonymous_telemetry}): {event_properties}")
diff --git a/continuedev/src/continuedev/plugins/steps/chat.py b/continuedev/src/continuedev/plugins/steps/chat.py
index 8e494ad0..ad19434f 100644
--- a/continuedev/src/continuedev/plugins/steps/chat.py
+++ b/continuedev/src/continuedev/plugins/steps/chat.py
@@ -6,12 +6,13 @@ from pydantic import Field
 
 from ...libs.util.strings import remove_quotes_and_escapes
 from .main import EditHighlightedCodeStep
-from .core.core import DisplayErrorStep, MessageStep
+from .core.core import MessageStep
 from ...core.main import FunctionCall, Models
 from ...core.main import ChatMessage, Step, step_to_json_schema
 from ...core.sdk import ContinueSDK
 from ...libs.llm.openai import OpenAI
 from ...libs.llm.maybe_proxy_openai import MaybeProxyOpenAI
+from ...libs.util.telemetry import posthog_logger
 import openai
 import os
 from dotenv import load_dotenv
@@ -59,6 +60,10 @@ class SimpleChatStep(Step):
         generator = sdk.models.default.stream_chat(
             messages, temperature=sdk.config.temperature)
 
+        posthog_logger.capture_event("model_use", {
+            "model": sdk.models.default.name
+        })
+
         async for chunk in generator:
             if sdk.current_step_was_deleted():
                 # So that the message doesn't disappear
diff --git a/continuedev/src/continuedev/plugins/steps/core/core.py b/continuedev/src/continuedev/plugins/steps/core/core.py
index 78174bf6..1ef201da 100644
--- a/continuedev/src/continuedev/plugins/steps/core/core.py
+++ b/continuedev/src/continuedev/plugins/steps/core/core.py
@@ -19,6 +19,7 @@ from ....core.observation import Observation, TextObservation, TracebackObservat
 from ....core.main import ChatMessage, ContinueCustomException, Step, SequentialStep
 from ....libs.util.count_tokens import DEFAULT_MAX_TOKENS
 from ....libs.util.strings import dedent_and_get_common_whitespace, remove_quotes_and_escapes
+from ....libs.util.telemetry import posthog_logger
 
 
 class ContinueSDK:
@@ -520,6 +521,10 @@ Please output the code to be inserted at the cursor in order to fulfill the user
         generator = model_to_use.stream_chat(
             messages, temperature=sdk.config.temperature, max_tokens=max_tokens)
 
+        posthog_logger.capture_event("model_use", {
+            "model": model_to_use.name
+        })
+
         try:
             async for chunk in generator:
                 # Stop early if it is repeating the file_suffix or the step was deleted
diff --git a/docs/docs/telemetry.md b/docs/docs/telemetry.md
index 1b9a8c4b..e0ea2158 100644
--- a/docs/docs/telemetry.md
+++ b/docs/docs/telemetry.md
@@ -11,6 +11,8 @@ We track
 - the steps that are run and their parameters
 - whether you accept or reject suggestions (not the code itself)
 - the traceback when an error occurs
+- the name of your OS
+- the name of the default model you configured
 
 All data is anonymous and cleaned of PII before being sent to PostHog.
 
diff --git a/extension/src/test-suite/environmentSetup.test.ts b/extension/src/test-suite/environmentSetup.test.ts
index d0406340..19d5470e 100644
--- a/extension/src/test-suite/environmentSetup.test.ts
+++ b/extension/src/test-suite/environmentSetup.test.ts
@@ -7,8 +7,8 @@ import fetch from "node-fetch";
 import fkill from "fkill";
 
 describe("Can start python server", () => {
-  test("Can start python server in under 10 seconds", async function () {
-    const allowedTime = 25_000;
+  test("Can start python server in under 35 seconds", async function () {
+    const allowedTime = 35_000;
     this.timeout(allowedTime + 10_000);
     try {
       fkill(65432, { force: true, silent: true });
-- 
cgit v1.2.3-70-g09d2