Load RAG database in memory on use.

codegate-profiling

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+
+import os
+import pstats
+import sys
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Error: file name required")
+        exit(1)
+    ps = pstats.Stats(*sys.argv[1:]).sort_stats("cumulative")
+    ps.print_stats(.03)

src/codegate/profiling/__init__.py

@@ -0,0 +1,28 @@
+import cProfile
+
+
+def profiled(prefix):
+    def wrapper(func):
+        import datetime
+        import os
+        from functools import wraps
+
+        @wraps(func)
+        async def inner(*args, **kwargs):
+
+            ts = datetime.datetime.utcnow()
+            fname = f"pstats-{prefix}-{ts.strftime('%Y%m%dT%H%M%S%f')}.prof"
+            fname = os.path.join(os.path.abspath("."), fname)
+
+            with cProfile.Profile() as pr:
+                res = await func(*args, **kwargs)
+                print(f"writing to {fname}")
+                pr.dump_stats(fname)
+                return res
+
+        if os.getenv(f"CODEGATE_PROFILE_{prefix.upper()}"):
+            return inner
+
+        return func
+
+    return wrapper

src/codegate/providers/anthropic/provider.py

@@ -8,6 +8,7 @@
 from codegate.clients.clients import ClientType
 from codegate.clients.detector import DetectClient
 from codegate.pipeline.factory import PipelineFactory
+from codegate.profiling import profiled
 from codegate.providers.anthropic.completion_handler import AnthropicCompletion
 from codegate.providers.base import BaseProvider, ModelFetchError
 from codegate.providers.fim_analyzer import FIMAnalyzer
@@ -68,6 +69,7 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
 
         return [model["id"] for model in respjson.get("data", [])]
 
+    @profiled("anthropic")
     async def process_request(
         self,
         data: dict,

src/codegate/providers/ollama/provider.py

@@ -9,6 +9,7 @@
 from codegate.clients.detector import DetectClient
 from codegate.config import Config
 from codegate.pipeline.factory import PipelineFactory
+from codegate.profiling import profiled
 from codegate.providers.base import BaseProvider, ModelFetchError
 from codegate.providers.fim_analyzer import FIMAnalyzer
 from codegate.providers.ollama.completion_handler import OllamaShim
@@ -59,6 +60,7 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
 
         return [model["name"] for model in jsonresp.get("models", [])]
 
+    @profiled("ollama")
     async def process_request(
         self,
         data: dict,

src/codegate/providers/openai/provider.py

@@ -7,6 +7,7 @@
 from codegate.clients.clients import ClientType
 from codegate.clients.detector import DetectClient
 from codegate.pipeline.factory import PipelineFactory
+from codegate.profiling import profiled
 from codegate.providers.base import BaseProvider, ModelFetchError
 from codegate.providers.completion import BaseCompletionHandler
 from codegate.providers.fim_analyzer import FIMAnalyzer
@@ -63,6 +64,7 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
 
         return [model["id"] for model in jsonresp.get("data", [])]
 
+    @profiled("openai")
     async def process_request(
         self,
         data: dict,

src/codegate/providers/openrouter/provider.py

@@ -5,6 +5,7 @@
 from codegate.clients.clients import ClientType
 from codegate.clients.detector import DetectClient
 from codegate.pipeline.factory import PipelineFactory
+from codegate.profiling import profiled
 from codegate.providers.fim_analyzer import FIMAnalyzer
 from codegate.providers.litellmshim import LiteLLmShim
 from codegate.providers.openai import OpenAIProvider
@@ -49,6 +50,7 @@ def __init__(self, pipeline_factory: PipelineFactory):
     def provider_route_name(self) -> str:
         return "openrouter"
 
+    @profiled("openrouter")
     async def process_request(
         self,
         data: dict,

src/codegate/storage/storage_engine.py

@@ -9,6 +9,7 @@
 
 from codegate.config import Config
 from codegate.inference.inference_engine import LlamaCppInferenceEngine
+from codegate.profiling import profiled
 
 logger = structlog.get_logger("codegate")
 VALID_ECOSYSTEMS = ["npm", "pypi", "crates", "maven", "go"]
@@ -78,7 +79,14 @@ def _get_connection(self):
             conn.enable_load_extension(True)
             sqlite_vec_sl_tmp.load(conn)
             conn.enable_load_extension(False)
-            return conn
+
+            dest = sqlite3.connect(":memory:")
+            dest.enable_load_extension(True)
+            sqlite_vec_sl_tmp.load(dest)
+            dest.enable_load_extension(False)
+            conn.backup(dest)
+
+            return dest
         except Exception as e:
             logger.error("Failed to initialize database connection", error=str(e))
             raise
@@ -136,6 +144,7 @@ async def search_by_property(self, name: str, properties: List[str]) -> list[dic
             logger.error(f"An error occurred during property search: {str(e)}")
             return []
 
+    @profiled("search")
     async def search(
         self,
         query: Optional[str] = None,

src/codegate/updates/client.py

@@ -1,8 +1,8 @@
+import os
 from enum import Enum
 
 import requests
 import structlog
-import os
 
 logger = structlog.get_logger("codegate")