feat: Add AI chat plugin, ai.py
Add ai.py plugin to support AI model interaction via LiteLLM. Requires litellm.
This commit is contained in:
262
src/ai.py
Normal file
262
src/ai.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user