From 000e9c0735c50bd6cb63866441d7f73083665902 Mon Sep 17 00:00:00 2001
From: Nate Sesti <sestinj@gmail.com>
Date: Mon, 24 Jul 2023 19:19:33 -0700
Subject: telemetry refactoring and switch config.json->py

---
 continuedev/.gitignore                             |  3 +-
 continuedev/src/continuedev/core/autopilot.py      | 14 ++---
 continuedev/src/continuedev/core/config.py         | 71 ----------------------
 continuedev/src/continuedev/core/policy.py         |  1 -
 continuedev/src/continuedev/core/sdk.py            | 15 +++--
 .../src/continuedev/libs/llm/proxy_server.py       |  6 +-
 .../src/continuedev/libs/util/create_async_task.py |  4 +-
 continuedev/src/continuedev/libs/util/telemetry.py | 38 +++++++-----
 .../src/continuedev/plugins/steps/feedback.py      |  5 +-
 continuedev/src/continuedev/plugins/steps/help.py  |  6 +-
 continuedev/src/continuedev/plugins/steps/main.py  | 14 -----
 .../src/continuedev/plugins/steps/open_config.py   | 27 ++++----
 continuedev/src/continuedev/server/gui.py          | 16 ++---
 continuedev/src/continuedev/server/gui_protocol.py |  4 --
 continuedev/src/continuedev/server/ide.py          | 31 +++++-----
 15 files changed, 82 insertions(+), 173 deletions(-)

(limited to 'continuedev')

diff --git a/continuedev/.gitignore b/continuedev/.gitignore
index 391ae7db..a9d34b8b 100644
--- a/continuedev/.gitignore
+++ b/continuedev/.gitignore
@@ -1,2 +1 @@
-notes.md
-config.json
\ No newline at end of file
+notes.md
\ No newline at end of file
diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 2ce7c1f9..0e355d78 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -2,7 +2,6 @@ from functools import cached_property
 import traceback
 import time
 from typing import Any, Callable, Coroutine, Dict, List, Union
-import os
 from aiohttp import ClientPayloadError
 from pydantic import root_validator
 
@@ -17,11 +16,11 @@ from ..libs.util.queue import AsyncSubscriptionQueue
 from ..models.main import ContinueBaseModel
 from .main import Context, ContinueCustomException, Policy, History, FullState, Step, HistoryNode
 from ..plugins.steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
-from ..libs.util.telemetry import capture_event
 from .sdk import ContinueSDK
 from ..libs.util.traceback_parsers import get_python_traceback, get_javascript_traceback
 from openai import error as openai_errors
 from ..libs.util.create_async_task import create_async_task
+from ..libs.util.telemetry import posthog_logger
 
 
 def get_error_title(e: Exception) -> str:
@@ -111,9 +110,6 @@ class Autopilot(ContinueBaseModel):
                               "name": x.name, "description": x.description}, self.continue_sdk.config.slash_commands)) or []
         return custom_commands + slash_commands
 
-    async def change_default_model(self, model: str):
-        self.continue_sdk.update_default_model(model)
-
     async def clear_history(self):
         # Reset history
         self.history = History.from_empty()
@@ -222,8 +218,8 @@ class Autopilot(ContinueBaseModel):
         #     last_depth = self.history.timeline[i].depth
         #     i -= 1
 
-        capture_event(self.continue_sdk.ide.unique_id, 'step run', {
-            'step_name': step.name, 'params': step.dict()})
+        posthog_logger.capture_event(
+            'step run', {'step_name': step.name, 'params': step.dict()})
 
         if not is_future_step:
             # Check manual edits buffer, clear out if needed by creating a ManualEditStep
@@ -264,8 +260,8 @@ class Autopilot(ContinueBaseModel):
             # Attach an InternalErrorObservation to the step and unhide it.
             print(
                 f"Error while running step: \n{error_string}\n{error_title}")
-            capture_event(self.continue_sdk.ide.unique_id, 'step error', {
-                'error_message': error_string, 'error_title': error_title, 'step_name': step.name, 'params': step.dict()})
+            posthog_logger.capture_event('step error', {
+                                         'error_message': error_string, 'error_title': error_title, 'step_name': step.name, 'params': step.dict()})
 
             observation = InternalErrorObservation(
                 error=error_string, title=error_title)
