9 files changed, 27 insertions, 22 deletions
diff --git a/continuedev/src/continuedev/core/abstract_sdk.py b/continuedev/src/continuedev/core/abstract_sdk.py
index 417971cd..3b85708d 100644
--- a/continuedev/src/continuedev/core/abstract_sdk.py
+++ b/continuedev/src/continuedev/core/abstract_sdk.py
@@ -88,6 +88,6 @@ class AbstractContinueSDK(ABC):
     def add_chat_context(self, content: str, role: ChatMessageRole = "assistent"):
         pass
 
-    @abstractproperty
-    def chat_context(self) -> List[ChatMessage]:
+    @abstractmethod
+    async def get_chat_context(self) -> List[ChatMessage]:
         pass
diff --git a/continuedev/src/continuedev/core/policy.py b/continuedev/src/continuedev/core/policy.py
index 6a183598..e71a1cb2 100644
--- a/continuedev/src/continuedev/core/policy.py
+++ b/continuedev/src/continuedev/core/policy.py
@@ -37,7 +37,6 @@ class DemoPolicy(Policy):
                 return WritePytestsRecipe(instructions=user_input)
             elif "/dlt" in user_input.lower() or " dlt" in user_input.lower():
                 return CreatePipelineRecipe()
-            elif "/comment" in user_input.lower():
             if "/pytest" in observation.user_input.lower():
                 return WritePytestsRecipe(instructions=observation.user_input)
             elif "/dlt" in observation.user_input.lower():
@@ -59,10 +58,10 @@ class DemoPolicy(Policy):
             # return EditHighlightedCodeStep(user_input=user_input)
             return NLDecisionStep(user_input=user_input, steps=[
                 EditHighlightedCodeStep(user_input=user_input),
-                AnswerQuestionChroma(question=user_input),
-                EditFileChroma(request=user_input),
+                # AnswerQuestionChroma(question=user_input),
+                # EditFileChroma(request=user_input),
                 SimpleChatStep(user_input=user_input)
-            ])
+            ], default_step=EditHighlightedCodeStep(user_input=user_input))
 
         state = history.get_current()
 
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index 59bfc0f2..a94b5026 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -141,6 +141,11 @@ class ContinueSDK(AbstractContinueSDK):
         self.history.timeline[self.history.current_index].step.chat_context.append(
             ChatMessage(content=content, role=role))
 
-    @property
-    def chat_context(self) -> List[ChatMessage]:
-        return self.history.to_chat_history()
+    async def get_chat_context(self) -> List[ChatMessage]:
+        history_context = self.history.to_chat_history()
+        highlighted_code = await self.ide.getHighlightedCode()
+        for rif in highlighted_code:
+            code = await self.ide.readRangeInFile(rif)
+            history_context.append(ChatMessage(
+                content=f"The following code is highlighted:\n```\n{code}\n```", role="user"))
+        return history_context
diff --git a/continuedev/src/continuedev/steps/chat.py b/continuedev/src/continuedev/steps/chat.py
index 817e10dd..80065c24 100644
--- a/continuedev/src/continuedev/steps/chat.py
+++ b/continuedev/src/continuedev/steps/chat.py
@@ -2,12 +2,11 @@ from textwrap import dedent
 from typing import List
 from ..core.main import Step
 from ..core.sdk import ContinueSDK
-from .main import MessageStep
+from .core.core import MessageStep
 
 
 class SimpleChatStep(Step):
     user_input: str
 
     async def run(self, sdk: ContinueSDK):
