diff --git a/migrations/versions/2025_03_05_2126-e4c05d7591a8_add_installation_table.py b/migrations/versions/2025_03_05_2126-e4c05d7591a8_add_installation_table.py index 775e3967b..9e2b6c130 100644 --- a/migrations/versions/2025_03_05_2126-e4c05d7591a8_add_installation_table.py +++ b/migrations/versions/2025_03_05_2126-e4c05d7591a8_add_installation_table.py @@ -9,8 +9,6 @@ from typing import Sequence, Union from alembic import op -import sqlalchemy as sa - # revision identifiers, used by Alembic. revision: str = "e4c05d7591a8" diff --git a/poetry.lock b/poetry.lock index c0d03cd85..28679f12a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -647,7 +647,7 @@ files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] -markers = {main = "sys_platform == \"win32\" or platform_system == \"Windows\"", dev = "platform_system == \"Windows\" or os_name == \"nt\" or sys_platform == \"win32\""} +markers = {main = "sys_platform == \"win32\" or platform_system == \"Windows\"", dev = "platform_system == \"Windows\" or sys_platform == \"win32\" or os_name == \"nt\""} [[package]] name = "coloredlogs" diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index 5db3ae549..b0f290e37 100644 --- a/src/codegate/api/v1.py +++ b/src/codegate/api/v1.py @@ -10,7 +10,7 @@ from codegate.config import API_DEFAULT_PAGE_SIZE, API_MAX_PAGE_SIZE import codegate.muxing.models as mux_models -from codegate import __version__ +from codegate import Config, __version__ from codegate.api import v1_models, v1_processing from codegate.db.connection import AlreadyExistsError, DbReader from codegate.db.models import AlertSeverity, AlertTriggerType, Persona, WorkspaceWithModel @@ -20,6 +20,7 @@ PersonaSimilarDescriptionError, ) from codegate.providers import crud as provendcrud +from codegate.updates.client import Origin, UpdateClient from codegate.workspaces import crud logger = structlog.get_logger("codegate") @@ -31,6 +32,7 @@ # This is a singleton object dbreader = DbReader() +update_client = UpdateClient(Config.get_config().update_service_url, __version__, dbreader) def uniq_name(route: APIRoute): @@ -724,10 +726,12 @@ async def stream_sse(): @v1.get("/version", tags=["Dashboard"], generate_unique_id_function=uniq_name) -def version_check(): +async def version_check(): try: - latest_version = v1_processing.fetch_latest_version() - + if Config.get_config().use_update_service: + latest_version = await update_client.get_latest_version(Origin.FrontEnd) + else: + latest_version = v1_processing.fetch_latest_version() # normalize the versions as github will return them with a 'v' prefix current_version = __version__.lstrip("v") latest_version_stripped = latest_version.lstrip("v") diff --git a/src/codegate/cli.py b/src/codegate/cli.py index 1ae3f9c22..5c08821c5 100644 --- a/src/codegate/cli.py +++ b/src/codegate/cli.py @@ -16,8 +16,8 @@ from codegate.config import Config, ConfigurationError from codegate.db.connection import ( init_db_sync, - init_session_if_not_exists, init_instance, + init_session_if_not_exists, ) from codegate.pipeline.factory import PipelineFactory from codegate.pipeline.sensitive_data.manager import SensitiveDataManager diff --git a/src/codegate/config.py b/src/codegate/config.py index ee5cb1689..e549314f5 100644 --- a/src/codegate/config.py +++ b/src/codegate/config.py @@ -59,6 +59,10 @@ class Config: server_key: str = "server.key" force_certs: bool = False + # Update configuration. + use_update_service: bool = False + update_service_url: str = "https://updates.codegate.ai/api/v1/version" + max_fim_hash_lifetime: int = 60 * 5 # Time in seconds. Default is 5 minutes. # Min value is 0 (max similarity), max value is 2 (orthogonal) @@ -165,6 +169,8 @@ def from_file(cls, config_path: Union[str, Path]) -> "Config": force_certs=config_data.get("force_certs", cls.force_certs), prompts=prompts_config, provider_urls=provider_urls, + use_update_service=config_data.get("use_update_service", cls.use_update_service), + update_service_url=config_data.get("update_service_url", cls.update_service_url), ) except yaml.YAMLError as e: raise ConfigurationError(f"Failed to parse config file: {e}") @@ -209,11 +215,17 @@ def from_env(cls) -> "Config": if "CODEGATE_SERVER_KEY" in os.environ: config.server_key = os.environ["CODEGATE_SERVER_KEY"] if "CODEGATE_FORCE_CERTS" in os.environ: - config.force_certs = os.environ["CODEGATE_FORCE_CERTS"] + config.force_certs = cls.__bool_from_string(os.environ["CODEGATE_FORCE_CERTS"]) if "CODEGATE_DB_PATH" in os.environ: config.db_path = os.environ["CODEGATE_DB_PATH"] if "CODEGATE_VEC_DB_PATH" in os.environ: config.vec_db_path = os.environ["CODEGATE_VEC_DB_PATH"] + if "CODEGATE_USE_UPDATE_SERVICE" in os.environ: + config.use_update_service = cls.__bool_from_string( + os.environ["CODEGATE_USE_UPDATE_SERVICE"] + ) + if "CODEGATE_UPDATE_SERVICE_URL" in os.environ: + config.update_service_url = os.environ["CODEGATE_UPDATE_SERVICE_URL"] # Load provider URLs from environment variables for provider in DEFAULT_PROVIDER_URLS.keys(): @@ -246,6 +258,8 @@ def load( force_certs: Optional[bool] = None, db_path: Optional[str] = None, vec_db_path: Optional[str] = None, + use_update_service: Optional[bool] = None, + update_service_url: Optional[str] = None, ) -> "Config": """Load configuration with priority resolution. @@ -274,6 +288,8 @@ def load( force_certs: Optional flag to force certificate generation db_path: Optional path to the main SQLite database file vec_db_path: Optional path to the vector SQLite database file + use_update_service: Optional flag to enable the update service + update_service_url: Optional URL for the update service Returns: Config: Resolved configuration @@ -326,6 +342,10 @@ def load( config.db_path = env_config.db_path if "CODEGATE_VEC_DB_PATH" in os.environ: config.vec_db_path = env_config.vec_db_path + if "CODEGATE_USE_UPDATE_SERVICE" in os.environ: + config.use_update_service = env_config.use_update_service + if "CODEGATE_UPDATE_SERVICE_URL" in os.environ: + config.update_service_url = env_config.update_service_url # Override provider URLs from environment for provider, url in env_config.provider_urls.items(): @@ -366,6 +386,10 @@ def load( config.vec_db_path = vec_db_path if force_certs is not None: config.force_certs = force_certs + if use_update_service is not None: + config.use_update_service = use_update_service + if update_service_url is not None: + config.update_service_url = update_service_url # Set the __config class attribute Config.__config = config @@ -375,3 +399,7 @@ def load( @classmethod def get_config(cls) -> "Config": return cls.__config + + @staticmethod + def __bool_from_string(raw_value) -> bool: + return raw_value.lower() == "true" diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 915c4251c..44f6ea1c3 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -619,7 +619,7 @@ async def init_instance(self) -> None: await self._execute_with_no_return(sql, instance.model_dump()) except IntegrityError as e: logger.debug(f"Exception type: {type(e)}") - raise AlreadyExistsError(f"Instance already initialized.") + raise AlreadyExistsError("Instance already initialized.") class DbReader(DbCodeGate): diff --git a/src/codegate/updates/client.py b/src/codegate/updates/client.py new file mode 100644 index 000000000..5ba21ee83 --- /dev/null +++ b/src/codegate/updates/client.py @@ -0,0 +1,54 @@ +from enum import Enum + +import cachetools.func +import requests +import structlog + +from codegate.db.connection import DbReader + +logger = structlog.get_logger("codegate") + + +# Enum representing whether the request is coming from the front-end or the back-end. +class Origin(Enum): + FrontEnd = "FE" + BackEnd = "BE" + + +class UpdateClient: + def __init__(self, update_url: str, current_version: str, db_reader: DbReader): + self.__update_url = update_url + self.__current_version = current_version + self.__db_reader = db_reader + self.__instance_id = None + + async def get_latest_version(self, origin: Origin) -> str: + """ + Retrieves the latest version of CodeGate from updates.codegate.ai + """ + logger.info(f"Fetching latest version from {self.__update_url}") + instance_id = await self.__get_instance_id() + return self.__fetch_latest_version(instance_id, origin) + + @cachetools.func.ttl_cache(maxsize=128, ttl=20 * 60) + def __fetch_latest_version(self, instance_id: str, origin: Origin) -> str: + headers = { + "X-Instance-ID": instance_id, + "User-Agent": f"codegate/{self.__current_version} {origin.value}", + } + + try: + response = requests.get(self.__update_url, headers=headers, timeout=10) + # Throw if the request was not successful. + response.raise_for_status() + return response.json()["version"] + except Exception as e: + logger.error(f"Error fetching latest version from f{self.__update_url}: {e}") + return "unknown" + + # Lazy load the instance ID from the DB. + async def __get_instance_id(self): + if self.__instance_id is None: + instance_data = await self.__db_reader.get_instance() + self.__instance_id = instance_data[0].id + return self.__instance_id