feat(vertex_ai_anthropic.py): add claude 3 on vertex ai support - working .completions call

.completions() call works

litellm/__init__.py

@@ -597,6 +597,7 @@ from .llms.nlp_cloud import NLPCloudConfig
 from .llms.aleph_alpha import AlephAlphaConfig
 from .llms.petals import PetalsConfig
 from .llms.vertex_ai import VertexAIConfig
+from .llms.vertex_ai_anthropic import VertexAIAnthropicConfig
 from .llms.sagemaker import SagemakerConfig
 from .llms.ollama import OllamaConfig
 from .llms.ollama_chat import OllamaChatConfig

litellm/llms/custom_httpx/http_handler.py

@@ -39,3 +39,40 @@ class AsyncHTTPHandler:
             asyncio.get_running_loop().create_task(self.close())
         except Exception:
             pass
+
+
+class HTTPHandler:
+    def __init__(self, concurrent_limit=1000):
+        # Create a client with a connection pool
+        self.client = httpx.Client(
+            limits=httpx.Limits(
+                max_connections=concurrent_limit,
+                max_keepalive_connections=concurrent_limit,
+            )
+        )
+
+    def close(self):
+        # Close the client when you're done with it
+        self.client.close()
+
+    def get(
+        self, url: str, params: Optional[dict] = None, headers: Optional[dict] = None
+    ):
+        response = self.client.get(url, params=params, headers=headers)
+        return response
+
+    def post(
+        self,
+        url: str,
+        data: Optional[dict] = None,
+        params: Optional[dict] = None,
+        headers: Optional[dict] = None,
+    ):
+        response = self.client.post(url, data=data, params=params, headers=headers)
+        return response
+
+    def __del__(self) -> None:
+        try:
+            self.close()
+        except Exception:
+            pass

litellm/llms/vertex_ai_anthropic.py

@@ -1,7 +1,36 @@
 # What is this?
 ## Handler file for calling claude-3 on vertex ai
-from typing import Callable, Optional, Any, Union, List
+import os, types
+import json
+from enum import Enum
+import requests, copy
+import time, uuid
+from typing import Callable, Optional, List
+from litellm.utils import ModelResponse, Usage, map_finish_reason, CustomStreamWrapper
 import litellm
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+from .prompt_templates.factory import (
+    contains_tag,
+    prompt_factory,
+    custom_prompt,
+    construct_tool_use_system_prompt,
+    extract_between_tags,
+    parse_xml_params,
+)
+import httpx
+
+
+class VertexAIError(Exception):
+    def __init__(self, status_code, message):
+        self.status_code = status_code
+        self.message = message
+        self.request = httpx.Request(
+            method="POST", url=" https://cloud.google.com/vertex-ai/"
+        )
+        self.response = httpx.Response(status_code=status_code, request=self.request)
+        super().__init__(
+            self.message
+        )  # Call the base class constructor with the parameters it needs
 
 
 class VertexAIAnthropicConfig:
