src/config.py

from pydantic import Field
from pydantic_settings import BaseSettings
from typing import Dict, Any

class MockStatus():
    # Required for headless mode
    def update(self, *args, **kwargs):
        print("MockStatus update called with args: ", args, " and kwargs: ", kwargs)

class Config(BaseSettings):
    hf_token: str = Field(...)
    hf_model: str = Field("InvestmentResearchAI/LLM-ADE-dev") # We need this because I can't get the model template out of the ollama model
    ollama_model: str = Field("llama3")
    headless: bool = Field(False, description="Run in headless mode.")

    status: Any = MockStatus()

    az_search_endpoint: str = Field("https://analysis-bank.search.windows.net")
    az_search_api_key: str = Field(...)
    az_search_idx_name: str = Field("analysis-index-2024-05-19-0003")
    az_search_top_k: int = Field(4, description="Max number of results to retrun")
    
    azure_openai_api_key: str = Field(...)
    azure_openai_endpoint: str = Field("https://irai-openai-eastus.openai.azure.com/")

    chat_template: str = Field("chatml", description="Chat template for prompt formatting")
    num_fewshot: int | None = Field(None, description="Option to use json mode examples")
    max_depth: int = Field(3, description="Maximum number of recursive iteration")

config = Config(_env_file=".env")