diff --git a/src/ai.py b/src/ai.py new file mode 100644 index 0000000..be6356f --- /dev/null +++ b/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 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 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 + /ai start [] + /ai set token + /ai clear + /ai correct + """ + 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 ") + 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 ", + "/ai start []", + "/ai set token ", + "/ai clear", + "/ai correct ", + ] + 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 ", "Set the default AI model (e.g., gpt-3.5-turbo)"], + ["start []", "Start a new AI chat with the specified or default model"], + ["set token ", "Set an API token for a specific company"], + ["clear", "Clear the current AI chat window"], + ["correct ", "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