diff options
Diffstat (limited to 'continuedev')
| -rw-r--r-- | continuedev/src/continuedev/core/context.py | 46 | ||||
| -rw-r--r-- | continuedev/src/continuedev/plugins/context_providers/filetree.py | 66 | ||||
| -rw-r--r-- | continuedev/src/continuedev/server/ide.py | 2 |
3 files changed, 80 insertions, 34 deletions
diff --git a/continuedev/src/continuedev/core/context.py b/continuedev/src/continuedev/core/context.py index 125ddc23..571e5dc8 100644 --- a/continuedev/src/continuedev/core/context.py +++ b/continuedev/src/continuedev/core/context.py @@ -285,8 +285,6 @@ class ContextManager: await globalSearchIndex.update_filterable_attributes(["workspace_dir"]) async def load_context_provider(provider: ContextProvider): - ti = time.time() - context_items = await provider.provide_context_items(workspace_dir) documents = [ { @@ -299,28 +297,34 @@ class ContextManager: for item in context_items ] if len(documents) > 0: - try: - await asyncio.wait_for( - globalSearchIndex.add_documents(documents), timeout=20 - ) - except asyncio.TimeoutError: - logger.warning( - f"Failed to add documents to meilisearch for context provider {provider.__class__.__name__} in 20 seconds" - ) - return - except Exception as e: - logger.warning( - f"Error adding documents to meilisearch for context provider {provider.__class__.__name__}: {e}" - ) - return + await globalSearchIndex.add_documents(documents) + + return len(documents) + + async def safe_load(provider: ContextProvider): + ti = time.time() + try: + num_documents = await asyncio.wait_for( + load_context_provider(provider), timeout=20 + ) + except asyncio.TimeoutError: + logger.warning( + f"Failed to add documents to meilisearch for context provider {provider.__class__.__name__} in 20 seconds" + ) + return + except Exception as e: + logger.warning( + f"Error adding documents to meilisearch for context provider {provider.__class__.__name__}: {e}" + ) + return tf = time.time() logger.debug( - f"Loaded {len(documents)} documents into meilisearch in {tf - ti} seconds for context provider {provider.title}" + f"Loaded {num_documents} documents into meilisearch in {tf - ti} seconds for context provider {provider.title}" ) tasks = [ - load_context_provider(provider) + safe_load(provider) for _, provider in self.context_providers.items() ] await asyncio.wait_for(asyncio.gather(*tasks), timeout=20) @@ -330,9 +334,11 @@ class ContextManager: if should_retry: await restart_meilisearch() try: - asyncio.wait_for(await poll_meilisearch_running(), timeout=20) + await asyncio.wait_for(poll_meilisearch_running(), timeout=20) except asyncio.TimeoutError: - logger.warning("Meilisearch did not restart in less than 20 seconds. Stopping polling.") + logger.warning( + "Meilisearch did not restart in less than 20 seconds. Stopping polling." + ) await self.load_index(workspace_dir, False) async def select_context_item(self, id: str, query: str): diff --git a/continuedev/src/continuedev/plugins/context_providers/filetree.py b/continuedev/src/continuedev/plugins/context_providers/filetree.py index ea86f214..959a0a66 100644 --- a/continuedev/src/continuedev/plugins/context_providers/filetree.py +++ b/continuedev/src/continuedev/plugins/context_providers/filetree.py @@ -1,31 +1,71 @@ -import os from typing import List +from pydantic import BaseModel + from ...core.context import ContextProvider from ...core.main import ContextItem, ContextItemDescription, ContextItemId -def format_file_tree(startpath) -> str: +class Directory(BaseModel): + name: str + files: List[str] + directories: List["Directory"] + + +def format_file_tree(tree: Directory, indentation: str = "") -> str: result = "" - for root, dirs, files in os.walk(startpath): - level = root.replace(startpath, "").count(os.sep) - indent = " " * 4 * (level) - result += "{}{}/".format(indent, os.path.basename(root)) + "\n" - subindent = " " * 4 * (level + 1) - for f in files: - result += "{}{}".format(subindent, f) + "\n" + for file in tree.files: + result += f"{indentation}{file}\n" + + for directory in tree.directories: + result += f"{indentation}{directory.name}/\n" + result += format_file_tree(directory, indentation + " ") return result +def split_path(path: str, with_root=None) -> List[str]: + parts = path.split("/") if "/" in path else path.split("\\") + if with_root is not None: + root_parts = split_path(with_root) + parts = parts[len(root_parts) - 1 :] + + return parts + + class FileTreeContextProvider(ContextProvider): title = "tree" workspace_dir: str = None - def _filetree_context_item(self): + async def _get_file_tree(self, directory: str) -> str: + contents = await self.sdk.ide.listDirectoryContents(directory, recursive=True) + + tree = Directory( + name=split_path(self.workspace_dir)[-1], files=[], directories=[] + ) + + for file in contents: + parts = split_path(file, with_root=self.workspace_dir) + + current_tree = tree + for part in parts[:-1]: + if part not in [d.name for d in current_tree.directories]: + current_tree.directories.append( + Directory(name=part, files=[], directories=[]) + ) + + current_tree = [d for d in current_tree.directories if d.name == part][ + 0 + ] + + current_tree.files.append(parts[-1]) + + return format_file_tree(tree) + + async def _filetree_context_item(self): return ContextItem( - content=format_file_tree(self.workspace_dir), + content=await self._get_file_tree(self.workspace_dir), description=ContextItemDescription( name="File Tree", description="Add a formatted file tree of this directory to the context", @@ -35,10 +75,10 @@ class FileTreeContextProvider(ContextProvider): async def provide_context_items(self, workspace_dir: str) -> List[ContextItem]: self.workspace_dir = workspace_dir - return [self._filetree_context_item()] + return [await self._filetree_context_item()] async def get_item(self, id: ContextItemId, query: str) -> ContextItem: if not id.item_id == self.title: raise Exception("Invalid item id") - return self._filetree_context_item() + return await self._filetree_context_item() diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py index 5d1d897e..bdb5817d 100644 --- a/continuedev/src/continuedev/server/ide.py +++ b/continuedev/src/continuedev/server/ide.py @@ -196,7 +196,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer): f"Tried to send message, but websocket is disconnected: {message_type}" ) return - logger.debug(f"Sending IDE message: {message_type}") + # logger.debug(f"Sending IDE message: {message_type}") await self.websocket.send_json({"messageType": message_type, "data": data}) except RuntimeError as e: logger.warning(f"Error sending IDE message, websocket probably closed: {e}") |