feat: Add AI chat plugin, ai.py

Add ai.py plugin to support AI model interaction via LiteLLM. Requires litellm.

src/ai.py

@@ -0,0 +1,262 @@
+"""AI Chat Plugin for CProof XMPP Client
+
+This plugin enables interaction with AI models via a dedicated chat window using the
+LiteLLM library, supporting providers like xAI (Grok) and OpenAI. Commands allow setting 
+default models, API tokens, starting chats, clearing windows, and correcting messages.
+
+- Requires the `litellm` library (`pip install litellm`).
+- Configure API tokens with zero-data-retention agreements for maximum privacy.
+- Use models like xAI's Grok for minimal data retention.
+- Chat history is not persisted locally beyond the active session.
+
+See Also:
+    - LiteLLM documentation: https://docs.litellm.ai/docs/
+    - xAI API: https://x.ai/api
+"""
+
+import threading
+import time
+import queue
+
+import litellm
+import prof
+
+from typing import Callable, Dict, List, Optional, Tuple
+
+# Global settings
+DEFAULT_MODEL_KEY = "default_model"  # Changed to avoid hardcoding model
+TOKEN_KEY = "ai_tokens"
+CHAT_WIN_ID: Optional[int] = None
+# Global message history: {win_id: [{"role": "user"|"assistant", "content": str}, ...]}
+CHAT_HISTORY: Dict[str, List[Dict[str, str]]] = {}
+# Output queue for safe display: (win_id, content)
+OUTPUT_QUEUE: queue.Queue[Tuple[str, str]] = queue.Queue()
+
+# Privacy settings for LiteLLM
+litellm.drop_params = True
+litellm.set_verbose = False
+
+# Disable printing in the console to avoid breaking profanity's outline
+litellm.suppress_debug_info = True
+
+
+def _get_tokens() -> Dict[str, str]:
+    """Retrieve tokens from settings as a dictionary."""
+    tokens: Dict[str, str] = {}
+    token_strings = prof.settings_string_list_get("ai_plugin", TOKEN_KEY)
+    if not token_strings:
+        return tokens
+    for ts in token_strings:
+        if ":" in ts:
+            company, token = ts.split(":", 1)
+            tokens[company] = token
+    return tokens
+
+
+def _save_tokens(tokens: Dict[str, str]) -> None:
+    """Save tokens to settings as a list of 'company:token' strings."""
+    prof.settings_string_list_clear("ai_plugin", TOKEN_KEY)
+    for company, token in tokens.items():
+        prof.settings_string_list_add("ai_plugin", TOKEN_KEY, f"{company}:{token}")
+
+
+def display_settings() -> None:
+    """Show the current default model and registered tokens."""
+    model = get_default_model()
+    tokens = _get_tokens()
+    token_info = ", ".join([f"{company}: {token[:4]}..." for company, token in tokens.items()])
+    prof.cons_show(f"AI Settings: Default Model: {model}, Tokens: {token_info or 'None set'}")
+
+
+def set_model(model) -> None:
+    """Set the default model to be used for AI chats."""
+    set_default_model(model)
+    prof.cons_show(f"Default model set to: {model}")
+
+def handler(win_id: str, message: str) -> None:
+    """Process messages in a chat window using the model from the window title."""
+    model = win_id.split(" - ", 2)[1]
+    prof.win_show(win_id, f"Me: {message}")
+    CHAT_HISTORY.setdefault(win_id, []).append({"role": "user", "content": message})
+
+    def run_completion():
+        tokens = _get_tokens()
+        try:
+            response = litellm.completion(
+                model=model,
+                messages=CHAT_HISTORY[win_id],
+                api_key=tokens.get(model.split("/")[0], None),
+            ).choices[0].message.content
+            CHAT_HISTORY[win_id].append({"role": "assistant", "content": response})
+            OUTPUT_QUEUE.put_nowait((win_id, f"AI: {response}"))
+        except Exception as e:
+            OUTPUT_QUEUE.put_nowait((win_id, f"Error: {str(e)}"))
+
+    thread = threading.Thread(target=run_completion)
+    thread.start()
+
+
+def process_queued_outputs() -> None:
+    """Process one output from the queue using prof.win_show."""
+    try:
+        win_id, content = OUTPUT_QUEUE.get_nowait()
+        prof.win_show(win_id, content)
+    except queue.Empty:
+        pass
+
+
+def create_chat_window(model: str) -> str:
+    """Create a new window for AI chat with the specified model."""
+    win_id = f"AI Chat - {model} - {int(time.time())}"
+    prof.win_create(win_id, handler)
+    prof.win_show(win_id, f"Chat started with {model}. Type a message to interact.")
+    prof.win_focus(win_id)
+    return win_id
+
+
+def start_chat(model: Optional[str] = None) -> None:
+    """Open a new chat window with the specified or default model."""
+    model = model if model else get_default_model()
+    create_chat_window(model)
+    prof.cons_show(f"Started AI chat with model: {model}")
+
+
+def set_token(company: str, token: str) -> None:
+    """Store an API token for a specific company."""
+    tokens = _get_tokens()
+    tokens[company] = token
+    _save_tokens(tokens)
+    prof.cons_show(f"Token set for {company}")
+
+
+def clear_chat() -> None:
+    """Notify that the current chat window is cleared."""
+    prof.cons_show("Sorry, current API doesn't support cleaning windows")
+    return
+    win_id = prof.get_current_win()
+    if win_id:
+        prof.win_show(win_id, "Chat cleared.")
+    else:
+        prof.cons_show("No active chat window.")
+
+
+def correct_message(corrected_text: str) -> None:
+    """Replace the latest user message in the current window's history and get AI response."""
+    prof.cons_show("Sorry, current API doesn't support correcting messages and getting current window")
+    return
+    # work in progress (get_current_win doesn't work)
+    win_id = prof.get_current_win()
+    if not win_id:
+        prof.cons_show("No active chat window. Use /ai start <model> first.")
+        return
+    title = prof.get_win_title(win_id)
+    if not title or not title.startswith("AI Chat - "):
+        prof.cons_show("Invalid chat window. Use /ai start <model> to create one.")
+        return
+    model = title.split(" - ", 2)[1]
+    history = CHAT_HISTORY.get(win_id, [])
+    user_messages = [msg for msg in history if msg["role"] == "user"]
+    if not user_messages:
+        prof.cons_show("No user messages in this chat to correct.")
+        return
+    # Replace the latest user message
+    for msg in history[::-1]:
+        if msg["role"] == "user":
+            msg["content"] = corrected_text
+            break
+    try:
+        response = litellm.completion(
+            model=model,
+            messages=[{"role": "user", "content": corrected_text}],
+            api_key=_get_tokens().get(model.split("/")[0], None),
+        ).choices[0].message.content
+        CHAT_HISTORY[win_id].append({"role": "assistant", "content": response})
+        prof.win_show(win_id, f"AI: {response}")
+    except Exception as e:
+        prof.cons_show(f"Error: {str(e)}")
+
+
+def get_default_model() -> str:
+    """Retrieve the default model from settings, defaulting to gpt-3.5-turbo."""
+    return prof.settings_string_get("ai_plugin", DEFAULT_MODEL_KEY, "gpt-3.5-turbo")
+
+
+def set_default_model(model: str) -> None:
+    """Save the default model to settings."""
+    prof.settings_string_set("ai_plugin", DEFAULT_MODEL_KEY, model)
+
+
+
+def _cmd_ai(*args: Tuple[Optional[str], ...]) -> None:
+    """Handle /ai commands for interacting with AI models.
+
+    Synopsis:
+        /ai
+        /ai set model <model>
+        /ai start [<model>]
+        /ai set token <company> <token>
+        /ai clear
+        /ai correct <message>
+    """
+    if not args:
+        display_settings()
+        return
+    if args[0] == "set" and len(args) >= 2:
+        if args[1] == "model" and len(args) == 3:
+            set_model(args[2])
+        elif args[1] == "token" and len(args) == 3:
+            try:
+                company, token = args[2].split(" ", 1)
+                set_token(company, token)
+            except ValueError:
+                prof.cons_show("Invalid format, use: /ai set token <company> <token>")
+        else:
+            prof.cons_show("Invalid command, use: /ai set model|token")
+    elif args[0] == "start" and len(args) <= 3:
+        start_chat(args[1] if len(args) == 2 else None)
+    elif args[0] == "clear" and len(args) == 1:
+        clear_chat()
+    elif args[0] == "correct" and len(args) == 3:
+        correct_message(args[2])
+    else:
+        prof.cons_bad_cmd_usage("/ai")
+
+
+def prof_init(version, status, account_name, fulljid):
+    """Initialize the AI chat plugin and register commands."""
+
+    synopsis = [
+        "/ai",
+        "/ai set model <model>",
+        "/ai start [<model>]",
+        "/ai set token <company> <token>",
+        "/ai clear",
+        "/ai correct <message>",
+    ]
+    description = """Interact with AI models via a chat interface using LiteLLM. 
+You can see the list of available models here: https://models.litellm.ai/"""
+    args = [
+        ["", "Display current AI plugin settings"],
+        ["set model <model>", "Set the default AI model (e.g., gpt-3.5-turbo)"],
+        ["start [<model>]", "Start a new AI chat with the specified or default model"],
+        ["set token <company> <token>", "Set an API token for a specific company"],
+        ["clear", "Clear the current AI chat window"],
+        ["correct <message>", "Correct a message using the current AI model"],
+    ]
+    examples = [
+        "/ai",
+        "/ai set token openai sk-xxx",
+        "/ai set model gpt-4",
+        "/ai start xai/grok",
+        "/ai clear",
+        '/ai correct I has a error',
+    ]
+
+    prof.register_command("/ai", 0, 3, synopsis, description, args, examples, _cmd_ai)
+    prof.completer_add("/ai", ["set", "start", "clear", "correct"])
+    prof.completer_add("/ai set", ["model", "token"])
+    prof.completer_add("/ai set model", ["openai/gpt-4o-mini", "xai/grok"])
+    prof.completer_add("/ai start", ["xai/grok", "openai/gpt-4o"])
+    prof.completer_add("/ai set token", ["openai", "xai"])
+
+    prof.register_timed(process_queued_outputs, 1)  # 1s interval to process AI message output