viento 0.2.0__py3-none-any.whl

viento/__init__.py

@@ -0,0 +1,34 @@
+"""
+Zephyr SDK — Distributed AI Inference Runtime
+
+Connect your local LLMs (Ollama, llama.cpp, vLLM) to the Zephyr Cloud
+mesh and serve distributed inference jobs via the OpenAI-compatible API.
+
+Quick Start:
+    >>> from viento.client.client import VientoClient
+    >>> client = VientoClient(api_key="zph_tmp_...")
+    >>> response = client.chat.completions.create(
+    ...     model="llama3:latest",
+    ...     messages=[{"role": "user", "content": "Hello!"}],
+    ... )
+
+Run as a node:
+    $ viento run
+
+Documentation: https://github.com/abhinav00anand/viento
+"""
+
+__version__ = "0.2.0"
+__author__ = "Zephyr Cloud Team"
+__email__ = "indrohelpdesk@gmail.com"
+__license__ = "MIT"
+__url__ = "https://github.com/abhinav00anand/viento"
+
+from viento.client.client import AsyncVientoClient, VientoClient
+
+__all__ = [
+    "__version__",
+    "__author__",
+    "VientoClient",
+    "AsyncVientoClient",
+]

viento/backends/__init__.py

@@ -0,0 +1,45 @@
+"""Inference backend adapters for local and cloud LLM runtime engines."""
+
+from typing import Optional
+from viento.backends.base import (
+    InferenceBackend,
+    BackendError,
+    BackendOfflineError,
+    ModelNotFoundError,
+    ContextOverflowError,
+    BackendTimeoutError,
+)
+from viento.backends.ollama import OllamaAdapter
+from viento.backends.llamacpp import LlamaCppAdapter
+from viento.backends.vllm import VLLMAdapter
+
+
+def get_backend_adapter(backend_name: str = "ollama", base_url: Optional[str] = None) -> InferenceBackend:
+    """
+    Factory function to instantiate the matching backend adapter class
+    based on backend_name string ("ollama", "vllm", "llamacpp").
+    """
+    name = (backend_name or "ollama").strip().lower()
+    if name in ("vllm", "vllm_adapter"):
+        url = base_url or "http://localhost:8000/v1"
+        return VLLMAdapter(base_url=url)
+    elif name in ("llamacpp", "llama_cpp", "llama.cpp"):
+        url = base_url or "http://localhost:8080"
+        return LlamaCppAdapter(base_url=url)
+    else:
+        url = base_url or "http://localhost:11434"
+        return OllamaAdapter(base_url=url)
+
+
+__all__ = [
+    "InferenceBackend",
+    "BackendError",
+    "BackendOfflineError",
+    "ModelNotFoundError",
+    "ContextOverflowError",
+    "BackendTimeoutError",
+    "OllamaAdapter",
+    "LlamaCppAdapter",
+    "VLLMAdapter",
+    "get_backend_adapter",
+]

viento/backends/base.py

