diff --git a/internal/tests/ai.go b/internal/tests/ai.go
new file mode 100644
index 0000000..2f200ce
--- /dev/null
+++ b/internal/tests/ai.go
@@ -0,0 +1,14 @@
+package tests
+
+import "github.com/lukaszraczylo/llm-testbench/internal/testkit"
+
+// registerAITests registers every ai-category test. Each subcategory -
+// vector-search, llm-integration, and rag - has its own register function
+// and source file (ai_vectorsearch.go, ai_llmintegration.go, ai_rag.go) to
+// keep any one file from growing past a few hundred lines, mirroring
+// agents.go's split across its three subcategories.
+func registerAITests(r *testkit.Registry) {
+ registerAIVectorSearchTests(r)
+ registerAILLMIntegrationTests(r)
+ registerAIRAGTests(r)
+}
diff --git a/internal/tests/ai_llmintegration.go b/internal/tests/ai_llmintegration.go
new file mode 100644
index 0000000..264cd44
--- /dev/null
+++ b/internal/tests/ai_llmintegration.go
@@ -0,0 +1,369 @@
+package tests
+
+import (
+ "github.com/lukaszraczylo/llm-testbench/internal/eval"
+ "github.com/lukaszraczylo/llm-testbench/internal/testkit"
+)
+
+func registerAILLMIntegrationTests(r *testkit.Registry) {
+ r.Register(llmFieldSemanticsTest())
+ r.Register(llmTokenBudgetReasoningTest())
+ r.Register(llm429RetryBackoffTest())
+ r.Register(llmContextOverflowStrategyTest())
+ r.Register(llmToolCallHandlingTest())
+ r.Register(llmRolePlacementTest())
+ r.Register(llmTemperatureZeroCaveatTest())
+ r.Register(llmSSEStreamDoneTest())
+ r.Register(llmEmbeddingBatchMathTest())
+ r.Register(llmStopSequenceTraceTest())
+}
+
+// aiCeilDiv returns ceil(total/batchSize) for positive total and
+// batchSize: the number of fixed-size batches needed to cover total items,
+// with a final partial batch counted as one whole batch.
+func aiCeilDiv(total, batchSize int) int {
+ return (total + batchSize - 1) / batchSize
+}
+
+// llmFieldSemanticsTest: match three OpenAI-compatible request fields
+// (temperature, max_tokens, stop) to their actual behavior, from an
+// enumerated multiple-choice roster given inline so the answer is
+// determined by the definitions in the prompt, not by memorized field
+// names.
+//
+// ground truth: temperature scales the randomness of token sampling from
+// the model's output distribution; max_tokens is a hard cap on how many
+// completion tokens the response may contain; stop is a list of strings
+// that, once generated, immediately end generation (the stop string itself
+// is excluded from the returned text).
+func llmFieldSemanticsTest() testkit.Test {
+ prompt := `An OpenAI-compatible chat completions request can set these
+fields: temperature, max_tokens, stop. For each field, pick the one option
+id below that correctly describes what it controls.
+
+temperature options:
+- sampling_randomness: how much randomness is applied when sampling the
+ next token from the model's output probability distribution.
+- reply_length_cap: a hard cap on how many tokens the reply may contain.
+- output_format: whether the reply is returned as JSON or plain text.
+
+max_tokens options:
+- completion_token_cap: a hard cap on how many completion tokens the reply
+ may contain; generation stops once this many tokens have been produced.
+- minimum_reply_length: the minimum number of tokens the reply must
+ contain before it can end.
+- request_timeout: how many seconds the client waits before giving up on
+ the request.
+
+stop options:
+- generation_end_strings: a list of strings that, once generated,
+ immediately end generation; the stop string itself is excluded from the
+ returned text.
+- required_reply_prefix: a list of strings the reply must begin with.
+- system_override: a list of strings that replace the system prompt.
+
+Respond with only a JSON object mapping each field to its correct option
+id: {"temperature":"...","max_tokens":"...","stop":"..."}`
+
+ evaluator := eval.Mean(
+ eval.JSONField("temperature", "sampling_randomness"),
+ eval.JSONField("max_tokens", "completion_token_cap"),
+ eval.JSONField("stop", "generation_end_strings"),
+ )
+
+ return testkit.Test{
+ ID: "llm-field-semantics",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Match temperature/max_tokens/stop to their correct behavior from an inline multiple-choice roster.",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// llmTokenBudgetReasoningWant is derived by plain subtraction, matching the
+// arithmetic ai_llmintegration_test.go independently recomputes.
+//
+// ground truth: this framework's own PLAN.md documents a live-run
+// incident where a reasoning model's hidden ... content
+// consumed most of a request's max_tokens budget before any visible answer
+// was produced, truncating the visible answer to empty - which is why
+// MaxTokens defaults to 0 (provider default) for new catalog tests rather
+// than a small fixed cap. Here: max_tokens = 500, and the model's
+// block consumes 437 completion tokens before the visible answer begins.
+// Tokens remaining for the visible answer = 500 - 437 = 63.
+const llmTokenBudgetReasoningWant = 500 - 437
+
+func llmTokenBudgetReasoningTest() testkit.Test {
+ prompt := `An OpenAI-compatible chat completion request sets
+max_tokens = 500. This caps the total number of completion tokens the
+response may contain, counting every token the model generates: a
+reasoning model's hidden ... content counts against this
+same budget before its visible answer, since both are completion tokens.
+
+For this request, the model's ... block consumes exactly
+437 completion tokens before it begins writing its visible answer.
+
+How many completion tokens remain in the max_tokens budget for the visible
+answer, once the block has been generated? Respond with only the
+number.`
+
+ return testkit.Test{
+ ID: "llm-token-budget-reasoning",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Compute completion tokens remaining for a visible answer after a reasoning model's block consumes part of max_tokens.",
+ Prompt: prompt,
+ Eval: eval.Numeric(eval.ExtractLastNumber[int], llmTokenBudgetReasoningWant, 0),
+ }
+}
+
+// llm429RetryBackoffTest: decide whether to retry an HTTP 429 response and
+// for how long to wait, given an explicit Retry-After header value.
+//
+// ground truth: HTTP 429 (Too Many Requests) signals a transient,
+// rate-limiting condition, not a permanent failure of the request itself -
+// the correct client behavior is to retry rather than give up immediately.
+// When the server supplies a Retry-After header, that value is the
+// authoritative wait time the server itself is asking for; the client
+// should honor it exactly rather than substitute its own default backoff
+// duration.
+func llm429RetryBackoffTest() testkit.Test {
+ prompt := `An OpenAI-compatible chat completions client sends a request
+and receives this response:
+
+HTTP/1.1 429 Too Many Requests
+Retry-After: 2
+
+This is the first attempt at a request that is otherwise valid and
+retryable. Should the client retry this request, and if so, how many
+seconds should it wait before retrying (using the server's own guidance
+rather than an arbitrary default)? Respond with only a JSON object:
+{"should_retry":"yes"|"no","wait_seconds":}`
+
+ evaluator := eval.Mean(
+ eval.JSONField("should_retry", "yes"),
+ eval.JSONField("wait_seconds", 2),
+ )
+
+ return testkit.Test{
+ ID: "llm-429-retry-backoff",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Decide to retry an HTTP 429 and honor its exact Retry-After: 2 header value rather than an arbitrary default.",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// llmContextOverflowStrategyTest: pick the only context-overflow strategy
+// that satisfies a stated "no altering any turn's text" constraint.
+//
+// ground truth: total required content is 500 (system) + 300 (new
+// message) + 1800 (2 most recent turn pairs, 900 tokens each, which must
+// stay verbatim) + 7200 (8 older turn pairs, 900 tokens each) = 10100
+// tokens, against an 8000-token window - an overflow of 2100 tokens.
+// Protected content (system, new message, 2 most recent pairs) totals
+// 2600 tokens, leaving 5400 tokens for the 7200 tokens of older turns: an
+// overflow of 1800 tokens among only the older turns. The stated
+// constraint that no text anywhere may be paraphrased or summarized rules
+// out summarizing older turns (summarizing alters their text) and the
+// system prompt is explicitly protected, so the only strategy that fits
+// the window without altering any retained turn's text is dropping whole
+// older turns until the remainder fits.
+func llmContextOverflowStrategyTest() testkit.Test {
+ prompt := `A chat client's context window holds 8000 tokens total.
+Current conversation content:
+- system prompt: 500 tokens, must remain byte-for-byte unaltered.
+- 10 prior user/assistant turn pairs, 900 tokens each (9000 tokens total).
+ The 2 most recent of these 10 pairs must remain completely unaltered,
+ verbatim (1800 of the 9000 tokens).
+- new user message: 300 tokens, must be included as-is.
+
+Hard constraint: no text anywhere in the assembled context may be
+paraphrased, summarized, or otherwise altered - every turn that is
+included must appear exactly as it was originally written.
+
+Total required content exceeds the 8000-token window. Given the hard
+constraint above, which strategy fits: drop-oldest-turns (remove whole
+older turn pairs entirely until it fits), summarize-oldest-turns
+(compress older turns into shorter summaries), or truncate-system-prompt
+(shorten the system prompt)? Respond with only a JSON object:
+{"strategy":"drop-oldest-turns"|"summarize-oldest-turns"|"truncate-system-prompt"}`
+
+ return testkit.Test{
+ ID: "llm-context-overflow-strategy",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Pick drop-oldest-turns as the only context-overflow strategy that satisfies a no-altered-text constraint.",
+ Prompt: prompt,
+ Eval: eval.JSONField("strategy", "drop-oldest-turns"),
+ }
+}
+
+// llmToolCallHandlingTest: trace OpenAI-compatible tool-call handling -
+// the message role and id field used to return a tool's result.
+//
+// ground truth: when an assistant message includes a tool_calls array,
+// each entry carries its own id. The result of executing that call is
+// returned as a new message with role "tool", whose tool_call_id field
+// must equal the originating tool_calls entry's id so the model can match
+// the result back to the specific call it made.
+func llmToolCallHandlingTest() testkit.Test {
+ prompt := `In the OpenAI-compatible chat completions tool-calling
+protocol, the assistant's message can include a tool_calls array; each
+entry has its own id and a function name/arguments to execute.
+
+Once your code has executed the requested function and has its result,
+which message role must the next message use to feed that result back
+into the conversation, and which field on that new message must be set to
+the originating tool_calls entry's id so the model can match the result to
+the call it made? Respond with only a JSON object:
+{"next_role":"...","id_field":"..."}`
+
+ evaluator := eval.Mean(
+ eval.JSONField("next_role", "tool"),
+ eval.JSONField("id_field", "tool_call_id"),
+ )
+
+ return testkit.Test{
+ ID: "llm-tool-call-handling",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Trace OpenAI-compatible tool-call result handling: role=tool, id field=tool_call_id.",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// llmRolePlacementTest: identify the system role as the correct carrier
+// for a fixed, non-negotiable operating instruction.
+func llmRolePlacementTest() testkit.Test {
+ prompt := `In an OpenAI-compatible chat completions request, which
+message role should carry a fixed operating instruction that must apply to
+every request regardless of what the user's own message says (for
+example, "never reveal API keys in your reply"): system or user? Respond
+with only one word.`
+
+ return testkit.Test{
+ ID: "llm-role-placement",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Identify system (not user) as the role that carries a fixed, non-negotiable operating instruction.",
+ Prompt: prompt,
+ Eval: eval.Equals("system"),
+ }
+}
+
+// llmTemperatureZeroCaveatTest: explain why temperature=0 is not a full
+// bit-for-bit determinism guarantee.
+//
+// ground truth: temperature=0 makes token selection greedy (always the
+// highest-probability token) rather than random, but it does not by
+// itself guarantee bit-identical output across repeated calls. Floating-
+// point arithmetic is not strictly associative, and provider-side
+// execution details - batching requests together, and which GPU/hardware
+// kernel path a given batch is routed through - can change the order
+// operations are summed in, producing tiny floating-point differences
+// that occasionally flip which token has the (near-)highest probability.
+func llmTemperatureZeroCaveatTest() testkit.Test {
+ prompt := `A test suite calls an OpenAI-compatible chat completions API
+with temperature=0 on every request, expecting bit-for-bit identical
+output across repeated calls with the identical prompt. In practice,
+repeated calls can still occasionally return slightly different text.
+
+Explain why temperature=0 does not fully guarantee deterministic,
+bit-for-bit identical output: name the numerical-computation reason and
+the provider-side execution reason that both contribute.`
+
+ evaluator := eval.Mean(
+ eval.ContainsAny("floating point", "floating-point", "floating-point precision", "numerical precision"),
+ eval.ContainsAny("batching", "batch", "hardware", "gpu", "kernel", "parallel execution"),
+ )
+
+ return testkit.Test{
+ ID: "llm-temperature-zero-caveat",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Explain why temperature=0 does not fully guarantee bit-identical output (floating-point non-associativity plus provider-side batching/hardware).",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// llmSSEStreamDoneTest: recall the exact text that closes an
+// OpenAI-compatible streaming response.
+func llmSSEStreamDoneTest() testkit.Test {
+ prompt := `An OpenAI-compatible chat completions request streams its
+response as Server-Sent Events: a sequence of lines of the form
+"data: ", each followed by a blank line, one JSON chunk per
+generated token or small group of tokens.
+
+The very last event before the stream closes is also a "data: " line, but
+its payload is not a JSON chunk. What exact text follows "data: " in that
+final event? Respond with only that exact text, nothing else.`
+
+ return testkit.Test{
+ ID: "llm-sse-stream-done",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Recall the exact [DONE] sentinel text of the final SSE event in an OpenAI-compatible stream.",
+ Prompt: prompt,
+ Eval: eval.Equals("[DONE]"),
+ }
+}
+
+// llmEmbeddingBatchMathWant is derived by calling aiCeilDiv, not
+// hardcoded.
+//
+// ground truth: ceil(10000/96) - 96*104 = 9984, leaving 16 documents that
+// still need one more (partial) request, so 104 full batches plus 1
+// partial batch = 105 requests total.
+var llmEmbeddingBatchMathWant = aiCeilDiv(10000, 96)
+
+func llmEmbeddingBatchMathTest() testkit.Test {
+ prompt := `You must generate embeddings for 10,000 documents. The
+embeddings API accepts a maximum of 96 items per request, and every
+request you send should be as full as possible (maximally batched) to
+minimize the number of requests.
+
+How many API requests are needed in total to embed all 10,000 documents?
+Respond with only the number.`
+
+ return testkit.Test{
+ ID: "llm-embedding-batch-math",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Compute the number of maximally-batched embedding requests needed for 10,000 documents at 96 items per request.",
+ Prompt: prompt,
+ Eval: eval.Numeric(eval.ExtractLastNumber[int], llmEmbeddingBatchMathWant, 0),
+ }
+}
+
+// llmStopSequenceTraceTest: trace what text an OpenAI-compatible completion
+// actually returns when a stop sequence is set - the reply is cut before
+// the stop sequence, and the stop sequence itself is excluded.
+func llmStopSequenceTraceTest() testkit.Test {
+ prompt := `A chat completions request sets stop = ["\n\n"] (generation
+ends the instant this exact two-newline sequence would be produced, and
+the stop sequence itself is excluded from the returned text).
+
+If the model were generating without any stop sequence, its full raw
+output would be exactly:
+
+Answer: 42
+
+Explanation: 6 times 7 is 42.
+
+Given stop = ["\n\n"], what text does the API actually return to the
+client? Respond with only that exact returned text, nothing else.`
+
+ return testkit.Test{
+ ID: "llm-stop-sequence-trace",
+ Category: "ai",
+ Subcategory: "llm-integration",
+ Description: "Trace the exact text an API returns when generation is cut by a stop=[\"\\n\\n\"] sequence, excluding the sequence itself.",
+ Prompt: prompt,
+ Eval: eval.Equals("Answer: 42"),
+ }
+}
diff --git a/internal/tests/ai_llmintegration_test.go b/internal/tests/ai_llmintegration_test.go
new file mode 100644
index 0000000..02a574f
--- /dev/null
+++ b/internal/tests/ai_llmintegration_test.go
@@ -0,0 +1,291 @@
+package tests
+
+import (
+ "context"
+ "testing"
+)
+
+func TestLlmFieldSemanticsTest_Eval(t *testing.T) {
+ tc := llmFieldSemanticsTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "all correct",
+ response: `{"temperature":"sampling_randomness","max_tokens":"completion_token_cap","stop":"generation_end_strings"}`,
+ want: 1,
+ },
+ {
+ name: "all correct fenced with prose",
+ response: "Here is my answer:\n```json\n{\"temperature\":\"sampling_randomness\",\"max_tokens\":\"completion_token_cap\",\"stop\":\"generation_end_strings\"}\n```",
+ want: 1,
+ },
+ {
+ name: "one wrong",
+ response: `{"temperature":"output_format","max_tokens":"completion_token_cap","stop":"generation_end_strings"}`,
+ want: 2.0 / 3.0,
+ },
+ {
+ name: "all wrong",
+ response: `{"temperature":"reply_length_cap","max_tokens":"minimum_reply_length","stop":"required_reply_prefix"}`,
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmTokenBudgetReasoningWant_GroundTruth(t *testing.T) {
+ // Independent re-derivation with plain subtraction.
+ const maxTokens, reasoningTokensUsed = 500, 437
+ want := maxTokens - reasoningTokensUsed
+
+ if want != 63 {
+ t.Fatalf("independently recomputed remaining budget = %d, want 63", want)
+ }
+ if llmTokenBudgetReasoningWant != want {
+ t.Errorf("llmTokenBudgetReasoningWant = %d, independently recomputed = %d", llmTokenBudgetReasoningWant, want)
+ }
+}
+
+func TestLlmTokenBudgetReasoningTest_Eval(t *testing.T) {
+ tc := llmTokenBudgetReasoningTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"exact", "63", 1},
+ {"prose wrapped", "63 completion tokens remain.", 1},
+ {"wrong: forgot to subtract", "500", 0},
+ {"wrong: reported the reasoning tokens used, not remaining", "437", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlm429RetryBackoffTest_Eval(t *testing.T) {
+ tc := llm429RetryBackoffTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"all correct", `{"should_retry":"yes","wait_seconds":2}`, 1},
+ {"all correct fenced", "```json\n{\"should_retry\":\"yes\",\"wait_seconds\":2}\n```", 1},
+ {"wrong: should not retry", `{"should_retry":"no","wait_seconds":2}`, 0.5},
+ {"wrong: ignored Retry-After value", `{"should_retry":"yes","wait_seconds":30}`, 0.5},
+ {"both wrong", `{"should_retry":"no","wait_seconds":30}`, 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmContextOverflowStrategyTest_Eval(t *testing.T) {
+ tc := llmContextOverflowStrategyTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: drop-oldest-turns", `{"strategy":"drop-oldest-turns"}`, 1},
+ {"correct, fenced with prose", "Given the no-altering constraint:\n```json\n{\"strategy\":\"drop-oldest-turns\"}\n```", 1},
+ {"wrong: summarize-oldest-turns", `{"strategy":"summarize-oldest-turns"}`, 0},
+ {"wrong: truncate-system-prompt", `{"strategy":"truncate-system-prompt"}`, 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmToolCallHandlingTest_Eval(t *testing.T) {
+ tc := llmToolCallHandlingTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"all correct", `{"next_role":"tool","id_field":"tool_call_id"}`, 1},
+ {"all correct fenced", "```json\n{\"next_role\":\"tool\",\"id_field\":\"tool_call_id\"}\n```", 1},
+ {"wrong role", `{"next_role":"assistant","id_field":"tool_call_id"}`, 0.5},
+ {"wrong field", `{"next_role":"tool","id_field":"id"}`, 0.5},
+ {"both wrong", `{"next_role":"user","id_field":"call_id"}`, 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmRolePlacementTest_Eval(t *testing.T) {
+ tc := llmRolePlacementTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: system", "system", 1},
+ {"correct, different case", "System", 1},
+ {"wrong: user", "user", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmTemperatureZeroCaveatTest_Eval(t *testing.T) {
+ tc := llmTemperatureZeroCaveatTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "both reasons present",
+ response: "Floating-point arithmetic is not strictly associative, and provider-side batching or GPU kernel routing changes summation order, which can flip the top token.",
+ want: 1,
+ },
+ {
+ name: "only the numerical reason",
+ response: "It is because of floating point rounding differences.",
+ want: 0.5,
+ },
+ {
+ name: "only the execution reason",
+ response: "It is because requests get batched together on different hardware.",
+ want: 0.5,
+ },
+ {
+ name: "neither reason",
+ response: "The model just has some inherent randomness left even at temperature zero.",
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmSSEStreamDoneTest_Eval(t *testing.T) {
+ tc := llmSSEStreamDoneTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct", "[DONE]", 1},
+ {"correct with whitespace", " [DONE] ", 1},
+ {"correct, case-insensitive", "[done]", 1},
+ {"wrong: missing brackets", "DONE", 0},
+ {"wrong: describes it instead", "the stream just closes", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmEmbeddingBatchMathWant_GroundTruth(t *testing.T) {
+ // Independent re-derivation with plain ceiling-division arithmetic,
+ // not via aiCeilDiv.
+ const total, batchSize = 10000, 96
+ full := total / batchSize
+ want := full
+ if total%batchSize != 0 {
+ want = full + 1
+ }
+
+ if want != 105 {
+ t.Fatalf("independently recomputed request count = %d, want 105", want)
+ }
+ if llmEmbeddingBatchMathWant != want {
+ t.Errorf("llmEmbeddingBatchMathWant = %d, independently recomputed = %d", llmEmbeddingBatchMathWant, want)
+ }
+}
+
+func TestLlmEmbeddingBatchMathTest_Eval(t *testing.T) {
+ tc := llmEmbeddingBatchMathTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"exact", "105", 1},
+ {"prose wrapped", "You need 105 requests.", 1},
+ {"wrong: dropped the partial batch", "104", 0},
+ {"wrong: off by a lot", "96", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestLlmStopSequenceTraceTest_Eval(t *testing.T) {
+ tc := llmStopSequenceTraceTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct", "Answer: 42", 1},
+ {"correct with whitespace", " Answer: 42 ", 1},
+ {"wrong: included text past the stop sequence", "Answer: 42\n\nExplanation: 6 times 7 is 42.", 0},
+ {"wrong: paraphrased instead of quoting exactly", "The answer is 42", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
diff --git a/internal/tests/ai_rag.go b/internal/tests/ai_rag.go
new file mode 100644
index 0000000..9505c1d
--- /dev/null
+++ b/internal/tests/ai_rag.go
@@ -0,0 +1,379 @@
+package tests
+
+import (
+ "github.com/lukaszraczylo/llm-testbench/internal/eval"
+ "github.com/lukaszraczylo/llm-testbench/internal/testkit"
+)
+
+func registerAIRAGTests(r *testkit.Registry) {
+ r.Register(ragChunkSizeTradeoffTest())
+ r.Register(ragRerankerPlacementTest())
+ r.Register(ragRetrievalFailureModeTest())
+ r.Register(ragCitationGroundingTest())
+ r.Register(ragContextSelectionVsStuffingTest())
+ r.Register(ragIndexStalenessTest())
+ r.Register(ragEvalMetricChoiceTest())
+ r.Register(ragHallucinationMitigationOrderingTest())
+ r.Register(ragMultihopDecompositionTest())
+ r.Register(ragPreAssemblyDedupTest())
+}
+
+// ragChunkSizeTradeoffTest: pick the smallest chunk size that reliably
+// holds a full fact without splitting it, given an explicit token-span
+// constraint.
+//
+// ground truth: facts span up to 450 tokens. The 128-token option is
+// smaller than 450, so it risks splitting a fact across two chunks. Both
+// the 512-token and 2048-token options are large enough to hold a 450-
+// token fact whole, but the larger option dilutes the embedded chunk with
+// more unrelated surrounding text. The smallest viable (non-splitting)
+// option is 512.
+func ragChunkSizeTradeoffTest() testkit.Test {
+ prompt := `You are choosing a chunk size for a RAG index over dense
+technical reference documents. A single self-contained fact in these
+documents commonly spans up to 450 tokens; a chunk boundary that falls
+inside a fact's span can leave a retrieved chunk missing part of that
+fact. A chunk much larger than a fact's span dilutes the embedding with
+unrelated surrounding text, lowering retrieval precision for queries about
+that fact.
+
+Three chunk-size options are available: 128 tokens, 512 tokens, and 2048
+tokens.
+
+Which of these is the smallest option that reliably holds a full 450-token
+fact without ever splitting it across two chunks? Respond with only a
+JSON object: {"chunk_size_tokens":128|512|2048}`
+
+ return testkit.Test{
+ ID: "rag-chunk-size-tradeoff",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Pick the smallest chunk size (512) that reliably holds a 450-token fact whole without excess dilution.",
+ Prompt: prompt,
+ Eval: eval.JSONField("chunk_size_tokens", 512),
+ }
+}
+
+// ragRerankerPlacementTest: order the three RAG pipeline stages -
+// retrieve, rerank, generate.
+//
+// ground truth: retrieval must run first to produce a broad candidate set
+// from the index; a reranker (e.g. a cross-encoder) can only reorder
+// candidates that retrieval already surfaced, so it must run second;
+// generation consumes the reranker's narrowed, reordered candidates last.
+func ragRerankerPlacementTest() testkit.Test {
+ prompt := `A RAG pipeline has these three stages, listed here in no
+particular order:
+- retrieve: run a broad search (vector and/or keyword) over the index to
+ produce an initial candidate set of chunks.
+- rerank: use a more expensive, more accurate model (e.g. a cross-encoder)
+ to reorder a candidate set by relevance and narrow it to the best few.
+- generate: have the LLM produce the final answer using the selected
+ chunks as context.
+
+Give the correct pipeline order as a JSON array of the three stage names,
+e.g. ["a","b","c"]. Respond with only the JSON array.`
+
+ return testkit.Test{
+ ID: "rag-reranker-placement",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Order the RAG pipeline stages: retrieve, rerank, generate.",
+ Prompt: prompt,
+ Eval: eval.JSONStringArrayEquals([]string{"retrieve", "rerank", "generate"}),
+ }
+}
+
+// ragRetrievalFailureModeCorpus is the inline single-document corpus for
+// ragRetrievalFailureModeTest, chosen so the query shares no exact term
+// with the relevant document.
+const ragRetrievalFailureModeCorpus = `Document 1: "Regular automobile maintenance, including oil changes and
+brake inspections, prevents most roadside breakdowns."`
+
+// ragRetrievalFailureModeTest: identify that a purely keyword/exact-term
+// search fails to retrieve a relevant document when the query uses
+// synonyms with no literal term overlap, while semantic (embedding-based)
+// search would succeed.
+//
+// ground truth: the query "car repair tips" shares zero exact terms with
+// Document 1's text ("automobile", "maintenance", "oil changes", "brake
+// inspections") - "car" vs "automobile" and "repair" vs "maintenance" are
+// synonym pairs, not shared tokens. A keyword/exact-term search over these
+// tokens finds no match and fails to retrieve this relevant document;
+// semantic (embedding-based) search, which captures the synonym
+// relationship, would retrieve it.
+func ragRetrievalFailureModeTest() testkit.Test {
+ prompt := `Here is a single-document corpus:
+
+` + ragRetrievalFailureModeCorpus + `
+
+Query: "car repair tips"
+
+Document 1 is genuinely relevant to this query, but one retrieval method
+fails to retrieve it because the query and the document share no exact,
+literal terms (only synonym pairs like "car"/"automobile" and
+"repair"/"maintenance"). Which retrieval method is expected to fail here:
+keyword-based (exact term matching, e.g. BM25) or semantic (embedding-
+based)? Respond with only one word: keyword or semantic.`
+
+ return testkit.Test{
+ ID: "rag-retrieval-failure-mode",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Identify keyword (exact-term) search as the method that fails on a synonym-only query/document pair.",
+ Prompt: prompt,
+ Eval: eval.Equals("keyword"),
+ }
+}
+
+// ragCitationGroundingTest: state the core RAG grounding requirement -
+// answers must be traceable to retrieved context, and the system must
+// decline (rather than guess from parametric knowledge) when the context
+// does not contain the answer.
+func ragCitationGroundingTest() testkit.Test {
+ prompt := `State the core grounding requirement a RAG system's answer
+generation must follow to avoid hallucination: what must every claim in
+the answer be traceable back to, and what must the system do instead of
+guessing when the retrieved context does not actually contain the
+information needed to answer the question?`
+
+ evaluator := eval.Mean(
+ eval.ContainsAny("cite", "citation", "source", "quote", "quoted"),
+ eval.ContainsAny("insufficient", "does not contain", "doesn't contain", "cannot answer", "can't answer", "no relevant", "unable to find", "not found in the context", "not present in the context"),
+ )
+
+ return testkit.Test{
+ ID: "rag-citation-grounding",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "State that answers must cite retrieved context and the system must decline when context lacks the answer, rather than guessing.",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// ragContextSelectionVsStuffingTest: choose relevance-filtered ("selective")
+// context assembly over cramming every chunk the budget allows
+// ("stuffed"), given a stated fact that low-relevance chunks degrade
+// answer quality.
+//
+// ground truth: the prompt states, as a given fact, that including
+// low-relevance chunks measurably degrades answer quality even when
+// budget remains. Under that stated fact, filling the entire available
+// budget with the 20 ranked chunks ("stuffed") would knowingly include
+// low-relevance chunks past the point where relevance drops off, so the
+// correct approach is "selective": include only chunks above a relevance
+// threshold, even when that uses less than the full budget.
+func ragContextSelectionVsStuffingTest() testkit.Test {
+ prompt := `A retriever ranks 20 chunks by relevance score for a query.
+The context window budget has room for 15 of those 20 chunks. It is an
+established fact for this system that including low-relevance chunks
+measurably degrades answer quality, even when there is still budget room
+left for them.
+
+Only the top 6 of the 20 ranked chunks score above the system's relevance
+threshold; the remaining 14 score below it.
+
+Given the stated fact above, should you assemble context by including all
+15 chunks the budget allows ("stuffed"), or by including only the 6
+chunks that clear the relevance threshold, leaving the rest of the budget
+unused ("selective")? Respond with only one word: stuffed or selective.`
+
+ return testkit.Test{
+ ID: "rag-context-selection-vs-stuffing",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Choose relevance-filtered ('selective') context assembly over budget-filling ('stuffed') given low-relevance chunks degrade quality.",
+ Prompt: prompt,
+ Eval: eval.Equals("selective"),
+ }
+}
+
+// ragIndexStalenessTest: diagnose stale chunks coexisting with fresh ones
+// after a naive re-embed pipeline, and pick the fix.
+//
+// ground truth: a pipeline that only inserts newly embedded chunks for an
+// updated document, without deleting that document's previous chunks
+// first, leaves both the old (stale) and new (fresh) chunks in the index
+// simultaneously. Retrieval can then surface the stale version alongside
+// or instead of the current one. The fix is to delete the document's old
+// chunks before (or atomically with) inserting the new ones, not to
+// change chunk overlap or top-k, which do not address stale chunks
+// remaining in the index at all.
+func ragIndexStalenessTest() testkit.Test {
+ prompt := `A document is updated from v1 to v2. The ingestion pipeline
+re-chunks and re-embeds only the new v2 text, then inserts those new
+chunks into the index. It never removes the v1 chunks that were inserted
+when the document was first indexed. After this update, queries about the
+document sometimes return v1's outdated chunks.
+
+What problem does this describe, and what is the correct fix? Respond
+with only a JSON object:
+{"problem":"stale-chunks-coexist"|"embedding-model-drift"|"query-cache-stale",
+ "fix":"delete-old-chunks-before-reinsert"|"increase-chunk-overlap"|"raise-top-k"}`
+
+ evaluator := eval.Mean(
+ eval.JSONField("problem", "stale-chunks-coexist"),
+ eval.JSONField("fix", "delete-old-chunks-before-reinsert"),
+ )
+
+ return testkit.Test{
+ ID: "rag-index-staleness",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Diagnose stale v1 chunks coexisting with v2 chunks after a naive re-embed, and pick delete-before-reinsert as the fix.",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// ragEvalMetricChoiceTest: map two scenarios to the RAG evaluation metric
+// that captures each, using definitions given entirely inline so the
+// answer follows from the prompt rather than memorized metric names.
+//
+// ground truth: scenario_a describes irrelevant chunks diluting the
+// retrieved set (3 relevant out of 10 retrieved) - exactly the definition
+// given for context_precision (fraction of retrieved chunks that are
+// relevant). scenario_b describes claims in the generated answer that are
+// not supported by the retrieved context (invented content) - exactly the
+// definition given for faithfulness (fraction of generated claims
+// supported by the retrieved context).
+func ragEvalMetricChoiceTest() testkit.Test {
+ prompt := `Three RAG evaluation metrics, defined here exactly:
+- context_precision: of the chunks that were retrieved, what fraction are
+ actually relevant to the query.
+- context_recall: of the chunks that are truly relevant to the query, what
+ fraction were retrieved.
+- faithfulness: of the individual factual claims made in the generated
+ answer, what fraction are directly supported by the retrieved context.
+
+scenario_a: "Retrieval returned 10 chunks for a query; only 3 of those 10
+were actually relevant to the query."
+
+scenario_b: "The generated answer made 5 factual claims; 2 of those 5
+claims were not supported by anything in the retrieved context - the
+model invented them."
+
+Using only the definitions given above, which single metric captures the
+specific problem in each scenario? Respond with only a JSON object:
+{"scenario_a":"context_precision"|"context_recall"|"faithfulness",
+ "scenario_b":"context_precision"|"context_recall"|"faithfulness"}`
+
+ evaluator := eval.Mean(
+ eval.JSONField("scenario_a", "context_precision"),
+ eval.JSONField("scenario_b", "faithfulness"),
+ )
+
+ return testkit.Test{
+ ID: "rag-eval-metric-choice",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Map two RAG failure scenarios to context_precision and faithfulness using inline metric definitions.",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// ragHallucinationMitigationOrderingTest: order 5 hallucination-mitigation
+// pipeline steps.
+//
+// ground truth: grounding context must be retrieved before the model can
+// be instructed to use it; the instruction to cite the context is part of
+// the prompt assembled before generation; generation itself must run
+// before its output can be verified; verification must run before
+// unsupported claims can be flagged/removed, since flagging depends on
+// verification's result.
+func ragHallucinationMitigationOrderingTest() testkit.Test {
+ prompt := `A RAG pipeline's hallucination-mitigation steps are listed
+here, identified by id, in no particular order:
+- step_retrieve: retrieve grounding context relevant to the query.
+- step_instruct: instruct the model, in the prompt, to answer using only
+ the retrieved context and to cite it.
+- step_generate: generate the answer.
+- step_verify: verify each claim in the generated answer against the
+ retrieved context.
+- step_flag: flag or remove any claim that verification could not support.
+
+Give the correct order these steps must run in, as a JSON array of their
+ids, e.g. ["step_a","step_b"]. Respond with only the JSON array.`
+
+ return testkit.Test{
+ ID: "rag-hallucination-mitigation-ordering",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Order 5 hallucination-mitigation pipeline steps: retrieve, instruct, generate, verify, flag.",
+ Prompt: prompt,
+ Eval: eval.JSONStringArrayEquals([]string{
+ "step_retrieve", "step_instruct", "step_generate", "step_verify", "step_flag",
+ }),
+ }
+}
+
+// ragMultihopDecompositionTest: decompose a multi-hop question into its
+// ordered sub-question fragments, selecting only the genuinely needed
+// fragments from a roster that includes one irrelevant distractor.
+//
+// ground truth: answering "what is the annual maintenance cost of the
+// storage tier used by the database that backs the service which handles
+// user authentication" requires, in strict dependency order: (1) which
+// service handles user authentication (frag_q), (2) which database backs
+// that service (frag_p, depends on frag_q's answer), (3) which storage
+// tier that database uses (frag_s, depends on frag_p's answer), (4) that
+// storage tier's annual maintenance cost (frag_r, depends on frag_s's
+// answer). frag_x (team headcount) answers a question never asked and
+// must be excluded.
+func ragMultihopDecompositionTest() testkit.Test {
+ prompt := `Multi-hop question: "What is the annual maintenance cost of
+the storage tier used by the database that backs the service which
+handles user authentication?"
+
+Candidate sub-question fragments, identified by id, in no particular
+order:
+- frag_q: "Which service handles user authentication?"
+- frag_p: "Which database backs that service?"
+- frag_s: "Which storage tier does that database use?"
+- frag_r: "What is the annual maintenance cost of that storage tier?"
+- frag_x: "What is the total headcount of the team that owns that
+ service?"
+
+Exactly one of these fragments (frag_x) does not help answer the
+multi-hop question and must be excluded. Give the remaining fragments, in
+the order they must be answered (each depends on the previous answer), as
+a JSON array of their ids. Respond with only the JSON array.`
+
+ return testkit.Test{
+ ID: "rag-multihop-decomposition",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Decompose a 4-hop question into its ordered sub-question fragments, excluding one irrelevant distractor.",
+ Prompt: prompt,
+ Eval: eval.JSONStringArrayEquals([]string{"frag_q", "frag_p", "frag_s", "frag_r"}),
+ }
+}
+
+// ragPreAssemblyDedupTest: explain why deduplicating near-identical
+// retrieved chunks before assembling context matters.
+func ragPreAssemblyDedupTest() testkit.Test {
+ prompt := `A query's top-8 retrieved chunks include 3 chunks that are
+near-identical copies of the same FAQ answer, mirrored across 3 different
+source pages, all of which scored highly enough to land in the top 8.
+
+Explain why deduplicating these near-identical chunks before assembling
+the final context (rather than passing all 3 copies straight through to
+the LLM) matters: name what assembling all 3 copies wastes, and what
+problem it risks introducing into the generated answer.`
+
+ evaluator := eval.Mean(
+ eval.ContainsAny("token budget", "context window", "context budget", "waste", "wastes", "wasted"),
+ eval.ContainsAny("redundant", "redundancy", "duplicate", "duplication", "repetition", "repetitive", "diversity", "biased", "bias", "skew"),
+ )
+
+ return testkit.Test{
+ ID: "rag-preassembly-dedup",
+ Category: "ai",
+ Subcategory: "rag",
+ Description: "Explain why deduplicating near-identical retrieved chunks before context assembly matters (wasted budget, redundancy/bias risk).",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
diff --git a/internal/tests/ai_rag_test.go b/internal/tests/ai_rag_test.go
new file mode 100644
index 0000000..c2f6bde
--- /dev/null
+++ b/internal/tests/ai_rag_test.go
@@ -0,0 +1,331 @@
+package tests
+
+import (
+ "context"
+ "testing"
+)
+
+func TestRagChunkSizeTradeoffTest_Eval(t *testing.T) {
+ tc := ragChunkSizeTradeoffTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: 512", `{"chunk_size_tokens":512}`, 1},
+ {"correct, fenced with prose", "The smallest viable option is:\n```json\n{\"chunk_size_tokens\":512}\n```", 1},
+ {"wrong: too small, would split a fact", `{"chunk_size_tokens":128}`, 0},
+ {"wrong: larger than needed, more dilution", `{"chunk_size_tokens":2048}`, 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagRerankerPlacementTest_Eval(t *testing.T) {
+ tc := ragRerankerPlacementTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct order", `["retrieve","rerank","generate"]`, 1},
+ {"correct order fenced", "```json\n[\"retrieve\",\"rerank\",\"generate\"]\n```", 1},
+ {"correct order different case", `["Retrieve","Rerank","Generate"]`, 1},
+ {"reranker before retrieval", `["rerank","retrieve","generate"]`, 0},
+ {"generate before rerank", `["retrieve","generate","rerank"]`, 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagRetrievalFailureModeTest_Eval(t *testing.T) {
+ tc := ragRetrievalFailureModeTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: keyword", "keyword", 1},
+ {"correct, different case", "Keyword", 1},
+ {"wrong: semantic", "semantic", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagCitationGroundingTest_Eval(t *testing.T) {
+ tc := ragCitationGroundingTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "both requirements present",
+ response: "Every claim must be traceable to a cited source in the retrieved context; if the context does not contain the answer, the system must say it cannot answer rather than guess.",
+ want: 1,
+ },
+ {
+ name: "only the citation requirement",
+ response: "Every claim in the answer must include a citation back to the retrieved context.",
+ want: 0.5,
+ },
+ {
+ name: "only the decline requirement",
+ response: "If the retrieved context does not contain the answer, the system must say it is unable to find it.",
+ want: 0.5,
+ },
+ {
+ name: "neither requirement",
+ response: "The answer should just sound confident and helpful.",
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagContextSelectionVsStuffingTest_Eval(t *testing.T) {
+ tc := ragContextSelectionVsStuffingTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: selective", "selective", 1},
+ {"correct, different case", "Selective", 1},
+ {"wrong: stuffed", "stuffed", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagIndexStalenessTest_Eval(t *testing.T) {
+ tc := ragIndexStalenessTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "all correct",
+ response: `{"problem":"stale-chunks-coexist","fix":"delete-old-chunks-before-reinsert"}`,
+ want: 1,
+ },
+ {
+ name: "all correct fenced",
+ response: "```json\n{\"problem\":\"stale-chunks-coexist\",\"fix\":\"delete-old-chunks-before-reinsert\"}\n```",
+ want: 1,
+ },
+ {
+ name: "wrong problem",
+ response: `{"problem":"embedding-model-drift","fix":"delete-old-chunks-before-reinsert"}`,
+ want: 0.5,
+ },
+ {
+ name: "wrong fix",
+ response: `{"problem":"stale-chunks-coexist","fix":"increase-chunk-overlap"}`,
+ want: 0.5,
+ },
+ {
+ name: "both wrong",
+ response: `{"problem":"query-cache-stale","fix":"raise-top-k"}`,
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagEvalMetricChoiceTest_Eval(t *testing.T) {
+ tc := ragEvalMetricChoiceTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "all correct",
+ response: `{"scenario_a":"context_precision","scenario_b":"faithfulness"}`,
+ want: 1,
+ },
+ {
+ name: "all correct fenced",
+ response: "```json\n{\"scenario_a\":\"context_precision\",\"scenario_b\":\"faithfulness\"}\n```",
+ want: 1,
+ },
+ {
+ name: "scenario_a wrong",
+ response: `{"scenario_a":"context_recall","scenario_b":"faithfulness"}`,
+ want: 0.5,
+ },
+ {
+ name: "scenario_b wrong",
+ response: `{"scenario_a":"context_precision","scenario_b":"context_recall"}`,
+ want: 0.5,
+ },
+ {
+ name: "both swapped",
+ response: `{"scenario_a":"faithfulness","scenario_b":"context_precision"}`,
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagHallucinationMitigationOrderingTest_Eval(t *testing.T) {
+ tc := ragHallucinationMitigationOrderingTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "correct order",
+ response: `["step_retrieve","step_instruct","step_generate","step_verify","step_flag"]`,
+ want: 1,
+ },
+ {
+ name: "correct order fenced",
+ response: "```json\n[\"step_retrieve\",\"step_instruct\",\"step_generate\",\"step_verify\",\"step_flag\"]\n```",
+ want: 1,
+ },
+ {
+ name: "verify before generate",
+ response: `["step_retrieve","step_instruct","step_verify","step_generate","step_flag"]`,
+ want: 0,
+ },
+ {
+ name: "flag before verify",
+ response: `["step_retrieve","step_instruct","step_generate","step_flag","step_verify"]`,
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagMultihopDecompositionTest_Eval(t *testing.T) {
+ tc := ragMultihopDecompositionTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "correct order, distractor excluded",
+ response: `["frag_q","frag_p","frag_s","frag_r"]`,
+ want: 1,
+ },
+ {
+ name: "correct order fenced",
+ response: "```json\n[\"frag_q\",\"frag_p\",\"frag_s\",\"frag_r\"]\n```",
+ want: 1,
+ },
+ {
+ name: "distractor incorrectly included",
+ response: `["frag_q","frag_p","frag_s","frag_r","frag_x"]`,
+ want: 0,
+ },
+ {
+ name: "hops reordered",
+ response: `["frag_p","frag_q","frag_s","frag_r"]`,
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestRagPreAssemblyDedupTest_Eval(t *testing.T) {
+ tc := ragPreAssemblyDedupTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {
+ name: "both reasons present",
+ response: "Passing all 3 copies wastes context window budget, and the redundant repetition risks biasing the answer toward that one duplicated point.",
+ want: 1,
+ },
+ {
+ name: "only the waste reason",
+ response: "It wastes token budget that could hold other useful chunks.",
+ want: 0.5,
+ },
+ {
+ name: "only the redundancy reason",
+ response: "It introduces redundant, repetitive information into the context.",
+ want: 0.5,
+ },
+ {
+ name: "neither reason",
+ response: "It just looks messy to a human reviewer.",
+ want: 0,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
diff --git a/internal/tests/ai_test.go b/internal/tests/ai_test.go
new file mode 100644
index 0000000..51baeaf
--- /dev/null
+++ b/internal/tests/ai_test.go
@@ -0,0 +1,49 @@
+package tests
+
+import (
+ "testing"
+
+ "github.com/lukaszraczylo/llm-testbench/internal/testkit"
+)
+
+// TestRegisterAITests_Wiring checks registerAITests wires up all three ai
+// subcategories - vector-search, llm-integration, rag - with 10 tests
+// each, all categorized "ai", and no duplicate IDs across the three
+// sibling files (testkit.Registry.Register panics on a duplicate ID,
+// which would fail this test outright). This is the integration point the
+// orchestrator's catalog wiring depends on; catalog.go itself is not
+// touched by this worktree.
+func TestRegisterAITests_Wiring(t *testing.T) {
+ r := testkit.NewRegistry()
+ registerAITests(r)
+
+ const wantTotal = 30
+ if r.Len() != wantTotal {
+ t.Fatalf("registerAITests: registry has %d tests, want %d", r.Len(), wantTotal)
+ }
+
+ counts := map[string]int{}
+ for _, tc := range r.All() {
+ if tc.Category != "ai" {
+ t.Errorf("test %q: Category = %q, want %q", tc.ID, tc.Category, "ai")
+ }
+ if tc.Prompt == "" {
+ t.Errorf("test %q: Prompt is empty", tc.ID)
+ }
+ if tc.Eval == nil {
+ t.Errorf("test %q: Eval is nil", tc.ID)
+ }
+ if tc.Description == "" {
+ t.Errorf("test %q: Description is empty", tc.ID)
+ }
+ if tc.MaxTokens != 0 {
+ t.Errorf("test %q: MaxTokens = %d, want 0 (omitted, per round-2 rule 7)", tc.ID, tc.MaxTokens)
+ }
+ counts[tc.Subcategory]++
+ }
+ for _, sub := range []string{"vector-search", "llm-integration", "rag"} {
+ if counts[sub] != 10 {
+ t.Errorf("subcategory %q: got %d tests, want 10", sub, counts[sub])
+ }
+ }
+}
diff --git a/internal/tests/ai_vectorsearch.go b/internal/tests/ai_vectorsearch.go
new file mode 100644
index 0000000..4a8473b
--- /dev/null
+++ b/internal/tests/ai_vectorsearch.go
@@ -0,0 +1,427 @@
+package tests
+
+import (
+ "strings"
+
+ "github.com/lukaszraczylo/llm-testbench/internal/eval"
+ "github.com/lukaszraczylo/llm-testbench/internal/testkit"
+)
+
+func registerAIVectorSearchTests(r *testkit.Registry) {
+ r.Register(vecCosineVsDotTest())
+ r.Register(vecRecallAtKTest())
+ r.Register(vecHNSWEfSearchTradeoffTest())
+ r.Register(vecPQMemoryMathTest())
+ r.Register(vecPreVsPostFilteringTest())
+ r.Register(vecRRFFusionTest())
+ r.Register(vecDistanceToSimilarityTest())
+ r.Register(vecNearDuplicateThresholdTest())
+ r.Register(vecIndexBuildQueryTradeoffTest())
+ r.Register(vecEmbeddingDimensionTradeoffTest())
+}
+
+// aiRRFScore computes the standard Reciprocal Rank Fusion score for one
+// document: the sum of 1/(k+rank) over every ranked list the document
+// appears in, using 1-indexed ranks. A document absent from a list
+// contributes nothing for that list (omit its rank from ranks).
+func aiRRFScore(k int, ranks ...int) float64 {
+ var sum float64
+ for _, rank := range ranks {
+ sum += 1 / float64(k+rank)
+ }
+ return sum
+}
+
+// aiUnitDistanceToCosineSimilarity converts a Euclidean distance d between
+// two unit-length (L2-normalized) vectors to their cosine similarity, using
+// ||a-b||^2 = |a|^2 + |b|^2 - 2*cos_sim = 2 - 2*cos_sim for |a|=|b|=1, so
+// cos_sim = 1 - d^2/2.
+func aiUnitDistanceToCosineSimilarity(d float64) float64 {
+ return 1 - d*d/2
+}
+
+// vecCosineVsDotTest: confirm that for two already-unit-normalized vectors,
+// the dot product equals their cosine similarity.
+//
+// ground truth: cosine similarity is dot(a,b) / (|a|*|b|). a = [0.6, 0.8]
+// has |a| = sqrt(0.36+0.64) = sqrt(1) = 1, and b = [0.8, 0.6] has the same
+// magnitude by symmetry. With both magnitudes exactly 1, the denominator is
+// 1, so dot(a,b) and cosine similarity are the same value. Answer: yes.
+func vecCosineVsDotTest() testkit.Test {
+ prompt := `Two vectors have already been L2-normalized to unit length
+(magnitude exactly 1): a = [0.6, 0.8] and b = [0.8, 0.6].
+
+Confirm the magnitudes: |a| = sqrt(0.6^2 + 0.8^2) = sqrt(0.36 + 0.64) =
+sqrt(1.0) = 1, and by the same arithmetic |b| = sqrt(0.8^2 + 0.6^2) = 1.
+
+Cosine similarity is defined as dot(a,b) / (|a| * |b|). Since |a| and |b|
+are both exactly 1 for this pair, is the plain dot product dot(a,b)
+mathematically equivalent to the cosine similarity, with no further
+division needed? Respond with only one word: yes or no.`
+
+ return testkit.Test{
+ ID: "vec-cosine-vs-dot",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Confirm that for two unit-normalized vectors, the plain dot product equals the cosine similarity.",
+ Prompt: prompt,
+ Eval: eval.Equals("yes"),
+ }
+}
+
+// vecRecallAtKRelevant and vecRecallAtKRetrieved are the inline
+// relevant-item set and ranked top-5 retrieval list for vecRecallAtKTest -
+// a fresh fixture (V-prefixed ids, recall@5 not @6) distinct from
+// paper-recall-at-k in whitepapers_vector.go.
+var (
+ vecRecallAtKRelevant = []string{"V3", "V8", "V11", "V14", "V20"}
+ vecRecallAtKRetrieved = []string{"V8", "V1", "V14", "V6", "V9"}
+)
+
+// vecRecallAtKWant is derived by calling wpRecallAtK (already defined in
+// whitepapers_compute.go for the same recall@k metric), not hardcoded.
+//
+// ground truth: of the 5 relevant items (V3, V8, V11, V14, V20), exactly 2
+// (V8 and V14) appear in the top-5 retrieved list; recall@5 = 2/5 = 0.4.
+// ai_vectorsearch_test.go independently recomputes this with a plain set
+// intersection.
+var vecRecallAtKWant = wpRecallAtK(vecRecallAtKRelevant, vecRecallAtKRetrieved)
+
+// vecRecallAtKTest: compute recall@k from a fresh relevant-item set and
+// ranked top-k retrieval list, distinct from the whitepapers recall@k
+// fixture.
+func vecRecallAtKTest() testkit.Test {
+ prompt := `Recall@k is a retrieval-quality metric: for one query, it is
+the fraction of that query's truly relevant items that appear anywhere
+within the system's top k retrieved results, regardless of exact rank.
+recall@k = |relevant intersect top-k| / |relevant|.
+
+A query's truly relevant vectors are: V3, V8, V11, V14, V20.
+
+A vector search system's ranked top-5 results for that query are, in rank
+order: V8, V1, V14, V6, V9.
+
+Compute recall@5 for this query. Respond with only the number, as a
+decimal (e.g. 0.6), rounded to 4 decimal places.`
+
+ return testkit.Test{
+ ID: "vec-recall-at-k",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Compute recall@5 from a fresh relevant-item set and ranked top-5 retrieval list.",
+ Prompt: prompt,
+ Eval: eval.Numeric(eval.ExtractLastNumber[float64], vecRecallAtKWant, 0.0005),
+ }
+}
+
+// vecHNSWEfSearchTradeoffTest: name the shared direction recall and query
+// latency move in when efSearch increases.
+//
+// ground truth: increasing efSearch (candidates explored per query) at
+// fixed efConstruction and M makes the search examine more candidate nodes
+// before stopping. That means more distance computations per query
+// (latency up) and a lower chance of missing a true nearest neighbor
+// (recall up) - both move in the same direction: up.
+func vecHNSWEfSearchTradeoffTest() testkit.Test {
+ prompt := `In an HNSW approximate nearest-neighbor index, efSearch
+controls how many candidate nodes the query-time search explores before
+stopping; efConstruction and M are fixed, index-build-time parameters that
+do not change per query.
+
+If you increase efSearch, keeping efConstruction and M fixed, query recall
+and query latency move in the same direction as each other. Do they both
+go up, or both go down? Respond with only one word: up or down.`
+
+ return testkit.Test{
+ ID: "vec-hnsw-efsearch-tradeoff",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "State that increasing HNSW's efSearch moves recall and query latency in the same direction (up).",
+ Prompt: prompt,
+ Eval: eval.Equals("up"),
+ }
+}
+
+// vecPQMemoryMathTest: derive total compressed index memory from PQ
+// parameters distinct from the whitepapers PQ-compression-ratio fixture
+// (different m, nbits, dimension, and a total-bytes framing rather than a
+// ratio).
+//
+// ground truth: codeBytes = m*nbits/8 = 16*4/8 = 8 bytes per compressed
+// vector (via wpQuantizedCodeBytes, already defined in
+// whitepapers_compute.go). Total for 1,000,000 vectors = 8 * 1,000,000 =
+// 8,000,000 bytes. ai_vectorsearch_test.go independently recomputes this
+// with plain arithmetic.
+var vecPQMemoryMathWant = wpQuantizedCodeBytes(16, 4) * 1000000
+
+func vecPQMemoryMathTest() testkit.Test {
+ prompt := `Product quantization (PQ) compresses each vector by splitting
+it into m equal-length subvectors and quantizing each one independently
+against its own codebook, storing only each subvector's nbits-bit centroid
+index. Concatenating the m per-subvector codes gives the full compressed
+representation, with no wasted bits (assume m*nbits is a multiple of 8).
+
+An index uses m = 16 subvectors and nbits = 4 bits per subvector.
+
+What is the total memory, in bytes, needed to store the compressed
+representations of 1,000,000 vectors indexed with these PQ parameters?
+Respond with only the integer number of bytes, with no commas, units, or
+other text (for example: 12345678).`
+
+ return testkit.Test{
+ ID: "vec-pq-memory-math",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Derive total compressed-index memory in bytes for 1M vectors from PQ parameters m=16, nbits=4.",
+ Prompt: prompt,
+ Eval: eval.Numeric(eval.ExtractLastNumber[int], vecPQMemoryMathWant, 0),
+ }
+}
+
+// vecPreVsPostFilteringTest: choose pre-filtering over post-filtering when
+// a metadata filter is highly selective relative to the requested result
+// count.
+//
+// ground truth: the filter matches only 0.5% of 10 million vectors (about
+// 50,000 candidates). Post-filtering runs ANN search first and discards
+// non-matching results afterward; with such low selectivity, a top-k ANN
+// search is very likely to return few or zero matching candidates unless k
+// is made enormous, so it cannot reliably return 10 matching results.
+// Pre-filtering (restricting the searched candidate set to only
+// filter-matching vectors before or during the ANN search) is required to
+// reliably surface 10 results that satisfy the filter.
+func vecPreVsPostFilteringTest() testkit.Test {
+ prompt := `An ANN vector index holds 10,000,000 vectors. A metadata
+filter (e.g. tenant_id = X) matches only 0.5% of those vectors. A query
+needs the top 10 nearest-neighbor results that also satisfy this filter.
+
+Should the filter be applied before/during the ANN search so only
+filter-matching vectors are ever considered (pre-filtering), or should the
+ANN search run first and non-matching results be discarded afterward
+(post-filtering), to reliably return 10 matching results? Respond with
+only one word: pre or post.`
+
+ return testkit.Test{
+ ID: "vec-pre-vs-post-filtering",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Choose pre-filtering over post-filtering ANN search when the metadata filter is highly selective (0.5%).",
+ Prompt: prompt,
+ Eval: eval.Equals("pre"),
+ }
+}
+
+// vecRRFFusionKeywordRanks and vecRRFFusionVectorRanks are the inline,
+// 1-indexed rank lists for vecRRFFusionTest, also used to build the
+// prompt's rank-list text via strings.Join so the displayed lists and the
+// computed ground truth can never drift apart. D5 sits at rank 3 in the
+// keyword list and rank 7 in the vector list.
+var (
+ vecRRFFusionKeywordRanks = []string{"D2", "D9", "D5", "D1", "D7", "D3", "D6", "D4"}
+ vecRRFFusionVectorRanks = []string{"D9", "D1", "D3", "D2", "D8", "D6", "D5", "D4"}
+)
+
+// vecRRFFusionK is the RRF constant used in vecRRFFusionTest's prompt,
+// stated explicitly so the answer is computable rather than memorized.
+const vecRRFFusionK = 60
+
+// vecRRFFusionDoc is the document vecRRFFusionTest asks the model to score.
+const vecRRFFusionDoc = "D5"
+
+// aiRankOf returns id's 1-indexed position in ranked, or 0 if id is
+// absent (matching aiRRFScore's convention that an absent document
+// contributes nothing for that list).
+func aiRankOf(id string, ranked []string) int {
+ for i, v := range ranked {
+ if v == id {
+ return i + 1
+ }
+ }
+ return 0
+}
+
+// vecRRFFusionWant is derived by calling aiRRFScore with vecRRFFusionDoc's
+// actual ranks in both lists, looked up via aiRankOf rather than
+// hardcoded.
+//
+// ground truth: RRF(d) = sum over each list of 1/(k + rank_list(d)). D5 is
+// rank 3 in the keyword list and rank 7 in the vector list, so RRF(D5) =
+// 1/(60+3) + 1/(60+7) = 1/63 + 1/67 ~= 0.030798. Rounded to 4dp: 0.0308.
+// ai_vectorsearch_test.go independently recomputes this with plain
+// arithmetic.
+var vecRRFFusionWant = round4dp(aiRRFScore(
+ vecRRFFusionK,
+ aiRankOf(vecRRFFusionDoc, vecRRFFusionKeywordRanks),
+ aiRankOf(vecRRFFusionDoc, vecRRFFusionVectorRanks),
+))
+
+func vecRRFFusionTest() testkit.Test {
+ prompt := `Reciprocal Rank Fusion (RRF) combines multiple ranked result
+lists into one fused score per document. For a document d, its RRF score is
+the sum, over every ranked list it appears in, of 1/(k + rank), where rank
+is the document's 1-indexed position in that list and k is a fixed
+constant. A document missing from a list contributes nothing for that list.
+Use k = 60.
+
+Keyword search results, ranked 1st to 8th: ` + strings.Join(vecRRFFusionKeywordRanks, ", ") + `
+
+Vector search results, ranked 1st to 8th: ` + strings.Join(vecRRFFusionVectorRanks, ", ") + `
+
+Using the formula above with k = 60, compute the RRF score for document ` + vecRRFFusionDoc + `.
+Respond with only the number, rounded to 4 decimal places.`
+
+ return testkit.Test{
+ ID: "vec-rrf-fusion",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Compute a document's Reciprocal Rank Fusion score (k=60) from its ranks in two inline result lists.",
+ Prompt: prompt,
+ Eval: eval.Numeric(eval.ExtractLastNumber[float64], vecRRFFusionWant, 0.0005),
+ }
+}
+
+// vecDistanceToSimilarityWant is derived by calling
+// aiUnitDistanceToCosineSimilarity, not hardcoded.
+//
+// ground truth: for unit-length vectors, ||a-b||^2 = 2 - 2*cos_sim, so
+// cos_sim = 1 - d^2/2. With d = 0.6, d^2 = 0.36, cos_sim = 1 - 0.18 = 0.82.
+// ai_vectorsearch_test.go independently recomputes this with plain
+// arithmetic.
+var vecDistanceToSimilarityWant = round4dp(aiUnitDistanceToCosineSimilarity(0.6))
+
+func vecDistanceToSimilarityTest() testkit.Test {
+ prompt := `Two vectors are both unit-length (L2-normalized to magnitude
+1). For unit-length vectors, squared Euclidean distance relates to cosine
+similarity by: ||a-b||^2 = |a|^2 + |b|^2 - 2*cos_sim(a,b) = 2 - 2*cos_sim(a,b),
+since |a|=|b|=1. Rearranging: cos_sim(a,b) = 1 - (||a-b||^2) / 2.
+
+The Euclidean distance between this pair is 0.6.
+
+Using the formula above, compute cos_sim(a,b). Respond with only the
+number, rounded to 4 decimal places.`
+
+ return testkit.Test{
+ ID: "vec-distance-to-similarity",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Convert a Euclidean distance between unit vectors to cosine similarity using the stated formula.",
+ Prompt: prompt,
+ Eval: eval.Numeric(eval.ExtractLastNumber[float64], vecDistanceToSimilarityWant, 0.0005),
+ }
+}
+
+// vecNearDuplicateA and vecNearDuplicateB are the inline 4-dim vectors for
+// vecNearDuplicateThresholdTest, distinct from py-cosine's 8-dim fixture in
+// python.go.
+var (
+ vecNearDuplicateA = []float64{1, 2, 2, 1}
+ vecNearDuplicateB = []float64{2, 2, 1, 1}
+)
+
+// vecNearDuplicateThreshold is the near-duplicate cosine similarity cutoff
+// stated in vecNearDuplicateThresholdTest's prompt.
+const vecNearDuplicateThreshold = 0.97
+
+// vecNearDuplicateThresholdWant is derived by calling cosineSimilarity
+// (mathutil.go), not hardcoded.
+//
+// ground truth: dot(a,b) = 1*2+2*2+2*1+1*1 = 9; |a|^2 = 1+4+4+1 = 10;
+// |b|^2 = 4+4+1+1 = 10; cosine similarity = 9/sqrt(10*10) = 9/10 = 0.9.
+// 0.9 - 0.97 = -0.07. ai_vectorsearch_test.go independently recomputes
+// this with plain arithmetic.
+var vecNearDuplicateThresholdWant = round4dp(cosineSimilarity(vecNearDuplicateA, vecNearDuplicateB) - vecNearDuplicateThreshold)
+
+func vecNearDuplicateThresholdTest() testkit.Test {
+ prompt := `A vector index deduplicates near-identical items: a candidate
+pair is treated as a near-duplicate only when its cosine similarity is at
+or above 0.97.
+
+Candidate pair: a = [1, 2, 2, 1], b = [2, 2, 1, 1].
+
+Compute the cosine similarity of a and b, then subtract the 0.97 threshold
+from it (similarity minus threshold). Respond with only the resulting
+number, rounded to 4 decimal places (it may be negative).`
+
+ return testkit.Test{
+ ID: "vec-near-duplicate-threshold",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Compute cosine similarity of a 4-dim vector pair and its signed distance from the 0.97 near-duplicate threshold.",
+ Prompt: prompt,
+ Eval: eval.Numeric(eval.ExtractLastNumber[float64], vecNearDuplicateThresholdWant, 0.0005),
+ }
+}
+
+// vecIndexBuildQueryTradeoffTest: pick the index type each stated
+// constraint forces, between HNSW (slower build, faster/more accurate
+// query) and IVF-Flat with few clusters (fast build, coarser query).
+//
+// ground truth: scenario_a's constraint - 50M vectors must be indexed in
+// under 10 minutes, with query recall/latency explicitly not critical -
+// rules out HNSW's comparatively slow graph-construction build and forces
+// IVF-Flat's fast single-pass clustering build. scenario_b's constraint -
+// query recall is critical, build time is unconstrained - forces HNSW,
+// since its graph gives materially better recall at query time than a
+// coarse IVF-Flat partition.
+func vecIndexBuildQueryTradeoffTest() testkit.Test {
+ prompt := `Two ANN index types are available:
+- hnsw: builds a navigable graph; build time is comparatively slow, but
+ queries are fast and highly accurate (high recall).
+- ivf-flat: clusters vectors into a small number of buckets via k-means;
+ build time is comparatively fast, but queries are coarser (lower recall)
+ than a well-tuned HNSW graph.
+
+scenario_a: "You must index 50,000,000 vectors in under 10 minutes on
+commodity hardware. Query recall and query latency are not critical for
+this use case."
+
+scenario_b: "Query recall is critical for this use case and build time is
+not constrained at all."
+
+For each scenario, which index type must you choose? Respond with only a
+JSON object: {"scenario_a":"hnsw"|"ivf-flat","scenario_b":"hnsw"|"ivf-flat"}`
+
+ evaluator := eval.Mean(
+ eval.JSONField("scenario_a", "ivf-flat"),
+ eval.JSONField("scenario_b", "hnsw"),
+ )
+
+ return testkit.Test{
+ ID: "vec-index-build-query-tradeoff",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Pick HNSW vs IVF-Flat for a build-time-constrained scenario and a recall-critical scenario.",
+ Prompt: prompt,
+ Eval: evaluator,
+ }
+}
+
+// vecEmbeddingDimensionTradeoffTest: pick the embedding dimension that fits
+// a fixed, uncompressed float32 storage budget.
+//
+// ground truth: bytes per vector = dims * 4 (float32). 128-dim: 128*4 = 512
+// bytes/vector, total for 100,000,000 vectors = 51,200,000,000 bytes =
+// 51.2 GB (using decimal GB = 1e9 bytes), which fits the 60 GB budget.
+// 768-dim: 768*4 = 3072 bytes/vector, total = 307,200,000,000 bytes =
+// 307.2 GB, which does not fit.
+func vecEmbeddingDimensionTradeoffTest() testkit.Test {
+ prompt := `You must store float32 (4 bytes per dimension), uncompressed
+embeddings for 100,000,000 vectors, with a hard storage budget of 60 GB
+(1 GB = 1,000,000,000 bytes). No quantization or compression is available.
+
+Two embedding models are candidates:
+- option 128: produces 128-dimensional embeddings.
+- option 768: produces 768-dimensional embeddings.
+
+Which option's total uncompressed storage fits within the 60 GB budget?
+Respond with only a JSON object: {"choice":"128"|"768"}`
+
+ return testkit.Test{
+ ID: "vec-embedding-dimension-tradeoff",
+ Category: "ai",
+ Subcategory: "vector-search",
+ Description: "Pick the embedding dimension (128 vs 768) whose uncompressed float32 storage fits a fixed 60 GB budget for 100M vectors.",
+ Prompt: prompt,
+ Eval: eval.JSONField("choice", "128"),
+ }
+}
diff --git a/internal/tests/ai_vectorsearch_test.go b/internal/tests/ai_vectorsearch_test.go
new file mode 100644
index 0000000..aecb869
--- /dev/null
+++ b/internal/tests/ai_vectorsearch_test.go
@@ -0,0 +1,321 @@
+package tests
+
+import (
+ "context"
+ "math"
+ "testing"
+)
+
+func TestVecCosineVsDotTest_Eval(t *testing.T) {
+ // Independent re-derivation: |a| = sqrt(0.6^2+0.8^2), |b| =
+ // sqrt(0.8^2+0.6^2), both exactly 1, so cosine similarity's
+ // denominator is 1 and the dot product equals the cosine similarity.
+ magA := math.Sqrt(0.6*0.6 + 0.8*0.8)
+ magB := math.Sqrt(0.8*0.8 + 0.6*0.6)
+ if math.Abs(magA-1) > 1e-9 || math.Abs(magB-1) > 1e-9 {
+ t.Fatalf("independently recomputed magnitudes = %v, %v, want both 1", magA, magB)
+ }
+
+ tc := vecCosineVsDotTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: yes", "yes", 1},
+ {"correct: yes, prose", "Yes, since both magnitudes are 1.", 0},
+ {"wrong: no", "no", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecRecallAtKWant_GroundTruth(t *testing.T) {
+ // Independent re-derivation with a plain set intersection, not via
+ // wpRecallAtK.
+ relevant := map[string]bool{"V3": true, "V8": true, "V11": true, "V14": true, "V20": true}
+ retrieved := []string{"V8", "V1", "V14", "V6", "V9"}
+ hit := 0
+ for _, id := range retrieved {
+ if relevant[id] {
+ hit++
+ }
+ }
+ want := float64(hit) / float64(len(relevant))
+
+ if want != 0.4 {
+ t.Fatalf("independently recomputed recall@5 = %v, want 0.4", want)
+ }
+ if vecRecallAtKWant != want {
+ t.Errorf("vecRecallAtKWant = %v, independently recomputed = %v", vecRecallAtKWant, want)
+ }
+}
+
+func TestVecRecallAtKTest_Eval(t *testing.T) {
+ tc := vecRecallAtKTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"exact", "0.4", 1},
+ {"prose wrapped", "recall@5 is 0.4.", 1},
+ {"wrong: counted an extra hit", "0.6", 0},
+ {"wrong: missed a hit", "0.2", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecHNSWEfSearchTradeoffTest_Eval(t *testing.T) {
+ tc := vecHNSWEfSearchTradeoffTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: up", "up", 1},
+ {"wrong: down", "down", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecPQMemoryMathWant_GroundTruth(t *testing.T) {
+ // Independent re-derivation with plain arithmetic, not via
+ // wpQuantizedCodeBytes.
+ const m, nbits, n = 16, 4, 1000000
+ codeBytes := m * nbits / 8
+ want := codeBytes * n
+
+ if want != 8000000 {
+ t.Fatalf("independently recomputed total bytes = %d, want 8000000", want)
+ }
+ if vecPQMemoryMathWant != want {
+ t.Errorf("vecPQMemoryMathWant = %d, independently recomputed = %d", vecPQMemoryMathWant, want)
+ }
+}
+
+func TestVecPQMemoryMathTest_Eval(t *testing.T) {
+ tc := vecPQMemoryMathTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"exact", "8000000", 1},
+ {"prose wrapped", "The total is 8000000 bytes.", 1},
+ {"wrong", "512", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecPreVsPostFilteringTest_Eval(t *testing.T) {
+ tc := vecPreVsPostFilteringTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: pre", "pre", 1},
+ {"wrong: post", "post", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecRRFFusionWant_GroundTruth(t *testing.T) {
+ // Independent re-derivation with plain arithmetic, not via aiRRFScore.
+ const k = 60
+ const keywordRank, vectorRank = 3, 7
+ want := math.Round((1.0/float64(k+keywordRank)+1.0/float64(k+vectorRank))*10000) / 10000
+
+ if math.Abs(want-0.0308) > 1e-9 {
+ t.Fatalf("independently recomputed RRF score = %v, want ~0.0308", want)
+ }
+ if vecRRFFusionWant != want {
+ t.Errorf("vecRRFFusionWant = %v, independently recomputed = %v", vecRRFFusionWant, want)
+ }
+}
+
+func TestVecRRFFusionTest_Eval(t *testing.T) {
+ tc := vecRRFFusionTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"exact", "0.0308", 1},
+ {"within tolerance", "0.0307", 1},
+ {"prose wrapped", "The RRF score for D5 is 0.0308.", 1},
+ {"wrong: only counted one list", "0.0159", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecDistanceToSimilarityWant_GroundTruth(t *testing.T) {
+ // Independent re-derivation with plain arithmetic, not via
+ // aiUnitDistanceToCosineSimilarity.
+ const d = 0.6
+ want := math.Round((1-d*d/2)*10000) / 10000
+
+ if math.Abs(want-0.82) > 1e-9 {
+ t.Fatalf("independently recomputed cosine similarity = %v, want 0.82", want)
+ }
+ if vecDistanceToSimilarityWant != want {
+ t.Errorf("vecDistanceToSimilarityWant = %v, independently recomputed = %v", vecDistanceToSimilarityWant, want)
+ }
+}
+
+func TestVecDistanceToSimilarityTest_Eval(t *testing.T) {
+ tc := vecDistanceToSimilarityTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"exact", "0.82", 1},
+ {"prose wrapped", "cos_sim(a,b) = 0.82", 1},
+ {"wrong: used d not d^2", "0.7", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecNearDuplicateThresholdWant_GroundTruth(t *testing.T) {
+ // Independent re-derivation with plain arithmetic, not via
+ // cosineSimilarity.
+ dot := 1.0*2 + 2.0*2 + 2.0*1 + 1.0*1
+ normA := math.Sqrt(1*1 + 2*2 + 2*2 + 1*1)
+ normB := math.Sqrt(2*2 + 2*2 + 1*1 + 1*1)
+ cos := dot / (normA * normB)
+ want := math.Round((cos-0.97)*10000) / 10000
+
+ if math.Abs(want-(-0.07)) > 1e-9 {
+ t.Fatalf("independently recomputed difference = %v, want -0.07", want)
+ }
+ if vecNearDuplicateThresholdWant != want {
+ t.Errorf("vecNearDuplicateThresholdWant = %v, independently recomputed = %v", vecNearDuplicateThresholdWant, want)
+ }
+}
+
+func TestVecNearDuplicateThresholdTest_Eval(t *testing.T) {
+ tc := vecNearDuplicateThresholdTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"exact", "-0.07", 1},
+ {"prose wrapped", "similarity minus threshold is -0.07.", 1},
+ {"wrong: forgot the sign", "0.07", 0},
+ {"wrong: reported similarity, not the difference", "0.9", 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecIndexBuildQueryTradeoffTest_Eval(t *testing.T) {
+ tc := vecIndexBuildQueryTradeoffTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"all correct", `{"scenario_a":"ivf-flat","scenario_b":"hnsw"}`, 1},
+ {"all correct fenced", "```json\n{\"scenario_a\":\"ivf-flat\",\"scenario_b\":\"hnsw\"}\n```", 1},
+ {"scenario_a wrong", `{"scenario_a":"hnsw","scenario_b":"hnsw"}`, 0.5},
+ {"scenario_b wrong", `{"scenario_a":"ivf-flat","scenario_b":"ivf-flat"}`, 0.5},
+ {"both swapped", `{"scenario_a":"hnsw","scenario_b":"ivf-flat"}`, 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}
+
+func TestVecEmbeddingDimensionTradeoffTest_Eval(t *testing.T) {
+ // Independent re-derivation confirming only 128-dim fits the budget.
+ const budgetBytes = 60_000_000_000
+ bytes128 := 128 * 4 * 100_000_000
+ bytes768 := 768 * 4 * 100_000_000
+ if bytes128 > budgetBytes {
+ t.Fatalf("128-dim total %d exceeds budget %d, expected it to fit", bytes128, budgetBytes)
+ }
+ if bytes768 <= budgetBytes {
+ t.Fatalf("768-dim total %d fits budget %d, expected it not to fit", bytes768, budgetBytes)
+ }
+
+ tc := vecEmbeddingDimensionTradeoffTest()
+ tests := []struct {
+ name string
+ response string
+ want float64
+ }{
+ {"correct: 128", `{"choice":"128"}`, 1},
+ {"correct, fenced with prose", "The 128-dim option fits:\n```json\n{\"choice\":\"128\"}\n```", 1},
+ {"wrong: 768", `{"choice":"768"}`, 0},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := tc.Eval.Evaluate(context.Background(), tt.response)
+ if got.Value != tt.want {
+ t.Errorf("Eval.Evaluate(%q) = %v, want %v (detail: %s)", tt.response, got.Value, tt.want, got.Detail)
+ }
+ })
+ }
+}