config.py

"""Configuration management for MuSEAgent."""

from dataclasses import dataclass
from typing import Optional


@dataclass
class AgentConfig:
    """Agent configuration settings."""

    max_iterations: int = 10

    temperature: float = 0.2
    max_tokens: Optional[int] = None  # None means no limit

    memory_dir: str = "memory"


@dataclass
class APIConfig:
    """API configuration settings."""

    base_url: str = "http://localhost:8000/v1"  # Replace with your API endpoint
    max_concurrent_per_key: int = 10
    max_retries: int = 5
    timeout: int = 120


@dataclass
class RetrievalConfig:
    """State-level retrieval-augmented generation settings. enable=False skips all retrieval."""

    enable: bool = False
    bank_memory_dir: str = ""
    bank_dir_name: str = ""
    embedding_model: str = ""
    embedding_base_url: str = ""  # Replace with your embedding API endpoint
    embedding_api_key: str = ""  # Replace with your embedding API key
    min_score: float = 0.1
    min_q_value: int = 7
    experience_top_n: int = 1
    max_epoch: int = 1


@dataclass
class Config:
    """Main configuration container."""
    
    agent: AgentConfig
    api: APIConfig
    retrieval: RetrievalConfig
    
    def __init__(
        self,
        agent: Optional[AgentConfig] = None,
        api: Optional[APIConfig] = None,
        retrieval: Optional[RetrievalConfig] = None,
    ):
        """Initialize config with default or provided values."""
        self.agent = agent or AgentConfig()
        self.api = api or APIConfig()
        self.retrieval = retrieval or RetrievalConfig()
    
    @classmethod
    def default(cls) -> "Config":
        """Get default configuration."""
        return cls()


default_config = Config.default()