check for codespaces, stream edits

4 files changed, 104 insertions, 36 deletions

diff --git a/continuedev/poetry.lock b/continuedev/poetry.lock
index 93aaf82b..017f12f9 100644
--- a/continuedev/poetry.lock
+++ b/continuedev/poetry.lock
@@ -586,6 +586,18 @@ files = [
 ]
 
 [[package]]
+name = "jsonref"
+version = "1.1.0"
+description = "jsonref is a library for automatic dereferencing of JSON Reference objects for Python."
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "jsonref-1.1.0-py3-none-any.whl", hash = "sha256:590dc7773df6c21cbf948b5dac07a72a251db28b0238ceecce0a2abfa8ec30a9"},
+    {file = "jsonref-1.1.0.tar.gz", hash = "sha256:32fe8e1d85af0fdefbebce950af85590b22b60f9e95443176adbde4e1ecea552"},
+]
+
+[[package]]
 name = "langchain"
 version = "0.0.171"
 description = "Building applications with LLMs through composability"
@@ -1737,4 +1749,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9"
-content-hash = "17910714e3ad780ae7222b62c98539489d198aea67e5c7e4a9fc7672207f500f"
+content-hash = "9406bc70d0463b354c294bd9548897a33270b8a04f55141a763d45af8d6928b8"
diff --git a/continuedev/pyproject.toml b/continuedev/pyproject.toml
index af6ff045..bbd8a687 100644
--- a/continuedev/pyproject.toml
+++ b/continuedev/pyproject.toml
@@ -21,6 +21,7 @@ urllib3 = "1.26.15"
 gpt-index = "^0.6.8"
 posthog = "^3.0.1"
 tiktoken = "^0.4.0"
+jsonref = "^1.1.0"
 
 [tool.poetry.scripts]
 typegen = "src.continuedev.models.generate_json_schema:main" 
diff --git a/continuedev/src/continuedev/libs/util/telemetry.py b/continuedev/src/continuedev/libs/util/telemetry.py
index 03ec93c6..bd9fde9d 100644
--- a/continuedev/src/continuedev/libs/util/telemetry.py
+++ b/continuedev/src/continuedev/libs/util/telemetry.py
@@ -1,12 +1,22 @@
+from typing import Any
 from posthog import Posthog
 from ...core.config import load_config
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+in_codespaces = os.getenv("CODESPACES") == "true"
 
 # 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')
 
 
-def capture_event(unique_id: str, event_name, event_properties):
+def capture_event(unique_id: str, event_name: str, event_properties: Any):
     config = load_config('.continue/config.json')
-    if config.allow_anonymous_telemetry:
-        posthog.capture(unique_id, event_name, event_properties)
+    if not config.allow_anonymous_telemetry:
+        return
+
+    if in_codespaces:
+        event_properties['codespaces'] = True
+    posthog.capture(unique_id, event_name, event_properties)
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index 6d6246fd..237171a6 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -235,45 +235,90 @@ class DefaultModelEditCodeStep(Step):
             prompt = self._prompt.format(
                 code=rif.contents, user_request=self.user_input, file_prefix=segs[0], file_suffix=segs[1])
 
