path: root/server/continuedev/core/autopilot.py

import json
import os
import time
import traceback
import uuid
from functools import cached_property
from typing import Callable, Coroutine, Dict, List, Optional

import redbaron
from aiohttp import ClientPayloadError
from openai import error as openai_errors
from pydantic import root_validator

from ..libs.llm.prompts.chat import template_alpaca_messages
from ..libs.util.create_async_task import create_async_task
from ..libs.util.devdata import dev_data_logger
from ..libs.util.edit_config import edit_config_property
from ..libs.util.logging import logger
from ..libs.util.paths import getSavedContextGroupsPath
from ..libs.util.queue import AsyncSubscriptionQueue
from ..libs.util.strings import remove_quotes_and_escapes
from ..libs.util.telemetry import posthog_logger
from ..libs.util.traceback.traceback_parsers import (
    get_javascript_traceback,
    get_python_traceback,
)
from ..models.filesystem import RangeInFileWithContents
from ..models.filesystem_edit import FileEditWithFullContents
from ..models.main import ContinueBaseModel
from ..plugins.context_providers.file import FileContextProvider
from ..plugins.context_providers.highlighted_code import HighlightedCodeContextProvider
from ..plugins.policies.default import DefaultPolicy
from ..plugins.steps.on_traceback import DefaultOnTracebackStep
from ..server.ide_protocol import AbstractIdeProtocolServer
from ..server.meilisearch_server import get_meilisearch_url, stop_meilisearch
from .config import ContinueConfig
from .context import ContextManager
from .main import (
    Context,
    ContextItem,
    ContinueCustomException,
    FullState,
    History,
    HistoryNode,
    Policy,
    SessionInfo,
    Step,
)
from .observation import InternalErrorObservation, Observation
from .sdk import ContinueSDK
from .steps import DisplayErrorStep, ManualEditStep, ReversibleStep, UserInputStep


def get_error_title(e: Exception) -> str:
    if isinstance(e, openai_errors.APIError):
        return "OpenAI is overloaded with requests. Please try again."
    elif isinstance(e, openai_errors.RateLimitError):
        return "This OpenAI API key has been rate limited. Please try again."
    elif isinstance(e, openai_errors.Timeout):
        return "OpenAI timed out. Please try again."
    elif (
        isinstance(e, openai_errors.InvalidRequestError)
        and e.code == "context_length_exceeded"
    ):
        return e._message
    elif isinstance(e, ClientPayloadError):
        return "The request failed. Please try again."
    elif isinstance(e, openai_errors.APIConnectionError):
        return 'The request failed. Please check your internet connection and try again. If this issue persists, you can use our API key for free by going to VS Code settings and changing the value of continue.OPENAI_API_KEY to ""'
    elif isinstance(e, openai_errors.InvalidRequestError):
        return "Invalid request sent to OpenAI. Please try again."
    elif "rate_limit_ip_middleware" in e.__str__():
        return "You have reached your limit for free usage of our token. You can continue using Continue by entering your own OpenAI API key in VS Code settings."
    elif e.__str__().startswith("Cannot connect to host"):
        return (
            "The request failed. Please check your internet connection and try again."
        )
    return e.__str__() or e.__repr__()


