mirror of
https://github.com/Paillat-dev/viralfactory.git
synced 2026-01-02 01:06:19 +00:00
97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
import anthropic
|
|
import gradio as gr
|
|
import fix_busted_json
|
|
import orjson
|
|
|
|
from .BaseLLMEngine import BaseLLMEngine
|
|
|
|
ANTHROPIC_POSSIBLE_MODELS = [
|
|
"claude-2.1",
|
|
"claude-3-opus-20240229",
|
|
"claude-3-sonnet-20240229",
|
|
]
|
|
|
|
|
|
class AnthropicLLMEngine(BaseLLMEngine):
|
|
num_options = 1
|
|
name = "Anthropic"
|
|
description = "Anthropic language model engine."
|
|
|
|
def __init__(self, options: list) -> None:
|
|
self.model = options[0]
|
|
api_key = self.retrieve_setting(identifier="anthropic_api_key")["api_key"]
|
|
self.client = anthropic.Anthropic(api_key=api_key)
|
|
super().__init__()
|
|
|
|
def generate(
|
|
self,
|
|
system_prompt: str,
|
|
chat_prompt: str,
|
|
max_tokens: int = 1024,
|
|
temperature: float = 1.0,
|
|
json_mode: bool = False,
|
|
top_p: float = 1,
|
|
frequency_penalty: float = 0,
|
|
presence_penalty: float = 0,
|
|
) -> str | dict:
|
|
tries = 0
|
|
while tries < 2:
|
|
messages = [
|
|
{"role": "user", "content": chat_prompt},
|
|
]
|
|
if json_mode:
|
|
# anthropic does not officially support JSON mode, but we can bias the output towards a JSON-like format
|
|
messages.append({"role": "assistant", "content": "{"})
|
|
response: anthropic.types.Message = self.client.messages.create(
|
|
max_tokens=max_tokens,
|
|
messages=messages,
|
|
model=self.model,
|
|
top_p=top_p,
|
|
temperature=temperature,
|
|
system=system_prompt,
|
|
)
|
|
|
|
content = response.content[0].text
|
|
if json_mode:
|
|
content = "{" + content
|
|
# we remove everything after the last closing curly brace
|
|
content = content[: content.rfind("}") + 1]
|
|
content = content.replace("\n", "")
|
|
try:
|
|
returnable = fix_busted_json.repair_json(content)
|
|
returnable = orjson.loads(returnable)
|
|
return returnable
|
|
except Exception as e: # noqa wait for library to imlement pep https://peps.python.org/pep-0352/ (Required Superclass for Exceptions
|
|
tries += 1
|
|
else:
|
|
return content
|
|
|
|
@classmethod
|
|
def get_options(cls) -> list:
|
|
return [
|
|
gr.Dropdown(
|
|
label="Model",
|
|
choices=ANTHROPIC_POSSIBLE_MODELS,
|
|
max_choices=1,
|
|
value=ANTHROPIC_POSSIBLE_MODELS[0],
|
|
)
|
|
]
|
|
|
|
@classmethod
|
|
def get_settings(cls):
|
|
current_api_key = cls.retrieve_setting(identifier="anthropic_api_key")
|
|
current_api_key = current_api_key["api_key"] if current_api_key else ""
|
|
api_key_input = gr.Textbox(
|
|
label="Anthropic API Key",
|
|
type="password",
|
|
value=current_api_key,
|
|
)
|
|
save = gr.Button("Save")
|
|
|
|
def save_api_key(api_key: str):
|
|
cls.store_setting(identifier="anthropic_api_key", data={"api_key": api_key})
|
|
gr.Info("API key saved successfully.")
|
|
return gr.update(value=api_key)
|
|
|
|
save.click(save_api_key, inputs=[api_key_input])
|