diff --git a/ai/brain_server.py b/ai/brain_server.py index 1979235..4f02a96 100755 --- a/ai/brain_server.py +++ b/ai/brain_server.py @@ -23,6 +23,7 @@ if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from codex.prompt_builder import build_prompt +from ai.worker_manager import WorkerManager from services.blender_brain_contract import ( BlenderBrainContractError, DEEP_MODEL as BLENDER_DEEP_MODEL, @@ -40,6 +41,22 @@ DEEP_MODEL = os.environ.get( PORT = int(os.environ.get("SHIRE_BRAIN_PORT", "8765")) HOST_OVERRIDE = os.environ.get("SHIRE_BRAIN_HOST", "").strip() OLLAMA_TIMEOUT = int(os.environ.get("SHIRE_OLLAMA_TIMEOUT", "300")) +PARAMETRIC_PYTHON = os.environ.get( + "SHIRE_PARAMETRIC_PYTHON", + "/home/shire3d/.venvs/shire-parametric-forge/bin/python", +) +WORKER_NODE_ID = os.environ.get( + "SHIRE_WORKER_NODE_ID", + "brain-node-01", +) + +WORKER_MANAGER = WorkerManager( + node_id=WORKER_NODE_ID, + ollama_url=OLLAMA_URL, + fast_model=FAST_MODEL, + deep_model=DEEP_MODEL, + parametric_python=PARAMETRIC_PYTHON, +) def tailscale_ip(): @@ -234,32 +251,52 @@ def ask_ollama(model, prompt, max_tokens, system_context=""): class BrainHandler(BaseHTTPRequestHandler): - server_version = "SHIREBrainAPI/0007" + server_version = "SHIREBrainAPI/0008" def log_message(self, fmt, *args): print("%s - - [%s] %s" % (self.address_string(), self.log_date_time_string(), fmt % args)) def do_GET(self): - if self.path not in {"/", "/health"}: - json_response(self, {"ok": False, "error": "Not found"}, status=404) + if self.path in {"/", "/health"}: + json_response( + self, + { + "ok": True, + "service": "SHIRE Brain API", + "host": socket.gethostname(), + "model": FAST_MODEL, + "fast_model": FAST_MODEL, + "deep_model": DEEP_MODEL, + "blender_fast_model": BLENDER_FAST_MODEL, + "blender_deep_model": BLENDER_DEEP_MODEL, + "blender_routing_contract": "shire.blender_brain.route.v1", + "knowledge_context": "ray-approved registry", + "ollama_url": OLLAMA_URL, + "message": "SHIRE brain online.", + }, + ) + return + + if self.path in { + "/capabilities", + "/api/capabilities", + "/worker/capabilities", + }: + json_response(self, WORKER_MANAGER.capabilities()) + return + + if self.path in { + "/heartbeat", + "/api/heartbeat", + "/worker/heartbeat", + }: + json_response(self, WORKER_MANAGER.heartbeat()) return json_response( self, - { - "ok": True, - "service": "SHIRE Brain API", - "host": socket.gethostname(), - "model": FAST_MODEL, - "fast_model": FAST_MODEL, - "deep_model": DEEP_MODEL, - "blender_fast_model": BLENDER_FAST_MODEL, - "blender_deep_model": BLENDER_DEEP_MODEL, - "blender_routing_contract": "shire.blender_brain.route.v1", - "knowledge_context": "ray-approved registry", - "ollama_url": OLLAMA_URL, - "message": "SHIRE brain online.", - }, + {"ok": False, "error": "Not found"}, + status=404, ) def do_POST(self):