config.py

"""
Configuration file for RS Studies MCP Server
Customize these settings based on your local setup.
"""

import os
import torch

class Config:
    """Configuration settings for the RS Studies MCP Server"""
    
    # Server Information
    SERVER_NAME = "rs-studies-mcp"
    SERVER_DESCRIPTION = "RS Studies Knowledge Base MCP Server"
    SERVER_VERSION = "1.0.0"
        
    # Model path - adjust if your model is in a different location
    MODEL_PATH = "google/embeddinggemma-300m"
    
    # ChromaDB path - relative to the base directory
    CHROMADB_PATH = os.path.join(BASE_DIR, "chromadb_data")
    
    # Collection name in ChromaDB
    COLLECTION_NAME = "rs_studies_embeddings"
    
    # Search Configuration
    DEFAULT_NUM_RESULTS = 5
    MAX_NUM_RESULTS = 50
    DEFAULT_CONTEXT_SIZE = 3
    MAX_CONTEXT_SIZE = 10
    
    # Valid source filters based on your data structure
    VALID_SOURCES = [
        'rs_stkege_01',        # RS trading system documentation
        'cheenai_meet_full',   # Chennai meetup transcripts
        'QnAYoutubeChannel'    # Q&A discussions
    ]
    
    # Device configuration
    @staticmethod
    def get_device():
        """Get the best available device for inference"""
        if torch.cuda.is_available():
            return "cuda"
        else:
            return "cpu"
    
    # Embedding configuration
    EMBEDDING_BATCH_SIZE = 32

# Create a default config instance
config = Config()

class ProdConfig(Config):
    """Production configuration with optimized settings"""
    MAX_NUM_RESULTS = 50
    EMBEDDING_BATCH_SIZE = 64

# Environment-based configuration selection
def get_config(env="default"):
    """Get configuration based on environment"""
    return ProdConfig()