From 3b3b9451759fcb0530a79542d500792b841d0526 Mon Sep 17 00:00:00 2001
From: Michelangelo Mori <328978+blkt@users.noreply.github.com>
Date: Tue, 8 Apr 2025 13:31:27 +0200
Subject: [PATCH] Load RAG database in memory on use.

This change uses a trick to force the RAG database into memory by
dumping the SQLite database from file to a temporary in-memory version
of it using SQLite's `backup` function.

Based on `percall` statistic, this almost halves execution times.

With on-disk database.

```
   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      911    0.017    0.000   47.983    0.053 .../src/codegate/storage/storage_engine.py:149(search)
      911   41.336    0.045   41.336    0.045 {method 'execute' of 'sqlite3.Cursor' objects}
      911    0.007    0.000    6.489    0.007 .../src/codegate/inference/inference_engine.py:90(embed)
```

With in-memory database.

```
ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   658    0.009    0.000   21.198    0.032 .../src/codegate/storage/storage_engine.py:147(search)
   658   16.458    0.025   16.458    0.025 {method 'execute' of 'sqlite3.Cursor' objects}
   658    0.004    0.000    4.638    0.007 .../src/codegate/inference/inference_engine.py:90(embed)
```

Also, added profiling to providers and search. Profiling can be
activated exporting `CODEGATE_PROFILE_<key>` where `<key>` is the
string passed to `@profiled` annotations. Profiling points are
statically defined throughout the codebase.
---
 codegate-profiling                            | 13 +++++++++
 src/codegate/profiling/__init__.py            | 28 +++++++++++++++++++
 src/codegate/providers/anthropic/provider.py  |  2 ++
 src/codegate/providers/ollama/provider.py     |  2 ++
 src/codegate/providers/openai/provider.py     |  2 ++
 src/codegate/providers/openrouter/provider.py |  2 ++
 src/codegate/storage/storage_engine.py        | 11 +++++++-
 src/codegate/updates/client.py                |  2 +-
 8 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100755 codegate-profiling
 create mode 100644 src/codegate/profiling/__init__.py

diff --git a/codegate-profiling b/codegate-profiling
new file mode 100755
index 000000000..186dc74fb
--- /dev/null
+++ b/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)
diff --git a/src/codegate/profiling/__init__.py b/src/codegate/profiling/__init__.py
new file mode 100644
index 000000000..0ee66ca3f
--- /dev/null
+++ b/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
diff --git a/src/codegate/providers/anthropic/provider.py b/src/codegate/providers/anthropic/provider.py
index 13741b85f..9f8e87912 100644
--- a/src/codegate/providers/anthropic/provider.py
+++ b/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,
diff --git a/src/codegate/providers/ollama/provider.py b/src/codegate/providers/ollama/provider.py
index b6de1d46b..91c92c119 100644
--- a/src/codegate/providers/ollama/provider.py
+++ b/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,
diff --git a/src/codegate/providers/openai/provider.py b/src/codegate/providers/openai/provider.py
index ef8c4b5b4..3e4c5dc98 100644
--- a/src/codegate/providers/openai/provider.py
+++ b/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,
diff --git a/src/codegate/providers/openrouter/provider.py b/src/codegate/providers/openrouter/provider.py
index beedb34bd..254379a45 100644
--- a/src/codegate/providers/openrouter/provider.py
+++ b/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,
diff --git a/src/codegate/storage/storage_engine.py b/src/codegate/storage/storage_engine.py
index a03cc935b..fb4e664b2 100644
--- a/src/codegate/storage/storage_engine.py
+++ b/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,
diff --git a/src/codegate/updates/client.py b/src/codegate/updates/client.py
index 7c958d8c1..fc1aa778b 100644
--- a/src/codegate/updates/client.py
+++ b/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")
 

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://clevelandohioweatherforecast.com/php-proxy/index.php?q=https://patch-diff.githubusercontent.com/raw/stacklok/codegate/pull/1344.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/index.php?u=https://patch-diff.githubusercontent.com/raw/stacklok/codegate/pull/1344.patch" target="_blank">pFad Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/v3index.php?u=https://patch-diff.githubusercontent.com/raw/stacklok/codegate/pull/1344.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/v4index.php?u=https://patch-diff.githubusercontent.com/raw/stacklok/codegate/pull/1344.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>