@@ -27,7 +56,6 @@ class VertexAIAnthropicConfig:
     """
 
     max_tokens: Optional[int] = litellm.max_tokens
-    anthropic_version: Optional[str] = "bedrock-2023-05-31"
     system: Optional[str] = None
     temperature: Optional[float] = None
     top_p: Optional[float] = None
@@ -88,3 +116,154 @@ class VertexAIAnthropicConfig:
             if param == "top_p":
                 optional_params["top_p"] = value
         return optional_params
+
+
+"""
+- Run client init 
+- Support async completion, streaming
+"""
+
+
+# makes headers for API call
+
+
+def completion(
+    model: str,
+    messages: list,
+    model_response: ModelResponse,
+    print_verbose: Callable,
+    encoding,
+    logging_obj,
+    vertex_project=None,
+    vertex_location=None,
+    optional_params=None,
+    litellm_params=None,
+    logger_fn=None,
+    acompletion: bool = False,
+    client=None,
+):
+    try:
+        import vertexai
+    except:
+        raise VertexAIError(
+            status_code=400,
+            message="""vertexai import failed please run `pip install -U google-cloud-aiplatform "anthropic[vertex]"`""",
+        )
+
+    if not (
+        hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")
+    ):
+        raise VertexAIError(
+            status_code=400,
+            message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
+        )
+    try:
+        import google.auth  # type: ignore
+        from google.auth.transport.requests import Request
+        from anthropic import AnthropicVertex
+
+        ## Load Config
+        config = litellm.VertexAIAnthropicConfig.get_config()
+        for k, v in config.items():
+            if k not in optional_params:
+                optional_params[k] = v
+
+        ## Format Prompt
+        _is_function_call = False
+        messages = copy.deepcopy(messages)
+        optional_params = copy.deepcopy(optional_params)
+        # Separate system prompt from rest of message
+        system_prompt_indices = []
+        system_prompt = ""
+        for idx, message in enumerate(messages):
+            if message["role"] == "system":
+                system_prompt += message["content"]
+                system_prompt_indices.append(idx)
+        if len(system_prompt_indices) > 0:
+            for idx in reversed(system_prompt_indices):
+                messages.pop(idx)
+        if len(system_prompt) > 0:
+            optional_params["system"] = system_prompt
+        # Format rest of message according to anthropic guidelines
+        try:
+            messages = prompt_factory(
+                model=model, messages=messages, custom_llm_provider="anthropic"
+            )
+        except Exception as e:
+            raise VertexAIError(status_code=400, message=str(e))
+
+        ## Handle Tool Calling
+        if "tools" in optional_params:
+            _is_function_call = True
+            tool_calling_system_prompt = construct_tool_use_system_prompt(
+                tools=optional_params["tools"]
+            )
+            optional_params["system"] = (
+                optional_params.get("system", "\n") + tool_calling_system_prompt
+            )  # add the anthropic tool calling prompt to the system prompt
+            optional_params.pop("tools")
+
+        stream = optional_params.pop("stream", None)
+
+        data = {
+            "model": model,
+            "messages": messages,
+            **optional_params,
+        }
+        print_verbose(f"_is_function_call: {_is_function_call}")
+
+        ## Completion Call
+
+        print_verbose(
+            f"VERTEX AI: vertex_project={vertex_project}; vertex_location={vertex_location}"
+        )
+        if client is None:
+            vertex_ai_client = AnthropicVertex(
+                project_id=vertex_project, region=vertex_location
+            )
+        else:
+            vertex_ai_client = client
+
+        message = vertex_ai_client.messages.create(**data)  # type: ignore
+        text_content = message.content[0].text
+        ## TOOL CALLING - OUTPUT PARSE
+        if text_content is not None and contains_tag("invoke", text_content):
+            function_name = extract_between_tags("tool_name", text_content)[0]
+            function_arguments_str = extract_between_tags("invoke", text_content)[
+                0
+            ].strip()
+            function_arguments_str = f"<invoke>{function_arguments_str}</invoke>"
+            function_arguments = parse_xml_params(function_arguments_str)
+            _message = litellm.Message(
+                tool_calls=[
+                    {
+                        "id": f"call_{uuid.uuid4()}",
+                        "type": "function",
+                        "function": {
+                            "name": function_name,
+                            "arguments": json.dumps(function_arguments),
+                        },
+                    }
+                ],
+                content=None,
+            )
+            model_response.choices[0].message = _message  # type: ignore
+        else:
+            model_response.choices[0].message.content = text_content  # type: ignore
+        model_response.choices[0].finish_reason = map_finish_reason(message.stop_reason)
+
+        ## CALCULATING USAGE
+        prompt_tokens = message.usage.input_tokens
+        completion_tokens = message.usage.output_tokens
+
+        model_response["created"] = int(time.time())
+        model_response["model"] = model
+        usage = Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+        )
+        model_response.usage = usage
+        return model_response
+    except Exception as e:
+        raise VertexAIError(status_code=500, message=str(e))

litellm/main.py

@@ -62,6 +62,7 @@ from .llms import (
     palm,
     gemini,
     vertex_ai,
+    vertex_ai_anthropic,
     maritalk,
 )
 from .llms.openai import OpenAIChatCompletion, OpenAITextCompletion
@@ -1627,20 +1628,36 @@ def completion(
                 or get_secret("VERTEXAI_LOCATION")
             )
 
-            model_response = vertex_ai.completion(
+            if "claude-3" in model:
-                model=model,
+                model_response = vertex_ai_anthropic.completion(
-                messages=messages,
+                    model=model,
-                model_response=model_response,
+                    messages=messages,
-                print_verbose=print_verbose,
+                    model_response=model_response,
-                optional_params=optional_params,
+                    print_verbose=print_verbose,
-                litellm_params=litellm_params,
+                    optional_params=optional_params,
-                logger_fn=logger_fn,
+                    litellm_params=litellm_params,
-                encoding=encoding,
+                    logger_fn=logger_fn,
-                vertex_location=vertex_ai_location,
+                    encoding=encoding,
-                vertex_project=vertex_ai_project,
+                    vertex_location=vertex_ai_location,
-                logging_obj=logging,
+                    vertex_project=vertex_ai_project,
-                acompletion=acompletion,
+                    logging_obj=logging,
-            )
+                    acompletion=acompletion,
+                )
+            else:
+                model_response = vertex_ai.completion(
+                    model=model,
+                    messages=messages,
+                    model_response=model_response,
+                    print_verbose=print_verbose,
+                    optional_params=optional_params,
+                    litellm_params=litellm_params,
+                    logger_fn=logger_fn,
+                    encoding=encoding,
+                    vertex_location=vertex_ai_location,
+                    vertex_project=vertex_ai_project,
+                    logging_obj=logging,
+                    acompletion=acompletion,
+                )
 
             if (
                 "stream" in optional_params

litellm/tests/test_amazing_vertex_completion.py

@@ -84,6 +84,24 @@ async def get_response():
         pytest.fail(f"An error occurred - {str(e)}")
 
 
+def test_vertex_ai_anthropic():
+    load_vertex_ai_credentials()
+
+    model = "claude-3-sonnet@20240229"
+
+    vertex_ai_project = "adroit-crow-413218"
+    vertex_ai_location = "asia-southeast1"
+
+    response = completion(
+        model="vertex_ai/" + model,
+        messages=[{"role": "user", "content": "hi"}],
+        temperature=0.7,
+        vertex_ai_project=vertex_ai_project,
+        vertex_ai_location=vertex_ai_location,
+    )
+    print("\nModel Response", response)
+
+
 def test_vertex_ai():
     import random
 

litellm/tests/vertex_key.json

@@ -1,13 +1,13 @@
 {
   "type": "service_account",
-  "project_id": "reliablekeys",
+  "project_id": "adroit-crow-413218",
   "private_key_id": "",
   "private_key": "",
-  "client_email": "73470430121-compute@developer.gserviceaccount.com",
+  "client_email": "test-adroit-crow@adroit-crow-413218.iam.gserviceaccount.com",
-  "client_id": "108560959659377334173",
+  "client_id": "104886546564708740969",
   "auth_uri": "https://accounts.google.com/o/oauth2/auth",
   "token_uri": "https://oauth2.googleapis.com/token",
   "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
-  "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/73470430121-compute%40developer.gserviceaccount.com",
+  "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-adroit-crow%40adroit-crow-413218.iam.gserviceaccount.com",
   "universe_domain": "googleapis.com"
-}
+}