Utils/llm/config.py (171 lines of code) (raw):

import os import subprocess from dotenv import load_dotenv load_dotenv() deployed_llm_base_url = os.getenv('LLM_PROVIDER_BASE_URL') deployed_llm_key = os.getenv('LLM_PROVIDER_KEY') open_api_key = os.getenv('OPENAI_API_KEY') xai_api_key = os.getenv('XAI_API_KEY') fireworks_api_key = os.getenv('FIREWORKS_API_KEY') google_ai_api_key = os.getenv('GOOGLE_AI_STUDIO_API_KEY') gcloud_path = os.getenv('GCLOUD_PATH') gcloud_project_id = os.getenv('GCLOUD_PROJECT_ID') temperature = 0 attempts_count = 1 def get_gcp_access_token(): return subprocess.check_output( [gcloud_path, 'auth', 'print-access-token']).decode('utf-8').strip() class Model: Gemini = "Gemini1_0" GPT4 = "GPT4_Turbo" GeminiPro = "Gemini1_5" GeminiPro_0801 = "GeminiPro_0801" Gemini_15_Pro_002 = "Gemini_15_Pro_002" GeminiPro_1114 = "GeminiPro_1114" GeminiPro_1121 = "GeminiPro_1121" Opus_3 = "Claude_Opus_3" Sonnet_35 = "Claude_Sonnet_35" Sonnet_35v2 = "Claude_Sonnet_35v2" Haiku_35 = "Claude_Haiku_35" GPT35_Turbo_0125 = "GPT35_Turbo_0125" GPT4_Turbo_0409 = "GPT4_Turbo_0409" GPT4o_0513 = "GPT4o_0513" GPT4o_0806 = "GPT4o_0806" GPT4o_1120 = "GPT4o_1120" ChatGPT4o = "ChatGPT4o" GPT4o_mini = "GPT4o_mini_0718" OpenAi_o1_0912 = "OpenAi_o1_0912" OpenAi_o1_mini_0912 = "OpenAi_o1_mini_0912" Llama3_70B = "Llama3_70B" Llama31_405B = "Llama31_405B" GrokBeta = "GrokBeta" Qwen25Coder32B = "Qwen25Coder32B" AmazonNovaPro = "AmazonNovaPro" def get_azure_config(model): def config(): return { "model_id": model, "api_key": deployed_llm_key, "url": f'{deployed_llm_base_url}/openai/deployments/{model}/chat/completions?api-version=2023-12-01-preview' } return config def get_open_ai_config(model, max_tokens=None, skip_system=False): def config(): return { "model_id": model, "api_key": open_api_key, "max_tokens": max_tokens, "skip_system": skip_system, "url": 'https://api.openai.com/v1/chat/completions' } return config def get_xai_config(model): def config(): return { "model_id": model, "api_key": xai_api_key, "url": 'https://api.x.ai/v1/chat/completions' } return config def get_fireworks_config(model, max_tokens): def config(): return { "model_id": model, "max_tokens": max_tokens, "api_key": fireworks_api_key, "url": "https://api.fireworks.ai/inference/v1/chat/completions" } return config def get_gemini_pro_config(): gcp_access_token = get_gcp_access_token() LOCATION_ID = "us-central1" PROJECT_ID = gcloud_project_id MODEL_ID = "gemini-1.5-pro-preview-0409" return { "model_id": MODEL_ID, "api_key": gcp_access_token, "url": f"https://{LOCATION_ID}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION_ID}/publishers/google/models/{MODEL_ID}:generateContent" } def get_gemini_ai_studio_config(model): def config(): return { "model_id": model, "url": f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={google_ai_api_key}" } return config def get_opus_3_config(): gcp_access_token = get_gcp_access_token() LOCATION_ID = "us-east5" PROJECT_ID = gcloud_project_id MODEL_ID = "claude-3-opus@20240229" return { "version": "vertex-2023-10-16", "model_id": MODEL_ID, "api_key": gcp_access_token, "url": f"https://{LOCATION_ID}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION_ID}/publishers/anthropic/models/{MODEL_ID}:streamRawPredict" } def get_sonnet_35_config(): gcp_access_token = get_gcp_access_token() LOCATION_ID = "europe-west1" PROJECT_ID = gcloud_project_id MODEL_ID = "claude-3-5-sonnet@20240620" return { "version": "vertex-2023-10-16", "model_id": MODEL_ID, "api_key": gcp_access_token, "url": f"https://{LOCATION_ID}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION_ID}/publishers/anthropic/models/{MODEL_ID}:streamRawPredict" } def get_sonnet_35_v2_config(): gcp_access_token = get_gcp_access_token() # LOCATION_ID = "us-east5" LOCATION_ID = "europe-west1" PROJECT_ID = gcloud_project_id MODEL_ID = "claude-3-5-sonnet-v2@20241022" return { "version": "vertex-2023-10-16", "model_id": MODEL_ID, "api_key": gcp_access_token, "url": f"https://{LOCATION_ID}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION_ID}/publishers/anthropic/models/{MODEL_ID}:streamRawPredict" } def get_haiku_35_config(): gcp_access_token = get_gcp_access_token() LOCATION_ID = "us-east5" PROJECT_ID = gcloud_project_id MODEL_ID = "claude-3-5-haiku@20241022" return { "version": "vertex-2023-10-16", "model_id": MODEL_ID, "api_key": gcp_access_token, "url": f"https://{LOCATION_ID}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION_ID}/publishers/anthropic/models/{MODEL_ID}:streamRawPredict" } def get_amazon_nova_pro_config(): MODEL_ID = "us.amazon.nova-pro-v1:0" return { "model_id": MODEL_ID } API = { Model.Gemini: get_azure_config('gemini-pro'), Model.GeminiPro: get_gemini_pro_config, Model.GeminiPro_0801: get_gemini_ai_studio_config('gemini-1.5-pro-exp-0801'), Model.Gemini_15_Pro_002: get_gemini_ai_studio_config("gemini-1.5-pro-002"), Model.GeminiPro_1114: get_gemini_ai_studio_config("gemini-exp-1114"), Model.GeminiPro_1121: get_gemini_ai_studio_config("gemini-exp-1121"), Model.GPT4: get_azure_config('gpt-4-0125-preview'), Model.GPT35_Turbo_0125: get_azure_config('gpt-35-turbo-0125'), Model.GPT4_Turbo_0409: get_open_ai_config('gpt-4-turbo-2024-04-09'), Model.GPT4o_0513: get_open_ai_config('gpt-4o-2024-05-13'), Model.GPT4o_0806: get_open_ai_config('gpt-4o-2024-08-06', 16384), Model.GPT4o_1120: get_open_ai_config('gpt-4o-2024-11-20'), Model.ChatGPT4o: get_open_ai_config('chatgpt-4o-latest', 16384), Model.GPT4o_mini: get_open_ai_config('gpt-4o-mini-2024-07-18'), Model.OpenAi_o1_0912: get_open_ai_config('o1-preview-2024-09-12', skip_system=True), Model.OpenAi_o1_mini_0912: get_open_ai_config('o1-mini-2024-09-12', skip_system=True), Model.Opus_3: get_opus_3_config, Model.Sonnet_35: get_sonnet_35_config, Model.Sonnet_35v2: get_sonnet_35_v2_config, Model.Haiku_35: get_haiku_35_config, Model.Llama3_70B: get_azure_config('llama-3-70b-instruct-awq'), Model.Llama31_405B: get_fireworks_config("accounts/fireworks/models/llama-v3p1-405b-instruct", 16384), Model.GrokBeta: get_xai_config('grok-beta'), Model.Qwen25Coder32B: get_fireworks_config("accounts/fireworks/models/qwen2p5-coder-32b-instruct", 4096), Model.AmazonNovaPro: get_amazon_nova_pro_config, }