diff --git a/continuedev/src/continuedev/core/config.py b/continuedev/src/continuedev/core/config.py
index 565c617d..cb9c8977 100644
--- a/continuedev/src/continuedev/core/config.py
+++ b/continuedev/src/continuedev/core/config.py
@@ -95,74 +95,3 @@ class ContinueConfig(BaseModel):
     @validator('temperature', pre=True)
     def temperature_validator(cls, v):
         return max(0.0, min(1.0, v))
-
-
-def load_config(config_file: str) -> ContinueConfig:
-    """
-    Load the config file and return a ContinueConfig object.
-    """
-    if not os.path.exists(config_file):
-        return ContinueConfig()
-
-    _, ext = os.path.splitext(config_file)
-    if ext == '.yaml':
-        with open(config_file, 'r') as f:
-            try:
-                config_dict = yaml.safe_load(f)
-            except:
-                return ContinueConfig()
-    elif ext == '.json':
-        with open(config_file, 'r') as f:
-            try:
-                config_dict = json.load(f)
-            except:
-                return ContinueConfig()
-    else:
-        raise ValueError(f'Unknown config file extension: {ext}')
-    return ContinueConfig(**config_dict)
-
-
-def load_global_config() -> ContinueConfig:
-    """
-    Load the global config file and return a ContinueConfig object.
-    """
-    global_dir = os.path.expanduser('~/.continue')
-    if not os.path.exists(global_dir):
-        os.mkdir(global_dir)
-
-    yaml_path = os.path.join(global_dir, 'config.yaml')
-    if os.path.exists(yaml_path):
-        with open(config_path, 'r') as f:
-            try:
-                config_dict = yaml.safe_load(f)
-            except:
-                return ContinueConfig()
-    else:
-        config_path = os.path.join(global_dir, 'config.json')
-        if not os.path.exists(config_path):
-            with open(config_path, 'w') as f:
-                json.dump(ContinueConfig().dict(), f, indent=4)
-        with open(config_path, 'r') as f:
-            try:
-                config_dict = json.load(f)
-            except:
-                return ContinueConfig()
-    return ContinueConfig(**config_dict)
-
-
-def update_global_config(config: ContinueConfig):
-    """
-    Update the config file with the given ContinueConfig object.
-    """
-    global_dir = os.path.expanduser('~/.continue')
-    if not os.path.exists(global_dir):
-        os.mkdir(global_dir)
-
-    yaml_path = os.path.join(global_dir, 'config.yaml')
-    if os.path.exists(yaml_path):
-        with open(config_path, 'w') as f:
-            yaml.dump(config.dict(), f, indent=4)
-    else:
-        config_path = os.path.join(global_dir, 'config.json')
-        with open(config_path, 'w') as f:
-            json.dump(config.dict(exclude_unset=False), f, indent=4)
diff --git a/continuedev/src/continuedev/core/policy.py b/continuedev/src/continuedev/core/policy.py
index 1c87cfeb..d90177b5 100644
--- a/continuedev/src/continuedev/core/policy.py
+++ b/continuedev/src/continuedev/core/policy.py
@@ -56,7 +56,6 @@ class DefaultPolicy(Policy):
                     - Use `cmd+m` (Mac) / `ctrl+m` (Windows) to open Continue
                     - Use `/help` to ask questions about how to use Continue""")) >>
                 WelcomeStep() >>
-                # SetupContinueWorkspaceStep() >>
                 # CreateCodebaseIndexChroma() >>
                 StepsOnStartupStep())
 
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index f925f20f..6668c8c3 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -1,13 +1,11 @@
-import asyncio
 from functools import cached_property
 from typing import Coroutine, Dict, Union
 import os
 
 from ..plugins.steps.core.core import DefaultModelEditCodeStep
 from ..models.main import Range
-from .context import ContextItem
 from .abstract_sdk import AbstractContinueSDK
-from .config import ContinueConfig, load_config, load_global_config, update_global_config
+from .config import ContinueConfig
 from ..models.filesystem_edit import FileEdit, FileSystemEdit, AddFile, DeleteFile, AddDirectory, DeleteDirectory
 from ..models.filesystem import RangeInFile
 from ..libs.llm.hf_inference_api import HuggingFaceInferenceAPI
@@ -19,6 +17,7 @@ from ..server.ide_protocol import AbstractIdeProtocolServer
 from .main import Context, ContinueCustomException, History, HistoryNode, Step, ChatMessage
 from ..plugins.steps.core.core import *
 from ..libs.llm.proxy_server import ProxyServer
+from ..libs.util.telemetry import posthog_logger
 
 
 class Autopilot:
@@ -266,6 +265,11 @@ class ContinueSDK(AbstractContinueSDK):
             config = importlib.util.module_from_spec(spec)
             spec.loader.exec_module(config)
             self._last_valid_config = config.config
+
+            # When the config is loaded, setup posthog logger
+            posthog_logger.setup(
+                self.ide.unique_id, config.config.allow_anonymous_telemetry or True)
+
             return config.config
         except Exception as e:
             print("Error loading config.py: ", e)
@@ -276,11 +280,6 @@ class ContinueSDK(AbstractContinueSDK):
                        ) if only_editing else self.__autopilot._highlighted_ranges
         return [c.range for c in context]
 
-    def update_default_model(self, model: str):
-        config = self.config
-        config.default_model = model
-        update_global_config(config)
-
     def set_loading_message(self, message: str):
         # self.__autopilot.set_loading_message(message)
         raise NotImplementedError()
diff --git a/continuedev/src/continuedev/libs/llm/proxy_server.py b/continuedev/src/continuedev/libs/llm/proxy_server.py
index 17694afe..f9e3fa01 100644
--- a/continuedev/src/continuedev/libs/llm/proxy_server.py
+++ b/continuedev/src/continuedev/libs/llm/proxy_server.py
@@ -3,9 +3,9 @@ import json
 import traceback
 from typing import Any, Callable, Coroutine, Dict, Generator, List, Literal, Union
 import aiohttp
-from ..util.telemetry import capture_event
 from ...core.main import ChatMessage
 from ..llm import LLM
+from ..util.telemetry import posthog_logger
 from ..util.count_tokens import DEFAULT_ARGS, compile_chat_messages, count_tokens, format_chat_messages
 import certifi
 import ssl
@@ -36,7 +36,7 @@ class ProxyServer(LLM):
 
     def count_tokens(self, text: str):
         return count_tokens(self.default_model, text)
-    
+
     def get_headers(self):
         # headers with unique id
         return {"unique_id": self.unique_id}
@@ -87,7 +87,7 @@ class ProxyServer(LLM):
                                     if "content" in loaded_chunk:
                                         completion += loaded_chunk["content"]
                         except Exception as e:
-                            capture_event(self.unique_id, "proxy_server_parse_error", {
+                            posthog_logger.capture_event(self.unique_id, "proxy_server_parse_error", {
                                 "error_title": "Proxy server stream_chat parsing failed", "error_message": '\n'.join(traceback.format_exception(e))})
                     else:
                         break
diff --git a/continuedev/src/continuedev/libs/util/create_async_task.py b/continuedev/src/continuedev/libs/util/create_async_task.py
index 354cea82..2473c638 100644
--- a/continuedev/src/continuedev/libs/util/create_async_task.py
+++ b/continuedev/src/continuedev/libs/util/create_async_task.py
@@ -1,6 +1,6 @@
 from typing import Coroutine, Union
 import traceback
-from .telemetry import capture_event
+from .telemetry import posthog_logger
 import asyncio
 import nest_asyncio
 nest_asyncio.apply()
@@ -16,7 +16,7 @@ def create_async_task(coro: Coroutine, unique_id: Union[str, None] = None):
         except Exception as e:
             print("Exception caught from async task: ",
                   '\n'.join(traceback.format_exception(e)))
-            capture_event(unique_id or "None", "async_task_error", {
+            posthog_logger.capture_event("async_task_error", {
                 "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))
             })
 
diff --git a/continuedev/src/continuedev/libs/util/telemetry.py b/continuedev/src/continuedev/libs/util/telemetry.py
index 17735dce..a967828e 100644
--- a/continuedev/src/continuedev/libs/util/telemetry.py
+++ b/continuedev/src/continuedev/libs/util/telemetry.py
@@ -1,27 +1,37 @@
 from typing import Any
 from posthog import Posthog
-from ...core.config import load_config
 import os
 from dotenv import load_dotenv
 from .commonregex import clean_pii_from_any
 
 load_dotenv()
 in_codespaces = os.getenv("CODESPACES") == "true"
+POSTHOG_API_KEY = 'phc_JS6XFROuNbhJtVCEdTSYk6gl5ArRrTNMpCcguAXlSPs'
 
-# The personal API key is necessary only if you want to use local evaluation of feature flags.
-posthog = Posthog('phc_JS6XFROuNbhJtVCEdTSYk6gl5ArRrTNMpCcguAXlSPs',
-                  host='https://app.posthog.com')
 
+class PostHogLogger:
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+        self.unique_id = None
+        self.allow_anonymous_telemetry = True
 
-def capture_event(unique_id: str, event_name: str, event_properties: Any):
-    # Return early if telemetry is disabled
-    config = load_config('.continue/config.json')
-    if not config.allow_anonymous_telemetry:
-        return
+    def setup(self, unique_id: str, allow_anonymous_telemetry: bool):
+        self.unique_id = unique_id
+        self.allow_anonymous_telemetry = allow_anonymous_telemetry
 
-    if in_codespaces:
-        event_properties['codespaces'] = True
+        # The personal API key is necessary only if you want to use local evaluation of feature flags.
+        self.posthog = Posthog(self.api_key, host='https://app.posthog.com')
 
-    # Send event to PostHog
-    posthog.capture(unique_id, event_name,
-                    clean_pii_from_any(event_properties))
+    def capture_event(self, event_name: str, event_properties: Any):
+        if not self.allow_anonymous_telemetry or self.unique_id is None:
+            return
+
+        if in_codespaces:
+            event_properties['codespaces'] = True
+
+        # Send event to PostHog
+        self.posthog.capture(self.unique_id, event_name,
+                             clean_pii_from_any(event_properties))
+
+
+posthog_logger = PostHogLogger(api_key=POSTHOG_API_KEY)
diff --git a/continuedev/src/continuedev/plugins/steps/feedback.py b/continuedev/src/continuedev/plugins/steps/feedback.py
index 119e3112..fa56a4d9 100644
--- a/continuedev/src/continuedev/plugins/steps/feedback.py
+++ b/continuedev/src/continuedev/plugins/steps/feedback.py
@@ -2,7 +2,7 @@ from typing import Coroutine
 from ...core.main import Models
 from ...core.main import Step
 from ...core.sdk import ContinueSDK
-from ...libs.util.telemetry import capture_event
+from ...libs.util.telemetry import posthog_logger
 
 
 class FeedbackStep(Step):
@@ -13,5 +13,4 @@ class FeedbackStep(Step):
         return f"`{self.user_input}`\n\nWe'll see your feedback and make improvements as soon as possible. If you'd like to directly email us, you can contact [nate@continue.dev](mailto:nate@continue.dev?subject=Feedback%20On%20Continue)."
 
     async def run(self, sdk: ContinueSDK):
-        capture_event(sdk.ide.unique_id, "feedback",
-                      {"feedback": self.user_input})
+        posthog_logger.capture_event("feedback", {"feedback": self.user_input})
diff --git a/continuedev/src/continuedev/plugins/steps/help.py b/continuedev/src/continuedev/plugins/steps/help.py
index 5111c7cf..d3807706 100644
--- a/continuedev/src/continuedev/plugins/steps/help.py
+++ b/continuedev/src/continuedev/plugins/steps/help.py
@@ -1,7 +1,7 @@
 from textwrap import dedent
 from ...core.main import ChatMessage, Step
 from ...core.sdk import ContinueSDK
-from ...libs.util.telemetry import capture_event
+from ...libs.util.telemetry import posthog_logger
 
 help = dedent("""\
         Continue is an open-source coding autopilot. It is a VS Code extension that brings the power of ChatGPT to your IDE.
