code-graph-rag 0.0.79__py3-none-any.whl

codebase_rag/config.py

@@ -0,0 +1,370 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import TypedDict, Unpack
+
+from dotenv import load_dotenv
+from loguru import logger
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from . import constants as cs
+from . import exceptions as ex
+from . import logs
+from .types_defs import CgrignorePatterns, ModelConfigKwargs
+
+load_dotenv()
+
+
+class ApiKeyInfoEntry(TypedDict):
+    env_var: str
+    url: str
+    name: str
+
+
+API_KEY_INFO: dict[str, ApiKeyInfoEntry] = {
+    cs.Provider.OPENAI: {
+        "env_var": "OPENAI_API_KEY",
+        "url": "https://platform.openai.com/api-keys",
+        "name": "OpenAI",
+    },
+    cs.Provider.ANTHROPIC: {
+        "env_var": "ANTHROPIC_API_KEY",
+        "url": "https://console.anthropic.com/settings/keys",
+        "name": "Anthropic",
+    },
+    cs.Provider.GOOGLE: {
+        "env_var": "GOOGLE_API_KEY",
+        "url": "https://console.cloud.google.com/apis/credentials",
+        "name": "Google AI",
+    },
+    cs.Provider.AZURE: {
+        "env_var": "AZURE_API_KEY",
+        "url": "https://portal.azure.com/",
+        "name": "Azure OpenAI",
+    },
+    cs.Provider.COHERE: {
+        "env_var": "COHERE_API_KEY",
+        "url": "https://dashboard.cohere.com/api-keys",
+        "name": "Cohere",
+    },
+}
+
+
+def format_missing_api_key_errors(
+    provider: str, role: str = cs.DEFAULT_MODEL_ROLE
+) -> str:
+    provider_lower = provider.lower()
+
+    if provider_lower in API_KEY_INFO:
+        info = API_KEY_INFO[provider_lower]
+        env_var = info["env_var"]
+        url = info["url"]
+        name = info["name"]
+    else:
+        env_var = f"{provider.upper()}_API_KEY"
+        url = f"your {provider} provider's website"
+        name = provider.capitalize()
+
+    role_msg = f" for {role}" if role != cs.DEFAULT_MODEL_ROLE else ""
+
+    error_msg = f"""
+─── API Key Missing ───────────────────────────────────────────────
+
+  Error: {env_var} environment variable is not set.
+         This is required to use {name}{role_msg}.
+
+  To fix this:
+
+  1. Get your API key from:
+     {url}
+
+  2. Set it in your environment:
+     export {env_var}='your-key-here'
+
+     Or add it to your .env file in the project root:
+     {env_var}=your-key-here
+
+  3. Alternatively, you can use a local model with Ollama:
+     (No API key required)
+
+───────────────────────────────────────────────────────────────────
+""".strip()  # noqa: W293
+    return error_msg
+
+
+LOCAL_PROVIDERS = frozenset({cs.Provider.OLLAMA, cs.Provider.LOCAL, cs.Provider.VLLM})
+
+
+@dataclass
+class ModelConfig:
+    provider: str
+    model_id: str
+    api_key: str | None = None
+    endpoint: str | None = None
+    project_id: str | None = None
+    region: str | None = None
+    provider_type: str | None = None
+    thinking_budget: int | None = None
+    service_account_file: str | None = None
+
+    def to_update_kwargs(self) -> ModelConfigKwargs:
+        result = asdict(self)
+        del result[cs.FIELD_PROVIDER]
+        del result[cs.FIELD_MODEL_ID]
+        return ModelConfigKwargs(**result)
+
+    def validate_api_key(self, role: str = cs.DEFAULT_MODEL_ROLE) -> None:
+        provider_lower = self.provider.lower()
+        if provider_lower in LOCAL_PROVIDERS or (
+            provider_lower == cs.Provider.GOOGLE
+            and self.provider_type == cs.GoogleProviderType.VERTEX
+        ):
+            return
+        if (
+            not self.api_key
+            or not self.api_key.strip()
+            or self.api_key == cs.DEFAULT_API_KEY
+        ):
+            error_msg = format_missing_api_key_errors(self.provider, role)
+            raise ValueError(error_msg)
+
+
+class AppConfig(BaseSettings):
+    """
+    (H) All settings are loaded from environment variables or a .env file.
+    """
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        case_sensitive=False,
+    )
+
+    MEMGRAPH_HOST: str = "localhost"
+    MEMGRAPH_PORT: int = 7687
+    MEMGRAPH_HTTP_PORT: int = 7444
+    LAB_PORT: int = 3000
+    MEMGRAPH_BATCH_SIZE: int = 1000
+    AGENT_RETRIES: int = 3
+    ORCHESTRATOR_OUTPUT_RETRIES: int = 100
+
+    ORCHESTRATOR_PROVIDER: str = ""
+    ORCHESTRATOR_MODEL: str = ""
+    ORCHESTRATOR_API_KEY: str | None = None
+    ORCHESTRATOR_ENDPOINT: str | None = None
+    ORCHESTRATOR_PROJECT_ID: str | None = None
+    ORCHESTRATOR_REGION: str = cs.DEFAULT_REGION
+    ORCHESTRATOR_PROVIDER_TYPE: cs.GoogleProviderType | None = None
+    ORCHESTRATOR_THINKING_BUDGET: int | None = None
+    ORCHESTRATOR_SERVICE_ACCOUNT_FILE: str | None = None
+
+    CYPHER_PROVIDER: str = ""
+    CYPHER_MODEL: str = ""
+    CYPHER_API_KEY: str | None = None
+    CYPHER_ENDPOINT: str | None = None
+    CYPHER_PROJECT_ID: str | None = None
+    CYPHER_REGION: str = cs.DEFAULT_REGION
+    CYPHER_PROVIDER_TYPE: cs.GoogleProviderType | None = None
+    CYPHER_THINKING_BUDGET: int | None = None
+    CYPHER_SERVICE_ACCOUNT_FILE: str | None = None
+
+    OLLAMA_BASE_URL: str = "http://localhost:11434"
+
+    @property
+    def ollama_endpoint(self) -> str:
+        return f"{self.OLLAMA_BASE_URL.rstrip('/')}/v1"
+
+    TARGET_REPO_PATH: str = "."
+    SHELL_COMMAND_TIMEOUT: int = 30
+    SHELL_COMMAND_ALLOWLIST: frozenset[str] = frozenset(
+        {
+            "ls",
+            "rg",
+            "cat",
+            "git",
+            "echo",
+            "pwd",
+            "pytest",
+            "mypy",
+            "ruff",
+            "uv",
+            "find",
+            "pre-commit",
+            "rm",
+            "cp",
+            "mv",
+            "mkdir",
+            "rmdir",
+            "wc",
+            "head",
+            "tail",
+            "sort",
+            "uniq",
+            "cut",
+            "tr",
+            "xargs",
+            "awk",
+            "sed",
+            "tee",
+        }
+    )
+    SHELL_READ_ONLY_COMMANDS: frozenset[str] = frozenset(
+        {
+            "ls",
+            "cat",
+            "find",
+            "pwd",
+            "rg",
+            "echo",
+            "wc",
+            "head",
+            "tail",
+            "sort",
+            "uniq",
+            "cut",
+            "tr",
+        }
+    )
+    SHELL_SAFE_GIT_SUBCOMMANDS: frozenset[str] = frozenset(
+        {
+            "status",
+            "log",
+            "diff",
+            "show",
+            "ls-files",
+            "remote",
+            "config",
+            "branch",
+        }
+    )
+
+    QDRANT_DB_PATH: str = "./.qdrant_code_embeddings"
+    QDRANT_COLLECTION_NAME: str = "code_embeddings"
+    QDRANT_VECTOR_DIM: int = 768
+    QDRANT_TOP_K: int = 5
+    EMBEDDING_MAX_LENGTH: int = 512
+    EMBEDDING_PROGRESS_INTERVAL: int = 10
+
+    CACHE_MAX_ENTRIES: int = 1000
+    CACHE_MAX_MEMORY_MB: int = 500
+    CACHE_EVICTION_DIVISOR: int = 10
+    CACHE_MEMORY_THRESHOLD_RATIO: float = 0.8
+
+    OLLAMA_HEALTH_TIMEOUT: float = 5.0
+
+    _active_orchestrator: ModelConfig | None = None
+    _active_cypher: ModelConfig | None = None
+
+    QUIET: bool = Field(False, validation_alias="CGR_QUIET")
+
+    def _get_default_config(self, role: str) -> ModelConfig:
+        role_upper = role.upper()
+
+        provider = getattr(self, f"{role_upper}_PROVIDER", None)
+        model = getattr(self, f"{role_upper}_MODEL", None)
+
+        if provider and model:
+            return ModelConfig(
+                provider=provider.lower(),
+                model_id=model,
+                api_key=getattr(self, f"{role_upper}_API_KEY", None),
+                endpoint=getattr(self, f"{role_upper}_ENDPOINT", None),
+                project_id=getattr(self, f"{role_upper}_PROJECT_ID", None),
+                region=getattr(self, f"{role_upper}_REGION", cs.DEFAULT_REGION),
+                provider_type=getattr(self, f"{role_upper}_PROVIDER_TYPE", None),
+                thinking_budget=getattr(self, f"{role_upper}_THINKING_BUDGET", None),
+                service_account_file=getattr(
+                    self, f"{role_upper}_SERVICE_ACCOUNT_FILE", None
+                ),
+            )
+
+        return ModelConfig(
+            provider=cs.Provider.OLLAMA,
+            model_id=cs.DEFAULT_MODEL,
+            endpoint=self.ollama_endpoint,
+            api_key=cs.DEFAULT_API_KEY,
+        )
+
+    def _get_default_orchestrator_config(self) -> ModelConfig:
+        return self._get_default_config(cs.ModelRole.ORCHESTRATOR)
+
+    def _get_default_cypher_config(self) -> ModelConfig:
+        return self._get_default_config(cs.ModelRole.CYPHER)
+
+    @property
+    def active_orchestrator_config(self) -> ModelConfig:
+        return self._active_orchestrator or self._get_default_orchestrator_config()
+
+    @property
+    def active_cypher_config(self) -> ModelConfig:
+        return self._active_cypher or self._get_default_cypher_config()
+
+    def set_orchestrator(
+        self, provider: str, model: str, **kwargs: Unpack[ModelConfigKwargs]
+    ) -> None:
+        config = ModelConfig(provider=provider.lower(), model_id=model, **kwargs)
+        self._active_orchestrator = config
+
+    def set_cypher(
+        self, provider: str, model: str, **kwargs: Unpack[ModelConfigKwargs]
+    ) -> None:
+        config = ModelConfig(provider=provider.lower(), model_id=model, **kwargs)
+        self._active_cypher = config
+
+    def parse_model_string(self, model_string: str) -> tuple[str, str]:
+        if ":" not in model_string:
+            return cs.Provider.OLLAMA, model_string
+        provider, model = model_string.split(":", 1)
+        if not provider:
+            raise ValueError(ex.PROVIDER_EMPTY)
+        return provider.lower(), model
+
+    def resolve_batch_size(self, batch_size: int | None) -> int:
+        resolved = self.MEMGRAPH_BATCH_SIZE if batch_size is None else batch_size
+        if resolved < 1:
+            raise ValueError(ex.BATCH_SIZE_POSITIVE)
+        return resolved
+
+
+settings = AppConfig()
+
+CGRIGNORE_FILENAME = ".cgrignore"
+
+
+EMPTY_CGRIGNORE = CgrignorePatterns(exclude=frozenset(), unignore=frozenset())
+
+
+def load_cgrignore_patterns(repo_path: Path) -> CgrignorePatterns:
+    ignore_file = repo_path / CGRIGNORE_FILENAME
+    if not ignore_file.is_file():
+        return EMPTY_CGRIGNORE
+
+    exclude: set[str] = set()
+    unignore: set[str] = set()
+    try:
+        with ignore_file.open(encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith("#"):
+                    continue
+                if line.startswith("!"):
+                    unignore.add(line[1:].strip())
+                else:
+                    exclude.add(line)
+        if exclude or unignore:
+            logger.info(
+                logs.CGRIGNORE_LOADED.format(
+                    exclude_count=len(exclude),
+                    unignore_count=len(unignore),
+                    path=ignore_file,
+                )
+            )
+        return CgrignorePatterns(
+            exclude=frozenset(exclude),
+            unignore=frozenset(unignore),
+        )
+    except OSError as e:
+        logger.warning(logs.CGRIGNORE_READ_FAILED.format(path=ignore_file, error=e))
+        return EMPTY_CGRIGNORE