Skip to content

Commit 2ed45bd

Browse files
committed
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.
1 parent 02d218b commit 2ed45bd

File tree

8 files changed

+60
-2
lines changed

8 files changed

+60
-2
lines changed

Diff for: codegate-profiling

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#!/usr/bin/env python3
2+
3+
import os
4+
import pstats
5+
import sys
6+
7+
8+
if __name__ == "__main__":
9+
if len(sys.argv) < 2:
10+
print("Error: file name required")
11+
exit(1)
12+
ps = pstats.Stats(*sys.argv[1:]).sort_stats("cumulative")
13+
ps.print_stats(.03)

Diff for: src/codegate/profiling/__init__.py

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import cProfile
2+
3+
4+
def profiled(prefix):
5+
def wrapper(func):
6+
import datetime
7+
import os
8+
from functools import wraps
9+
10+
@wraps(func)
11+
async def inner(*args, **kwargs):
12+
13+
ts = datetime.datetime.utcnow()
14+
fname = f"pstats-{prefix}-{ts.strftime('%Y%m%dT%H%M%S%f')}.prof"
15+
fname = os.path.join(os.path.abspath("."), fname)
16+
17+
with cProfile.Profile() as pr:
18+
res = await func(*args, **kwargs)
19+
print(f"writing to {fname}")
20+
pr.dump_stats(fname)
21+
return res
22+
23+
if os.getenv(f"CODEGATE_PROFILE_{prefix.upper()}"):
24+
return inner
25+
26+
return func
27+
28+
return wrapper

Diff for: src/codegate/providers/anthropic/provider.py

+2
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from codegate.clients.clients import ClientType
99
from codegate.clients.detector import DetectClient
1010
from codegate.pipeline.factory import PipelineFactory
11+
from codegate.profiling import profiled
1112
from codegate.providers.anthropic.completion_handler import AnthropicCompletion
1213
from codegate.providers.base import BaseProvider, ModelFetchError
1314
from codegate.providers.fim_analyzer import FIMAnalyzer
@@ -68,6 +69,7 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
6869

6970
return [model["id"] for model in respjson.get("data", [])]
7071

72+
@profiled("anthropic")
7173
async def process_request(
7274
self,
7375
data: dict,

Diff for: src/codegate/providers/ollama/provider.py

+2
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from codegate.clients.detector import DetectClient
1010
from codegate.config import Config
1111
from codegate.pipeline.factory import PipelineFactory
12+
from codegate.profiling import profiled
1213
from codegate.providers.base import BaseProvider, ModelFetchError
1314
from codegate.providers.fim_analyzer import FIMAnalyzer
1415
from codegate.providers.ollama.completion_handler import OllamaShim
@@ -59,6 +60,7 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
5960

6061
return [model["name"] for model in jsonresp.get("models", [])]
6162

63+
@profiled("ollama")
6264
async def process_request(
6365
self,
6466
data: dict,

Diff for: src/codegate/providers/openai/provider.py

+2
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from codegate.clients.clients import ClientType
88
from codegate.clients.detector import DetectClient
99
from codegate.pipeline.factory import PipelineFactory
10+
from codegate.profiling import profiled
1011
from codegate.providers.base import BaseProvider, ModelFetchError
1112
from codegate.providers.completion import BaseCompletionHandler
1213
from codegate.providers.fim_analyzer import FIMAnalyzer
@@ -63,6 +64,7 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
6364

6465
return [model["id"] for model in jsonresp.get("data", [])]
6566

67+
@profiled("openai")
6668
async def process_request(
6769
self,
6870
data: dict,

Diff for: src/codegate/providers/openrouter/provider.py

+2
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from codegate.clients.clients import ClientType
66
from codegate.clients.detector import DetectClient
77
from codegate.pipeline.factory import PipelineFactory
8+
from codegate.profiling import profiled
89
from codegate.providers.fim_analyzer import FIMAnalyzer
910
from codegate.providers.litellmshim import LiteLLmShim
1011
from codegate.providers.openai import OpenAIProvider
@@ -49,6 +50,7 @@ def __init__(self, pipeline_factory: PipelineFactory):
4950
def provider_route_name(self) -> str:
5051
return "openrouter"
5152

53+
@profiled("openrouter")
5254
async def process_request(
5355
self,
5456
data: dict,

Diff for: src/codegate/storage/storage_engine.py

+10-1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
from codegate.config import Config
1111
from codegate.inference.inference_engine import LlamaCppInferenceEngine
12+
from codegate.profiling import profiled
1213

1314
logger = structlog.get_logger("codegate")
1415
VALID_ECOSYSTEMS = ["npm", "pypi", "crates", "maven", "go"]
@@ -78,7 +79,14 @@ def _get_connection(self):
7879
conn.enable_load_extension(True)
7980
sqlite_vec_sl_tmp.load(conn)
8081
conn.enable_load_extension(False)
81-
return conn
82+
83+
dest = sqlite3.connect(":memory:")
84+
dest.enable_load_extension(True)
85+
sqlite_vec_sl_tmp.load(dest)
86+
dest.enable_load_extension(False)
87+
conn.backup(dest)
88+
89+
return dest
8290
except Exception as e:
8391
logger.error("Failed to initialize database connection", error=str(e))
8492
raise
@@ -136,6 +144,7 @@ async def search_by_property(self, name: str, properties: List[str]) -> list[dic
136144
logger.error(f"An error occurred during property search: {str(e)}")
137145
return []
138146

147+
@profiled("search")
139148
async def search(
140149
self,
141150
query: Optional[str] = None,

Diff for: src/codegate/updates/client.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1+
import os
12
from enum import Enum
23

34
import requests
45
import structlog
5-
import os
66

77
logger = structlog.get_logger("codegate")
88

0 commit comments

Comments
 (0)