Source code for ollama_classifier.backends.llamacpp

"""llama.cpp inference backend.

Supports both local and remote ``llama-server`` instances via the OpenAI-compatible API.

llama.cpp supports GBNF (GGML BNF) grammar constraints, which can express bare
label alternatives directly::

    root ::= "positive" | "negative" | "neutral"

This generates **bare label text** — no JSON wrapper — so logprob
reconstruction is clean. ``llama-server`` accepts a non-standard ``grammar``
field on the ``/v1/chat/completions`` endpoint.

Local server::

    ./llama-server -m model.gguf --host 0.0.0.0 --port 8080 -c 4096

Connect::

    backend = LlamaCppBackend(model="model", base_url="http://localhost:8080/v1")
"""

from typing import Any, Dict, List, Optional

import httpx

from .base import ChatMessage, ChatResponse, LLMBackend, ScoringResponse, Token, TokenLogprob


[docs] class LlamaCppBackend(LLMBackend): """Backend for llama.cpp server (``llama-server``). llama.cpp provides a lightweight inference server with an OpenAI-compatible API. GBNF grammar constraints and logprobs are supported when compiled with the appropriate flags. Logprobs are pre-mask (model's raw distribution before grammar masking). Note: ``response_format`` with JSON schema on ``/v1/chat/completions`` is buggy in llama.cpp (GitHub issues #11988, #11847). This backend uses the non-standard ``grammar`` field with GBNF instead, which works reliably for bare label generation. """
[docs] def __init__( self, model: str, base_url: str = "http://localhost:8080/v1", *, api_key: Optional[str] = None, timeout: float = 120.0, max_tokens: int = 256, extra_body: Optional[Dict[str, Any]] = None, ): super().__init__( model=model, base_url=base_url, api_key=api_key, timeout=timeout, max_tokens=max_tokens, extra_body=extra_body, )
@property def supports_bare_label_constraint(self) -> bool: """True — GBNF grammar generates bare label text.""" return True def _apply_constraint(self, body: Dict[str, Any], labels: List[str]) -> None: """Apply GBNF grammar constraint for bare-label generation. Builds a grammar rule that allows exactly one of the provided labels:: root ::= "label1" | "label2" | "label3" """ # Escape quotes in labels for GBNF quoted = [f'"{l}"' for l in labels] body["grammar"] = f"root ::= {' | '.join(quoted)}" def _render_prompt(self, messages: List[ChatMessage]) -> str: """Render messages to a plain text prompt for the completions endpoint.""" parts: list[str] = [] for m in messages: if m.role == "system": parts.append(f"<|system|>\n{m.content}") elif m.role == "user": parts.append(f"<|user|>\n{m.content}") return "\n\n".join(parts) + "\n\n<|assistant|>\n" # ------------------------------------------------------------------ # Sync # ------------------------------------------------------------------
[docs] def chat( self, messages: List[ChatMessage], *, temperature: float = 0.0, constrain_labels: Optional[List[str]] = None, logprobs: bool = False, top_logprobs: int = 5, ) -> ChatResponse: """Synchronous constrained chat completion via llama.cpp server.""" url = f"{self._base_url}/chat/completions" body = self._build_chat_body( messages, temperature=temperature, constrain_labels=constrain_labels, logprobs=logprobs, top_logprobs=top_logprobs, ) with httpx.Client(timeout=self._timeout) as client: resp = client.post(url, headers=self._build_headers(), json=body) resp.raise_for_status() result = self._parse_chat_response(resp.json()) result.label = result.content.strip() return result
[docs] def score( self, messages: List[ChatMessage], completion: str, ) -> ScoringResponse: """Score a completion using llama.cpp's completions endpoint with suffix. Uses ``/v1/completions`` with ``suffix`` to compute the per-token logprobs of the completion given the prompt context (fill-in-the-middle style scoring). """ prompt = self._render_prompt(messages) url = f"{self._base_url}/completions" body: Dict[str, Any] = { "model": self._model, "prompt": prompt, "suffix": completion, "max_tokens": 0, "temperature": 0.0, "logprobs": 1, "echo": True, **self._extra_body, } with httpx.Client(timeout=self._timeout) as client: resp = client.post(url, headers=self._build_headers(), json=body) resp.raise_for_status() data = resp.json() choice = data["choices"][0] all_logprobs = choice.get("logprobs", {}) tokens_list = all_logprobs.get("tokens", []) token_lps_list = all_logprobs.get("token_logprobs", []) top_lps_list = all_logprobs.get("top_logprobs", []) # Find where the completion starts in the token list # With suffix + echo, llama.cpp returns prompt tokens followed by completion tokens # The completion tokens are at the end # We identify them by matching the completion text completion_start = self._find_completion_start(tokens_list, completion) token_logprobs: list[TokenLogprob] = [] for i in range(completion_start, len(tokens_list)): top: dict[str, float] = {} if i < len(top_lps_list) and top_lps_list[i]: for t, lp in top_lps_list[i].items(): top[t] = lp lp = token_lps_list[i] if i < len(token_lps_list) else 0.0 token_logprobs.append( TokenLogprob(token=tokens_list[i], logprob=lp or 0.0, top_logprobs=top) ) return ScoringResponse(completion=completion, logprobs=token_logprobs, raw=data)
@staticmethod def _find_completion_start(tokens: list[str], completion: str) -> int: """Heuristically find where the completion tokens begin. With suffix + echo, llama.cpp returns [prompt_tokens, completion_tokens]. We search for the suffix position by accumulating token text until it matches the start of the completion string. """ if not tokens: return 0 accumulated = "" for i, tok in enumerate(tokens): accumulated += tok if completion.startswith(accumulated.lstrip()): return i return len(tokens)
[docs] def tokenize( self, text: str, *, context: Optional[str] = None, ) -> List[Token]: """Tokenize text using llama.cpp's tokenize endpoint. llama-server provides a ``/tokenize`` endpoint. """ full_text = (context or "") + text url = f"{self._base_url}/tokenize" body = {"content": full_text} with httpx.Client(timeout=self._timeout) as client: resp = client.post(url, headers=self._build_headers(), json=body) resp.raise_for_status() data = resp.json() tokens_data: list = data.get("tokens", []) if context: ctx_body = {"content": context} with httpx.Client(timeout=self._timeout) as client: ctx_resp = client.post(url, headers=self._build_headers(), json=ctx_body) ctx_resp.raise_for_status() ctx_data = ctx_resp.json() tokens_data = tokens_data[len(ctx_data.get("tokens", [])):] return [ Token( text=t if isinstance(t, str) else "", id=t if isinstance(t, int) else -1, ) for t in tokens_data ]
# ------------------------------------------------------------------ # Async # ------------------------------------------------------------------
[docs] async def achat( self, messages: List[ChatMessage], *, temperature: float = 0.0, constrain_labels: Optional[List[str]] = None, logprobs: bool = False, top_logprobs: int = 5, ) -> ChatResponse: """Async constrained chat completion via llama.cpp server.""" url = f"{self._base_url}/chat/completions" body = self._build_chat_body( messages, temperature=temperature, constrain_labels=constrain_labels, logprobs=logprobs, top_logprobs=top_logprobs, ) async with httpx.AsyncClient(timeout=self._timeout) as client: resp = await client.post(url, headers=self._build_headers(), json=body) resp.raise_for_status() result = self._parse_chat_response(resp.json()) result.label = result.content.strip() return result
[docs] async def ascore( self, messages: List[ChatMessage], completion: str, ) -> ScoringResponse: """Async completion scoring via llama.cpp server.""" prompt = self._render_prompt(messages) url = f"{self._base_url}/completions" body: Dict[str, Any] = { "model": self._model, "prompt": prompt, "suffix": completion, "max_tokens": 0, "temperature": 0.0, "logprobs": 1, "echo": True, **self._extra_body, } async with httpx.AsyncClient(timeout=self._timeout) as client: resp = await client.post(url, headers=self._build_headers(), json=body) resp.raise_for_status() data = resp.json() choice = data["choices"][0] all_logprobs = choice.get("logprobs", {}) tokens_list = all_logprobs.get("tokens", []) token_lps_list = all_logprobs.get("token_logprobs", []) top_lps_list = all_logprobs.get("top_logprobs", []) completion_start = self._find_completion_start(tokens_list, completion) token_logprobs: list[TokenLogprob] = [] for i in range(completion_start, len(tokens_list)): top: dict[str, float] = {} if i < len(top_lps_list) and top_lps_list[i]: for t, lp in top_lps_list[i].items(): top[t] = lp lp = token_lps_list[i] if i < len(token_lps_list) else 0.0 token_logprobs.append( TokenLogprob(token=tokens_list[i], logprob=lp or 0.0, top_logprobs=top) ) return ScoringResponse(completion=completion, logprobs=token_logprobs, raw=data)
[docs] async def atokenize( self, text: str, *, context: Optional[str] = None, ) -> List[Token]: """Async tokenization via llama.cpp server.""" full_text = (context or "") + text url = f"{self._base_url}/tokenize" body = {"content": full_text} async with httpx.AsyncClient(timeout=self._timeout) as client: resp = await client.post(url, headers=self._build_headers(), json=body) resp.raise_for_status() data = resp.json() tokens_data: list = data.get("tokens", []) if context: ctx_body = {"content": context} async with httpx.AsyncClient(timeout=self._timeout) as client: ctx_resp = await client.post(url, headers=self._build_headers(), json=ctx_body) ctx_resp.raise_for_status() ctx_data = ctx_resp.json() tokens_data = tokens_data[len(ctx_data.get("tokens", [])):] return [ Token( text=t if isinstance(t, str) else "", id=t if isinstance(t, int) else -1, ) for t in tokens_data ]