Refactoring and shortcut functions to common steps

21 files changed, 270 insertions, 240 deletions

diff --git a/continuedev/src/continuedev/libs/agent.py b/continuedev/src/continuedev/core/agent.py
index 466aa234..509a54b3 100644
--- a/continuedev/src/continuedev/libs/agent.py
+++ b/continuedev/src/continuedev/core/agent.py
@@ -2,13 +2,13 @@ import traceback
 import time
 from typing import Callable, Coroutine, List
 from ..models.filesystem_edit import FileEditWithFullContents
-from .llm import LLM
+from ..libs.llm import LLM
 from .observation import Observation
 from ..server.ide_protocol import AbstractIdeProtocolServer
-from .util.queue import AsyncSubscriptionQueue
+from ..libs.util.queue import AsyncSubscriptionQueue
 from ..models.main import ContinueBaseModel
-from .core import Policy, History, FullState, Step, HistoryNode
-from .steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
+from .main import Policy, History, FullState, Step, HistoryNode
+from ..libs.steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
 from .sdk import ContinueSDK
 
 
diff --git a/continuedev/src/continuedev/libs/env.py b/continuedev/src/continuedev/core/env.py
index d7275b41..d7275b41 100644
--- a/continuedev/src/continuedev/libs/env.py
+++ b/continuedev/src/continuedev/core/env.py
diff --git a/continuedev/src/continuedev/libs/core.py b/continuedev/src/continuedev/core/main.py
index 9d432c4b..51fcd299 100644
--- a/continuedev/src/continuedev/libs/core.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -2,7 +2,7 @@ from typing import Callable, Coroutine, Dict, Generator, List, Tuple, Union
 
 from ..models.main import ContinueBaseModel
 from pydantic import validator
-from .llm import LLM
+from ..libs.llm import LLM
 from .observation import Observation
 
 
@@ -68,14 +68,6 @@ class FullState(ContinueBaseModel):
     user_input_queue: List[str]
 
 
-class Policy(ContinueBaseModel):
-    """A rule that determines which step to take next"""
-
-    # Note that history is mutable, kinda sus
-    def next(self, history: History = History.from_empty()) -> "Step":
-        raise NotImplementedError
-
-
 class ContinueSDK:
     pass
 
@@ -84,6 +76,14 @@ class SequentialStep:
     pass
 
 
+class Policy(ContinueBaseModel):
+    """A rule that determines which step to take next"""
+
+    # Note that history is mutable, kinda sus
+    def next(self, history: History = History.from_empty()) -> "Step":
+        raise NotImplementedError
+
+
 class Step(ContinueBaseModel):
     name: str = None
     hide: bool = False
diff --git a/continuedev/src/continuedev/libs/observation.py b/continuedev/src/continuedev/core/observation.py
index fef04311..fef04311 100644
--- a/continuedev/src/continuedev/libs/observation.py
+++ b/continuedev/src/continuedev/core/observation.py
diff --git a/continuedev/src/continuedev/libs/policy.py b/continuedev/src/continuedev/core/policy.py
index 586eaebe..9e6abf14 100644
--- a/continuedev/src/continuedev/libs/policy.py
+++ b/continuedev/src/continuedev/core/policy.py
@@ -1,11 +1,13 @@
 from typing import List, Tuple, Type
 
-from .steps.ty import CreatePipelineStep
-from .core import Step, Validator, Policy, History
+from ..models.main import ContinueBaseModel
+
+from ..libs.steps.ty import CreatePipelineStep
+from .main import Step, Validator, History, Policy
 from .observation import Observation, TracebackObservation, UserInputObservation