@@ -0,0 +1,122 @@
+"""
+Abstract Base Classes for Inference Backends and Execution Handles.
+
+Defines the execution handle contract for non-destructive HTTP stream cancellation,
+backend exceptions, and thread-safe compute resource release.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, List, Optional, Tuple
+from pydantic import BaseModel
+
+
+class BackendError(Exception):
+    """Base exception for inference backend errors."""
+    pass
+
+
+class BackendOfflineError(BackendError):
+    """Raised when the backend server is unreachable or offline."""
+    pass
+
+
+class BackendTimeoutError(BackendError):
+    """Raised when request times out on backend."""
+    pass
+
+
+class ModelNotFoundError(BackendError):
+    """Raised when a requested model is not found on the backend."""
+    pass
+
+
+class ContextLengthExceededError(BackendError):
+    """Raised when request prompt tokens exceed backend context length limits."""
+    pass
+
+
+ContextOverflowError = ContextLengthExceededError
+
+
+class ExecutionHandle(ABC):
+    """Abstract execution handle wrapping an active inference HTTP response stream."""
+
+    @abstractmethod
+    def cancel(self) -> None:
+        """Cancel execution and close underlying HTTP socket connection immediately."""
+        pass
+
+    @abstractmethod
+    def is_done(self) -> bool:
+        """Check if execution has completed or aborted."""
+        pass
+
+
+class GenerationChunk(BaseModel):
+    delta: str
+    finish_reason: Optional[str] = None
+    index: int = 0
+
+
+class GenerationResult(BaseModel):
+    full_text: str
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    finish_reason: str = "stop"
+    is_estimated: bool = False
+
+
+class EmbeddingResult(BaseModel):
+    embeddings: List[List[float]]
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    is_estimated: bool = False
+
+
+class InferenceBackend(ABC):
+    """Abstract base class for all local inference engine adapters (Ollama, vLLM, llama.cpp)."""
+
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+    @abstractmethod
+    def capabilities(self) -> List[str]:
+        pass
+
+    @abstractmethod
+    def health(self) -> bool:
+        pass
+
+    @abstractmethod
+    def list_models(self) -> List[Dict[str, Any]]:
+        pass
+
+    @abstractmethod
+    def generate(
+        self,
+        job_id: str,
+        model: str,
+        messages: List[Dict[str, str]],
+        temperature: float = 0.7,
+        max_tokens: int = 512,
+        callback: Optional[Callable[[GenerationChunk], None]] = None,
+        stop: Optional[List[str]] = None,
+        handle_callback: Optional[Callable[[ExecutionHandle], None]] = None,
+    ) -> Tuple[GenerationResult, ExecutionHandle]:
+        pass
+
+    @abstractmethod
+    def embeddings(
+        self,
+        model: str,
+        inputs: List[str],
+        job_id: Optional[str] = None,
+        handle_callback: Optional[Callable[[ExecutionHandle], None]] = None,
+    ) -> EmbeddingResult:
+        pass
+
+    @abstractmethod
+    def cancel(self, job_id: str) -> None:
+        pass

viento/backends/llamacpp.py

@@ -0,0 +1,242 @@
+"""Production llama.cpp REST backend adapter for local llama.cpp server instances."""
+
+import json
+import logging
+import threading
+import time
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import httpx
+
+from viento.backends.base import (
+    BackendError,
+    BackendOfflineError,
+    BackendTimeoutError,
+    ContextOverflowError,
+    EmbeddingResult,
+    ExecutionHandle,
+    GenerationChunk,
+    GenerationResult,
+    InferenceBackend,
+    ModelNotFoundError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class LlamaCppExecutionHandle(ExecutionHandle):
+    def __init__(self, job_id: str, response: httpx.Response):
+        self.job_id = job_id
+        self.response = response
+        self._is_done = False
+        self._cancelled = False
+        self._lock = threading.Lock()
+
+    def cancel(self) -> None:
+        with self._lock:
+            if not self._is_done and not self._cancelled:
+                self._cancelled = True
+                try:
+                    self.response.close()
+                    logger.info("Closed HTTP response stream for cancelled job %s", self.job_id)
+                except Exception:
+                    pass
+
+    def is_done(self) -> bool:
+        with self._lock:
+            return self._is_done
+
+    def mark_done(self):
+        with self._lock:
+            self._is_done = True
+
+
+class LlamaCppAdapter(InferenceBackend):
+    """Adapter connecting to llama.cpp server (default endpoint: http://localhost:8080)."""
+
+    def __init__(self, base_url: str = "http://localhost:8080", timeout: float = 120.0):
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self._client = httpx.Client(base_url=self.base_url, timeout=self.timeout)
+
+    def name(self) -> str:
+        return "llamacpp"
+
+    def capabilities(self) -> List[str]:
+        return ["chat", "embeddings", "streaming"]
+
+    def health(self) -> bool:
+        try:
+            resp = self._client.get("/health")
+            if resp.status_code == 200:
+                return True
+            resp_root = self._client.get("/")
+            return resp_root.status_code == 200
+        except Exception:
+            return False
+
+    def list_models(self) -> List[Dict[str, Any]]:
+        try:
+            resp = self._client.get("/v1/models")
+            if resp.status_code == 200:
+                data = resp.json()
+                return data.get("data", [])
+
+            resp_props = self._client.get("/props")
+            if resp_props.status_code == 200:
+                props = resp_props.json()
+                default_name = props.get("default_generation_settings", {}).get("model", "llama.cpp-model")
+                return [{"id": default_name, "object": "model", "owned_by": "llama.cpp"}]
+
+            return [{"id": "llama.cpp-default", "object": "model"}]
+        except Exception as e:
+            logger.error("Failed to list llama.cpp models: %s", e)
+            return []
+
+    def generate(
+        self,
+        job_id: str,
+        model: str,
+        messages: List[Dict[str, str]],
+        temperature: float = 0.7,
+        max_tokens: int = 512,
+        callback: Optional[Callable[[GenerationChunk], None]] = None,
+        stop: Optional[List[str]] = None,
+        handle_callback: Optional[Callable[[ExecutionHandle], None]] = None,
+    ) -> Tuple[GenerationResult, ExecutionHandle]:
+        payload = {
+            "model": model,
+            "messages": messages,
+            "stream": True,
+            "temperature": temperature,
+            "max_tokens": max_tokens,
+        }
+        if stop:
+            payload["stop"] = stop
+
+        req = self._client.build_request("POST", "/v1/chat/completions", json=payload)
+        res = self._client.send(req, stream=True)
+
+        if res.status_code != 200:
+            res.close()
+            raise RuntimeError(f"llama.cpp returned HTTP {res.status_code}: {res.text}")
+
+        handle = LlamaCppExecutionHandle(job_id, res)
+        if handle_callback:
+            try:
+                handle_callback(handle)
+            except Exception:
+                pass
+
+        full_text = ""
+        prompt_tokens = 0
+        completion_tokens = 0
+        total_tokens = 0
+        finish_reason = "stop"
+        chunk_idx = 0
+        is_estimated = False
+
+        try:
+            for line in res.iter_lines():
+                if handle._cancelled:
+                    finish_reason = "cancelled"
+                    break
+                line_str = line.strip()
+                if not line_str or not line_str.startswith("data:"):
+                    continue
+                data_str = line_str[5:].strip()
+                if data_str == "[DONE]":
+                    break
+
+                try:
+                    chunk = json.loads(data_str)
+                except Exception:
+                    continue
+
+                choices = chunk.get("choices", [])
+                if choices:
+                    delta = choices[0].get("delta", {}).get("content", "")
+                    if delta:
+                        full_text += delta
+                        if callback:
+                            callback(GenerationChunk(delta=delta, index=chunk_idx))
+                        chunk_idx += 1
+
+                    if choices[0].get("finish_reason"):
+                        finish_reason = choices[0].get("finish_reason")
+
+                usage = chunk.get("usage")
+                if usage:
+                    prompt_tokens = usage.get("prompt_tokens", 0)
+                    completion_tokens = usage.get("completion_tokens", 0)
+
+        finally:
+            handle.mark_done()
+            res.close()
+
+        if prompt_tokens == 0 and completion_tokens == 0:
+            completion_tokens = chunk_idx
+            total_tokens = completion_tokens
+            is_estimated = True
+
+        result = GenerationResult(
+            full_text=full_text,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+            finish_reason=finish_reason,
+            is_estimated=is_estimated,
+        )
+        return result, handle
+
+    def embeddings(
+        self,
+        model: str,
+        inputs: List[str],
+        job_id: Optional[str] = None,
+        handle_callback: Optional[Callable[[ExecutionHandle], None]] = None,
+    ) -> EmbeddingResult:
+        req = self._client.build_request(
+            "POST",
+            "/v1/embeddings",
+            json={"model": model, "input": inputs},
+        )
+        res = self._client.send(req, stream=True)
+        if res.status_code != 200:
+            res.close()
+            raise RuntimeError(f"llama.cpp embeddings failed HTTP {res.status_code}")
+
+        handle = LlamaCppExecutionHandle(job_id or "embedding", res)
+        if handle_callback:
+            try:
+                handle_callback(handle)
+            except Exception:
+                pass
+
+        try:
+            body_bytes = res.read()
+            if handle._cancelled:
+                return EmbeddingResult(embeddings=[], prompt_tokens=0, total_tokens=0, is_estimated=False)
+            data = json.loads(body_bytes)
+            items = data.get("data", [])
+            embeddings_list = [item.get("embedding", []) for item in items]
+            usage = data.get("usage", {})
+            prompt_tokens = usage.get("prompt_tokens", 0)
+            total_tokens = usage.get("total_tokens", prompt_tokens)
+
+            is_estimated = (prompt_tokens == 0)
+            if is_estimated:
+                prompt_tokens = sum(len(t) // 4 for t in inputs)
+                total_tokens = prompt_tokens
+
+            return EmbeddingResult(
+                embeddings=embeddings_list,
+                prompt_tokens=prompt_tokens,
+                total_tokens=total_tokens,
+                is_estimated=is_estimated,
+            )
+        finally:
+            handle.mark_done()
+            res.close()
+
+    def cancel(self, job_id: str) -> None:
+        pass

viento/backends/ollama.py

@@ -0,0 +1,249 @@
+"""
+Production Adapter for Ollama REST Backend (http://localhost:11434).
+
+Supports NDJSON streaming token extraction, ExecutionHandle TCP socket abort cancellation,
+and accurate token count reporting.
+"""
+
+import json
+import logging
+import threading
+from typing import Any, Callable, Dict, List, Optional, Tuple
+import httpx
+
+from viento.backends.base import (
+    EmbeddingResult,
+    ExecutionHandle,
+    GenerationChunk,
+    GenerationResult,
+    InferenceBackend,
+)
+
+logger = logging.getLogger("viento.backends.ollama")
+
+
+class OllamaExecutionHandle(ExecutionHandle):
+    """Execution handle wrapping an active Ollama HTTP streaming response."""
+
+    def __init__(self, job_id: str, response: httpx.Response):
+        self.job_id = job_id
+        self.response = response
+        self._is_done = False
+        self._cancelled = False
+        self._lock = threading.Lock()
+
+    def cancel(self) -> None:
+        """Close response stream instantly, aborting underlying TCP socket."""
+        with self._lock:
+            if not self._is_done and not self._cancelled:
+                self._cancelled = True
+                try:
+                    self.response.close()
+                    logger.info("Closed HTTP response stream for cancelled job %s", self.job_id)
+                except Exception as exc:
+                    logger.warning("Error closing response stream for job %s: %s", self.job_id, exc)
+
+    def is_done(self) -> bool:
+        with self._lock:
+            return self._is_done
+
+    def mark_done(self):
+        with self._lock:
+            self._is_done = True
+
+
+class OllamaAdapter(InferenceBackend):
+    """Ollama backend adapter interfacing with local Ollama service."""
+
+    def __init__(self, base_url: str = "http://localhost:11434", timeout: float = 120.0):
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self._client = httpx.Client(base_url=self.base_url, timeout=self.timeout)
+
+    def name(self) -> str:
+        return "ollama"
+
+    def capabilities(self) -> List[str]:
+        return ["chat", "embeddings", "streaming", "model_pull"]
+
+    def health(self) -> bool:
+        try:
+            res = self._client.get("/")
+            return res.status_code == 200
+        except Exception:
+            return False
+
+    def list_models(self) -> List[Dict[str, Any]]:
+        try:
+            res = self._client.get("/api/tags")
+            if res.status_code != 200:
+                return []
+            data = res.json()
+            models = []
+            for item in data.get("models", []):
+                models.append({
+                    "id": item.get("name"),
+                    "name": item.get("name"),
+                    "status": "ready",
+                    "backend": "ollama",
+                    "context_length": 8192,
+                    "quantization": item.get("details", {}).get("quantization_level", "unknown"),
+                    "capabilities": ["chat", "streaming", "embeddings"],
+                    "max_concurrency": 2,
+                    "active_jobs": 0,
+                })
+            return models
+        except Exception as exc:
+            logger.error("Failed to list Ollama models: %s", exc)
+            return []
+
+    def generate(
+        self,
+        job_id: str,
+        model: str,
+        messages: List[Dict[str, str]],
+        temperature: float = 0.7,
+        max_tokens: int = 512,
+        callback: Optional[Callable[[GenerationChunk], None]] = None,
+        stop: Optional[List[str]] = None,
+        handle_callback: Optional[Callable[[ExecutionHandle], None]] = None,
+    ) -> Tuple[GenerationResult, ExecutionHandle]:
+        payload = {
+            "model": model,
+            "messages": messages,
+            "stream": True,
+            "options": {
+                "temperature": temperature,
+                "num_predict": max_tokens,
+            },
+        }
+        if stop:
+            payload["options"]["stop"] = stop
+
+        req = self._client.build_request("POST", "/api/chat", json=payload)
+        res = self._client.send(req, stream=True)
+
+        if res.status_code != 200:
+            res.close()
+            raise RuntimeError(f"Ollama returned HTTP {res.status_code}: {res.text}")
+
+        handle = OllamaExecutionHandle(job_id, res)
+        if handle_callback:
+            try:
+                handle_callback(handle)
+            except Exception:
+                pass
+
+        full_text = ""
+        prompt_tokens = 0
+        completion_tokens = 0
+        total_tokens = 0
+        finish_reason = "stop"
+        chunk_idx = 0
+
+        try:
+            for line in res.iter_lines():
+                if handle._cancelled:
+                    finish_reason = "cancelled"
+                    break
+                if not line.strip():
+                    continue
+
+                try:
+                    data = json.loads(line)
+                except Exception:
+                    continue
+
+                delta = data.get("message", {}).get("content", "")
+                if delta:
+                    full_text += delta
+                    if callback:
+                        callback(GenerationChunk(delta=delta, index=chunk_idx))
+                    chunk_idx += 1
+
+                if data.get("done"):
+                    prompt_tokens = data.get("prompt_eval_count", 0)
+                    completion_tokens = data.get("eval_count", 0)
+                    total_tokens = prompt_tokens + completion_tokens
+                    finish_reason = data.get("done_reason", "stop")
+                    break
+
+        finally:
+            handle.mark_done()
+            res.close()
+
+        is_estimated = (prompt_tokens == 0 and completion_tokens == 0)
+        if is_estimated:
+            completion_tokens = chunk_idx
+            total_tokens = completion_tokens
+
+        result = GenerationResult(
+            full_text=full_text,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+            finish_reason=finish_reason,
+            is_estimated=is_estimated,
+        )
+        return result, handle
+
+    def embeddings(
+        self,
+        model: str,
+        inputs: List[str],
+        job_id: Optional[str] = None,
+        handle_callback: Optional[Callable[[ExecutionHandle], None]] = None,
+    ) -> EmbeddingResult:
+        all_embeddings: List[List[float]] = []
+        total_prompt_tokens = 0
+        is_estimated = False
+
+        for input_text in inputs:
+            payload = {"model": model, "prompt": input_text}
+            req = self._client.build_request("POST", "/api/embeddings", json=payload)
+            res = self._client.send(req, stream=True)
+
+            if res.status_code != 200:
+                res.close()
+                req = self._client.build_request("POST", "/api/embed", json={"model": model, "input": input_text})
+                res = self._client.send(req, stream=True)
+                if res.status_code != 200:
+                    res.close()
+                    raise RuntimeError(f"Ollama embedding failed HTTP {res.status_code}")
+
+            handle = OllamaExecutionHandle(job_id or "embedding", res)
+            if handle_callback:
+                try:
+                    handle_callback(handle)
+                except Exception:
+                    pass
+
+            try:
+                body_bytes = res.read()
+                if handle._cancelled:
+                    break
+                data = json.loads(body_bytes)
+                emb = data.get("embedding") or (data.get("embeddings", [[]])[0])
+                all_embeddings.append(emb)
+
+                tokens = data.get("prompt_eval_count", 0)
+                if tokens == 0:
+                    tokens = max(1, len(input_text) // 4)
+                    is_estimated = True
+                total_prompt_tokens += tokens
+            finally:
+                handle.mark_done()
+                res.close()
+
+            if handle._cancelled:
+                break
+
+        return EmbeddingResult(
+            embeddings=all_embeddings,
+            prompt_tokens=total_prompt_tokens,
+            total_tokens=total_prompt_tokens,
+            is_estimated=is_estimated,
+        )
+
+    def cancel(self, job_id: str) -> None:
+        pass