import time
from stark_qa.tools.api_lib import anthropic_client
from typing import Any, Union
[docs]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.
"""
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")