inject data from weaviate in context

Author: yrobla

scripts/import_packages.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 
+from src.codegate.utils.utils import generate_vector_string
 import weaviate
 from weaviate.classes.config import DataType, Property
 from weaviate.embedded import EmbeddedOptions
@@ -37,33 +38,8 @@ def setup_schema(self):
                 ],
             )
 
-    def generate_vector_string(self, package):
-        vector_str = f"{package['name']}"
-        package_url = ""
-        type_map = {
-            "pypi": "Python package available on PyPI",
-            "npm": "JavaScript package available on NPM",
-            "go": "Go package",
-            "crates": "Rust package available on Crates",
-            "java": "Java package",
-        }
-        status_messages = {
-            "archived": "However, this package is found to be archived and no longer maintained.",
-            "deprecated": "However, this package is found to be deprecated and no longer "
-            "recommended for use.",
-            "malicious": "However, this package is found to be malicious.",
-        }
-        vector_str += f" is a {type_map.get(package['type'], 'unknown type')} "
-        package_url = f"https://trustypkg.dev/{package['type']}/{package['name']}"
-
-        # Add extra status
-        status_suffix = status_messages.get(package["status"], "")
-        if status_suffix:
-            vector_str += f"{status_suffix} For additional information refer to {package_url}"
-        return vector_str
-
     async def process_package(self, batch, package):
-        vector_str = self.generate_vector_string(package)
+        vector_str = generate_vector_string(package)
         vector = await self.inference_engine.embed(self.model_path, [vector_str])
         # This is where the synchronous call is made
         batch.add_object(properties=package, vector=vector[0])

src/codegate/pipeline/codegate_context_retriever/__init__.py

@@ -0,0 +1,3 @@
+from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
+
+__all__ = ["CodegateContextRetriever"]

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -0,0 +1,80 @@
+from typing import Optional
+
+from litellm import ChatCompletionRequest, ChatCompletionSystemMessage
+
+from codegate.pipeline.base import (
+    PipelineContext,
+    PipelineResult,
+    PipelineStep,
+)
+from src.codegate.storage.storage_engine import StorageEngine
+from src.codegate.utils.utils import generate_vector_string
+
+
+class CodegateContextRetriever(PipelineStep):
+    """
+    Pipeline step that adds a context message to the completion request when it detects
+    the word "codegate" in the user message.
+    """
+
+    def __init__(self, system_prompt_message: Optional[str] = None):
+        self._system_message = ChatCompletionSystemMessage(
+            content=system_prompt_message,
+            role="system"
+        )
+        self.storage_engine = StorageEngine()
+
+    @property
+    def name(self) -> str:
+        """
+        Returns the name of this pipeline step.
+        """
+        return "codegate-context-retriever"
+
+    async def get_objects_from_search(self, search: str) -> list[object]:
+        objects = await self.storage_engine.search(search)
+        return objects
+
+    def generate_context_str(self, objects: list[object]) -> str:
+        context_str = "Please use the information about related packages to influence your answer:\n"
+        for obj in objects:
+            # generate dictionary from object
+            package_obj = {
+                "name": obj.properties["name"],
+                "type": obj.properties["type"],
+                "status": obj.properties["status"],
+                "description": obj.properties["description"],
+            }
+            package_str = generate_vector_string(package_obj)
+            context_str += package_str + "\n"
+        return context_str
+
+    async def process(
+            self, request: ChatCompletionRequest, context: PipelineContext
+    ) -> PipelineResult:
+        """
+        Process the completion request and add a system prompt if the user message contains
+        the word "codegate".
+        """
+        # no prompt configured
+        if not self._system_message["content"]:
+            return PipelineResult(request=request)
+
+        last_user_message = self.get_last_user_message(request)
+
+        if last_user_message is not None:
+            last_user_message_str, last_user_idx = last_user_message
+            if "codegate" in last_user_message_str.lower():
+                # strip codegate from prompt and trim it
+                last_user_message_str = last_user_message_str.lower().replace("codegate", "").strip()
+                searched_objects = await self.get_objects_from_search(last_user_message_str)
+                context_str = self.generate_context_str(searched_objects)
+                # Add a system prompt to the completion request
+                new_request = request.copy()
+                new_request["messages"].insert(last_user_idx, context_str)
+                return PipelineResult(
+                    request=new_request,
+                )
+
+        # Fall through
+        return PipelineResult(request=request)

src/codegate/server.py

@@ -6,6 +6,7 @@
 from codegate.config import Config
 from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
 from codegate.pipeline.codegate_system_prompt.codegate import CodegateSystemPrompt
+from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
 from codegate.pipeline.version.version import CodegateVersion
 from codegate.providers.anthropic.provider import AnthropicProvider
 from codegate.providers.llamacpp.provider import LlamaCppProvider
@@ -24,6 +25,7 @@ def init_app() -> FastAPI:
     steps: List[PipelineStep] = [
         CodegateVersion(),
         CodegateSystemPrompt(Config.get_config().prompts.codegate_chat),
+        CodegateContextRetriever(Config.get_config().prompts.codegate_chat),
         # CodegateSecrets(),
     ]
     # Leaving the pipeline empty for now

src/codegate/storage/storage_engine.py

@@ -1,6 +1,6 @@
 from codegate.codegate_logging import setup_logging
 from codegate.inference.inference_engine import LlamaCppInferenceEngine
-from weaviate.classes.config import DataType, Property
+from weaviate.classes.config import DataType
 from weaviate.classes.query import MetadataQuery
 import weaviate
 
@@ -45,7 +45,7 @@ def setup_schema(self):
                 self.client.collections.create(class_config['name'], properties=class_config['properties'])
                 self.__logger.info(f"Weaviate schema for class {class_config['name']} setup complete.")
 
-    async def search(self, query, limit=5, distance=0.1):
+    async def search(self, query: str, limit=5, distance=0.3) -> list[object]:
         """
         Search the 'Package' collection based on a query string.
 
@@ -62,7 +62,7 @@ async def search(self, query, limit=5, distance=0.1):
         # Perform the vector search
         try:
             collection = self.client.collections.get("Package")
-            response = collection.query.near_vector(query_vector, limit=limit, distance=distance, return_metadata=MetadataQuery(distance=True))
+            response = collection.query.near_vector(query_vector[0], limit=limit, distance=distance, return_metadata=MetadataQuery(distance=True))
             if not response:
                 return []
             return response.objects

src/codegate/utils/utils.py

@@ -0,0 +1,27 @@
+def generate_vector_string(package) -> str:
+    vector_str = f"{package['name']}"
+    package_url = ""
+    type_map = {
+        "pypi": "Python package available on PyPI",
+        "npm": "JavaScript package available on NPM",
+        "go": "Go package",
+        "crates": "Rust package available on Crates",
+        "java": "Java package",
+    }
+    status_messages = {
+        "archived": "However, this package is found to be archived and no longer maintained.",
+        "deprecated": "However, this package is found to be deprecated and no longer "
+        "recommended for use.",
+        "malicious": "However, this package is found to be malicious.",
+    }
+    vector_str += f" is a {type_map.get(package['type'], 'unknown type')} "
+    package_url = f"https://trustypkg.dev/{package['type']}/{package['name']}"
+
+    # Add extra status
+    status_suffix = status_messages.get(package["status"], "")
+    if status_suffix:
+        vector_str += f"{status_suffix} For additional information refer to {package_url}"
+
+    # add description
+    vector_str += f" - Package offers this functionality: {package['description']}"
+    return vector_str