diff --git a/codegate-profiling b/codegate-profiling new file mode 100755 index 00000000..186dc74f --- /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 00000000..0ee66ca3 --- /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 13741b85..9f8e8791 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 b6de1d46..91c92c11 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 ef8c4b5b..3e4c5dc9 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 beedb34b..254379a4 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 a03cc935..fb4e664b 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 7c958d8c..fc1aa778 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")