From 97eafa67e7691f334e8ab883e51fdc2628bd5fe9 Mon Sep 17 00:00:00 2001 From: BeaconCat Date: Sun, 12 Apr 2026 15:57:14 +0800 Subject: [PATCH] membench: add LLM-as-Judge evaluation mode Add --eval-mode=llm to membench for LLM-based answer generation and semantic scoring via an OpenAI-compatible API endpoint. New files: - llm_client.go: generic OpenAI-compatible chat completion client with support for API key, configurable timeout, and optional chat_template_kwargs (for llama.cpp thinking models) - eval_llm.go: LLM answer generation + LLM-as-Judge scoring for both legacy and seahorse retrieval modes Changes to main.go: - --eval-mode flag (token|llm) to select evaluation strategy - --api-base, --api-key, --model flags with env var fallback (MEMBENCH_API_BASE, MEMBENCH_API_KEY, MEMBENCH_MODEL) - --no-thinking flag for llama.cpp + Qwen thinking models - --limit flag to cap QA questions per sample for quick testing --- cmd/membench/eval_llm.go | 271 +++++++++++++++++++++++++++++++++++++ cmd/membench/llm_client.go | 131 ++++++++++++++++++ cmd/membench/main.go | 103 ++++++++++++-- 3 files changed, 495 insertions(+), 10 deletions(-) create mode 100644 cmd/membench/eval_llm.go create mode 100644 cmd/membench/llm_client.go diff --git a/cmd/membench/eval_llm.go b/cmd/membench/eval_llm.go new file mode 100644 index 000000000..430548003 --- /dev/null +++ b/cmd/membench/eval_llm.go @@ -0,0 +1,271 @@ +package main + +import ( + "context" + "fmt" + "log" + "strconv" + "strings" + + "github.com/sipeed/picoclaw/pkg/seahorse" +) + +const answerSystemPrompt = `You are a helpful assistant. Given conversation context, answer the question concisely and accurately. If the answer is not in the context, say "I don't know". Answer in 1-3 sentences maximum.` + +const judgeSystemPrompt = `You are an impartial judge evaluating answer quality. +Compare the candidate answer against the reference answer. +Consider semantic equivalence — different wording expressing the same meaning should score high. + +Output ONLY a single integer score from 1 to 5: +1 = completely wrong or irrelevant +2 = partially related but mostly incorrect +3 = partially correct, missing key details +4 = mostly correct with minor omissions +5 = fully correct, semantically equivalent + +Output ONLY the number, nothing else.` + +// generateAnswer asks the LLM to answer a question given retrieved context. +func generateAnswer(ctx context.Context, client *LLMClient, contextText, question string) (string, error) { + // Truncate context to avoid exceeding model limits + if len(contextText) > 6000 { + contextText = contextText[:6000] + "\n... [truncated]" + } + + userPrompt := fmt.Sprintf("## Conversation Context\n\n%s\n\n## Question\n\n%s", contextText, question) + return client.Complete(ctx, answerSystemPrompt, userPrompt) +} + +// judgeAnswer asks the LLM to score the candidate answer vs the gold answer. +// Returns a score from 0.0 to 1.0. +func judgeAnswer(ctx context.Context, client *LLMClient, question, goldAnswer, candidateAnswer string) (float64, error) { + userPrompt := fmt.Sprintf( + "Question: %s\n\nReference Answer: %s\n\nCandidate Answer: %s\n\nScore:", + question, goldAnswer, candidateAnswer, + ) + + response, err := client.Complete(ctx, judgeSystemPrompt, userPrompt) + if err != nil { + return 0.0, err + } + + // Parse score from response + response = strings.TrimSpace(response) + // Extract first digit found + for _, ch := range response { + if ch >= '1' && ch <= '5' { + score, _ := strconv.Atoi(string(ch)) + return float64(score-1) / 4.0, nil // Normalize 1-5 to 0.0-1.0 + } + } + log.Printf("WARNING: could not parse judge score from: %q, defaulting to 0.0", response) + return 0.0, nil +} + +// EvalLegacyLLM evaluates legacy store using LLM generation + LLM-as-Judge. +func EvalLegacyLLM( + ctx context.Context, + samples []LocomoSample, + legacy *LegacyStore, + budgetTokens int, + client *LLMClient, +) []EvalResult { + results := make([]EvalResult, 0, len(samples)) + total := 0 + for si := range samples { + sample := &samples[si] + history := legacy.GetHistory(sample.SampleID) + + allContent := make([]string, 0, len(history)) + for _, msg := range history { + allContent = append(allContent, msg.Content) + } + + qaResults := make([]QAResult, 0, len(sample.QA)) + for qi := range sample.QA { + qa := &sample.QA[qi] + total++ + truncated, _ := BudgetTruncate(allContent, budgetTokens) + contextText := StringListToContent(truncated) + + // Generate answer with LLM + llmAnswer, err := generateAnswer(ctx, client, contextText, qa.Question) + if err != nil { + log.Printf("WARN: LLM generation failed for sample %s Q%d: %v", sample.SampleID, qi, err) + llmAnswer = "" + } + + // Judge the answer + score := 0.0 + if llmAnswer != "" { + score, err = judgeAnswer(ctx, client, qa.Question, qa.AnswerString(), llmAnswer) + if err != nil { + log.Printf("WARN: LLM judge failed for sample %s Q%d: %v", sample.SampleID, qi, err) + } + } + + hitRate := RecallHitRate(qa.Evidence, sample, contextText) + + qaResults = append(qaResults, QAResult{ + Question: qa.Question, + Category: qa.Category, + GoldAnswer: qa.AnswerString(), + TokenF1: score, + HitRate: hitRate, + }) + + log.Printf("[legacy-llm] sample=%s q=%d/%d score=%.2f answer=%q", + sample.SampleID, total, countTotalQA(samples), score, truncateStr(llmAnswer, 80)) + } + + results = append(results, EvalResult{ + Mode: "legacy-llm", + SampleID: sample.SampleID, + QAResults: qaResults, + Agg: aggregateMetrics(qaResults), + }) + } + return results +} + +// EvalSeahorseLLM evaluates seahorse retrieval using LLM generation + LLM-as-Judge. +func EvalSeahorseLLM( + ctx context.Context, + samples []LocomoSample, + ir *SeahorseIngestResult, + budgetTokens int, + client *LLMClient, +) []EvalResult { + store := ir.Engine.GetRetrieval().Store() + retrieval := ir.Engine.GetRetrieval() + + results := make([]EvalResult, 0, len(samples)) + total := 0 + for si := range samples { + sample := &samples[si] + convID, ok := ir.ConvMap[sample.SampleID] + if !ok { + log.Printf("WARN: no conversation ID for sample %s", sample.SampleID) + continue + } + + qaResults := make([]QAResult, 0, len(sample.QA)) + for qi := range sample.QA { + qa := &sample.QA[qi] + total++ + keywords := ExtractKeywords(qa.Question) + + // Search and rank + bestRank := map[int64]float64{} + for _, kw := range keywords { + searchResults, err := store.SearchMessages(ctx, seahorse.SearchInput{ + Pattern: kw, + ConversationID: convID, + Limit: 20, + }) + if err != nil { + continue + } + for _, sr := range searchResults { + if sr.MessageID > 0 { + if prev, ok := bestRank[sr.MessageID]; !ok || sr.Rank < prev { + bestRank[sr.MessageID] = sr.Rank + } + } + } + } + + messageIDs := make([]int64, 0, len(bestRank)) + for id := range bestRank { + messageIDs = append(messageIDs, id) + } + sortByRank(messageIDs, bestRank) + + var contentParts []string + if len(messageIDs) > 0 { + expandResult, err := retrieval.ExpandMessages(ctx, messageIDs) + if err == nil { + for _, msg := range expandResult.Messages { + contentParts = append(contentParts, msg.Content) + } + } + } + + contextText := "" + if len(contentParts) > 0 { + truncated, _ := BudgetTruncate(contentParts, budgetTokens) + contextText = StringListToContent(truncated) + } + + // Generate answer with LLM + llmAnswer := "" + score := 0.0 + if contextText != "" { + var err error + llmAnswer, err = generateAnswer(ctx, client, contextText, qa.Question) + if err != nil { + log.Printf("WARN: LLM generation failed for sample %s Q%d: %v", sample.SampleID, qi, err) + } + } + + // Judge the answer + if llmAnswer != "" { + var err error + score, err = judgeAnswer(ctx, client, qa.Question, qa.AnswerString(), llmAnswer) + if err != nil { + log.Printf("WARN: LLM judge failed for sample %s Q%d: %v", sample.SampleID, qi, err) + } + } + + hitRate := RecallHitRate(qa.Evidence, sample, contextText) + + qaResults = append(qaResults, QAResult{ + Question: qa.Question, + Category: qa.Category, + GoldAnswer: qa.AnswerString(), + TokenF1: score, + HitRate: hitRate, + }) + + log.Printf("[seahorse-llm] sample=%s q=%d/%d score=%.2f answer=%q", + sample.SampleID, total, countTotalQA(samples), score, truncateStr(llmAnswer, 80)) + } + + results = append(results, EvalResult{ + Mode: "seahorse-llm", + SampleID: sample.SampleID, + QAResults: qaResults, + Agg: aggregateMetrics(qaResults), + }) + } + return results +} + +func countTotalQA(samples []LocomoSample) int { + n := 0 + for i := range samples { + n += len(samples[i].QA) + } + return n +} + +func truncateStr(s string, maxLen int) string { + s = strings.ReplaceAll(s, "\n", " ") + if len(s) > maxLen { + return s[:maxLen] + "..." + } + return s +} + +// sortByRank sorts message IDs by BM25 rank (more negative = better). +func sortByRank(ids []int64, ranks map[int64]float64) { + for i := 1; i < len(ids); i++ { + key := ids[i] + j := i - 1 + for j >= 0 && ranks[ids[j]] > ranks[key] { + ids[j+1] = ids[j] + j-- + } + ids[j+1] = key + } +} diff --git a/cmd/membench/llm_client.go b/cmd/membench/llm_client.go new file mode 100644 index 000000000..57734ae5b --- /dev/null +++ b/cmd/membench/llm_client.go @@ -0,0 +1,131 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +// LLMClient wraps an OpenAI-compatible chat completion endpoint. +type LLMClient struct { + BaseURL string + Model string + APIKey string + NoThinking bool // send chat_template_kwargs to disable thinking (llama.cpp specific) + Client *http.Client +} + +// LLMClientOptions configures the LLM client. +type LLMClientOptions struct { + BaseURL string + Model string + APIKey string + Timeout time.Duration + NoThinking bool +} + +// NewLLMClient creates a client for an OpenAI-compatible chat completion API. +func NewLLMClient(opts LLMClientOptions) *LLMClient { + if opts.Timeout == 0 { + opts.Timeout = 120 * time.Second + } + return &LLMClient{ + BaseURL: strings.TrimRight(opts.BaseURL, "/"), + Model: opts.Model, + APIKey: opts.APIKey, + NoThinking: opts.NoThinking, + Client: &http.Client{ + Timeout: opts.Timeout, + }, + } +} + +type chatRequest struct { + Model string `json:"model"` + Messages []chatMessage `json:"messages"` + Temperature float64 `json:"temperature"` + MaxTokens int `json:"max_tokens"` + ChatTemplateKwargs map[string]interface{} `json:"chat_template_kwargs,omitempty"` +} + +type chatMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type chatResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` +} + +// Complete sends a chat completion request and returns the assistant's reply. +func (c *LLMClient) Complete(ctx context.Context, systemPrompt, userPrompt string) (string, error) { + messages := []chatMessage{} + if systemPrompt != "" { + messages = append(messages, chatMessage{Role: "system", Content: systemPrompt}) + } + messages = append(messages, chatMessage{Role: "user", Content: userPrompt}) + + body := chatRequest{ + Model: c.Model, + Messages: messages, + Temperature: 0.1, + MaxTokens: 512, + } + if c.NoThinking { + body.ChatTemplateKwargs = map[string]interface{}{ + "enable_thinking": false, + } + } + + jsonBody, err := json.Marshal(body) + if err != nil { + return "", fmt.Errorf("marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/v1/chat/completions", bytes.NewReader(jsonBody)) + if err != nil { + return "", fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + if c.APIKey != "" { + req.Header.Set("Authorization", "Bearer "+c.APIKey) + } + + resp, err := c.Client.Do(req) + if err != nil { + return "", fmt.Errorf("http request: %w", err) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return "", fmt.Errorf("read response: %w", err) + } + + if resp.StatusCode != 200 { + return "", fmt.Errorf("API error %d: %s", resp.StatusCode, string(respBody)) + } + + var chatResp chatResponse + if err := json.Unmarshal(respBody, &chatResp); err != nil { + return "", fmt.Errorf("parse response: %w", err) + } + if len(chatResp.Choices) == 0 { + return "", fmt.Errorf("no choices in response") + } + content := strings.TrimSpace(chatResp.Choices[0].Message.Content) + // Strip any residual ... blocks + if idx := strings.Index(content, ""); idx >= 0 { + content = strings.TrimSpace(content[idx+len(""):]) + } + return content, nil +} diff --git a/cmd/membench/main.go b/cmd/membench/main.go index 0c5a9387a..ff1970d52 100644 --- a/cmd/membench/main.go +++ b/cmd/membench/main.go @@ -15,10 +15,16 @@ import ( ) var ( - flagData string - flagOut string - flagMode string - flagBudget int + flagData string + flagOut string + flagMode string + flagBudget int + flagEvalMode string + flagAPIBase string + flagAPIKey string + flagModel string + flagNoThinking bool + flagLimit int ) func main() { @@ -48,6 +54,12 @@ func main() { evalCmd.Flags().StringVar(&flagOut, "out", "./bench-out", "output working directory") evalCmd.Flags().StringVar(&flagMode, "mode", "all", "modes to evaluate: legacy, seahorse, or all") evalCmd.Flags().IntVar(&flagBudget, "budget", 4000, "token budget for retrieval") + evalCmd.Flags().StringVar(&flagEvalMode, "eval-mode", "token", "evaluation mode: token (direct match) or llm (LLM-as-Judge)") + evalCmd.Flags().StringVar(&flagAPIBase, "api-base", "", "OpenAI-compatible API base URL (env: MEMBENCH_API_BASE)") + evalCmd.Flags().StringVar(&flagAPIKey, "api-key", "", "API key for the LLM endpoint (env: MEMBENCH_API_KEY)") + evalCmd.Flags().StringVar(&flagModel, "model", "", "model name for LLM eval (env: MEMBENCH_MODEL)") + evalCmd.Flags().BoolVar(&flagNoThinking, "no-thinking", false, "disable thinking mode via chat_template_kwargs (llama.cpp + Qwen)") + evalCmd.Flags().IntVar(&flagLimit, "limit", 0, "max QA questions per sample (0 = all)") reportCmd := &cobra.Command{ Use: "report", @@ -65,6 +77,12 @@ func main() { runCmd.Flags().StringVar(&flagOut, "out", "./bench-out", "output working directory") runCmd.Flags().StringVar(&flagMode, "mode", "all", "modes to run: legacy, seahorse, or all") runCmd.Flags().IntVar(&flagBudget, "budget", 4000, "token budget for retrieval") + runCmd.Flags().StringVar(&flagEvalMode, "eval-mode", "token", "evaluation mode: token (direct match) or llm (LLM-as-Judge)") + runCmd.Flags().StringVar(&flagAPIBase, "api-base", "", "OpenAI-compatible API base URL (env: MEMBENCH_API_BASE)") + runCmd.Flags().StringVar(&flagAPIKey, "api-key", "", "API key for the LLM endpoint (env: MEMBENCH_API_KEY)") + runCmd.Flags().StringVar(&flagModel, "model", "", "model name for LLM eval (env: MEMBENCH_MODEL)") + runCmd.Flags().BoolVar(&flagNoThinking, "no-thinking", false, "disable thinking mode via chat_template_kwargs (llama.cpp + Qwen)") + runCmd.Flags().IntVar(&flagLimit, "limit", 0, "max QA questions per sample (0 = all)") rootCmd.AddCommand(ingestCmd, evalCmd, reportCmd, runCmd) @@ -136,6 +154,26 @@ func runEval(cmd *cobra.Command, args []string) error { } log.Printf("Loaded %d samples", len(samples)) + if flagLimit > 0 { + for i := range samples { + if len(samples[i].QA) > flagLimit { + samples[i].QA = samples[i].QA[:flagLimit] + } + } + log.Printf("Limited to %d QA per sample", flagLimit) + } + + useLLM := strings.ToLower(flagEvalMode) == "llm" + var llmClient *LLMClient + if useLLM { + opts, err := buildLLMOptions() + if err != nil { + return err + } + llmClient = NewLLMClient(opts) + log.Printf("LLM eval mode: model=%s base=%s no-thinking=%v", opts.Model, opts.BaseURL, opts.NoThinking) + } + var allResults []EvalResult for _, mode := range modes { @@ -145,18 +183,30 @@ func runEval(cmd *cobra.Command, args []string) error { for i := range samples { legacy.IngestSample(&samples[i]) } - results := EvalLegacy(ctx, samples, legacy, flagBudget) - allResults = append(allResults, results...) - log.Printf("legacy: evaluated %d samples", len(results)) + if useLLM { + results := EvalLegacyLLM(ctx, samples, legacy, flagBudget, llmClient) + allResults = append(allResults, results...) + log.Printf("legacy-llm: evaluated %d samples", len(results)) + } else { + results := EvalLegacy(ctx, samples, legacy, flagBudget) + allResults = append(allResults, results...) + log.Printf("legacy: evaluated %d samples", len(results)) + } case "seahorse": dbPath := filepath.Join(flagOut, "seahorse.db") ir, err := IngestSeahorse(ctx, samples, dbPath) if err != nil { return fmt.Errorf("ingest seahorse: %w", err) } - results := EvalSeahorse(ctx, samples, ir, flagBudget) - allResults = append(allResults, results...) - log.Printf("seahorse: evaluated %d samples", len(results)) + if useLLM { + results := EvalSeahorseLLM(ctx, samples, ir, flagBudget, llmClient) + allResults = append(allResults, results...) + log.Printf("seahorse-llm: evaluated %d samples", len(results)) + } else { + results := EvalSeahorse(ctx, samples, ir, flagBudget) + allResults = append(allResults, results...) + log.Printf("seahorse: evaluated %d samples", len(results)) + } } } @@ -206,3 +256,36 @@ func runReport(cmd *cobra.Command, args []string) error { func runAll(cmd *cobra.Command, args []string) error { return runEval(cmd, args) } + +// envOrFlag returns the flag value if non-empty, otherwise falls back to the +// environment variable. +func envOrFlag(flag, envKey string) string { + if flag != "" { + return flag + } + return os.Getenv(envKey) +} + +// buildLLMOptions resolves LLM client configuration from flags and environment +// variables. Flag values take precedence over environment variables. +// +// Environment variables: +// +// MEMBENCH_API_BASE – OpenAI-compatible base URL (default http://127.0.0.1:8080) +// MEMBENCH_API_KEY – Bearer token for the endpoint +// MEMBENCH_MODEL – Model name to send in the request +func buildLLMOptions() (LLMClientOptions, error) { + base := envOrFlag(flagAPIBase, "MEMBENCH_API_BASE") + if base == "" { + base = "http://127.0.0.1:8080" + } + model := envOrFlag(flagModel, "MEMBENCH_MODEL") + apiKey := envOrFlag(flagAPIKey, "MEMBENCH_API_KEY") + + return LLMClientOptions{ + BaseURL: base, + Model: model, + APIKey: apiKey, + NoThinking: flagNoThinking, + }, nil +}