memory_agent/config.py

from __future__ import annotations

import os
from dataclasses import dataclass

from dotenv import load_dotenv

load_dotenv()


@dataclass(frozen=True)
class AppConfig:
    mistral_api_key: str
    memory_data_dir: str = "./data/faiss"
    mistral_model: str = "mistral-small-latest"
    embedding_model: str = "mistral-embed"
    mistral_max_retries: int = 6
    mistral_max_concurrent_requests: int = 1
    mistral_timeout_seconds: int = 120
    rate_limit_base_delay_seconds: float = 1.0
    rate_limit_max_delay_seconds: float = 16.0
    dense_weight: float = 0.55
    sparse_weight: float = 0.45
    top_k: int = 6

    @classmethod
    def from_env(cls) -> "AppConfig":
        mistral_api_key = os.getenv("MISTRAL_API_KEY") or os.getenv("API_KEY")

        missing = []
        if not mistral_api_key:
            missing.append("API_KEY or MISTRAL_API_KEY")

        if missing:
            formatted = ", ".join(missing)
            raise ValueError(f"Missing required environment variable(s): {formatted}")

        memory_data_dir = os.getenv("MEMORY_DATA_DIR", "./data/faiss")
        mistral_model = os.getenv("MISTRAL_MODEL", "mistral-small-latest")
        embedding_model = os.getenv("EMBEDDING_MODEL", "mistral-embed")
        if embedding_model == "mistrall-embed":
            embedding_model = "mistral-embed"
        mistral_max_retries = int(os.getenv("MISTRAL_MAX_RETRIES", "6"))
        mistral_max_concurrent_requests = int(os.getenv("MISTRAL_MAX_CONCURRENT_REQUESTS", "1"))
        mistral_timeout_seconds = int(os.getenv("MISTRAL_TIMEOUT_SECONDS", "120"))
        rate_limit_base_delay_seconds = float(os.getenv("RATE_LIMIT_BASE_DELAY_SECONDS", "1.0"))
        rate_limit_max_delay_seconds = float(os.getenv("RATE_LIMIT_MAX_DELAY_SECONDS", "16.0"))
        dense_weight = float(os.getenv("DENSE_WEIGHT", "0.55"))
        sparse_weight = float(os.getenv("SPARSE_WEIGHT", "0.45"))
        top_k = int(os.getenv("TOP_K", "6"))

        return cls(
            mistral_api_key=mistral_api_key,
            memory_data_dir=memory_data_dir,
            mistral_model=mistral_model,
            embedding_model=embedding_model,
            mistral_max_retries=mistral_max_retries,
            mistral_max_concurrent_requests=mistral_max_concurrent_requests,
            mistral_timeout_seconds=mistral_timeout_seconds,
            rate_limit_base_delay_seconds=rate_limit_base_delay_seconds,
            rate_limit_max_delay_seconds=rate_limit_max_delay_seconds,
            dense_weight=dense_weight,
            sparse_weight=sparse_weight,
            top_k=top_k,
        )