class Autopilot(ContinueBaseModel):
    ide: AbstractIdeProtocolServer

    policy: Policy = DefaultPolicy()
    history: History = History.from_empty()
    context: Context = Context()
    full_state: Optional[FullState] = None
    session_info: Optional[SessionInfo] = None
    context_manager: ContextManager = ContextManager()
    continue_sdk: ContinueSDK = None

    _on_update_callbacks: List[Callable[[FullState], None]] = []

    _active: bool = False
    _should_halt: bool = False
    _main_user_input_queue: List[str] = []

    _user_input_queue = AsyncSubscriptionQueue()
    _retry_queue = AsyncSubscriptionQueue()

    started: bool = False

    async def load(
        self, config: Optional[ContinueConfig] = None, only_reloading: bool = False
    ):
        self.continue_sdk = await ContinueSDK.create(self, config=config)
        if override_policy := self.continue_sdk.config.policy_override:
            self.policy = override_policy

        # Load documents into the search index
        logger.debug("Starting context manager")
        await self.context_manager.start(
            self.continue_sdk.config.context_providers
            + [
                HighlightedCodeContextProvider(ide=self.ide),
                FileContextProvider(workspace_dir=self.ide.workspace_directory),
            ],
            self.continue_sdk,
            only_reloading=only_reloading,
        )

    async def start(
        self,
        full_state: Optional[FullState] = None,
        config: Optional[ContinueConfig] = None,
    ):
        await self.load(config=config, only_reloading=False)

        if full_state is not None:
            self.history = full_state.history
            self.session_info = full_state.session_info

        # Load saved context groups
        context_groups_file = getSavedContextGroupsPath()
        try:
            with open(context_groups_file, "r") as f:
                json_ob = json.load(f)
                for title, context_group in json_ob.items():
                    self._saved_context_groups[title] = [
                        ContextItem(**item) for item in context_group
                    ]
        except Exception as e:
            logger.warning(
                f"Failed to load saved_context_groups.json: {e}. Reverting to empty list."
            )
            self._saved_context_groups = {}

        self.started = True

    async def reload_config(self):
        await self.load(config=None, only_reloading=True)
        await self.update_subscribers()

    async def cleanup(self):
        stop_meilisearch()

    class Config:
        arbitrary_types_allowed = True
        keep_untouched = (cached_property,)

    @root_validator(pre=True)
    def fill_in_values(cls, values):
        full_state: FullState = values.get("full_state")
        if full_state is not None:
            values["history"] = full_state.history
        return values

    async def get_full_state(self) -> FullState:
        full_state = FullState(
            history=self.history,
            active=self._active,
            user_input_queue=self._main_user_input_queue,
            slash_commands=self.get_available_slash_commands(),
            adding_highlighted_code=self.context_manager.context_providers[
                "code"
            ].adding_highlighted_code
            if "code" in self.context_manager.context_providers
            else False,
            selected_context_items=await self.context_manager.get_selected_items()
            if self.context_manager is not None
            else [],
            session_info=self.session_info,
            config=self.continue_sdk.config,
            saved_context_groups=self._saved_context_groups,
            context_providers=self.context_manager.get_provider_descriptions(),
            meilisearch_url=get_meilisearch_url(),
        )
        self.full_state = full_state
        return full_state

    def get_available_slash_commands(self) -> List[Dict]:
        custom_commands = (
            list(
                map(
                    lambda x: {"name": x.name, "description": x.description},
                    self.continue_sdk.config.custom_commands,
                )
            )
            or []
        )
        slash_commands = (
            list(
                map(
                    lambda x: {"name": x.name, "description": x.description},
                    self.continue_sdk.config.slash_commands,
                )
            )
            or []
        )
        cmds = custom_commands + slash_commands
        cmds.sort(key=lambda x: x["name"] == "edit", reverse=True)
        return cmds

    async def clear_history(self):
        # Reset history
        self.history = History.from_empty()
        self._main_user_input_queue = []
        self._active = False

        # Clear context
        # await self.context_manager.clear_context()

        await self.update_subscribers()

    def on_update(self, callback: Coroutine["FullState", None, None]):
        """Subscribe to changes to state"""
        self._on_update_callbacks.append(callback)

    async def update_subscribers(self):
        full_state = await self.get_full_state()
        for callback in self._on_update_callbacks:
            await callback(full_state)

    def give_user_input(self, input: str, index: int):
        self._user_input_queue.post(str(index), input)

    async def wait_for_user_input(self) -> str:
        self._active = False
        await self.update_subscribers()
        user_input = await self._user_input_queue.get(str(self.history.current_index))
        self._active = True
        await self.update_subscribers()
        return user_input

    _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.continue_sdk)

                await self.update_subscribers()
        except Exception as e:
            logger.debug(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 unnecessary data here. Can compress into EditDiffs on the spot, and merge.
            # self._manual_edits_buffer = merge_file_edit(self._manual_edits_buffer, edit)
            # Note that this is being overridden to do nothing in DemoAgent

    async def handle_command_output(self, output: str):
        get_traceback_funcs = [get_python_traceback, get_javascript_traceback]
        for get_tb_func in get_traceback_funcs:
            traceback = get_tb_func(output)
            if (
                traceback is not None
                and self.continue_sdk.config.on_traceback is not None
            ):
                step = self.continue_sdk.config.on_traceback(output=output)
                await self._run_singular_step(step)

    async def handle_debug_terminal(self, content: str):
        """Run the debug terminal step"""
        # step = self.continue_sdk.config.on_traceback(output=content)
        step = DefaultOnTracebackStep(output=content)
        await self._run_singular_step(step)

    async def handle_highlighted_code(
        self,
        range_in_files: List[RangeInFileWithContents],
        edit: Optional[bool] = False,
    ):
        if "code" not in self.context_manager.context_providers:
            return

        # Add to context manager
        await self.context_manager.context_providers["code"].handle_highlighted_code(
            range_in_files, edit
        )

        await self.update_subscribers()

    _step_depth: int = 0

    async def retry_at_index(self, index: int):
        self.history.timeline[index].step.hide = True
        self._retry_queue.post(str(index), None)

    async def delete_at_index(self, index: int):
        if not self.history.timeline[index].active:
            self.history.timeline[index].step.hide = True

        self.history.timeline[index].deleted = True
        self.history.timeline[index].active = False

        await self.update_subscribers()

    async def edit_step_at_index(self, user_input: str, index: int):
        node_to_rerun = self.history.timeline[index].copy()
        step_to_rerun = node_to_rerun.step
        step_to_rerun.user_input = user_input
        step_to_rerun.description = user_input

        # Halt the agent's currently running jobs (delete them)
        while len(self.history.timeline) > index:
            # Remove from timeline
            node_to_delete = self.history.timeline.pop()
            # Delete so it is stopped if in the middle of running
            node_to_delete.deleted = True

        self.history.current_index = index - 1

        # Set the context to the context used by that step
        await self.context_manager.clear_context()
        for context_item in node_to_rerun.context_used:
            await self.context_manager.manually_add_context_item(context_item)

        await self.update_subscribers()

        # Rerun from the current step
        await self.run_from_step(step_to_rerun)

    async def delete_context_with_ids(
        self, ids: List[str], index: Optional[int] = None
    ):
        if index is None:
            await self.context_manager.delete_context_with_ids(ids)
        else:
            self.history.timeline[index].context_used = list(
                filter(
                    lambda item: item.description.id.to_string() not in ids,
                    self.history.timeline[index].context_used,
                )
            )
        await self.update_subscribers()

    async def toggle_adding_highlighted_code(self):
        if "code" not in self.context_manager.context_providers:
            return

        self.context_manager.context_providers[
            "code"
        ].adding_highlighted_code = not self.context_manager.context_providers[
            "code"
        ].adding_highlighted_code
        await self.update_subscribers()

    async def set_editing_at_ids(self, ids: List[str]):
        if "code" not in self.context_manager.context_providers:
            return

        await self.context_manager.context_providers["code"].set_editing_at_ids(ids)
        await self.update_subscribers()

    async def _run_singular_step(
        self, step: "Step", is_future_step: bool = False
    ) -> Coroutine[Observation, None, None]:
        # Allow config to set disallowed steps
        if step.__class__.__name__ in self.continue_sdk.config.disallowed_steps:
            return None

        # If a parent step is deleted/cancelled, don't run this step
        # TODO: This was problematic because when running a step after deleting one, it seemed to think that was the parent
        # last_depth = self._step_depth
        # i = self.history.current_index
        # while i >= 0 and self.history.timeline[i].depth == last_depth - 1:
        #     if self.history.timeline[i].deleted:
        #         return None
        #     last_depth = self.history.timeline[i].depth
        #     i -= 1

        # Log the context and step to dev data
        context_used = await self.context_manager.get_selected_items()
        posthog_logger.capture_event(
            "step run", {"step_name": step.name, "params": step.dict()}
        )
        step_id = uuid.uuid4().hex
        dev_data_logger.capture(
            "step_run",
            {"step_name": step.name, "params": step.dict(), "step_id": step_id},
        )
        dev_data_logger.capture(
            "context_used",
            {
                "context": list(
                    map(
                        lambda item: item.dict(),
                        context_used,
                    )
                ),
                "step_id": step_id,
            },
        )

        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
        index_of_history_node = self.history.add_node(
            HistoryNode(
                step=step,
                observation=None,
                depth=self._step_depth,
                context_used=context_used,
            )
        )

        # Call all subscribed callbacks
        await self.update_subscribers()

        # Try to run step and handle errors
        self._step_depth += 1

        caught_error = False
        try:
            observation = await step(self.continue_sdk)
        except Exception as e:
            if (
                index_of_history_node >= len(self.history.timeline)
                or self.history.timeline[index_of_history_node].deleted
            ):
                # If step was deleted/cancelled, don't show error or allow retry
                return None

            caught_error = True

            is_continue_custom_exception = (
                issubclass(e.__class__, ContinueCustomException)
                or e.__class__.__name__ == ContinueCustomException.__name__
            )

            error_string = (
                e.message
                if is_continue_custom_exception
                else "\n".join(traceback.format_exception(e))
            )
            error_title = (
                e.title if is_continue_custom_exception else get_error_title(e)
            )

            # Attach an InternalErrorObservation to the step and unhide it.
            logger.error(f"Error while running step: \n{error_string}\n{error_title}")
            posthog_logger.capture_event(
                "step error",
                {
                    "error_message": error_string,
                    "error_title": error_title,
                    "step_name": step.name,
                    "params": step.dict(),
                },
            )

            observation = InternalErrorObservation(
                error=error_string, title=error_title
            )

            # Reveal this step, but hide all of the following steps (its substeps)
            step_was_hidden = step.hide

            step.hide = False
            i = self.history.get_current_index()
            while self.history.timeline[i].step.name != step.name:
                self.history.timeline[i].step.hide = True
                i -= 1

            # i is now the index of the step that we want to show/rerun
            self.history.timeline[i].observation = observation
            self.history.timeline[i].active = False

            await self.update_subscribers()

            # ContinueCustomException can optionally specify a step to run on the error
            if is_continue_custom_exception and e.with_step is not None:
                await self._run_singular_step(e.with_step)

            # Wait for a retry signal and then resume the step
            self._active = False
            await self._retry_queue.get(str(i))
            self._active = True
            # You might consider a "ignore and continue" button
            # want it to have same step depth, so have to decrement
            self._step_depth -= 1
            copy_step = step.copy()
            copy_step.hide = step_was_hidden
            observation = await self._run_singular_step(copy_step)
            self._step_depth += 1

        self._step_depth -= 1

        # Add observation to history, unless already attached error observation
        if not caught_error and index_of_history_node < len(self.history.timeline):
            self.history.timeline[index_of_history_node].observation = observation
            self.history.timeline[index_of_history_node].active = False
            await self.update_subscribers()

        # Update its description
        async def update_description():
            if self.continue_sdk.config.disable_summaries:
                return

            description = await step.describe(self.continue_sdk.models)
            if description is not None:
                step.description = description
            # Update subscribers with new description
            await self.update_subscribers()

        create_async_task(
            update_description(),
            on_error=lambda e: self.continue_sdk.run_step(
                DisplayErrorStep.from_exception(e)
            ),
        )

        # Create the session title if not done yet
        if self.session_info is None or self.session_info.title is None:
            visible_nodes = list(
                filter(lambda node: not node.step.hide, self.history.timeline)
            )

            user_input = None
            should_create_title = False
            for visible_node in visible_nodes:
                if isinstance(visible_node.step, UserInputStep):
                    if user_input is None:
                        user_input = visible_node.step.user_input
                    else:
                        # More than one user input, so don't create title
                        should_create_title = False
                        break
                elif user_input is None:
                    continue
                else:
                    # Already have user input, now have the next step
                    should_create_title = True
                    break

            # Only create the title if the step after the first input is done
            if should_create_title:
                create_async_task(
                    self.create_title(backup=user_input),
                    on_error=lambda e: self.continue_sdk.run_step(
                        DisplayErrorStep.from_exception(e)
                    ),
                )

        return observation

    async def run_from_step(self, step: "Step"):
        # if self._active:
        #     raise RuntimeError("Autopilot is already running")
        self._active = True

        next_step = step
        is_future_step = False
        while not (next_step is None or self._should_halt):
            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()

            await self._run_singular_step(next_step, is_future_step)

            if next_step := self.policy.next(self.continue_sdk.config, self.history):
                is_future_step = False
            elif next_step := self.history.take_next_step():
                is_future_step = True
            else:
                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
        await self.update_subscribers()

    async def run_from_observation(self, observation: Observation):
        next_step = self.policy.next(self.continue_sdk.config, self.history)
        await self.run_from_step(next_step)

    async def run_policy(self):
        first_step = self.policy.next(self.continue_sdk.config, 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

    def set_current_session_title(self, title: str):
        self.session_info = SessionInfo(
            title=title,
            session_id=self.ide.session_id,
            date_created=str(time.time()),
            workspace_directory=self.ide.workspace_directory,
        )

    async def create_title(self, backup: str = None):
        # Use the first input and first response to create title for session info, and make the session saveable
        if self.session_info is not None and self.session_info.title is not None:
            return

        if self.continue_sdk.config.disable_summaries:
            if backup is not None:
                title = backup
            else:
                title = "New Session"
        else:
            chat_history = list(
                map(lambda x: x.dict(), await self.continue_sdk.get_chat_context())
            )
            chat_history_str = template_alpaca_messages(chat_history)
            title = await self.continue_sdk.models.summarize.complete(
                f"{chat_history_str}\n\nGive a short title to describe the above chat session. Do not put quotes around the title. Do not use more than 6 words. The title is: ",
                max_tokens=20,
                log=False,
            )
            title = remove_quotes_and_escapes(title)

        self.set_current_session_title(title)
        await self.update_subscribers()
        dev_data_logger.capture("new_session", self.session_info.dict())

    async def accept_user_input(self, user_input: str):
        self._main_user_input_queue.append(user_input)
        # await 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)
        # await 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))

    async def reject_diff(self, step_index: int):
        # Hide the edit step and the UserInputStep before it
        self.history.timeline[step_index].step.hide = True
        for i in range(step_index - 1, -1, -1):
            if isinstance(self.history.timeline[i].step, UserInputStep):
                self.history.timeline[i].step.hide = True
                break
        await self.update_subscribers()

    async def select_context_item(self, id: str, query: str):
        await self.context_manager.select_context_item(id, query)
        await self.update_subscribers()

    async def select_context_item_at_index(self, id: str, query: str, index: int):
        # TODO: This is different from how it works for the main input
        # Ideally still tracked through the ContextProviders
        # so they can watch for duplicates
        context_item = await self.context_manager.get_context_item(id, query)
        if context_item is None:
            return
        self.history.timeline[index].context_used.append(context_item)
        await self.update_subscribers()

    async def set_config_attr(self, key_path: List[str], value: redbaron.RedBaron):
        edit_config_property(key_path, value)
        await self.update_subscribers()

    _saved_context_groups: Dict[str, List[ContextItem]] = {}

    def _persist_context_groups(self):
        context_groups_file = getSavedContextGroupsPath()
        if os.path.exists(context_groups_file):
            with open(context_groups_file, "w") as f:
                dict_to_save = {
                    title: [item.dict() for item in context_items]
                    for title, context_items in self._saved_context_groups.items()
                }
                json.dump(dict_to_save, f)

    async def save_context_group(self, title: str, context_items: List[ContextItem]):
        self._saved_context_groups[title] = context_items
        await self.update_subscribers()

        # Update saved context groups
        self._persist_context_groups()

        posthog_logger.capture_event(
            "save_context_group", {"title": title, "length": len(context_items)}
        )

    async def select_context_group(self, id: str):
        if id not in self._saved_context_groups:
            logger.warning(f"Context group {id} not found")
            return
        context_group = self._saved_context_groups[id]
        await self.context_manager.clear_context()
        for item in context_group:
            await self.context_manager.manually_add_context_item(item)
        await self.update_subscribers()

        posthog_logger.capture_event(
            "select_context_group", {"title": id, "length": len(context_group)}
        )
        dev_data_logger.capture(
            "select_context_group", {"title": id, "items": context_group}
        )

    async def delete_context_group(self, id: str):
        if id not in self._saved_context_groups:
            logger.warning(f"Context group {id} not found")
            return
        del self._saved_context_groups[id]
        await self.update_subscribers()

        # Update saved context groups
        self._persist_context_groups()

        posthog_logger.capture_event("delete_context_group", {"title": id})