refactor: async I/O, input sanitisation, and shared utilities cleanup
This commit is contained in:
+35
-102
@@ -1,40 +1,18 @@
|
||||
"""
|
||||
Plugin for generating text using Infermatic AI API and sending it to a Matrix chat room.
|
||||
"""
|
||||
|
||||
import os
|
||||
import requests
|
||||
import argparse
|
||||
import aiohttp
|
||||
import json
|
||||
import simplematrixbotlib as botlib
|
||||
from asyncio import Queue
|
||||
from dotenv import load_dotenv
|
||||
import re
|
||||
from plugins.common import html_escape
|
||||
|
||||
# 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
|
||||
# No load_dotenv – handled centrally by funguy.py
|
||||
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)
|
||||
@@ -42,29 +20,20 @@ async def handle_command(room, message, bot, prefix, config):
|
||||
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."
|
||||
)
|
||||
await bot.api.send_text_message(room.room_id, "Infermatic API key not configured. Set INFERMATIC_API in .env.")
|
||||
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 = 512
|
||||
custom_model = None
|
||||
@@ -86,13 +55,11 @@ async def handle_command(room, message, bot, prefix, config):
|
||||
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:
|
||||
@@ -101,7 +68,6 @@ async def handle_command(room, message, bot, prefix, config):
|
||||
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 = """
|
||||
<strong>📄 Infermatic Text Generation Usage:</strong>
|
||||
|
||||
@@ -119,75 +85,57 @@ async def show_usage(room, bot):
|
||||
<strong>Examples:</strong>
|
||||
• <code>!text write a python function to calculate fibonacci</code>
|
||||
• <code>!text --list-models</code>
|
||||
• <code>!text --use-model llama-v3-8b-instruct explain quantum computing</code>
|
||||
• <code>!text --temperature 0.7 write a haiku about AI</code>
|
||||
"""
|
||||
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"
|
||||
}
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, headers=headers, timeout=30) as response:
|
||||
response.raise_for_status()
|
||||
data = await response.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.")
|
||||
await bot.api.send_text_message(room.room_id, "No models found.")
|
||||
return
|
||||
|
||||
# Format the model list
|
||||
output = "<strong>🔧 Available Models:</strong><br><br>"
|
||||
|
||||
for model in models:
|
||||
model_id = model.get('id', 'Unknown')
|
||||
model_name = model.get('name', model_id)
|
||||
model_id = html_escape(model.get('id', 'Unknown'))
|
||||
model_name = html_escape(model.get('name', model_id))
|
||||
context_length = model.get('context_length', 'Unknown')
|
||||
pricing = model.get('pricing', {})
|
||||
|
||||
output += f"<strong>• {model_name}</strong><br>"
|
||||
output += f" └─ ID: <code>{model_id}</code><br>"
|
||||
output += f" └─ Context: {context_length}<br>"
|
||||
|
||||
if pricing:
|
||||
prompt_price = pricing.get('prompt', '0')
|
||||
completion_price = pricing.get('completion', '0')
|
||||
output += f" └─ Price: ${prompt_price}/${completion_price} per 1M tokens<br>"
|
||||
|
||||
output += f" └─ <strong>Usage:</strong> <code>!text --use-model {model_id} <prompt></code><br><br>"
|
||||
|
||||
# Wrap in collapsible details since list can be long
|
||||
output = f"<details><summary><strong>🔧 Available Models (Click to expand)</strong></summary>{output}</details>"
|
||||
# Wrap in collapsible (from common)
|
||||
from plugins.common import collapsible_summary
|
||||
msg = collapsible_summary("🔧 Available Models (Click to expand)", output)
|
||||
await bot.api.send_markdown_message(room.room_id, msg)
|
||||
|
||||
await bot.api.send_markdown_message(room.room_id, output)
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
await bot.api.send_text_message(room.room_id, f"❌ Error fetching models: {str(e)}")
|
||||
except aiohttp.ClientError as e:
|
||||
await bot.api.send_text_message(room.room_id, f"❌ API error: {e}")
|
||||
except Exception as e:
|
||||
await bot.api.send_text_message(room.room_id, f"❌ Unexpected error: {str(e)}")
|
||||
|
||||
import re # add at the top of the file
|
||||
await bot.api.send_text_message(room.room_id, f"❌ Error: {e}")
|
||||
|
||||
async def generate_text(room, bot, prompt, model, temperature, max_tokens):
|
||||
"""Generate text using the Infermatic AI API."""
|
||||
safe_prompt = html_escape(prompt)
|
||||
safe_model = html_escape(model)
|
||||
try:
|
||||
await bot.api.send_text_message(room.room_id, f"📝 Generating text...")
|
||||
|
||||
await bot.api.send_text_message(room.room_id, "📝 Generating text...")
|
||||
url = f"{INFERMATIC_API_BASE}/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {INFERMATIC_API_KEY}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
@@ -197,49 +145,34 @@ async def generate_text(room, bot, prompt, model, temperature, max_tokens):
|
||||
"max_tokens": max_tokens
|
||||
}
|
||||
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=120)
|
||||
response.raise_for_status()
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, headers=headers, json=payload, timeout=120) as response:
|
||||
response.raise_for_status()
|
||||
data = await response.json()
|
||||
|
||||
data = response.json()
|
||||
generated_text = data.get('choices', [{}])[0].get('message', {}).get('content', '').strip()
|
||||
|
||||
if not generated_text:
|
||||
await bot.api.send_text_message(room.room_id, "No response generated.")
|
||||
return
|
||||
|
||||
# ---- Clean up blank lines that break list rendering ----
|
||||
# Remove blank lines directly before a list item (number‐dot or hyphen).
|
||||
# Clean up blank lines that break list rendering
|
||||
generated_text = re.sub(r'\n\n(\d+\.)', r'\n\1', generated_text)
|
||||
generated_text = re.sub(r'\n\n(- )', r'\n\1', generated_text)
|
||||
|
||||
# Build a pure Markdown message (no HTML)
|
||||
output = f"**Model:** `{model}`\n\n**Prompt:** {prompt}\n\n**Response:**\n\n{generated_text}"
|
||||
# Escape any stray HTML inside the generated text before embedding
|
||||
generated_text = html_escape(generated_text)
|
||||
|
||||
output = f"<strong>Model:</strong> <code>{safe_model}</code><br><strong>Prompt:</strong> {safe_prompt}<br><br><strong>Response:</strong><br><br>{generated_text}"
|
||||
await bot.api.send_markdown_message(room.room_id, output)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
await bot.api.send_text_message(room.room_id, "❌ Request timed out. The model is taking too long to respond.")
|
||||
except requests.exceptions.HTTPError as e:
|
||||
if e.response.status_code == 401:
|
||||
await bot.api.send_text_message(room.room_id, "❌ Authentication failed. Please check your INFERMATIC_API key.")
|
||||
elif e.response.status_code == 429:
|
||||
await bot.api.send_text_message(room.room_id, "❌ Rate limit exceeded. Please try again later.")
|
||||
else:
|
||||
await bot.api.send_text_message(room.room_id, f"❌ API error: HTTP {e.response.status_code}")
|
||||
except aiohttp.ClientError as e:
|
||||
await bot.api.send_text_message(room.room_id, f"❌ API error: {e}")
|
||||
except Exception as e:
|
||||
await bot.api.send_text_message(room.room_id, f"❌ Error generating text: {str(e)}")
|
||||
finally:
|
||||
if not command_queue.empty():
|
||||
next_command = await command_queue.get()
|
||||
await handle_command(*next_command)
|
||||
await bot.api.send_text_message(room.room_id, f"❌ Error: {e}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Plugin Metadata
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
__version__ = "1.0.2"
|
||||
__version__ = "1.0.3"
|
||||
__author__ = "Funguy Bot"
|
||||
__description__ = "AI text generation via Infermatic API (pure Markdown output)"
|
||||
__description__ = "AI text generation via Infermatic API (async, safe)"
|
||||
__help__ = """
|
||||
<details>
|
||||
<summary><strong>!text</strong> – AI text generation (Infermatic)</summary>
|
||||
|
||||
Reference in New Issue
Block a user