102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
import json
|
|
import uuid
|
|
import threading
|
|
import subprocess
|
|
import sys
|
|
import os
|
|
from typing import Any, Optional
|
|
|
|
|
|
class AIProxyClient:
|
|
def __init__(self):
|
|
self._process: Optional[subprocess.Popen] = None
|
|
self._status: str = "disconnected"
|
|
self._pending: dict[str, Any] = {}
|
|
self._reader_thread: Optional[threading.Thread] = None
|
|
self._pending_lock: threading.Lock = threading.Lock()
|
|
|
|
@property
|
|
def status(self) -> str:
|
|
return self._status
|
|
|
|
def start_server(self):
|
|
if self._process is not None:
|
|
return
|
|
self._process = subprocess.Popen(
|
|
[sys.executable, "-m", "src.ai_server"],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
self._status = "init"
|
|
self._reader_thread = threading.Thread(target=self._read_loop, daemon=True)
|
|
self._reader_thread.start()
|
|
|
|
def _read_loop(self):
|
|
if self._process is None or self._process.stdout is None:
|
|
return
|
|
try:
|
|
for line in self._process.stdout:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
response = json.loads(line)
|
|
if response.get("type") == "ready" and self._status == "init":
|
|
self._status = "ready"
|
|
continue
|
|
rid = response.get("id")
|
|
if rid in self._pending:
|
|
self._pending[rid] = response
|
|
event_key = rid + "_event"
|
|
if event_key in self._pending:
|
|
self._pending[event_key].set()
|
|
except json.JSONDecodeError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
def send_command(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
|
|
if self._process is None or self._process.stdin is None:
|
|
return {"error": "server not started"}
|
|
|
|
request_id = str(uuid.uuid4())
|
|
event = threading.Event()
|
|
self._pending[request_id] = None
|
|
self._pending[request_id + "_event"] = event
|
|
|
|
command = {"id": request_id, "method": method, "params": params}
|
|
try:
|
|
self._process.stdin.write(json.dumps(command) + "\n")
|
|
self._process.stdin.flush()
|
|
except Exception as e:
|
|
self._pending.pop(request_id, None)
|
|
self._pending.pop(request_id + "_event", None)
|
|
return {"error": str(e)}
|
|
|
|
if not event.wait(timeout=60):
|
|
self._pending.pop(request_id, None)
|
|
self._pending.pop(request_id + "_event", None)
|
|
return {"error": "timeout"}
|
|
|
|
result = self._pending.pop(request_id, {"error": "response not found"})
|
|
self._pending.pop(request_id + "_event", None)
|
|
return result if result else {"error": "no response"}
|
|
|
|
def stop(self):
|
|
if self._process:
|
|
try:
|
|
self._process.stdin.close()
|
|
self._process.stdout.close()
|
|
self._process.stderr.close()
|
|
self._process.terminate()
|
|
self._process.wait(timeout=5)
|
|
except Exception:
|
|
try:
|
|
self._process.kill()
|
|
except Exception:
|
|
pass
|
|
self._process = None
|
|
self._status = "disconnected"
|
|
self._pending.clear() |