|
|
|
@@ -2,21 +2,32 @@ import base64 |
|
|
|
import json |
|
|
|
import logging |
|
|
|
from collections.abc import Generator |
|
|
|
from typing import Optional, Union |
|
|
|
from typing import Optional, Union, cast |
|
|
|
|
|
|
|
import google.api_core.exceptions as exceptions |
|
|
|
import vertexai.generative_models as glm |
|
|
|
from anthropic import AnthropicVertex, Stream |
|
|
|
from anthropic.types import ( |
|
|
|
ContentBlockDeltaEvent, |
|
|
|
Message, |
|
|
|
MessageDeltaEvent, |
|
|
|
MessageStartEvent, |
|
|
|
MessageStopEvent, |
|
|
|
MessageStreamEvent, |
|
|
|
) |
|
|
|
from google.cloud import aiplatform |
|
|
|
from google.oauth2 import service_account |
|
|
|
from vertexai.generative_models import HarmBlockThreshold, HarmCategory |
|
|
|
|
|
|
|
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta |
|
|
|
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage |
|
|
|
from core.model_runtime.entities.message_entities import ( |
|
|
|
AssistantPromptMessage, |
|
|
|
ImagePromptMessageContent, |
|
|
|
PromptMessage, |
|
|
|
PromptMessageContentType, |
|
|
|
PromptMessageTool, |
|
|
|
SystemPromptMessage, |
|
|
|
TextPromptMessageContent, |
|
|
|
ToolPromptMessage, |
|
|
|
UserPromptMessage, |
|
|
|
) |
|
|
|
@@ -63,9 +74,287 @@ class VertexAiLargeLanguageModel(LargeLanguageModel): |
|
|
|
:param user: unique user id |
|
|
|
:return: full response or stream response chunk generator result |
|
|
|
""" |
|
|
|
# invoke model |
|
|
|
# invoke anthropic models via anthropic official SDK |
|
|
|
if "claude" in model: |
|
|
|
return self._generate_anthropic(model, credentials, prompt_messages, model_parameters, stop, stream, user) |
|
|
|
# invoke Gemini model |
|
|
|
return self._generate(model, credentials, prompt_messages, model_parameters, tools, stop, stream, user) |
|
|
|
|
|
|
|
|
|
|
|
def _generate_anthropic(self, model: str, credentials: dict, prompt_messages: list[PromptMessage], model_parameters: dict, |
|
|
|
stop: Optional[list[str]] = None, stream: bool = True, user: Optional[str] = None) -> Union[LLMResult, Generator]: |
|
|
|
""" |
|
|
|
Invoke Anthropic large language model |
|
|
|
|
|
|
|
:param model: model name |
|
|
|
:param credentials: model credentials |
|
|
|
:param prompt_messages: prompt messages |
|
|
|
:param model_parameters: model parameters |
|
|
|
:param stop: stop words |
|
|
|
:param stream: is stream response |
|
|
|
:return: full response or stream response chunk generator result |
|
|
|
""" |
|
|
|
# use Anthropic official SDK references |
|
|
|
# - https://github.com/anthropics/anthropic-sdk-python |
|
|
|
project_id = credentials["vertex_project_id"] |
|
|
|
|
|
|
|
if 'opus' in model: |
|
|
|
location = 'us-east5' |
|
|
|
else: |
|
|
|
location = 'us-central1' |
|
|
|
|
|
|
|
client = AnthropicVertex( |
|
|
|
region=location, |
|
|
|
project_id=project_id |
|
|
|
) |
|
|
|
|
|
|
|
extra_model_kwargs = {} |
|
|
|
if stop: |
|
|
|
extra_model_kwargs['stop_sequences'] = stop |
|
|
|
|
|
|
|
system, prompt_message_dicts = self._convert_claude_prompt_messages(prompt_messages) |
|
|
|
|
|
|
|
if system: |
|
|
|
extra_model_kwargs['system'] = system |
|
|
|
|
|
|
|
response = client.messages.create( |
|
|
|
model=model, |
|
|
|
messages=prompt_message_dicts, |
|
|
|
stream=stream, |
|
|
|
**model_parameters, |
|
|
|
**extra_model_kwargs |
|
|
|
) |
|
|
|
|
|
|
|
if stream: |
|
|
|
return self._handle_claude_stream_response(model, credentials, response, prompt_messages) |
|
|
|
|
|
|
|
return self._handle_claude_response(model, credentials, response, prompt_messages) |
|
|
|
|
|
|
|
def _handle_claude_response(self, model: str, credentials: dict, response: Message, |
|
|
|
prompt_messages: list[PromptMessage]) -> LLMResult: |
|
|
|
""" |
|
|
|
Handle llm chat response |
|
|
|
|
|
|
|
:param model: model name |
|
|
|
:param credentials: credentials |
|
|
|
:param response: response |
|
|
|
:param prompt_messages: prompt messages |
|
|
|
:return: full response chunk generator result |
|
|
|
""" |
|
|
|
|
|
|
|
# transform assistant message to prompt message |
|
|
|
assistant_prompt_message = AssistantPromptMessage( |
|
|
|
content=response.content[0].text |
|
|
|
) |
|
|
|
|
|
|
|
# calculate num tokens |
|
|
|
if response.usage: |
|
|
|
# transform usage |
|
|
|
prompt_tokens = response.usage.input_tokens |
|
|
|
completion_tokens = response.usage.output_tokens |
|
|
|
else: |
|
|
|
# calculate num tokens |
|
|
|
prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages) |
|
|
|
completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message]) |
|
|
|
|
|
|
|
# transform usage |
|
|
|
usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens) |
|
|
|
|
|
|
|
# transform response |
|
|
|
response = LLMResult( |
|
|
|
model=response.model, |
|
|
|
prompt_messages=prompt_messages, |
|
|
|
message=assistant_prompt_message, |
|
|
|
usage=usage |
|
|
|
) |
|
|
|
|
|
|
|
return response |
|
|
|
|
|
|
|
def _handle_claude_stream_response(self, model: str, credentials: dict, response: Stream[MessageStreamEvent], |
|
|
|
prompt_messages: list[PromptMessage], ) -> Generator: |
|
|
|
""" |
|
|
|
Handle llm chat stream response |
|
|
|
|
|
|
|
:param model: model name |
|
|
|
:param credentials: credentials |
|
|
|
:param response: response |
|
|
|
:param prompt_messages: prompt messages |
|
|
|
:return: full response or stream response chunk generator result |
|
|
|
""" |
|
|
|
|
|
|
|
try: |
|
|
|
full_assistant_content = '' |
|
|
|
return_model = None |
|
|
|
input_tokens = 0 |
|
|
|
output_tokens = 0 |
|
|
|
finish_reason = None |
|
|
|
index = 0 |
|
|
|
|
|
|
|
for chunk in response: |
|
|
|
if isinstance(chunk, MessageStartEvent): |
|
|
|
return_model = chunk.message.model |
|
|
|
input_tokens = chunk.message.usage.input_tokens |
|
|
|
elif isinstance(chunk, MessageDeltaEvent): |
|
|
|
output_tokens = chunk.usage.output_tokens |
|
|
|
finish_reason = chunk.delta.stop_reason |
|
|
|
elif isinstance(chunk, MessageStopEvent): |
|
|
|
usage = self._calc_response_usage(model, credentials, input_tokens, output_tokens) |
|
|
|
yield LLMResultChunk( |
|
|
|
model=return_model, |
|
|
|
prompt_messages=prompt_messages, |
|
|
|
delta=LLMResultChunkDelta( |
|
|
|
index=index + 1, |
|
|
|
message=AssistantPromptMessage( |
|
|
|
content='' |
|
|
|
), |
|
|
|
finish_reason=finish_reason, |
|
|
|
usage=usage |
|
|
|
) |
|
|
|
) |
|
|
|
elif isinstance(chunk, ContentBlockDeltaEvent): |
|
|
|
chunk_text = chunk.delta.text if chunk.delta.text else '' |
|
|
|
full_assistant_content += chunk_text |
|
|
|
assistant_prompt_message = AssistantPromptMessage( |
|
|
|
content=chunk_text if chunk_text else '', |
|
|
|
) |
|
|
|
index = chunk.index |
|
|
|
yield LLMResultChunk( |
|
|
|
model=model, |
|
|
|
prompt_messages=prompt_messages, |
|
|
|
delta=LLMResultChunkDelta( |
|
|
|
index=index, |
|
|
|
message=assistant_prompt_message, |
|
|
|
) |
|
|
|
) |
|
|
|
except Exception as ex: |
|
|
|
raise InvokeError(str(ex)) |
|
|
|
|
|
|
|
def _calc_claude_response_usage(self, model: str, credentials: dict, prompt_tokens: int, completion_tokens: int) -> LLMUsage: |
|
|
|
""" |
|
|
|
Calculate response usage |
|
|
|
|
|
|
|
:param model: model name |
|
|
|
:param credentials: model credentials |
|
|
|
:param prompt_tokens: prompt tokens |
|
|
|
:param completion_tokens: completion tokens |
|
|
|
:return: usage |
|
|
|
""" |
|
|
|
# get prompt price info |
|
|
|
prompt_price_info = self.get_price( |
|
|
|
model=model, |
|
|
|
credentials=credentials, |
|
|
|
price_type=PriceType.INPUT, |
|
|
|
tokens=prompt_tokens, |
|
|
|
) |
|
|
|
|
|
|
|
# get completion price info |
|
|
|
completion_price_info = self.get_price( |
|
|
|
model=model, |
|
|
|
credentials=credentials, |
|
|
|
price_type=PriceType.OUTPUT, |
|
|
|
tokens=completion_tokens |
|
|
|
) |
|
|
|
|
|
|
|
# transform usage |
|
|
|
usage = LLMUsage( |
|
|
|
prompt_tokens=prompt_tokens, |
|
|
|
prompt_unit_price=prompt_price_info.unit_price, |
|
|
|
prompt_price_unit=prompt_price_info.unit, |
|
|
|
prompt_price=prompt_price_info.total_amount, |
|
|
|
completion_tokens=completion_tokens, |
|
|
|
completion_unit_price=completion_price_info.unit_price, |
|
|
|
completion_price_unit=completion_price_info.unit, |
|
|
|
completion_price=completion_price_info.total_amount, |
|
|
|
total_tokens=prompt_tokens + completion_tokens, |
|
|
|
total_price=prompt_price_info.total_amount + completion_price_info.total_amount, |
|
|
|
currency=prompt_price_info.currency, |
|
|
|
latency=time.perf_counter() - self.started_at |
|
|
|
) |
|
|
|
|
|
|
|
return usage |
|
|
|
|
|
|
|
def _convert_claude_prompt_messages(self, prompt_messages: list[PromptMessage]) -> tuple[str, list[dict]]: |
|
|
|
""" |
|
|
|
Convert prompt messages to dict list and system |
|
|
|
""" |
|
|
|
|
|
|
|
system = "" |
|
|
|
first_loop = True |
|
|
|
for message in prompt_messages: |
|
|
|
if isinstance(message, SystemPromptMessage): |
|
|
|
message.content=message.content.strip() |
|
|
|
if first_loop: |
|
|
|
system=message.content |
|
|
|
first_loop=False |
|
|
|
else: |
|
|
|
system+="\n" |
|
|
|
system+=message.content |
|
|
|
|
|
|
|
prompt_message_dicts = [] |
|
|
|
for message in prompt_messages: |
|
|
|
if not isinstance(message, SystemPromptMessage): |
|
|
|
prompt_message_dicts.append(self._convert_claude_prompt_message_to_dict(message)) |
|
|
|
|
|
|
|
return system, prompt_message_dicts |
|
|
|
|
|
|
|
def _convert_claude_prompt_message_to_dict(self, message: PromptMessage) -> dict: |
|
|
|
""" |
|
|
|
Convert PromptMessage to dict |
|
|
|
""" |
|
|
|
if isinstance(message, UserPromptMessage): |
|
|
|
message = cast(UserPromptMessage, message) |
|
|
|
if isinstance(message.content, str): |
|
|
|
message_dict = {"role": "user", "content": message.content} |
|
|
|
else: |
|
|
|
sub_messages = [] |
|
|
|
for message_content in message.content: |
|
|
|
if message_content.type == PromptMessageContentType.TEXT: |
|
|
|
message_content = cast(TextPromptMessageContent, message_content) |
|
|
|
sub_message_dict = { |
|
|
|
"type": "text", |
|
|
|
"text": message_content.data |
|
|
|
} |
|
|
|
sub_messages.append(sub_message_dict) |
|
|
|
elif message_content.type == PromptMessageContentType.IMAGE: |
|
|
|
message_content = cast(ImagePromptMessageContent, message_content) |
|
|
|
if not message_content.data.startswith("data:"): |
|
|
|
# fetch image data from url |
|
|
|
try: |
|
|
|
image_content = requests.get(message_content.data).content |
|
|
|
mime_type, _ = mimetypes.guess_type(message_content.data) |
|
|
|
base64_data = base64.b64encode(image_content).decode('utf-8') |
|
|
|
except Exception as ex: |
|
|
|
raise ValueError(f"Failed to fetch image data from url {message_content.data}, {ex}") |
|
|
|
else: |
|
|
|
data_split = message_content.data.split(";base64,") |
|
|
|
mime_type = data_split[0].replace("data:", "") |
|
|
|
base64_data = data_split[1] |
|
|
|
|
|
|
|
if mime_type not in ["image/jpeg", "image/png", "image/gif", "image/webp"]: |
|
|
|
raise ValueError(f"Unsupported image type {mime_type}, " |
|
|
|
f"only support image/jpeg, image/png, image/gif, and image/webp") |
|
|
|
|
|
|
|
sub_message_dict = { |
|
|
|
"type": "image", |
|
|
|
"source": { |
|
|
|
"type": "base64", |
|
|
|
"media_type": mime_type, |
|
|
|
"data": base64_data |
|
|
|
} |
|
|
|
} |
|
|
|
sub_messages.append(sub_message_dict) |
|
|
|
|
|
|
|
message_dict = {"role": "user", "content": sub_messages} |
|
|
|
elif isinstance(message, AssistantPromptMessage): |
|
|
|
message = cast(AssistantPromptMessage, message) |
|
|
|
message_dict = {"role": "assistant", "content": message.content} |
|
|
|
elif isinstance(message, SystemPromptMessage): |
|
|
|
message = cast(SystemPromptMessage, message) |
|
|
|
message_dict = {"role": "system", "content": message.content} |
|
|
|
else: |
|
|
|
raise ValueError(f"Got unknown type {message}") |
|
|
|
|
|
|
|
return message_dict |
|
|
|
|
|
|
|
def get_num_tokens(self, model: str, credentials: dict, prompt_messages: list[PromptMessage], |
|
|
|
tools: Optional[list[PromptMessageTool]] = None) -> int: |
|
|
|
""" |
Pan, Wen-Ming
1 year ago