@@ -55,5 +55,5 @@ class HelpStep(Step):
                 self.description += chunk["content"]
                 await sdk.update_ui()
 
-        capture_event(sdk.ide.unique_id, "help", {
-                      "question": question, "answer": self.description})
+        posthog_logger.capture_event(
+            "help", {"question": question, "answer": self.description})
diff --git a/continuedev/src/continuedev/plugins/steps/main.py b/continuedev/src/continuedev/plugins/steps/main.py
index 30117c55..a8752df2 100644
--- a/continuedev/src/continuedev/plugins/steps/main.py
+++ b/continuedev/src/continuedev/plugins/steps/main.py
@@ -15,20 +15,6 @@ from .core.core import DefaultModelEditCodeStep
 from ...libs.util.calculate_diff import calculate_diff2
 
 
-class SetupContinueWorkspaceStep(Step):
-    async def describe(self, models: Models) -> Coroutine[str, None, None]:
-        return "Set up Continue workspace by adding a .continue directory"
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        if not os.path.exists(os.path.join(await sdk.ide.getWorkspaceDirectory(), ".continue")):
-            await sdk.add_directory(".continue")
-            if not os.path.exists(os.path.join(await sdk.ide.getWorkspaceDirectory(), ".continue", "config.json")):
-                await sdk.add_file(".continue/config.json", dedent("""\
-                    {
-                        "allow_anonymous_telemetry": true
-                    }"""))
-
-
 class Policy(BaseModel):
     pass
 
