Merge pull request #272 from stacklok/issue-257

feat: infer ecosystem from the user's query

Author: yrobla

prompts/default.yaml

@@ -31,6 +31,15 @@ lookup_packages: |
   Assume that a package can be any named entity.
   You MUST RESPOND with a list of packages in JSON FORMAT: {"packages": ["pkg1", "pkg2", ...]}.
 
+lookup_ecosystem: |
+  You are a software expert with knowledge of various programming languages ecosystems.
+  When given a query related to coding or programming tasks, your job is to determine
+  the associated programming language and then infer the corresponding language ecosystem
+  based on the context provided in the query.
+  Valid ecosystems are: pypi (Python), npm (Node.js), maven (Java), crates (Rust), go (golang).
+  If you are not sure or you cannot infer it, please respond with an empty value.
+  You MUST RESPOND with a JSON dictionary on this format: {"ecosystem": "ecosystem_name"}.
+
 secrets_redacted: |
   The files in the context contain sensitive information that has been redacted. Do not warn the user
   about any tokens, passwords or similar sensitive information in the context whose value begins with

src/codegate/llm_utils/extractor.py

@@ -41,3 +41,29 @@ async def extract_packages(
         packages = result if isinstance(result, list) else result.get("packages", [])
         logger.info(f"Extracted packages: {packages}")
         return packages
+
+    @staticmethod
+    async def extract_ecosystem(
+        content: str,
+        provider: str,
+        model: str = None,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ) -> List[str]:
+        """Extract ecosystem from the given content."""
+        system_prompt = Config.get_config().prompts.lookup_ecosystem
+
+        result = await LLMClient.complete(
+            content=content,
+            system_prompt=system_prompt,
+            provider=provider,
+            model=model,
+            api_key=api_key,
+            base_url=base_url,
+        )
+
+        ecosystem = result if isinstance(result, str) else result.get("ecosystem")
+        if ecosystem:
+            ecosystem = ecosystem.lower()
+        logger.info(f"Extracted ecosystem: {ecosystem}")
+        return ecosystem

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -30,10 +30,11 @@ def name(self) -> str:
         return "codegate-context-retriever"
 
     async def get_objects_from_search(
-        self, search: str, packages: list[str] = None
+        self, search: str, ecosystem, packages: list[str] = None
     ) -> list[object]:
         storage_engine = StorageEngine()
-        objects = await storage_engine.search(search, distance=0.8, packages=packages)
+        objects = await storage_engine.search(
+            search, distance=0.8, ecosystem=ecosystem, packages=packages)
         return objects
 
     def generate_context_str(self, objects: list[object], context: PipelineContext) -> str:
@@ -69,6 +70,19 @@ async def __lookup_packages(self, user_query: str, context: PipelineContext):
         logger.info(f"Packages in user query: {packages}")
         return packages
 
+    async def __lookup_ecosystem(self, user_query: str, context: PipelineContext):
+        # Use PackageExtractor to extract ecosystem from the user query
+        ecosystem = await PackageExtractor.extract_ecosystem(
+            content=user_query,
+            provider=context.sensitive.provider,
+            model=context.sensitive.model,
+            api_key=context.sensitive.api_key,
+            base_url=context.sensitive.api_base,
+        )
+
+        logger.info(f"Ecosystem in user query: {ecosystem}")
+        return ecosystem
+
     async def process(
         self, request: ChatCompletionRequest, context: PipelineContext
     ) -> PipelineResult:
@@ -85,6 +99,7 @@ async def process(
 
         # Extract packages from the user message
         last_user_message_str, last_user_idx = last_user_message
+        ecosystem = await self.__lookup_ecosystem(last_user_message_str, context)
         packages = await self.__lookup_packages(last_user_message_str, context)
         packages = [pkg.lower() for pkg in packages]
 
@@ -93,7 +108,8 @@ async def process(
             return PipelineResult(request=request)
 
         # Look for matches in vector DB using list of packages as filter
-        searched_objects = await self.get_objects_from_search(last_user_message_str, packages)
+        searched_objects = await self.get_objects_from_search(
+            last_user_message_str, ecosystem, packages)
 
         logger.info(
             f"Found {len(searched_objects)} matches in the database",

src/codegate/storage/storage_engine.py

@@ -11,6 +11,7 @@
 from codegate.inference.inference_engine import LlamaCppInferenceEngine
 
 logger = structlog.get_logger("codegate")
+VALID_ECOSYSTEMS = ["npm", "pypi", "crates", "maven", "go"]
 
 schema_config = [
     {
@@ -120,6 +121,7 @@ async def search_by_property(self, name: str, properties: List[str]) -> list[obj
             return []
 
         try:
+
             packages = self.weaviate_client.collections.get("Package")
             response = packages.query.fetch_objects(
                 filters=Filter.by_property(name).contains_any(properties),
@@ -142,13 +144,17 @@ async def search_by_property(self, name: str, properties: List[str]) -> list[obj
             logger.error(f"An error occurred: {str(e)}")
             return []
 
-    async def search(self, query: str, limit=5, distance=0.3, packages=None) -> list[object]:
+    async def search(self, query: str, limit=5, distance=0.3,
+                     ecosystem=None, packages=None) -> list[object]:
         """
         Search the 'Package' collection based on a query string.
 
         Args:
             query (str): The text query for which to search.
             limit (int): The number of results to return.
+            distance (float): The distance threshold for the search.
+            ecosystem (str): The ecosystem to search in.
+            packages (list): The list of packages to filter the search.
 
         Returns:
             list: A list of matching results with their properties and distances.
@@ -160,11 +166,16 @@ async def search(self, query: str, limit=5, distance=0.3, packages=None) -> list
         try:
             collection = self.weaviate_client.collections.get("Package")
             if packages:
+                # filter by packages and ecosystem if present
+                filters = []
+                if ecosystem and ecosystem in VALID_ECOSYSTEMS:
+                    filters.append(wvc.query.Filter.by_property("type").equal(ecosystem))
+                filters.append(wvc.query.Filter.by_property("name").contains_any(packages))
                 response = collection.query.near_vector(
                     query_vector[0],
                     limit=limit,
                     distance=distance,
-                    filters=wvc.query.Filter.by_property("name").contains_any(packages),
+                    filters=wvc.query.Filter.all_of(filters),
                     return_metadata=MetadataQuery(distance=True),
                 )
             else:

tests/test_cli.py

@@ -74,7 +74,7 @@ def test_serve_default_options(
             "port": 8989,
             "log_level": "INFO",
             "log_format": "JSON",
-            "prompts_loaded": 6,
+            "prompts_loaded": 7,
             "provider_urls": DEFAULT_PROVIDER_URLS,
         }
 
@@ -123,7 +123,7 @@ def test_serve_custom_options(
             "port": 8989,
             "log_level": "DEBUG",
             "log_format": "TEXT",
-            "prompts_loaded": 6,  # Default prompts are loaded
+            "prompts_loaded": 7,  # Default prompts are loaded
             "provider_urls": DEFAULT_PROVIDER_URLS,
         }
 
@@ -170,7 +170,7 @@ def test_serve_with_config_file(
             "port": 8989,
             "log_level": "DEBUG",
             "log_format": "JSON",
-            "prompts_loaded": 6,  # Default prompts are loaded
+            "prompts_loaded": 7,  # Default prompts are loaded
             "provider_urls": DEFAULT_PROVIDER_URLS,
         }
 
@@ -229,7 +229,7 @@ def test_serve_priority_resolution(
             "port": 8080,
             "log_level": "ERROR",
             "log_format": "TEXT",
-            "prompts_loaded": 6,  # Default prompts are loaded
+            "prompts_loaded": 7,  # Default prompts are loaded
             "provider_urls": DEFAULT_PROVIDER_URLS,
         }