feat(ai-server): Add AIProxyClient queue communication layer
This commit is contained in:
@@ -0,0 +1,98 @@
|
||||
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
|
||||
|
||||
@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)
|
||||
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()
|
||||
Reference in New Issue
Block a user