diff options
Diffstat (limited to 'continuedev/src')
| -rw-r--r-- | continuedev/src/continuedev/core/autopilot.py | 138 | ||||
| -rw-r--r-- | continuedev/src/continuedev/core/config.py | 3 | ||||
| -rw-r--r-- | continuedev/src/continuedev/core/context.py | 205 | ||||
| -rw-r--r-- | continuedev/src/continuedev/core/context_manager.py | 119 | ||||
| -rw-r--r-- | continuedev/src/continuedev/core/main.py | 59 | ||||
| -rw-r--r-- | continuedev/src/continuedev/core/sdk.py | 31 | ||||
| -rw-r--r-- | continuedev/src/continuedev/libs/context_providers/highlighted_code_context_provider.py | 191 | ||||
| -rw-r--r-- | continuedev/src/continuedev/libs/util/paths.py | 20 | ||||
| -rw-r--r-- | continuedev/src/continuedev/models/generate_json_schema.py | 3 | ||||
| -rw-r--r-- | continuedev/src/continuedev/server/gui.py | 37 | ||||
| -rw-r--r-- | continuedev/src/continuedev/server/gui_protocol.py | 10 | ||||
| -rw-r--r-- | continuedev/src/continuedev/server/main.py | 28 | ||||
| -rw-r--r-- | continuedev/src/continuedev/server/meilisearch_server.py | 56 | ||||
| -rw-r--r-- | continuedev/src/continuedev/server/session_manager.py | 6 |
14 files changed, 608 insertions, 298 deletions
diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py index abda50b0..c0f95414 100644 --- a/continuedev/src/continuedev/core/autopilot.py +++ b/continuedev/src/continuedev/core/autopilot.py @@ -9,10 +9,12 @@ from pydantic import root_validator from ..models.filesystem import RangeInFileWithContents from ..models.filesystem_edit import FileEditWithFullContents from .observation import Observation, InternalErrorObservation +from .context import ContextItem, ContextItemDescription, ContextItemId, ContextManager +from ..libs.context_providers.highlighted_code_context_provider import HighlightedCodeContextProvider from ..server.ide_protocol import AbstractIdeProtocolServer from ..libs.util.queue import AsyncSubscriptionQueue from ..models.main import ContinueBaseModel -from .main import Context, ContinueCustomException, HighlightedRangeContext, Policy, History, FullState, Step, HistoryNode +from .main import Context, ContinueCustomException, Policy, History, FullState, Step, HistoryNode from ..steps.core.core import ReversibleStep, ManualEditStep, UserInputStep from ..libs.util.telemetry import capture_event from .sdk import ContinueSDK @@ -47,10 +49,11 @@ class Autopilot(ContinueBaseModel): history: History = History.from_empty() context: Context = Context() full_state: Union[FullState, None] = None - _on_update_callbacks: List[Callable[[FullState], None]] = [] - + context_manager: Union[ContextManager, None] = None continue_sdk: ContinueSDK = None + _on_update_callbacks: List[Callable[[FullState], None]] = [] + _active: bool = False _should_halt: bool = False _main_user_input_queue: List[str] = [] @@ -62,6 +65,14 @@ class Autopilot(ContinueBaseModel): async def create(cls, policy: Policy, ide: AbstractIdeProtocolServer, full_state: FullState) -> "Autopilot": autopilot = cls(ide=ide, policy=policy) autopilot.continue_sdk = await ContinueSDK.create(autopilot) + + # Load documents into the search index + autopilot.context_manager = ContextManager( + autopilot.continue_sdk.config.context_providers + [ + HighlightedCodeContextProvider(ide=ide) + ]) + await autopilot.context_manager.load_index() + return autopilot class Config: @@ -75,15 +86,16 @@ class Autopilot(ContinueBaseModel): values['history'] = full_state.history return values - def get_full_state(self) -> FullState: + async def get_full_state(self) -> FullState: full_state = FullState( history=self.history, active=self._active, user_input_queue=self._main_user_input_queue, default_model=self.continue_sdk.config.default_model, - highlighted_ranges=self._highlighted_ranges, slash_commands=self.get_available_slash_commands(), - adding_highlighted_code=self._adding_highlighted_code, + adding_highlighted_code=self.context_manager.context_providers[ + "code"].adding_highlighted_code, + selected_context_items=await self.context_manager.get_selected_items() ) self.full_state = full_state return full_state @@ -104,8 +116,8 @@ class Autopilot(ContinueBaseModel): self._main_user_input_queue = [] self._active = False - # Also remove all context - self._highlighted_ranges = [] + # Clear context + await self.context_manager.clear_context() await self.update_subscribers() @@ -114,7 +126,7 @@ class Autopilot(ContinueBaseModel): self._on_update_callbacks.append(callback) async def update_subscribers(self): - full_state = self.get_full_state() + full_state = await self.get_full_state() for callback in self._on_update_callbacks: await callback(full_state) @@ -159,81 +171,10 @@ class Autopilot(ContinueBaseModel): step = tb_step.step({"output": output, **tb_step.params}) await self._run_singular_step(step) - _highlighted_ranges: List[HighlightedRangeContext] = [] - _adding_highlighted_code: bool = False - - 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.editing, self._highlighted_ranges)): - self._highlighted_ranges[0].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 rif in self._highlighted_ranges: - basename = os.path.basename(rif.range.filepath) - if basename in name_status: - name_status[basename].add(rif.range.filepath) - else: - name_status[basename] = {rif.range.filepath} - - for rif in self._highlighted_ranges: - basename = os.path.basename(rif.range.filepath) - if len(name_status[basename]) > 1: - rif.display_name = os.path.join( - os.path.basename(os.path.dirname(rif.range.filepath)), basename) - else: - rif.display_name = basename - async def handle_highlighted_code(self, range_in_files: List[RangeInFileWithContents]): - # 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")] - - # Make sure all filepaths are relative to workspace - workspace_path = self.continue_sdk.ide.workspace_directory - - # 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 = [] - await self.update_subscribers() - 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 = [HighlightedRangeContext( - range=range_in_files[0], editing=True, pinned=False, display_name=os.path.basename(range_in_files[0].filepath))] - await self.update_subscribers() - return - - # If current range overlaps with any others, delete them and only keep the new range - new_ranges = [] - for i, rif in enumerate(self._highlighted_ranges): - found_overlap = False - for new_rif in range_in_files: - if rif.range.filepath == new_rif.filepath and rif.range.range.overlaps_with(new_rif.range): - found_overlap = True - break - - # Also don't allow multiple ranges in same file with same content. This is useless to the model, and avoids - # the bug where cmd+f causes repeated highlights - if rif.range.filepath == new_rif.filepath and rif.range.contents == new_rif.contents: - found_overlap = True - break - - if not found_overlap: - new_ranges.append(rif) - - self._highlighted_ranges = new_ranges + [HighlightedRangeContext( - range=rif, editing=False, pinned=False, display_name=os.path.basename(rif.filepath) - ) for rif in range_in_files] - - self._make_sure_is_editing_range() - self._disambiguate_highlighted_ranges() + # Add to context manager + await self.context_manager.context_providers["code"].handle_highlighted_code( + range_in_files) await self.update_subscribers() @@ -250,29 +191,16 @@ class Autopilot(ContinueBaseModel): await self.update_subscribers() - async def delete_context_at_indices(self, indices: List[int]): - kept_ranges = [] - for i, rif in enumerate(self._highlighted_ranges): - if i not in indices: - kept_ranges.append(rif) - self._highlighted_ranges = kept_ranges - - self._make_sure_is_editing_range() - + async def delete_context_with_ids(self, ids: List[str]): + await self.context_manager.delete_context_with_ids(ids) await self.update_subscribers() async def toggle_adding_highlighted_code(self): - self._adding_highlighted_code = not self._adding_highlighted_code - await self.update_subscribers() - - async def set_editing_at_indices(self, indices: List[int]): - for i in range(len(self._highlighted_ranges)): - self._highlighted_ranges[i].editing = i in indices + 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_pinned_at_indices(self, indices: List[int]): - for i in range(len(self._highlighted_ranges)): - self._highlighted_ranges[i].pinned = i in indices + async def set_editing_at_ids(self, ids: List[str]): + 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]: @@ -437,10 +365,6 @@ class Autopilot(ContinueBaseModel): if len(self._main_user_input_queue) > 1: return - # Remove context unless pinned - # self._highlighted_ranges = [ - # hr for hr in self._highlighted_ranges if hr.pinned] - # 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. @@ -456,3 +380,7 @@ class Autopilot(ContinueBaseModel): await self._request_halt() await self.reverse_to_index(index) await self.run_from_step(UserInputStep(user_input=user_input)) + + async def select_context_item(self, id: str, query: str): + await self.context_manager.select_context_item(id, query) + await self.update_subscribers() diff --git a/continuedev/src/continuedev/core/config.py b/continuedev/src/continuedev/core/config.py index 54f15143..bb9ca323 100644 --- a/continuedev/src/continuedev/core/config.py +++ b/continuedev/src/continuedev/core/config.py @@ -1,6 +1,7 @@ import json import os from .main import Step +from .context import ContextProvider from pydantic import BaseModel, validator from typing import List, Literal, Optional, Dict, Type, Union import yaml @@ -50,6 +51,8 @@ class ContinueConfig(BaseModel): system_message: Optional[str] = None azure_openai_info: Optional[AzureInfo] = None + context_providers: List[ContextProvider] = [] + # Want to force these to be the slash commands for now @validator('slash_commands', pre=True) def default_slash_commands_validator(cls, v): diff --git a/continuedev/src/continuedev/core/context.py b/continuedev/src/continuedev/core/context.py new file mode 100644 index 00000000..67bba651 --- /dev/null +++ b/continuedev/src/continuedev/core/context.py @@ -0,0 +1,205 @@ + +from abc import abstractmethod +from typing import Dict, List +import meilisearch +from pydantic import BaseModel + + +from .main import ChatMessage, ContextItem, ContextItemDescription, ContextItemId +from ..server.meilisearch_server import check_meilisearch_running + + +SEARCH_INDEX_NAME = "continue_context_items" + + +class ContextProvider(BaseModel): + """ + 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: str + + selected_items: List[ContextItem] = [] + + async def get_selected_items(self) -> List[ContextItem]: + """ + Returns all of the selected ContextItems. + + Default implementation simply returns self.selected_items. + + Other implementations may add an async processing step. + """ + return self.selected_items + + @abstractmethod + async def provide_context_items(self) -> List[ContextItem]: + """ + Provide documents for search index. This is run on startup. + + This is the only method that must be implemented. + """ + + async def get_chat_messages(self) -> List[ChatMessage]: + """ + Returns all of the chat messages for the context provider. + + Default implementation has a string template. + """ + return [ChatMessage(role="user", content=f"{item.description.name}: {item.description.description}\n\n{item.content}", summary=item.description.description) for item in await self.get_selected_items()] + + async def get_item(self, id: ContextItemId, query: str, search_client: meilisearch.Client) -> ContextItem: + """ + Returns the ContextItem with the given id. + + Default implementation uses the search index to get the item. + """ + result = search_client.index( + SEARCH_INDEX_NAME).get_document(id.to_string()) + return ContextItem( + description=ContextItemDescription( + name=result.name, + description=result.description, + id=id + ), + content=result.content + ) + + async def delete_context_with_ids(self, ids: List[ContextItemId]): + """ + Deletes the ContextItems with the given IDs, lets ContextProviders recalculate. + + Default implementation simply deletes those with the given ids. + """ + id_strings = {id.to_string() for id in ids} + self.selected_items = list( + filter(lambda item: item.description.id.to_string() not in id_strings, self.selected_items)) + + async def clear_context(self): + """ + Clears all context. + + Default implementation simply clears the selected items. + """ + self.selected_items = [] + + async def add_context_item(self, id: ContextItemId, query: str, search_client: meilisearch.Client): + """ + Adds the given ContextItem to the list of ContextItems. + + Default implementation simply appends the item, not allowing duplicates. + + This method also allows you not to have to load all of the information until an item is selected. + """ + + # Don't add duplicate context + for item in self.selected_items: + if item.description.id.item_id == id.item_id: + return + + new_item = await self.get_item(id, query, search_client) + self.selected_items.append(new_item) + + +class ContextManager: + """ + The context manager is responsible for storing the context to be passed to the LLM, including + - ContextItems (highlighted code, GitHub Issues, etc.) + - ChatMessages in the history + - System Message + - Functions + + It is responsible for compiling all of this information into a single prompt without exceeding the token limit. + """ + + async def get_selected_items(self) -> List[ContextItem]: + """ + Returns all of the selected ContextItems. + """ + return sum([await provider.get_selected_items() for provider in self.context_providers.values()], []) + + async def get_chat_messages(self) -> List[ChatMessage]: + """ + Returns chat messages from each provider. + """ + return sum([await provider.get_chat_messages() for provider in self.context_providers.values()], []) + + search_client: meilisearch.Client + + def __init__(self, context_providers: List[ContextProvider]): + self.search_client = meilisearch.Client('http://localhost:7700') + + # If meilisearch isn't running, don't use any ContextProviders that might depend on it + if not check_meilisearch_running(): + context_providers = list( + filter(lambda cp: cp.title == "code", context_providers)) + + self.context_providers = { + prov.title: prov for prov in context_providers} + self.provider_titles = { + provider.title for provider in context_providers} + + async def load_index(self): + for _, provider in self.context_providers.items(): + context_items = await provider.provide_context_items() + documents = [ + { + "id": item.description.id.to_string(), + "name": item.description.name, + "description": item.description.description, + "content": item.content + } + for item in context_items + ] + if len(documents) > 0: + self.search_client.index( + SEARCH_INDEX_NAME).add_documents(documents) + + # def compile_chat_messages(self, max_tokens: int) -> List[Dict]: + # """ + # Compiles the chat prompt into a single string. + # """ + # return compile_chat_messages(self.model, self.chat_history, max_tokens, self.prompt, self.functions, self.system_message) + + async def select_context_item(self, id: str, query: str): + """ + Selects the ContextItem with the given id. + """ + id: ContextItemId = ContextItemId.from_string(id) + if id.provider_title not in self.provider_titles: + raise ValueError( + f"Context provider with title {id.provider_title} not found") + + await self.context_providers[id.provider_title].add_context_item(id, query, self.search_client) + + async def delete_context_with_ids(self, ids: List[str]): + """ + Deletes the ContextItems with the given IDs, lets ContextProviders recalculate. + """ + + # Group by provider title + provider_title_to_ids: Dict[str, List[ContextItemId]] = {} + for id in ids: + id: ContextItemId = ContextItemId.from_string(id) + if id.provider_title not in provider_title_to_ids: + provider_title_to_ids[id.provider_title] = [] + provider_title_to_ids[id.provider_title].append(id) + + # Recalculate context for each updated provider + for provider_title, ids in provider_title_to_ids.items(): + await self.context_providers[provider_title].delete_context_with_ids(ids) + + async def clear_context(self): + """ + Clears all context. + """ + for provider in self.context_providers.values(): + await self.context_providers[provider.title].clear_context() + + +""" +Should define "ArgsTransformer" and "PromptTransformer" classes for the different LLMs. A standard way for them to ingest the +same format of prompts so you don't have to redo all of this logic. +""" diff --git a/continuedev/src/continuedev/core/context_manager.py b/continuedev/src/continuedev/core/context_manager.py deleted file mode 100644 index 37905535..00000000 --- a/continuedev/src/continuedev/core/context_manager.py +++ /dev/null @@ -1,119 +0,0 @@ - -from abc import ABC, abstractmethod, abstractproperty -from ast import List -from pydantic import BaseModel - -from ..libs.util.count_tokens import compile_chat_messages - - -class ContextItemDescription(BaseModel): - """ - A ContextItemDescription is a description of a ContextItem that is displayed to the user when they type '@'. - - The id can be used to retrieve the ContextItem from the ContextManager. - """ - name: str - description: str - id: str - - -class ContextItem(BaseModel): - """ - A ContextItem is a single item that is stored in the ContextManager. - """ - description: ContextItemDescription - content: str - - -class ContextManager(ABC): - """ - The context manager is responsible for storing the context to be passed to the LLM, including - - ContextItems (highlighted code, GitHub Issues, etc.) - - ChatMessages in the history - - System Message - - Functions - - It is responsible for compiling all of this information into a single prompt without exceeding the token limit. - """ - - def compile_chat_messages(self, max_tokens: int) -> List[Dict]: - """ - Compiles the chat prompt into a single string. - """ - return compile_chat_messages(self.model, self.chat_history, max_tokens, self.prompt, self.functions, self.system_message) - - -""" -Should define "ArgsTransformer" and "PromptTransformer" classes for the different LLMs. A standard way for them to ingest the -same format of prompts so you don't have to redo all of this logic. -""" - - -class ContextProvider(ABC): - """ - 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: str - - @abstractmethod - async def load(self): - """ - Loads the ContextProvider, possibly reading persisted data from disk. This will be called on startup. - """ - - @abstractmethod - async def save(self): - """ - Saves the ContextProvider, possibly writing persisted data to disk. This will be called upon cache refresh. - """ - - @abstractmethod - async def refresh_cache(self): - """ - Refreshes the cache of items. This will be called on startup and periodically. - """ - - @abstractmethod - async def get_item_descriptions(self, query: str) -> List[ContextItemDescription]: - """ - Returns a list of options that should be displayed to the user. - """ - - @abstractmethod - async def get_item(self, id: str) -> ContextItem: - """ - Returns the ContextItem with the given id. This allows you not to have to load all of the information until an item is selected. - """ - - @abstractmethod - async def should_refresh(self) -> bool: - """ - Returns whether the ContextProvider should be refreshed. - - For example, embeddings might need to be recalculated after commits, - or GitHub issues might need to be refreshed after a new issue is created. - - This method will be called every startup? Every once in a while? Every hour? - User defined? Maybe just have a schedule instead of this method. - """ - - -class GitHubIssuesContextProvider(ContextProvider): - """ - The GitHubIssuesContextProvider is a ContextProvider that allows you to search GitHub issues in a repo. - """ - - title = "issues" - - def __init__(self, repo: str): - self.repo = repo - - async def get_item_descriptions(self, query: str) -> List[ContextItemDescription]: - pass - - async def get_item(self, id: str) -> ContextItem: - pass diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py index 50d01f8d..6c6adccc 100644 --- a/continuedev/src/continuedev/core/main.py +++ b/continuedev/src/continuedev/core/main.py @@ -1,12 +1,11 @@ import json -from textwrap import dedent -from typing import Callable, Coroutine, Dict, Generator, List, Literal, Tuple, Union +from typing import Coroutine, Dict, List, Literal, Union +from pydantic.schema import schema + -from ..models.filesystem import RangeInFileWithContents from ..models.main import ContinueBaseModel -from pydantic import validator +from pydantic import BaseModel, validator from .observation import Observation -from pydantic.schema import schema ChatMessageRole = Literal["assistant", "user", "system", "function"] @@ -201,12 +200,48 @@ class SlashCommandDescription(ContinueBaseModel): description: str -class HighlightedRangeContext(ContinueBaseModel): - """Context for a highlighted range""" - range: RangeInFileWithContents - editing: bool - pinned: bool - display_name: str +class ContextItemId(BaseModel): + """ + A ContextItemId is a unique identifier for a ContextItem. + """ + provider_title: str + item_id: str + + def to_string(self) -> str: + return f"{self.provider_title}-{self.item_id}" + + @staticmethod + def from_string(string: str) -> 'ContextItemId': + provider_title, item_id = string.split('-') + return ContextItemId(provider_title=provider_title, item_id=item_id) + + +class ContextItemDescription(BaseModel): + """ + A ContextItemDescription is a description of a ContextItem that is displayed to the user when they type '@'. + + The id can be used to retrieve the ContextItem from the ContextManager. + """ + name: str + description: str + id: ContextItemId + + +class ContextItem(BaseModel): + """ + A ContextItem is a single item that is stored in the ContextManager. + """ + description: ContextItemDescription + content: str + + @validator('content', pre=True) + def content_must_be_string(cls, v): + if v is None: + return '' + return v + + editing: bool = False + editable: bool = False class FullState(ContinueBaseModel): @@ -215,9 +250,9 @@ class FullState(ContinueBaseModel): active: bool user_input_queue: List[str] default_model: str - highlighted_ranges: List[HighlightedRangeContext] slash_commands: List[SlashCommandDescription] adding_highlighted_code: bool + selected_context_items: List[ContextItem] class ContinueSDK: diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py index 4100efa6..59f33707 100644 --- a/continuedev/src/continuedev/core/sdk.py +++ b/continuedev/src/continuedev/core/sdk.py @@ -3,8 +3,10 @@ from functools import cached_property from typing import Coroutine, Dict, Union import os + from ..steps.core.core import DefaultModelEditCodeStep from ..models.main import Range +from .context import ContextItem from .abstract_sdk import AbstractContinueSDK from .config import ContinueConfig, load_config, load_global_config, update_global_config from ..models.filesystem_edit import FileEdit, FileSystemEdit, AddFile, DeleteFile, AddDirectory, DeleteDirectory @@ -289,28 +291,13 @@ class ContinueSDK(AbstractContinueSDK): async def get_chat_context(self) -> List[ChatMessage]: history_context = self.history.to_chat_history() - highlighted_code = [ - hr.range for hr in self.__autopilot._highlighted_ranges] - - preface = "The following code is highlighted" - - # If no higlighted ranges, use first file as context - if len(highlighted_code) == 0: - preface = "The following file is open" - visible_files = await self.ide.getVisibleFiles() - if len(visible_files) > 0: - content = await self.ide.readFile(visible_files[0]) - highlighted_code = [ - RangeInFileWithContents.from_entire_file(visible_files[0], content)] - - for rif in highlighted_code: - msg = ChatMessage(content=f"{preface} ({rif.filepath}):\n```\n{rif.contents}\n```", - role="user", summary=f"{preface}: {rif.filepath}") - - # Don't insert after latest user message or function call - i = -1 - if len(history_context) > 0 and (history_context[i].role == "user" or history_context[i].role == "function"): - i -= 1 + + context_messages: List[ChatMessage] = await self.__autopilot.context_manager.get_chat_messages() + + # Insert at the end, but don't insert after latest user message or function call + i = -2 if (len(history_context) > 0 and ( + history_context[-1].role == "user" or history_context[-1].role == "function")) else -1 + for msg in context_messages: history_context.insert(i, msg) return history_context diff --git a/continuedev/src/continuedev/libs/context_providers/highlighted_code_context_provider.py b/continuedev/src/continuedev/libs/context_providers/highlighted_code_context_provider.py new file mode 100644 index 00000000..23d4fc86 --- /dev/null +++ b/continuedev/src/continuedev/libs/context_providers/highlighted_code_context_provider.py @@ -0,0 +1,191 @@ +import os +from typing import Any, Dict, List + +import meilisearch +from ...core.main import ChatMessage +from ...models.filesystem import RangeInFile, RangeInFileWithContents +from ...core.context import ContextItem, ContextItemDescription, ContextItemId +from pydantic import BaseModel + + +class HighlightedRangeContextItem(BaseModel): + rif: RangeInFileWithContents + item: ContextItem + + +class HighlightedCodeContextProvider(BaseModel): + """ + 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" + + ide: Any # IdeProtocolServer + + highlighted_ranges: List[HighlightedRangeContextItem] = [] + adding_highlighted_code: bool = False + + should_get_fallback_context_item: bool = True + last_added_fallback: bool = False + + async def _get_fallback_context_item(self) -> HighlightedRangeContextItem: + 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: + 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) -> List[ContextItem]: + return [] + + async def delete_context_with_ids(self, ids: List[ContextItemId]) -> List[ContextItem]: + indices_to_delete = [ + int(id.item_id) for id in ids + ] + + kept_ranges = [] + for i, hr in enumerate(self.highlighted_ranges): + if i not in indices_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, + editable=True + ) + + async def handle_highlighted_code(self, range_in_files: List[RangeInFileWithContents]): + 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, True))] + + return + + # If current range overlaps with any others, delete them and only keep the new range + new_ranges = [] + for i, hr in enumerate(self.highlighted_ranges): + found_overlap = False + for new_rif in range_in_files: + if hr.rif.filepath == new_rif.filepath and hr.rif.range.overlaps_with(new_rif.range): + found_overlap = True + break + + # Also don't allow multiple ranges in same file with same content. This is useless to the model, and avoids + # the bug where cmd+f causes repeated highlights + if hr.rif.filepath == new_rif.filepath and hr.rif.contents == new_rif.contents: + found_overlap = True + break + + if not found_overlap: + new_ranges.append(HighlightedRangeContextItem(rif=hr.rif, item=self._rif_to_context_item( + hr.rif, len(new_ranges), False))) + + self.highlighted_ranges = new_ranges + [HighlightedRangeContextItem(rif=rif, item=self._rif_to_context_item( + rif, len(new_ranges) + idx, False)) for idx, rif in enumerate(range_in_files)] + + self._make_sure_is_editing_range() + self._disambiguate_highlighted_ranges() + + async def set_editing_at_ids(self, ids: List[str]): + for hr in self.highlighted_ranges: + hr.item.editing = hr.item.description.id.to_string() in ids + + async def add_context_item(self, id: ContextItemId, query: str, search_client: meilisearch.Client, prev: List[ContextItem] = None) -> List[ContextItem]: + raise NotImplementedError() diff --git a/continuedev/src/continuedev/libs/util/paths.py b/continuedev/src/continuedev/libs/util/paths.py index fddef887..d6ce13b3 100644 --- a/continuedev/src/continuedev/libs/util/paths.py +++ b/continuedev/src/continuedev/libs/util/paths.py @@ -2,16 +2,26 @@ import os from ..constants.main import CONTINUE_SESSIONS_FOLDER, CONTINUE_GLOBAL_FOLDER, CONTINUE_SERVER_FOLDER -def getGlobalFolderPath(): - return os.path.join(os.path.expanduser("~"), CONTINUE_GLOBAL_FOLDER) +def getGlobalFolderPath(): + path = os.path.join(os.path.expanduser("~"), CONTINUE_GLOBAL_FOLDER) + os.makedirs(path, exist_ok=True) + return path def getSessionsFolderPath(): - return os.path.join(getGlobalFolderPath(), CONTINUE_SESSIONS_FOLDER) + path = os.path.join(getGlobalFolderPath(), CONTINUE_SESSIONS_FOLDER) + os.makedirs(path, exist_ok=True) + return path + def getServerFolderPath(): - return os.path.join(getGlobalFolderPath(), CONTINUE_SERVER_FOLDER) + path = os.path.join(getGlobalFolderPath(), CONTINUE_SERVER_FOLDER) + os.makedirs(path, exist_ok=True) + return path + def getSessionFilePath(session_id: str): - return os.path.join(getSessionsFolderPath(), f"{session_id}.json")
\ No newline at end of file + path = os.path.join(getSessionsFolderPath(), f"{session_id}.json") + os.makedirs(os.path.dirname(path), exist_ok=True) + return path diff --git a/continuedev/src/continuedev/models/generate_json_schema.py b/continuedev/src/continuedev/models/generate_json_schema.py index 6cebf429..06614984 100644 --- a/continuedev/src/continuedev/models/generate_json_schema.py +++ b/continuedev/src/continuedev/models/generate_json_schema.py @@ -2,6 +2,7 @@ from .main import * from .filesystem import RangeInFile, FileEdit from .filesystem_edit import FileEditWithFullContents from ..core.main import History, HistoryNode, FullState +from ..core.context import ContextItem from pydantic import schema_json_of import os @@ -13,6 +14,8 @@ MODELS_TO_GENERATE = [ FileEditWithFullContents ] + [ History, HistoryNode, FullState +] + [ + ContextItem ] RENAMES = { diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py index ae57c0b6..36b2f3fa 100644 --- a/continuedev/src/continuedev/server/gui.py +++ b/continuedev/src/continuedev/server/gui.py @@ -91,25 +91,19 @@ class GUIProtocolServer(AbstractGUIProtocolServer): self.on_clear_history() elif message_type == "delete_at_index": self.on_delete_at_index(data["index"]) - elif message_type == "delete_context_at_indices": - self.on_delete_context_at_indices(data["indices"]) + elif message_type == "delete_context_with_ids": + self.on_delete_context_with_ids(data["ids"]) elif message_type == "toggle_adding_highlighted_code": self.on_toggle_adding_highlighted_code() elif message_type == "set_editing_at_indices": self.on_set_editing_at_indices(data["indices"]) - elif message_type == "set_pinned_at_indices": - self.on_set_pinned_at_indices(data["indices"]) elif message_type == "show_logs_at_index": self.on_show_logs_at_index(data["index"]) + elif message_type == "select_context_item": + self.select_context_item(data["id"], data["query"]) except Exception as e: print(e) - async def send_state_update(self): - state = self.session.autopilot.get_full_state().dict() - await self._send_json("state_update", { - "state": state - }) - def on_main_input(self, input: str): # Do something with user input create_async_task(self.session.autopilot.accept_user_input( @@ -144,10 +138,10 @@ class GUIProtocolServer(AbstractGUIProtocolServer): create_async_task(self.session.autopilot.delete_at_index( index), self.session.autopilot.continue_sdk.ide.unique_id) - def on_delete_context_at_indices(self, indices: List[int]): + def on_delete_context_with_ids(self, ids: List[str]): create_async_task( - self.session.autopilot.delete_context_at_indices( - indices), self.session.autopilot.continue_sdk.ide.unique_id + self.session.autopilot.delete_context_with_ids( + ids), self.session.autopilot.continue_sdk.ide.unique_id ) def on_toggle_adding_highlighted_code(self): @@ -162,18 +156,17 @@ class GUIProtocolServer(AbstractGUIProtocolServer): indices), self.session.autopilot.continue_sdk.ide.unique_id ) - def on_set_pinned_at_indices(self, indices: List[int]): - create_async_task( - self.session.autopilot.set_pinned_at_indices( - indices), self.session.autopilot.continue_sdk.ide.unique_id - ) - def on_show_logs_at_index(self, index: int): name = f"continue_logs.txt" logs = "\n\n############################################\n\n".join( ["This is a log of the exact prompt/completion pairs sent/received from the LLM during this step"] + self.session.autopilot.continue_sdk.history.timeline[index].logs) create_async_task( - self.session.autopilot.ide.showVirtualFile(name, logs)) + self.session.autopilot.ide.showVirtualFile(name, logs), self.session.autopilot.continue_sdk.ide.unique_id) + + def select_context_item(self, id: str, query: str): + """Called when user selects an item from the dropdown""" + create_async_task( + self.session.autopilot.select_context_item(id, query), self.session.autopilot.continue_sdk.ide.unique_id) @router.websocket("/ws") @@ -188,7 +181,7 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we protocol.websocket = websocket # Update any history that may have happened before connection - await protocol.send_state_update() + await protocol.session.autopilot.update_subscribers() while AppStatus.should_exit is False: message = await websocket.receive_text() @@ -214,5 +207,5 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close() - session_manager.persist_session(session.session_id) + await session_manager.persist_session(session.session_id) session_manager.remove_session(session.session_id) diff --git a/continuedev/src/continuedev/server/gui_protocol.py b/continuedev/src/continuedev/server/gui_protocol.py index 9766fcd0..fb230216 100644 --- a/continuedev/src/continuedev/server/gui_protocol.py +++ b/continuedev/src/continuedev/server/gui_protocol.py @@ -1,6 +1,8 @@ from typing import Any, Dict, List from abc import ABC, abstractmethod +from ..core.context import ContextItem + class AbstractGUIProtocolServer(ABC): @abstractmethod @@ -24,10 +26,6 @@ class AbstractGUIProtocolServer(ABC): """Called when the user inputs a step""" @abstractmethod - async def send_state_update(self, state: dict): - """Send a state update to the client""" - - @abstractmethod def on_retry_at_index(self, index: int): """Called when the user requests a retry at a previous index""" @@ -42,3 +40,7 @@ class AbstractGUIProtocolServer(ABC): @abstractmethod def on_delete_at_index(self, index: int): """Called when the user requests to delete a step at a given index""" + + @abstractmethod + def select_context_item(self, id: str, query: str): + """Called when user selects an item from the dropdown""" diff --git a/continuedev/src/continuedev/server/main.py b/continuedev/src/continuedev/server/main.py index 42dc0cc1..7ee64041 100644 --- a/continuedev/src/continuedev/server/main.py +++ b/continuedev/src/continuedev/server/main.py @@ -1,15 +1,20 @@ +import asyncio +import subprocess import time +import meilisearch import psutil import os from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware -from .ide import router as ide_router -from .gui import router as gui_router -from .session_manager import session_manager import atexit import uvicorn import argparse +from .ide import router as ide_router +from .gui import router as gui_router +from .session_manager import session_manager +from .meilisearch_server import start_meilisearch + app = FastAPI() app.include_router(ide_router) @@ -41,15 +46,20 @@ args = parser.parse_args() # log_file = open('output.log', 'a') # sys.stdout = log_file - def run_server(): uvicorn.run(app, host="0.0.0.0", port=args.port) -def cleanup(): +async def cleanup_coroutine(): print("Cleaning up sessions") for session_id in session_manager.sessions: - session_manager.persist_session(session_id) + await session_manager.persist_session(session_id) + + +def cleanup(): + loop = asyncio.new_event_loop() + loop.run_until_complete(cleanup_coroutine()) + loop.close() def cpu_usage_report(): @@ -77,6 +87,12 @@ if __name__ == "__main__": # cpu_thread = threading.Thread(target=cpu_usage_loop) # cpu_thread.start() + try: + start_meilisearch() + except Exception as e: + print("Failed to start MeiliSearch") + print(e) + run_server() except Exception as e: cleanup() diff --git a/continuedev/src/continuedev/server/meilisearch_server.py b/continuedev/src/continuedev/server/meilisearch_server.py new file mode 100644 index 00000000..419f081f --- /dev/null +++ b/continuedev/src/continuedev/server/meilisearch_server.py @@ -0,0 +1,56 @@ +import os +import subprocess + +import meilisearch +from ..libs.util.paths import getServerFolderPath + + +def check_meilisearch_installed() -> bool: + """ + Checks if MeiliSearch is installed. + """ + + serverPath = getServerFolderPath() + meilisearchPath = os.path.join(serverPath, "meilisearch") + + return os.path.exists(meilisearchPath) + + +def check_meilisearch_running() -> bool: + """ + Checks if MeiliSearch is running. + """ + + try: + client = meilisearch.Client('http://localhost:7700') + resp = client.health() + if resp["status"] != "available": + return False + return True + except Exception: + return False + + +def start_meilisearch(): + """ + Starts the MeiliSearch server, wait for it. + """ + + # Doesn't work on windows for now + if not os.name == "posix": + return + + serverPath = getServerFolderPath() + + # Check if MeiliSearch is installed + if not check_meilisearch_installed(): + # Download MeiliSearch + print("Downloading MeiliSearch...") + subprocess.run( + f"curl -L https://install.meilisearch.com | sh", shell=True, check=True, cwd=serverPath) + + # Check if MeiliSearch is running + if not check_meilisearch_running(): + print("Starting MeiliSearch...") + subprocess.Popen(["./meilisearch"], cwd=serverPath, stdout=subprocess.DEVNULL, + stderr=subprocess.STDOUT, close_fds=True, start_new_session=True) diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py index 90172a4e..96daf92c 100644 --- a/continuedev/src/continuedev/server/session_manager.py +++ b/continuedev/src/continuedev/server/session_manager.py @@ -74,7 +74,7 @@ class SessionManager: async def on_update(state: FullState): await session_manager.send_ws_data(session_id, "state_update", { - "state": autopilot.get_full_state().dict() + "state": state.dict() }) autopilot.on_update(on_update) @@ -84,9 +84,9 @@ class SessionManager: def remove_session(self, session_id: str): del self.sessions[session_id] - def persist_session(self, session_id: str): + async def persist_session(self, session_id: str): """Save the session's FullState as a json file""" - full_state = self.sessions[session_id].autopilot.get_full_state() + full_state = await self.sessions[session_id].autopilot.get_full_state() if not os.path.exists(getSessionsFolderPath()): os.mkdir(getSessionsFolderPath()) with open(getSessionFilePath(session_id), "w") as f: |