diff options
Diffstat (limited to 'continuedev/src')
28 files changed, 1077 insertions, 435 deletions
diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py index 9dbced32..42a58423 100644 --- a/continuedev/src/continuedev/core/autopilot.py +++ b/continuedev/src/continuedev/core/autopilot.py @@ -2,24 +2,25 @@ from functools import cached_property import traceback import time from typing import Any, Callable, Coroutine, Dict, List, Union -import os from aiohttp import ClientPayloadError from pydantic import root_validator from ..models.filesystem import RangeInFileWithContents from ..models.filesystem_edit import FileEditWithFullContents from .observation import Observation, InternalErrorObservation +from .context import ContextManager +from ..plugins.context_providers.file import FileContextProvider +from ..plugins.context_providers.highlighted_code 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 ..plugins.steps.core.core import ReversibleStep, ManualEditStep, UserInputStep -from ..libs.util.telemetry import capture_event from .sdk import ContinueSDK -from ..libs.util.step_name_to_steps import get_step_from_name from ..libs.util.traceback_parsers import get_python_traceback, get_javascript_traceback from openai import error as openai_errors from ..libs.util.create_async_task import create_async_task +from ..libs.util.telemetry import posthog_logger def get_error_title(e: Exception) -> str: @@ -50,10 +51,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] = [] @@ -65,6 +67,15 @@ 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 = await ContextManager.create( + autopilot.continue_sdk.config.context_providers + [ + HighlightedCodeContextProvider(ide=ide), + FileContextProvider(workspace_dir=ide.workspace_directory) + ]) + await autopilot.context_manager.load_index() + return autopilot class Config: @@ -78,15 +89,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 @@ -98,17 +110,14 @@ class Autopilot(ContinueBaseModel): "name": x.name, "description": x.description}, self.continue_sdk.config.slash_commands)) or [] return custom_commands + slash_commands - async def change_default_model(self, model: str): - self.continue_sdk.update_default_model(model) - async def clear_history(self): # Reset history self.history = History.from_empty() 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() @@ -117,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,85 +168,13 @@ class Autopilot(ContinueBaseModel): traceback = get_tb_func(output) if traceback is not None: for tb_step in self.continue_sdk.config.on_traceback: - step = get_step_from_name( - tb_step.step_name, {"output": output, **tb_step.params}) + 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() @@ -254,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]: @@ -294,8 +218,8 @@ class Autopilot(ContinueBaseModel): # last_depth = self.history.timeline[i].depth # i -= 1 - capture_event(self.continue_sdk.ide.unique_id, 'step run', { - 'step_name': step.name, 'params': step.dict()}) + posthog_logger.capture_event( + 'step run', {'step_name': step.name, 'params': step.dict()}) if not is_future_step: # Check manual edits buffer, clear out if needed by creating a ManualEditStep @@ -336,8 +260,8 @@ class Autopilot(ContinueBaseModel): # Attach an InternalErrorObservation to the step and unhide it. print( f"Error while running step: \n{error_string}\n{error_title}") - capture_event(self.continue_sdk.ide.unique_id, 'step error', { - 'error_message': error_string, 'error_title': error_title, 'step_name': step.name, 'params': step.dict()}) + 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) @@ -441,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. @@ -460,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 70c4876e..cb9c8977 100644 --- a/continuedev/src/continuedev/core/config.py +++ b/continuedev/src/continuedev/core/config.py @@ -1,14 +1,16 @@ import json import os +from .main import Step +from .context import ContextProvider from pydantic import BaseModel, validator -from typing import List, Literal, Optional, Dict +from typing import List, Literal, Optional, Dict, Type, Union import yaml class SlashCommand(BaseModel): name: str description: str - step_name: str + step: Type[Step] params: Optional[Dict] = {} @@ -19,54 +21,10 @@ class CustomCommand(BaseModel): class OnTracebackSteps(BaseModel): - step_name: str + step: Type[Step] params: Optional[Dict] = {} -DEFAULT_SLASH_COMMANDS = [ - # SlashCommand( - # name="pytest", - # description="Write pytest unit tests for the current file", - # step_name="WritePytestsRecipe", - # params=??) - SlashCommand( - name="edit", - description="Edit code in the current file or the highlighted code", - step_name="EditHighlightedCodeStep", - ), - # SlashCommand( - # name="explain", - # description="Reply to instructions or a question with previous steps and the highlighted code or current file as context", - # step_name="SimpleChatStep", - # ), - SlashCommand( - name="config", - description="Open the config file to create new and edit existing slash commands", - step_name="OpenConfigStep", - ), - SlashCommand( - name="help", - description="Ask a question like '/help what is given to the llm as context?'", - step_name="HelpStep", - ), - SlashCommand( - name="comment", - description="Write comments for the current file or highlighted code", - step_name="CommentCodeStep", - ), - SlashCommand( - name="feedback", - description="Send feedback to improve Continue", - step_name="FeedbackStep", - ), - SlashCommand( - name="clear", - description="Clear step history", - step_name="ClearHistoryStep", - ) -] - - class AzureInfo(BaseModel): endpoint: str engine: str @@ -77,7 +35,7 @@ class ContinueConfig(BaseModel): """ A pydantic class for the continue config file. """ - steps_on_startup: Optional[Dict[str, Dict]] = {} + steps_on_startup: List[Step] = [] disallowed_steps: Optional[List[str]] = [] allow_anonymous_telemetry: Optional[bool] = True default_model: Literal["gpt-3.5-turbo", "gpt-3.5-turbo-16k", @@ -88,88 +46,52 @@ class ContinueConfig(BaseModel): description="This is an example custom command. Use /config to edit it and create more", prompt="Write a comprehensive set of unit tests for the selected code. It should setup, run tests that check for correctness including important edge cases, and teardown. Ensure that the tests are complete and sophisticated. Give the tests just as chat output, don't edit any file.", )] - slash_commands: Optional[List[SlashCommand]] = DEFAULT_SLASH_COMMANDS - on_traceback: Optional[List[OnTracebackSteps]] = [ - OnTracebackSteps(step_name="DefaultOnTracebackStep")] + slash_commands: Optional[List[SlashCommand]] = [] + on_traceback: Optional[List[OnTracebackSteps]] = [] 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): - return DEFAULT_SLASH_COMMANDS + from ..plugins.steps.open_config import OpenConfigStep + from ..plugins.steps.clear_history import ClearHistoryStep + from ..plugins.steps.feedback import FeedbackStep + from ..plugins.steps.comment_code import CommentCodeStep + from ..plugins.steps.main import EditHighlightedCodeStep + + DEFAULT_SLASH_COMMANDS = [ + SlashCommand( + name="edit", + description="Edit code in the current file or the highlighted code", + step=EditHighlightedCodeStep, + ), + SlashCommand( + name="config", + description="Open the config file to create new and edit existing slash commands", + step=OpenConfigStep, + ), + SlashCommand( + name="comment", + description="Write comments for the current file or highlighted code", + step=CommentCodeStep, + ), + SlashCommand( + name="feedback", + description="Send feedback to improve Continue", + step=FeedbackStep, + ), + SlashCommand( + name="clear", + description="Clear step history", + step=ClearHistoryStep, + ) + ] + + return DEFAULT_SLASH_COMMANDS + v @validator('temperature', pre=True) def temperature_validator(cls, v): return max(0.0, min(1.0, v)) - - -def load_config(config_file: str) -> ContinueConfig: - """ - Load the config file and return a ContinueConfig object. - """ - if not os.path.exists(config_file): - return ContinueConfig() - - _, ext = os.path.splitext(config_file) - if ext == '.yaml': - with open(config_file, 'r') as f: - try: - config_dict = yaml.safe_load(f) - except: - return ContinueConfig() - elif ext == '.json': - with open(config_file, 'r') as f: - try: - config_dict = json.load(f) - except: - return ContinueConfig() - else: - raise ValueError(f'Unknown config file extension: {ext}') - return ContinueConfig(**config_dict) - - -def load_global_config() -> ContinueConfig: - """ - Load the global config file and return a ContinueConfig object. - """ - global_dir = os.path.expanduser('~/.continue') - if not os.path.exists(global_dir): - os.mkdir(global_dir) - - yaml_path = os.path.join(global_dir, 'config.yaml') - if os.path.exists(yaml_path): - with open(config_path, 'r') as f: - try: - config_dict = yaml.safe_load(f) - except: - return ContinueConfig() - else: - config_path = os.path.join(global_dir, 'config.json') - if not os.path.exists(config_path): - with open(config_path, 'w') as f: - json.dump(ContinueConfig().dict(), f, indent=4) - with open(config_path, 'r') as f: - try: - config_dict = json.load(f) - except: - return ContinueConfig() - return ContinueConfig(**config_dict) - - -def update_global_config(config: ContinueConfig): - """ - Update the config file with the given ContinueConfig object. - """ - global_dir = os.path.expanduser('~/.continue') - if not os.path.exists(global_dir): - os.mkdir(global_dir) - - yaml_path = os.path.join(global_dir, 'config.yaml') - if os.path.exists(yaml_path): - with open(config_path, 'w') as f: - yaml.dump(config.dict(), f, indent=4) - else: - config_path = os.path.join(global_dir, 'config.json') - with open(config_path, 'w') as f: - json.dump(config.dict(exclude_unset=False), f, indent=4) diff --git a/continuedev/src/continuedev/core/context.py b/continuedev/src/continuedev/core/context.py new file mode 100644 index 00000000..7d302656 --- /dev/null +++ b/continuedev/src/continuedev/core/context.py @@ -0,0 +1,209 @@ + +from abc import abstractmethod +from typing import Dict, List +from meilisearch_python_async import Client +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: Client) -> ContextItem: + """ + Returns the ContextItem with the given id. + + Default implementation uses the search index to get the item. + """ + result = await 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: 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: Client + + def __init__(self, context_providers: List[ContextProvider], search_client: Client): + self.search_client = search_client + self.context_providers = { + prov.title: prov for prov in context_providers} + self.provider_titles = { + provider.title for provider in context_providers} + + @classmethod + async def create(cls, context_providers: List[ContextProvider]): + search_client = Client('http://localhost:7700') + health = await search_client.health() + if not health.status == "available": + print("MeiliSearch not running, avoiding any dependent context providers") + context_providers = list( + filter(lambda cp: cp.title == "code", context_providers)) + + return cls(context_providers, search_client) + + 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: + await 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/main.py b/continuedev/src/continuedev/core/main.py index 50d01f8d..df9b98ef 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,57 @@ 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 + + @validator('provider_title', 'item_id') + def must_be_valid_id(cls, v): + import re + if not re.match(r'^[0-9a-zA-Z_-]*$', v): + raise ValueError( + "Both provider_title and item_id can only include characters 0-9, a-z, A-Z, -, and _") + return v + + def to_string(self) -> str: + return f"{self.provider_title}-{self.item_id}" + + @staticmethod + def from_string(string: str) -> 'ContextItemId': + provider_title, *rest = string.split('-') + item_id = '-'.join(rest) + 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 +259,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/policy.py b/continuedev/src/continuedev/core/policy.py index dfa0e7f9..d90177b5 100644 --- a/continuedev/src/continuedev/core/policy.py +++ b/continuedev/src/continuedev/core/policy.py @@ -8,8 +8,8 @@ from ..plugins.steps.steps_on_startup import StepsOnStartupStep from .main import Step, History, Policy from .observation import UserInputObservation from ..plugins.steps.core.core import MessageStep -from ..libs.util.step_name_to_steps import get_step_from_name from ..plugins.steps.custom_command import CustomCommandStep +from ..plugins.steps.main import EditHighlightedCodeStep def parse_slash_command(inp: str, config: ContinueConfig) -> Union[None, Step]: @@ -24,7 +24,11 @@ def parse_slash_command(inp: str, config: ContinueConfig) -> Union[None, Step]: if slash_command.name == command_name[1:]: params = slash_command.params params["user_input"] = after_command - return get_step_from_name(slash_command.step_name, params) + try: + return slash_command.step(**params) + except TypeError as e: + raise Exception( + f"Incorrect params used for slash command '{command_name}': {e}") return None @@ -52,7 +56,6 @@ class DefaultPolicy(Policy): - Use `cmd+m` (Mac) / `ctrl+m` (Windows) to open Continue - Use `/help` to ask questions about how to use Continue""")) >> WelcomeStep() >> - # SetupContinueWorkspaceStep() >> # CreateCodebaseIndexChroma() >> StepsOnStartupStep()) @@ -69,6 +72,9 @@ class DefaultPolicy(Policy): if custom_command is not None: return custom_command + if user_input.startswith("/edit"): + return EditHighlightedCodeStep(user_input=user_input[5:]) + return SimpleChatStep() return None diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py index 9d1025e3..e9aefa76 100644 --- a/continuedev/src/continuedev/core/sdk.py +++ b/continuedev/src/continuedev/core/sdk.py @@ -1,4 +1,3 @@ -import asyncio from functools import cached_property from typing import Coroutine, Dict, Union import os @@ -6,7 +5,7 @@ import os from ..plugins.steps.core.core import DefaultModelEditCodeStep from ..models.main import Range from .abstract_sdk import AbstractContinueSDK -from .config import ContinueConfig, load_config, load_global_config, update_global_config +from .config import ContinueConfig from ..models.filesystem_edit import FileEdit, FileSystemEdit, AddFile, DeleteFile, AddDirectory, DeleteDirectory from ..models.filesystem import RangeInFile from ..libs.llm.hf_inference_api import HuggingFaceInferenceAPI @@ -18,6 +17,8 @@ from ..server.ide_protocol import AbstractIdeProtocolServer from .main import Context, ContinueCustomException, History, HistoryNode, Step, ChatMessage from ..plugins.steps.core.core import * from ..libs.llm.proxy_server import ProxyServer +from ..libs.util.telemetry import posthog_logger +from ..libs.util.paths import getConfigFilePath class Autopilot: @@ -144,20 +145,20 @@ class ContinueSDK(AbstractContinueSDK): ide: AbstractIdeProtocolServer models: Models context: Context + config: ContinueConfig __autopilot: Autopilot def __init__(self, autopilot: Autopilot): self.ide = autopilot.ide self.__autopilot = autopilot self.context = autopilot.context - self.config = self._load_config() @classmethod async def create(cls, autopilot: Autopilot) -> "ContinueSDK": sdk = ContinueSDK(autopilot) try: - config = sdk._load_config() + config = sdk._load_config_dot_py() sdk.config = config except Exception as e: print(e) @@ -175,19 +176,6 @@ class ContinueSDK(AbstractContinueSDK): sdk.models = await Models.create(sdk) return sdk - config: ContinueConfig - - def _load_config(self) -> ContinueConfig: - dir = self.ide.workspace_directory - yaml_path = os.path.join(dir, '.continue', 'config.yaml') - json_path = os.path.join(dir, '.continue', 'config.json') - if os.path.exists(yaml_path): - return load_config(yaml_path) - elif os.path.exists(json_path): - return load_config(json_path) - else: - return load_global_config() - @property def history(self) -> History: return self.__autopilot.history @@ -267,16 +255,32 @@ class ContinueSDK(AbstractContinueSDK): async def get_user_secret(self, env_var: str, prompt: str) -> str: return await self.ide.getUserSecret(env_var) + _last_valid_config: ContinueConfig = None + + def _load_config_dot_py(self) -> ContinueConfig: + # Use importlib to load the config file config.py at the given path + path = getConfigFilePath() + try: + import importlib.util + spec = importlib.util.spec_from_file_location("config", path) + config = importlib.util.module_from_spec(spec) + spec.loader.exec_module(config) + self._last_valid_config = config.config + + # When the config is loaded, setup posthog logger + posthog_logger.setup( + self.ide.unique_id, config.config.allow_anonymous_telemetry or True) + + return config.config + except Exception as e: + print("Error loading config.py: ", e) + return ContinueConfig() if self._last_valid_config is None else self._last_valid_config + def get_code_context(self, only_editing: bool = False) -> List[RangeInFileWithContents]: context = list(filter(lambda x: x.editing, self.__autopilot._highlighted_ranges) ) if only_editing else self.__autopilot._highlighted_ranges return [c.range for c in context] - def update_default_model(self, model: str): - config = self.config - config.default_model = model - update_global_config(config) - def set_loading_message(self, message: str): # self.__autopilot.set_loading_message(message) raise NotImplementedError() @@ -286,28 +290,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/constants/default_config.py.txt b/continuedev/src/continuedev/libs/constants/default_config.py.txt new file mode 100644 index 00000000..f80a9ff0 --- /dev/null +++ b/continuedev/src/continuedev/libs/constants/default_config.py.txt @@ -0,0 +1,87 @@ +""" +This is the Continue configuration file. + +If you aren't getting strong typing on these imports, +be sure to select the Python interpreter in ~/.continue/server/env. +""" + +import subprocess + +from continuedev.src.continuedev.core.main import Step +from continuedev.src.continuedev.core.sdk import ContinueSDK +from continuedev.src.continuedev.core.config import CustomCommand, SlashCommand, ContinueConfig +from continuedev.src.continuedev.plugins.context_providers.github import GitHubIssuesContextProvider +from continuedev.src.continuedev.plugins.context_providers.google import GoogleContextProvider + + +class CommitMessageStep(Step): + """ + This is a Step, the building block of Continue. + It can be used below as a slash command, so that + run will be called when you type '/commit'. + """ + async def run(self, sdk: ContinueSDK): + + # Get the root directory of the workspace + dir = sdk.ide.workspace_directory + + # Run git diff in that directory + diff = subprocess.check_output( + ["git", "diff"], cwd=dir).decode("utf-8") + + # Ask gpt-3.5-16k to write a commit message, + # and set it as the description of this step + self.description = await sdk.models.gpt3516k.complete( + f"{diff}\n\nWrite a short, specific (less than 50 chars) commit message about the above changes:") + + +config = ContinueConfig( + + # If set to False, we will not collect any usage data + # See here to learn what anonymous data we collect: https://continue.dev/docs/telemetry + allow_anonymous_telemetry=True, + + # GPT-4 is recommended for best results + # See options here: https://continue.dev/docs/customization#change-the-default-llm + default_model="gpt-4", + + # Set a system message with information that the LLM should always keep in mind + # E.g. "Please give concise answers. Always respond in Spanish." + system_message=None, + + # Set temperature to any value between 0 and 1. Higher values will make the LLM + # more creative, while lower values will make it more predictable. + temperature=0.5, + + # Custom commands let you map a prompt to a shortened slash command + # They are like slash commands, but more easily defined - write just a prompt instead of a Step class + # Their output will always be in chat form + custom_commands=[CustomCommand( + name="test", + description="This is an example custom command. Use /config to edit it and create more", + prompt="Write a comprehensive set of unit tests for the selected code. It should setup, run tests that check for correctness including important edge cases, and teardown. Ensure that the tests are complete and sophisticated. Give the tests just as chat output, don't edit any file.", + )], + + # Slash commands let you run a Step from a slash command + slash_commands=[ + # SlashCommand( + # name="commit", + # description="This is an example slash command. Use /config to edit it and create more", + # step=CommitMessageStep, + # ) + ], + + # Context providers let you quickly select context by typing '@' + # Uncomment the following to + # - quickly reference GitHub issues + # - show Google search results to the LLM + context_providers=[ + # GitHubIssuesContextProvider( + # repo_name="<your github username or organization>/<your repo name>", + # auth_token="<your github auth token>" + # ), + # GoogleContextProvider( + # serper_api_key="<your serper.dev api key>" + # ) + ] +) diff --git a/continuedev/src/continuedev/libs/llm/proxy_server.py b/continuedev/src/continuedev/libs/llm/proxy_server.py index 75c91c4e..f9e3fa01 100644 --- a/continuedev/src/continuedev/libs/llm/proxy_server.py +++ b/continuedev/src/continuedev/libs/llm/proxy_server.py @@ -3,10 +3,10 @@ import json import traceback from typing import Any, Callable, Coroutine, Dict, Generator, List, Literal, Union import aiohttp -from ..util.telemetry import capture_event from ...core.main import ChatMessage from ..llm import LLM -from ..util.count_tokens import DEFAULT_ARGS, DEFAULT_MAX_TOKENS, compile_chat_messages, CHAT_MODELS, count_tokens, format_chat_messages +from ..util.telemetry import posthog_logger +from ..util.count_tokens import DEFAULT_ARGS, compile_chat_messages, count_tokens, format_chat_messages import certifi import ssl @@ -36,7 +36,7 @@ class ProxyServer(LLM): def count_tokens(self, text: str): return count_tokens(self.default_model, text) - + def get_headers(self): # headers with unique id return {"unique_id": self.unique_id} @@ -87,7 +87,7 @@ class ProxyServer(LLM): if "content" in loaded_chunk: completion += loaded_chunk["content"] except Exception as e: - capture_event(self.unique_id, "proxy_server_parse_error", { + posthog_logger.capture_event(self.unique_id, "proxy_server_parse_error", { "error_title": "Proxy server stream_chat parsing failed", "error_message": '\n'.join(traceback.format_exception(e))}) else: break diff --git a/continuedev/src/continuedev/libs/util/create_async_task.py b/continuedev/src/continuedev/libs/util/create_async_task.py index 354cea82..2473c638 100644 --- a/continuedev/src/continuedev/libs/util/create_async_task.py +++ b/continuedev/src/continuedev/libs/util/create_async_task.py @@ -1,6 +1,6 @@ from typing import Coroutine, Union import traceback -from .telemetry import capture_event +from .telemetry import posthog_logger import asyncio import nest_asyncio nest_asyncio.apply() @@ -16,7 +16,7 @@ def create_async_task(coro: Coroutine, unique_id: Union[str, None] = None): except Exception as e: print("Exception caught from async task: ", '\n'.join(traceback.format_exception(e))) - capture_event(unique_id or "None", "async_task_error", { + posthog_logger.capture_event("async_task_error", { "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e)) }) diff --git a/continuedev/src/continuedev/libs/util/paths.py b/continuedev/src/continuedev/libs/util/paths.py index fddef887..14a97f57 100644 --- a/continuedev/src/continuedev/libs/util/paths.py +++ b/continuedev/src/continuedev/libs/util/paths.py @@ -2,16 +2,45 @@ 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 + + +def getDefaultConfigFile() -> str: + current_path = os.path.dirname(os.path.realpath(__file__)) + config_path = os.path.join( + current_path, "..", "constants", "default_config.py.txt") + with open(config_path, 'r') as f: + return f.read() + + +def getConfigFilePath() -> str: + path = os.path.join(getGlobalFolderPath(), "config.py") + os.makedirs(os.path.dirname(path), exist_ok=True) + + if not os.path.exists(path): + with open(path, 'w') as f: + f.write(getDefaultConfigFile()) + + return path diff --git a/continuedev/src/continuedev/libs/util/telemetry.py b/continuedev/src/continuedev/libs/util/telemetry.py index 17735dce..a967828e 100644 --- a/continuedev/src/continuedev/libs/util/telemetry.py +++ b/continuedev/src/continuedev/libs/util/telemetry.py @@ -1,27 +1,37 @@ from typing import Any from posthog import Posthog -from ...core.config import load_config import os from dotenv import load_dotenv from .commonregex import clean_pii_from_any load_dotenv() in_codespaces = os.getenv("CODESPACES") == "true" +POSTHOG_API_KEY = 'phc_JS6XFROuNbhJtVCEdTSYk6gl5ArRrTNMpCcguAXlSPs' -# The personal API key is necessary only if you want to use local evaluation of feature flags. -posthog = Posthog('phc_JS6XFROuNbhJtVCEdTSYk6gl5ArRrTNMpCcguAXlSPs', - host='https://app.posthog.com') +class PostHogLogger: + def __init__(self, api_key: str): + self.api_key = api_key + self.unique_id = None + self.allow_anonymous_telemetry = True -def capture_event(unique_id: str, event_name: str, event_properties: Any): - # Return early if telemetry is disabled - config = load_config('.continue/config.json') - if not config.allow_anonymous_telemetry: - return + def setup(self, unique_id: str, allow_anonymous_telemetry: bool): + self.unique_id = unique_id + self.allow_anonymous_telemetry = allow_anonymous_telemetry - if in_codespaces: - event_properties['codespaces'] = True + # The personal API key is necessary only if you want to use local evaluation of feature flags. + self.posthog = Posthog(self.api_key, host='https://app.posthog.com') - # Send event to PostHog - posthog.capture(unique_id, event_name, - clean_pii_from_any(event_properties)) + def capture_event(self, event_name: str, event_properties: Any): + if not self.allow_anonymous_telemetry or self.unique_id is None: + return + + if in_codespaces: + event_properties['codespaces'] = True + + # Send event to PostHog + self.posthog.capture(self.unique_id, event_name, + clean_pii_from_any(event_properties)) + + +posthog_logger = PostHogLogger(api_key=POSTHOG_API_KEY) 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/plugins/context_providers/file.py b/continuedev/src/continuedev/plugins/context_providers/file.py new file mode 100644 index 00000000..6222ec6a --- /dev/null +++ b/continuedev/src/continuedev/plugins/context_providers/file.py @@ -0,0 +1,62 @@ +import os +import re +from typing import List +from ...core.main import ContextItem, ContextItemDescription, ContextItemId +from ...core.context import ContextProvider +from fnmatch import fnmatch + + +def get_file_contents(filepath: str) -> str: + try: + with open(filepath, "r") as f: + return f.read() + except UnicodeDecodeError: + return "" + + +class FileContextProvider(ContextProvider): + """ + The FileContextProvider is a ContextProvider that allows you to search files in the open workspace. + """ + + title = "file" + workspace_dir: str + ignore_patterns: List[str] = [ + ".git", + ".vscode", + ".idea", + ".vs", + ".venv", + "env", + ".env", + "node_modules", + "dist", + "build", + "target", + "out", + "bin", + ".pytest_cache", + ".vscode-test", + ".continue", + ] + + async def provide_context_items(self) -> List[ContextItem]: + filepaths = [] + for root, dir_names, file_names in os.walk(self.workspace_dir): + dir_names[:] = [d for d in dir_names if not any( + fnmatch(d, pattern) for pattern in self.ignore_patterns)] + for file_name in file_names: + filepaths.append(os.path.join(root, file_name)) + + return [ContextItem( + content=get_file_contents(file)[:min( + 2000, len(get_file_contents(file)))], + description=ContextItemDescription( + name=os.path.basename(file), + description=file, + id=ContextItemId( + provider_title=self.title, + item_id=re.sub(r'[^0-9a-zA-Z_-]', '', file) + ) + ) + ) for file in filepaths] diff --git a/continuedev/src/continuedev/plugins/context_providers/github.py b/continuedev/src/continuedev/plugins/context_providers/github.py new file mode 100644 index 00000000..765a534d --- /dev/null +++ b/continuedev/src/continuedev/plugins/context_providers/github.py @@ -0,0 +1,35 @@ +from typing import List +from github import Github +from github import Auth + +from ...core.context import ContextProvider, ContextItemDescription, ContextItem, ContextItemId + + +class GitHubIssuesContextProvider(ContextProvider): + """ + The GitHubIssuesContextProvider is a ContextProvider + that allows you to search GitHub issues in a repo. + """ + + title = "issues" + repo_name: str + auth_token: str + + async def provide_context_items(self) -> List[ContextItem]: + auth = Auth.Token(self.auth_token) + gh = Github(auth=auth) + + repo = gh.get_repo(self.repo_name) + issues = repo.get_issues().get_page(0) + + return [ContextItem( + content=issue.body, + description=ContextItemDescription( + name=f"Issue #{issue.number}", + description=issue.title, + id=ContextItemId( + provider_title=self.title, + item_id=issue.id + ) + ) + ) for issue in issues] diff --git a/continuedev/src/continuedev/plugins/context_providers/google.py b/continuedev/src/continuedev/plugins/context_providers/google.py new file mode 100644 index 00000000..64954833 --- /dev/null +++ b/continuedev/src/continuedev/plugins/context_providers/google.py @@ -0,0 +1,64 @@ +import json +from typing import List + +import aiohttp +from ...core.main import ContextItem, ContextItemDescription, ContextItemId +from ...core.context import ContextProvider + + +class GoogleContextProvider(ContextProvider): + title = "google" + + serper_api_key: str + + GOOGLE_CONTEXT_ITEM_ID = "google_search" + + @property + def BASE_CONTEXT_ITEM(self): + return ContextItem( + content="", + description=ContextItemDescription( + name="Google Search", + description="Enter a query to search google", + id=ContextItemId( + provider_title=self.title, + item_id=self.GOOGLE_CONTEXT_ITEM_ID + ) + ) + ) + + async def _google_search(self, query: str) -> str: + url = "https://google.serper.dev/search" + + payload = json.dumps({ + "q": query + }) + headers = { + 'X-API-KEY': self.serper_api_key, + 'Content-Type': 'application/json' + } + + async with aiohttp.ClientSession() as session: + async with session.post(url, headers=headers, data=payload) as response: + return await response.text() + + async def provide_context_items(self) -> List[ContextItem]: + return [self.BASE_CONTEXT_ITEM] + + async def get_item(self, id: ContextItemId, query: str, _) -> ContextItem: + if not id.item_id == self.GOOGLE_CONTEXT_ITEM_ID: + raise Exception("Invalid item id") + + results = await self._google_search(query) + json_results = json.loads(results) + content = f"Google Search: {query}\n\n" + if answerBox := json_results.get("answerBox"): + content += f"Answer Box ({answerBox['title']}): {answerBox['answer']}\n\n" + + for result in json_results["organic"]: + content += f"{result['title']}\n{result['link']}\n{result['snippet']}\n\n" + + ctx_item = self.BASE_CONTEXT_ITEM.copy() + ctx_item.content = content + ctx_item.description.id.item_id = query + return ctx_item diff --git a/continuedev/src/continuedev/plugins/context_providers/highlighted_code.py b/continuedev/src/continuedev/plugins/context_providers/highlighted_code.py new file mode 100644 index 00000000..426c0804 --- /dev/null +++ b/continuedev/src/continuedev/plugins/context_providers/highlighted_code.py @@ -0,0 +1,191 @@ +import os +from typing import Any, Dict, List + +from meilisearch_python_async import Client +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: Client, prev: List[ContextItem] = None) -> List[ContextItem]: + raise NotImplementedError() diff --git a/continuedev/src/continuedev/plugins/steps/custom_command.py b/continuedev/src/continuedev/plugins/steps/custom_command.py index d5b6e48b..419b3c3d 100644 --- a/continuedev/src/continuedev/plugins/steps/custom_command.py +++ b/continuedev/src/continuedev/plugins/steps/custom_command.py @@ -1,6 +1,6 @@ from ...libs.util.templating import render_templated_string from ...core.main import Step -from ...core.sdk import ContinueSDK +from ...core.sdk import ContinueSDK, Models from ..steps.chat import SimpleChatStep @@ -11,7 +11,7 @@ class CustomCommandStep(Step): slash_command: str hide: bool = True - async def describe(self): + async def describe(self, models: Models): return self.prompt async def run(self, sdk: ContinueSDK): diff --git a/continuedev/src/continuedev/plugins/steps/feedback.py b/continuedev/src/continuedev/plugins/steps/feedback.py index 119e3112..fa56a4d9 100644 --- a/continuedev/src/continuedev/plugins/steps/feedback.py +++ b/continuedev/src/continuedev/plugins/steps/feedback.py @@ -2,7 +2,7 @@ from typing import Coroutine from ...core.main import Models from ...core.main import Step from ...core.sdk import ContinueSDK -from ...libs.util.telemetry import capture_event +from ...libs.util.telemetry import posthog_logger class FeedbackStep(Step): @@ -13,5 +13,4 @@ class FeedbackStep(Step): return f"`{self.user_input}`\n\nWe'll see your feedback and make improvements as soon as possible. If you'd like to directly email us, you can contact [nate@continue.dev](mailto:nate@continue.dev?subject=Feedback%20On%20Continue)." async def run(self, sdk: ContinueSDK): - capture_event(sdk.ide.unique_id, "feedback", - {"feedback": self.user_input}) + posthog_logger.capture_event("feedback", {"feedback": self.user_input}) diff --git a/continuedev/src/continuedev/plugins/steps/help.py b/continuedev/src/continuedev/plugins/steps/help.py index 5111c7cf..d3807706 100644 --- a/continuedev/src/continuedev/plugins/steps/help.py +++ b/continuedev/src/continuedev/plugins/steps/help.py @@ -1,7 +1,7 @@ from textwrap import dedent from ...core.main import ChatMessage, Step from ...core.sdk import ContinueSDK -from ...libs.util.telemetry import capture_event +from ...libs.util.telemetry import posthog_logger help = dedent("""\ Continue is an open-source coding autopilot. It is a VS Code extension that brings the power of ChatGPT to your IDE. @@ -55,5 +55,5 @@ class HelpStep(Step): self.description += chunk["content"] await sdk.update_ui() - capture_event(sdk.ide.unique_id, "help", { - "question": question, "answer": self.description}) + posthog_logger.capture_event( + "help", {"question": question, "answer": self.description}) diff --git a/continuedev/src/continuedev/plugins/steps/main.py b/continuedev/src/continuedev/plugins/steps/main.py index 30117c55..a8752df2 100644 --- a/continuedev/src/continuedev/plugins/steps/main.py +++ b/continuedev/src/continuedev/plugins/steps/main.py @@ -15,20 +15,6 @@ from .core.core import DefaultModelEditCodeStep from ...libs.util.calculate_diff import calculate_diff2 -class SetupContinueWorkspaceStep(Step): - async def describe(self, models: Models) -> Coroutine[str, None, None]: - return "Set up Continue workspace by adding a .continue directory" - - async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]: - if not os.path.exists(os.path.join(await sdk.ide.getWorkspaceDirectory(), ".continue")): - await sdk.add_directory(".continue") - if not os.path.exists(os.path.join(await sdk.ide.getWorkspaceDirectory(), ".continue", "config.json")): - await sdk.add_file(".continue/config.json", dedent("""\ - { - "allow_anonymous_telemetry": true - }""")) - - class Policy(BaseModel): pass diff --git a/continuedev/src/continuedev/plugins/steps/open_config.py b/continuedev/src/continuedev/plugins/steps/open_config.py index d950c26f..64ead547 100644 --- a/continuedev/src/continuedev/plugins/steps/open_config.py +++ b/continuedev/src/continuedev/plugins/steps/open_config.py @@ -1,6 +1,7 @@ from textwrap import dedent from ...core.main import Step from ...core.sdk import ContinueSDK +from ...libs.util.paths import getConfigFilePath import os @@ -9,21 +10,20 @@ class OpenConfigStep(Step): async def describe(self, models): return dedent("""\ - `\"config.json\"` is now open. You can add a custom slash command in the `\"custom_commands\"` section, like in this example: - ```json - "custom_commands": [ - { - "name": "test", - "description": "Write unit tests like I do for the highlighted code", - "prompt": "Write a comprehensive set of unit tests for the selected code. It should setup, run tests that check for correctness including important edge cases, and teardown. Ensure that the tests are complete and sophisticated." - } - ] + `\"config.py\"` is now open. You can add a custom slash command in the `\"custom_commands\"` section, like in this example: + ```python + config = ContinueConfig( + ... + custom_commands=[CustomCommand( + name="test", + description="Write unit tests like I do for the highlighted code", + prompt="Write a comprehensive set of unit tests for the selected code. It should setup, run tests that check for correctness including important edge cases, and teardown. Ensure that the tests are complete and sophisticated.", + )] + ) ``` - `"name"` is the command you will type. - `"description"` is the description displayed in the slash command menu. - `"prompt"` is the instruction given to the model. The overall prompt becomes "Task: {prompt}, Additional info: {user_input}". For example, if you entered "/test exactly 5 assertions", the overall prompt would become "Task: Write a comprehensive...and sophisticated, Additional info: exactly 5 assertions".""") + `name` is the command you will type. + `description` is the description displayed in the slash command menu. + `prompt` is the instruction given to the model. The overall prompt becomes "Task: {prompt}, Additional info: {user_input}". For example, if you entered "/test exactly 5 assertions", the overall prompt would become "Task: Write a comprehensive...and sophisticated, Additional info: exactly 5 assertions".""") async def run(self, sdk: ContinueSDK): - global_dir = os.path.expanduser('~/.continue') - config_path = os.path.join(global_dir, 'config.json') - await sdk.ide.setFileOpen(config_path) + await sdk.ide.setFileOpen(getConfigFilePath()) diff --git a/continuedev/src/continuedev/plugins/steps/steps_on_startup.py b/continuedev/src/continuedev/plugins/steps/steps_on_startup.py index 19d62d30..489cada3 100644 --- a/continuedev/src/continuedev/plugins/steps/steps_on_startup.py +++ b/continuedev/src/continuedev/plugins/steps/steps_on_startup.py @@ -1,6 +1,5 @@ from ...core.main import Step from ...core.sdk import Models, ContinueSDK -from ...libs.util.step_name_to_steps import get_step_from_name class StepsOnStartupStep(Step): @@ -12,6 +11,6 @@ class StepsOnStartupStep(Step): async def run(self, sdk: ContinueSDK): steps_on_startup = sdk.config.steps_on_startup - for step_name, step_params in steps_on_startup.items(): - step = get_step_from_name(step_name, step_params) + for step_type in steps_on_startup: + step = step_type() await sdk.run_step(step) diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py index ae57c0b6..c0957395 100644 --- a/continuedev/src/continuedev/server/gui.py +++ b/continuedev/src/continuedev/server/gui.py @@ -2,15 +2,15 @@ import asyncio import json from fastapi import Depends, Header, WebSocket, APIRouter from starlette.websockets import WebSocketState, WebSocketDisconnect -from typing import Any, List, Type, TypeVar, Union +from typing import Any, List, Type, TypeVar from pydantic import BaseModel import traceback from uvicorn.main import Server -from .session_manager import SessionManager, session_manager, Session +from .session_manager import session_manager, Session from .gui_protocol import AbstractGUIProtocolServer from ..libs.util.queue import AsyncSubscriptionQueue -from ..libs.util.telemetry import capture_event +from ..libs.util.telemetry import posthog_logger from ..libs.util.create_async_task import create_async_task router = APIRouter(prefix="/gui", tags=["gui"]) @@ -61,12 +61,12 @@ class GUIProtocolServer(AbstractGUIProtocolServer): "data": data }) - async def _receive_json(self, message_type: str, timeout: int = 5) -> Any: + async def _receive_json(self, message_type: str, timeout: int = 20) -> Any: try: return await asyncio.wait_for(self.sub_queue.get(message_type), timeout=timeout) except asyncio.TimeoutError: raise Exception( - "GUI Protocol _receive_json timed out after 5 seconds") + "GUI Protocol _receive_json timed out after 20 seconds") async def _send_and_receive_json(self, data: Any, resp_model: Type[T], message_type: str) -> T: await self._send_json(message_type, data) @@ -85,31 +85,23 @@ class GUIProtocolServer(AbstractGUIProtocolServer): self.on_reverse_to_index(data["index"]) elif message_type == "retry_at_index": self.on_retry_at_index(data["index"]) - elif message_type == "change_default_model": - self.on_change_default_model(data["model"]) elif message_type == "clear_history": 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( @@ -132,10 +124,6 @@ class GUIProtocolServer(AbstractGUIProtocolServer): create_async_task( self.session.autopilot.retry_at_index(index), self.session.autopilot.continue_sdk.ide.unique_id) - def on_change_default_model(self, model: str): - create_async_task(self.session.autopilot.change_default_model( - model), self.session.autopilot.continue_sdk.ide.unique_id) - def on_clear_history(self): create_async_task(self.session.autopilot.clear_history( ), self.session.autopilot.continue_sdk.ide.unique_id) @@ -144,10 +132,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 +150,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,11 +175,11 @@ 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() - print("Received message", message) + print("Received GUI message", message) if type(message) is str: message = json.loads(message) @@ -206,13 +193,13 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we print("GUI websocket disconnected") except Exception as e: print("ERROR in gui websocket: ", e) - capture_event(session.autopilot.continue_sdk.ide.unique_id, "gui_error", { - "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))}) + posthog_logger.capture_event("gui_error", { + "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))}) raise e finally: print("Closing gui websocket") 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..990833be 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,21 +26,17 @@ 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""" @abstractmethod - def on_change_default_model(self): - """Called when the user requests to change the default model""" - - @abstractmethod def on_clear_history(self): """Called when the user requests to clear the history""" @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/ide.py b/continuedev/src/continuedev/server/ide.py index aeff5623..cf8b32a1 100644 --- a/continuedev/src/continuedev/server/ide.py +++ b/continuedev/src/continuedev/server/ide.py @@ -1,23 +1,25 @@ # This is a separate server from server/main.py -from functools import cached_property import json import os -from typing import Any, Dict, List, Type, TypeVar, Union +from typing import Any, List, Type, TypeVar, Union import uuid -from fastapi import WebSocket, Body, APIRouter +from fastapi import WebSocket, APIRouter from starlette.websockets import WebSocketState, WebSocketDisconnect from uvicorn.main import Server +from pydantic import BaseModel import traceback +import asyncio -from ..libs.util.telemetry import capture_event +from .meilisearch_server import start_meilisearch +from ..libs.util.telemetry import posthog_logger from ..libs.util.queue import AsyncSubscriptionQueue from ..models.filesystem import FileSystem, RangeInFile, EditDiff, RangeInFileWithContents, RealFileSystem from ..models.filesystem_edit import AddDirectory, AddFile, DeleteDirectory, DeleteFile, FileSystemEdit, FileEdit, FileEditWithFullContents, RenameDirectory, RenameFile, SequentialFileSystemEdit -from pydantic import BaseModel -from .gui import SessionManager, session_manager +from .gui import session_manager from .ide_protocol import AbstractIdeProtocolServer -import asyncio from ..libs.util.create_async_task import create_async_task +from .session_manager import SessionManager + import nest_asyncio nest_asyncio.apply() @@ -138,6 +140,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer): continue message_type = message["messageType"] data = message["data"] + print("Received message while initializing", message_type) if message_type == "workspaceDirectory": self.workspace_directory = data["workspaceDirectory"] elif message_type == "uniqueId": @@ -152,17 +155,18 @@ class IdeProtocolServer(AbstractIdeProtocolServer): async def _send_json(self, message_type: str, data: Any): if self.websocket.application_state == WebSocketState.DISCONNECTED: return + print("Sending IDE message: ", message_type) await self.websocket.send_json({ "messageType": message_type, "data": data }) - async def _receive_json(self, message_type: str, timeout: int = 5) -> Any: + async def _receive_json(self, message_type: str, timeout: int = 20) -> Any: try: return await asyncio.wait_for(self.sub_queue.get(message_type), timeout=timeout) except asyncio.TimeoutError: raise Exception( - "IDE Protocol _receive_json timed out after 5 seconds") + "IDE Protocol _receive_json timed out after 20 seconds", message_type) async def _send_and_receive_json(self, data: Any, resp_model: Type[T], message_type: str) -> T: await self._send_json(message_type, data) @@ -273,12 +277,12 @@ class IdeProtocolServer(AbstractIdeProtocolServer): # like file changes, tracebacks, etc... def onAcceptRejectSuggestion(self, accepted: bool): - capture_event(self.unique_id, "accept_reject_suggestion", { + posthog_logger.capture_event("accept_reject_suggestion", { "accepted": accepted }) def onAcceptRejectDiff(self, accepted: bool): - capture_event(self.unique_id, "accept_reject_diff", { + posthog_logger.capture_event("accept_reject_diff", { "accepted": accepted }) @@ -431,6 +435,13 @@ class IdeProtocolServer(AbstractIdeProtocolServer): @router.websocket("/ws") async def websocket_endpoint(websocket: WebSocket, session_id: str = None): try: + # Start meilisearch + try: + await start_meilisearch() + except Exception as e: + print("Failed to start MeiliSearch") + print(e) + await websocket.accept() print("Accepted websocket connection from, ", websocket.client) await websocket.send_json({"messageType": "connected", "data": {}}) @@ -443,6 +454,7 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None): message_type = message["messageType"] data = message["data"] + print("Received IDE message: ", message_type) create_async_task( ideProtocolServer.handle_json(message_type, data)) @@ -450,8 +462,8 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None): if session_id is not None: session_manager.registered_ides[session_id] = ideProtocolServer other_msgs = await ideProtocolServer.initialize(session_id) - capture_event(ideProtocolServer.unique_id, "session_started", { - "session_id": ideProtocolServer.session_id}) + posthog_logger.capture_event("session_started", { + "session_id": ideProtocolServer.session_id}) for other_msg in other_msgs: handle_msg(other_msg) @@ -465,13 +477,14 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None): print("IDE wbsocket disconnected") except Exception as e: print("Error in ide websocket: ", e) - capture_event(ideProtocolServer.unique_id, "gui_error", { - "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))}) + posthog_logger.capture_event("gui_error", { + "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))}) raise e finally: if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close() - capture_event(ideProtocolServer.unique_id, "session_ended", { - "session_id": ideProtocolServer.session_id}) - session_manager.registered_ides.pop(ideProtocolServer.session_id) + posthog_logger.capture_event("session_ended", { + "session_id": ideProtocolServer.session_id}) + if ideProtocolServer.session_id in session_manager.registered_ides: + session_manager.registered_ides.pop(ideProtocolServer.session_id) diff --git a/continuedev/src/continuedev/server/main.py b/continuedev/src/continuedev/server/main.py index 42dc0cc1..0b59d4fe 100644 --- a/continuedev/src/continuedev/server/main.py +++ b/continuedev/src/continuedev/server/main.py @@ -1,15 +1,17 @@ +import asyncio import time 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 + app = FastAPI() app.include_router(ide_router) @@ -41,15 +43,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(): @@ -79,5 +86,6 @@ if __name__ == "__main__": run_server() except Exception as e: + print("Error starting Continue server: ", e) cleanup() raise e diff --git a/continuedev/src/continuedev/server/meilisearch_server.py b/continuedev/src/continuedev/server/meilisearch_server.py new file mode 100644 index 00000000..286019e1 --- /dev/null +++ b/continuedev/src/continuedev/server/meilisearch_server.py @@ -0,0 +1,77 @@ +import os +import shutil +import subprocess + +from meilisearch_python_async import Client +from ..libs.util.paths import getServerFolderPath + + +def ensure_meilisearch_installed(): + """ + Checks if MeiliSearch is installed. + """ + serverPath = getServerFolderPath() + meilisearchPath = os.path.join(serverPath, "meilisearch") + dumpsPath = os.path.join(serverPath, "dumps") + dataMsPath = os.path.join(serverPath, "data.ms") + + paths = [meilisearchPath, dumpsPath, dataMsPath] + + existing_paths = set() + non_existing_paths = set() + for path in paths: + if os.path.exists(path): + existing_paths.add(path) + else: + non_existing_paths.add(path) + + if len(non_existing_paths) > 0: + # Clear the meilisearch binary + if meilisearchPath in existing_paths: + os.remove(meilisearchPath) + non_existing_paths.remove(meilisearchPath) + + # Clear the existing directories + for p in existing_paths: + shutil.rmtree(p, ignore_errors=True) + + # Download MeiliSearch + print("Downloading MeiliSearch...") + subprocess.run( + f"curl -L https://install.meilisearch.com | sh", shell=True, check=True, cwd=serverPath) + + +async def check_meilisearch_running() -> bool: + """ + Checks if MeiliSearch is running. + """ + + try: + client = Client('http://localhost:7700') + resp = await client.health() + if resp["status"] != "available": + return False + return True + except Exception: + return False + + +async 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 download + ensure_meilisearch_installed() + + # Check if MeiliSearch is running + if not await 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 20219273..3136f1bf 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: |