-from .steps.main import EditCodeStep, EditHighlightedCodeStep, SolveTracebackStep, RunCodeStep, FasterEditHighlightedCodeStep
-from .steps.nate import WritePytestsStep, CreateTableStep
-from .steps.chroma import AnswerQuestionChroma, EditFileChroma
+from ..libs.steps.main import EditHighlightedCodeStep, SolveTracebackStep, RunCodeStep
+from ..libs.steps.nate import WritePytestsStep, CreateTableStep
+from ..libs.steps.chroma import AnswerQuestionChroma, EditFileChroma
 
 
 class DemoPolicy(Policy):
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
new file mode 100644
index 00000000..ff62a2b1
--- /dev/null
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -0,0 +1,62 @@
+from typing import Coroutine, Union
+from ..models.filesystem_edit import FileSystemEdit
+from ..models.filesystem import RangeInFile
+from ..libs.llm import LLM
+from .observation import Observation
+from ..server.ide_protocol import AbstractIdeProtocolServer
+from .main import History, Step
+from ..libs.steps.core.core import *
+
+
+class Agent:
+    pass
+
+
+class ContinueSDKSteps:
+    def __init__(self, sdk: "ContinueSDK"):
+        self.sdk = sdk
+
+
+class ContinueSDK:
+    """The SDK provided as parameters to a step"""
+    llm: LLM
+    ide: AbstractIdeProtocolServer
+    steps: ContinueSDKSteps
+    __agent: Agent
+
+    def __init__(self, agent: Agent, llm: Union[LLM, None] = None):
+        if llm is None:
+            self.llm = agent.llm
+        else:
+            self.llm = llm
+        self.ide = agent.ide
+        self.__agent = agent
+        self.steps = ContinueSDKSteps(self)
+
+    @property
+    def history(self) -> History:
+        return self.__agent.history
+
+    async def run_step(self, step: Step) -> Coroutine[Observation, None, None]:
+        return await self.__agent._run_singular_step(step)
+
+    async def apply_filesystem_edit(self, edit: FileSystemEdit):
+        await self.run_step(FileSystemEditStep(edit=edit))
+
+    async def wait_for_user_input(self) -> str:
+        return await self.__agent.wait_for_user_input()
+
+    async def wait_for_user_confirmation(self, prompt: str):
+        return await self.run_step(WaitForUserConfirmationStep(prompt=prompt))
+
+    async def run(self, commands: List[str] | str, cwd: str = None):
+        commands = commands if isinstance(commands, List) else [commands]
+        return self.run_step(ShellCommandsStep(commands=commands, cwd=cwd))
+
+    async def edit_file(self, filename: str, prompt: str):
+        await self.ide.setFileOpen(filename)
+        contents = await self.ide.readFile(filename)
+        await self.run_step(EditCodeStep(
+            range_in_files=[RangeInFile.from_entire_file(filename, contents)],
+            prompt=f'Here is the code before:\n\n{{code}}\n\nHere is the user request:\n\n{prompt}\n\nHere is the code edited to perfectly solve the user request:\n\n'
+        ))
diff --git a/continuedev/src/continuedev/libs/sdk.py b/continuedev/src/continuedev/libs/sdk.py
deleted file mode 100644
index 0295bf35..00000000
--- a/continuedev/src/continuedev/libs/sdk.py
+++ /dev/null
@@ -1,39 +0,0 @@
-from typing import Coroutine, Union
-from ..models.filesystem_edit import FileSystemEdit
-from .llm import LLM
-from .observation import Observation
-from ..server.ide_protocol import AbstractIdeProtocolServer
-from .core import History, Step
-from .steps.core.core import *
-
-
-class Agent:
-    pass
-
-
-class ContinueSDK:
-    """The SDK provided as parameters to a step"""
-    llm: LLM
-    ide: AbstractIdeProtocolServer
-    __agent: Agent
-
-    def __init__(self, agent: Agent, llm: Union[LLM, None] = None):
-        if llm is None:
-            self.llm = agent.llm
-        else:
-            self.llm = llm
-        self.ide = agent.ide
-        self.__agent = agent
-
-    @property
-    def history(self) -> History:
-        return self.__agent.history
-
-    async def run_step(self, step: Step) -> Coroutine[Observation, None, None]:
-        return await self.__agent._run_singular_step(step)
-
-    async def apply_filesystem_edit(self, edit: FileSystemEdit):
-        await self.run_step(FileSystemEditStep(edit=edit))
-
-    async def wait_for_user_input(self) -> str:
-        return await self.__agent.wait_for_user_input()
diff --git a/continuedev/src/continuedev/libs/steps/chroma.py b/continuedev/src/continuedev/libs/steps/chroma.py
index 2d8742e8..f13a2bab 100644
--- a/continuedev/src/continuedev/libs/steps/chroma.py
+++ b/continuedev/src/continuedev/libs/steps/chroma.py
@@ -1,11 +1,10 @@
 from textwrap import dedent
 from typing import Coroutine, Union
-from ...models.filesystem_edit import AddDirectory, AddFile
-from ..observation import Observation, TextObservation
-from ..core import Step, ContinueSDK
-from .main import EditCodeStep, EditFileStep, RunCommandStep, WaitForUserConfirmationStep
+from ...core.observation import Observation, TextObservation
+from ...core.main import Step, ContinueSDK
+from .core.core import EditFileStep
 from ..chroma.query import query_codebase_index
-from .main import EditFileStep
+from .core.core import EditFileStep
 
 
 class AnswerQuestionChroma(Step):
diff --git a/continuedev/src/continuedev/libs/steps/core/core.py b/continuedev/src/continuedev/libs/steps/core/core.py
index 23086b53..0338d635 100644
--- a/continuedev/src/continuedev/libs/steps/core/core.py
+++ b/continuedev/src/continuedev/libs/steps/core/core.py
@@ -1,11 +1,16 @@
 # These steps are depended upon by ContinueSDK
-from typing import Callable, Coroutine, Dict, Generator, List, Tuple, Union
+import subprocess
+from textwrap import dedent
+from typing import Coroutine, List, Union
+from ...llm.prompt_utils import MarkdownStyleEncoderDecoder
+
+from ...util.traceback_parsers import parse_python_traceback
 
 from ....models.filesystem_edit import EditDiff, FileEditWithFullContents, FileSystemEdit
-from ....models.filesystem import FileSystem
+from ....models.filesystem import FileSystem, RangeInFile, RangeInFileWithContents
 from ...llm import LLM
-from ...observation import Observation, UserInputObservation
-from ...core import Step
+from ....core.observation import Observation, TextObservation, TracebackObservation, UserInputObservation
+from ....core.main import Step
 
 
 class ContinueSDK:
@@ -42,6 +47,98 @@ class FileSystemEditStep(ReversibleStep):
         # Where and when should file saves happen?
 
 
+def ShellCommandsStep(Step):
+    cmds: List[str]
+    cwd: str | None = None
+    name: str = "Run Shell Commands"
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        return "\n".join(self.cmds)
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        cwd = await sdk.ide.getWorkspaceDirectory() if self.cwd is None else self.cwd
+
+        process = subprocess.Popen(
+            '/bin/bash', stdin=subprocess.PIPE, stdout=subprocess.PIPE, cwd=cwd)
+
+        stdin_input = "\n".join(self.cmds)
+        out, err = process.communicate(stdin_input.encode())
+
+        # If it fails, return the error
+        if err is not None and err != "":
+            return TextObservation(text=err)
+
+        return None
+
+
+class EditCodeStep(Step):
+    # Might make an even more specific atomic step, which is "apply file edit"
+    range_in_files: List[RangeInFile]
+    prompt: str  # String with {code} somewhere
+    name: str = "Edit code"
+
+    _edit_diffs: Union[List[EditDiff], None] = None
+    _prompt: Union[str, None] = None
+    _completion: Union[str, None] = None
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        if self._edit_diffs is None:
+            return "Editing files: " + ", ".join(map(lambda rif: rif.filepath, self.range_in_files))
+        elif len(self._edit_diffs) == 0:
+            return "No edits made"
+        else:
+            return llm.complete(dedent(f"""{self._prompt}{self._completion}
+
+                Maximally concise summary of changes in bullet points (can use markdown):
+            """))
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        rif_with_contents = []
+        for range_in_file in self.range_in_files:
+            file_contents = await sdk.ide.readRangeInFile(range_in_file)
+            rif_with_contents.append(
+                RangeInFileWithContents.from_range_in_file(range_in_file, file_contents))
+        enc_dec = MarkdownStyleEncoderDecoder(rif_with_contents)
+        code_string = enc_dec.encode()
+        prompt = self.prompt.format(code=code_string)
+
+        completion = sdk.llm.complete(prompt)
+
+        # Temporarily doing this to generate description.
+        self._prompt = prompt
+        self._completion = completion
+
+        file_edits = enc_dec.decode(completion)
+
+        self._edit_diffs = []
+        for file_edit in file_edits:
+            diff = await sdk.apply_filesystem_edit(file_edit)
+            self._edit_diffs.append(diff)
+
+        for filepath in set([file_edit.filepath for file_edit in file_edits]):
+            await sdk.ide.saveFile(filepath)
+            await sdk.ide.setFileOpen(filepath)
+
+        return None
+
+
+class EditFileStep(Step):
+    filepath: str
+    prompt: str
+    hide: bool = True
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        return "Editing file: " + self.filepath
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        file_contents = await sdk.ide.readFile(self.filepath)
+        await sdk.run_step(EditCodeStep(
+            range_in_files=[RangeInFile.from_entire_file(
+                self.filepath, file_contents)],
+            prompt=self.prompt
+        ))
+
+
 class ManualEditStep(ReversibleStep):
     edit_diff: EditDiff
     hide: bool = True
@@ -89,3 +186,31 @@ class UserInputStep(Step):
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[UserInputObservation, None, None]:
         return UserInputObservation(user_input=self.user_input)
+
+
+class WaitForUserInputStep(Step):
+    prompt: str
+    name: str = "Waiting for user input"
+
+    _description: Union[str, None] = None
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        return self.prompt
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        self._description = self.prompt
+        resp = await sdk.wait_for_user_input()
+        return TextObservation(text=resp)
+
+
+class WaitForUserConfirmationStep(Step):
+    prompt: str
+    name: str = "Waiting for user confirmation"
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        return self.prompt
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        self._description = self.prompt
+        resp = await sdk.wait_for_user_input()
+        return TextObservation(text=resp)
diff --git a/continuedev/src/continuedev/libs/steps/draft/abstract_method.py b/continuedev/src/continuedev/libs/steps/draft/abstract_method.py
index 927d93fd..f3131c4b 100644
--- a/continuedev/src/continuedev/libs/steps/draft/abstract_method.py
+++ b/continuedev/src/continuedev/libs/steps/draft/abstract_method.py
@@ -1,4 +1,5 @@
-from ...core import ContinueSDK, Step
+from ....core.sdk import ContinueSDK
+from ....core.main import Step
 
 
 class ImplementAbstractMethodStep(Step):
diff --git a/continuedev/src/continuedev/libs/steps/draft/dlt.py b/continuedev/src/continuedev/libs/steps/draft/dlt.py
index 608f089a..5ba5692a 100644
--- a/continuedev/src/continuedev/libs/steps/draft/dlt.py
+++ b/continuedev/src/continuedev/libs/steps/draft/dlt.py
@@ -1,6 +1,7 @@
 from textwrap import dedent
 from ....models.filesystem_edit import AddFile
-from ...core import Step, ContinueSDK
+from ....core.main import Step
+from ....core.sdk import ContinueSDK
 from ..main import WaitForUserInputStep
 
 
diff --git a/continuedev/src/continuedev/libs/steps/draft/redux.py b/continuedev/src/continuedev/libs/steps/draft/redux.py
index 52a8fbd8..efaa9ba4 100644
--- a/continuedev/src/continuedev/libs/steps/draft/redux.py
+++ b/continuedev/src/continuedev/libs/steps/draft/redux.py
@@ -1,7 +1,6 @@
-from textwrap import dedent
-from ....models.filesystem_edit import AddFile
-from ...core import Step, ContinueSDK
-from ..main import WaitForUserInputStep, EditFileStep
+from ....core.main import Step
+from ....core.sdk import ContinueSDK
+from ..core.core import EditFileStep
 
 
 class EditReduxStateStep(Step):
diff --git a/continuedev/src/continuedev/libs/steps/draft/typeorm.py b/continuedev/src/continuedev/libs/steps/draft/typeorm.py
index 9d058f1e..d06a6fb4 100644
--- a/continuedev/src/continuedev/libs/steps/draft/typeorm.py
+++ b/continuedev/src/continuedev/libs/steps/draft/typeorm.py
@@ -1,5 +1,6 @@
 from textwrap import dedent
-from ...core import Step, ContinueSDK
+from ....core.main import Step
+from ....core.sdk import ContinueSDK
 
 
 class CreateTableStep(Step):
diff --git a/continuedev/src/continuedev/libs/steps/main.py b/continuedev/src/continuedev/libs/steps/main.py
index 555c1180..f28cb23f 100644
--- a/continuedev/src/continuedev/libs/steps/main.py
+++ b/continuedev/src/continuedev/libs/steps/main.py
@@ -1,18 +1,39 @@
-import time
 from typing import Callable, Coroutine, List, Union
 
+from ..util.traceback_parsers import parse_python_traceback
 from ..llm import LLM
 from ...models.main import Traceback, Range
 from ...models.filesystem_edit import EditDiff, FileEdit
 from ...models.filesystem import RangeInFile, RangeInFileWithContents
-from ..observation import Observation, TextObservation
+from ...core.observation import Observation, TextObservation, TracebackObservation
 from ..llm.prompt_utils import MarkdownStyleEncoderDecoder
 from textwrap import dedent
-from ..core import History, Policy, Step, ContinueSDK, Observation
+from ...core.main import History, Policy, Step, ContinueSDK, Observation
 import subprocess
-from ..util.traceback_parsers import parse_python_traceback
-from ..observation import TracebackObservation
 import json
