297 lines
12 KiB
Python
297 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import shlex
|
|
import subprocess
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from collections.abc import Iterator
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from control_plane.resources.models import Resource
|
|
from model_router.router import ModelChunk, ModelRequestContract, ModelResponseContract
|
|
|
|
|
|
class ProviderError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _extract_json(text: str) -> dict[str, Any]:
|
|
try:
|
|
value = json.loads(text)
|
|
except json.JSONDecodeError as exc:
|
|
raise ProviderError("Provider returned malformed JSON") from exc
|
|
if not isinstance(value, dict):
|
|
raise ProviderError("Provider JSON response must be an object")
|
|
return value
|
|
|
|
|
|
def extract_json_object(text: str) -> dict[str, Any]:
|
|
try:
|
|
return _extract_json(text)
|
|
except ProviderError:
|
|
start = text.find("{")
|
|
end = text.rfind("}")
|
|
if start == -1 or end == -1 or end <= start:
|
|
raise
|
|
return _extract_json(text[start : end + 1])
|
|
|
|
|
|
def _clean_opencode_output(text: str) -> str:
|
|
text = re.sub(r"\x1b\[[0-9;?]*[ -/]*[@-~]", "", text)
|
|
lines = []
|
|
for line in text.splitlines():
|
|
stripped = line.strip()
|
|
if not stripped or stripped.startswith("> build"):
|
|
continue
|
|
lines.append(line)
|
|
return "\n".join(lines).strip()
|
|
|
|
|
|
@dataclass
|
|
class SolProvider:
|
|
resource: Resource
|
|
provider_name: str = "opencode"
|
|
|
|
def complete(self, request: ModelRequestContract) -> ModelResponseContract:
|
|
config = self.resource.config
|
|
compute = self.resource.compute
|
|
timeout = int(config.get("timeout_seconds", 120))
|
|
command = str(config.get("command", "opencode run"))
|
|
transport = str(config.get("transport", "ssh"))
|
|
use_pty = bool(config.get("use_pty", False))
|
|
working_directory = str(config.get("working_directory") or "").strip() or None
|
|
if transport == "local":
|
|
argv = shlex.split(command, posix=os.name != "nt")
|
|
if use_pty:
|
|
shell_command = " ".join(shlex.quote(part) for part in argv)
|
|
argv = ["script", "-q", "-e", "-c", shell_command, "/dev/null"]
|
|
completed = subprocess.run(
|
|
argv,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout,
|
|
check=False,
|
|
cwd=working_directory,
|
|
input=request.prompt,
|
|
)
|
|
else:
|
|
ssh_alias = (compute.config if compute else {}).get("ssh_alias", config.get("ssh_alias", "spark"))
|
|
remote_command = " ".join(shlex.quote(part) for part in shlex.split(command))
|
|
if working_directory:
|
|
remote_command = f"cd {shlex.quote(working_directory)} && {remote_command}"
|
|
completed = subprocess.run(
|
|
["ssh", "-T", str(ssh_alias), remote_command],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout,
|
|
check=False,
|
|
input=request.prompt,
|
|
)
|
|
if completed.returncode != 0:
|
|
raise ProviderError(completed.stderr.strip() or "Sol provider failed")
|
|
content = _clean_opencode_output(completed.stdout)
|
|
if not content:
|
|
raise ProviderError("Sol provider response missing content")
|
|
return ModelResponseContract(
|
|
model=self.resource.name,
|
|
content=content,
|
|
metadata={"provider": self.provider_name, "usage": {}},
|
|
)
|
|
|
|
def stream(self, request: ModelRequestContract) -> Iterator[ModelChunk]:
|
|
config = self.resource.config
|
|
command = str(config.get("command", "opencode run"))
|
|
transport = str(config.get("transport", "ssh"))
|
|
working_directory = str(config.get("working_directory") or "").strip() or None
|
|
if transport == "local":
|
|
argv = shlex.split(command, posix=os.name != "nt")
|
|
else:
|
|
compute = self.resource.compute
|
|
ssh_alias = (compute.config if compute else {}).get(
|
|
"ssh_alias", config.get("ssh_alias", "spark")
|
|
)
|
|
remote_command = " ".join(shlex.quote(part) for part in shlex.split(command))
|
|
if working_directory:
|
|
remote_command = f"cd {shlex.quote(working_directory)} && {remote_command}"
|
|
argv = ["ssh", "-T", str(ssh_alias), remote_command]
|
|
process = subprocess.Popen(
|
|
argv,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
bufsize=1,
|
|
cwd=working_directory if transport == "local" else None,
|
|
stdin=subprocess.PIPE,
|
|
)
|
|
assert process.stdin is not None
|
|
process.stdin.write(request.prompt)
|
|
process.stdin.close()
|
|
emitted = False
|
|
assert process.stdout is not None
|
|
for line in process.stdout:
|
|
cleaned = _clean_opencode_output(line)
|
|
if not cleaned:
|
|
continue
|
|
emitted = True
|
|
yield ModelChunk(cleaned + "\n", {"provider": self.provider_name})
|
|
stderr = process.stderr.read().strip() if process.stderr is not None else ""
|
|
return_code = process.wait()
|
|
if return_code != 0:
|
|
raise ProviderError(stderr or "Sol streaming provider failed")
|
|
if not emitted:
|
|
raise ProviderError("Sol streaming provider response missing content")
|
|
|
|
def health(self) -> str:
|
|
config = self.resource.config
|
|
try:
|
|
if config.get("transport") == "local":
|
|
command = shlex.split(str(config.get("command", "opencode")), posix=os.name != "nt")[0]
|
|
completed = subprocess.run([command, "--version"], capture_output=True, text=True, timeout=10, check=False)
|
|
else:
|
|
compute = self.resource.compute
|
|
ssh_alias = (compute.config if compute else {}).get("ssh_alias", self.resource.config.get("ssh_alias", "spark"))
|
|
completed = subprocess.run(["ssh", str(ssh_alias), "true"], capture_output=True, text=True, timeout=10, check=False)
|
|
except Exception:
|
|
return "UNAVAILABLE"
|
|
return "AVAILABLE" if completed.returncode == 0 else "UNAVAILABLE"
|
|
|
|
|
|
@dataclass
|
|
class QwenProvider:
|
|
resource: Resource
|
|
provider_name: str = "local_inference"
|
|
|
|
def complete(self, request: ModelRequestContract) -> ModelResponseContract:
|
|
config = self.resource.config
|
|
url = str(config.get("endpoint_url", "http://localhost:8000/v1/chat/completions"))
|
|
timeout = int(config.get("timeout_seconds", 120))
|
|
attempts = max(1, int(config.get("retry_attempts", 1)))
|
|
backoff = float(config.get("retry_backoff_seconds", 1.0))
|
|
body = {
|
|
"model": config.get("model", self.resource.name),
|
|
"messages": [{"role": "user", "content": request.prompt}],
|
|
"max_tokens": request.token_budget,
|
|
"temperature": config.get("temperature", 0),
|
|
}
|
|
if config.get("response_format"):
|
|
body["response_format"] = config["response_format"]
|
|
if config.get("extra_body"):
|
|
body.update(config["extra_body"])
|
|
failures = []
|
|
for attempt in range(1, attempts + 1):
|
|
http_request = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(body).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(http_request, timeout=timeout) as response:
|
|
data = json.loads(response.read().decode("utf-8"))
|
|
break
|
|
except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
|
|
failures.append(f"attempt {attempt}/{attempts}: {exc}")
|
|
if attempt == attempts:
|
|
raise ProviderError("Qwen provider failed after retries: " + " | ".join(failures)) from exc
|
|
time.sleep(backoff * attempt)
|
|
choices = data.get("choices", [])
|
|
content = ""
|
|
if choices:
|
|
message = choices[0].get("message", {})
|
|
content = message.get("content", "")
|
|
if not isinstance(content, str) or not content:
|
|
raise ProviderError("Qwen provider response missing content")
|
|
return ModelResponseContract(
|
|
model=str(data.get("model", self.resource.name)),
|
|
content=content,
|
|
metadata={
|
|
"provider": self.provider_name,
|
|
"usage": data.get("usage", {}),
|
|
"attempts": len(failures) + 1,
|
|
},
|
|
)
|
|
|
|
def stream(self, request: ModelRequestContract) -> Iterator[ModelChunk]:
|
|
config = self.resource.config
|
|
url = str(config.get("endpoint_url", "http://localhost:8000/v1/chat/completions"))
|
|
body = {
|
|
"model": config.get("model", self.resource.name),
|
|
"messages": [{"role": "user", "content": request.prompt}],
|
|
"max_tokens": request.token_budget,
|
|
"temperature": config.get("temperature", 0),
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
}
|
|
if config.get("extra_body"):
|
|
body.update(config["extra_body"])
|
|
http_request = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(body).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
emitted = False
|
|
usage: dict[str, object] = {}
|
|
with urllib.request.urlopen(
|
|
http_request, timeout=int(config.get("timeout_seconds", 240))
|
|
) as response:
|
|
for raw_line in response:
|
|
line = raw_line.decode("utf-8").strip()
|
|
if not line.startswith("data:"):
|
|
continue
|
|
payload = line[5:].strip()
|
|
if payload == "[DONE]":
|
|
break
|
|
try:
|
|
event = json.loads(payload)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if event.get("usage"):
|
|
usage = event["usage"]
|
|
yield ModelChunk("", {"provider": self.provider_name, "usage": usage})
|
|
choices = event.get("choices") or []
|
|
if not choices:
|
|
continue
|
|
content = (choices[0].get("delta") or {}).get("content")
|
|
if content:
|
|
emitted = True
|
|
yield ModelChunk(
|
|
content,
|
|
{"provider": self.provider_name, "usage": usage},
|
|
)
|
|
if not emitted:
|
|
raise ProviderError("Qwen streaming provider response missing content")
|
|
|
|
def health(self) -> str:
|
|
base_url = str(self.resource.config.get("health_url", self.resource.config.get("endpoint_url", ""))).replace(
|
|
"/v1/chat/completions", "/health"
|
|
)
|
|
if not base_url:
|
|
return "UNAVAILABLE"
|
|
try:
|
|
with urllib.request.urlopen(base_url, timeout=5) as response:
|
|
return "AVAILABLE" if 200 <= response.status < 500 else "DEGRADED"
|
|
except Exception:
|
|
return "UNAVAILABLE"
|
|
|
|
|
|
def providers_from_resources() -> dict[str, object]:
|
|
providers: dict[str, object] = {}
|
|
opencode_resources = Resource.objects.filter(is_active=True, provider="opencode")
|
|
qwen = Resource.objects.filter(is_active=True, provider="local_inference").first()
|
|
for resource in opencode_resources:
|
|
model_key = str(resource.config.get("model_key") or "").strip().lower()
|
|
if model_key:
|
|
providers[model_key] = SolProvider(resource)
|
|
if "sol" not in providers:
|
|
sol = opencode_resources.first()
|
|
if sol is not None:
|
|
providers["sol"] = SolProvider(sol)
|
|
if qwen is not None:
|
|
providers["qwen"] = QwenProvider(qwen)
|
|
return providers
|