Spaces:
Running
Running
import os | |
import gradio as gr | |
from cactus.agent import Cactus | |
from huggingface_hub import InferenceClient | |
MODEL_API_KEY_MAP = { | |
"gpt-3.5-turbo": "OPENAI_API_KEY", | |
"gpt-4": "OPENAI_API_KEY", | |
"claude-3-haiku-20240307": "ANTHROPIC_API_KEY", | |
"claude-3-opus-20240229": "ANTHROPIC_API_KEY", | |
"claude-3-sonnet-20240229": "ANTHROPIC_API_KEY", | |
} | |
def agent_fn(prompt, api_key=None): | |
env_key_name = MODEL_API_KEY_MAP.get(model_name) | |
if env_key_name: | |
if api_key: | |
os.environ[env_key_name] = api_key | |
else: | |
return f"Error: API key not found for {model_name}. Please provide it." | |
try: | |
cactus_model = Cactus(model_name, model_type="api") | |
response = cactus_model.run(prompt) | |
return response | |
except Exception as e: | |
return f"Error: {e}" | |
with gr.Blocks() as iface: | |
with gr.Row(): | |
model_name_dropdown = gr.Dropdown( | |
list(MODEL_API_KEY_MAP.keys()), label="Select Model" | |
) | |
api_key_input = gr.Textbox(label="Enter API Key", type="password") | |
prompt_input = gr.Textbox( | |
lines=2, placeholder="Enter your cheminformatics question..." | |
) | |
response_output = gr.Textbox(label="Agent Response") | |
# Event Triggers | |
model_name_dropdown.change( | |
lambda x: gr.update(visible=MODEL_API_KEY_MAP.get(x) is not None), | |
model_name_dropdown, | |
api_key_input, | |
) | |
api_key_input.change( | |
agent_fn, [prompt_input, model_name_dropdown, api_key_input], response_output | |
) | |
prompt_input.change( | |
agent_fn, [prompt_input, model_name_dropdown, api_key_input], response_output | |
) | |
iface.launch() | |