File size: 2,408 Bytes
7bf4b88 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import time
import anthropic
from typing import Any, Union
def complete_text_claude(message: Union[str, list],
model: str = "claude-2.1",
json_object: bool = False,
max_tokens: int = 2048,
temperature: float = 1.0,
max_retry: int = 1,
sleep_time: int = 0,
tools: list = [],
**kwargs: Any) -> str:
"""
Call the Claude API to complete a prompt.
Args:
message (Union[str, list]): The input message or a list of message dicts.
model (str): The model to use for completion. Default is "claude-2.1".
json_object (bool): Whether to output in JSON format. Default is False.
max_tokens (int): Maximum number of tokens to generate. Default is 2048.
temperature (float): Sampling temperature. Default is 1.0.
max_retry (int): Maximum number of retries in case of an error. Default is 1.
sleep_time (int): Sleep time between retries in seconds. Default is 0.
tools (list): List of tools to use for the completion. Default is an empty list.
**kwargs (Any): Additional keyword arguments to pass to the API.
Returns:
str: The completed text generated by the Claude model.
Raises:
Exception: If the completion fails after the maximum number of retries.
"""
anthropic_client = anthropic.Anthropic()
if isinstance(message, str):
if json_object:
message = "You are a helpful assistant designed to output in JSON format." + message
messages = [{"role": "user", "content": message}]
else:
messages = message
for cnt in range(max_retry):
try:
response = anthropic_client.beta.tools.messages.create(
messages=messages,
model=model,
temperature=temperature,
max_tokens=max_tokens,
tools=tools,
**kwargs
)
completion = response.to_dict()
return completion["content"][0]['text']
except Exception as e:
print(f"Attempt {cnt} failed: {e}. Retrying after {sleep_time} seconds...")
time.sleep(sleep_time)
raise Exception("Failed to complete text after maximum retries")
|