fix(AI plugin): bufferize response with file

Bufferizing response with a file allows to avoid crash caused by data transfer, based on plugins API issue in the CProof.

Minor changes: change default model and set correct xai models in the docs
This commit is contained in:
2025-09-08 13:05:52 +02:00
parent 94e9b155ef
commit 39781d0149

View File

@@ -16,7 +16,8 @@ See Also:
import threading import threading
import time import time
import queue import json
import os
import litellm import litellm
import prof import prof
@@ -29,8 +30,10 @@ TOKEN_KEY = "ai_tokens"
CHAT_WIN_ID: Optional[int] = None CHAT_WIN_ID: Optional[int] = None
# Global message history: {win_id: [{"role": "user"|"assistant", "content": str}, ...]} # Global message history: {win_id: [{"role": "user"|"assistant", "content": str}, ...]}
CHAT_HISTORY: Dict[str, List[Dict[str, str]]] = {} CHAT_HISTORY: Dict[str, List[Dict[str, str]]] = {}
# Output queue for safe display: (win_id, content) # Output file for safe display
OUTPUT_QUEUE: queue.Queue[Tuple[str, str]] = queue.Queue() OUTPUT_FILE = "/tmp/ai_output.txt"
FILE_LOCK = threading.Lock()
# Privacy settings for LiteLLM # Privacy settings for LiteLLM
litellm.drop_params = True litellm.drop_params = True
@@ -73,36 +76,56 @@ def set_model(model) -> None:
set_default_model(model) set_default_model(model)
prof.cons_show(f"Default model set to: {model}") prof.cons_show(f"Default model set to: {model}")
def run_completion(window_id: str, model: str, history: list, outfile: str, tokens: dict, file_lock) -> None:
"""Run AI completion and write response to file."""
try:
response = litellm.completion(
model=model,
messages=history,
api_key=tokens.get(model.split("/")[0], None),
).choices[0].message.content
with file_lock:
with open(outfile, "a") as f:
json.dump({"win_id": window_id, "content": f"AI: {response}"}, f)
f.write("\n")
except Exception as e:
with file_lock:
with open(outfile, "a") as f:
json.dump({"win_id": window_id, "content": f"Error: {str(e)}"}, f)
f.write("\n")
def handler(win_id: str, message: str) -> None: def handler(win_id: str, message: str) -> None:
"""Process messages in a chat window using the model from the window title.""" """Process messages in a chat window using the model from the window title.
Args:
win_id: Identifier for the chat window, used to extract the model name.
message: The user's message to process.
"""
model = win_id.split(" - ", 2)[1] model = win_id.split(" - ", 2)[1]
prof.win_show(win_id, f"Me: {message}") prof.win_show(win_id, f"Me: {message}")
CHAT_HISTORY.setdefault(win_id, []).append({"role": "user", "content": message})
# Create a new history list for this call to avoid modifying shared state
current_history = CHAT_HISTORY.get(win_id, []) + [{"role": "user", "content": message}]
def run_completion(): thread = threading.Thread(target=run_completion, args=(win_id, model, current_history, OUTPUT_FILE, _get_tokens(), FILE_LOCK))
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() thread.start()
def process_queued_outputs() -> None: def process_queued_outputs() -> None:
"""Process one output from the queue using prof.win_show.""" """Process outputs from the file using prof.win_show."""
try: if os.path.exists(OUTPUT_FILE):
win_id, content = OUTPUT_QUEUE.get_nowait() with open(OUTPUT_FILE, "r") as f:
prof.win_show(win_id, content) lines = f.readlines()
except queue.Empty: outputs = []
pass for line in lines:
if line.strip():
try:
data = json.loads(line)
outputs.append((data["win_id"], data["content"]))
except json.JSONDecodeError:
pass
for win_id, content in outputs:
prof.win_show(win_id, content)
open(OUTPUT_FILE, "w").close() # truncate the file
def create_chat_window(model: str) -> str: def create_chat_window(model: str) -> str:
@@ -164,21 +187,12 @@ def correct_message(corrected_text: str) -> None:
if msg["role"] == "user": if msg["role"] == "user":
msg["content"] = corrected_text msg["content"] = corrected_text
break break
try: handler(win_id, corrected_text)
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: def get_default_model() -> str:
"""Retrieve the default model from settings, defaulting to gpt-3.5-turbo.""" """Retrieve the default model from settings, defaulting to gpt-5.0."""
return prof.settings_string_get("ai_plugin", DEFAULT_MODEL_KEY, "gpt-3.5-turbo") return prof.settings_string_get("ai_plugin", DEFAULT_MODEL_KEY, "openai/gpt-5.0")
def set_default_model(model: str) -> None: def set_default_model(model: str) -> None:
@@ -247,7 +261,7 @@ You can see the list of available models here: https://models.litellm.ai/"""
"/ai", "/ai",
"/ai set token openai sk-xxx", "/ai set token openai sk-xxx",
"/ai set model gpt-4", "/ai set model gpt-4",
"/ai start xai/grok", "/ai start xai/grok-4",
"/ai clear", "/ai clear",
'/ai correct I has a error', '/ai correct I has a error',
] ]
@@ -255,8 +269,8 @@ You can see the list of available models here: https://models.litellm.ai/"""
prof.register_command("/ai", 0, 3, synopsis, description, args, examples, _cmd_ai) 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", "start", "clear", "correct"])
prof.completer_add("/ai set", ["model", "token"]) prof.completer_add("/ai set", ["model", "token"])
prof.completer_add("/ai set model", ["openai/gpt-4o-mini", "xai/grok"]) prof.completer_add("/ai set model", ["openai/gpt-5", "xai/grok-3-mini"])
prof.completer_add("/ai start", ["xai/grok", "openai/gpt-4o"]) prof.completer_add("/ai start", ["xai/grok-4", "openai/gpt-4o"])
prof.completer_add("/ai set token", ["openai", "xai"]) prof.completer_add("/ai set token", ["openai", "xai"])
prof.register_timed(process_queued_outputs, 1) # 1s interval to process AI message output prof.register_timed(process_queued_outputs, 1) # 1s interval to process AI message output