File size: 2,968 Bytes
fe52ef9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import os
from typing import Dict, Any
from dotenv import load_dotenv
from openai import OpenAI, AzureOpenAI
load_dotenv()
# Initialize clients only if their required environment variables are present
engine_map = {}
# Import here to avoid circular imports - will be populated in code below
from agent_config import BaseAgentConfig, SITAASAgentConfig
# Azure OpenAI client setup
if all(os.getenv(var) for var in ["AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_API_VERSION"]):
azure_client = AzureOpenAI(
api_key=os.getenv("AZURE_OPENAI_API_KEY"),
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
api_version=os.getenv("AZURE_OPENAI_API_VERSION")
)
engine_map.update({
"azure-gpt-4o-mini": {
"client": azure_client,
"model": "gpt-4o-mini",
"config": BaseAgentConfig("azure-gpt-4o-mini")
},
"azure-gpt-4o": {
"client": azure_client,
"model": "gpt-4o",
"config": BaseAgentConfig("azure-gpt-4o")
},
})
# OpenAI client setup
if os.getenv("OPENAI_API_KEY"):
openai_client = OpenAI(
api_key=os.getenv("OPENAI_API_KEY"),
)
engine_map.update({
"openai-gpt-4o-mini": {
"client": openai_client,
"model": "gpt-4o-mini",
"config": BaseAgentConfig("openai-gpt-4o-mini")
},
"openai-gpt-4o": {
"client": openai_client,
"model": "gpt-4o",
"config": BaseAgentConfig("openai-gpt-4o")
},
})
# x.ai client setup
if os.getenv("XAI_API_KEY"):
xai_client = OpenAI(
api_key=os.getenv("XAI_API_KEY"),
base_url="https://api.x.ai/v1",
)
engine_map["grok-2-latest"] = {
"client": xai_client,
"model": "grok-2-latest",
"config": BaseAgentConfig("grok-2-latest")
}
# Add the SITAAS agent to the engine_map if required env vars are available
if all(os.getenv(var) for var in ["SEARCH_ENDPOINT", "SEARCH_KEY"]):
engine_map["sitaas-agent"] = {
"client": azure_client,
"model": "gpt-4.1-mini",
"config": SITAASAgentConfig("gpt-4.1-mini") # Use the SITAAS agent configuration
}
def get_client(engine: str) -> OpenAI:
"""Get the appropriate client for the given engine"""
try:
return engine_map[engine]["client"]
except KeyError:
raise ValueError(f"Unsupported engine in get_client: {engine}")
def get_model(engine: str) -> str:
"""Get the model name for the given engine"""
try:
return engine_map[engine]["model"]
except KeyError:
raise ValueError(f"Unsupported engine in get_model: {engine}")
def get_config(engine: str) -> Any:
"""Get the agent configuration for the given engine"""
try:
return engine_map[engine]["config"]
except KeyError:
raise ValueError(f"Unsupported engine in get_config: {engine}")
|