846 lines
30 KiB
Go
846 lines
30 KiB
Go
// Package vllmkit talks to the production vLLM OpenAI-compatible
|
|
// inference servers that host Qwen 3.5 (chat + tool calling) and a
|
|
// separate embedding model. Shared by the agent service (ChatStream)
|
|
// and the knowledge service (Embed) — the same internal/llm.Client the
|
|
// monolith had, split out since both of those bounded contexts need it
|
|
// and neither should depend on the other.
|
|
//
|
|
// This package previously spoke Ollama's native /api/chat and
|
|
// /api/embed protocol directly against a single local daemon. It now
|
|
// speaks the OpenAI /v1/chat/completions and /v1/embeddings wire
|
|
// format that `vllm serve` exposes, against two independently scaled,
|
|
// independently deployed vLLM servers — one per model, since a single
|
|
// vLLM engine process serves exactly one model. See
|
|
// ENKLY_VLLM_CHAT_URL / ENKLY_VLLM_EMBED_URL in cmd/server/main.go and
|
|
// k8s/11-vllm.yaml for how those servers are configured and deployed.
|
|
package vllmkit
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
// DefaultChatModel is the served model name the chat vLLM
|
|
// deployment must be started with (`vllm serve ... --served-model-name
|
|
// qwen3.5`). vLLM matches the request body's "model" field against
|
|
// that flag's value, not against a Hugging Face repo id, so this
|
|
// is a deployment-chosen alias, not a model identifier — keep it
|
|
// in sync with k8s/11-vllm.yaml's --served-model-name.
|
|
DefaultChatModel = "qwen3.5"
|
|
// DefaultEmbedModel is the served model name for the embedding
|
|
// vLLM deployment. vLLM's pooling/embedding model support covers
|
|
// nomic-ai/nomic-embed-text-v1.5 (what this deployment actually
|
|
// runs; it needs `--trust-remote-code` since Nomic's config uses
|
|
// custom modeling code) — small, fast, and a good CPU/GPU-both
|
|
// fit. Update the `vector(768)` dimension in schema.sql first if
|
|
// you point this at a model with a different output size.
|
|
DefaultEmbedModel = "nomic-embed-text"
|
|
)
|
|
|
|
// Options mirrors the subset of the OpenAI-compatible chat-completions
|
|
// request body that's worth tuning per call. Unlike Ollama's
|
|
// num_ctx/num_gpu/num_thread, context length and GPU/tensor-parallel
|
|
// layout are not per-request knobs under vLLM — they're fixed at
|
|
// server start (--max-model-len, --tensor-parallel-size,
|
|
// --gpu-memory-utilization; see k8s/11-vllm.yaml) and apply to every
|
|
// request the deployment serves. Zero values are omitted from the wire
|
|
// request, so the server's own defaults apply unless explicitly
|
|
// overridden — see the ENKLY_CHAT_* env vars wired up in
|
|
// cmd/server/main.go.
|
|
type Options struct {
|
|
Temperature float64 // sampling temperature (temperature)
|
|
TopP float64 // nucleus sampling cutoff (top_p)
|
|
MaxTokens int // max tokens to generate this turn (max_tokens); 0 = server default
|
|
}
|
|
|
|
// Client is a thin wrapper around two vLLM OpenAI-compatible HTTP
|
|
// APIs: one serving the chat model, one serving the embedding model.
|
|
type Client struct {
|
|
// ChatBaseURL / EmbedBaseURL point at each vLLM deployment's
|
|
// OpenAI-compatible router, e.g. "https://vllm-chat.enkly.org/v1" and
|
|
// "https://vllm-embed.enkly.org/v1". They are deliberately separate
|
|
// fields (rather than one shared host) because a chat model and
|
|
// an embedding model are different vLLM engine processes in
|
|
// production — usually on different Services/Deployments with
|
|
// different GPU shapes — not two endpoints behind the same daemon
|
|
// the way Ollama modeled it.
|
|
ChatBaseURL string
|
|
EmbedBaseURL string
|
|
ChatModel string
|
|
EmbedModel string
|
|
HTTP *http.Client
|
|
|
|
// APIKey, when set, is sent as "Authorization: Bearer <APIKey>" on
|
|
// every request — vLLM's --api-key flag. Leave empty for a
|
|
// cluster-internal deployment with no auth in front of it (the
|
|
// Service is only reachable from inside the mesh/namespace); set
|
|
// it for anything exposed beyond that boundary. Set via
|
|
// ENKLY_VLLM_API_KEY.
|
|
APIKey string
|
|
|
|
// Think controls extended "thinking" on chat requests via vLLM's
|
|
// chat_template_kwargs extra-body extension: "" leaves it to the
|
|
// chat template's own default (Qwen3's thinking-capable template
|
|
// defaults to thinking on), "true"/"false" (and the usual
|
|
// on/off/yes/no spellings) force it on/off. A "low"/"medium"/
|
|
// "high" effort string (see pkg/effort) is also accepted for
|
|
// compatibility with the Anthropic-style effort vocabulary the
|
|
// rest of the agent uses — vLLM/Qwen3's reasoning toggle itself is
|
|
// a single on/off switch with no server-side notion of graded
|
|
// effort, so any non-empty, non-"false" level here just turns
|
|
// thinking on; depth is instead shaped by Options.MaxTokens and
|
|
// agentrun's own round/compaction budgets (see effort.Profile).
|
|
// When on, the reasoning trace streams back separately from the
|
|
// answer, in OpenAI's "reasoning_content" delta field
|
|
// (ChatEvent.Thinking) — requires the chat deployment to be
|
|
// started with a reasoning parser, e.g. `--reasoning-parser
|
|
// qwen3`. Set via ENKLY_CHAT_THINK.
|
|
Think string
|
|
// ChatOptions is sent as the top-level temperature/top_p/max_tokens
|
|
// fields of the chat-completions request.
|
|
ChatOptions Options
|
|
|
|
// embedCache memoizes Embed() results within this process's
|
|
// lifetime, keyed by (model, text) hash. Embedding the exact same
|
|
// string twice is common — repeated search_codebase queries within
|
|
// a conversation, or an unchanged chunk re-indexed after a
|
|
// no-op edit that rag.Store's content-hash check didn't already
|
|
// short-circuit — and skipping the round trip to vLLM for a
|
|
// repeat is a pure latency win with no correctness downside, since
|
|
// embeddings are a deterministic function of (model, text).
|
|
embedCache *embedCache
|
|
}
|
|
|
|
// New builds a Client pointed at the given chat and embed vLLM
|
|
// deployments, with connection reuse tuned for the request pattern
|
|
// here: a handful of long-lived streaming chat requests plus frequent,
|
|
// short embedding calls, each against its own vLLM host. Either URL
|
|
// may be left empty if this Client will only ever be used for the
|
|
// other call (e.g. the knowledge service never calls ChatStream, so it
|
|
// passes "" for chatBaseURL).
|
|
func New(chatBaseURL, embedBaseURL, chatModel, embedModel string) *Client {
|
|
if chatBaseURL == "" {
|
|
chatBaseURL = "http://localhost:8000/v1"
|
|
}
|
|
if embedBaseURL == "" {
|
|
embedBaseURL = "http://localhost:8001/v1"
|
|
}
|
|
if chatModel == "" {
|
|
chatModel = DefaultChatModel
|
|
}
|
|
if embedModel == "" {
|
|
embedModel = DefaultEmbedModel
|
|
}
|
|
transport := &http.Transport{
|
|
Proxy: http.ProxyFromEnvironment,
|
|
DialContext: (&net.Dialer{
|
|
Timeout: 10 * time.Second,
|
|
KeepAlive: 60 * time.Second,
|
|
}).DialContext,
|
|
// Two vLLM hosts, but chat + embed + concurrent requests (see
|
|
// agent.go's parallel tool execution, which can itself fan out
|
|
// concurrent search_codebase embed calls) can be in flight
|
|
// together — a small pool avoids a fresh TCP+HTTP handshake on
|
|
// every call, a real chunk of latency in practice even against
|
|
// a fast in-cluster backend.
|
|
MaxIdleConns: 64,
|
|
MaxIdleConnsPerHost: 32,
|
|
IdleConnTimeout: 90 * time.Second,
|
|
}
|
|
return &Client{
|
|
ChatBaseURL: chatBaseURL,
|
|
EmbedBaseURL: embedBaseURL,
|
|
ChatModel: chatModel,
|
|
EmbedModel: embedModel,
|
|
HTTP: &http.Client{
|
|
Transport: transport,
|
|
// No overall timeout: chat streams can legitimately run for
|
|
// minutes while the agent loop executes several tool calls.
|
|
// Callers should thread a context deadline through ctx instead.
|
|
},
|
|
embedCache: newEmbedCache(500),
|
|
}
|
|
}
|
|
|
|
func (c *Client) authHeader(req *http.Request) {
|
|
if c.APIKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------
|
|
// Chat (streaming, with tool calling) — OpenAI /v1/chat/completions
|
|
// ---------------------------------------------------------------------
|
|
|
|
// wireToolCall is the OpenAI wire shape for a tool call, both on a
|
|
// complete assistant message (echoed back into history) and,
|
|
// incrementally, on a streamed delta (see toolCallAccumulator).
|
|
type wireToolCall struct {
|
|
Index *int `json:"index,omitempty"` // present only on streamed deltas
|
|
ID string `json:"id,omitempty"`
|
|
Type string `json:"type,omitempty"` // always "function"
|
|
Function wireToolCallFunction `json:"function"`
|
|
}
|
|
|
|
type wireToolCallFunction struct {
|
|
Name string `json:"name,omitempty"`
|
|
// Arguments is a JSON-encoded string on the wire (OpenAI's
|
|
// function-calling contract), not a nested object the way Ollama's
|
|
// native API sent it — toWireMessages/fromWireToolCalls marshal
|
|
// and unmarshal it at this package's boundary so the rest of the
|
|
// codebase still deals in map[string]any.
|
|
Arguments string `json:"arguments,omitempty"`
|
|
}
|
|
|
|
type wireMessage struct {
|
|
Role string `json:"role"`
|
|
Content *string `json:"content"` // nil (omitted) for an assistant message that is tool_calls-only
|
|
ToolCalls []wireToolCall `json:"tool_calls,omitempty"`
|
|
ToolCallID string `json:"tool_call_id,omitempty"` // role="tool" only
|
|
Name string `json:"name,omitempty"` // role="tool" only; some parsers want the tool name too
|
|
}
|
|
|
|
type wireToolDef struct {
|
|
Type string `json:"type"` // always "function"
|
|
Function wireToolFunction `json:"function"`
|
|
}
|
|
|
|
type wireToolFunction struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
Parameters map[string]any `json:"parameters"`
|
|
}
|
|
|
|
type chatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []wireMessage `json:"messages"`
|
|
Tools []wireToolDef `json:"tools,omitempty"`
|
|
Stream bool `json:"stream"`
|
|
// ChatTemplateKwargs is vLLM's extra-body extension point for
|
|
// per-request chat-template flags; "enable_thinking" is how
|
|
// Qwen3's template (and vLLM's qwen3 reasoning parser) take the
|
|
// thinking on/off toggle. Opaque to the OpenAI spec itself, so
|
|
// this is marshaled as a bare map rather than a typed field.
|
|
ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"`
|
|
|
|
Temperature float64 `json:"temperature,omitempty"`
|
|
TopP float64 `json:"top_p,omitempty"`
|
|
MaxTokens int `json:"max_tokens,omitempty"`
|
|
}
|
|
|
|
// chatStreamChunk is one Server-Sent-Events "data:" line of an OpenAI
|
|
// streaming chat-completions response.
|
|
type chatStreamChunk struct {
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
ReasoningContent string `json:"reasoning_content"`
|
|
ToolCalls []wireToolCall `json:"tool_calls"`
|
|
} `json:"delta"`
|
|
FinishReason *string `json:"finish_reason"`
|
|
} `json:"choices"`
|
|
// Error is set when a vLLM server streams a mid-response failure
|
|
// as a normal 200 SSE line rather than an HTTP error status (seen
|
|
// e.g. when the tool-call parser can't make sense of the model's
|
|
// raw output partway through generation).
|
|
Error *wireError `json:"error"`
|
|
}
|
|
|
|
// ErrUpstreamUnavailable wraps a transport-level failure to reach a
|
|
// vLLM deployment at all — connection refused, DNS failure, dial/TLS
|
|
// timeout — as opposed to vLLM responding with an error status or a
|
|
// malformed body. This is exactly the signature of a vllm-chat/
|
|
// vllm-embed pod that's Pending, still cold-starting (model load, see
|
|
// k8s/11-vllm.yaml's startupProbe), or mid-rollout (Recreate strategy
|
|
// briefly has zero replicas) — not a request vLLM itself rejected.
|
|
//
|
|
// Callers use IsUpstreamUnavailable to tell the two apart: the agent
|
|
// service turns this into a short "model is starting up" message for
|
|
// the end user instead of a raw dial error (see
|
|
// services/agent/internal/application/service.go), and the knowledge
|
|
// service's gRPC layer maps it onto codes.Unavailable (safe to retry)
|
|
// instead of codes.Internal (opaque failure).
|
|
var ErrUpstreamUnavailable = errors.New("vllm: upstream unreachable")
|
|
|
|
// IsUpstreamUnavailable reports whether err is, or wraps,
|
|
// ErrUpstreamUnavailable.
|
|
func IsUpstreamUnavailable(err error) bool { return errors.Is(err, ErrUpstreamUnavailable) }
|
|
|
|
// looksTransient reports whether err is a transport-level failure
|
|
// (connection refused, DNS, dial/handshake timeout, connection reset)
|
|
// worth a couple of quick retries, as opposed to an error vLLM itself
|
|
// returned after accepting the connection.
|
|
func looksTransient(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
var netErr net.Error
|
|
if errors.As(err, &netErr) {
|
|
return true // covers dial timeout/refused, DNS failure, TLS handshake failure, etc.
|
|
}
|
|
msg := strings.ToLower(err.Error())
|
|
return strings.Contains(msg, "connection refused") ||
|
|
strings.Contains(msg, "no such host") ||
|
|
strings.Contains(msg, "connection reset")
|
|
}
|
|
|
|
// retryDelays bounds how long doWithRetry rides out a transient
|
|
// failure before giving up: short and few on purpose. This exists to
|
|
// smooth over a brief reconnect window — e.g. a Recreate rollout's old
|
|
// pod having just terminated — not to mask a genuinely Pending pod,
|
|
// which is a scheduling problem no amount of HTTP retrying fixes (see
|
|
// k8s/scripts/gpu-preflight.sh for that diagnosis).
|
|
var retryDelays = []time.Duration{200 * time.Millisecond, 600 * time.Millisecond}
|
|
|
|
// doWithRetry POSTs buf as JSON to url, retrying only on a transient
|
|
// transport failure (see looksTransient) — never after a response has
|
|
// actually been received, since a non-2xx status or a stream already
|
|
// in progress isn't safe or meaningful to blindly resend.
|
|
func (c *Client) doWithRetry(ctx context.Context, url string, buf []byte, sse bool) (*http.Response, error) {
|
|
for attempt := 0; ; attempt++ {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(buf))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if sse {
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
}
|
|
c.authHeader(req)
|
|
|
|
resp, err := c.HTTP.Do(req)
|
|
if err == nil {
|
|
return resp, nil
|
|
}
|
|
if attempt >= len(retryDelays) || !looksTransient(err) {
|
|
return nil, err
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-time.After(retryDelays[attempt]):
|
|
}
|
|
}
|
|
}
|
|
|
|
type wireError struct {
|
|
Message string `json:"message"`
|
|
Type string `json:"type"`
|
|
Code string `json:"code"`
|
|
}
|
|
|
|
func (e *wireError) Error() string {
|
|
if e == nil {
|
|
return ""
|
|
}
|
|
if e.Type != "" {
|
|
return fmt.Sprintf("%s: %s", e.Type, e.Message)
|
|
}
|
|
return e.Message
|
|
}
|
|
|
|
// errMalformedToolArgs is wrapped into the error ChatStream reports
|
|
// when vLLM's tool-call parser hands back an "arguments" string that
|
|
// doesn't parse as JSON. This is this package's own client-side
|
|
// detection of a broken tool-call round (vLLM never surfaces this as a
|
|
// distinct HTTP/wire error itself) — see IsToolCallParseError.
|
|
var errMalformedToolArgs = errors.New("vllm: tool call arguments were not valid JSON")
|
|
|
|
// toolCallParseErrPattern matches the family of messages vLLM's
|
|
// OpenAI-compatible server reports when a configured tool-call parser
|
|
// (hermes, llama3_json, qwen3, mistral, …) fails to extract a
|
|
// structured call from the model's raw generation — e.g. "Failed to
|
|
// parse tool calls", "could not parse function call", "invalid tool
|
|
// call arguments produced by model". Server-side wording differs by
|
|
// parser plugin and vLLM version, so this matches on the stable,
|
|
// shared vocabulary ("tool call" / "function call" alongside "parse"
|
|
// or "invalid") rather than one exact phrasing.
|
|
func toolCallParseErrPattern(msg string) bool {
|
|
msg = strings.ToLower(msg)
|
|
mentionsCall := strings.Contains(msg, "tool call") || strings.Contains(msg, "tool_call") ||
|
|
strings.Contains(msg, "function call") || strings.Contains(msg, "function_call")
|
|
mentionsFailure := strings.Contains(msg, "parse") || strings.Contains(msg, "invalid") || strings.Contains(msg, "malformed")
|
|
return mentionsCall && mentionsFailure
|
|
}
|
|
|
|
// IsToolCallParseError reports whether err is a transient failure to
|
|
// extract a structured tool call from the model's output — either
|
|
// vLLM's own tool-call parser rejecting the raw generation server-side
|
|
// (see toolCallParseErrPattern), or this client failing to decode the
|
|
// "arguments" JSON string vLLM did return (errMalformedToolArgs). Both
|
|
// are almost always a one-off decoding hiccup — something in the
|
|
// model's raw output didn't match what the configured parser expected
|
|
// — rather than a reason to give up the whole agent turn. Callers
|
|
// (agentrun.Agent.Run) use this to retry the round, and if it keeps
|
|
// recurring, to simplify the most recent tool results before one final
|
|
// retry.
|
|
//
|
|
// This replaces IsTemplateParseError, which detected the equivalent
|
|
// failure mode under Ollama's native chat-template/XML tool-call
|
|
// encoding — a protocol vLLM's OpenAI-compatible server doesn't use,
|
|
// so that detection logic (XML syntax errors, angle-bracket tag
|
|
// sniffing) no longer applies, but the role it plays in the retry loop
|
|
// is unchanged.
|
|
func IsToolCallParseError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
if errors.Is(err, errMalformedToolArgs) {
|
|
return true
|
|
}
|
|
return toolCallParseErrPattern(err.Error())
|
|
}
|
|
|
|
// thinkWire maps a Think setting onto vLLM's chat_template_kwargs
|
|
// extra body. It returns nil when thinking should be left to the
|
|
// template's own default (omitting chat_template_kwargs entirely), and
|
|
// otherwise a map with "enable_thinking" set to true/false.
|
|
//
|
|
// override, when non-empty, takes precedence over c.Think for this one
|
|
// call — see ChatStream's think parameter. This is what lets a caller
|
|
// (agentrun.Agent's ThinkPolicy) dial reasoning on or off on a
|
|
// per-round basis — e.g. thinking on for the first, planning round of
|
|
// an agentic tool-use turn, thinking off for the rounds that just
|
|
// digest a tool result.
|
|
func thinkWire(setting, override string) map[string]any {
|
|
s := override
|
|
if s == "" {
|
|
s = setting
|
|
}
|
|
switch strings.ToLower(strings.TrimSpace(s)) {
|
|
case "":
|
|
return nil
|
|
case "false", "off", "0", "no":
|
|
return map[string]any{"enable_thinking": false}
|
|
default:
|
|
// "true"/"on"/"1"/"yes" and the effort levels ("low", "medium",
|
|
// "high", …) all just mean "thinking on" — see Think's doc
|
|
// comment for why vLLM/Qwen3 has no graded-effort equivalent.
|
|
return map[string]any{"enable_thinking": true}
|
|
}
|
|
}
|
|
|
|
// ChatEvent is one item streamed back from ChatStream: either a token of
|
|
// assistant content, a batch of tool calls the model wants executed, or
|
|
// the terminal done/error signal.
|
|
type ChatEvent struct {
|
|
Token string
|
|
Thinking string // a chunk of the model's reasoning trace (when thinking is on)
|
|
ToolCalls []ToolCallRequest
|
|
Done bool
|
|
Err error
|
|
}
|
|
|
|
// ChatStream sends the full message history (plus tool definitions) to
|
|
// the chat vLLM deployment and streams the reply back over the
|
|
// returned channel. The channel is closed after a Done or Err event.
|
|
//
|
|
// think overrides c.Think for this one call when non-empty (see
|
|
// thinkWire); pass "" to just use the Client's configured default.
|
|
func (c *Client) ChatStream(ctx context.Context, history []Message, tools []ToolDef, think string) (<-chan ChatEvent, error) {
|
|
reqBody := chatRequest{
|
|
Model: c.ChatModel,
|
|
Messages: toWireMessages(history),
|
|
Tools: toWireTools(tools),
|
|
Stream: true,
|
|
ChatTemplateKwargs: thinkWire(c.Think, think),
|
|
}
|
|
reqBody.Temperature = c.ChatOptions.Temperature
|
|
reqBody.TopP = c.ChatOptions.TopP
|
|
reqBody.MaxTokens = c.ChatOptions.MaxTokens
|
|
|
|
buf, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("llm: marshal chat request: %w", err)
|
|
}
|
|
|
|
resp, err := c.doWithRetry(ctx, c.ChatBaseURL+"/chat/completions", buf, true)
|
|
if err != nil {
|
|
if looksTransient(err) {
|
|
return nil, fmt.Errorf("llm: chat request failed (is vLLM running at %s?): %w: %w", c.ChatBaseURL, err, ErrUpstreamUnavailable)
|
|
}
|
|
return nil, fmt.Errorf("llm: chat request failed (is vLLM running at %s?): %w", c.ChatBaseURL, err)
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
defer resp.Body.Close()
|
|
body := readErrorBody(resp)
|
|
if resp.StatusCode == http.StatusServiceUnavailable {
|
|
return nil, fmt.Errorf("llm: vllm returned %s: %s: %w", resp.Status, body, ErrUpstreamUnavailable)
|
|
}
|
|
return nil, fmt.Errorf("llm: vllm returned %s: %s", resp.Status, body)
|
|
}
|
|
|
|
events := make(chan ChatEvent, 16)
|
|
|
|
go func() {
|
|
defer resp.Body.Close()
|
|
defer close(events)
|
|
|
|
// send delivers an event unless the consumer has gone away (ctx
|
|
// cancelled), so this goroutine can never block forever on a
|
|
// channel nobody is reading.
|
|
send := func(ev ChatEvent) bool {
|
|
select {
|
|
case events <- ev:
|
|
return true
|
|
case <-ctx.Done():
|
|
return false
|
|
}
|
|
}
|
|
|
|
acc := newToolCallAccumulator()
|
|
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
|
|
|
sawDone := false
|
|
for scanner.Scan() {
|
|
line := bytes.TrimSpace(scanner.Bytes())
|
|
if len(line) == 0 {
|
|
continue
|
|
}
|
|
data, ok := bytes.CutPrefix(line, []byte("data:"))
|
|
if !ok {
|
|
// Non-"data:" SSE lines (comments, "event:", blank
|
|
// keep-alives) carry nothing this client needs.
|
|
continue
|
|
}
|
|
data = bytes.TrimSpace(data)
|
|
if string(data) == "[DONE]" {
|
|
sawDone = true
|
|
break
|
|
}
|
|
|
|
var chunk chatStreamChunk
|
|
if err := json.Unmarshal(data, &chunk); err != nil {
|
|
send(ChatEvent{Err: fmt.Errorf("llm: decode stream chunk: %w", err)})
|
|
return
|
|
}
|
|
if chunk.Error != nil {
|
|
send(ChatEvent{Err: fmt.Errorf("llm: vllm error: %w", chunk.Error)})
|
|
return
|
|
}
|
|
if len(chunk.Choices) == 0 {
|
|
continue
|
|
}
|
|
choice := chunk.Choices[0]
|
|
|
|
if choice.Delta.ReasoningContent != "" {
|
|
if !send(ChatEvent{Thinking: choice.Delta.ReasoningContent}) {
|
|
return
|
|
}
|
|
}
|
|
if choice.Delta.Content != "" {
|
|
if !send(ChatEvent{Token: choice.Delta.Content}) {
|
|
return
|
|
}
|
|
}
|
|
if len(choice.Delta.ToolCalls) > 0 {
|
|
acc.add(choice.Delta.ToolCalls)
|
|
}
|
|
|
|
if choice.FinishReason != nil {
|
|
if acc.len() > 0 {
|
|
calls, err := acc.finish()
|
|
if err != nil {
|
|
send(ChatEvent{Err: fmt.Errorf("llm: %w: %w", errMalformedToolArgs, err)})
|
|
return
|
|
}
|
|
if !send(ChatEvent{ToolCalls: calls}) {
|
|
return
|
|
}
|
|
}
|
|
sawDone = true
|
|
send(ChatEvent{Done: true})
|
|
return
|
|
}
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
send(ChatEvent{Err: fmt.Errorf("llm: reading stream: %w", err)})
|
|
return
|
|
}
|
|
// A [DONE] sentinel with no finish_reason chunk (or EOF with
|
|
// neither) means vLLM dropped the stream before completing the
|
|
// reply (crash, OOM-kill, network cut) — surface it rather
|
|
// than treating a truncated reply as a complete one.
|
|
if !sawDone && ctx.Err() == nil {
|
|
send(ChatEvent{Err: fmt.Errorf("llm: vllm closed the stream before finishing the reply")})
|
|
return
|
|
}
|
|
if sawDone {
|
|
send(ChatEvent{Done: true})
|
|
}
|
|
}()
|
|
|
|
return events, nil
|
|
}
|
|
|
|
func readErrorBody(resp *http.Response) string {
|
|
body := make([]byte, 2048)
|
|
n, _ := resp.Body.Read(body)
|
|
raw := body[:n]
|
|
|
|
var envelope struct {
|
|
Error wireError `json:"error"`
|
|
}
|
|
if err := json.Unmarshal(raw, &envelope); err == nil && envelope.Error.Message != "" {
|
|
return envelope.Error.Error()
|
|
}
|
|
return string(raw)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------
|
|
// toolCallAccumulator: reassembles a streamed, index-keyed sequence of
|
|
// partial tool_calls deltas into complete calls. The OpenAI streaming
|
|
// contract sends a tool call's id/name once (on the chunk that first
|
|
// introduces that index) and its "arguments" JSON string in arbitrary-
|
|
// length fragments across however many subsequent chunks it takes to
|
|
// finish generating them — chunk boundaries do not align with JSON
|
|
// token boundaries at all, so nothing is parseable until finish().
|
|
// ---------------------------------------------------------------------
|
|
|
|
type toolCallAccumulator struct {
|
|
order []int
|
|
byIdx map[int]*accumulatingCall
|
|
}
|
|
|
|
type accumulatingCall struct {
|
|
id, name string
|
|
args strings.Builder
|
|
}
|
|
|
|
func newToolCallAccumulator() *toolCallAccumulator {
|
|
return &toolCallAccumulator{byIdx: map[int]*accumulatingCall{}}
|
|
}
|
|
|
|
func (a *toolCallAccumulator) len() int { return len(a.byIdx) }
|
|
|
|
func (a *toolCallAccumulator) add(deltas []wireToolCall) {
|
|
for _, d := range deltas {
|
|
idx := 0
|
|
if d.Index != nil {
|
|
idx = *d.Index
|
|
}
|
|
call, ok := a.byIdx[idx]
|
|
if !ok {
|
|
call = &accumulatingCall{}
|
|
a.byIdx[idx] = call
|
|
a.order = append(a.order, idx)
|
|
}
|
|
if d.ID != "" {
|
|
call.id = d.ID
|
|
}
|
|
if d.Function.Name != "" {
|
|
call.name = d.Function.Name
|
|
}
|
|
call.args.WriteString(d.Function.Arguments)
|
|
}
|
|
}
|
|
|
|
// finish parses every accumulated call's argument string as JSON and
|
|
// returns them in the order their indices first appeared. An empty
|
|
// arguments string (a tool with no parameters) is treated as "{}",
|
|
// matching what every vLLM tool-call parser actually emits for a
|
|
// no-arg call.
|
|
func (a *toolCallAccumulator) finish() ([]ToolCallRequest, error) {
|
|
sort.Ints(a.order)
|
|
out := make([]ToolCallRequest, 0, len(a.order))
|
|
for _, idx := range a.order {
|
|
call := a.byIdx[idx]
|
|
raw := strings.TrimSpace(call.args.String())
|
|
if raw == "" {
|
|
raw = "{}"
|
|
}
|
|
var args map[string]any
|
|
if err := json.Unmarshal([]byte(raw), &args); err != nil {
|
|
return nil, fmt.Errorf("tool %q (id %s): %w", call.name, call.id, err)
|
|
}
|
|
out = append(out, ToolCallRequest{ID: call.id, Name: call.name, Arguments: args})
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func toWireMessages(history []Message) []wireMessage {
|
|
out := make([]wireMessage, 0, len(history))
|
|
for _, m := range history {
|
|
content := m.Content
|
|
wm := wireMessage{
|
|
Role: string(m.Role),
|
|
Content: &content,
|
|
ToolCallID: m.ToolCallID,
|
|
Name: m.ToolName,
|
|
}
|
|
// An assistant message that only carries tool calls (no reply
|
|
// text) sends content as null, not "" — some vLLM tool-call
|
|
// parsers treat an empty-string content as "the model also
|
|
// said something" and choke on it.
|
|
if m.Role == RoleAssistant && len(m.ToolCalls) > 0 && m.Content == "" {
|
|
wm.Content = nil
|
|
}
|
|
for _, tc := range m.ToolCalls {
|
|
args := tc.Arguments
|
|
if args == nil {
|
|
args = map[string]any{}
|
|
}
|
|
argsJSON, err := json.Marshal(args)
|
|
if err != nil {
|
|
// Arguments is always built from values this process
|
|
// itself constructed (see toolcatalog/executor), so a
|
|
// marshal failure here would be a programming error,
|
|
// not a runtime condition — fall back to "{}" rather
|
|
// than drop the whole request.
|
|
argsJSON = []byte("{}")
|
|
}
|
|
wm.ToolCalls = append(wm.ToolCalls, wireToolCall{
|
|
ID: tc.ID,
|
|
Type: "function",
|
|
Function: wireToolCallFunction{Name: tc.Name, Arguments: string(argsJSON)},
|
|
})
|
|
}
|
|
out = append(out, wm)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func toWireTools(tools []ToolDef) []wireToolDef {
|
|
out := make([]wireToolDef, 0, len(tools))
|
|
for _, t := range tools {
|
|
out = append(out, wireToolDef{
|
|
Type: "function",
|
|
Function: wireToolFunction{
|
|
Name: t.Name,
|
|
Description: t.Description,
|
|
Parameters: t.Parameters,
|
|
},
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
// ---------------------------------------------------------------------
|
|
// Embeddings — OpenAI /v1/embeddings
|
|
// ---------------------------------------------------------------------
|
|
|
|
type embedRequest struct {
|
|
Model string `json:"model"`
|
|
Input string `json:"input"`
|
|
}
|
|
|
|
type embedResponse struct {
|
|
Data []struct {
|
|
Embedding []float32 `json:"embedding"`
|
|
Index int `json:"index"`
|
|
} `json:"data"`
|
|
}
|
|
|
|
// Embed returns the embedding vector for a single string, using the
|
|
// embed vLLM deployment's /v1/embeddings endpoint. Identical (model,
|
|
// text) pairs are served from an in-memory cache instead of
|
|
// round-tripping to vLLM again — see the embedCache doc comment on
|
|
// Client.
|
|
func (c *Client) Embed(ctx context.Context, text string) ([]float32, error) {
|
|
if c.embedCache != nil {
|
|
if v, ok := c.embedCache.get(c.EmbedModel, text); ok {
|
|
return v, nil
|
|
}
|
|
}
|
|
|
|
reqBody := embedRequest{Model: c.EmbedModel, Input: text}
|
|
buf, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("llm: marshal embed request: %w", err)
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
|
defer cancel()
|
|
|
|
resp, err := c.doWithRetry(ctx, c.EmbedBaseURL+"/embeddings", buf, false)
|
|
if err != nil {
|
|
if looksTransient(err) {
|
|
return nil, fmt.Errorf("llm: embed request failed (is vLLM running at %s?): %w: %w", c.EmbedBaseURL, err, ErrUpstreamUnavailable)
|
|
}
|
|
return nil, fmt.Errorf("llm: embed request failed (is vLLM running at %s?): %w", c.EmbedBaseURL, err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body := readErrorBody(resp)
|
|
if resp.StatusCode == http.StatusServiceUnavailable {
|
|
return nil, fmt.Errorf("llm: vllm embed returned %s: %s: %w", resp.Status, body, ErrUpstreamUnavailable)
|
|
}
|
|
return nil, fmt.Errorf("llm: vllm embed returned %s: %s", resp.Status, body)
|
|
}
|
|
|
|
var parsed embedResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
|
|
return nil, fmt.Errorf("llm: decode embed response: %w", err)
|
|
}
|
|
if len(parsed.Data) == 0 {
|
|
return nil, fmt.Errorf("llm: embed response contained no vectors")
|
|
}
|
|
|
|
vec := parsed.Data[0].Embedding
|
|
if c.embedCache != nil {
|
|
c.embedCache.put(c.EmbedModel, text, vec)
|
|
}
|
|
return vec, nil
|
|
}
|
|
|
|
// ---------------------------------------------------------------------
|
|
// embedCache: small bounded in-memory cache for Embed() results
|
|
// ---------------------------------------------------------------------
|
|
|
|
type embedCache struct {
|
|
mu sync.Mutex
|
|
cap int
|
|
order []string // insertion order, for cheap FIFO eviction
|
|
values map[string][]float32
|
|
}
|
|
|
|
func newEmbedCache(capacity int) *embedCache {
|
|
return &embedCache{cap: capacity, values: make(map[string][]float32, capacity)}
|
|
}
|
|
|
|
func embedCacheKey(model, text string) string {
|
|
h := sha256.Sum256([]byte(model + "\x00" + text))
|
|
return hex.EncodeToString(h[:])
|
|
}
|
|
|
|
func (c *embedCache) get(model, text string) ([]float32, bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
v, ok := c.values[embedCacheKey(model, text)]
|
|
return v, ok
|
|
}
|
|
|
|
func (c *embedCache) put(model, text string, vec []float32) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
key := embedCacheKey(model, text)
|
|
if _, exists := c.values[key]; !exists {
|
|
if len(c.order) >= c.cap {
|
|
// Evict the oldest entry (simple FIFO — good enough for a
|
|
// latency optimization, not a correctness-critical cache).
|
|
oldest := c.order[0]
|
|
c.order = c.order[1:]
|
|
delete(c.values, oldest)
|
|
}
|
|
c.order = append(c.order, key)
|
|
}
|
|
c.values[key] = vec
|
|
}
|