Spaces:

chrizefan
/

digichat

Runtime error

digichat / client_utils.py

Upload folder using huggingface_hub

fe52ef9 verified over 1 year ago

2.97 kB

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}")

	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}")