axnwork-cli 0.2.0__py3-none-any.whl

axon/backends/litellm_backend.py

@@ -0,0 +1,51 @@
+"""LiteLLM backend — wraps existing call_llm() logic."""
+from __future__ import annotations
+
+import logging
+import time
+from datetime import datetime
+
+from axon.backends.base import BackendResult
+from axon.backends.registry import register
+
+log = logging.getLogger("axon.backend.litellm")
+
+
+def _now_iso() -> str:
+    return datetime.now().astimezone().isoformat(timespec="seconds")
+
+
+@register("litellm")
+class LiteLLMBackend:
+    name = "litellm"
+
+    def __init__(self, config: dict):
+        self._model = config.get("default_model", "anthropic/claude-sonnet-4-20250514")
+        self._api_base = config.get("api_base", "")
+
+    def call(self, prompt: str, task: dict) -> BackendResult:
+        from axon.llm import call_llm
+        started_at = _now_iso()
+        started_mono = time.monotonic()
+        log.info(
+            "LiteLLM start started_at=%s model=%s prompt_chars=%d api_base=%s",
+            started_at,
+            self._model,
+            len(prompt),
+            self._api_base or "",
+        )
+        thinking, answer, usage = call_llm(prompt, self._model, self._api_base)
+        log.info(
+            "LiteLLM finished started_at=%s finished_at=%s duration_s=%.2f answer_chars=%d thinking_chars=%d total_tokens=%s cost=%s",
+            started_at,
+            _now_iso(),
+            time.monotonic() - started_mono,
+            len(answer or ""),
+            len(thinking or ""),
+            usage.get("total_tokens", 0),
+            usage.get("cost", 0.0),
+        )
+        return BackendResult(thinking=thinking, answer=answer, usage=usage)
+
+    def display_name(self) -> str:
+        return self._model

axon/backends/registry.py

@@ -0,0 +1,61 @@
+"""Backend registry and factory function."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from axon.backends.base import Backend
+
+_REGISTRY: dict[str, type] = {}
+
+
+def register(name: str):
+    """Decorator to register a backend class by name."""
+    def wrapper(cls):
+        _REGISTRY[name] = cls
+        return cls
+    return wrapper
+
+
+def auto_detect_backend() -> str:
+    """Pick the best available backend: claude-cli > codex-cli > litellm."""
+    import shutil
+    if shutil.which("claude"):
+        return "claude-cli"
+    if shutil.which("codex"):
+        return "codex-cli"
+    return "litellm"
+
+
+def create_backend(backend_name: str, config: dict) -> Backend:
+    """Create a backend instance by name.
+
+    If backend_name is "auto", detect the best available CLI tool.
+    Lazily imports backend modules so we don't pull in dependencies
+    (litellm, subprocess, etc.) until actually needed.
+    """
+    if backend_name == "auto":
+        backend_name = auto_detect_backend()
+
+    # Ensure all backends are registered
+    _ensure_loaded()
+
+    if backend_name not in _REGISTRY:
+        available = ", ".join(sorted(_REGISTRY.keys()))
+        raise ValueError(f"Unknown backend '{backend_name}'. Available: {available}")
+
+    return _REGISTRY[backend_name](config)
+
+
+_loaded = False
+
+
+def _ensure_loaded():
+    global _loaded
+    if _loaded:
+        return
+    _loaded = True
+    # Import backend modules to trigger @register decorators
+    import axon.backends.litellm_backend  # noqa: F401
+    import axon.backends.claude_cli  # noqa: F401
+    import axon.backends.codex_cli  # noqa: F401