77 lines
3.1 KiB
Python
77 lines
3.1 KiB
Python
from typing import Union
|
|
|
|
from anthropic import AsyncAnthropic
|
|
from nio import RoomMessageImage
|
|
|
|
from matrix_gpt.chat_functions import download_mxc
|
|
from matrix_gpt.generate_clients.api_client import ApiClient
|
|
from matrix_gpt.generate_clients.command_info import CommandInfo
|
|
from matrix_gpt.image import process_image
|
|
|
|
|
|
class AnthropicApiClient(ApiClient):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def _create_client(self, base_url: str = None):
|
|
return AsyncAnthropic(
|
|
api_key=self._api_key
|
|
)
|
|
|
|
def assemble_context(self, messages: Union[str, list], system_prompt: str = None, injected_system_prompt: str = None):
|
|
if isinstance(messages, list):
|
|
messages = messages
|
|
else:
|
|
messages = [{"role": self._HUMAN_NAME, "content": [{"type": "text", "text": str(messages)}]}]
|
|
self._context = messages
|
|
return messages
|
|
|
|
def verify_context(self):
|
|
"""
|
|
Verify that the context alternates between the human and assistant, inserting the opposite
|
|
user type if it does not alternate correctly.
|
|
"""
|
|
i = 0
|
|
while i < len(self._context) - 1:
|
|
if self._context[i]['role'] == self._context[i + 1]['role']:
|
|
dummy = self.generate_text_msg(f'<{self._BOT_NAME} did not respond>', self._BOT_NAME) if self._context[i]['role'] == self._HUMAN_NAME else self.generate_text_msg(f'<{self._HUMAN_NAME} did not respond>', self._HUMAN_NAME)
|
|
self._context.insert(i + 1, dummy)
|
|
i += 1
|
|
# if self._context[-1]['role'] == self._HUMAN_NAME:
|
|
# self._context.append(self.generate_text_msg(f'<{self._BOT_NAME} did not respond>', self._BOT_NAME))
|
|
|
|
def generate_text_msg(self, content: str, role: str):
|
|
assert role in [self._HUMAN_NAME, self._BOT_NAME]
|
|
return {"role": role, "content": [{"type": "text", "text": str(content)}]}
|
|
|
|
def append_msg(self, content: str, role: str):
|
|
assert role in [self._HUMAN_NAME, self._BOT_NAME]
|
|
self._context.append(self.generate_text_msg(content, role))
|
|
|
|
async def append_img(self, img_event: RoomMessageImage, role: str):
|
|
assert role in [self._HUMAN_NAME, self._BOT_NAME]
|
|
img_bytes = await download_mxc(img_event.url, self._client_helper.client)
|
|
encoded_image = process_image(img_bytes, resize_px=784)
|
|
self._context.append({
|
|
"role": role,
|
|
'content': [{
|
|
'type': 'image',
|
|
'source': {
|
|
'type': 'base64',
|
|
'media_type': 'image/png',
|
|
'data': encoded_image
|
|
}
|
|
}]
|
|
})
|
|
|
|
async def generate(self, command_info: CommandInfo):
|
|
self.verify_context()
|
|
r = await self._create_client().messages.create(
|
|
model=command_info.model,
|
|
max_tokens=None if command_info.max_tokens == 0 else command_info.max_tokens,
|
|
temperature=command_info.temperature,
|
|
system='' if not command_info.system_prompt else command_info.system_prompt,
|
|
messages=self.context
|
|
)
|
|
return r.content[0].text
|