"""
Plugin for generating text using Infermatic AI API and sending it to a Matrix chat room.
"""
import os
import requests
import argparse
import json
import simplematrixbotlib as botlib
from asyncio import Queue
from dotenv import load_dotenv
# Load environment variables from .env file in the parent directory
plugin_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(plugin_dir)
dotenv_path = os.path.join(parent_dir, '.env')
load_dotenv(dotenv_path)
# Infermatic AI API configuration
INFERMATIC_API_KEY = os.getenv("INFERMATIC_API", "")
DEFAULT_MODEL = os.getenv("INFERMATIC_MODEL", "Sao10K-L3.1-70B-Hanami-x1")
INFERMATIC_API_BASE = "https://api.totalgpt.ai/v1"
# Queue to store pending commands
command_queue = Queue()
async def process_command(room, message, bot, prefix, config):
"""Queue and process !text commands sequentially."""
match = botlib.MessageMatch(room, message, bot, prefix)
if match.prefix() and match.command("text"):
if command_queue.empty():
await handle_command(room, message, bot, prefix, config)
else:
await command_queue.put((room, message, bot, prefix, config))
await bot.api.send_text_message(room.room_id, "Command queued. Please wait for the current request to finish.")
async def handle_command(room, message, bot, prefix, config):
"""Handle !text command: generate text using Infermatic AI API."""
match = botlib.MessageMatch(room, message, bot, prefix)
if not (match.prefix() and match.command("text")):
return
# Check if API key is configured
if not INFERMATIC_API_KEY:
await bot.api.send_text_message(
room.room_id,
"Infermatic API key not configured. Please set INFERMATIC_API environment variable."
)
return
# Parse command arguments
args = match.args()
if len(args) < 1:
await show_usage(room, bot)
return
# Check if it's a --list-models command
if args[0] == "--list-models":
await list_models(room, bot)
return
# Parse other arguments
try:
# Extract options manually since argparse doesn't handle mixed positional/optional well
temperature = 0.9
max_tokens = 2048
custom_model = None
prompt_parts = []
i = 0
while i < len(args):
if args[i] == "--temperature" and i + 1 < len(args):
temperature = float(args[i + 1])
i += 2
elif args[i] == "--max-tokens" and i + 1 < len(args):
max_tokens = int(args[i + 1])
i += 2
elif args[i] == "--use-model" and i + 1 < len(args):
custom_model = args[i + 1]
i += 2
else:
prompt_parts.append(args[i])
i += 1
prompt = ' '.join(prompt_parts).strip()
if not prompt:
await show_usage(room, bot)
return
model = custom_model or DEFAULT_MODEL
await generate_text(room, bot, prompt, model, temperature, max_tokens)
except ValueError as e:
await bot.api.send_text_message(room.room_id, f"Invalid parameter value: {e}")
except Exception as e:
await bot.api.send_text_message(room.room_id, f"Error processing command: {str(e)}")
async def show_usage(room, bot):
"""Display command usage information."""
usage = """
📄 Infermatic Text Generation Usage:
Basic:
• !text <prompt> - Generate text using default model
Commands:
• !text --list-models - List all available models
• !text --use-model <model> <prompt> - Use specific model
Parameters:
• --temperature <0.0-1.0> - Set temperature (default: 0.9)
• --max-tokens <number> - Set max tokens (default: 2048)
Examples:
• !text write a python function to calculate fibonacci
• !text --list-models
• !text --use-model llama-v3-8b-instruct explain quantum computing
• !text --temperature 0.7 write a haiku about AI
"""
await bot.api.send_markdown_message(room.room_id, usage)
async def list_models(room, bot):
"""List all available models from Infermatic AI."""
try:
await bot.api.send_text_message(room.room_id, "🔍 Fetching available models...")
url = f"{INFERMATIC_API_BASE}/models"
headers = {
"Authorization": f"Bearer {INFERMATIC_API_KEY}",
"Content-Type": "application/json"
}
response = requests.get(url, headers=headers, timeout=30)
response.raise_for_status()
data = response.json()
models = data.get('data', [])
if not models:
await bot.api.send_text_message(room.room_id, "No models found or error in response.")
return
# Format the model list
output = "🔧 Available Models:
"
for model in models:
model_id = model.get('id', 'Unknown')
model_name = model.get('name', model_id)
context_length = model.get('context_length', 'Unknown')
pricing = model.get('pricing', {})
output += f"• {model_name}
"
output += f" └─ ID: {model_id}
"
output += f" └─ Context: {context_length}
"
if pricing:
prompt_price = pricing.get('prompt', '0')
completion_price = pricing.get('completion', '0')
output += f" └─ Price: ${prompt_price}/${completion_price} per 1M tokens
"
output += f" └─ Usage: !text --use-model {model_id} <prompt>
"
# Wrap in collapsible details since list can be long
output = f"🔧 Available Models (Click to expand)
{output}📝 Generated Text (Click to expand)
"
output += f"Model: {model}
"
output += f"Prompt: {prompt}
"
output += f"Response:
"
output += f"{generated_text}"
output += f"