-            completion = str(await model_to_use.complete(prompt, with_history=await sdk.get_chat_context()))
-
-            eot_token = "<|endoftext|>"
-            completion = completion.removesuffix(eot_token)
-
-            # Remove tags and If it accidentally includes prefix or suffix, remove it
-            if completion.strip().startswith("```"):
-                completion = completion.strip().removeprefix("```").removesuffix("```")
-            completion = completion.replace("<file_prefix>", "").replace("<file_suffix>", "").replace(
-                "<commit_before>", "").replace("<commit_msg>", "").replace("<commit_after>", "")
-            completion = completion.removeprefix(segs[0])
-            completion = completion.removesuffix(segs[1])
+            lines = []
+            unfinished_line = ""
+            i = 0
+            original_lines = rif.contents.split("\n")
+            lines_to_highlight = []
+
+            async def add_line(i: int, line: str):
+                range = Range.from_shorthand(
+                    rif.range.start.line + i, rif.range.start.character if i == 0 else 0, rif.range.start.line + i + 1, 0)
+                await sdk.ide.applyFileSystemEdit(FileEdit(
+                    filepath=rif.filepath,
+                    range=range,
+                    replacement=line + "\n"
+                ))
+                lines_to_highlight.append(rif.range.start.line + i)
+                # await sdk.ide.highlightCode(RangeInFile(
+                #     filepath=rif.filepath,
+                #     range=range
+                # ))
+
+            async for chunk in model_to_use.stream_chat(prompt, with_history=await sdk.get_chat_context()):
+                chunk_lines = chunk.split("\n")
+                chunk_lines[0] = unfinished_line + chunk_lines[0]
+                if chunk.endswith("\n"):
+                    unfinished_line = ""
+                    chunk_lines.pop()  # because this will be an empty string
+                else:
+                    unfinished_line = chunk_lines.pop()
+                lines.extend(chunk_lines)
+
+                for line in chunk_lines:
+                    if i < len(original_lines) and line == original_lines[i]:
+                        i += 1
+                        continue
+
+                    await add_line(i, line)
+                    i += 1
+
+            # Add the unfinished line
+            if unfinished_line != "":
+                if not i < len(original_lines) or not unfinished_line == original_lines[i]:
+                    await add_line(i, unfinished_line)
+                lines.append(unfinished_line)
+                i += 1
+
+            # Remove the leftover original lines
+            while i < len(original_lines):
+                range = Range.from_shorthand(
+                    rif.range.start.line + i, rif.range.start.character, rif.range.start.line + i, len(original_lines[i]))
+                await sdk.ide.applyFileSystemEdit(FileEdit(
+                    filepath=rif.filepath,
+                    range=range,
+                    replacement=""
+                ))
+                # await sdk.ide.highlightCode(RangeInFile(
+                #     filepath=rif.filepath,
+                #     range=range
+                # ))
+                i += 1
+
+            completion = "\n".join(lines)
+            # eot_token = "<|endoftext|>"
+            # completion = completion.removesuffix(eot_token)
+
+            # # Remove tags and If it accidentally includes prefix or suffix, remove it
+            # if completion.strip().startswith("```"):
+            #     completion = completion.strip().removeprefix("```").removesuffix("```")
+            # completion = completion.replace("<file_prefix>", "").replace("<file_suffix>", "").replace(
+            #     "<commit_before>", "").replace("<commit_msg>", "").replace("<commit_after>", "")
+            # completion = completion.removeprefix(segs[0])
+            # completion = completion.removesuffix(segs[1])
 
             self._prompt_and_completion += prompt + completion
 
-            diff = list(difflib.ndiff(rif.contents.splitlines(
-                keepends=True), completion.splitlines(keepends=True)))
-
-            lines_to_highlight = set()
-            index = 0
-            for line in diff:
-                if line.startswith("-"):
-                    pass
-                elif line.startswith("+"):
-                    lines_to_highlight.add(index + rif.range.start.line)
-                    index += 1
-                elif line.startswith(" "):
-                    index += 1
-
-            await sdk.ide.applyFileSystemEdit(FileEdit(
-                filepath=rif.filepath,
-                range=rif.range,
-                replacement=completion
-            ))
+            # await sdk.ide.applyFileSystemEdit(FileEdit(
+            #     filepath=rif.filepath,
+            #     range=rif.range,
+            #     replacement=completion
+            # ))
 
             current_hl_start = None
             last_hl = None
             rifs_to_highlight = []
-            for line in sorted(list(lines_to_highlight)):
+            for line in lines_to_highlight:
                 if current_hl_start is None:
                     current_hl_start = line
                 elif line != last_hl + 1: