Merge pull request #11 from continuedev/syntactic-sugar

Syntactic sugar

23 files changed, 687 insertions, 585 deletions

diff --git a/continuedev/src/continuedev/core/agent.py b/continuedev/src/continuedev/core/agent.py
new file mode 100644
index 00000000..509a54b3
--- /dev/null
+++ b/continuedev/src/continuedev/core/agent.py
@@ -0,0 +1,180 @@
+import traceback
+import time
+from typing import Callable, Coroutine, List
+from ..models.filesystem_edit import FileEditWithFullContents
+from ..libs.llm import LLM
+from .observation import Observation
+from ..server.ide_protocol import AbstractIdeProtocolServer
+from ..libs.util.queue import AsyncSubscriptionQueue
+from ..models.main import ContinueBaseModel
+from .main import Policy, History, FullState, Step, HistoryNode
+from ..libs.steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
+from .sdk import ContinueSDK
+
+
+class Agent(ContinueBaseModel):
+    llm: LLM
+    policy: Policy
+    ide: AbstractIdeProtocolServer
+    history: History = History.from_empty()
+    continue_sdk: "ContinueSDK"
+    _on_update_callbacks: List[Callable[[FullState], None]] = []
+
+    _active: bool = False
+    _should_halt: bool = False
+    _main_user_input_queue: List[str] = []
+
+    _user_input_queue = AsyncSubscriptionQueue()
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def get_full_state(self) -> FullState:
+        return FullState(history=self.history, active=self._active, user_input_queue=self._main_user_input_queue)
+
+    def on_update(self, callback: Callable[["FullState"], None]):
+        """Subscribe to changes to state"""
+        self._on_update_callbacks.append(callback)
+
+    def update_subscribers(self):
+        full_state = self.get_full_state()
+        for callback in self._on_update_callbacks:
+            callback(full_state)
+
+    def __get_step_params(self, step: "Step"):
+        return ContinueSDK(agent=self, llm=self.llm.with_system_message(step.system_message))
+
+    def give_user_input(self, input: str, index: int):
+        self._user_input_queue.post(index, input)
+
+    async def wait_for_user_input(self) -> str:
+        self._active = False
+        self.update_subscribers()
+        await self._user_input_queue.get(self.history.current_index)
+        self._active = True
+        self.update_subscribers()
+
+    _manual_edits_buffer: List[FileEditWithFullContents] = []
+
+    async def reverse_to_index(self, index: int):
+        try:
+            while self.history.get_current_index() >= index:
+                current_step = self.history.get_current().step
+                self.history.step_back()
+                if issubclass(current_step.__class__, ReversibleStep):
+                    await current_step.reverse(self.__get_step_params(current_step))
+
+                self.update_subscribers()
+        except Exception as e:
+            print(e)
+
+    def handle_manual_edits(self, edits: List[FileEditWithFullContents]):
+        for edit in edits:
+            self._manual_edits_buffer.append(edit)
+            # TODO: You're storing a lot of unecessary data here. Can compress into EditDiffs on the spot, and merge.
+            # self._manual_edits_buffer = merge_file_edit(self._manual_edits_buffer, edit)
+
+    def handle_traceback(self, traceback: str):
+        raise NotImplementedError
+
+    _step_depth: int = 0
+
+    async def _run_singular_step(self, step: "Step", is_future_step: bool = False) -> Coroutine[Observation, None, None]:
+        if not is_future_step:
+            # Check manual edits buffer, clear out if needed by creating a ManualEditStep
+            if len(self._manual_edits_buffer) > 0:
+                manualEditsStep = ManualEditStep.from_sequence(
+                    self._manual_edits_buffer)
+                self._manual_edits_buffer = []
+                await self._run_singular_step(manualEditsStep)
+
+        # Update history - do this first so we get top-first tree ordering
+        self.history.add_node(HistoryNode(
+            step=step, observation=None, depth=self._step_depth))
+
+        # Run step
+        self._step_depth += 1
+        observation = await step(self.__get_step_params(step))
+        self._step_depth -= 1
+
+        # Add observation to history
+        self.history.get_current().observation = observation
+
+        # Update its description
+        step._set_description(await step.describe(self.llm))
+
+        # Call all subscribed callbacks
+        self.update_subscribers()
+
+        return observation
+
+    async def run_from_step(self, step: "Step"):
+        # if self._active:
+        #     raise RuntimeError("Agent is already running")
+        self._active = True
+
+        next_step = step
+        is_future_step = False
+        while not (next_step is None or self._should_halt):
+            try:
+                if is_future_step:
+                    # If future step, then we are replaying and need to delete the step from history so it can be replaced
+                    self.history.remove_current_and_substeps()
+
+                observation = await self._run_singular_step(next_step, is_future_step)
+                if next_step := self.policy.next(self.history):
+                    is_future_step = False
+                elif next_step := self.history.take_next_step():
+                    is_future_step = True
+                else:
+                    next_step = None
+
+            except Exception as e:
+                print(
+                    f"Error while running step: \n{''.join(traceback.format_tb(e.__traceback__))}\n{e}")
+                next_step = None
+
+        self._active = False
+
+        # Doing this so active can make it to the frontend after steps are done. But want better state syncing tools
+        for callback in self._on_update_callbacks:
+            callback(None)
+
+    async def run_from_observation(self, observation: Observation):
+        next_step = self.policy.next(self.history)
+        await self.run_from_step(next_step)
+
+    async def run_policy(self):
+        first_step = self.policy.next(self.history)
+        await self.run_from_step(first_step)
+
+    async def _request_halt(self):
+        if self._active:
+            self._should_halt = True
+            while self._active:
+                time.sleep(0.1)
+        self._should_halt = False
+        return None
+
+    async def accept_user_input(self, user_input: str):
+        self._main_user_input_queue.append(user_input)
+        self.update_subscribers()
+
+        if len(self._main_user_input_queue) > 1:
+            return
+
+        # await self._request_halt()
+        # Just run the step that takes user input, and
+        # then up to the policy to decide how to deal with it.
+        self._main_user_input_queue.pop(0)
+        self.update_subscribers()
+        await self.run_from_step(UserInputStep(user_input=user_input))
+
+        while len(self._main_user_input_queue) > 0:
+            await self.run_from_step(UserInputStep(
+                user_input=self._main_user_input_queue.pop(0)))
+
+    async def accept_refinement_input(self, user_input: str, index: int):
+        await self._request_halt()
+        await self.reverse_to_index(index)
+        await self.run_from_step(UserInputStep(user_input=user_input))
diff --git a/continuedev/src/continuedev/core/env.py b/continuedev/src/continuedev/core/env.py
new file mode 100644
index 00000000..d7275b41
--- /dev/null
+++ b/continuedev/src/continuedev/core/env.py
@@ -0,0 +1,7 @@
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+
+openai_api_key = os.getenv("OPENAI_API_KEY")
diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
new file mode 100644
index 00000000..51fcd299
--- /dev/null
+++ b/continuedev/src/continuedev/core/main.py
@@ -0,0 +1,148 @@
+from typing import Callable, Coroutine, Dict, Generator, List, Tuple, Union
+
+from ..models.main import ContinueBaseModel
+from pydantic import validator
+from ..libs.llm import LLM
+from .observation import Observation
+
+
+class HistoryNode(ContinueBaseModel):
+    """A point in history, a list of which make up History"""
+    step: "Step"
+    observation: Union[Observation, None]
+    depth: int
+
+
+class History(ContinueBaseModel):
+    """A history of steps taken and their results"""
+    timeline: List[HistoryNode]
+    current_index: int
+
+    def add_node(self, node: HistoryNode):
+        self.timeline.insert(self.current_index + 1, node)
+        self.current_index += 1
+
+    def get_current(self) -> Union[HistoryNode, None]:
+        if self.current_index < 0:
+            return None
+        return self.timeline[self.current_index]
+
+    def remove_current_and_substeps(self):
+        self.timeline.pop(self.current_index)
+        while self.get_current() is not None and self.get_current().depth > 0:
+            self.timeline.pop(self.current_index)
+
+    def take_next_step(self) -> Union["Step", None]:
+        if self.has_future():
+            self.current_index += 1
+            current_state = self.get_current()
+            if current_state is None:
+                return None
+            return current_state.step
+        return None
+
+    def get_current_index(self) -> int:
+        return self.current_index
+
+    def has_future(self) -> bool:
+        return self.current_index < len(self.timeline) - 1
+
+    def step_back(self):
+        self.current_index -= 1
+
+    def last_observation(self) -> Union[Observation, None]:
+        state = self.get_current()
+        if state is None:
+            return None
+        return state.observation
+
+    @classmethod
+    def from_empty(cls):
+        return cls(timeline=[], current_index=-1)
+
+
+class FullState(ContinueBaseModel):
+    """A full state of the program, including the history"""
+    history: History
+    active: bool
+    user_input_queue: List[str]
+
+
+class ContinueSDK:
+    pass
+
+
+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
+    _description: Union[str, None] = None
+
+    system_message: Union[str, None] = None
+
+    class Config:
+        copy_on_model_validation = False
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        if self._description is not None:
+            return self._description
+        return "Running step: " + self.name
+
+    def _set_description(self, description: str):
+        self._description = description
+
+    def dict(self, *args, **kwargs):
+        d = super().dict(*args, **kwargs)
+        if self._description is not None:
+            d["description"] = self._description
+        else:
+            d["description"] = self.name
+        return d
+
+    @validator("name", pre=True, always=True)
+    def name_is_class_name(cls, name):
+        if name is None:
+            return cls.__name__
+        return name
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        raise NotImplementedError
+
+    async def __call__(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        return await self.run(sdk)
+
+    def __rshift__(self, other: "Step"):
+        steps = []
+        if isinstance(self, SequentialStep):
+            steps = self.steps
+        else:
+            steps.append(self)
+        if isinstance(other, SequentialStep):
+            steps += other.steps
+        else:
+            steps.append(other)
+        return SequentialStep(steps=steps)
+
+
+class ValidatorObservation(Observation):
+    passed: bool
+    observation: Observation
+
+
+class Validator(Step):
+    def run(self, sdk: ContinueSDK) -> ValidatorObservation:
+        raise NotImplementedError
+
+
+HistoryNode.update_forward_refs()
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/core.py b/continuedev/src/continuedev/libs/core.py
deleted file mode 100644
index 6a8a83ba..00000000
--- a/continuedev/src/continuedev/libs/core.py
+++ /dev/null
@@ -1,423 +0,0 @@
-import traceback
-import time
-from typing import Callable, Coroutine, Dict, Generator, List, Tuple, Union
-from ..models.filesystem_edit import EditDiff, FileEdit, FileEditWithFullContents, FileSystemEdit
-from ..models.filesystem import FileSystem
-from pydantic import BaseModel, parse_file_as, validator
-from .llm import LLM
-from .observation import Observation, UserInputObservation
-from ..server.ide_protocol import AbstractIdeProtocolServer
-from .util.queue import AsyncSubscriptionQueue
-
-
-class ContinueBaseModel(BaseModel):
-    class Config:
-        underscore_attrs_are_private = True
-
-
-class HistoryNode(ContinueBaseModel):
-    """A point in history, a list of which make up History"""
-    step: "Step"
-    observation: Union[Observation, None]
-    depth: int
-
-
-class History(ContinueBaseModel):
-    """A history of steps taken and their results"""
-    timeline: List[HistoryNode]
-    current_index: int
-
-    def add_node(self, node: HistoryNode):
-        self.timeline.insert(self.current_index + 1, node)
-        self.current_index += 1
-
-    def get_current(self) -> Union[HistoryNode, None]:
-        if self.current_index < 0:
-            return None
-        return self.timeline[self.current_index]
-
-    def remove_current_and_substeps(self):
-        self.timeline.pop(self.current_index)
-        while self.get_current() is not None and self.get_current().depth > 0:
-            self.timeline.pop(self.current_index)
-
-    def take_next_step(self) -> Union["Step", None]:
-        if self.has_future():
-            self.current_index += 1
-            current_state = self.get_current()
-            if current_state is None:
-                return None
-            return current_state.step
-        return None
-
-    def get_current_index(self) -> int:
-        return self.current_index
-
-    def has_future(self) -> bool:
-        return self.current_index < len(self.timeline) - 1
-
-    def step_back(self):
-        self.current_index -= 1
-
-    def last_observation(self) -> Union[Observation, None]:
-        state = self.get_current()
-        if state is None:
-            return None
-        return state.observation
-
-    @classmethod
-    def from_empty(cls):
-        return cls(timeline=[], current_index=-1)
-
-
-class FullState(ContinueBaseModel):
-    """A full state of the program, including the history"""
-    history: History
-    active: bool
-    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:
-    """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()
-
-
-class Agent(ContinueBaseModel):
-    llm: LLM
-    policy: Policy
-    ide: AbstractIdeProtocolServer
-    history: History = History.from_empty()
-    _on_update_callbacks: List[Callable[["FullState"], None]] = []
-
-    _active: bool = False
-    _should_halt: bool = False
-    _main_user_input_queue: List[str] = []
-
-    _user_input_queue = AsyncSubscriptionQueue()
-
-    class Config:
-        arbitrary_types_allowed = True
-
-    def get_full_state(self) -> FullState:
-        return FullState(history=self.history, active=self._active, user_input_queue=self._main_user_input_queue)
-
-    def on_update(self, callback: Callable[["FullState"], None]):
-        """Subscribe to changes to state"""
-        self._on_update_callbacks.append(callback)
-
-    def update_subscribers(self):
-        full_state = self.get_full_state()
-        for callback in self._on_update_callbacks:
-            callback(full_state)
-
-    def __get_step_params(self, step: "Step"):
-        return ContinueSDK(agent=self, llm=self.llm.with_system_message(step.system_message))
-
-    def give_user_input(self, input: str, index: int):
-        self._user_input_queue.post(index, input)
-
-    async def wait_for_user_input(self) -> str:
-        self._active = False
-        self.update_subscribers()
-        await self._user_input_queue.get(self.history.current_index)
-        self._active = True
-        self.update_subscribers()
-
-    _manual_edits_buffer: List[FileEditWithFullContents] = []
-
-    async def reverse_to_index(self, index: int):
-        try:
-            while self.history.get_current_index() >= index:
-                current_step = self.history.get_current().step
-                self.history.step_back()
-                if issubclass(current_step.__class__, ReversibleStep):
-                    await current_step.reverse(self.__get_step_params(current_step))
-
-                self.update_subscribers()
-        except Exception as e:
-            print(e)
-
-    def handle_manual_edits(self, edits: List[FileEditWithFullContents]):
-        for edit in edits:
-            self._manual_edits_buffer.append(edit)
-            # TODO: You're storing a lot of unecessary data here. Can compress into EditDiffs on the spot, and merge.
-            # self._manual_edits_buffer = merge_file_edit(self._manual_edits_buffer, edit)
-
-    def handle_traceback(self, traceback: str):
-        raise NotImplementedError
-
-    _step_depth: int = 0
-
-    async def _run_singular_step(self, step: "Step", is_future_step: bool = False) -> Coroutine[Observation, None, None]:
-        if not is_future_step:
-            # Check manual edits buffer, clear out if needed by creating a ManualEditStep
-            if len(self._manual_edits_buffer) > 0:
-                manualEditsStep = ManualEditStep.from_sequence(
-                    self._manual_edits_buffer)
-                self._manual_edits_buffer = []
-                await self._run_singular_step(manualEditsStep)
-
-        # Update history - do this first so we get top-first tree ordering
-        self.history.add_node(HistoryNode(
-            step=step, observation=None, depth=self._step_depth))
-
-        # Run step
-        self._step_depth += 1
-        observation = await step(self.__get_step_params(step))
-        self._step_depth -= 1
-
-        # Add observation to history
-        self.history.get_current().observation = observation
-
-        # Update its description
-        step._set_description(await step.describe(self.llm))
-
-        # Call all subscribed callbacks
-        self.update_subscribers()
-
-        return observation
-
-    async def run_from_step(self, step: "Step"):
-        # if self._active:
-        #     raise RuntimeError("Agent is already running")
-        self._active = True
-
-        next_step = step
-        is_future_step = False
-        while not (next_step is None or self._should_halt):
-            try:
-                if is_future_step:
-                    # If future step, then we are replaying and need to delete the step from history so it can be replaced
-                    self.history.remove_current_and_substeps()
-
-                observation = await self._run_singular_step(next_step, is_future_step)
-                if next_step := self.policy.next(self.history):
-                    is_future_step = False
-                elif next_step := self.history.take_next_step():
-                    is_future_step = True
-                else:
-                    next_step = None
-
-            except Exception as e:
-                print(
-                    f"Error while running step: \n{''.join(traceback.format_tb(e.__traceback__))}\n{e}")
-                next_step = None
-
-        self._active = False
-
-        # Doing this so active can make it to the frontend after steps are done. But want better state syncing tools
-        for callback in self._on_update_callbacks:
-            callback(None)
-
-    async def run_from_observation(self, observation: Observation):
-        next_step = self.policy.next(self.history)
-        await self.run_from_step(next_step)
-
-    async def run_policy(self):
-        first_step = self.policy.next(self.history)
-        await self.run_from_step(first_step)
-
-    async def _request_halt(self):
-        if self._active:
-            self._should_halt = True
-            while self._active:
-                time.sleep(0.1)
-        self._should_halt = False
-        return None
-
-    async def accept_user_input(self, user_input: str):
-        self._main_user_input_queue.append(user_input)
-        self.update_subscribers()
-
-        if len(self._main_user_input_queue) > 1:
-            return
-
-        # await self._request_halt()
-        # Just run the step that takes user input, and
-        # then up to the policy to decide how to deal with it.
-        self._main_user_input_queue.pop(0)
-        self.update_subscribers()
-        await self.run_from_step(UserInputStep(user_input=user_input))
-
-        while len(self._main_user_input_queue) > 0:
-            await self.run_from_step(UserInputStep(
-                user_input=self._main_user_input_queue.pop(0)))
-
-    async def accept_refinement_input(self, user_input: str, index: int):
-        await self._request_halt()
-        await self.reverse_to_index(index)
-        await self.run_from_step(UserInputStep(user_input=user_input))
-
-
-class Step(ContinueBaseModel):
-    name: str = None
-    hide: bool = False
-    _description: Union[str, None] = None
-
-    system_message: Union[str, None] = None
-
-    class Config:
-        copy_on_model_validation = False
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        if self._description is not None:
-            return self._description
-        return "Running step: " + self.name
-
-    def _set_description(self, description: str):
-        self._description = description
-
-    def dict(self, *args, **kwargs):
-        d = super().dict(*args, **kwargs)
-        if self._description is not None:
-            d["description"] = self._description
-        else:
-            d["description"] = self.name
-        return d
-
-    @validator("name", pre=True, always=True)
-    def name_is_class_name(cls, name):
-        if name is None:
-            return cls.__name__
-        return name
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        raise NotImplementedError
-
-    async def __call__(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        return await self.run(sdk)
-
-    def __rshift__(self, other: "Step"):
-        steps = []
-        if isinstance(self, SequentialStep):
-            steps = self.steps
-        else:
-            steps.append(self)
-        if isinstance(other, SequentialStep):
-            steps += other.steps
-        else:
-            steps.append(other)
-        return SequentialStep(steps=steps)
-
-
-class SequentialStep(Step):
-    steps: list[Step]
-    hide: bool = True
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        for step in self.steps:
-            observation = await sdk.run_step(step)
-        return observation
-
-
-class ReversibleStep(Step):
-    async def reverse(self, sdk: ContinueSDK):
-        raise NotImplementedError
-
-
-class FileSystemEditStep(ReversibleStep):
-    edit: FileSystemEdit
-    _diff: Union[EditDiff, None] = None
-
-    hide: bool = True
-
-    async def run(self, sdk: "ContinueSDK") -> Coroutine[Observation, None, None]:
-        self._diff = await sdk.ide.applyFileSystemEdit(self.edit)
-        return None
-
-    async def reverse(self, sdk: "ContinueSDK"):
-        await sdk.ide.applyFileSystemEdit(self._diff.backward)
-        # Where and when should file saves happen?
-
-
-class ManualEditStep(ReversibleStep):
-    edit_diff: EditDiff
-    hide: bool = True
-
-    hide: bool = True
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        return "Manual edit step"
-        # TODO - only handling FileEdit here, but need all other types of FileSystemEdits
-        # Also requires the merge_file_edit function
-        # return llm.complete(dedent(f"""This code was replaced:
-
-        #     {self.edit_diff.backward.replacement}
-
-        #     With this code:
-
-        #     {self.edit_diff.forward.replacement}
-
-        #     Maximally concise summary of changes in bullet points (can use markdown):
-        # """))
-
-    @classmethod
-    def from_sequence(cls, edits: List[FileEditWithFullContents]) -> "ManualEditStep":
-        diffs = []
-        for edit in edits:
-            _, diff = FileSystem.apply_edit_to_str(
-                edit.fileContents, edit.fileEdit)
-            diffs.append(diff)
-        return cls(edit_diff=EditDiff.from_sequence(diffs))
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        return None
-
-    async def reverse(self, sdk: ContinueSDK):
-        await sdk.ide.applyFileSystemEdit(self.edit_diff.backward)
-
-
-class UserInputStep(Step):
-    user_input: str
-    name: str = "User Input"
-    hide: bool = True
-
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
-        return self.user_input
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[UserInputObservation, None, None]:
-        return UserInputObservation(user_input=self.user_input)
-
-
-class ValidatorObservation(Observation):
-    passed: bool
-    observation: Observation
-
-
-class Validator(Step):
-    def run(self, sdk: ContinueSDK) -> ValidatorObservation:
-        raise NotImplementedError
-
-
-HistoryNode.update_forward_refs()
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
new file mode 100644
index 00000000..0338d635
--- /dev/null
+++ b/continuedev/src/continuedev/libs/steps/core/core.py
@@ -0,0 +1,216 @@
+# These steps are depended upon by ContinueSDK
+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, RangeInFile, RangeInFileWithContents
+from ...llm import LLM
+from ....core.observation import Observation, TextObservation, TracebackObservation, UserInputObservation
+from ....core.main import Step
+
+
+class ContinueSDK:
+    pass
+
+
+class SequentialStep(Step):
+    steps: list[Step]
+    hide: bool = True
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        for step in self.steps:
+            observation = await sdk.run_step(step)
+        return observation
+
+
+class ReversibleStep(Step):
+    async def reverse(self, sdk: ContinueSDK):
+        raise NotImplementedError
+
+
+class FileSystemEditStep(ReversibleStep):
+    edit: FileSystemEdit
+    _diff: Union[EditDiff, None] = None
+
+    hide: bool = True
+
+    async def run(self, sdk: "ContinueSDK") -> Coroutine[Observation, None, None]:
+        self._diff = await sdk.ide.applyFileSystemEdit(self.edit)
+        return None
+
+    async def reverse(self, sdk: "ContinueSDK"):
+        await sdk.ide.applyFileSystemEdit(self._diff.backward)
+        # 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
+
+    hide: bool = True
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        return "Manual edit step"
+        # TODO - only handling FileEdit here, but need all other types of FileSystemEdits
+        # Also requires the merge_file_edit function
+        # return llm.complete(dedent(f"""This code was replaced:
+
+        #     {self.edit_diff.backward.replacement}
+
+        #     With this code:
+
+        #     {self.edit_diff.forward.replacement}
+
+        #     Maximally concise summary of changes in bullet points (can use markdown):
+        # """))
+
+    @classmethod
+    def from_sequence(cls, edits: List[FileEditWithFullContents]) -> "ManualEditStep":
+        diffs = []
+        for edit in edits:
+            _, diff = FileSystem.apply_edit_to_str(
+                edit.fileContents, edit.fileEdit)
+            diffs.append(diff)
+        return cls(edit_diff=EditDiff.from_sequence(diffs))
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        return None
+
+    async def reverse(self, sdk: ContinueSDK):
+        await sdk.ide.applyFileSystemEdit(self.edit_diff.backward)
+
+
+class UserInputStep(Step):
+    user_input: str
+    name: str = "User Input"
+    hide: bool = True
+
+    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+        return self.user_input
+
+    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 70953e95..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,124 +72,6 @@ class RunCommandStep(Step):
             return TextObservation(text=stdout)
 
 
-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/models/generate_json_schema.py b/continuedev/src/continuedev/models/generate_json_schema.py
index da78dfac..07337029 100644
--- a/continuedev/src/continuedev/models/generate_json_schema.py
+++ b/continuedev/src/continuedev/models/generate_json_schema.py
@@ -1,6 +1,7 @@
 from .main import *
 from .filesystem import RangeInFile, FileEdit
 from .filesystem_edit import FileEditWithFullContents
+from ..core.main import History, HistoryNode
 from pydantic import schema_json_of
 import os
 
@@ -10,6 +11,8 @@ MODELS_TO_GENERATE = [
     RangeInFile, FileEdit
 ] + [
     FileEditWithFullContents
+] + [
+    History, HistoryNode
 ]
 
 RENAMES = {
diff --git a/continuedev/src/continuedev/models/main.py b/continuedev/src/continuedev/models/main.py
index 081ec4af..7986b30c 100644
--- a/continuedev/src/continuedev/models/main.py
+++ b/continuedev/src/continuedev/models/main.py
@@ -4,6 +4,11 @@ from pydantic import BaseModel, root_validator
 from functools import total_ordering
 
 
+class ContinueBaseModel(BaseModel):
+    class Config:
+        underscore_attrs_are_private = True
+
+
 @total_ordering
 class Position(BaseModel):
     line: int
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 c9d4edc5..bfd7a09c 100644
--- a/continuedev/src/continuedev/server/notebook.py
+++ b/continuedev/src/continuedev/server/notebook.py
@@ -5,21 +5,18 @@ from pydantic import BaseModel
 from uvicorn.main import Server
 
 from ..models.filesystem_edit import FileEditWithFullContents
-from ..libs.policy import DemoPolicy
-from ..libs.core import Agent, FullState, History, Step
+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 dotenv import load_dotenv
+from ..core.observation import Observation
 from ..libs.llm.openai import OpenAI
 from .ide_protocol import AbstractIdeProtocolServer
-import os
+from ..core.env import openai_api_key
 import asyncio
 import nest_asyncio
 nest_asyncio.apply()
 
-load_dotenv()
-openai_api_key = os.getenv("OPENAI_API_KEY")
-
 router = APIRouter(prefix="/notebook", tags=["notebook"])
 
 # Graceful shutdown by closing websockets