-        # TODO: With history
-        self.description = await sdk.models.gpt35.complete(self.user_input)
+        self.description = sdk.models.gpt35.complete(self.user_input, with_history=await sdk.get_chat_context())
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index 53df65cc..57689f19 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -85,7 +85,7 @@ class ShellCommandsStep(Step):
                     {output}
                     ```
 
-                    This is a brief summary of the error followed by a suggestion on how it can be fixed:"""), with_history=sdk.chat_context)
+                    This is a brief summary of the error followed by a suggestion on how it can be fixed:"""), with_history=await sdk.get_chat_context())
 
                 sdk.raise_exception(
                     title="Error while running query", message=output, with_step=MessageStep(name=f"Suggestion to solve error {AI_ASSISTED_STRING}", message=suggestion)
diff --git a/continuedev/src/continuedev/steps/draft/migration.py b/continuedev/src/continuedev/steps/draft/migration.py
index f3b36b5e..7c4b7eb5 100644
--- a/continuedev/src/continuedev/steps/draft/migration.py
+++ b/continuedev/src/continuedev/steps/draft/migration.py
@@ -13,7 +13,7 @@ class MigrationStep(Step):
         recent_edits = await sdk.ide.get_recent_edits(self.edited_file)
         recent_edits_string = "\n\n".join(
             map(lambda x: x.to_string(), recent_edits))
-        description = await sdk.models.gpt35.complete(f"{recent_edits_string}\n\nGenerate a short description of the migration made in the above changes:\n")
+        description = sdk.models.gpt35.complete(f"{recent_edits_string}\n\nGenerate a short description of the migration made in the above changes:\n")
         await sdk.run([
             "cd libs",
             "poetry run alembic revision --autogenerate -m " + description,
diff --git a/continuedev/src/continuedev/steps/find_and_replace.py b/continuedev/src/continuedev/steps/find_and_replace.py
index fec33997..690872c0 100644
--- a/continuedev/src/continuedev/steps/find_and_replace.py
+++ b/continuedev/src/continuedev/steps/find_and_replace.py
@@ -10,7 +10,7 @@ class FindAndReplaceStep(Step):
     replacement: str
 
     async def describe(self, models: Models):
-        return f"Replace all instances of `{self.pattern}` with `{self.replacement}` in `{self.filepath}`"
+        return f"Replaced all instances of `{self.pattern}` with `{self.replacement}` in `{self.filepath}`"
 
     async def run(self, sdk: ContinueSDK):
         file_content = await sdk.ide.readFile(self.filepath)
diff --git a/continuedev/src/continuedev/steps/input/nl_multiselect.py b/continuedev/src/continuedev/steps/input/nl_multiselect.py
index c3c832f5..36c489c7 100644
--- a/continuedev/src/continuedev/steps/input/nl_multiselect.py
+++ b/continuedev/src/continuedev/steps/input/nl_multiselect.py
@@ -23,5 +23,6 @@ class NLMultiselectStep(Step):
         if first_try is not None:
             return first_try
 
-        gpt_parsed = await sdk.models.gpt35.complete(f"These are the available options are: [{', '.join(self.options)}]. The user requested {user_response}. This is the exact string from the options array that they selected:")
+        gpt_parsed = sdk.models.gpt35.complete(
+            f"These are the available options are: [{', '.join(self.options)}]. The user requested {user_response}. This is the exact string from the options array that they selected:")
         return extract_option(gpt_parsed) or self.options[0]
diff --git a/continuedev/src/continuedev/steps/react.py b/continuedev/src/continuedev/steps/react.py
index 411adc87..6b6024ce 100644
--- a/continuedev/src/continuedev/steps/react.py
+++ b/continuedev/src/continuedev/steps/react.py
@@ -1,13 +1,15 @@
 from textwrap import dedent
-from typing import List
+from typing import List, Union
 from ..core.main import Step
 from ..core.sdk import ContinueSDK
-from .main import MessageStep
+from .core.core import MessageStep
 
 
 class NLDecisionStep(Step):
     user_input: str
     steps: List[Step]
+    hide: bool = True
+    default_step: Union[Step, None] = None
 
     async def run(self, sdk: ContinueSDK):
         step_descriptions = "\n".join([
@@ -24,14 +26,13 @@ class NLDecisionStep(Step):
                         
                         Select the step which should be taken next. Say only the name of the selected step:""")
 
-        resp = (await sdk.models.gpt35.complete(prompt)).lower()
+        resp = sdk.models.gpt35.complete(prompt).lower()
 
         step_to_run = None
         for step in self.steps:
-            if step.name in resp:
+            if step.name.lower() in resp:
                 step_to_run = step
 
-        step_to_run = step_to_run or MessageStep(
-            message="Unable to decide the next step")
+        step_to_run = step_to_run or self.default_step or self.steps[0]
 
         await sdk.run_step(step_to_run)