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
2026-04-12 15:57:14 +08:00 · 2026-04-12 15:57:14 +08:00 · 97eafa67e7
commit 97eafa67e7
parent 748ac58dd1
3 changed files with 495 additions and 10 deletions
--- a/cmd/membench/eval_llm.go
+++ 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
+	}
+}
--- a/cmd/membench/llm_client.go
+++ 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 <think>...</think> blocks
+	if idx := strings.Index(content, "</think>"); idx >= 0 {
+		content = strings.TrimSpace(content[idx+len("</think>"):])
+	}
+	return content, nil
+}
--- 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
+}