path: root/server/continuedev/plugins/context_providers/highlighted_code.py

import os
from typing import Any, Dict, List, Optional

from pydantic import BaseModel

from ...core.context import (
    ContextItem,
    ContextItemDescription,
    ContextItemId,
    ContextProvider,
)
from ...core.main import ChatMessage
from ...models.filesystem import RangeInFileWithContents
from ...models.main import Range


class HighlightedRangeContextItem(BaseModel):
    rif: RangeInFileWithContents
    item: ContextItem


class HighlightedCodeContextProvider(ContextProvider):
    """
    The ContextProvider class is a plugin that lets you provide new information to the LLM by typing '@'.
    When you type '@', the context provider will be asked to populate a list of options.
    These options will be updated on each keystroke.
    When you hit enter on an option, the context provider will add that item to the autopilot's list of context (which is all stored in the ContextManager object).
    """

    title = "code"
    display_title = "Highlighted Code"
    description = "Highlight code"
    dynamic = True

    ide: Any  # IdeProtocolServer

    highlighted_ranges: List[HighlightedRangeContextItem] = []
    adding_highlighted_code: bool = True
    # Controls whether you can have more than one highlighted range. Now always True.

    should_get_fallback_context_item: bool = True
    last_added_fallback: bool = False

    async def _get_fallback_context_item(self) -> HighlightedRangeContextItem:
        # Used to automatically include the currently open file. Disabled for now.
        return None

        if not self.should_get_fallback_context_item:
            return None

        visible_files = await self.ide.getVisibleFiles()
        if len(visible_files) > 0:
            content = await self.ide.readFile(visible_files[0])
            rif = RangeInFileWithContents.from_entire_file(visible_files[0], content)

            item = self._rif_to_context_item(rif, 0, True)
            item.description.name = self._rif_to_name(rif, show_line_nums=False)

            self.last_added_fallback = True
            return HighlightedRangeContextItem(rif=rif, item=item)

        return None

    async def get_selected_items(self) -> List[ContextItem]:
        items = [hr.item for hr in self.highlighted_ranges]

        if len(items) == 0 and (
            fallback_item := await self._get_fallback_context_item()
        ):
            items = [fallback_item.item]

        return items

    async def get_chat_messages(self) -> List[ContextItem]:
        ranges = self.highlighted_ranges
        if len(ranges) == 0 and (
            fallback_item := await self._get_fallback_context_item()
        ):
            ranges = [fallback_item]

        return [
            ChatMessage(
                role="user",
                content=f"Code in this file is highlighted ({r.rif.filepath}):\n```\n{r.rif.contents}\n```",
                summary=f"Code in this file is highlighted: {r.rif.filepath}",
            )
            for r in ranges
        ]

    def _make_sure_is_editing_range(self):
        """If none of the highlighted ranges are currently being edited, the first should be selected"""
        if len(self.highlighted_ranges) == 0:
            return
        if not any(map(lambda x: x.item.editing, self.highlighted_ranges)):
            self.highlighted_ranges[0].item.editing = True

    def _disambiguate_highlighted_ranges(self):
        """If any files have the same name, also display their folder name"""
        name_status: Dict[
            str, set
        ] = {}  # basename -> set of full paths with that basename
        for hr in self.highlighted_ranges:
            basename = os.path.basename(hr.rif.filepath)
            if basename in name_status:
                name_status[basename].add(hr.rif.filepath)
            else:
                name_status[basename] = {hr.rif.filepath}

        for hr in self.highlighted_ranges:
            basename = os.path.basename(hr.rif.filepath)
            if len(name_status[basename]) > 1:
                hr.item.description.name = self._rif_to_name(
                    hr.rif,
                    display_filename=os.path.join(
                        os.path.basename(os.path.dirname(hr.rif.filepath)), basename
                    ),
                )
            else:
                hr.item.description.name = self._rif_to_name(
                    hr.rif, display_filename=basename
                )

    async def provide_context_items(self, workspace_dir: str) -> List[ContextItem]:
        return []

    async def get_item(self, id: ContextItemId, query: str) -> ContextItem:
        raise NotImplementedError()

    async def clear_context(self):
        self.highlighted_ranges = []
        self.adding_highlighted_code = False
        self.should_get_fallback_context_item = True
        self.last_added_fallback = False

    async def delete_context_with_ids(
        self, ids: List[ContextItemId]
    ) -> List[ContextItem]:
        ids_to_delete = [id.item_id for id in ids]

        kept_ranges = []
        for hr in self.highlighted_ranges:
            if hr.item.description.id.item_id not in ids_to_delete:
                kept_ranges.append(hr)
        self.highlighted_ranges = kept_ranges

        self._make_sure_is_editing_range()

        if len(self.highlighted_ranges) == 0 and self.last_added_fallback:
            self.should_get_fallback_context_item = False

        return [hr.item for hr in self.highlighted_ranges]

    def _rif_to_name(
        self,
        rif: RangeInFileWithContents,
        display_filename: str = None,
        show_line_nums: bool = True,
    ) -> str:
        line_nums = (
            f" ({rif.range.start.line + 1}-{rif.range.end.line + 1})"
            if show_line_nums
            else ""
        )
        return f"{display_filename or os.path.basename(rif.filepath)}{line_nums}"

    def _rif_to_context_item(
        self, rif: RangeInFileWithContents, idx: int, editing: bool
    ) -> ContextItem:
        return ContextItem(
            description=ContextItemDescription(
                name=self._rif_to_name(rif),
                description=rif.filepath,
                id=ContextItemId(provider_title=self.title, item_id=str(idx)),
            ),
            content=rif.contents,
            editing=editing if editing is not None else False,
            editable=True,
        )

    async def handle_highlighted_code(
        self,
        range_in_files: List[RangeInFileWithContents],
        edit: Optional[bool] = False,
    ):
        self.should_get_fallback_context_item = True
        self.last_added_fallback = False

        # Filter out rifs from ~/.continue/diffs folder
        range_in_files = [
            rif
            for rif in range_in_files
            if not os.path.dirname(rif.filepath)
            == os.path.expanduser("~/.continue/diffs")
        ]

        # If not adding highlighted code
        if not self.adding_highlighted_code:
            if (
                len(self.highlighted_ranges) == 1
                and len(range_in_files) <= 1
                and (
                    len(range_in_files) == 0
                    or range_in_files[0].range.start == range_in_files[0].range.end
                )
            ):
                # If un-highlighting the range to edit, then remove the range
                self.highlighted_ranges = []
            elif len(range_in_files) > 0:
                # Otherwise, replace the current range with the new one
                # This is the first range to be highlighted
                self.highlighted_ranges = [
                    HighlightedRangeContextItem(
                        rif=range_in_files[0],
                        item=self._rif_to_context_item(range_in_files[0], 0, edit),
                    )
                ]

            return

        # If editing, make sure none of the other ranges are editing
        if edit:
            for hr in self.highlighted_ranges:
                hr.item.editing = False

        # If new range overlaps with any existing, keep the existing but merged
        new_ranges = []
        for i, new_hr in enumerate(range_in_files):
            found_overlap_with = None
            for existing_rif in self.highlighted_ranges:
                if (
                    new_hr.filepath == existing_rif.rif.filepath
                    and new_hr.range.overlaps_with(existing_rif.rif.range)
                ):
                    existing_rif.rif.range = existing_rif.rif.range.merge_with(
                        new_hr.range
                    )
                    found_overlap_with = existing_rif
                    break

            if found_overlap_with is None:
                new_ranges.append(
                    HighlightedRangeContextItem(
                        rif=new_hr,
                        item=self._rif_to_context_item(
                            new_hr, len(self.highlighted_ranges) + i, edit
                        ),
                    )
                )
            elif edit:
                # Want to update the range so it's only the newly selected portion
                found_overlap_with.rif.range = new_hr.range
                found_overlap_with.item.editing = True

        self.highlighted_ranges = self.highlighted_ranges + new_ranges

        self._make_sure_is_editing_range()
        self._disambiguate_highlighted_ranges()

    async def set_editing_at_ids(self, ids: List[str]):
        # Don't do anything if there are no valid ids here
        count = 0
        for hr in self.highlighted_ranges:
            if hr.item.description.id.item_id in ids:
                count += 1

        if count == 0:
            return

        for hr in self.highlighted_ranges:
            hr.item.editing = hr.item.description.id.item_id in ids

    async def add_context_item(
        self, id: ContextItemId, query: str, prev: List[ContextItem] = None
    ) -> List[ContextItem]:
        raise NotImplementedError()

    async def manually_add_context_item(self, context_item: ContextItem):
        full_file_content = await self.ide.readFile(
            context_item.description.description
        )
        self.highlighted_ranges.append(
            HighlightedRangeContextItem(
                rif=RangeInFileWithContents(
                    filepath=context_item.description.description,
                    range=Range.from_lines_snippet_in_file(
                        content=full_file_content,
                        snippet=context_item.content,
                    ),
                    contents=context_item.content,
                ),
                item=context_item,
            )
        )