diff --git a/src/ai.py b/src/ai.py index be6356f..949be4e 100644 --- a/src/ai.py +++ b/src/ai.py @@ -16,7 +16,8 @@ See Also: import threading import time -import queue +import json +import os import litellm import prof @@ -29,8 +30,10 @@ 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() +# Output file for safe display +OUTPUT_FILE = "/tmp/ai_output.txt" + +FILE_LOCK = threading.Lock() # Privacy settings for LiteLLM litellm.drop_params = True @@ -73,36 +76,56 @@ def set_model(model) -> None: set_default_model(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: - """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] 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(): - 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 = threading.Thread(target=run_completion, args=(win_id, model, current_history, OUTPUT_FILE, _get_tokens(), FILE_LOCK)) 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 + """Process outputs from the file using prof.win_show.""" + if os.path.exists(OUTPUT_FILE): + with open(OUTPUT_FILE, "r") as f: + lines = f.readlines() + outputs = [] + 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: @@ -164,21 +187,12 @@ def correct_message(corrected_text: str) -> None: 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)}") + handler(win_id, corrected_text) 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") + """Retrieve the default model from settings, defaulting to gpt-5.0.""" + return prof.settings_string_get("ai_plugin", DEFAULT_MODEL_KEY, "openai/gpt-5.0") 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 set token openai sk-xxx", "/ai set model gpt-4", - "/ai start xai/grok", + "/ai start xai/grok-4", "/ai clear", '/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.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 model", ["openai/gpt-5", "xai/grok-3-mini"]) + prof.completer_add("/ai start", ["xai/grok-4", "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 + prof.register_timed(process_queued_outputs, 1) # 1s interval to process AI message output \ No newline at end of file