diff --git a/README.md b/README.md index 4c97d4b..0aefe00 100644 --- a/README.md +++ b/README.md @@ -64,10 +64,10 @@ Once your project and credentials are configured, export environment variables t ```zsh export VERTEX_AI_PROJECT_ID= -export VERTEX_AI_LOCATION= +export VERTEX_AI_LOCATION=global ``` -The project location can be located under the **Region** on the [Vertex AI](https://console.cloud.google.com/vertex-ai) dashboard, as well as more details about available Gemini models. +The included Gemini models are served from the `global` location or the `us` and `eu` multi-regions. More details about available Gemini models and their supported regions are found in the [Google Cloud documentation](https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models). ##### OpenAI Setup diff --git a/ai/providers/anthropic.py b/ai/providers/anthropic.py index 7d0df54..0853056 100644 --- a/ai/providers/anthropic.py +++ b/ai/providers/anthropic.py @@ -12,25 +12,20 @@ class AnthropicAPI(BaseAPIProvider): MODELS: ClassVar[dict] = { - "claude-3-5-sonnet-20240620": { - "name": "Claude 3.5 Sonnet", + "claude-haiku-5-5": { + "name": "Claude Haiku 5.5", "provider": "Anthropic", - "max_tokens": 4096, # or 8192 with the header anthropic-beta: max-tokens-3-5-sonnet-2024-07-15 + "max_tokens": 16000, }, - "claude-3-sonnet-20240229": { - "name": "Claude 3 Sonnet", + "claude-opus-5-5": { + "name": "Claude Opus 5.5", "provider": "Anthropic", - "max_tokens": 4096, + "max_tokens": 16000, }, - "claude-3-haiku-20240307": { - "name": "Claude 3 Haiku", + "claude-sonnet-5-5": { + "name": "Claude Sonnet 5.5", "provider": "Anthropic", - "max_tokens": 4096, - }, - "claude-3-opus-20240229": { - "name": "Claude 3 Opus", - "provider": "Anthropic", - "max_tokens": 4096, + "max_tokens": 16000, }, } @@ -59,7 +54,9 @@ def generate_response(self, prompt: str, system_content: str) -> str: ], max_tokens=self.MODELS[self.current_model]["max_tokens"], ) - return response.content[0].text + return "".join( + block.text for block in response.content if block.type == "text" + ) except anthropic.APIConnectionError as e: logger.error(f"Server could not be reached: {e.__cause__}") raise diff --git a/ai/providers/openai.py b/ai/providers/openai.py index fba886a..3371e86 100644 --- a/ai/providers/openai.py +++ b/ai/providers/openai.py @@ -12,18 +12,21 @@ class OpenAI_API(BaseAPIProvider): MODELS: ClassVar[dict] = { - "gpt-4.1": {"name": "GPT-4.1", "provider": "OpenAI", "max_tokens": 10000}, - "gpt-4.1-mini": { - "name": "GPT-4.1 Mini", + "gpt-6-astra": { + "name": "GPT-6 Astra", "provider": "OpenAI", - "max_tokens": 10000, + "max_tokens": 16000, }, - "gpt-4.1-nano": { - "name": "GPT-4.1 Nano", + "gpt-6-luna": { + "name": "GPT-6 Luna", "provider": "OpenAI", - "max_tokens": 10000, + "max_tokens": 16000, + }, + "gpt-6.1-sol": { + "name": "GPT-6.1 Sol", + "provider": "OpenAI", + "max_tokens": 16000, }, - "o4-mini": {"name": "o4-mini", "provider": "OpenAI", "max_tokens": 50000}, } def __init__(self): diff --git a/ai/providers/vertexai.py b/ai/providers/vertexai.py index c3b1b85..34625b0 100644 --- a/ai/providers/vertexai.py +++ b/ai/providers/vertexai.py @@ -2,8 +2,8 @@ import os from typing import ClassVar -import google.api_core.exceptions -import vertexai.generative_models +from google import genai +from google.genai import errors, types from .base_provider import BaseAPIProvider @@ -14,58 +14,22 @@ class VertexAPI(BaseAPIProvider): VERTEX_AI_PROVIDER = "VertexAI" MODELS: ClassVar[dict] = { - "gemini-1.5-flash-001": { - "name": "Gemini 1.5 Flash 001", + "gemini-3.5-flash-lite": { + "name": "Gemini 3.5 Flash-Lite", "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, + "max_tokens": 16384, "system_instruction_supported": True, }, - "gemini-1.5-flash-002": { - "name": "Gemini 1.5 Flash 002", + "gemini-3.8-flash": { + "name": "Gemini 3.8 Flash", "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, + "max_tokens": 16384, "system_instruction_supported": True, }, - "gemini-1.5-pro-002": { - "name": "Gemini 1.5 Pro 002", + "gemini-4-argon": { + "name": "Gemini 4 Argon", "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, - "system_instruction_supported": True, - }, - "gemini-1.5-pro-001": { - "name": "Gemini 1.5 Pro 001", - "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, - "system_instruction_supported": True, - }, - "gemini-1.0-pro-002": { - "name": "Gemini 1.0 Pro 002", - "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, - "system_instruction_supported": True, - }, - "gemini-1.0-pro-001": { - "name": "Gemini 1.0 Pro 001", - "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, - "system_instruction_supported": False, - }, - "gemini-flash-experimental": { - "name": "Gemini Flash Experimental", - "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, - "system_instruction_supported": True, - }, - "gemini-pro-experimental": { - "name": "Gemini Pro Experimental", - "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, - "system_instruction_supported": True, - }, - "gemini-experimental": { - "name": "Gemini Experimental", - "provider": VERTEX_AI_PROVIDER, - "max_tokens": 8192, + "max_tokens": 16384, "system_instruction_supported": True, }, } @@ -73,7 +37,8 @@ class VertexAPI(BaseAPIProvider): def __init__(self): self.enabled = bool(os.environ.get("VERTEX_AI_PROJECT_ID", "")) if self.enabled: - vertexai.init( + self.client = genai.Client( + enterprise=True, project=os.environ.get("VERTEX_AI_PROJECT_ID"), location=os.environ.get("VERTEX_AI_LOCATION"), ) @@ -97,36 +62,22 @@ def generate_response(self, prompt: str, system_content: str) -> str: prompt = system_content + "\n" + prompt try: - self.client = vertexai.generative_models.GenerativeModel( - model_name=self.current_model, - generation_config={ - "max_output_tokens": self.MODELS[self.current_model]["max_tokens"], - }, - system_instruction=system_instruction, - ) - response = self.client.generate_content( + response = self.client.models.generate_content( + model=self.current_model, contents=prompt, + config=types.GenerateContentConfig( + max_output_tokens=self.MODELS[self.current_model]["max_tokens"], + system_instruction=system_instruction, + ), ) - return "".join(part.text for part in response.candidates[0].content.parts) + return response.text - except google.api_core.exceptions.Unauthorized as e: - logger.error(f"Client is not Authorized. {e.reason}, {e.message}") - raise - except google.api_core.exceptions.Forbidden as e: - logger.error(f"Client Forbidden. {e.reason}, {e.message}") - raise - except google.api_core.exceptions.TooManyRequests as e: - logger.error(f"Too many requests. {e.reason}, {e.message}") - raise - except google.api_core.exceptions.ClientError as e: - logger.error(f"Client error: {e.reason}, {e.message}") - raise - except google.api_core.exceptions.ServerError as e: - logger.error(f"Server error: {e.reason}, {e.message}") + except errors.ClientError as e: + logger.error(f"Client error: {e.code} {e.status}, {e.message}") raise - except google.api_core.exceptions.GoogleAPICallError as e: - logger.error(f"Error: {e.reason}, {e.message}") + except errors.ServerError as e: + logger.error(f"Server error: {e.code} {e.status}, {e.message}") raise - except google.api_core.exceptions.GoogleAPIError as e: - logger.error(f"Unknown error. {e}") + except errors.APIError as e: + logger.error(f"Error: {e.code} {e.status}, {e.message}") raise diff --git a/requirements.txt b/requirements.txt index ed1e645..7bfe6d8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,4 @@ ruff==0.16.10 slack-cli-hooks==0.3.0 openai==3.26.0 anthropic==1.13.0 -google-cloud-aiplatform==2.4.0 +google-genai==2.30.1