xintongsong commented on code in PR #139: URL: https://github.com/apache/flink-agents/pull/139#discussion_r2324091508
########## python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py: ########## @@ -0,0 +1,245 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +import uuid +from typing import Any, Dict, List, Optional, Sequence + +from anthropic import Anthropic +from anthropic._types import NOT_GIVEN +from anthropic.types import MessageParam, TextBlockParam, ToolParam +from pydantic import Field, PrivateAttr + +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_models.chat_model import ( + BaseChatModelConnection, + BaseChatModelSetup, +) +from flink_agents.api.tools.tool import BaseTool, ToolMetadata + + +def to_anthropic_tool(*, metadata: ToolMetadata, skip_length_check: bool = False) -> ToolParam: + """Convert to Anthropic tool: https://docs.anthropic.com/en/api/messages#body-tools."""; + if not skip_length_check and len(metadata.description) > 1024: + msg = ( + "Tool description exceeds maximum length of 1024 characters. " + "Please shorten your description or move it to the prompt." + ) + raise ValueError(msg) + return { + "name": metadata.name, + "description": metadata.description, + "input_schema": metadata.get_parameters_dict() + } + + +def convert_to_anthropic_message(message: ChatMessage) -> MessageParam: + """Convert ChatMessage to Anthropic MessageParam format.""" + if message.role == MessageRole.TOOL: + return { + "role": MessageRole.USER.value, + "content": [ + { + "type": "tool_result", + "tool_use_id": message.extra_args.get("external_id"), + "content": message.content, + } + ], + } + else: + return { + "role": message.role.value, + "content": message.content, + } + + +def convert_to_anthropic_messages(messages: Sequence[ChatMessage]) -> List[MessageParam]: + """Convert user/assistant messages to Anthropic input messages. + + See: https://docs.anthropic.com/en/api/messages#body-messages + """ + return [convert_to_anthropic_message(message) for message in messages if + message.role in [MessageRole.USER, MessageRole.ASSISTANT, MessageRole.TOOL]] + + +def convert_to_anthropic_system_prompts(messages: Sequence[ChatMessage]) -> List[TextBlockParam]: + """Convert system messages to Anthropic system prompts. + + See: https://docs.anthropic.com/en/api/messages#body-system + """ + system_messages = [message for message in messages if message.role == MessageRole.SYSTEM] + return [ + TextBlockParam( + type="text", + text=message.content + ) + for message in system_messages + ] + + +class AnthropicChatModelConnection(BaseChatModelConnection): + """Manages the connection to the Anthropic AI models for chat interactions. + + Attributes: + ---------- + api_key : str + The Anthropic API key. + max_retries : int + The number of times to retry the API call upon failure. + timeout : float + The number of seconds to wait for an API call before it times out. + reuse_client : bool + Whether to reuse the Anthropic client between requests. + """ + + api_key: str = Field(default=None, description="The Anthropic API key.") + + max_retries: int = Field( + default=3, + description="The number of times to retry the API call upon failure.", + ge=0, + ) + timeout: float = Field( + default=60.0, + description="The number of seconds to wait for an API call before it times out.", + ge=0, + ) + + def __init__( + self, + api_key: Optional[str] = None, + max_retries: int = 3, + timeout: float = 60.0, + **kwargs: Any, + ) -> None: + """Initialize the Anthropic chat model connection.""" + super().__init__( + api_key=api_key, + max_retries=max_retries, + timeout=timeout, + **kwargs, + ) + + _client: Optional[Anthropic] = PrivateAttr(default=None) + + @property + def client(self) -> Anthropic: + """Get or create the Anthropic client instance.""" + if self._client is None: + self._client = Anthropic(api_key=self.api_key, max_retries=self.max_retries, timeout=self.timeout) + return self._client + + def chat(self, messages: Sequence[ChatMessage], tools: Optional[List[BaseTool]] = None, + **kwargs: Any) -> ChatMessage: + """Direct communication with Anthropic model service for chat conversation.""" + anthropic_tools = None + if tools is not None: + anthropic_tools = [to_anthropic_tool(metadata=tool.metadata) for tool in tools] + + anthropic_system = convert_to_anthropic_system_prompts(messages) + anthropic_messages = convert_to_anthropic_messages(messages) + + message = self.client.messages.create( + messages=anthropic_messages, + tools=anthropic_tools or NOT_GIVEN, + system=anthropic_system or NOT_GIVEN, + **kwargs, + ) + + if message.stop_reason == "tool_use": + tool_calls = [ + { + "id": uuid.uuid4(), + "type": "function", + "function": { + "name": content_block.name, + "arguments": content_block.input, + }, + "original_id": content_block.id, + } + for content_block in message.content + if content_block.type == 'tool_use' + ] + + return ChatMessage( + role=MessageRole(message.role), + content=message.content, Review Comment: I see. In that case, I'd suggest to put the context in `ChatMessage.extra_args`, rather than the `content`. `role`, `content` and `tool_calls` are common fields of chat message that we should not assume only the anthropic chat model will access them. On the other hand, `extra_args` is a dict that accept arbitrary key-values, which is probably more suitable for storing implementation-specific information. WDYT? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: issues-unsubscr...@flink.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
