Notdiamond integration

agent/smart_model_routing.py

@@ -2,12 +2,17 @@
 
 from __future__ import annotations
 
+import logging
 import os
 import re
 from typing import Any, Dict, Optional
 
 from utils import is_truthy_value
 
+logger = logging.getLogger(__name__)
+
+_ND_API_URL = "https://api.notdiamond.ai/v2/modelRouter/modelSelect"
+
 _COMPLEX_KEYWORDS = {
     "debug",
     "debugging",
@@ -107,13 +112,89 @@ def choose_cheap_model_route(user_message: str, routing_config: Optional[Dict[st
     return route
 
 
+def _notdiamond_select_model(
+    messages: list,
+    nd_config: dict,
+) -> Optional[Dict[str, Any]]:
+    """Call NotDiamond to select the optimal model from candidates.
+
+    Returns {"provider": ..., "model": ..., "routing_reason": "notdiamond",
+    "session_id": ...} on success, or None on any failure.
+    """
+    import httpx
+
+    api_key_env = str(nd_config.get("api_key_env") or "NOTDIAMOND_API_KEY").strip()
+    api_key = os.getenv(api_key_env, "").strip()
+    if not api_key:
+        logger.debug("NotDiamond routing skipped: no API key in %s", api_key_env)
+        return None
+
+    candidates = nd_config.get("candidates") or []
+    if not candidates:
+        logger.debug("NotDiamond routing skipped: no candidates configured")
+        return None
+
+    timeout = _coerce_int(nd_config.get("timeout"), 5)
+    body: Dict[str, Any] = {
+        "messages": messages,
+        "llm_providers": [
+            {"provider": str(c.get("provider", "")), "model": str(c.get("model", ""))}
+            for c in candidates
+        ],
+    }
+    tradeoff = str(nd_config.get("tradeoff") or "").strip().lower()
+    if tradeoff in ("cost", "latency"):
+        body["tradeoff"] = tradeoff
+    if nd_config.get("hash_content"):
+        body["hash_content"] = True
+    preference_id = str(nd_config.get("preference_id") or "").strip()
+    if preference_id:
+        body["preference_id"] = preference_id
+
+    try:
+        with httpx.Client(timeout=timeout) as client:
+            resp = client.post(
+                _ND_API_URL,
+                json=body,
+                headers={
+                    "Authorization": f"Bearer {api_key}",
+                    "Content-Type": "application/json",
+                },
+            )
+            resp.raise_for_status()
+            data = resp.json()
+    except Exception as exc:
+        logger.debug("NotDiamond API call failed: %s", exc)
+        return None
+
+    providers = data.get("providers")
+    if not providers or not isinstance(providers, list):
+        logger.debug("NotDiamond returned no providers: %s", data)
+        return None
+
+    selected = providers[0]
+    provider = str(selected.get("provider", "")).strip()
+    model = str(selected.get("model", "")).strip()
+    if not provider or not model:
+        logger.debug("NotDiamond returned empty provider/model: %s", selected)
+        return None
+
+    return {
+        "provider": provider,
+        "model": model,
+        "routing_reason": "notdiamond",
+        "session_id": data.get("session_id", ""),
+    }
+
+
 def resolve_turn_route(user_message: str, routing_config: Optional[Dict[str, Any]], primary: Dict[str, Any]) -> Dict[str, Any]:
     """Resolve the effective model/runtime for one turn.
 
     Returns a dict with model/runtime/signature/label fields.
     """
-    route = choose_cheap_model_route(user_message, routing_config)
-    if not route:
+    cfg = routing_config or {}
+
+    def _primary_route():
         return {
             "model": primary.get("model"),
             "runtime": {
@@ -136,22 +217,50 @@ def resolve_turn_route(user_message: str, routing_config: Optional[Dict[str, Any
             ),
         }
 
-    from hermes_cli.runtime_provider import resolve_runtime_provider
+    # Determine strategy.
+    # Backward compat: no "strategy" key but enabled=True + cheap_model → heuristic.
+    strategy = str(cfg.get("strategy") or "").strip().lower()
+    logger.info("[routing] cfg=%s", {k: v for k, v in cfg.items() if k != "notdiamond"})
+    logger.info("[routing] strategy=%r enabled=%r", strategy, cfg.get("enabled"))
+    if not strategy or strategy == "off":
+        if _coerce_bool(cfg.get("enabled"), False) and cfg.get("cheap_model"):
+            strategy = "heuristic"
+        else:
+            logger.info("[routing] → primary (routing off)")
+            return _primary_route()
 
-    explicit_api_key = None
-    api_key_env = str(route.get("api_key_env") or "").strip()
-    if api_key_env:
-        explicit_api_key = os.getenv(api_key_env) or None
+    # ── NotDiamond ML-based routing ───────────────────────────────────
+    if strategy == "notdiamond":
+        if not _coerce_bool(cfg.get("enabled"), False):
+            logger.info("[routing] → primary (notdiamond strategy but enabled=False)")
+            return _primary_route()
+
+        nd_config = cfg.get("notdiamond") or {}
+        logger.info("[routing] calling NotDiamond with %d candidates", len(nd_config.get("candidates", [])))
+        messages = [{"role": "user", "content": user_message}]
+        nd_route = _notdiamond_select_model(messages, nd_config)
+        if not nd_route:
+            logger.info("[routing] → primary (NotDiamond returned None)")
+            return _primary_route()
+        logger.info("[routing] → NotDiamond selected: %s/%s", nd_route["provider"], nd_route["model"])
+
+        # Use the primary runtime (e.g. OpenRouter) and just swap the model.
+        # Aggregators like OpenRouter serve all models — we don't need to
+        # resolve a new provider for the NotDiamond-selected model.
+        # Model format: "provider/model" for aggregators (e.g. "openai/gpt-5"),
+        # or just "model" for direct providers.
+        primary_provider = str(primary.get("provider") or "").lower()
+        nd_model = nd_route["model"]
+
+        # For aggregators, use "provider/model" format (e.g. "openai/gpt-5")
+        _AGGREGATORS = {"openrouter", "nous", "opencode-zen", "opencode-go", "vercel", "huggingface", "kilocode"}
+        if primary_provider in _AGGREGATORS:
+            nd_model = f"{nd_route['provider']}/{nd_route['model']}"
+
+        logger.info("[routing] → using model %s via primary provider %s", nd_model, primary_provider)
 
-    try:
-        runtime = resolve_runtime_provider(
-            requested=route.get("provider"),
-            explicit_api_key=explicit_api_key,
-            explicit_base_url=route.get("base_url"),
-        )
-    except Exception:
         return {
-            "model": primary.get("model"),
+            "model": nd_model,
             "runtime": {
                 "api_key": primary.get("api_key"),
                 "base_url": primary.get("base_url"),
@@ -161,9 +270,9 @@ def resolve_turn_route(user_message: str, routing_config: Optional[Dict[str, Any
                 "args": list(primary.get("args") or []),
                 "credential_pool": primary.get("credential_pool"),
             },
-            "label": None,
+            "label": f"notdiamond → {nd_model} ({primary_provider})",
             "signature": (
-                primary.get("model"),
+                nd_model,
                 primary.get("provider"),
                 primary.get("base_url"),
                 primary.get("api_mode"),
@@ -172,6 +281,27 @@ def resolve_turn_route(user_message: str, routing_config: Optional[Dict[str, Any
             ),
         }
 
+    # ── Heuristic cheap-vs-strong routing ─────────────────────────────
+    route = choose_cheap_model_route(user_message, cfg)
+    if not route:
+        return _primary_route()
+
+    from hermes_cli.runtime_provider import resolve_runtime_provider
+
+    explicit_api_key = None
+    api_key_env = str(route.get("api_key_env") or "").strip()
+    if api_key_env:
+        explicit_api_key = os.getenv(api_key_env) or None
+
+    try:
+        runtime = resolve_runtime_provider(
+            requested=route.get("provider"),
+            explicit_api_key=explicit_api_key,
+            explicit_base_url=route.get("base_url"),
+        )
+    except Exception:
+        return _primary_route()
+
     return {
         "model": route.get("model"),
         "runtime": {