FunguyBot/plugins/infermatic-text.py

"""
Plugin for generating text using Infermatic AI API and sending it to a Matrix chat room.
"""
import os
import aiohttp
import json
import simplematrixbotlib as botlib
import re
from plugins.common import html_escape

# 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"

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

    if not INFERMATIC_API_KEY:
        await bot.api.send_text_message(room.room_id, "Infermatic API key not configured. Set INFERMATIC_API in .env.")
        return

    args = match.args()
    if len(args) < 1:
        await show_usage(room, bot)
        return

    if args[0] == "--list-models":
        await list_models(room, bot)
        return

    try:
        temperature = 0.9
        max_tokens = 512
        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):
    usage = """
<strong>📄 Infermatic Text Generation Usage:</strong>

<strong>Basic:</strong>
• <code>!text &lt;prompt&gt;</code> - Generate text using default model

<strong>Commands:</strong>
• <code>!text --list-models</code> - List all available models
• <code>!text --use-model &lt;model&gt; &lt;prompt&gt;</code> - Use specific model

<strong>Parameters:</strong>
• <code>--temperature &lt;0.0-1.0&gt;</code> - Set temperature (default: 0.9)
• <code>--max-tokens &lt;number&gt;</code> - Set max tokens (default: 2048)

<strong>Examples:</strong>
• <code>!text write a python function to calculate fibonacci</code>
• <code>!text --list-models</code>
"""
    await bot.api.send_markdown_message(room.room_id, usage)

async def list_models(room, bot):
    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()

        models = data.get('data', [])
        if not models:
            await bot.api.send_text_message(room.room_id, "No models found.")
            return

        output = "<strong>🔧 Available Models:</strong><br><br>"
        for model in models:
            model_id = html_escape(model.get('id', 'Unknown'))
            model_name = html_escape(model.get('name', model_id))
            context_length = model.get('context_length', 'Unknown')
            output += f"<strong>• {model_name}</strong><br>"
            output += f"  └─ ID: <code>{model_id}</code><br>"
            output += f"  └─ Context: {context_length}<br>"
            output += f"  └─ <strong>Usage:</strong> <code>!text --use-model {model_id} &lt;prompt&gt;</code><br><br>"

        # 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)

    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: {e}")

async def generate_text(room, bot, prompt, model, temperature, max_tokens):
    safe_prompt = html_escape(prompt)
    safe_model = html_escape(model)
    try:
        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": [
                {"role": "user", "content": prompt}
            ],
            "temperature": temperature,
            "max_tokens": max_tokens
        }

        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()

        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
        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)

        # 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 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: {e}")

__version__ = "1.0.3"
__author__ = "Funguy Bot"
__description__ = "AI text generation via Infermatic API"
__help__ = """
<details>
<summary><strong>!text</strong> – AI text generation (Infermatic)</summary>
<ul>
<li><code>!text &lt;prompt&gt;</code> – Generate text using default model</li>
<li><code>!text --list-models</code> – List available models</li>
<li><code>!text --use-model &lt;model&gt; &lt;prompt&gt;</code> – Specific model</li>
<li><code>--temperature &lt;0.0-1.0&gt;</code> – Set creativity (default 0.9)</li>
<li><code>--max-tokens &lt;number&gt;</code> – Max output length (default 512)</li>
</ul>
<p>Requires <strong>INFERMATIC_API</strong> env var.</p>
</details>
"""