diff --git a/cmd/engine/main.go b/cmd/engine/main.go index 3b2a887..39e97a2 100644 --- a/cmd/engine/main.go +++ b/cmd/engine/main.go @@ -20,6 +20,8 @@ import ( "syscall" "time" + "github.com/hallelx2/vectorless-engine/pkg/models" + "github.com/hallelx2/llmgate" "github.com/hallelx2/llmgate/judge/typesafe" "github.com/hallelx2/llmgate/middleware/limit" @@ -415,6 +417,11 @@ func modelFor(c config.LLMConfig) string { return c.OpenAI.Model case "gemini": return c.Gemini.Model + case "vertex": + if c.Vertex.Model != "" { + return c.Vertex.Model + } + return models.DefaultModels[models.ProviderVertex] } return "" } @@ -517,6 +524,10 @@ func buildLLM(c config.LLMConfig) (llmgate.Client, error) { Model: c.Gemini.Model, ReasoningModel: c.Gemini.ReasoningModel, }) + case "vertex": + return models.BuildClient(models.Generative{Provider: models.ProviderVertex, IngestModel: c.Vertex.Model, + Project: c.Vertex.Project, Location: c.Vertex.Location, + ThinkingBudget: c.Vertex.ThinkingBudget, ThinkingLevel: c.Vertex.ThinkingLevel}) default: // Config.Validate rejects unknown drivers; this is defensive. return nil, fmt.Errorf("unknown llm driver: %s", c.Driver) diff --git a/cmd/server/main.go b/cmd/server/main.go index 629b35d..a51c889 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -25,9 +25,12 @@ import ( "net/http" "os" "os/signal" + "runtime/debug" "syscall" "time" + "github.com/hallelx2/vectorless-engine/pkg/reqtiming" + "github.com/hallelx2/llmgate" "github.com/hallelx2/llmgate/judge/typesafe" "github.com/hallelx2/llmgate/middleware/limit" @@ -57,6 +60,37 @@ import ( // version is set at build time via -ldflags "-X main.version=..." var version = "dev" +// buildVersion is version, or — for a build without -ldflags, such as +// go run or go build in a checkout — the commit the binary was built from, +// so /v1/version always names the build a benchmark measured (HAL-1582). +func buildVersion() string { + if version != "dev" { + return version + } + info, ok := debug.ReadBuildInfo() + if !ok { + return version + } + var rev, dirty string + for _, s := range info.Settings { + switch s.Key { + case "vcs.revision": + rev = s.Value + case "vcs.modified": + if s.Value == "true" { + dirty = "-dirty" + } + } + } + if rev == "" { + return version + } + if len(rev) > 12 { + rev = rev[:12] + } + return "dev-" + rev + dirty +} + func main() { if err := run(); err != nil { fmt.Fprintf(os.Stderr, "fatal: %v\n", err) @@ -76,7 +110,7 @@ func run() error { logger := newLogger(cfg.Engine.Log) logger.Info("starting vectorless-server", - "version", version, + "version", buildVersion(), "role", *role, "addr", cfg.Server.Addr, "auth_mode", cfg.Auth.Mode, @@ -103,7 +137,7 @@ func run() error { Endpoint: cfg.Tracing.Endpoint, Insecure: cfg.Tracing.Insecure, ServiceName: cfg.Tracing.ServiceName, - Version: version, + Version: buildVersion(), SampleRate: cfg.Tracing.SampleRate, }) if err != nil { @@ -135,6 +169,7 @@ func run() error { if err != nil { return fmt.Errorf("init storage: %w", err) } + store = reqtiming.WrapStorage(store) // ── Queue ───────────────────────────────────────────────────── q, err := buildQueue(cfg.Engine.Queue, cfg.Engine.Database.URL) @@ -177,11 +212,14 @@ func run() error { logModelStatus(logger, registry.Status()) // ── LLM + retrieval strategy ────────────────────────────────── - llmClient := limit.Client(newLimiter("llm", cfg.Engine.LLM.Concurrency, logger))(registry.Client()) + // reqtiming wraps outermost, outside the limiter and retries: a call + // queued at the limiter is the engine's time, and every attempt of a + // retried call is one wait on the provider. + llmClient := reqtiming.WrapClient(limit.Client(newLimiter("llm", cfg.Engine.LLM.Concurrency, logger))(registry.Client())) // judge is always the registry's proxy. Code that takes a different // path without a Judge resolves it per job or request through // models.ActiveJudge; judgewalk falls back to treewalk. - judge := registry.Judge() + judge := reqtiming.WrapJudge(registry.Judge()) strategy := buildStrategy(cfg.Engine.Retrieval, llmClient, judge, store, pool) // Wrap with caching if enabled in engine config. @@ -346,7 +384,7 @@ func run() error { Queue: q, Strategy: strategy, MultiDoc: multiDoc, - Version: version, + Version: buildVersion(), Config: cfg, Strategies: strategies, LLM: llmClient, @@ -487,6 +525,11 @@ func modelFor(c enginecfg.LLMConfig) string { return c.OpenAI.Model case "gemini": return c.Gemini.Model + case "vertex": + if c.Vertex.Model != "" { + return c.Vertex.Model + } + return models.DefaultModels[models.ProviderVertex] } return "" } @@ -592,6 +635,12 @@ func envModelSettings(cfg config.Config) models.Settings { if l.Gemini.APIKey != "" { g = &models.Generative{Provider: models.ProviderGemini, APIKey: l.Gemini.APIKey, IngestModel: l.Gemini.Model} } + case "vertex": + // No key to wait for: choosing the driver is the configuration. + // Credentials are checked when the client is built. + v := l.Vertex + g = &models.Generative{Provider: models.ProviderVertex, IngestModel: v.Model, Project: v.Project, Location: v.Location, + ThinkingBudget: v.ThinkingBudget, ThinkingLevel: v.ThinkingLevel} } if g != nil { g.AnswerModel = cfg.Models.AnswerModel @@ -643,6 +692,10 @@ func buildLLM(c enginecfg.LLMConfig) (llmgate.Client, error) { Model: c.Gemini.Model, ReasoningModel: c.Gemini.ReasoningModel, }) + case "vertex": + return models.BuildClient(models.Generative{Provider: models.ProviderVertex, IngestModel: c.Vertex.Model, + Project: c.Vertex.Project, Location: c.Vertex.Location, + ThinkingBudget: c.Vertex.ThinkingBudget, ThinkingLevel: c.Vertex.ThinkingLevel}) default: return nil, fmt.Errorf("unknown llm driver: %s", c.Driver) } diff --git a/config.server.example.yaml b/config.server.example.yaml index 488d7d7..8283554 100644 --- a/config.server.example.yaml +++ b/config.server.example.yaml @@ -77,7 +77,7 @@ engine: # concurrency: 20 llm: - driver: "anthropic" # "anthropic", "openai", or "gemini" + driver: "anthropic" # "anthropic", "openai", "gemini" or "vertex" anthropic: api_key: "" # set via VLS_ANTHROPIC_API_KEY model: "claude-sonnet-4-20250514" @@ -89,6 +89,11 @@ engine: # gemini: # api_key: "" # model: "gemini-2.0-flash" + # vertex: # Gemini on Vertex AI: no key, ADC bills the project + # project: "" # VLS_LLM_VERTEX_PROJECT (or GOOGLE_CLOUD_PROJECT) + # location: "global" # VLS_LLM_VERTEX_LOCATION + # model: "gemini-2.5-flash-lite" + # thinking_budget: 0 # VLS_LLM_VERTEX_THINKING_BUDGET; omit for the model default # reasoning_model: "" # Adaptive per-provider concurrency (HAL-1372): starts at initial, # halves on 429/transport failure, widens on success, capped at max. diff --git a/go.mod b/go.mod index 1e6a1f3..3ab3102 100644 --- a/go.mod +++ b/go.mod @@ -14,7 +14,7 @@ require ( github.com/aws/smithy-go v1.25.0 github.com/go-chi/chi/v5 v5.2.5 github.com/google/uuid v1.6.0 - github.com/hallelx2/llmgate v0.5.1-0.20260928015141-010e07c576e6 + github.com/hallelx2/llmgate v0.5.1-0.20261016012949-4e674f5d6726 github.com/hallelx2/pdftable v0.4.0 github.com/hibiken/asynq v0.26.0 github.com/jackc/pgx/v5 v5.9.2 @@ -30,6 +30,7 @@ require ( go.opentelemetry.io/otel/trace v1.44.0 golang.org/x/net v0.57.0 golang.org/x/sync v0.22.0 + google.golang.org/api v0.274.0 google.golang.org/grpc v1.83.1 google.golang.org/protobuf v1.36.11 gopkg.in/yaml.v3 v3.0.1 @@ -126,7 +127,6 @@ require ( golang.org/x/sys v0.47.0 // indirect golang.org/x/text v0.40.0 // indirect golang.org/x/time v0.15.0 // indirect - google.golang.org/api v0.274.0 // indirect google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260729162451-8efbd57d26e0 // indirect diff --git a/go.sum b/go.sum index beda7ed..a82eeaf 100644 --- a/go.sum +++ b/go.sum @@ -132,14 +132,8 @@ github.com/googleapis/gax-go/v2 v2.21.0 h1:h45NjjzEO3faG9Lg/cFrBh2PgegVVgzqKzuZl github.com/googleapis/gax-go/v2 v2.21.0/go.mod h1:But/NJU6TnZsrLai/xBAQLLz+Hc7fHZJt/hsCz3Fih4= github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 h1:8Tjv8EJ+pM1xP8mK6egEbD1OgnVTyacbefKhmbLhIhU= github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2/go.mod h1:pkJQ2tZHJ0aFOVEEot6oZmaVEZcRme73eIFmhiVuRWs= -github.com/hallelx2/llmgate v0.4.0 h1:LvaRt1PWEiR0DcNd3j8qX0GEO0oh/rKPHqfL1RABzII= -github.com/hallelx2/llmgate v0.4.0/go.mod h1:WpKwV/utKOmb+G5DwezSGjmGtz9D5I6MuS4MqAs7rkA= -github.com/hallelx2/llmgate v0.4.1-0.20260918182842-b2ec96425ecc h1:4b6tUdYQANqf5iCTnDZzY5dWHbRElVznCH4LiAovBkM= -github.com/hallelx2/llmgate v0.4.1-0.20260918182842-b2ec96425ecc/go.mod h1:WpKwV/utKOmb+G5DwezSGjmGtz9D5I6MuS4MqAs7rkA= -github.com/hallelx2/llmgate v0.5.0 h1:yNBm7NOtDrCnpeMPz2BqbNcCJeUTw0WN6oiIuEFankg= -github.com/hallelx2/llmgate v0.5.0/go.mod h1:WpKwV/utKOmb+G5DwezSGjmGtz9D5I6MuS4MqAs7rkA= -github.com/hallelx2/llmgate v0.5.1-0.20260928015141-010e07c576e6 h1:26w9OIBo0Cfu5vAVyjZgr0Uyh25B+p4jSg2dit6BuBI= -github.com/hallelx2/llmgate v0.5.1-0.20260928015141-010e07c576e6/go.mod h1:WpKwV/utKOmb+G5DwezSGjmGtz9D5I6MuS4MqAs7rkA= +github.com/hallelx2/llmgate v0.5.1-0.20261016012949-4e674f5d6726 h1:n0KEoQX4ytgC0ZrqfC5K9lapEcOlfVlcnHBsnCNqekM= +github.com/hallelx2/llmgate v0.5.1-0.20261016012949-4e674f5d6726/go.mod h1:lvdLI9zL40oeU6rwN0GIUFE5QMJ+QEekdUNLoLjFYCI= github.com/hallelx2/pdftable v0.4.0 h1:ldF8qQrUbejWsbB/JSIqliEXF2jF8z8kjEsbKS8r2BU= github.com/hallelx2/pdftable v0.4.0/go.mod h1:pxNlc4D43wjzis7M6EfgQZvHOsQ4okggm+xqUu+OokI= github.com/hhrutter/lzw v1.0.0 h1:laL89Llp86W3rRs83LvKbwYRx6INE8gDn0XNb1oXtm0= diff --git a/internal/config/config.go b/internal/config/config.go index 0cb8a65..0324349 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -344,6 +344,7 @@ func applyEnvOverrides(c *Config) { if v := firstEnv("VLS_GEMINI_API_KEY", "VLE_GEMINI_API_KEY", "VLS_LLM_GEMINI_API_KEY", "VLE_LLM_GEMINI_API_KEY"); v != "" { c.Engine.LLM.Gemini.APIKey = v } + enginecfg.ApplyVertexEnv(&c.Engine.LLM.Vertex, "VLS_", "VLE_") if v := firstEnv("VLS_STORAGE_DRIVER", "VLE_STORAGE_DRIVER"); v != "" { c.Engine.Storage.Driver = v } diff --git a/internal/handler/answer.go b/internal/handler/answer.go index ec9a717..48d88c7 100644 --- a/internal/handler/answer.go +++ b/internal/handler/answer.go @@ -12,6 +12,8 @@ import ( "sync" "time" + "github.com/hallelx2/vectorless-engine/pkg/reqtiming" + "github.com/hallelx2/llmgate" enginecfg "github.com/hallelx2/vectorless-engine/pkg/config" @@ -144,7 +146,9 @@ func (h *AnswerHandler) resolveBYOK(w http.ResponseWriter, r *http.Request) (*by writeErr(w, http.StatusBadRequest, "model key: "+err.Error()) return nil, false } - return &byokClient{client: c, model: model, provider: provider}, true + // Timed like the registry's client, so the caller's own model counts + // as an answer-model wait rather than engine time. + return &byokClient{client: reqtiming.WrapClient(c), model: model, provider: provider}, true } func (b *byokClient) or(c llmgate.Client, model string) (llmgate.Client, string) { @@ -302,6 +306,7 @@ func (h *AnswerHandler) HandleAnswer(w http.ResponseWriter, r *http.Request) { return } + withTiming(r, resp) raw, err := marshalJSONForReplay(resp) entry := retrieval.ReplayEntry{DocumentID: body.DocumentID, Query: body.Query, Model: model, SelectedIDs: finalIDs} if body.Stream { diff --git a/internal/handler/answer_treewalk.go b/internal/handler/answer_treewalk.go index 0587cce..0628317 100644 --- a/internal/handler/answer_treewalk.go +++ b/internal/handler/answer_treewalk.go @@ -227,6 +227,7 @@ func (h *AnswerTreeWalkHandler) HandleAnswerTreeWalk(w http.ResponseWriter, r *h } finalIDs := append([]tree.SectionID(nil), res.SelectedIDs...) + withTiming(r, resp) setTokenHeaders(w.Header(), resp) raw, err := marshalJSONForReplay(resp) if err != nil { @@ -309,6 +310,7 @@ func (h *AnswerTreeWalkHandler) serveStream(w http.ResponseWriter, r *http.Reque "trace_token": res.TraceToken, "pages_read": res.PagesRead, } + withTiming(r, final) emitSSE("answer", final) setTokenHeaders(w.Header(), final) } diff --git a/internal/handler/helpers.go b/internal/handler/helpers.go index 0ef1a96..65a1f7a 100644 --- a/internal/handler/helpers.go +++ b/internal/handler/helpers.go @@ -7,6 +7,8 @@ import ( "strings" "time" + "github.com/hallelx2/vectorless-engine/pkg/reqtiming" + "github.com/hallelx2/vectorless-engine/pkg/retrieval" ) @@ -41,6 +43,17 @@ func marshalJSONForReplay(v any) ([]byte, error) { // MUST see the same bytes; this is the single point where that // invariant is enforced. When store is nil or token is empty the // replay write is skipped silently. +// withTiming adds the request's timing split to a response: how much of +// the wall clock was the engine's own work and how much was waiting on +// the Judge, the answer model, the database and document storage +// (pkg/reqtiming). A request served without the timing middleware gets +// no field rather than a zero one. +func withTiming(r *http.Request, resp map[string]any) { + if rec := reqtiming.From(r.Context()); rec != nil { + resp["timing"] = rec.Snapshot() + } +} + func writeJSONWithReplay(w http.ResponseWriter, store retrieval.ReplayStore, status int, raw []byte, token string, entry retrieval.ReplayEntry) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(status) diff --git a/internal/handler/query.go b/internal/handler/query.go index 291d736..3ac34ff 100644 --- a/internal/handler/query.go +++ b/internal/handler/query.go @@ -231,6 +231,7 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) { "llm_calls": usage.LLMCalls, } } + withTiming(r, resp) setTokenHeaders(w.Header(), resp) writeJSON(w, http.StatusOK, resp) } diff --git a/internal/handler/query_multi.go b/internal/handler/query_multi.go index 0c35373..61ccdc6 100644 --- a/internal/handler/query_multi.go +++ b/internal/handler/query_multi.go @@ -141,7 +141,7 @@ func (h *QueryMultiHandler) HandleQueryMulti(w http.ResponseWriter, r *http.Requ errs[string(docID)] = e.Error() } - writeJSON(w, http.StatusOK, map[string]any{ + resp := map[string]any{ "query": body.Query, "strategy": retrieval.PublicName(h.strategy.Name()), "model": body.Model, @@ -155,5 +155,7 @@ func (h *QueryMultiHandler) HandleQueryMulti(w http.ResponseWriter, r *http.Requ "cost_usd": result.TotalUsage.CostUSD, "llm_calls": result.TotalUsage.LLMCalls, }, - }) + } + withTiming(r, resp) + writeJSON(w, http.StatusOK, resp) } diff --git a/internal/handler/query_stream.go b/internal/handler/query_stream.go index 63d7751..d7a7e3d 100644 --- a/internal/handler/query_stream.go +++ b/internal/handler/query_stream.go @@ -9,6 +9,8 @@ import ( "net/http" "time" + "github.com/hallelx2/vectorless-engine/pkg/reqtiming" + "github.com/hallelx2/vectorless-engine/pkg/db" "github.com/hallelx2/vectorless-engine/pkg/retrieval" "github.com/hallelx2/vectorless-engine/pkg/storage" @@ -49,6 +51,8 @@ type sseEvent struct { DocumentID string `json:"document_id,omitempty"` Query string `json:"query,omitempty"` ElapsedMS int64 `json:"elapsed_ms,omitempty"` + // Timing is the request\'s split, on the completed event only. + Timing *reqtiming.Timing `json:"timing,omitempty"` } // HandleQueryStream serves POST /v1/query/stream as Server-Sent Events. @@ -155,6 +159,10 @@ func (h *QueryStreamHandler) HandleQueryStream(w http.ResponseWriter, r *http.Re case retrieval.EventCompleted: sse.ElapsedMS = time.Since(started).Milliseconds() + if rec := reqtiming.From(r.Context()); rec != nil { + t := rec.Snapshot() + sse.Timing = &t + } case retrieval.EventError: if evt.Error != nil { diff --git a/internal/handler/query_stream_multi.go b/internal/handler/query_stream_multi.go index cfa1063..51e6763 100644 --- a/internal/handler/query_stream_multi.go +++ b/internal/handler/query_stream_multi.go @@ -7,6 +7,8 @@ import ( "net/http" "time" + "github.com/hallelx2/vectorless-engine/pkg/reqtiming" + "github.com/hallelx2/vectorless-engine/pkg/retrieval" "github.com/hallelx2/vectorless-engine/pkg/storage" "github.com/hallelx2/vectorless-engine/pkg/tree" @@ -45,6 +47,8 @@ type multiDocSSEEvent struct { Strategy string `json:"strategy,omitempty"` Query string `json:"query,omitempty"` ElapsedMS int64 `json:"elapsed_ms,omitempty"` + // Timing is the request\'s split, on the completed event only. + Timing *reqtiming.Timing `json:"timing,omitempty"` } // HandleQueryStreamMulti serves POST /v1/query/multi/stream as SSE. @@ -119,6 +123,10 @@ func (h *QueryStreamMultiHandler) HandleQueryStreamMulti(w http.ResponseWriter, case retrieval.EventCompleted: sse.ElapsedMS = time.Since(started).Milliseconds() + if rec := reqtiming.From(r.Context()); rec != nil { + t := rec.Snapshot() + sse.Timing = &t + } case retrieval.EventError: if evt.Error != nil { diff --git a/internal/handler/router.go b/internal/handler/router.go index 0626ff3..03e5b1c 100644 --- a/internal/handler/router.go +++ b/internal/handler/router.go @@ -4,6 +4,8 @@ import ( "log/slog" "net/http" + "github.com/hallelx2/vectorless-engine/pkg/reqtiming" + "github.com/go-chi/chi/v5" "github.com/hallelx2/llmgate" "github.com/prometheus/client_golang/prometheus/promhttp" @@ -119,6 +121,10 @@ func Router(d Deps) http.Handler { } r.Use(middleware.RequestID) + // Every request records where its time went (engine work vs waiting + // on the Judge, the answer model, the database, storage); the query + // and answer handlers return it as "timing". + r.Use(reqtiming.Middleware) r.Use(middleware.RealIP) r.Use(middleware.Recovery(d.Logger)) r.Use(middleware.AccessLog(d.Logger)) diff --git a/openapi.yaml b/openapi.yaml index d7a9a52..1c8b1b1 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -609,6 +609,31 @@ components: example: document not found schemas: + TimingWait: + type: object + description: One thing a request waited on. + properties: + calls: {type: integer, description: Calls that reached the network.} + attempts: {type: integer, description: HTTP requests, retries included.} + failed: {type: integer} + wait_ms: {type: number, description: Wall-clock time at least one call was outstanding.} + busy_ms: {type: number, description: Sum of call durations; above wait_ms when calls overlapped.} + first_byte_ms: {type: number, description: Sum over calls of time from going on the wire to the first response byte.} + first_byte_max_ms: {type: number} + local_ms: {type: number, description: Time calls spent in the engine before going on the wire (limiter queue, prompt building).} + Timing: + type: object + description: >- + Where this request's time went. engine_ms + waiting_ms == total_ms; + waiting_ms is the union of the categories' waits. + properties: + total_ms: {type: number} + engine_ms: {type: number, description: Time not spent waiting on any provider, the database or storage.} + waiting_ms: {type: number} + judge: {$ref: '#/components/schemas/TimingWait'} + answer_model: {$ref: '#/components/schemas/TimingWait'} + db: {$ref: '#/components/schemas/TimingWait'} + storage: {$ref: '#/components/schemas/TimingWait'} IngestMode: type: string enum: [toc, full, minimal] @@ -804,6 +829,8 @@ components: $ref: "#/components/schemas/Plan" elapsed_ms: type: integer + timing: + $ref: "#/components/schemas/Timing" trace_token: type: string description: | @@ -1062,6 +1089,8 @@ components: $ref: "#/components/schemas/Plan" elapsed_ms: type: integer + timing: + $ref: "#/components/schemas/Timing" trace_token: type: string description: | @@ -1328,6 +1357,8 @@ components: llm_calls: {type: integer} elapsed_ms: type: integer + timing: + $ref: "#/components/schemas/Timing" trace_token: type: string description: | diff --git a/pkg/config/config.go b/pkg/config/config.go index befbcac..26b2fe7 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -376,6 +376,7 @@ type LLMConfig struct { Anthropic AnthropicBlock `yaml:"anthropic"` OpenAI OpenAIBlock `yaml:"openai"` Gemini GeminiBlock `yaml:"gemini"` + Vertex VertexBlock `yaml:"vertex"` // Concurrency bounds calls in flight per provider — one adaptive // limiter for the chat model, one for the Judge. The limit starts at @@ -456,6 +457,27 @@ type GeminiBlock struct { ReasoningModel string `yaml:"reasoning_model"` } +// VertexBlock configures Gemini on Google Cloud Vertex AI (driver +// "vertex"). It needs no API key: calls authenticate with Application +// Default Credentials (a VM's service account, or a gcloud login) and +// bill the project, so a project's Google Cloud credits pay for them. +type VertexBlock struct { + // Project is the billed Google Cloud project. Empty falls back to + // GOOGLE_CLOUD_PROJECT, then to the credentials' own project. + Project string `yaml:"project"` + // Location is "global" (default, and the cheaper Gemini 3 price) or + // a region such as us-central1. + Location string `yaml:"location"` + Model string `yaml:"model"` + // ThinkingBudget caps thinking tokens per call; 0 turns thinking off + // on models that allow it. Nil leaves the model's default. Hold it + // fixed across a benchmark so every system gets the same model. + ThinkingBudget *int `yaml:"thinking_budget"` + // ThinkingLevel is the Gemini 3 alternative to a budget: minimal, + // low, medium or high. Set at most one of the two. + ThinkingLevel string `yaml:"thinking_level"` +} + // RetrievalConfig configures the retrieval strategy. type RetrievalConfig struct { Strategy string `yaml:"strategy"` @@ -947,6 +969,7 @@ func applyEnvOverrides(c *Config) { if v := os.Getenv("VLE_GEMINI_API_KEY"); v != "" { c.LLM.Gemini.APIKey = v } + applyVertexEnv(&c.LLM.Vertex, "VLE_") // The Judge key accepts the deploy layer's VLS_ prefix and the bare // TYPESAFE_API_KEY — what the llmgate live tests and the bench // commands already read — so one export enables it everywhere. @@ -1315,7 +1338,7 @@ func (c Config) Validate() error { } switch c.LLM.Driver { - case "anthropic", "openai", "gemini", "openai-compatible": + case "anthropic", "openai", "gemini", "openai-compatible", "vertex": // API keys are checked lazily at first call so the engine can boot // in dev without all providers configured. default: @@ -1437,3 +1460,37 @@ func (c Config) Validate() error { return nil } + +// applyVertexEnv reads the Vertex block from LLM_VERTEX_* — +// PROJECT, LOCATION, MODEL, THINKING_BUDGET, THINKING_LEVEL. A thinking +// budget that is not a non-negative integer is ignored. +func applyVertexEnv(v *VertexBlock, prefixes ...string) { + get := func(name string) string { + names := make([]string, 0, len(prefixes)) + for _, p := range prefixes { + names = append(names, p+"LLM_VERTEX_"+name) + } + return firstEnv(names...) + } + if s := get("PROJECT"); s != "" { + v.Project = s + } + if s := get("LOCATION"); s != "" { + v.Location = s + } + if s := get("MODEL"); s != "" { + v.Model = s + } + if s := get("THINKING_BUDGET"); s != "" { + if n, err := strconv.Atoi(s); err == nil && n >= 0 { + v.ThinkingBudget = &n + } + } + if s := get("THINKING_LEVEL"); s != "" { + v.ThinkingLevel = s + } +} + +// ApplyVertexEnv is applyVertexEnv for the server's config layer, which +// reads VLS_ before VLE_. +func ApplyVertexEnv(v *VertexBlock, prefixes ...string) { applyVertexEnv(v, prefixes...) } diff --git a/pkg/config/vertex_env_test.go b/pkg/config/vertex_env_test.go new file mode 100644 index 0000000..8aded01 --- /dev/null +++ b/pkg/config/vertex_env_test.go @@ -0,0 +1,39 @@ +package config + +import "testing" + +func TestVertexEnv(t *testing.T) { + t.Setenv("VLS_LLM_VERTEX_PROJECT", "proj-a") + t.Setenv("VLE_LLM_VERTEX_PROJECT", "proj-b") + t.Setenv("VLE_LLM_VERTEX_LOCATION", "global") + t.Setenv("VLS_LLM_VERTEX_MODEL", "gemini-2.5-flash-lite") + t.Setenv("VLS_LLM_VERTEX_THINKING_BUDGET", "0") + + var v VertexBlock + ApplyVertexEnv(&v, "VLS_", "VLE_") + if v.Project != "proj-a" { + t.Errorf("project = %q, want VLS_ to win", v.Project) + } + if v.Location != "global" || v.Model != "gemini-2.5-flash-lite" { + t.Errorf("location/model = %q/%q", v.Location, v.Model) + } + if v.ThinkingBudget == nil || *v.ThinkingBudget != 0 { + t.Errorf("thinking budget = %v, want 0 (an explicit zero turns thinking off)", v.ThinkingBudget) + } + + t.Setenv("VLS_LLM_VERTEX_THINKING_BUDGET", "lots") + var w VertexBlock + ApplyVertexEnv(&w, "VLS_") + if w.ThinkingBudget != nil { + t.Errorf("a garbled budget must be ignored, got %v", *w.ThinkingBudget) + } +} + +func TestVertexDriverValidates(t *testing.T) { + c := Default() + c.LLM.Driver = "vertex" + c.Database.URL = "postgres://localhost/engine" + if err := c.Validate(); err != nil { + t.Errorf("Validate with driver vertex = %v", err) + } +} diff --git a/pkg/db/db.go b/pkg/db/db.go index a8f007d..b86ed44 100644 --- a/pkg/db/db.go +++ b/pkg/db/db.go @@ -15,6 +15,8 @@ import ( "strings" "time" + "github.com/hallelx2/vectorless-engine/pkg/reqtiming" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" @@ -44,6 +46,12 @@ func Open(ctx context.Context, url string, maxConns int32) (*Pool, error) { if maxConns > 0 { cfg.MaxConns = maxConns } + // Queries made with a request's context count as that request's + // database wait (pkg/reqtiming). Without a recorder on the context + // the tracer does nothing. + if cfg.ConnConfig.Tracer == nil { + cfg.ConnConfig.Tracer = reqtiming.QueryTracer{} + } pool, err := pgxpool.NewWithConfig(ctx, cfg) if err != nil { return nil, fmt.Errorf("connect pool: %w", err) diff --git a/pkg/models/build.go b/pkg/models/build.go index 85d5b5b..6829e80 100644 --- a/pkg/models/build.go +++ b/pkg/models/build.go @@ -12,6 +12,7 @@ import ( "github.com/hallelx2/llmgate/provider/anthropic" "github.com/hallelx2/llmgate/provider/gemini" "github.com/hallelx2/llmgate/provider/openai" + "github.com/hallelx2/llmgate/provider/vertex" ) // keylessPlaceholder is sent as the bearer token to an openai-compatible @@ -39,6 +40,17 @@ func BuildClient(g Generative) (llmgate.Client, error) { return openai.New(openai.Config{APIKey: key, Model: model, BaseURL: g.BaseURL}) case ProviderGemini: return gemini.New(gemini.Config{APIKey: g.APIKey, Model: model, AuthTimeout: 15 * time.Second}) + case ProviderVertex: + // Bounded like gemini's AuthTimeout: finding credentials can ask + // the metadata server, which hangs off GCP. + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + return vertex.New(ctx, vertex.Config{ + Project: g.Project, + Location: g.Location, + Model: model, + Thinking: vertex.Thinking{Budget: g.ThinkingBudget, Level: g.ThinkingLevel}, + }) } return nil, fmt.Errorf("unknown provider %q", g.Provider) } diff --git a/pkg/models/models.go b/pkg/models/models.go index 8c73a7e..d854ebc 100644 --- a/pkg/models/models.go +++ b/pkg/models/models.go @@ -31,11 +31,15 @@ const ( ProviderOpenAI = "openai" ProviderGemini = "gemini" ProviderOpenAICompatible = "openai-compatible" + // ProviderVertex is Gemini through Google Cloud Vertex AI. It takes no + // API key: Application Default Credentials authenticate it, and the + // project it names is billed. + ProviderVertex = "vertex" ) // Providers lists the generative providers in the order the dashboard // offers them. -var Providers = []string{ProviderAnthropic, ProviderOpenAI, ProviderGemini, ProviderOpenAICompatible} +var Providers = []string{ProviderAnthropic, ProviderOpenAI, ProviderGemini, ProviderVertex, ProviderOpenAICompatible} // DefaultModels is the model used when a provider is chosen and no model // is named. openai-compatible has none: its server decides what exists. @@ -43,6 +47,7 @@ var DefaultModels = map[string]string{ ProviderAnthropic: "claude-sonnet-4-5", ProviderOpenAI: "gpt-4.1-mini", ProviderGemini: "gemini-2.5-flash", + ProviderVertex: "gemini-2.5-flash-lite", } // Source says where an effective setting came from. @@ -68,6 +73,16 @@ type Generative struct { // IngestModel is the default for every other call: ingest stages and // navigation. Empty uses the provider's default. IngestModel string `json:"ingest_model,omitempty"` + + // Project and Location address Vertex (provider "vertex" only). + // Empty Project falls back to GOOGLE_CLOUD_PROJECT, then to the + // credentials' project; empty Location is "global". + Project string `json:"project,omitempty"` + Location string `json:"location,omitempty"` + // ThinkingBudget and ThinkingLevel cap Gemini's thinking on Vertex. + // At most one is set; neither leaves the model's default. + ThinkingBudget *int `json:"thinking_budget,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` } // JudgeSettings configures the Judge (TypeSafe Jev). @@ -107,6 +122,22 @@ func (g Generative) Validate() error { if strings.TrimSpace(g.APIKey) == "" { return fmt.Errorf("%s needs an API key", g.Provider) } + case ProviderVertex: + if strings.TrimSpace(g.BaseURL) != "" { + return errors.New("vertex does not take a base URL; set project and location instead") + } + if g.ThinkingBudget != nil && strings.TrimSpace(g.ThinkingLevel) != "" { + return errors.New("vertex takes a thinking budget or a thinking level, not both") + } + if g.ThinkingBudget != nil && *g.ThinkingBudget < 0 { + return errors.New("vertex thinking budget cannot be negative") + } + switch strings.TrimSpace(strings.ToLower(g.ThinkingLevel)) { + case "", "minimal", "low", "medium", "high": + default: + return fmt.Errorf("vertex thinking level %q is not one of minimal, low, medium, high", g.ThinkingLevel) + } + return nil case ProviderOpenAICompatible: if strings.TrimSpace(g.BaseURL) == "" { return errors.New("openai-compatible needs a base URL (for example http://ollama:11434/v1)") @@ -297,7 +328,21 @@ func trimGenerative(g Generative) Generative { BaseURL: strings.TrimRight(strings.TrimSpace(g.BaseURL), "/"), AnswerModel: strings.TrimSpace(g.AnswerModel), IngestModel: strings.TrimSpace(g.IngestModel), + // Vertex: dropping these would silently bill a different project + // or think by default. + Project: strings.TrimSpace(g.Project), + Location: strings.TrimSpace(strings.ToLower(g.Location)), + ThinkingBudget: copyInt(g.ThinkingBudget), + ThinkingLevel: strings.TrimSpace(strings.ToLower(g.ThinkingLevel)), + } +} + +func copyInt(p *int) *int { + if p == nil { + return nil } + v := *p + return &v } func trimJudge(j JudgeSettings) JudgeSettings { diff --git a/pkg/models/vertex_test.go b/pkg/models/vertex_test.go new file mode 100644 index 0000000..4744f83 --- /dev/null +++ b/pkg/models/vertex_test.go @@ -0,0 +1,41 @@ +package models + +import ( + "strings" + "testing" +) + +func intp(n int) *int { return &n } + +func TestVertexValidate(t *testing.T) { + ok := []Generative{ + {Provider: ProviderVertex}, + {Provider: ProviderVertex, Project: "p", Location: "us-central1", IngestModel: "gemini-3.5-flash-lite", ThinkingBudget: intp(0)}, + {Provider: ProviderVertex, ThinkingLevel: "low"}, + } + for _, g := range ok { + if err := g.Validate(); err != nil { + t.Errorf("Validate(%+v) = %v, want nil (vertex needs no API key)", g, err) + } + } + bad := map[string]Generative{ + "base URL": {Provider: ProviderVertex, BaseURL: "https://example.com"}, + "both": {Provider: ProviderVertex, ThinkingBudget: intp(0), ThinkingLevel: "low"}, + "negative cap": {Provider: ProviderVertex, ThinkingBudget: intp(-1)}, + } + for name, g := range bad { + if err := g.Validate(); err == nil { + t.Errorf("%s: Validate = nil, want an error", name) + } + } +} + +func TestVertexDefaultsAndListing(t *testing.T) { + g := Generative{Provider: ProviderVertex} + if got := g.ResolvedIngestModel(); got != "gemini-2.5-flash-lite" { + t.Errorf("default model = %q", got) + } + if !strings.Contains(strings.Join(Providers, ","), ProviderVertex) { + t.Error("vertex missing from Providers") + } +} diff --git a/pkg/models/vertex_trim_test.go b/pkg/models/vertex_trim_test.go new file mode 100644 index 0000000..5989205 --- /dev/null +++ b/pkg/models/vertex_trim_test.go @@ -0,0 +1,27 @@ +package models + +import "testing" + +func TestTrimKeepsVertexFields(t *testing.T) { + b := 0 + g := trimGenerative(Generative{Provider: " Vertex ", Project: " p ", Location: "US-CENTRAL1", ThinkingBudget: &b, ThinkingLevel: ""}) + if g.Provider != ProviderVertex || g.Project != "p" || g.Location != "us-central1" { + t.Errorf("trimmed = %+v", g) + } + if g.ThinkingBudget == nil || *g.ThinkingBudget != 0 { + t.Fatalf("thinking budget lost: %v", g.ThinkingBudget) + } + b = 9 + if *g.ThinkingBudget != 0 { + t.Error("trim must copy the budget, not alias it") + } +} + +func TestVertexThinkingLevelIsClosed(t *testing.T) { + if err := (Generative{Provider: ProviderVertex, ThinkingLevel: "extreme"}).Validate(); err == nil { + t.Error("an unknown thinking level validated") + } + if err := (Generative{Provider: ProviderVertex, ThinkingLevel: "High"}).Validate(); err != nil { + t.Errorf("High should validate: %v", err) + } +} diff --git a/pkg/reqtiming/reqtiming.go b/pkg/reqtiming/reqtiming.go new file mode 100644 index 0000000..06482af --- /dev/null +++ b/pkg/reqtiming/reqtiming.go @@ -0,0 +1,299 @@ +// Package reqtiming records where one API request's wall-clock time went: +// the engine's own work, or waiting on the Judge, the answer model, the +// database or document storage. +// +// A round-trip number cannot tell the engine's CPU from a provider's. The +// 37–50 s queries of September were the engine rebuilding a tokenizer on +// every count, not Jev (HAL-1708), and two sessions tuned the wrong thing +// against that number. So every query and answer response carries a +// timing object built here, per request, from the request's own context, +// rather than from log lines a benchmark would have to scrape and join. +// +// How the split is measured: +// +// - A provider call's wait starts at its first connection attempt +// (httptrace GetConn) and ends when the call returns. Work before the +// request goes on the wire — building the prompt, counting tokens, +// queueing at the concurrency limiter — is the engine's. +// - Calls run in parallel, so waits are merged as intervals: a category's +// wait_ms is the wall-clock time at least one of its calls was on the +// wire, and busy_ms is the plain sum. +// - engine_ms (prepare) is the total minus the union of every wait, so +// engine_ms + waiting_ms == total_ms exactly, and the categories' +// wait_ms add up to waiting_ms unless two categories overlapped. +package reqtiming + +import ( + "context" + "crypto/tls" + "math" + "net/http/httptrace" + "sort" + "sync" + "time" +) + +// Category names a thing a request can wait on. +type Category string + +const ( + // Judge is the System One model (TypeSafe Jev). + Judge Category = "judge" + // AnswerModel is the generative model: answers, and any generative + // navigation or extraction call. + AnswerModel Category = "answer_model" + // DB is Postgres. + DB Category = "db" + // Storage is the document store (local disk, S3, GCS). + Storage Category = "storage" +) + +var categories = []Category{Judge, AnswerModel, DB, Storage} + +type interval struct{ start, end time.Time } + +type call struct { + interval + local time.Duration // entry to first connection attempt + firstByte time.Duration // first connection attempt to first response byte + attempts int + failed bool +} + +// Recorder collects one request's waits. Safe for concurrent use. +type Recorder struct { + start time.Time + now func() time.Time + + mu sync.Mutex + calls map[Category][]call +} + +type ctxKey struct{} + +// Start returns a context carrying a new Recorder, started now. +func Start(ctx context.Context) (context.Context, *Recorder) { + r := newRecorder(time.Now) + return context.WithValue(ctx, ctxKey{}, r), r +} + +func newRecorder(now func() time.Time) *Recorder { + return &Recorder{start: now(), now: now, calls: map[Category][]call{}} +} + +// From returns the context's Recorder, or nil. Every method on a nil +// Recorder is a no-op, so callers need not check. +func From(ctx context.Context) *Recorder { + r, _ := ctx.Value(ctxKey{}).(*Recorder) + return r +} + +// Observe records a wait that is not an HTTP call (a database query, a +// storage read) from start until now. +func (r *Recorder) Observe(cat Category, start time.Time, err error) { + if r == nil { + return + } + end := r.now() + r.add(cat, call{interval: interval{start, end}, attempts: 1, failed: err != nil}) +} + +func (r *Recorder) add(cat Category, c call) { + if c.end.Before(c.start) { + c.end = c.start + } + r.mu.Lock() + r.calls[cat] = append(r.calls[cat], c) + r.mu.Unlock() +} + +// Span measures one provider call made with the returned context. Call +// done with the call's error when it returns. +// +// The wait is timed from the first connection attempt the call makes, +// observed through net/http/httptrace on the context, so any local work +// the client does first is not counted as waiting. A call that never +// reached the network (a cache hit, an error before sending) records no +// wait at all. +func (r *Recorder) Span(ctx context.Context, cat Category) (context.Context, func(error)) { + if r == nil { + return ctx, func(error) {} + } + entered := r.now() + var ( + mu sync.Mutex + wireStart time.Time + firstByte time.Time + attempts int + ) + trace := &httptrace.ClientTrace{ + GetConn: func(string) { + mu.Lock() + if wireStart.IsZero() { + wireStart = r.now() + } + attempts++ + mu.Unlock() + }, + GotFirstResponseByte: func() { + mu.Lock() + if firstByte.IsZero() { + firstByte = r.now() + } + mu.Unlock() + }, + // Present so a TLS handshake on a fresh connection is part of + // the wait (it follows GetConn); nothing to record. + TLSHandshakeDone: func(tls.ConnectionState, error) {}, + } + ctx = httptrace.WithClientTrace(ctx, trace) + return ctx, func(err error) { + end := r.now() + mu.Lock() + defer mu.Unlock() + if wireStart.IsZero() { + return + } + c := call{interval: interval{wireStart, end}, local: wireStart.Sub(entered), attempts: attempts, failed: err != nil} + if !firstByte.IsZero() { + c.firstByte = firstByte.Sub(wireStart) + } + r.add(cat, c) + } +} + +// Wait is one category's share of a request. +type Wait struct { + // Calls is the number of calls that reached the network. + Calls int `json:"calls"` + // Attempts counts HTTP requests, retries included. + Attempts int `json:"attempts,omitempty"` + // Failed is the number of calls that returned an error. + Failed int `json:"failed,omitempty"` + // WaitMS is wall-clock time at least one call was outstanding. + WaitMS float64 `json:"wait_ms"` + // BusyMS is the sum of every call's duration; above WaitMS when + // calls overlapped. + BusyMS float64 `json:"busy_ms"` + // FirstByteMS is the time from a call going on the wire to the first + // byte of its response: the provider's own latency. Sum and max. + FirstByteMS float64 `json:"first_byte_ms,omitempty"` + FirstByteMaxMS float64 `json:"first_byte_max_ms,omitempty"` + // LocalMS is the time calls spent in the engine before going on the + // wire (limiter queue, prompt building, token counting). It is part + // of EngineMS, reported here so a slow client is visible. + LocalMS float64 `json:"local_ms,omitempty"` +} + +// Timing is the wire shape of a request's split. +type Timing struct { + TotalMS float64 `json:"total_ms"` + // EngineMS is time not spent waiting on anything below: parsing the + // request, walking the tree, ranking, assembling the response. + EngineMS float64 `json:"engine_ms"` + // WaitingMS is the union of every category's waits. + WaitingMS float64 `json:"waiting_ms"` + Judge Wait `json:"judge"` + AnswerModel Wait `json:"answer_model"` + DB Wait `json:"db"` + Storage Wait `json:"storage"` +} + +// Snapshot computes the split from the start until now. +func (r *Recorder) Snapshot() Timing { + if r == nil { + return Timing{} + } + end := r.now() + total := end.Sub(r.start) + + r.mu.Lock() + defer r.mu.Unlock() + + var all []interval + out := Timing{TotalMS: ms(total)} + for _, cat := range categories { + calls := r.calls[cat] + var w Wait + ivs := make([]interval, 0, len(calls)) + for _, c := range calls { + w.Calls++ + w.Attempts += c.attempts + if c.failed { + w.Failed++ + } + d := c.end.Sub(c.start) + w.BusyMS += ms(d) + w.FirstByteMS += ms(c.firstByte) + if fb := ms(c.firstByte); fb > w.FirstByteMaxMS { + w.FirstByteMaxMS = fb + } + w.LocalMS += ms(c.local) + ivs = append(ivs, clip(c.interval, r.start, end)) + } + w.WaitMS = ms(unionLen(ivs)) + w.BusyMS, w.FirstByteMS, w.LocalMS = round2(w.BusyMS), round2(w.FirstByteMS), round2(w.LocalMS) + all = append(all, ivs...) + switch cat { + case Judge: + out.Judge = w + case AnswerModel: + out.AnswerModel = w + case DB: + out.DB = w + case Storage: + out.Storage = w + } + } + // Round the total and the waiting time once each and derive the engine + // share from the rounded pair, so the documented identity holds exactly + // rather than to within a rounding step. + out.WaitingMS = ms(unionLen(all)) + out.EngineMS = round2(out.TotalMS - out.WaitingMS) + return out +} + +func clip(iv interval, lo, hi time.Time) interval { + if iv.start.Before(lo) { + iv.start = lo + } + if iv.end.After(hi) { + iv.end = hi + } + if iv.end.Before(iv.start) { + iv.end = iv.start + } + return iv +} + +// unionLen is the length of the union of intervals. +func unionLen(ivs []interval) time.Duration { + if len(ivs) == 0 { + return 0 + } + s := append([]interval(nil), ivs...) + sort.Slice(s, func(i, j int) bool { return s[i].start.Before(s[j].start) }) + var total time.Duration + cur := s[0] + for _, iv := range s[1:] { + if !iv.start.After(cur.end) { + if iv.end.After(cur.end) { + cur.end = iv.end + } + continue + } + total += cur.end.Sub(cur.start) + cur = iv + } + return total + cur.end.Sub(cur.start) +} + +// round2 rounds a sum of millisecond values to two decimals, so adding +// many per-call figures does not print float noise (4.5599999). +func round2(v float64) float64 { return math.Round(v*100) / 100 } + +// ms rounds a duration to milliseconds with two decimals, so the parts of +// a fast request still add up. +func ms(d time.Duration) float64 { + return float64(d.Round(10*time.Microsecond)) / float64(time.Millisecond) +} diff --git a/pkg/reqtiming/reqtiming_test.go b/pkg/reqtiming/reqtiming_test.go new file mode 100644 index 0000000..836d45c --- /dev/null +++ b/pkg/reqtiming/reqtiming_test.go @@ -0,0 +1,262 @@ +package reqtiming + +import ( + "context" + "errors" + "io" + "math" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/hallelx2/llmgate" + + "github.com/hallelx2/vectorless-engine/pkg/storage" +) + +// fakeClock hands out times the test controls. +type fakeClock struct { + mu sync.Mutex + now time.Time +} + +func (c *fakeClock) Now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.now +} + +func (c *fakeClock) Advance(d time.Duration) { + c.mu.Lock() + c.now = c.now.Add(d) + c.mu.Unlock() +} + +func TestSnapshotSplitsEngineFromWaits(t *testing.T) { + clk := &fakeClock{now: time.Unix(0, 0)} + r := newRecorder(clk.Now) + + // 0–10ms engine work, then a 30ms judge call, then 5ms engine, then a + // 20ms answer call, then 5ms engine: total 70, waiting 50, engine 20. + clk.Advance(10 * time.Millisecond) + start := clk.Now() + clk.Advance(30 * time.Millisecond) + r.add(Judge, call{interval: interval{start, clk.Now()}, attempts: 1}) + clk.Advance(5 * time.Millisecond) + start = clk.Now() + clk.Advance(20 * time.Millisecond) + r.add(AnswerModel, call{interval: interval{start, clk.Now()}, attempts: 1}) + clk.Advance(5 * time.Millisecond) + + s := r.Snapshot() + if s.TotalMS != 70 || s.WaitingMS != 50 || s.EngineMS != 20 { + t.Fatalf("total/waiting/engine = %v/%v/%v, want 70/50/20", s.TotalMS, s.WaitingMS, s.EngineMS) + } + if s.Judge.WaitMS != 30 || s.AnswerModel.WaitMS != 20 || s.Judge.Calls != 1 { + t.Errorf("judge/answer wait = %v/%v", s.Judge.WaitMS, s.AnswerModel.WaitMS) + } + if s.EngineMS+s.WaitingMS != s.TotalMS { + t.Errorf("engine + waiting must equal total exactly") + } +} + +func TestParallelCallsMergeIntoOneWait(t *testing.T) { + clk := &fakeClock{now: time.Unix(0, 0)} + r := newRecorder(clk.Now) + t0 := clk.Now() + // Two judge calls overlapping: [0,40] and [10,50] → wall 50, busy 80. + r.add(Judge, call{interval: interval{t0, t0.Add(40 * time.Millisecond)}}) + r.add(Judge, call{interval: interval{t0.Add(10 * time.Millisecond), t0.Add(50 * time.Millisecond)}}) + clk.Advance(60 * time.Millisecond) + + s := r.Snapshot() + if s.Judge.WaitMS != 50 || s.Judge.BusyMS != 80 || s.Judge.Calls != 2 { + t.Errorf("wait/busy/calls = %v/%v/%d, want 50/80/2", s.Judge.WaitMS, s.Judge.BusyMS, s.Judge.Calls) + } + if s.EngineMS != 10 { + t.Errorf("engine = %v, want 10", s.EngineMS) + } +} + +func TestNilRecorderIsANoOp(t *testing.T) { + var r *Recorder + ctx, done := r.Span(context.Background(), Judge) + done(nil) + r.Observe(DB, time.Now(), nil) + if got := r.Snapshot(); got != (Timing{}) { + t.Errorf("nil Snapshot = %+v", got) + } + if From(ctx) != nil { + t.Error("no recorder expected on a plain context") + } +} + +// A real HTTP call through Span: the wait starts when the request goes on +// the wire, so time spent before that (here, a sleep standing in for a +// limiter queue) is local, not waiting. +func TestSpanTimesTheWireNotTheQueue(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(40 * time.Millisecond) + _, _ = io.WriteString(w, "ok") + })) + defer srv.Close() + + ctx, r := Start(context.Background()) + sctx, done := r.Span(ctx, AnswerModel) + time.Sleep(30 * time.Millisecond) // queued before sending + req, _ := http.NewRequestWithContext(sctx, http.MethodGet, srv.URL, nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + done(nil) + + s := r.Snapshot() + a := s.AnswerModel + if a.Calls != 1 || a.Attempts != 1 { + t.Fatalf("calls/attempts = %d/%d", a.Calls, a.Attempts) + } + // A floor only: a loaded test host can stretch the wait, never shrink + // it below the server's sleep. + if a.WaitMS < 35 { + t.Errorf("wait = %vms, want at least the server's 40ms", a.WaitMS) + } + if a.LocalMS < 25 { + t.Errorf("local = %vms, want the ~30ms spent before sending", a.LocalMS) + } + if a.FirstByteMS < 35 { + t.Errorf("first byte = %vms, want ~40", a.FirstByteMS) + } + if math.Abs(s.EngineMS+s.WaitingMS-s.TotalMS) > 0.02 { + t.Errorf("parts %v+%v do not sum to total %v", s.EngineMS, s.WaitingMS, s.TotalMS) + } +} + +func TestSpanWithoutNetworkRecordsNothing(t *testing.T) { + ctx, r := Start(context.Background()) + _, done := r.Span(ctx, Judge) + done(errors.New("rejected before sending")) + if s := r.Snapshot(); s.Judge.Calls != 0 { + t.Errorf("a call that never reached the network recorded %d calls", s.Judge.Calls) + } +} + +type stubClient struct{ fn func(ctx context.Context) } + +func (s stubClient) Complete(ctx context.Context, _ llmgate.Request) (*llmgate.Response, error) { + s.fn(ctx) + return &llmgate.Response{Content: "ok"}, nil +} +func (stubClient) CountTokens(context.Context, string) (int, error) { return 1, nil } + +type stubJudge struct{ fn func(ctx context.Context) } + +func (s stubJudge) Judge(ctx context.Context, _ llmgate.JudgeRequest) (*llmgate.Judgment, error) { + s.fn(ctx) + return &llmgate.Judgment{}, nil +} + +func hit(t *testing.T, url string) func(ctx context.Context) { + return func(ctx context.Context) { + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Error(err) + return + } + _ = resp.Body.Close() + } +} + +func TestWrappersAttributeToTheirCategory(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(10 * time.Millisecond) + })) + defer srv.Close() + + ctx, r := Start(context.Background()) + if _, err := WrapClient(stubClient{fn: hit(t, srv.URL)}).Complete(ctx, llmgate.Request{}); err != nil { + t.Fatal(err) + } + if _, err := WrapJudge(stubJudge{fn: hit(t, srv.URL)}).Judge(ctx, llmgate.JudgeRequest{}); err != nil { + t.Fatal(err) + } + s := r.Snapshot() + if s.AnswerModel.Calls != 1 || s.Judge.Calls != 1 { + t.Errorf("answer/judge calls = %d/%d, want 1/1", s.AnswerModel.Calls, s.Judge.Calls) + } + if WrapClient(nil) != nil || WrapJudge(nil) != nil || WrapStorage(nil) != nil { + t.Error("wrapping nil must stay nil so 'not configured' checks keep working") + } +} + +type memStore struct { + storage.Storage + data string +} + +func (m memStore) Get(context.Context, string) (io.ReadCloser, storage.Metadata, error) { + return io.NopCloser(strings.NewReader(m.data)), storage.Metadata{}, nil +} + +func TestStorageReadCountsUntilClose(t *testing.T) { + ctx, r := Start(context.Background()) + rc, _, err := WrapStorage(memStore{data: "hello"}).Get(ctx, "k") + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(rc) + _ = rc.Close() + if string(b) != "hello" { + t.Fatalf("read %q", b) + } + if s := r.Snapshot(); s.Storage.Calls != 1 { + t.Errorf("storage calls = %d, want exactly 1 (EOF and Close must not double count)", s.Storage.Calls) + } +} + +func TestMiddlewareInstallsARecorder(t *testing.T) { + var got *Recorder + h := Middleware(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { got = From(r.Context()) })) + h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/v1/query", nil)) + if got == nil { + t.Fatal("handler saw no recorder") + } +} + +type activeProxy struct{ inner llmgate.Judge } + +func (a activeProxy) Judge(ctx context.Context, r llmgate.JudgeRequest) (*llmgate.Judgment, error) { + return a.inner.Judge(ctx, r) +} +func (a activeProxy) Active() llmgate.Judge { return a.inner } + +func TestTimedJudgeForwardsActive(t *testing.T) { + none := WrapJudge(activeProxy{inner: nil}).(interface{ Active() llmgate.Judge }) + if none.Active() != nil { + t.Error("an unconfigured Judge must stay nil through the wrapper") + } + some := WrapJudge(activeProxy{inner: stubJudge{fn: func(context.Context) {}}}).(interface{ Active() llmgate.Judge }) + if _, ok := some.Active().(timedJudge); !ok { + t.Error("an active Judge must come back timed") + } +} + +func TestSumIsExactAfterRounding(t *testing.T) { + clk := &fakeClock{now: time.Unix(0, 0)} + r := newRecorder(clk.Now) + t0 := clk.Now() + r.add(Judge, call{interval: interval{t0.Add(3333 * time.Microsecond), t0.Add(7777 * time.Microsecond)}}) + clk.Advance(12345 * time.Microsecond) + s := r.Snapshot() + // Equal to the hundredth of a millisecond the fields carry; float + // addition of two-decimal values may differ by an ulp. + if math.Abs(s.EngineMS+s.WaitingMS-s.TotalMS) > 1e-9 { + t.Errorf("%v + %v != %v", s.EngineMS, s.WaitingMS, s.TotalMS) + } +} diff --git a/pkg/reqtiming/wrap.go b/pkg/reqtiming/wrap.go new file mode 100644 index 0000000..0c57a53 --- /dev/null +++ b/pkg/reqtiming/wrap.go @@ -0,0 +1,168 @@ +package reqtiming + +import ( + "context" + "io" + "net/http" + "sync" + "time" + + "github.com/hallelx2/llmgate" + "github.com/jackc/pgx/v5" + + "github.com/hallelx2/vectorless-engine/pkg/storage" +) + +// Middleware starts a Recorder for every request, so any handler below it +// can report its split with From(r.Context()).Snapshot(). +func Middleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx, _ := Start(r.Context()) + next.ServeHTTP(w, r.WithContext(ctx)) + }) +} + +// WrapClient times every Complete call made with a request context as an +// AnswerModel wait. CountTokens is local work and is not wrapped. +// +// Wrap outermost — outside retry and the concurrency limiter — so a call +// queued at the limiter counts as the engine's time (LocalMS), and every +// retry attempt is part of one wait. +func WrapClient(c llmgate.Client) llmgate.Client { + if c == nil { + return nil + } + return timedClient{next: c} +} + +type timedClient struct{ next llmgate.Client } + +func (t timedClient) Complete(ctx context.Context, req llmgate.Request) (*llmgate.Response, error) { + ctx, done := From(ctx).Span(ctx, AnswerModel) + resp, err := t.next.Complete(ctx, req) + done(err) + return resp, err +} + +func (t timedClient) CountTokens(ctx context.Context, text string) (int, error) { + return t.next.CountTokens(ctx, text) +} + +// WrapJudge times every Judge call made with a request context as a Judge +// wait. Same placement rule as WrapClient. +func WrapJudge(j llmgate.Judge) llmgate.Judge { + if j == nil { + return nil + } + return timedJudge{next: j} +} + +type timedJudge struct{ next llmgate.Judge } + +// Active forwards the registry proxy's Active (models.ActiveJudge): nil +// when no Judge is configured, so callers still take their no-Judge path, +// and the active Judge, timed, otherwise. +func (t timedJudge) Active() llmgate.Judge { + a, ok := t.next.(interface{ Active() llmgate.Judge }) + if !ok { + return t + } + inner := a.Active() + if inner == nil { + return nil + } + return timedJudge{next: inner} +} + +func (t timedJudge) Judge(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { + ctx, done := From(ctx).Span(ctx, Judge) + res, err := t.next.Judge(ctx, req) + done(err) + return res, err +} + +// QueryTracer is a pgx tracer that records each query made with a request +// context as a DB wait. Install it on the pool's ConnConfig.Tracer. +type QueryTracer struct{} + +type queryStartKey struct{} + +// TraceQueryStart implements pgx.QueryTracer. +func (QueryTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryStartData) context.Context { + if From(ctx) == nil { + return ctx + } + return context.WithValue(ctx, queryStartKey{}, time.Now()) +} + +// TraceQueryEnd implements pgx.QueryTracer. +func (QueryTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) { + r := From(ctx) + if r == nil { + return + } + if start, ok := ctx.Value(queryStartKey{}).(time.Time); ok { + r.Observe(DB, start, data.Err) + } +} + +// WrapStorage records document-store reads made with a request context as +// Storage waits. A Get's wait runs until its reader is closed (or read to +// the end), because the bytes arrive while the caller reads. +func WrapStorage(s storage.Storage) storage.Storage { + if s == nil { + return nil + } + return timedStorage{Storage: s} +} + +type timedStorage struct{ storage.Storage } + +func (t timedStorage) Get(ctx context.Context, key string) (io.ReadCloser, storage.Metadata, error) { + r := From(ctx) + start := time.Now() + rc, meta, err := t.Storage.Get(ctx, key) + if r == nil { + return rc, meta, err + } + if err != nil { + r.Observe(Storage, start, err) + return rc, meta, err + } + return &timedReader{ReadCloser: rc, rec: r, start: start}, meta, nil +} + +func (t timedStorage) Exists(ctx context.Context, key string) (bool, error) { + r := From(ctx) + start := time.Now() + ok, err := t.Storage.Exists(ctx, key) + r.Observe(Storage, start, err) + return ok, err +} + +type timedReader struct { + io.ReadCloser + rec *Recorder + start time.Time + once sync.Once +} + +func (t *timedReader) Read(p []byte) (int, error) { + n, err := t.ReadCloser.Read(p) + if err == io.EOF { + t.finish(nil) + } else if err != nil { + t.finish(err) + } + return n, err +} + +func (t *timedReader) Close() error { + err := t.ReadCloser.Close() + t.finish(nil) + return err +} + +func (t *timedReader) finish(err error) { + t.once.Do(func() { t.rec.Observe(Storage, t.start, err) }) +} diff --git a/selfhost/.env.example b/selfhost/.env.example index 41d1c25..9e62363 100644 --- a/selfhost/.env.example +++ b/selfhost/.env.example @@ -11,6 +11,17 @@ VECTORLESS_ANTHROPIC_API_KEY= # VECTORLESS_ANTHROPIC_MODEL=glm-4.6 VECTORLESS_OPENAI_API_KEY= VECTORLESS_GEMINI_API_KEY= +# Gemini on Vertex AI (VECTORLESS_LLM_PROVIDER=vertex): no key. On a GCP VM the +# attached service account authenticates; elsewhere mount a credentials file +# and set VECTORLESS_GOOGLE_CREDENTIALS_PATH to its path inside the container. +VECTORLESS_VERTEX_PROJECT= +VECTORLESS_VERTEX_LOCATION=global +VECTORLESS_VERTEX_MODEL=gemini-2.5-flash-lite +# Set a thinking budget (0 = off where the model allows) OR a level, not both. +VECTORLESS_VERTEX_THINKING_BUDGET= +VECTORLESS_VERTEX_THINKING_LEVEL= +# Off GCP: put a service-account key at ./gcp/key.json and uncomment. +# VECTORLESS_GOOGLE_CREDENTIALS_PATH=/etc/vectorless/gcp/key.json # Jev (TypeSafe's System One model). Without it the engine still runs, but uses the # slower generative walk and page-level citations instead of line-level pincites. diff --git a/selfhost/README.md b/selfhost/README.md index 9ad76e0..4f90afb 100644 --- a/selfhost/README.md +++ b/selfhost/README.md @@ -44,10 +44,11 @@ Behind a control plane, each request names its tenant in `X-Vectorless-Org`. Her | Variable | What it does | |---|---| -| `VECTORLESS_LLM_PROVIDER` | `anthropic`, `openai` or `gemini`: the model that writes answers. | +| `VECTORLESS_LLM_PROVIDER` | `anthropic`, `openai`, `gemini` or `vertex`: the model that writes answers. | | `VECTORLESS_ANTHROPIC_API_KEY`, `VECTORLESS_ANTHROPIC_BASE_URL`, `VECTORLESS_ANTHROPIC_MODEL` | Anthropic, or an Anthropic-compatible gateway. For GLM on z.ai the base URL is `https://api.z.ai/api/anthropic/v1`, and the `/v1` is required. | | `VECTORLESS_OPENAI_API_KEY`, `VECTORLESS_OPENAI_BASE_URL`, `VECTORLESS_OPENAI_MODEL` | OpenAI, or any OpenAI-compatible endpoint. | | `VECTORLESS_GEMINI_API_KEY`, `VECTORLESS_GEMINI_MODEL` | Gemini. | +| `VECTORLESS_VERTEX_PROJECT`, `VECTORLESS_VERTEX_LOCATION`, `VECTORLESS_VERTEX_MODEL`, `VECTORLESS_VERTEX_THINKING_BUDGET` | Gemini on Google Cloud Vertex AI. No key: a GCP VM's service account (or a credentials file at `VECTORLESS_GOOGLE_CREDENTIALS_PATH`) authenticates, and the project is billed. Set a thinking budget (`0` turns thinking off on models that allow it) or a level, not both. Off GCP, put a service-account key in `selfhost/gcp/key.json` (gitignored) and set `VECTORLESS_GOOGLE_CREDENTIALS_PATH=/etc/vectorless/gcp/key.json`. | | `VECTORLESS_ANSWER_MODEL` | Optional: a different model for answers only. | | `VECTORLESS_TYPESAFE_API_KEY` | Jev, TypeSafe's System One model, which makes the contents, page selection and pincite decisions. It's a hosted API, so document text goes to TypeSafe when it's set. | diff --git a/selfhost/docker-compose.yml b/selfhost/docker-compose.yml index bee0407..781236c 100644 --- a/selfhost/docker-compose.yml +++ b/selfhost/docker-compose.yml @@ -45,6 +45,15 @@ services: VLS_LLM_OPENAI_MODEL: ${VECTORLESS_OPENAI_MODEL:-} VLS_GEMINI_API_KEY: ${VECTORLESS_GEMINI_API_KEY:-} VLS_LLM_GEMINI_MODEL: ${VECTORLESS_GEMINI_MODEL:-} + # Gemini on Google Cloud Vertex AI (VECTORLESS_LLM_PROVIDER=vertex): no + # key; Application Default Credentials (a VM's service account) bill + # the project. Thinking budget 0 turns thinking off where allowed. + VLS_LLM_VERTEX_PROJECT: ${VECTORLESS_VERTEX_PROJECT:-} + VLS_LLM_VERTEX_LOCATION: ${VECTORLESS_VERTEX_LOCATION:-global} + VLS_LLM_VERTEX_MODEL: ${VECTORLESS_VERTEX_MODEL:-} + VLS_LLM_VERTEX_THINKING_BUDGET: ${VECTORLESS_VERTEX_THINKING_BUDGET:-} + VLS_LLM_VERTEX_THINKING_LEVEL: ${VECTORLESS_VERTEX_THINKING_LEVEL:-} + GOOGLE_APPLICATION_CREDENTIALS: ${VECTORLESS_GOOGLE_CREDENTIALS_PATH:-} VLS_LLM_ANSWER_MODEL: ${VECTORLESS_ANSWER_MODEL:-} # Jev, TypeSafe's System One model: contents, page selection, pincites. VLS_TYPESAFE_API_KEY: ${VECTORLESS_TYPESAFE_API_KEY:-} @@ -57,6 +66,9 @@ services: volumes: - documents:/data - ./engine.yaml:/etc/vectorless/engine.yaml:ro + # Vertex credentials off GCP (see .env.example). On a GCP VM the + # attached service account is used and this directory can stay empty. + - ./gcp:/etc/vectorless/gcp:ro depends_on: postgres: condition: service_healthy diff --git a/selfhost/gcp/.gitignore b/selfhost/gcp/.gitignore new file mode 100644 index 0000000..6d44674 --- /dev/null +++ b/selfhost/gcp/.gitignore @@ -0,0 +1,3 @@ +# Service-account keys go here and are never committed. +* +!.gitignore