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:
98
src/ai.py
98
src/ai.py
@@ -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
|
||||||
Reference in New Issue
Block a user