Spaces:
Sleeping
Sleeping
File size: 4,075 Bytes
43c34cc |
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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import os
import re
from typing import List
from tenacity import retry, stop_after_attempt, wait_random_exponential
from ..message import SYSTEM_NAME as SYSTEM
from ..message import Message
from .base import IntelligenceBackend
try:
import anthropic
except ImportError:
is_anthropic_available = False
# logging.warning("anthropic package is not installed")
else:
anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
if anthropic_api_key is None:
# logging.warning("Anthropic API key is not set. Please set the environment variable ANTHROPIC_API_KEY")
is_anthropic_available = False
else:
is_anthropic_available = True
DEFAULT_MAX_TOKENS = 256
DEFAULT_MODEL = "claude-v1"
class Claude(IntelligenceBackend):
"""Interface to the Claude offered by Anthropic."""
stateful = False
type_name = "claude"
def __init__(
self, max_tokens: int = DEFAULT_MAX_TOKENS, model: str = DEFAULT_MODEL, **kwargs
):
assert (
is_anthropic_available
), "anthropic package is not installed or the API key is not set"
super().__init__(max_tokens=max_tokens, model=model, **kwargs)
self.max_tokens = max_tokens
self.model = model
self.client = anthropic.Client(os.environ["ANTHROPIC_API_KEY"])
@retry(stop=stop_after_attempt(6), wait=wait_random_exponential(min=1, max=60))
def _get_response(self, prompt: str):
response = self.client.completion(
prompt=prompt,
stop_sequences=[anthropic.HUMAN_PROMPT],
model=self.model,
max_tokens_to_sample=self.max_tokens,
)
response = response["completion"].strip()
return response
def query(
self,
agent_name: str,
role_desc: str,
history_messages: List[Message],
global_prompt: str = None,
request_msg: Message = None,
*args,
**kwargs,
) -> str:
"""
Format the input and call the Claude API.
args:
agent_name: the name of the agent
role_desc: the description of the role of the agent
env_desc: the description of the environment
history_messages: the history of the conversation, or the observation for the agent
request_msg: the request from the system to guide the agent's next response
"""
all_messages = (
[(SYSTEM, global_prompt), (SYSTEM, role_desc)]
if global_prompt
else [(SYSTEM, role_desc)]
)
for message in history_messages:
all_messages.append((message.agent_name, message.content))
if request_msg:
all_messages.append((SYSTEM, request_msg.content))
prompt = ""
prev_is_human = False # Whether the previous message is from human (in anthropic, the human is the user)
for i, message in enumerate(all_messages):
if i == 0:
assert (
message[0] == SYSTEM
) # The first message should be from the system
if message[0] == agent_name:
if prev_is_human:
prompt = f"{prompt}{anthropic.AI_PROMPT} {message[1]}"
else:
prompt = f"{prompt}\n\n{message[1]}"
prev_is_human = False
else:
if prev_is_human:
prompt = f"{prompt}\n\n[{message[0]}]: {message[1]}"
else:
prompt = f"{prompt}{anthropic.HUMAN_PROMPT}\n[{message[0]}]: {message[1]}"
prev_is_human = True
assert prev_is_human # The last message should be from the human
# Add the AI prompt for Claude to generate the response
prompt = f"{prompt}{anthropic.AI_PROMPT}"
response = self._get_response(prompt, *args, **kwargs)
# Remove the agent name if the response starts with it
response = re.sub(rf"^\s*\[{agent_name}]:?", "", response).strip()
return response
|