109 lines
5 KiB
Python
109 lines
5 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from django.core.management.base import BaseCommand
|
|
|
|
from control_plane.resources.models import Resource, ResourceKind
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Seed DGX Spark compute and model resources for Sol and Qwen."
|
|
|
|
def opencode_command(self, model_key: str) -> str:
|
|
env_name = f"ARTIFEX_{model_key.upper()}_OPENCODE_COMMAND"
|
|
if os.environ.get(env_name):
|
|
return os.environ[env_name]
|
|
if os.environ.get("ARTIFEX_OPENCODE_COMMAND"):
|
|
return os.environ["ARTIFEX_OPENCODE_COMMAND"]
|
|
executable = shutil.which("opencode")
|
|
if executable is None:
|
|
default_user_install = Path.home() / ".opencode" / "bin" / "opencode"
|
|
executable = str(default_user_install) if default_user_install.exists() else "opencode"
|
|
return f"{executable} run"
|
|
|
|
def opencode_config(self, model_key: str) -> dict[str, object]:
|
|
return {
|
|
"model_key": model_key,
|
|
"transport": os.environ.get(f"ARTIFEX_{model_key.upper()}_TRANSPORT", os.environ.get("ARTIFEX_OPENCODE_TRANSPORT", "local")),
|
|
"command": self.opencode_command(model_key),
|
|
"use_pty": os.environ.get(f"ARTIFEX_{model_key.upper()}_USE_PTY", os.environ.get("ARTIFEX_OPENCODE_USE_PTY", "0")) == "1",
|
|
"timeout_seconds": int(os.environ.get(f"ARTIFEX_{model_key.upper()}_TIMEOUT_SECONDS", os.environ.get("ARTIFEX_OPENCODE_TIMEOUT_SECONDS", "120"))),
|
|
}
|
|
|
|
def update_opencode_resource(self, *, model_key: str, name: str, roles: list[str], compute: Resource) -> None:
|
|
defaults = {
|
|
"name": name,
|
|
"kind": ResourceKind.MODEL,
|
|
"provider": "opencode",
|
|
"compute": compute,
|
|
"roles": roles,
|
|
"config": self.opencode_config(model_key),
|
|
}
|
|
existing = next((resource for resource in Resource.objects.filter(provider="opencode") if resource.config.get("model_key") == model_key), None)
|
|
if existing is not None:
|
|
for field, value in defaults.items():
|
|
setattr(existing, field, value)
|
|
existing.save(update_fields=[*defaults.keys(), "updated_at"])
|
|
return
|
|
Resource.objects.update_or_create(name=name, defaults={key: value for key, value in defaults.items() if key != "name"})
|
|
|
|
def handle(self, *args, **options):
|
|
spark, _ = Resource.objects.update_or_create(
|
|
name=os.environ.get("ARTIFEX_SPARK_RESOURCE_NAME", "DGX Spark"),
|
|
defaults={
|
|
"kind": ResourceKind.MODEL_COMPUTE,
|
|
"provider": "ssh",
|
|
"config": {"ssh_alias": os.environ.get("ARTIFEX_SPARK_SSH_ALIAS", "spark")},
|
|
"roles": [],
|
|
},
|
|
)
|
|
self.update_opencode_resource(
|
|
model_key="sol",
|
|
name=os.environ.get("ARTIFEX_SOL_RESOURCE_NAME", "GPT-5.6 Sol"),
|
|
roles=["PROJECT_BRAIN", "PLANNING", "ARCHAEOLOGY_INTERPRETATION", "AGENT_DESIGN", "ESCALATION"],
|
|
compute=spark,
|
|
)
|
|
self.update_opencode_resource(
|
|
model_key="terra",
|
|
name=os.environ.get("ARTIFEX_TERRA_RESOURCE_NAME", "GPT-5.6 Terra"),
|
|
roles=["REASONING", "REVIEW", "PORTFOLIO_IC", "STRATEGY"],
|
|
compute=spark,
|
|
)
|
|
self.update_opencode_resource(
|
|
model_key="luna",
|
|
name=os.environ.get("ARTIFEX_LUNA_RESOURCE_NAME", "GPT-5.6 Luna"),
|
|
roles=["REASONING", "MARKET_RESEARCH", "WEB_RESEARCH", "SYNTHESIS"],
|
|
compute=spark,
|
|
)
|
|
Resource.objects.update_or_create(
|
|
name=os.environ.get("ARTIFEX_QWEN_RESOURCE_NAME", "Qwen"),
|
|
defaults={
|
|
"kind": ResourceKind.MODEL,
|
|
"provider": "local_inference",
|
|
"compute": spark,
|
|
"roles": ["CODING", "REVIEW", "REASONING"],
|
|
"config": {
|
|
"endpoint_url": os.environ.get("ARTIFEX_QWEN_ENDPOINT_URL", "http://192.168.1.162:8002/v1/chat/completions"),
|
|
"health_url": os.environ.get("ARTIFEX_QWEN_HEALTH_URL", "http://192.168.1.162:8002/health"),
|
|
"model": os.environ.get("ARTIFEX_QWEN_MODEL", "qwen38"),
|
|
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
|
|
},
|
|
},
|
|
)
|
|
Resource.objects.update_or_create(
|
|
name=os.environ.get("ARTIFEX_SEARXNG_RESOURCE_NAME", "SearXNG Search"),
|
|
defaults={
|
|
"kind": ResourceKind.COMPUTE,
|
|
"provider": "searxng",
|
|
"compute": spark,
|
|
"roles": ["WEB_RESEARCH", "MARKET_RESEARCH", "SEARCH"],
|
|
"config": {
|
|
"endpoint_url": os.environ.get("ARTIFEX_SEARXNG_ENDPOINT_URL", "http://127.0.0.1:8080"),
|
|
"timeout_seconds": int(os.environ.get("ARTIFEX_SEARXNG_TIMEOUT_SECONDS", "20")),
|
|
},
|
|
},
|
|
)
|
|
self.stdout.write(self.style.SUCCESS("Seeded Spark model resources."))
|