Merge pull request #360 from bra1nDump/fixup-search-cross-workspace-polution

Fixup search cross workspace polution

4 files changed, 58 insertions, 34 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 256f3439..71971ddd 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -81,11 +81,7 @@ class Autopilot(ContinueBaseModel):
             self.continue_sdk.config.context_providers + [
                 HighlightedCodeContextProvider(ide=self.ide),
                 FileContextProvider(workspace_dir=self.ide.workspace_directory)
-            ])
-
-        logger.debug("Loading index")
-        create_async_task(self.context_manager.load_index(
-            self.ide.workspace_directory))
+            ], self.ide.workspace_directory)
 
         if full_state is not None:
             self.history = full_state.history
diff --git a/continuedev/src/continuedev/core/context.py b/continuedev/src/continuedev/core/context.py
index b1f68b50..db1c770a 100644
--- a/continuedev/src/continuedev/core/context.py
+++ b/continuedev/src/continuedev/core/context.py
@@ -8,9 +8,10 @@ from pydantic import BaseModel
 
 
 from .main import ChatMessage, ContextItem, ContextItemDescription, ContextItemId
-from ..server.meilisearch_server import check_meilisearch_running
+from ..server.meilisearch_server import poll_meilisearch_running
 from ..libs.util.logging import logger
 from ..libs.util.telemetry import posthog_logger
+from ..libs.util.create_async_task import create_async_task
 
 SEARCH_INDEX_NAME = "continue_context_items"
 
@@ -140,31 +141,32 @@ class ContextManager:
         self.context_providers = {}
         self.provider_titles = set()
 
-    async def start(self, context_providers: List[ContextProvider]):
+    async def start(self, context_providers: List[ContextProvider], workspace_directory: str):
         """
         Starts the context manager.
         """
+        # Use only non-meilisearch-dependent providers until it is loaded
         self.context_providers = {
-            prov.title: prov for prov in context_providers}
+            title: provider for title, provider in self.context_providers.items() if title == "code"
+        }
         self.provider_titles = {
             provider.title for provider in context_providers}
 
-        async with Client('http://localhost:7700') as search_client:
-            meilisearch_running = True
+        # Start MeiliSearch in the background without blocking
+        async def start_meilisearch(context_providers):
             try:
-
-                health = await search_client.health()
-                if not health.status == "available":
-                    meilisearch_running = False
-            except:
-                meilisearch_running = False
-
-            if not meilisearch_running:
+                await asyncio.wait_for(poll_meilisearch_running(), timeout=20)
+                self.context_providers = {
+                    prov.title: prov for prov in context_providers}
+                logger.debug("Loading Meilisearch index...")
+                await self.load_index(workspace_directory)
+                logger.debug("Loaded Meilisearch index")
+            except asyncio.TimeoutError:
+                logger.warning("MeiliSearch did not start within 5 seconds")
                 logger.warning(
                     "MeiliSearch not running, avoiding any dependent context providers")
-                self.context_providers = {
-                    title: provider for title, provider in self.context_providers.items() if title == "code"
-                }
+
+        create_async_task(start_meilisearch(context_providers))
 
     async def load_index(self, workspace_dir: str):
         for _, provider in self.context_providers.items():
@@ -176,14 +178,24 @@ class ContextManager:
                     "id": item.description.id.to_string(),
                     "name": item.description.name,
                     "description": item.description.description,
-                    "content": item.content
+                    "content": item.content,
+                    "workspace_dir": workspace_dir,
                 }
                 for item in context_items
             ]
             if len(documents) > 0:
                 try:
                     async with Client('http://localhost:7700') as search_client:
-                        await asyncio.wait_for(search_client.index(SEARCH_INDEX_NAME).add_documents(documents), timeout=5)
+                        # First, create the index if it doesn't exist
+                        await search_client.create_index(SEARCH_INDEX_NAME)
+                        # The index is currently shared by all workspaces
+                        globalSearchIndex = await search_client.get_index(SEARCH_INDEX_NAME)
+                        await asyncio.wait_for(asyncio.gather(
+                            # Ensure that the index has the correct filterable attributes
+                            globalSearchIndex.update_filterable_attributes(
+                                ["workspace_dir"]),
+                            globalSearchIndex.add_documents(documents)
+                        ), timeout=5)
                 except Exception as e:
                     logger.debug(f"Error loading meilisearch index: {e}")
 
diff --git a/continuedev/src/continuedev/plugins/context_providers/file.py b/continuedev/src/continuedev/plugins/context_providers/file.py
index 31aa5423..b40092af 100644
--- a/continuedev/src/continuedev/plugins/context_providers/file.py
+++ b/continuedev/src/continuedev/plugins/context_providers/file.py
@@ -54,33 +54,37 @@ class FileContextProvider(ContextProvider):
         list(filter(lambda d: f"**/{d}", DEFAULT_IGNORE_DIRS))
 
     async def provide_context_items(self, workspace_dir: str) -> List[ContextItem]:
-        filepaths = []
+        absolute_filepaths: List[str] = []
         for root, dir_names, file_names in os.walk(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))
+                absolute_filepaths.append(os.path.join(root, file_name))
 
-                if len(filepaths) > 1000:
+                if len(absolute_filepaths) > 1000:
                     break
 
-            if len(filepaths) > 1000:
+            if len(absolute_filepaths) > 1000:
                 break
 
         items = []
-        for file in filepaths:
-            content = get_file_contents(file)
+        for absolute_filepath in absolute_filepaths:
+            content = get_file_contents(absolute_filepath)
             if content is None:
                 continue  # no pun intended
+            
+            relative_to_workspace = os.path.relpath(absolute_filepath, workspace_dir)
 
             items.append(ContextItem(
                 content=content[:min(2000, len(content))],
                 description=ContextItemDescription(
-                    name=os.path.basename(file),
-                    description=file,
+                    name=os.path.basename(absolute_filepath),
+                    # We should add the full path to the ContextItem
+                    # It warrants a data modeling discussion and has no immediate use case
+                    description=relative_to_workspace,
                     id=ContextItemId(
                         provider_title=self.title,
-                        item_id=remove_meilisearch_disallowed_chars(file)
+                        item_id=remove_meilisearch_disallowed_chars(absolute_filepath)
                     )
                 )
             ))
diff --git a/continuedev/src/continuedev/server/meilisearch_server.py b/continuedev/src/continuedev/server/meilisearch_server.py
index 7f460afc..f47c08ca 100644
--- a/continuedev/src/continuedev/server/meilisearch_server.py
+++ b/continuedev/src/continuedev/server/meilisearch_server.py
@@ -1,3 +1,4 @@
+import asyncio
 import os
 import shutil
 import subprocess
@@ -58,15 +59,26 @@ async def check_meilisearch_running() -> bool:
         async with Client('http://localhost:7700') as client:
             try:
                 resp = await client.health()
-                if resp["status"] != "available":
+                if resp.status != "available":
                     return False
                 return True
-            except:
+            except Exception as e:
+                logger.debug(e)
                 return False
     except Exception:
         return False
 
 
+async def poll_meilisearch_running(frequency: int = 0.1) -> bool:
+    """
+    Polls MeiliSearch to see if it is running.
+    """
+    while True:
+        if await check_meilisearch_running():
+            return True
+        await asyncio.sleep(frequency)
+
+
 async def start_meilisearch():
     """
     Starts the MeiliSearch server, wait for it.