1 files changed, 86 insertions, 0 deletions
diff --git a/server/continuedev/plugins/steps/chroma.py b/server/continuedev/plugins/steps/chroma.py
new file mode 100644
index 00000000..f357a872
--- /dev/null
+++ b/server/continuedev/plugins/steps/chroma.py
@@ -0,0 +1,86 @@
+from textwrap import dedent
+from typing import Coroutine, Union
+
+from ...core.main import Step
+from ...core.observation import Observation
+from ...core.sdk import ContinueSDK
+from ...core.steps import EditFileStep
+from ...libs.chroma.query import ChromaIndexManager
+
+
+class CreateCodebaseIndexChroma(Step):
+    name: str = "Create Codebase Index"
+    hide: bool = True
+
+    async def describe(self, llm) -> Coroutine[str, None, None]:
+        return "Indexing the codebase..."
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        index = ChromaIndexManager(await sdk.ide.getWorkspaceDirectory())
+        if not index.check_index_exists():
+            self.hide = False
+        index.create_codebase_index()
+
+
+class AnswerQuestionChroma(Step):
+    question: str
+    _answer: Union[str, None] = None
+    name: str = "Answer Question"
+
+    async def describe(self, llm) -> Coroutine[str, None, None]:
+        if self._answer is None:
+            return f"Answering the question: {self.question}"
+        else:
+            return self._answer
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        index = ChromaIndexManager(await sdk.ide.getWorkspaceDirectory())
+        results = index.query_codebase_index(self.question)
+
+        code_snippets = ""
+
+        files = []
+        for node in results.source_nodes:
+            resource_name = list(node.node.relationships.values())[0]
+            filepath = resource_name[: resource_name.index("::")]
+            files.append(filepath)
+            code_snippets += f"""{filepath}```\n{node.node.text}\n```\n\n"""
+
+        prompt = dedent(
+            f"""Here are a few snippets of code that might be useful in answering the question:
+
+            {code_snippets}
+
+            Here is the question to answer:
+
+            {self.question}
+
+            Here is the answer:"""
+        )
+
+        answer = await sdk.models.summarize.complete(prompt)
+        # Make paths relative to the workspace directory
+        answer = answer.replace(await sdk.ide.getWorkspaceDirectory(), "")
+
+        self._answer = answer
+
+        await sdk.ide.setFileOpen(files[0])
+
+
+class EditFileChroma(Step):
+    request: str
+    hide: bool = True
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        index = ChromaIndexManager(await sdk.ide.getWorkspaceDirectory())
+        results = index.query_codebase_index(self.request)
+
+        resource_name = list(results.source_nodes[0].node.relationships.values())[0]
+        filepath = resource_name[: resource_name.index("::")]
+
+        await sdk.run_step(
+            EditFileStep(
+                filepath=filepath,
+                prompt=f"Here is the code:\n\n{{code}}\n\nHere is the user request:\n\n{self.request}\n\nHere is the code after making the requested changes:\n",
+            )
+        )