Files
viralfactory/src/engines/LLMEngine/OpenaiLLMEngine.py

44 lines
1.5 KiB
Python
Raw Normal View History

2024-02-15 14:11:16 +01:00
import anthropic
2024-02-14 17:49:51 +01:00
import gradio as gr
from .BaseLLMEngine import BaseLLMEngine
2024-02-15 14:11:16 +01:00
# Assuming these are the models supported by Anthropics that you wish to include
ANTHROPIC_POSSIBLE_MODELS = [
"claude-2.1",
# Add more models as needed
2024-02-14 17:49:51 +01:00
]
2024-02-15 14:11:16 +01:00
class AnthropicsLLMEngine(BaseLLMEngine):
2024-02-14 17:49:51 +01:00
num_options = 1
2024-02-15 14:11:16 +01:00
name = "Anthropics"
description = "Anthropics language model engine."
2024-02-14 17:49:51 +01:00
def __init__(self, options: list) -> None:
self.model = options[0]
2024-02-15 14:11:16 +01:00
self.client = anthropic.Anthropic(api_key="YourAnthropicAPIKeyHere") # Ensure API key is securely managed
super().__init__()
2024-02-15 14:11:16 +01:00
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:
# Note: Adjust the parameters as per Anthropics API capabilities
message = self.client.messages.create(
max_tokens=max_tokens,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": chat_prompt},
],
2024-02-15 14:11:16 +01:00
model=self.model,
)
2024-02-15 14:11:16 +01:00
return message.content
@classmethod
def get_options(cls) -> list:
2024-02-14 17:49:51 +01:00
return [
gr.Dropdown(
label="Model",
2024-02-15 14:11:16 +01:00
choices=ANTHROPIC_POSSIBLE_MODELS,
2024-02-14 17:49:51 +01:00
max_choices=1,
2024-02-15 14:11:16 +01:00
value=ANTHROPIC_POSSIBLE_MODELS[0]
2024-02-14 17:49:51 +01:00
)
2024-02-15 14:11:16 +01:00
]