+from .core.core import EditCodeStep
+
+
+class RunCodeStep(Step):
+    cmd: str
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        return f"Ran command: `{self.cmd}`"
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        result = subprocess.run(
+            self.cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        stdout = result.stdout.decode("utf-8")
+        stderr = result.stderr.decode("utf-8")
+        print(stdout, stderr)
+
+        # If it fails, return the error
+        tb = parse_python_traceback(stdout) or parse_python_traceback(stderr)
+        if tb:
+            return TracebackObservation(traceback=tb)
+        else:
+            self.hide = True
+            return None
 
 
 class RunPolicyUntilDoneStep(Step):
@@ -51,149 +72,6 @@ class RunCommandStep(Step):
             return TextObservation(text=stdout)
 
 
-def ShellCommandsStep(Step):
-    cmds: List[str]
-    name: str = "Run Shell Commands"
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        return "\n".join(self.cmds)
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        cwd = await sdk.ide.getWorkspaceDirectory()
-
-        process = subprocess.Popen(
-            '/bin/bash', stdin=subprocess.PIPE, stdout=subprocess.PIPE, cwd=cwd)
-
-        stdin_input = "\n".join(self.cmds)
-        out, err = process.communicate(stdin_input.encode())
-
-        # TODO: How to await??
-
-        # If it fails, return the error
-        if err is not None and err != "":
-            return TextObservation(text=err)
-
-        return None
-
-
-class WaitForUserInputStep(Step):
-    prompt: str
-    name: str = "Waiting for user input"
-
-    _description: Union[str, None] = None
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        return self.prompt
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        self._description = self.prompt
-        resp = await sdk.wait_for_user_input()
-        return TextObservation(text=resp)
-
-
-class WaitForUserConfirmationStep(Step):
-    prompt: str
-    name: str = "Waiting for user confirmation"
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        return self.prompt
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        self._description = self.prompt
-        resp = await sdk.wait_for_user_input()
-        return TextObservation(text=resp)
-
-
-class RunCodeStep(Step):
-    cmd: str
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        return f"Ran command: `{self.cmd}`"
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        result = subprocess.run(
-            self.cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        stdout = result.stdout.decode("utf-8")
-        stderr = result.stderr.decode("utf-8")
-        print(stdout, stderr)
-
-        # If it fails, return the error
-        tb = parse_python_traceback(stdout) or parse_python_traceback(stderr)
-        if tb:
-            return TracebackObservation(traceback=tb)
-        else:
-            self.hide = True
-            return None
-
-
-class EditCodeStep(Step):
-    # Might make an even more specific atomic step, which is "apply file edit"
-    range_in_files: List[RangeInFile]
-    prompt: str  # String with {code} somewhere
-    name: str = "Edit code"
-
-    _edit_diffs: Union[List[EditDiff], None] = None
-    _prompt: Union[str, None] = None
-    _completion: Union[str, None] = None
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        if self._edit_diffs is None:
-            return "Editing files: " + ", ".join(map(lambda rif: rif.filepath, self.range_in_files))
-        elif len(self._edit_diffs) == 0:
-            return "No edits made"
-        else:
-            return llm.complete(dedent(f"""{self._prompt}{self._completion}
-
-                Maximally concise summary of changes in bullet points (can use markdown):
-            """))
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        rif_with_contents = []
-        for range_in_file in self.range_in_files:
-            file_contents = await sdk.ide.readRangeInFile(range_in_file)
-            rif_with_contents.append(
-                RangeInFileWithContents.from_range_in_file(range_in_file, file_contents))
-        enc_dec = MarkdownStyleEncoderDecoder(rif_with_contents)
-        code_string = enc_dec.encode()
-        prompt = self.prompt.format(code=code_string)
-
-        completion = sdk.llm.complete(prompt)
-
-        # Temporarily doing this to generate description.
-        self._prompt = prompt
-        self._completion = completion
-
-        file_edits = enc_dec.decode(completion)
-
-        self._edit_diffs = []
-        for file_edit in file_edits:
-            diff = await sdk.apply_filesystem_edit(file_edit)
-            self._edit_diffs.append(diff)
-
-        for filepath in set([file_edit.filepath for file_edit in file_edits]):
-            await sdk.ide.saveFile(filepath)
-            await sdk.ide.setFileOpen(filepath)
-
-        return None
-
-
-class EditFileStep(Step):
-    filepath: str
-    prompt: str
-    hide: bool = True
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        return "Editing file: " + self.filepath
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        file_contents = await sdk.ide.readFile(self.filepath)
-        await sdk.run_step(EditCodeStep(
-            range_in_files=[RangeInFile.from_entire_file(
-                self.filepath, file_contents)],
-            prompt=self.prompt
-        ))
-
-
 class FasterEditHighlightedCodeStep(Step):
     user_input: str
     hide = True
diff --git a/continuedev/src/continuedev/libs/steps/migration.py b/continuedev/src/continuedev/libs/steps/migration.py
index 04296836..f044a60f 100644
--- a/continuedev/src/continuedev/libs/steps/migration.py
+++ b/continuedev/src/continuedev/libs/steps/migration.py
@@ -3,7 +3,7 @@
 
 from ...models.filesystem import RangeInFile
 from .main import EditCodeStep, RunCommandStep
-from ..core import Step
+from ...core.main import Step
 
 
 class MigrationStep(Step):
diff --git a/continuedev/src/continuedev/libs/steps/nate.py b/continuedev/src/continuedev/libs/steps/nate.py
index 80436fa4..a0e728e5 100644
--- a/continuedev/src/continuedev/libs/steps/nate.py
+++ b/continuedev/src/continuedev/libs/steps/nate.py
@@ -1,14 +1,13 @@
-import asyncio
 from textwrap import dedent
 import time
 from typing import Coroutine, Union
 
-from ...models.main import Range
 from ...models.filesystem import RangeInFile
 from ...models.filesystem_edit import AddDirectory, AddFile
-from ..observation import Observation, TextObservation
-from ..core import Step, ContinueSDK
-from .main import EditCodeStep, EditFileStep, RunCommandStep, WaitForUserConfirmationStep
+from ...core.observation import Observation, TextObservation
+from ...core.main import Step, ContinueSDK
+from .main import RunCommandStep
+from .core.core import WaitForUserConfirmationStep, EditCodeStep, EditFileStep
 import os
 
 
diff --git a/continuedev/src/continuedev/libs/steps/pytest.py b/continuedev/src/continuedev/libs/steps/pytest.py
index e53eb465..b4e6dfd2 100644
--- a/continuedev/src/continuedev/libs/steps/pytest.py
+++ b/continuedev/src/continuedev/libs/steps/pytest.py
@@ -1,6 +1,6 @@
 from textwrap import dedent
 from ...models.filesystem_edit import AddDirectory, AddFile
-from ..core import Step, ContinueSDK
+from ...core.main import Step, ContinueSDK
 import os
 
 
diff --git a/continuedev/src/continuedev/libs/steps/ty.py b/continuedev/src/continuedev/libs/steps/ty.py
index 1eb6271d..5ff03f04 100644
--- a/continuedev/src/continuedev/libs/steps/ty.py
+++ b/continuedev/src/continuedev/libs/steps/ty.py
@@ -2,9 +2,11 @@ import subprocess
 from ...models.main import Position, Range
 from ...models.filesystem import RangeInFile
 from ...models.filesystem_edit import AddDirectory, AddFile, FileEdit
-from ..observation import DictObservation
-from ..core import History, Step, ContinueSDK, Policy
-from .main import EditCodeStep, RunCommandStep, WaitForUserInputStep, WaitForUserConfirmationStep
+from ...core.observation import DictObservation
+from ...core.main import History, Step, Policy
+from ...core.sdk import ContinueSDK
+from .main import RunCommandStep
+from ..steps.core.core import EditCodeStep, WaitForUserConfirmationStep, WaitForUserInputStep
 
 source_name = "weather_api"
 
diff --git a/continuedev/src/continuedev/plugins/policy/libs/alternate.py b/continuedev/src/continuedev/plugins/policy/libs/alternate.py
index 5bfbb821..3087c059 100644
--- a/continuedev/src/continuedev/plugins/policy/libs/alternate.py
+++ b/continuedev/src/continuedev/plugins/policy/libs/alternate.py
@@ -1,7 +1,5 @@
 from plugins import policy
-from ....libs.observation import Observation
-from ....libs.steps import Step
-from ....libs.core import History
+from ....core.main import History, Step
 
 
 class AlternatingPolicy:
diff --git a/continuedev/src/continuedev/plugins/step/hookspecs.py b/continuedev/src/continuedev/plugins/step/hookspecs.py
index 4309bad3..a5714fc5 100644
--- a/continuedev/src/continuedev/plugins/step/hookspecs.py
+++ b/continuedev/src/continuedev/plugins/step/hookspecs.py
@@ -1,6 +1,8 @@
 from typing import Coroutine
 import pluggy
-from ...libs.core import ContinueSDK, Step, Observation
+from ...core.main import Step
+from ...core.observation import Observation
+from ...core.sdk import ContinueSDK
 
 hookspec = pluggy.HookspecMarker("continue.step")
 
diff --git a/continuedev/src/continuedev/server/notebook.py b/continuedev/src/continuedev/server/notebook.py
index 5eb151d7..bfd7a09c 100644
--- a/continuedev/src/continuedev/server/notebook.py
+++ b/continuedev/src/continuedev/server/notebook.py
@@ -5,14 +5,14 @@ from pydantic import BaseModel
 from uvicorn.main import Server
 
 from ..models.filesystem_edit import FileEditWithFullContents
-from ..libs.policy import DemoPolicy
-from ..libs.core import FullState, History, Step
-from ..libs.agent import Agent
+from ..core.policy import DemoPolicy
+from ..core.main import FullState, History, Step
+from ..core.agent import Agent
 from ..libs.steps.nate import ImplementAbstractMethodStep
-from ..libs.observation import Observation
+from ..core.observation import Observation
 from ..libs.llm.openai import OpenAI
 from .ide_protocol import AbstractIdeProtocolServer
-from ..libs.env import openai_api_key
+from ..core.env import openai_api_key
 import asyncio
 import nest_asyncio
 nest_asyncio.apply()