diff --git a/continuedev/src/continuedev/plugins/steps/open_config.py b/continuedev/src/continuedev/plugins/steps/open_config.py
index d950c26f..7264a59b 100644
--- a/continuedev/src/continuedev/plugins/steps/open_config.py
+++ b/continuedev/src/continuedev/plugins/steps/open_config.py
@@ -9,21 +9,22 @@ class OpenConfigStep(Step):
 
     async def describe(self, models):
         return dedent("""\
-            `\"config.json\"` is now open. You can add a custom slash command in the `\"custom_commands\"` section, like in this example:
-            ```json
-            "custom_commands": [
-                {
-                    "name": "test",
-                    "description": "Write unit tests like I do for the highlighted code",
-                    "prompt": "Write a comprehensive set of unit tests for the selected code. It should setup, run tests that check for correctness including important edge cases, and teardown. Ensure that the tests are complete and sophisticated."
-                }
-            ]
+            `\"config.py\"` is now open. You can add a custom slash command in the `\"custom_commands\"` section, like in this example:
+            ```python
+            config = ContinueConfig(
+                ...
+                custom_commands=[CustomCommand(
+                    name="test",
+                    description="Write unit tests like I do for the highlighted code",
+                    prompt="Write a comprehensive set of unit tests for the selected code. It should setup, run tests that check for correctness including important edge cases, and teardown. Ensure that the tests are complete and sophisticated.",
+                )]
+            )
             ```
-            `"name"` is the command you will type.
-            `"description"` is the description displayed in the slash command menu.
-            `"prompt"` is the instruction given to the model. The overall prompt becomes "Task: {prompt}, Additional info: {user_input}". For example, if you entered "/test exactly 5 assertions", the overall prompt would become "Task: Write a comprehensive...and sophisticated, Additional info: exactly 5 assertions".""")
+            `name` is the command you will type.
+            `description` is the description displayed in the slash command menu.
+            `prompt` is the instruction given to the model. The overall prompt becomes "Task: {prompt}, Additional info: {user_input}". For example, if you entered "/test exactly 5 assertions", the overall prompt would become "Task: Write a comprehensive...and sophisticated, Additional info: exactly 5 assertions".""")
 
     async def run(self, sdk: ContinueSDK):
         global_dir = os.path.expanduser('~/.continue')
-        config_path = os.path.join(global_dir, 'config.json')
+        config_path = os.path.join(global_dir, 'config.py')
         await sdk.ide.setFileOpen(config_path)
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 36b2f3fa..70778dfe 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -2,15 +2,15 @@ import asyncio
 import json
 from fastapi import Depends, Header, WebSocket, APIRouter
 from starlette.websockets import WebSocketState, WebSocketDisconnect
-from typing import Any, List, Type, TypeVar, Union
+from typing import Any, List, Type, TypeVar
 from pydantic import BaseModel
 import traceback
 from uvicorn.main import Server
 
-from .session_manager import SessionManager, session_manager, Session
+from .session_manager import session_manager, Session
 from .gui_protocol import AbstractGUIProtocolServer
 from ..libs.util.queue import AsyncSubscriptionQueue
-from ..libs.util.telemetry import capture_event
+from ..libs.util.telemetry import posthog_logger
 from ..libs.util.create_async_task import create_async_task
 
 router = APIRouter(prefix="/gui", tags=["gui"])
@@ -85,8 +85,6 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
                 self.on_reverse_to_index(data["index"])
             elif message_type == "retry_at_index":
                 self.on_retry_at_index(data["index"])
-            elif message_type == "change_default_model":
-                self.on_change_default_model(data["model"])
             elif message_type == "clear_history":
                 self.on_clear_history()
             elif message_type == "delete_at_index":
@@ -126,10 +124,6 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
         create_async_task(
             self.session.autopilot.retry_at_index(index), self.session.autopilot.continue_sdk.ide.unique_id)
 
-    def on_change_default_model(self, model: str):
-        create_async_task(self.session.autopilot.change_default_model(
-            model), self.session.autopilot.continue_sdk.ide.unique_id)
-
     def on_clear_history(self):
         create_async_task(self.session.autopilot.clear_history(
         ), self.session.autopilot.continue_sdk.ide.unique_id)
@@ -199,8 +193,8 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we
         print("GUI websocket disconnected")
     except Exception as e:
         print("ERROR in gui websocket: ", e)
-        capture_event(session.autopilot.continue_sdk.ide.unique_id, "gui_error", {
-                      "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))})
+        posthog_logger.capture_event("gui_error", {
+            "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))})
         raise e
     finally:
         print("Closing gui websocket")
diff --git a/continuedev/src/continuedev/server/gui_protocol.py b/continuedev/src/continuedev/server/gui_protocol.py
index fb230216..990833be 100644
--- a/continuedev/src/continuedev/server/gui_protocol.py
+++ b/continuedev/src/continuedev/server/gui_protocol.py
@@ -29,10 +29,6 @@ class AbstractGUIProtocolServer(ABC):
     def on_retry_at_index(self, index: int):
         """Called when the user requests a retry at a previous index"""
 
-    @abstractmethod
-    def on_change_default_model(self):
-        """Called when the user requests to change the default model"""
-
     @abstractmethod
     def on_clear_history(self):
         """Called when the user requests to clear the history"""
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index aeff5623..aa69a6cb 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -1,23 +1,24 @@
 # This is a separate server from server/main.py
-from functools import cached_property
 import json
 import os
-from typing import Any, Dict, List, Type, TypeVar, Union
+from typing import Any, List, Type, TypeVar, Union
 import uuid
-from fastapi import WebSocket, Body, APIRouter
+from fastapi import WebSocket, APIRouter
 from starlette.websockets import WebSocketState, WebSocketDisconnect
 from uvicorn.main import Server
+from pydantic import BaseModel
 import traceback
+import asyncio
 
-from ..libs.util.telemetry import capture_event
+from ..libs.util.telemetry import posthog_logger
 from ..libs.util.queue import AsyncSubscriptionQueue
 from ..models.filesystem import FileSystem, RangeInFile, EditDiff, RangeInFileWithContents, RealFileSystem
 from ..models.filesystem_edit import AddDirectory, AddFile, DeleteDirectory, DeleteFile, FileSystemEdit, FileEdit, FileEditWithFullContents, RenameDirectory, RenameFile, SequentialFileSystemEdit
-from pydantic import BaseModel
-from .gui import SessionManager, session_manager
+from .gui import session_manager
 from .ide_protocol import AbstractIdeProtocolServer
-import asyncio
 from ..libs.util.create_async_task import create_async_task
+from .session_manager import SessionManager
+
 import nest_asyncio
 nest_asyncio.apply()
 
@@ -273,12 +274,12 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
     # like file changes, tracebacks, etc...
 
     def onAcceptRejectSuggestion(self, accepted: bool):
-        capture_event(self.unique_id, "accept_reject_suggestion", {
+        posthog_logger.capture_event("accept_reject_suggestion", {
             "accepted": accepted
         })
 
     def onAcceptRejectDiff(self, accepted: bool):
-        capture_event(self.unique_id, "accept_reject_diff", {
+        posthog_logger.capture_event("accept_reject_diff", {
             "accepted": accepted
         })
 
@@ -450,8 +451,8 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
         if session_id is not None:
             session_manager.registered_ides[session_id] = ideProtocolServer
         other_msgs = await ideProtocolServer.initialize(session_id)
-        capture_event(ideProtocolServer.unique_id, "session_started", {
-                      "session_id": ideProtocolServer.session_id})
+        posthog_logger.capture_event("session_started", {
+            "session_id": ideProtocolServer.session_id})
 
         for other_msg in other_msgs:
             handle_msg(other_msg)
@@ -465,13 +466,13 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
         print("IDE wbsocket disconnected")
     except Exception as e:
         print("Error in ide websocket: ", e)
-        capture_event(ideProtocolServer.unique_id, "gui_error", {
-                      "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))})
+        posthog_logger.capture_event("gui_error", {
+            "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))})
         raise e
     finally:
         if websocket.client_state != WebSocketState.DISCONNECTED:
             await websocket.close()
 
-        capture_event(ideProtocolServer.unique_id, "session_ended", {
-                      "session_id": ideProtocolServer.session_id})
+        posthog_logger.capture_event("session_ended", {
+            "session_id": ideProtocolServer.session_id})
         session_manager.registered_ides.pop(ideProtocolServer.session_id)
-- 
cgit v1.2.3-70-g09d2