fix - use anthropic class for clients

This commit is contained in:
Ishaan Jaff 2024-04-06 18:19:28 -07:00
parent 9be6b7ec7c
commit fcf5aa278b
2 changed files with 389 additions and 378 deletions

View file

@ -8,7 +8,7 @@ from litellm.utils import ModelResponse, Usage, map_finish_reason, CustomStreamW
import litellm import litellm
from .prompt_templates.factory import prompt_factory, custom_prompt from .prompt_templates.factory import prompt_factory, custom_prompt
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from .base import BaseLLM
import httpx import httpx
@ -19,9 +19,6 @@ class AnthropicConstants(Enum):
# constants from https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/_constants.py # constants from https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/_constants.py
async_handler = AsyncHTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0))
class AnthropicError(Exception): class AnthropicError(Exception):
def __init__(self, status_code, message): def __init__(self, status_code, message):
self.status_code = status_code self.status_code = status_code
@ -105,7 +102,15 @@ def validate_environment(api_key, user_headers):
return headers return headers
class AnthropicChatCompletion(BaseLLM):
def __init__(self) -> None:
super().__init__()
self.async_handler = AsyncHTTPHandler(
timeout=httpx.Timeout(timeout=600.0, connect=5.0)
)
def process_response( def process_response(
self,
model, model,
response, response,
model_response, model_response,
@ -129,7 +134,9 @@ def process_response(
try: try:
completion_response = response.json() completion_response = response.json()
except: except:
raise AnthropicError(message=response.text, status_code=response.status_code) raise AnthropicError(
message=response.text, status_code=response.status_code
)
if "error" in completion_response: if "error" in completion_response:
raise AnthropicError( raise AnthropicError(
message=str(completion_response["error"]), message=str(completion_response["error"]),
@ -232,8 +239,8 @@ def process_response(
model_response.usage = usage model_response.usage = usage
return model_response return model_response
async def acompletion_stream_function( async def acompletion_stream_function(
self,
model: str, model: str,
messages: list, messages: list,
api_base: str, api_base: str,
@ -251,12 +258,14 @@ async def acompletion_stream_function(
logger_fn=None, logger_fn=None,
headers={}, headers={},
): ):
response = await async_handler.post( response = await self.async_handler.post(
api_base, headers=headers, data=json.dumps(data) api_base, headers=headers, data=json.dumps(data)
) )
if response.status_code != 200: if response.status_code != 200:
raise AnthropicError(status_code=response.status_code, message=response.text) raise AnthropicError(
status_code=response.status_code, message=response.text
)
completion_stream = response.aiter_lines() completion_stream = response.aiter_lines()
@ -268,8 +277,8 @@ async def acompletion_stream_function(
) )
return streamwrapper return streamwrapper
async def acompletion_function( async def acompletion_function(
self,
model: str, model: str,
messages: list, messages: list,
api_base: str, api_base: str,
@ -287,10 +296,10 @@ async def acompletion_function(
logger_fn=None, logger_fn=None,
headers={}, headers={},
): ):
response = await async_handler.post( response = await self.async_handler.post(
api_base, headers=headers, data=json.dumps(data) api_base, headers=headers, data=json.dumps(data)
) )
return process_response( return self.process_response(
model=model, model=model,
response=response, response=response,
model_response=model_response, model_response=model_response,
@ -303,8 +312,8 @@ async def acompletion_function(
print_verbose=print_verbose, print_verbose=print_verbose,
) )
def completion( def completion(
self,
model: str, model: str,
messages: list, messages: list,
api_base: str, api_base: str,
@ -400,7 +409,7 @@ def completion(
): # if function call - fake the streaming (need complete blocks for output parsing in openai format) ): # if function call - fake the streaming (need complete blocks for output parsing in openai format)
print_verbose("makes async anthropic streaming POST request") print_verbose("makes async anthropic streaming POST request")
data["stream"] = stream data["stream"] = stream
return acompletion_stream_function( return self.acompletion_stream_function(
model=model, model=model,
messages=messages, messages=messages,
data=data, data=data,
@ -419,7 +428,7 @@ def completion(
headers=headers, headers=headers,
) )
else: else:
return acompletion_function( return self.acompletion_function(
model=model, model=model,
messages=messages, messages=messages,
data=data, data=data,
@ -466,12 +475,14 @@ def completion(
return streaming_response return streaming_response
else: else:
response = requests.post(api_base, headers=headers, data=json.dumps(data)) response = requests.post(
api_base, headers=headers, data=json.dumps(data)
)
if response.status_code != 200: if response.status_code != 200:
raise AnthropicError( raise AnthropicError(
status_code=response.status_code, message=response.text status_code=response.status_code, message=response.text
) )
return process_response( return self.process_response(
model=model, model=model,
response=response, response=response,
model_response=model_response, model_response=model_response,
@ -484,6 +495,10 @@ def completion(
print_verbose=print_verbose, print_verbose=print_verbose,
) )
def embedding(self):
# logic for parsing in - calling - parsing out model embedding calls
pass
class ModelResponseIterator: class ModelResponseIterator:
def __init__(self, model_response): def __init__(self, model_response):
@ -509,8 +524,3 @@ class ModelResponseIterator:
raise StopAsyncIteration raise StopAsyncIteration
self.is_done = True self.is_done = True
return self.model_response return self.model_response
def embedding():
# logic for parsing in - calling - parsing out model embedding calls
pass

View file

@ -39,7 +39,6 @@ from litellm.utils import (
get_optional_params_image_gen, get_optional_params_image_gen,
) )
from .llms import ( from .llms import (
anthropic,
anthropic_text, anthropic_text,
together_ai, together_ai,
ai21, ai21,
@ -68,6 +67,7 @@ from .llms import (
from .llms.openai import OpenAIChatCompletion, OpenAITextCompletion from .llms.openai import OpenAIChatCompletion, OpenAITextCompletion
from .llms.azure import AzureChatCompletion from .llms.azure import AzureChatCompletion
from .llms.azure_text import AzureTextCompletion from .llms.azure_text import AzureTextCompletion
from .llms.anthropic import AnthropicChatCompletion
from .llms.huggingface_restapi import Huggingface from .llms.huggingface_restapi import Huggingface
from .llms.prompt_templates.factory import ( from .llms.prompt_templates.factory import (
prompt_factory, prompt_factory,
@ -99,6 +99,7 @@ from litellm.utils import (
dotenv.load_dotenv() # Loading env variables using dotenv dotenv.load_dotenv() # Loading env variables using dotenv
openai_chat_completions = OpenAIChatCompletion() openai_chat_completions = OpenAIChatCompletion()
openai_text_completions = OpenAITextCompletion() openai_text_completions = OpenAITextCompletion()
anthropic_chat_completions = AnthropicChatCompletion()
azure_chat_completions = AzureChatCompletion() azure_chat_completions = AzureChatCompletion()
azure_text_completions = AzureTextCompletion() azure_text_completions = AzureTextCompletion()
huggingface = Huggingface() huggingface = Huggingface()
@ -1181,7 +1182,7 @@ def completion(
or get_secret("ANTHROPIC_API_BASE") or get_secret("ANTHROPIC_API_BASE")
or "https://api.anthropic.com/v1/messages" or "https://api.anthropic.com/v1/messages"
) )
response = anthropic.completion( response = anthropic_chat_completions.completion(
model=model, model=model,
messages=messages, messages=messages,
api_base=api_base, api_base=api_base,