mirror of
https://github.com/mudler/LocalAI.git
synced 2026-10-28 17:18:49 -04:00
pos-ei-don
5392eb813e
With a reasoning parser configured, the streaming path sent every raw text delta as ChatDelta.content. LocalAI uses the backend's ChatDeltas as they are, so with thinking enabled the reasoning arrived as delta.content, and the final chunk then repeated the whole answer and the reasoning a second time. Non-streaming requests were split correctly. Use vLLM's composed parser (ParserManager.get_parser, vLLM >= 0.30), the one vLLM's own OpenAI server streams with: one instance per request, built with the request's chat_template_kwargs (enable_thinking), fed each delta via parse_delta. Reasoning, content and tool calls go out in their own fields as they are produced; the final chunk carries only token counts and logprobs. If the parser raises mid-stream, the rest of the output is buffered and the final chunk sends what the full-text parse has beyond what was already streamed (suffixes per field and per tool-call index). Where the full parse disagrees with what went out, nothing is guessed and the mismatch is logged. Only active when a reasoning_parser is set; without one, or on a vLLM without ParserManager, streaming is unchanged. Signed-off-by: pos-ei-don <1822533+pos-ei-don@users.noreply.github.com>
1273 lines
56 KiB
Python
1273 lines
56 KiB
Python
#!/usr/bin/env python3
|
|
import asyncio
|
|
import dataclasses
|
|
import difflib
|
|
from concurrent import futures
|
|
import argparse
|
|
import signal
|
|
import sys
|
|
import os
|
|
import json
|
|
import time
|
|
import gc
|
|
from typing import List
|
|
from PIL import Image
|
|
|
|
import backend_pb2
|
|
import backend_pb2_grpc
|
|
|
|
import grpc
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
|
from python_utils import attach_media_parts
|
|
from grpc_auth import get_auth_interceptors
|
|
from model_utils import resolve_model_reference
|
|
from temp_utils import materialize_base64
|
|
from vllm_utils import apply_options_to_engine_args, normalize_option_key
|
|
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
from vllm.engine.async_llm_engine import AsyncLLMEngine
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.utils import random_uuid
|
|
try:
|
|
from vllm.tokenizers import get_tokenizer # vLLM >= 0.22
|
|
except ImportError:
|
|
from vllm.transformers_utils.tokenizer import get_tokenizer # vLLM < 0.22
|
|
from vllm.multimodal.utils import fetch_image
|
|
from vllm.assets.video import VideoAsset
|
|
import base64
|
|
import io
|
|
|
|
# Version-compat imports — wrap in try/except for older vLLM versions
|
|
try:
|
|
from vllm.tool_parsers import ToolParserManager
|
|
HAS_TOOL_PARSERS = True
|
|
except ImportError:
|
|
HAS_TOOL_PARSERS = False
|
|
|
|
try:
|
|
from vllm.reasoning import ReasoningParserManager
|
|
HAS_REASONING_PARSERS = True
|
|
except ImportError:
|
|
HAS_REASONING_PARSERS = False
|
|
|
|
# vLLM >= 0.30 composes reasoning and tool parsing into one per-delta parser
|
|
# (ParserManager / Parser.parse_delta) — the one its own OpenAI server uses
|
|
# for streaming.
|
|
try:
|
|
from vllm.parser.parser_manager import ParserManager
|
|
from vllm.entrypoints.openai.chat_completion.protocol import (
|
|
ChatCompletionRequest as _ParserRequest,
|
|
)
|
|
HAS_STREAM_PARSER = True
|
|
except ImportError:
|
|
HAS_STREAM_PARSER = False
|
|
|
|
# vLLM >= 0.23 renamed GuidedDecodingParams -> StructuredOutputsParams and the
|
|
# SamplingParams field guided_decoding -> structured_outputs.
|
|
try:
|
|
from vllm.sampling_params import StructuredOutputsParams
|
|
HAS_GUIDED_DECODING = True
|
|
except ImportError:
|
|
HAS_GUIDED_DECODING = False
|
|
|
|
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
|
|
|
# proto3 has no field presence, so an explicit 0 is indistinguishable from
|
|
# "unset". These two fields have a meaningful zero a caller can intend:
|
|
# temperature 0 is greedy decoding, and 0 is a valid seed.
|
|
_EXPLICIT_ZERO_FIELDS = ("Temperature", "Seed")
|
|
|
|
|
|
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
|
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
|
|
|
# Implement the BackendServicer class with the service methods
|
|
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
|
"""
|
|
A gRPC servicer that implements the Backend service defined in backend.proto.
|
|
"""
|
|
def generate(self,prompt, max_new_tokens):
|
|
"""
|
|
Generates text based on the given prompt and maximum number of new tokens.
|
|
|
|
Args:
|
|
prompt (str): The prompt to generate text from.
|
|
max_new_tokens (int): The maximum number of new tokens to generate.
|
|
|
|
Returns:
|
|
str: The generated text.
|
|
"""
|
|
self.generator.end_beam_search()
|
|
|
|
# Tokenizing the input
|
|
ids = self.generator.tokenizer.encode(prompt)
|
|
|
|
self.generator.gen_begin_reuse(ids)
|
|
initial_len = self.generator.sequence[0].shape[0]
|
|
has_leading_space = False
|
|
decoded_text = ''
|
|
for i in range(max_new_tokens):
|
|
token = self.generator.gen_single_token()
|
|
if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('▁'):
|
|
has_leading_space = True
|
|
|
|
decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])
|
|
if has_leading_space:
|
|
decoded_text = ' ' + decoded_text
|
|
|
|
if token.item() == self.generator.tokenizer.eos_token_id:
|
|
break
|
|
return decoded_text
|
|
|
|
def _parse_options(self, options_list):
|
|
"""Parse Options[] key:value string list into a dict.
|
|
|
|
Keys are normalized to their field spelling so the CLI form users copy
|
|
from vLLM's docs (``--reasoning-parser:qwen3``) selects the same parser
|
|
as LocalAI's own (``reasoning_parser:qwen3``).
|
|
"""
|
|
opts = {}
|
|
for opt in options_list:
|
|
if ":" not in opt:
|
|
continue
|
|
key, value = opt.split(":", 1)
|
|
opts[normalize_option_key(key)] = value.strip()
|
|
return opts
|
|
|
|
def _apply_engine_args(self, engine_args, engine_args_json):
|
|
"""Apply user-supplied engine_args (JSON object) onto an AsyncEngineArgs.
|
|
|
|
Returns a new AsyncEngineArgs with the typed fields preserved and the
|
|
user's overrides layered on top. Uses ``dataclasses.replace`` so vLLM's
|
|
``__post_init__`` re-runs and auto-converts dict-valued fields like
|
|
``compilation_config`` / ``attention_config`` into their dataclass form.
|
|
``speculative_config`` and ``kv_transfer_config`` are accepted as dicts
|
|
directly (vLLM converts them at engine init).
|
|
|
|
Unknown keys raise ValueError with the closest valid field as a hint.
|
|
"""
|
|
if not engine_args_json:
|
|
return engine_args
|
|
try:
|
|
extra = json.loads(engine_args_json)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"engine_args is not valid JSON: {e}") from e
|
|
if not isinstance(extra, dict):
|
|
raise ValueError(
|
|
f"engine_args must be a JSON object, got {type(extra).__name__}"
|
|
)
|
|
valid = {f.name for f in dataclasses.fields(type(engine_args))}
|
|
for key in extra:
|
|
if key not in valid:
|
|
suggestion = difflib.get_close_matches(key, valid, n=1)
|
|
hint = f" did you mean {suggestion[0]!r}?" if suggestion else ""
|
|
raise ValueError(f"unknown engine_args key {key!r}.{hint}")
|
|
return dataclasses.replace(engine_args, **extra)
|
|
|
|
def _messages_to_dicts(self, messages):
|
|
"""Convert proto Messages to list of dicts suitable for apply_chat_template()."""
|
|
result = []
|
|
for msg in messages:
|
|
d = {"role": msg.role, "content": msg.content or ""}
|
|
if msg.name:
|
|
d["name"] = msg.name
|
|
if msg.tool_call_id:
|
|
d["tool_call_id"] = msg.tool_call_id
|
|
if msg.reasoning_content:
|
|
d["reasoning_content"] = msg.reasoning_content
|
|
if msg.tool_calls:
|
|
try:
|
|
tool_calls = json.loads(msg.tool_calls)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
else:
|
|
# OpenAI wire format carries function.arguments as a
|
|
# JSON-encoded string, but chat templates (e.g. Qwen3)
|
|
# iterate over it as a mapping. vLLM's own OpenAI server
|
|
# parses arguments before applying the template, so do
|
|
# the same here.
|
|
if isinstance(tool_calls, list):
|
|
for tc in tool_calls:
|
|
func = tc.get("function") if isinstance(tc, dict) else None
|
|
if isinstance(func, dict) and isinstance(func.get("arguments"), str):
|
|
try:
|
|
func["arguments"] = json.loads(func["arguments"])
|
|
except json.JSONDecodeError:
|
|
pass
|
|
d["tool_calls"] = tool_calls
|
|
result.append(d)
|
|
return result
|
|
|
|
def Health(self, request, context):
|
|
"""
|
|
Returns a health check message.
|
|
|
|
Args:
|
|
request: The health check request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Reply: The health check reply.
|
|
"""
|
|
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
|
|
|
|
async def LoadModel(self, request, context):
|
|
"""
|
|
Loads a language model.
|
|
|
|
Args:
|
|
request: The load model request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Result: The load model result.
|
|
"""
|
|
model_ref, _local_only = resolve_model_reference(request)
|
|
engine_args = AsyncEngineArgs(model=model_ref)
|
|
|
|
if request.Quantization != "":
|
|
engine_args.quantization = request.Quantization
|
|
if request.LoadFormat != "":
|
|
engine_args.load_format = request.LoadFormat
|
|
if request.GPUMemoryUtilization != 0:
|
|
engine_args.gpu_memory_utilization = request.GPUMemoryUtilization
|
|
if request.TrustRemoteCode:
|
|
engine_args.trust_remote_code = request.TrustRemoteCode
|
|
if request.EnforceEager:
|
|
engine_args.enforce_eager = request.EnforceEager
|
|
if request.TensorParallelSize:
|
|
engine_args.tensor_parallel_size = request.TensorParallelSize
|
|
if request.SwapSpace != 0:
|
|
engine_args.swap_space = request.SwapSpace
|
|
if request.MaxModelLen != 0:
|
|
engine_args.max_model_len = request.MaxModelLen
|
|
if request.DisableLogStatus:
|
|
engine_args.disable_log_status = request.DisableLogStatus
|
|
if request.DType != "":
|
|
engine_args.dtype = request.DType
|
|
if request.LimitImagePerPrompt != 0 or request.LimitVideoPerPrompt != 0 or request.LimitAudioPerPrompt != 0:
|
|
# limit-mm-per-prompt defaults to 1 per modality, based on vLLM docs
|
|
engine_args.limit_mm_per_prompt = {
|
|
"image": max(request.LimitImagePerPrompt, 1),
|
|
"video": max(request.LimitVideoPerPrompt, 1),
|
|
"audio": max(request.LimitAudioPerPrompt, 1)
|
|
}
|
|
|
|
# CLI-style flags in options: (--quantization:gptq_marlin,
|
|
# --enable-prefix-caching, ...) land on the engine args too - they must
|
|
# be applied *before* the engine is created or they do nothing.
|
|
# engine_args: is applied after this, so it stays the last word.
|
|
engine_args = apply_options_to_engine_args(engine_args, request.Options)
|
|
|
|
# engine_args from YAML overrides typed fields above so operators can
|
|
# tune anything the AsyncEngineArgs dataclass exposes without waiting
|
|
# on protobuf changes.
|
|
try:
|
|
engine_args = self._apply_engine_args(engine_args, request.EngineArgs)
|
|
except ValueError as err:
|
|
print(f"engine_args error: {err}", file=sys.stderr)
|
|
return backend_pb2.Result(success=False, message=str(err))
|
|
|
|
try:
|
|
self.llm = AsyncLLMEngine.from_engine_args(engine_args)
|
|
except Exception as err:
|
|
print(f"Unexpected {err=}, {type(err)=}", file=sys.stderr)
|
|
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
|
|
|
try:
|
|
# vLLM >= 0.14 removed get_model_config() on AsyncLLM; the tokenizer
|
|
# is either already loaded on the engine or can be built from the
|
|
# Model name directly.
|
|
tokenizer = None
|
|
if hasattr(self.llm, "get_tokenizer"):
|
|
try:
|
|
tokenizer = await self.llm.get_tokenizer()
|
|
except TypeError:
|
|
tokenizer = self.llm.get_tokenizer()
|
|
except Exception:
|
|
tokenizer = None
|
|
if tokenizer is None and hasattr(self.llm, "tokenizer"):
|
|
tokenizer = self.llm.tokenizer
|
|
if tokenizer is None:
|
|
tokenizer = get_tokenizer(
|
|
model_ref,
|
|
trust_remote_code=bool(request.TrustRemoteCode),
|
|
truncation_side="left",
|
|
)
|
|
self.tokenizer = tokenizer
|
|
except Exception as err:
|
|
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
|
|
|
# Parse options for parser selection
|
|
opts = self._parse_options(request.Options)
|
|
|
|
# Instantiate tool/reasoning parser classes (they'll be instantiated per-request with tokenizer)
|
|
self.tool_parser_cls = None
|
|
self.reasoning_parser_cls = None
|
|
if HAS_TOOL_PARSERS and opts.get("tool_parser"):
|
|
try:
|
|
self.tool_parser_cls = ToolParserManager.get_tool_parser(opts["tool_parser"])
|
|
print(f"Loaded tool_parser: {opts['tool_parser']}", file=sys.stderr)
|
|
except Exception as e:
|
|
print(f"Failed to load tool_parser {opts.get('tool_parser')}: {e}", file=sys.stderr)
|
|
|
|
if HAS_REASONING_PARSERS and opts.get("reasoning_parser"):
|
|
try:
|
|
self.reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(opts["reasoning_parser"])
|
|
print(f"Loaded reasoning_parser: {opts['reasoning_parser']}", file=sys.stderr)
|
|
except Exception as e:
|
|
print(f"Failed to load reasoning_parser {opts.get('reasoning_parser')}: {e}", file=sys.stderr)
|
|
|
|
# With a reasoning parser, streaming has to split reasoning from content
|
|
# per delta. Use vLLM's composed parser for that (reasoning + tools).
|
|
self.stream_parser_cls = None
|
|
if HAS_STREAM_PARSER and self.reasoning_parser_cls is not None:
|
|
try:
|
|
self.stream_parser_cls = ParserManager.get_parser(
|
|
tool_parser_name=opts.get("tool_parser") if self.tool_parser_cls else None,
|
|
reasoning_parser_name=opts.get("reasoning_parser"),
|
|
enable_auto_tools=True,
|
|
)
|
|
except Exception as e:
|
|
print(f"Streaming parser unavailable, streaming raw deltas: {e!r}", file=sys.stderr)
|
|
self.stream_parser_cls = None
|
|
|
|
print("Model loaded successfully", file=sys.stderr)
|
|
return backend_pb2.Result(message="Model loaded successfully", success=True)
|
|
|
|
async def Predict(self, request, context):
|
|
"""
|
|
Generates text based on the given prompt and sampling parameters.
|
|
|
|
Args:
|
|
request: The predict request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Reply: The predict result.
|
|
"""
|
|
gen = self._predict(request, context, streaming=False)
|
|
res = await gen.__anext__()
|
|
return res
|
|
|
|
def Embedding(self, request, context):
|
|
"""
|
|
A gRPC method that calculates embeddings for a given sentence.
|
|
|
|
Args:
|
|
request: An EmbeddingRequest object that contains the request parameters.
|
|
context: A grpc.ServicerContext object that provides information about the RPC.
|
|
|
|
Returns:
|
|
An EmbeddingResult object that contains the calculated embeddings.
|
|
"""
|
|
print("Calculated embeddings for: " + request.Embeddings, file=sys.stderr)
|
|
outputs = self.model.encode(request.Embeddings)
|
|
# Check if we have one result at least
|
|
if len(outputs) == 0:
|
|
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
|
context.set_details("No embeddings were calculated.")
|
|
return backend_pb2.EmbeddingResult()
|
|
return backend_pb2.EmbeddingResult(
|
|
embeddings=outputs[0].outputs.embedding,
|
|
layout=backend_pb2.EMBEDDING_LAYOUT_FINAL,
|
|
)
|
|
|
|
async def PredictStream(self, request, context):
|
|
"""
|
|
Generates text based on the given prompt and sampling parameters, and streams the results.
|
|
|
|
Args:
|
|
request: The predict stream request.
|
|
context: The gRPC context.
|
|
|
|
Returns:
|
|
backend_pb2.Result: The predict stream result.
|
|
"""
|
|
iterations = self._predict(request, context, streaming=True)
|
|
try:
|
|
async for iteration in iterations:
|
|
yield iteration
|
|
finally:
|
|
await iterations.aclose()
|
|
|
|
async def TokenizeString(self, request, context):
|
|
if not hasattr(self, 'tokenizer') or self.tokenizer is None:
|
|
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
|
context.set_details("Model/tokenizer not loaded")
|
|
return backend_pb2.TokenizationResponse()
|
|
try:
|
|
tokens = self.tokenizer.encode(request.Prompt)
|
|
return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
|
|
except Exception as e:
|
|
context.set_code(grpc.StatusCode.INTERNAL)
|
|
context.set_details(str(e))
|
|
return backend_pb2.TokenizationResponse()
|
|
|
|
async def Free(self, request, context):
|
|
try:
|
|
if hasattr(self, 'llm'):
|
|
del self.llm
|
|
if hasattr(self, 'tokenizer'):
|
|
del self.tokenizer
|
|
self.tool_parser_cls = None
|
|
self.reasoning_parser_cls = None
|
|
self.stream_parser_cls = None
|
|
gc.collect()
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except ImportError:
|
|
pass
|
|
return backend_pb2.Result(success=True, message="Model freed")
|
|
except Exception as e:
|
|
return backend_pb2.Result(success=False, message=str(e))
|
|
|
|
async def Score(self, request, context):
|
|
"""
|
|
Joint log-probability of each candidate continuation given the
|
|
shared prompt. Used by routing-policy multi-label classification
|
|
(read the distribution rather than asking the model to emit a
|
|
single argmax label), reranking, and reward-model scoring.
|
|
|
|
Implementation uses vLLM's `prompt_logprobs` to recover the
|
|
per-token log P(token_i | tokens_<i) for the full concatenated
|
|
sequence; the candidate's tokens are the suffix whose logprobs
|
|
get summed. max_tokens=1 because vLLM requires at least one
|
|
generated token; the generated token is discarded.
|
|
"""
|
|
if not hasattr(self, 'llm') or self.llm is None:
|
|
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
|
context.set_details("Model not loaded")
|
|
return backend_pb2.ScoreResponse()
|
|
if not hasattr(self, 'tokenizer') or self.tokenizer is None:
|
|
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
|
context.set_details("Tokenizer not available")
|
|
return backend_pb2.ScoreResponse()
|
|
if len(request.candidates) == 0:
|
|
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
|
context.set_details("candidates must be non-empty")
|
|
return backend_pb2.ScoreResponse()
|
|
|
|
try:
|
|
prompt = request.prompt or ""
|
|
prompt_token_ids = self.tokenizer.encode(prompt)
|
|
prompt_len = len(prompt_token_ids)
|
|
results = []
|
|
|
|
for candidate in request.candidates:
|
|
# Tokenise the concatenated sequence. We can't naively
|
|
# use len(prompt_tokens) + len(tokenizer.encode(candidate))
|
|
# because BPE merges at the boundary may produce a
|
|
# different tokenisation. Encoding the joined text and
|
|
# walking the divergence point is the correct primitive.
|
|
full_text = prompt + candidate
|
|
full_token_ids = self.tokenizer.encode(full_text)
|
|
|
|
divergence = prompt_len
|
|
min_len = min(prompt_len, len(full_token_ids))
|
|
for i in range(min_len):
|
|
if prompt_token_ids[i] != full_token_ids[i]:
|
|
divergence = i
|
|
break
|
|
|
|
candidate_token_ids = full_token_ids[divergence:]
|
|
num_candidate_tokens = len(candidate_token_ids)
|
|
if num_candidate_tokens == 0:
|
|
results.append(backend_pb2.CandidateScore(
|
|
log_prob=0.0,
|
|
length_normalized_log_prob=0.0,
|
|
num_tokens=0,
|
|
))
|
|
continue
|
|
|
|
sampling = SamplingParams(
|
|
max_tokens=1,
|
|
temperature=0.0,
|
|
prompt_logprobs=1,
|
|
detokenize=False,
|
|
)
|
|
|
|
request_id = random_uuid()
|
|
last_output = None
|
|
outputs_iter = self.llm.generate(
|
|
{"prompt": full_text},
|
|
sampling_params=sampling,
|
|
request_id=request_id,
|
|
)
|
|
try:
|
|
async for out in outputs_iter:
|
|
last_output = out
|
|
finally:
|
|
try:
|
|
await outputs_iter.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
_pl = getattr(last_output, "prompt_logprobs", None) if last_output is not None else None
|
|
# Some engines accept the prompt_logprobs request but return a
|
|
# list of all-None entries instead of computing them (observed
|
|
# with vllm-metal's MLX backend on macOS). Treat that as
|
|
# unsupported rather than silently scoring every candidate as 0.
|
|
if not _pl or all(e is None for e in _pl):
|
|
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
|
context.set_details("This backend did not return prompt_logprobs; scoring is unsupported on this engine (e.g. vllm-metal / MLX on macOS).")
|
|
return backend_pb2.ScoreResponse()
|
|
|
|
prompt_logprobs = last_output.prompt_logprobs
|
|
total = 0.0
|
|
tokens_proto = []
|
|
for offset, tok_id in enumerate(candidate_token_ids):
|
|
position = divergence + offset
|
|
if position >= len(prompt_logprobs) or prompt_logprobs[position] is None:
|
|
continue
|
|
entry = prompt_logprobs[position]
|
|
lp_obj = entry.get(tok_id)
|
|
if lp_obj is not None:
|
|
lp = lp_obj.logprob
|
|
else:
|
|
# Token not in top-K; vLLM's top-1 may miss it.
|
|
# Fall back to the lowest available logprob in the
|
|
# entry — a conservative lower-bound on the true
|
|
# log P, biased against this candidate.
|
|
lp = min(v.logprob for v in entry.values())
|
|
total += lp
|
|
if request.include_token_logprobs:
|
|
tokens_proto.append(backend_pb2.TokenLogProb(
|
|
token=self.tokenizer.decode([tok_id]),
|
|
log_prob=lp,
|
|
))
|
|
|
|
cs = backend_pb2.CandidateScore(
|
|
log_prob=total,
|
|
num_tokens=num_candidate_tokens,
|
|
)
|
|
if request.length_normalize and num_candidate_tokens > 0:
|
|
cs.length_normalized_log_prob = total / num_candidate_tokens
|
|
if tokens_proto:
|
|
cs.tokens.extend(tokens_proto)
|
|
results.append(cs)
|
|
|
|
return backend_pb2.ScoreResponse(candidates=results)
|
|
except Exception as e:
|
|
print(f"Score error: {e}", file=sys.stderr)
|
|
context.set_code(grpc.StatusCode.INTERNAL)
|
|
context.set_details(str(e))
|
|
return backend_pb2.ScoreResponse()
|
|
|
|
def _build_sampling_params(self, request):
|
|
request_to_sampling_params = {
|
|
"N": "n",
|
|
"PresencePenalty": "presence_penalty",
|
|
"FrequencyPenalty": "frequency_penalty",
|
|
"RepetitionPenalty": "repetition_penalty",
|
|
"Temperature": "temperature",
|
|
"TopP": "top_p",
|
|
"TopK": "top_k",
|
|
"MinP": "min_p",
|
|
"Seed": "seed",
|
|
"StopPrompts": "stop",
|
|
"StopTokenIds": "stop_token_ids",
|
|
"BadWords": "bad_words",
|
|
"IncludeStopStrInOutput": "include_stop_str_in_output",
|
|
"IgnoreEOS": "ignore_eos",
|
|
"Tokens": "max_tokens",
|
|
"MinTokens": "min_tokens",
|
|
"Logprobs": "logprobs",
|
|
"PromptLogprobs": "prompt_logprobs",
|
|
"SkipSpecialTokens": "skip_special_tokens",
|
|
"SpacesBetweenSpecialTokens": "spaces_between_special_tokens",
|
|
"TruncatePromptTokens": "truncate_prompt_tokens",
|
|
}
|
|
|
|
sampling_params = SamplingParams(top_p=0.9, max_tokens=200)
|
|
|
|
for request_field, param_field in request_to_sampling_params.items():
|
|
if hasattr(request, request_field):
|
|
value = getattr(request, request_field)
|
|
# See _EXPLICIT_ZERO_FIELDS: temperature 0 is greedy decoding
|
|
# and 0 is a valid seed, so neither may be filtered out.
|
|
if request_field in _EXPLICIT_ZERO_FIELDS or value not in (None, 0, [], False):
|
|
setattr(sampling_params, param_field, value)
|
|
|
|
return sampling_params
|
|
|
|
def _new_reasoning_parser(self, chat_template_kwargs):
|
|
"""Build the reasoning parser, telling it whether thinking is on.
|
|
|
|
vLLM's newer parser engines decide their *initial state* from
|
|
``chat_template_kwargs``: ``Qwen3Parser`` reads
|
|
``chat_template_kwargs["enable_thinking"]`` and defaults to ``True``,
|
|
starting in the REASONING state. Constructed without it, a completion
|
|
produced with thinking disabled is classified as reasoning end to end,
|
|
and the answer is reported in both ``reasoning_content`` and
|
|
``content``.
|
|
|
|
vLLM's own OpenAI server forwards the request's chat template kwargs
|
|
here; this backend renders the template itself, so it forwards the
|
|
same dict. Older parsers do not accept the argument — fall back to the
|
|
plain constructor for those.
|
|
"""
|
|
try:
|
|
return self.reasoning_parser_cls(
|
|
self.tokenizer, chat_template_kwargs=chat_template_kwargs or {},
|
|
)
|
|
except TypeError:
|
|
return self.reasoning_parser_cls(self.tokenizer)
|
|
|
|
@staticmethod
|
|
def _split_reasoning(rp, generated_text, prompt, reasoning, content):
|
|
"""Decide what the reasoning parser's output actually means.
|
|
|
|
Covers the *older* parser shape, which has no initial state to set:
|
|
``BaseThinkingReasoningParser.extract_reasoning`` documents its own
|
|
fallback — "For models that may not generate start token, assume the
|
|
reasoning content is always at the start." When no end token is
|
|
present it returns *everything* as reasoning and ``None`` as content,
|
|
which is right for a truncated reasoning run and wrong for a
|
|
completion that never contained reasoning at all.
|
|
|
|
Taking ``None`` content to mean "keep the raw text" then duplicates
|
|
the answer into both fields.
|
|
|
|
The prompt says which case it is. A template with thinking on leaves
|
|
the reasoning block open (the prompt ends with the start token); with
|
|
thinking off it closes the block in the prompt, so the completion is
|
|
plain content. Parsers that expose no token pair (the engine-based
|
|
adapters, which take the ``chat_template_kwargs`` route above) keep
|
|
the parser's verdict unchanged.
|
|
"""
|
|
start = getattr(rp, "start_token", None)
|
|
end = getattr(rp, "end_token", None)
|
|
|
|
if end and end in generated_text:
|
|
# The parser split on the end token. Empty content here means the
|
|
# model stopped right after it, not that parsing failed.
|
|
return reasoning or "", content or ""
|
|
|
|
if not start:
|
|
# Unknown token layout — keep the previous behaviour rather than
|
|
# guess.
|
|
return reasoning or "", content if content is not None else generated_text
|
|
|
|
if not (start in generated_text or (prompt or "").rstrip().endswith(start)):
|
|
# No end token and the block was never open: the "reasoning starts
|
|
# at the beginning" fallback does not apply to this completion.
|
|
return "", generated_text
|
|
|
|
# Block was open and the end token never arrived — reasoning ran out of
|
|
# budget. It is all reasoning, and there is no answer to report.
|
|
return reasoning or "", content or ""
|
|
|
|
def _new_stream_parser(self, request, template_kwargs):
|
|
"""Build vLLM's per-delta parser for one streaming request.
|
|
|
|
Same construction as vLLM's OpenAI server: the request's
|
|
``chat_template_kwargs`` (``enable_thinking``) set the parser's initial
|
|
state, and the prompt token ids passed to the first ``parse_delta``
|
|
tell it whether the reasoning block is open. Returns ``(None, None)``
|
|
when that is not possible; the caller then streams raw deltas as
|
|
before.
|
|
"""
|
|
cls = getattr(self, "stream_parser_cls", None)
|
|
if cls is None:
|
|
return None, None
|
|
try:
|
|
tools = None
|
|
if request.Tools:
|
|
try:
|
|
tools = json.loads(request.Tools) or None
|
|
except json.JSONDecodeError:
|
|
tools = None
|
|
parser_request = _ParserRequest(
|
|
model="local",
|
|
messages=[{"role": "user", "content": ""}],
|
|
tools=tools,
|
|
)
|
|
chat_template_kwargs = {
|
|
k: v for k, v in (template_kwargs or {}).items()
|
|
if k not in ("tokenize", "add_generation_prompt", "tools")
|
|
}
|
|
parser = cls(self.tokenizer, parser_request.tools,
|
|
chat_template_kwargs=chat_template_kwargs)
|
|
return parser, parser_request
|
|
except Exception as e:
|
|
print(f"Could not build streaming parser, streaming raw deltas: {e!r}",
|
|
file=sys.stderr)
|
|
return None, None
|
|
|
|
@staticmethod
|
|
def _delta_message_to_reply(msg):
|
|
"""vLLM ``DeltaMessage`` -> ``Reply``, or ``None`` if it carries nothing."""
|
|
tc_protos = []
|
|
for tc in (msg.tool_calls or []):
|
|
fn = tc.function or None
|
|
tc_protos.append(backend_pb2.ToolCallDelta(
|
|
index=tc.index,
|
|
id=tc.id or "",
|
|
name=(fn.name if fn and fn.name else "") or "",
|
|
arguments=(fn.arguments if fn and fn.arguments else "") or "",
|
|
))
|
|
cd_kwargs = {}
|
|
if msg.content:
|
|
cd_kwargs["content"] = msg.content
|
|
if msg.reasoning:
|
|
cd_kwargs["reasoning_content"] = msg.reasoning
|
|
if tc_protos:
|
|
cd_kwargs["tool_calls"] = tc_protos
|
|
if not cd_kwargs:
|
|
return None
|
|
return backend_pb2.Reply(
|
|
message=bytes(msg.content or "", encoding='utf-8'),
|
|
chat_deltas=[backend_pb2.ChatDelta(**cd_kwargs)],
|
|
)
|
|
|
|
@staticmethod
|
|
def _unsent_remainder(sent_content, sent_reasoning, sent_tools,
|
|
content, reasoning, tool_calls):
|
|
"""After the streaming parser failed mid-stream: what is still unsent.
|
|
|
|
``content``/``reasoning``/``tool_calls`` come from parsing the full
|
|
text; ``sent_tools`` maps a tool-call index to the name and arguments
|
|
already streamed for it. Only suffixes of what was sent are added. If
|
|
the full parse disagrees with what went out, nothing is guessed: the
|
|
field is left as streamed and the mismatch is logged.
|
|
"""
|
|
def unsent(field, full, sent):
|
|
full = full or ""
|
|
if full.startswith(sent):
|
|
return full[len(sent):]
|
|
print(f"Streamed {field} does not match the final parse; "
|
|
f"not sending the rest", file=sys.stderr)
|
|
return ""
|
|
|
|
remainder = {}
|
|
rest = unsent("reasoning", reasoning, sent_reasoning)
|
|
if rest:
|
|
remainder["reasoning_content"] = rest
|
|
if sent_tools and not tool_calls:
|
|
# A tool call was streamed but the full text does not parse as
|
|
# one: the rest of the content would be raw tool-call markup.
|
|
print("Streamed tool call not found in the final parse; "
|
|
"not sending the rest", file=sys.stderr)
|
|
else:
|
|
rest = unsent("content", content, sent_content)
|
|
if rest:
|
|
remainder["content"] = rest
|
|
tc_protos = []
|
|
for i, (call_id, name, arguments) in enumerate(tool_calls or []):
|
|
sent = sent_tools.get(i)
|
|
if sent is None:
|
|
tc_protos.append(backend_pb2.ToolCallDelta(
|
|
index=i, id=call_id or "", name=name or "", arguments=arguments or "",
|
|
))
|
|
continue
|
|
if sent["name"] and name and sent["name"] != name:
|
|
print(f"Streamed tool call {i} does not match the final parse; "
|
|
f"not sending the rest", file=sys.stderr)
|
|
continue
|
|
rest = unsent(f"tool call {i} arguments", arguments, sent["arguments"])
|
|
if rest:
|
|
tc_protos.append(backend_pb2.ToolCallDelta(index=i, arguments=rest))
|
|
if tc_protos:
|
|
remainder["tool_calls"] = tc_protos
|
|
return remainder
|
|
|
|
async def _predict(self, request, context, streaming=False):
|
|
# Build the sampling parameters
|
|
sampling_params = self._build_sampling_params(request)
|
|
|
|
# Structured-output decoding: use Grammar field to pass JSON schema or BNF
|
|
if HAS_GUIDED_DECODING and request.Grammar:
|
|
try:
|
|
json.loads(request.Grammar) # valid JSON = JSON schema
|
|
sampling_params.structured_outputs = StructuredOutputsParams(json=request.Grammar)
|
|
except json.JSONDecodeError:
|
|
sampling_params.structured_outputs = StructuredOutputsParams(grammar=request.Grammar)
|
|
|
|
# Extract image paths and process images
|
|
prompt = request.Prompt
|
|
# Kept in scope: the reasoning parser needs to know which chat
|
|
# template kwargs produced this prompt.
|
|
template_kwargs = {}
|
|
|
|
image_paths = request.Images
|
|
image_data = [self.load_image(img_path) for img_path in image_paths]
|
|
|
|
videos_path = request.Videos
|
|
video_data = [self.load_video(video_path) for video_path in videos_path]
|
|
|
|
# If tokenizer template is enabled and messages are provided instead of prompt, apply the tokenizer template
|
|
if not request.Prompt and request.UseTokenizerTemplate and request.Messages:
|
|
messages_dicts = self._messages_to_dicts(request.Messages)
|
|
template_kwargs.update({"tokenize": False, "add_generation_prompt": True})
|
|
|
|
# Pass tools for tool calling
|
|
if request.Tools:
|
|
try:
|
|
template_kwargs["tools"] = json.loads(request.Tools)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
_thinking = request.Metadata.get("enable_thinking", "").lower()
|
|
if _thinking in ("true", "false"):
|
|
template_kwargs["enable_thinking"] = (_thinking == "true")
|
|
|
|
# vLLM substitutes multi_modal_data into the model's own media
|
|
# token, so the template has to be given content *parts* - string
|
|
# content renders a prompt with no placeholder and the media are
|
|
# dropped without a word (#11621).
|
|
prompt = None
|
|
media_dicts = attach_media_parts(
|
|
messages_dicts, len(image_data), len(video_data)
|
|
)
|
|
if media_dicts is not None:
|
|
try:
|
|
prompt = self.tokenizer.apply_chat_template(media_dicts, **template_kwargs)
|
|
except Exception as e:
|
|
# A text-only template cannot iterate content parts; fall
|
|
# through to the text-only prompt instead of failing.
|
|
print(
|
|
f"chat template rejected multimodal content parts: {e!r}",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
if prompt is None:
|
|
try:
|
|
prompt = self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
|
except TypeError:
|
|
# Some tokenizers don't support tools/enable_thinking kwargs — retry without them
|
|
prompt = self.tokenizer.apply_chat_template(
|
|
messages_dicts, tokenize=False, add_generation_prompt=True
|
|
)
|
|
|
|
# Generate text using the LLM engine
|
|
request_id = random_uuid()
|
|
print(f"Generating text with request_id: {request_id}", file=sys.stderr)
|
|
multi_modal_data = {}
|
|
if image_data:
|
|
multi_modal_data["image"] = image_data
|
|
if video_data:
|
|
multi_modal_data["video"] = video_data
|
|
outputs = self.llm.generate(
|
|
{
|
|
"prompt": prompt,
|
|
"multi_modal_data": multi_modal_data if multi_modal_data else None,
|
|
},
|
|
sampling_params=sampling_params,
|
|
request_id=request_id,
|
|
)
|
|
|
|
# Stream the results
|
|
generated_text = ""
|
|
generated_token_ids: list[int] = []
|
|
last_output = None
|
|
|
|
# Tool-parsing strategy decision (made once, before the loop):
|
|
#
|
|
# When a tool parser is active, the model's raw tool-call markup
|
|
# (e.g. <tool_call>...) must not be streamed verbatim as delta.content
|
|
# — clients would see the unparsed syntax. Two paths:
|
|
#
|
|
# (A) native streaming via parser.extract_tool_calls_streaming. All
|
|
# concrete tool parsers shipped with vLLM 0.23+ implement this
|
|
# (Granite4, Qwen3Coder, DeepSeekV31, Jamba, Ernie45, Hermes,
|
|
# llama3_json, mistral, …). The parser decides per-delta whether
|
|
# to emit content or suppress tool-call markup, and emits a
|
|
# structured DeltaMessage(tool_calls=[...]) when a call is ready.
|
|
# (B) buffer fallback — used only when the parser surprisingly lacks
|
|
# the streaming method or it raises mid-stream. The post-loop
|
|
# extract_tool_calls assembles the final chat_delta. Same correctness
|
|
# guarantee as a non-streaming response, at the cost of a delayed
|
|
# final chunk.
|
|
has_tool_parser = bool(self.tool_parser_cls and request.Tools)
|
|
tp_instance = None
|
|
tp_request = None
|
|
native_streaming = False
|
|
native_streaming_error = False
|
|
if has_tool_parser:
|
|
try:
|
|
tools_for_parser = json.loads(request.Tools)
|
|
except json.JSONDecodeError:
|
|
tools_for_parser = []
|
|
try:
|
|
tp_instance = self.tool_parser_cls(self.tokenizer, tools=tools_for_parser)
|
|
except TypeError:
|
|
tp_instance = self.tool_parser_cls(self.tokenizer)
|
|
# Build a minimal ChatCompletionRequest so the streaming method
|
|
# sees the tools list. We do not need any other request fields —
|
|
# parsers only read .tools (and sometimes .tool_choice, which we
|
|
# leave at default).
|
|
try:
|
|
from vllm.entrypoints.openai.chat_completion.protocol import (
|
|
ChatCompletionRequest as _CCR,
|
|
)
|
|
tp_request = _CCR(
|
|
model="local",
|
|
messages=[{"role": "user", "content": ""}],
|
|
tools=tools_for_parser or None,
|
|
)
|
|
except Exception as e:
|
|
print(f"Could not build ChatCompletionRequest for streaming parser: {e}",
|
|
file=sys.stderr)
|
|
tp_request = None
|
|
native_streaming = (
|
|
tp_request is not None
|
|
and hasattr(tp_instance, "extract_tool_calls_streaming")
|
|
)
|
|
|
|
# With a reasoning parser, split reasoning, content and tool calls per
|
|
# delta, as vLLM's OpenAI server does. Without this the raw text —
|
|
# reasoning included — was streamed as content, and LocalAI takes the
|
|
# backend's ChatDeltas as they are.
|
|
stream_parser, parser_request = (
|
|
self._new_stream_parser(request, template_kwargs) if streaming else (None, None)
|
|
)
|
|
stream_parser_error = False
|
|
sent_content = ""
|
|
sent_reasoning = ""
|
|
sent_tools = {} # tool-call index -> {"name", "arguments"} already streamed
|
|
|
|
try:
|
|
async for request_output in outputs:
|
|
iteration_text = request_output.outputs[0].text
|
|
last_output = request_output
|
|
|
|
if streaming and stream_parser is not None and not stream_parser_error:
|
|
delta_iteration_text = iteration_text.removeprefix(generated_text)
|
|
new_token_ids = list(request_output.outputs[0].token_ids)
|
|
delta_token_ids = new_token_ids[len(generated_token_ids):]
|
|
try:
|
|
msg = stream_parser.parse_delta(
|
|
delta_text=delta_iteration_text,
|
|
delta_token_ids=delta_token_ids,
|
|
request=parser_request,
|
|
prompt_token_ids=list(request_output.prompt_token_ids or []),
|
|
finished=bool(
|
|
getattr(request_output, "finished", False)
|
|
or getattr(request_output.outputs[0], "finish_reason", None) is not None
|
|
),
|
|
)
|
|
except Exception as e:
|
|
print(f"Streaming parser error (sending the rest with the "
|
|
f"final chunk): {e!r}", file=sys.stderr)
|
|
stream_parser_error = True
|
|
msg = None
|
|
if msg is not None:
|
|
reply = self._delta_message_to_reply(msg)
|
|
if reply is not None:
|
|
sent_content += msg.content or ""
|
|
sent_reasoning += msg.reasoning or ""
|
|
for tc in (msg.tool_calls or []):
|
|
fn = tc.function or None
|
|
sent = sent_tools.setdefault(tc.index, {"name": "", "arguments": ""})
|
|
if fn and fn.name and not sent["name"]:
|
|
sent["name"] = fn.name # sent once, as in OpenAI deltas
|
|
sent["arguments"] += (fn.arguments if fn and fn.arguments else "") or ""
|
|
yield reply
|
|
elif streaming and stream_parser is not None:
|
|
pass # parser failed: buffer, the final chunk carries the rest
|
|
elif streaming:
|
|
delta_iteration_text = iteration_text.removeprefix(generated_text)
|
|
new_token_ids = list(request_output.outputs[0].token_ids)
|
|
delta_token_ids = new_token_ids[len(generated_token_ids):]
|
|
|
|
if not has_tool_parser:
|
|
# Plain streaming — unchanged from pre-tool-parser path.
|
|
yield backend_pb2.Reply(
|
|
message=bytes(delta_iteration_text, encoding='utf-8'),
|
|
chat_deltas=[backend_pb2.ChatDelta(content=delta_iteration_text)],
|
|
)
|
|
elif native_streaming and not native_streaming_error:
|
|
# (A) Native vLLM extract_tool_calls_streaming.
|
|
try:
|
|
msg = tp_instance.extract_tool_calls_streaming(
|
|
previous_text=generated_text,
|
|
current_text=iteration_text,
|
|
delta_text=delta_iteration_text,
|
|
previous_token_ids=generated_token_ids,
|
|
current_token_ids=new_token_ids,
|
|
delta_token_ids=delta_token_ids,
|
|
request=tp_request,
|
|
)
|
|
except Exception as e:
|
|
print(f"Streaming tool parser error (falling back to "
|
|
f"buffer for the rest of the stream): {e}",
|
|
file=sys.stderr)
|
|
native_streaming_error = True
|
|
msg = None
|
|
if msg is not None:
|
|
tc_protos = []
|
|
for tc in (msg.tool_calls or []):
|
|
fn = tc.function or None
|
|
tc_protos.append(backend_pb2.ToolCallDelta(
|
|
index=tc.index,
|
|
id=tc.id or "",
|
|
name=(fn.name if fn and fn.name else "") or "",
|
|
arguments=(fn.arguments if fn and fn.arguments else "") or "",
|
|
))
|
|
cd_kwargs = {}
|
|
if msg.content:
|
|
cd_kwargs["content"] = msg.content
|
|
if msg.reasoning:
|
|
cd_kwargs["reasoning_content"] = msg.reasoning
|
|
if tc_protos:
|
|
cd_kwargs["tool_calls"] = tc_protos
|
|
if cd_kwargs:
|
|
yield backend_pb2.Reply(
|
|
message=bytes(msg.content or "", encoding='utf-8'),
|
|
chat_deltas=[backend_pb2.ChatDelta(**cd_kwargs)],
|
|
)
|
|
# (B) buffer fallback — emit nothing during the stream.
|
|
# The post-loop extract_tool_calls block builds the final chunk.
|
|
|
|
# Keep track of text + token_ids generated
|
|
generated_text = iteration_text
|
|
generated_token_ids = list(request_output.outputs[0].token_ids)
|
|
finally:
|
|
await outputs.aclose()
|
|
|
|
# Remove the image files from /tmp folder
|
|
for img_path in image_paths:
|
|
try:
|
|
os.remove(img_path)
|
|
except Exception as e:
|
|
print(f"Error removing image file: {img_path}, {e}", file=sys.stderr)
|
|
|
|
# Parse reasoning and tool calls from final text using vLLM's native parsers
|
|
content = generated_text
|
|
reasoning_content = ""
|
|
tool_calls_proto = []
|
|
|
|
if self.reasoning_parser_cls:
|
|
try:
|
|
rp = self._new_reasoning_parser(template_kwargs)
|
|
r, c = rp.extract_reasoning(generated_text, request=None)
|
|
reasoning_content, content = self._split_reasoning(
|
|
rp, generated_text, prompt, r, c,
|
|
)
|
|
except Exception as e:
|
|
print(f"Reasoning parser error: {e}", file=sys.stderr)
|
|
|
|
# Content after the reasoning split, before tool-call extraction.
|
|
split_content = content
|
|
|
|
# When (A) native streaming ran cleanly, per-delta yields above already
|
|
# delivered everything — do NOT extract again on the full text or we'd
|
|
# duplicate content/tool_calls into the final chunk.
|
|
# NOTE: `native_streaming` is a capability flag ("streaming parser is
|
|
# available"), not a state flag ("streaming actually ran"). For
|
|
# non-streaming requests it is still True but the per-delta loop was
|
|
# never entered, so we MUST still run extract_tool_calls here. Hence
|
|
# the explicit `streaming and …` guard on both branches.
|
|
if has_tool_parser and not (streaming and native_streaming and not native_streaming_error):
|
|
try:
|
|
tp = tp_instance
|
|
if tp is None:
|
|
# Defensive: tp_instance build failed earlier; reconstruct.
|
|
tools = json.loads(request.Tools)
|
|
try:
|
|
tp = self.tool_parser_cls(self.tokenizer, tools=tools)
|
|
except TypeError:
|
|
tp = self.tool_parser_cls(self.tokenizer)
|
|
info = tp.extract_tool_calls(content, request=None)
|
|
if info.tools_called:
|
|
content = info.content or ""
|
|
for i, tc in enumerate(info.tool_calls):
|
|
tool_calls_proto.append(backend_pb2.ToolCallDelta(
|
|
index=i,
|
|
id=tc.id,
|
|
name=tc.function.name,
|
|
arguments=tc.function.arguments,
|
|
))
|
|
except Exception as e:
|
|
print(f"Tool parser error: {e}", file=sys.stderr)
|
|
elif streaming and native_streaming and not native_streaming_error:
|
|
# Per-delta path already emitted content + tool_calls; the final
|
|
# chat_delta should carry only metadata (token counts, logprobs).
|
|
content = ""
|
|
|
|
# Extract token counts
|
|
prompt_tokens = 0
|
|
completion_tokens = 0
|
|
if last_output is not None:
|
|
try:
|
|
prompt_tokens = len(last_output.prompt_token_ids or [])
|
|
except Exception:
|
|
pass
|
|
try:
|
|
completion_tokens = len(last_output.outputs[0].token_ids or [])
|
|
except Exception:
|
|
pass
|
|
|
|
# Extract logprobs if requested
|
|
logprobs_bytes = b""
|
|
if last_output is not None and request.Logprobs > 0:
|
|
try:
|
|
lp = last_output.outputs[0].logprobs
|
|
if lp:
|
|
logprobs_data = {"content": []}
|
|
for token_lp_dict in lp:
|
|
if token_lp_dict:
|
|
first_tok_id, first_lp = next(iter(token_lp_dict.items()))
|
|
logprobs_data["content"].append({
|
|
"token": getattr(first_lp, "decoded_token", str(first_tok_id)),
|
|
"logprob": first_lp.logprob,
|
|
})
|
|
logprobs_bytes = json.dumps(logprobs_data).encode("utf-8")
|
|
except Exception as e:
|
|
print(f"Logprobs extraction error: {e}", file=sys.stderr)
|
|
|
|
chat_delta = backend_pb2.ChatDelta(
|
|
content=content,
|
|
reasoning_content=reasoning_content,
|
|
tool_calls=tool_calls_proto,
|
|
)
|
|
|
|
if streaming and stream_parser is not None:
|
|
# Everything went out per delta; the final chunk carries only token
|
|
# counts and logprobs, or LocalAI would repeat answer and
|
|
# reasoning. If the parser failed mid-stream, send what is left.
|
|
remainder = {}
|
|
if stream_parser_error:
|
|
rest_content = split_content
|
|
rest_tool_calls = []
|
|
if has_tool_parser and rest_content:
|
|
try:
|
|
tp = tp_instance or self.tool_parser_cls(self.tokenizer)
|
|
info = tp.extract_tool_calls(rest_content, request=None)
|
|
if info.tools_called:
|
|
rest_content = info.content or ""
|
|
rest_tool_calls = [
|
|
(tc.id, tc.function.name, tc.function.arguments)
|
|
for tc in info.tool_calls
|
|
]
|
|
except Exception as e:
|
|
print(f"Tool parser error: {e}", file=sys.stderr)
|
|
remainder = self._unsent_remainder(
|
|
sent_content, sent_reasoning, sent_tools,
|
|
rest_content, reasoning_content, rest_tool_calls,
|
|
)
|
|
yield backend_pb2.Reply(
|
|
message=bytes(remainder.get("content", ""), encoding='utf-8'),
|
|
prompt_tokens=prompt_tokens,
|
|
tokens=completion_tokens,
|
|
chat_deltas=[backend_pb2.ChatDelta(**remainder)],
|
|
logprobs=logprobs_bytes,
|
|
)
|
|
return
|
|
|
|
if streaming:
|
|
# Final chunk with structured data.
|
|
#
|
|
# If we used the buffer fallback (has_tool_parser=True AND native
|
|
# streaming did NOT run cleanly) and the parser found no tool call,
|
|
# flush the buffered content as ONE content delta — and clear the
|
|
# final chat_delta's content so the metadata chunk does not repeat
|
|
# what we just sent. This is the plain-text-with-tool-parser path.
|
|
buffered_fallback = (
|
|
has_tool_parser
|
|
and not (native_streaming and not native_streaming_error)
|
|
)
|
|
if buffered_fallback and not tool_calls_proto and content:
|
|
yield backend_pb2.Reply(
|
|
message=bytes(content, encoding='utf-8'),
|
|
chat_deltas=[backend_pb2.ChatDelta(content=content)],
|
|
)
|
|
chat_delta = backend_pb2.ChatDelta(
|
|
reasoning_content=reasoning_content,
|
|
tool_calls=tool_calls_proto,
|
|
)
|
|
yield backend_pb2.Reply(
|
|
message=b"",
|
|
prompt_tokens=prompt_tokens,
|
|
tokens=completion_tokens,
|
|
chat_deltas=[chat_delta],
|
|
logprobs=logprobs_bytes,
|
|
)
|
|
return
|
|
|
|
# Non-streaming: single Reply with everything
|
|
yield backend_pb2.Reply(
|
|
message=bytes(content, encoding='utf-8'),
|
|
prompt_tokens=prompt_tokens,
|
|
tokens=completion_tokens,
|
|
chat_deltas=[chat_delta],
|
|
logprobs=logprobs_bytes,
|
|
)
|
|
|
|
def load_image(self, image_path: str):
|
|
"""
|
|
Load an image from the given file path or base64 encoded data.
|
|
|
|
Args:
|
|
image_path (str): The path to the image file or base64 encoded data.
|
|
|
|
Returns:
|
|
Image: The loaded image.
|
|
"""
|
|
try:
|
|
|
|
image_data = base64.b64decode(image_path)
|
|
image = Image.open(io.BytesIO(image_data))
|
|
return image
|
|
except Exception as e:
|
|
print(f"Error loading image {image_path}: {e}", file=sys.stderr)
|
|
return None
|
|
|
|
def load_video(self, video_path: str):
|
|
"""
|
|
Load a video from the given file path.
|
|
|
|
Args:
|
|
video_path (str): The path to the image file.
|
|
|
|
Returns:
|
|
Video: The loaded video.
|
|
"""
|
|
try:
|
|
with materialize_base64(video_path, suffix=".data") as path:
|
|
return VideoAsset(name=path).np_ndarrays
|
|
except Exception as e:
|
|
print(f"Error loading video {video_path}: {e}", file=sys.stderr)
|
|
return None
|
|
|
|
async def serve(address):
|
|
# Start asyncio gRPC server
|
|
server = grpc.aio.server(migration_thread_pool=futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
|
options=[
|
|
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
|
|
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
|
|
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
|
|
],
|
|
interceptors=get_auth_interceptors(aio=True),
|
|
)
|
|
# Add the servicer to the server
|
|
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
|
# Bind the server to the address
|
|
server.add_insecure_port(address)
|
|
|
|
# Gracefully shutdown the server on SIGTERM or SIGINT
|
|
loop = asyncio.get_event_loop()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
loop.add_signal_handler(
|
|
sig, lambda: asyncio.ensure_future(server.stop(5))
|
|
)
|
|
|
|
# Start the server
|
|
await server.start()
|
|
print("Server started. Listening on: " + address, file=sys.stderr)
|
|
# Wait for the server to be terminated
|
|
await server.wait_for_termination()
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run the gRPC server.")
|
|
parser.add_argument(
|
|
"--addr", default="localhost:50051", help="The address to bind the server to."
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
asyncio.run(serve(args.addr))
|