Files
Jay Brown 72de690894 Merged in feature/chatbot1 (pull request #223)
Chatbot functionality

* baseline working

* missing test file

* more tests
2026-05-07 20:56:18 +00:00

595 lines
21 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Package bot — Milestone 3 payload assembly tests.
//
// These tests cover the four payload-assembly helpers backend-eng adds
// in M3 (per plan §7 mapping table):
//
// - Conversation history: last RecentTurns completed turns × 2 messages,
// ascending ordinal.
// - Scope payload assembly: stored doc ids folder-expanded doc ids,
// with the original folder ids preserved in the payload.
// - Cached-results stripping on updated_session_state.
// - 64KB session-state cap (chatbot.StateMaxBytes).
// - Metadata summarization (sum input/output tokens across the
// metadata.agents map; wall-clock fallback when metadata is absent).
//
// SQL-driven helpers (conversation history, scope expansion) use the
// real Postgres testcontainer. Pure functions (strip, cap, summarize)
// run in-memory. No httptest.NewServer here — that lives in
// fastapi_test.go for transport-layer cases only.
//
// Compile contract: this file references bot.BuildConversationHistory,
// bot.BuildChatScope, bot.StripCachedResults, bot.SummarizeMetadata,
// bot.SessionStateOversize, and the previously-declared types from
// fastapi_test.go (ChatScope, ConversationMessage, Metadata,
// AgentTelemetry). Backend-eng adds payload.go / state_strip.go /
// metadata_summarise.go and the package compiles.
package bot_test
import (
"context"
"encoding/json"
"os"
"strings"
"testing"
"time"
"queryorchestration/internal/bot"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/serviceconfig"
"queryorchestration/internal/serviceconfig/chatbot"
"queryorchestration/internal/test"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// payloadDefaultCfg returns a chatbot.ChatbotConfig with the plan §8
// defaults so the test does not have to repeat the literal values. The
// SERVICE_URL/API_KEY values are placeholders; payload tests do not
// dial out.
func payloadDefaultCfg() chatbot.ChatbotConfig {
return chatbot.ChatbotConfig{
ServiceURL: "http://chatbot.test:8000",
APIKey: "test-api-key",
RequestTimeoutSeconds: 60,
MaxTurnChars: 2000,
RecentTurns: 5,
}
}
// resetClientForPayload mirrors resetClientForBotTests but lives in the
// bot package's test space so this file does not pull in queryapi_test
// helpers. The shared testcontainer Postgres reuses volumes across runs,
// so leftover bot rows must be wiped before CreateClient.
func resetClientForPayload(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, clientID, name string) {
t.Helper()
pool := cfg.GetDBPool()
stmts := []string{
`DELETE FROM document_custom_metadata WHERE document_id IN (SELECT id FROM documents WHERE clientId = $1)`,
`UPDATE documents SET custom_schema_id = NULL WHERE clientId = $1`,
`DELETE FROM client_metadata_schemas WHERE client_id = $1`,
`DELETE FROM bot_session_documents WHERE document_id IN (SELECT id FROM documents WHERE clientId = $1)`,
`DELETE FROM bot_session_folders WHERE folder_id IN (SELECT id FROM folders WHERE clientId = $1)`,
`DELETE FROM bot_sessions WHERE client_id = $1`,
`DELETE FROM documents WHERE clientId = $1`,
`DELETE FROM folders WHERE clientId = $1`,
`DELETE FROM clients WHERE clientId = $1`,
}
for _, stmt := range stmts {
_, err := pool.Exec(ctx, stmt, clientID)
require.NoErrorf(t, err, "reset stmt failed: %s", stmt)
}
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
Name: name,
Clientid: clientID,
})
require.NoError(t, err)
}
// seedPayloadSession inserts a bot_sessions row owned by (clientID, actor).
func seedPayloadSession(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, clientID, actor string) uuid.UUID {
t.Helper()
var id uuid.UUID
err := cfg.GetDBPool().QueryRow(ctx, `
INSERT INTO bot_sessions (client_id, created_by) VALUES ($1, $2) RETURNING id
`, clientID, actor).Scan(&id)
require.NoError(t, err)
return id
}
// seedCompletedTurn inserts one completed bot_turns row at the given
// ordinal, with prompt and completion populated.
func seedCompletedTurn(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, sessionID uuid.UUID, ordinal int, prompt, completion string, createdAt time.Time) {
t.Helper()
_, err := cfg.GetDBPool().Exec(ctx, `
INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id, completion, attempt_started_at, created_at)
VALUES ($1, $2, $3, 'completed', $4, $5, $6, $6)
`, sessionID, ordinal, prompt, uuid.New(), completion, createdAt)
require.NoError(t, err)
}
// ---------- conversation history ----------
// TestConversationHistoryConversion: given the last RecentTurns
// completed turns, BuildConversationHistory produces N×2 messages
// (user prompt + assistant completion) in ascending ordinal order.
// Plan §7 mapping table.
func TestConversationHistoryConversion(t *testing.T) {
if testing.Short() {
t.SkipNow()
}
ctx := t.Context()
dbCfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, dbCfg)
const (
clientID = "BOT_M3_CONV"
actor = "history@example.com"
)
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-conv")
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
// Seed 7 completed turns. With cfg.RecentTurns=5 the helper must
// pick the last 5 ordered ascending: ordinals 3..7. createdAt is
// staggered by one second so any future ordering by created_at
// (rather than ordinal) is also deterministic.
now := time.Now().UTC().Truncate(time.Second)
for ord := 1; ord <= 7; ord++ {
seedCompletedTurn(t, ctx, dbCfg, sessionID, ord,
"prompt-"+string(rune('0'+ord)),
"answer-"+string(rune('0'+ord)),
now.Add(time.Duration(ord)*time.Second))
}
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
require.NoError(t, err)
history, err := svc.BuildConversationHistory(ctx, sessionID)
require.NoError(t, err)
require.Len(t, history, 5*2,
"plan §7: RecentTurns=5 -> 5 turns x 2 messages = 10")
// Expect ordinals 3,4,5,6,7 in ascending order. Each turn produces
// (user prompt, assistant completion) in that order.
for i := 0; i < 5; i++ {
ord := i + 3
userMsg := history[i*2]
assistantMsg := history[i*2+1]
assert.Equal(t, "user", userMsg.Role,
"index %d must be user; ordinal=%d", i*2, ord)
assert.Equal(t, "prompt-"+string(rune('0'+ord)), userMsg.Content,
"user content for ordinal %d", ord)
assert.Equal(t, "assistant", assistantMsg.Role,
"index %d must be assistant; ordinal=%d", i*2+1, ord)
assert.Equal(t, "answer-"+string(rune('0'+ord)), assistantMsg.Content,
"assistant content for ordinal %d", ord)
assert.Falsef(t, userMsg.Timestamp.IsZero(),
"user message ordinal %d must carry a non-zero timestamp", ord)
}
}
// TestConversationHistoryConversion_EmptySession: a session with zero
// completed turns yields a zero-length slice (NOT nil), so the FastAPI
// client emits the JSON `"conversation_history": []` per the swagger
// requirement that the field be a non-null array.
func TestConversationHistoryConversion_EmptySession(t *testing.T) {
if testing.Short() {
t.SkipNow()
}
ctx := t.Context()
dbCfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, dbCfg)
const (
clientID = "BOT_M3_CONV_E"
actor = "empty@example.com"
)
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-conv-e")
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
require.NoError(t, err)
history, err := svc.BuildConversationHistory(ctx, sessionID)
require.NoError(t, err)
require.NotNil(t, history,
"empty result must be a non-nil slice so json.Marshal emits []")
assert.Empty(t, history)
}
// ---------- scope payload assembly ----------
// TestScopePayloadAssembly_FoldersExpanded: stored scope has 2 doc ids
// + 1 folder id; folder contains 3 unrelated docs. Resulting payload
// has documents = the union of 5 ids; folders = the original 1 folder.
// Plan §7 mapping table: "scope": "Stored document ids union expanded
// folder document ids; stored folder ids preserved".
func TestScopePayloadAssembly_FoldersExpanded(t *testing.T) {
if testing.Short() {
t.SkipNow()
}
ctx := t.Context()
dbCfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, dbCfg)
pool := dbCfg.GetDBPool()
const (
clientID = "BOT_M3_SCOPE_FX"
actor = "scope@example.com"
)
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-fx")
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
queries := dbCfg.GetDBQueries()
// 2 stored documents (no folder).
docA, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{Clientid: clientID, Hash: "scope-fx-a"})
require.NoError(t, err)
docB, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{Clientid: clientID, Hash: "scope-fx-b"})
require.NoError(t, err)
// 1 folder containing 3 documents (different from the stored docs).
folder, err := queries.CreateFolder(ctx, &repository.CreateFolderParams{
Path: "/scope-fx-folder", Parentid: nil, Clientid: clientID, Createdby: "tester",
})
require.NoError(t, err)
folderDocs := make([]uuid.UUID, 3)
for i := 0; i < 3; i++ {
var id uuid.UUID
err := pool.QueryRow(ctx, `
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
`, clientID, "scope-fx-folderdoc-"+string(rune('0'+i)), folder.ID).Scan(&id)
require.NoError(t, err)
folderDocs[i] = id
}
// Persist the scope: 2 stored docs + 1 folder.
for _, d := range []uuid.UUID{docA, docB} {
_, err := pool.Exec(ctx,
`INSERT INTO bot_session_documents (session_id, document_id) VALUES ($1, $2)`,
sessionID, d)
require.NoError(t, err)
}
_, err = pool.Exec(ctx,
`INSERT INTO bot_session_folders (session_id, folder_id) VALUES ($1, $2)`,
sessionID, folder.ID)
require.NoError(t, err)
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
require.NoError(t, err)
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
require.NoError(t, err)
wantDocs := []string{
docA.String(), docB.String(),
folderDocs[0].String(), folderDocs[1].String(), folderDocs[2].String(),
}
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
"plan §7: documents = stored docs folder-expanded docs")
assert.ElementsMatch(t, []string{folder.ID.String()}, scope.FolderIDs,
"plan §7: stored folder ids preserved")
}
// TestScopePayloadAssembly_DocsAndFolderOverlap: stored doc id is also
// inside the folder being expanded. Payload documents must dedup so the
// shared id appears exactly once. Folder id preserved.
func TestScopePayloadAssembly_DocsAndFolderOverlap(t *testing.T) {
if testing.Short() {
t.SkipNow()
}
ctx := t.Context()
dbCfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, dbCfg)
pool := dbCfg.GetDBPool()
const (
clientID = "BOT_M3_SCOPE_OL"
actor = "overlap@example.com"
)
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-ol")
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
queries := dbCfg.GetDBQueries()
folder, err := queries.CreateFolder(ctx, &repository.CreateFolderParams{
Path: "/scope-ol-folder", Parentid: nil, Clientid: clientID, Createdby: "tester",
})
require.NoError(t, err)
// docA is the overlap: stored AND inside the folder.
var docA uuid.UUID
err = pool.QueryRow(ctx, `
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
`, clientID, "scope-ol-overlap", folder.ID).Scan(&docA)
require.NoError(t, err)
// docB is folder-only.
var docB uuid.UUID
err = pool.QueryRow(ctx, `
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
`, clientID, "scope-ol-folder-only", folder.ID).Scan(&docB)
require.NoError(t, err)
// Stored scope: docA (also in folder) + folder.
_, err = pool.Exec(ctx,
`INSERT INTO bot_session_documents (session_id, document_id) VALUES ($1, $2)`,
sessionID, docA)
require.NoError(t, err)
_, err = pool.Exec(ctx,
`INSERT INTO bot_session_folders (session_id, folder_id) VALUES ($1, $2)`,
sessionID, folder.ID)
require.NoError(t, err)
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
require.NoError(t, err)
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
require.NoError(t, err)
wantDocs := []string{docA.String(), docB.String()}
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
"overlap document must appear exactly once after dedup")
assert.ElementsMatch(t, []string{folder.ID.String()}, scope.FolderIDs)
}
// TestScopePayloadAssembly_EmptyScopeReturnsAllClientDocs: with no
// stored doc/folder rows, the payload documents fall back to "all
// documents for the client" per plan §4 ("Empty scope tables mean all
// documents for the client"). Folders array is empty.
func TestScopePayloadAssembly_EmptyScopeReturnsAllClientDocs(t *testing.T) {
if testing.Short() {
t.SkipNow()
}
ctx := t.Context()
dbCfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, dbCfg)
const (
clientID = "BOT_M3_SCOPE_E"
actor = "empty-scope@example.com"
)
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-e")
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
queries := dbCfg.GetDBQueries()
wantDocs := []string{}
for i := 0; i < 4; i++ {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientID, Hash: "scope-e-" + string(rune('0'+i)),
})
require.NoError(t, err)
wantDocs = append(wantDocs, id.String())
}
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
require.NoError(t, err)
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
require.NoError(t, err)
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
"empty scope tables -> all documents for the client")
assert.Empty(t, scope.FolderIDs)
}
// ---------- cached_results stripping ----------
// TestStripCachedResults_RemovesKey: input state with a cached_results
// entry returns the same object minus that one key. Other keys are
// preserved verbatim.
func TestStripCachedResults_RemovesKey(t *testing.T) {
t.Parallel()
in := json.RawMessage(`{"foo":1,"cached_results":{"abc":{"k":"v"}},"bar":"hello"}`)
out, err := bot.StripCachedResults(in)
require.NoError(t, err)
var parsed map[string]interface{}
require.NoError(t, json.Unmarshal(out, &parsed))
_, hasCached := parsed["cached_results"]
assert.False(t, hasCached, "cached_results must be removed")
assert.EqualValues(t, 1, parsed["foo"], "other keys must be preserved")
assert.Equal(t, "hello", parsed["bar"])
}
// TestStripCachedResults_NoOpWhenAbsent: input without cached_results is
// returned with the same logical shape.
func TestStripCachedResults_NoOpWhenAbsent(t *testing.T) {
t.Parallel()
in := json.RawMessage(`{"foo":1,"bar":"hello"}`)
out, err := bot.StripCachedResults(in)
require.NoError(t, err)
var got map[string]interface{}
require.NoError(t, json.Unmarshal(out, &got))
var want map[string]interface{}
require.NoError(t, json.Unmarshal(in, &want))
assert.Equal(t, want, got)
}
// TestStripCachedResults_NullInput: a nil/empty input passes through
// unchanged so the caller can hand it null state without a special case.
func TestStripCachedResults_NullInput(t *testing.T) {
t.Parallel()
out, err := bot.StripCachedResults(nil)
require.NoError(t, err)
assert.Nil(t, out, "nil input must round-trip as nil")
}
// ---------- 64KB session-state cap ----------
// TestSessionStateSizeCap_64KB_Boundary: the cap is plan §8's
// chatbot.StateMaxBytes. A stripped state at exactly the cap reports
// as not-oversize; one byte over reports oversize. The helper returns
// (oversize bool) without mutating the input.
func TestSessionStateSizeCap_64KB_Boundary(t *testing.T) {
t.Parallel()
// Build a JSON object whose length is exactly chatbot.StateMaxBytes.
// We use a single-key string-valued object and pad the value so the
// total raw byte length is exact.
const head = `{"k":"`
const tail = `"}`
overhead := len(head) + len(tail)
require.Greater(t, chatbot.StateMaxBytes, overhead)
value := strings.Repeat("a", chatbot.StateMaxBytes-overhead)
stateAtCap := json.RawMessage(head + value + tail)
require.Equal(t, chatbot.StateMaxBytes, len(stateAtCap),
"test setup: stateAtCap must equal the cap byte-for-byte")
assert.False(t, bot.SessionStateOversize(stateAtCap),
"plan §8: state at exactly StateMaxBytes is not oversize")
// One byte over the cap.
overValue := strings.Repeat("a", chatbot.StateMaxBytes-overhead+1)
stateOver := json.RawMessage(head + overValue + tail)
require.Equal(t, chatbot.StateMaxBytes+1, len(stateOver))
assert.True(t, bot.SessionStateOversize(stateOver),
"plan §8: state above StateMaxBytes is oversize")
}
// ---------- metadata summarization ----------
// TestMetadataSummarization: response metadata.agents = {a: {input:10,
// output:20}, b: {input:5, output:30}} -> tokens_in=15, tokens_out=50,
// latency = metadata.total_latency_ms (NOT wall clock when metadata is
// present).
func TestMetadataSummarization(t *testing.T) {
t.Parallel()
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
receivedAt := sentAt.Add(7 * time.Second)
metadata := &bot.Metadata{
TotalLatencyMs: 6000,
RetryCount: 0,
Agents: map[string]bot.AgentTelemetry{
"a": {LatencyMs: 1000, InputTokens: 10, OutputTokens: 20},
"b": {LatencyMs: 5000, InputTokens: 5, OutputTokens: 30},
},
}
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(metadata, sentAt, receivedAt)
assert.Equal(t, 6000, latencyMs,
"plan §7: latency = metadata.total_latency_ms when metadata is present")
assert.Equal(t, 15, tokensIn,
"plan §7: tokens_in = sum of agents[].input_tokens")
assert.Equal(t, 50, tokensOut,
"plan §7: tokens_out = sum of agents[].output_tokens")
}
// TestMetadataSummarization_NilMetadataUsesWallClock: nil metadata ->
// latency = ReceivedAt - SentAt in ms, tokens default to 0.
func TestMetadataSummarization_NilMetadataUsesWallClock(t *testing.T) {
t.Parallel()
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
receivedAt := sentAt.Add(2500 * time.Millisecond)
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(nil, sentAt, receivedAt)
assert.Equal(t, 2500, latencyMs,
"nil metadata -> wall-clock latency in milliseconds")
assert.Equal(t, 0, tokensIn)
assert.Equal(t, 0, tokensOut)
}
// TestMetadataSummarization_EmptyAgentsMap: metadata present but agents
// map is empty -> tokens default to 0; latency uses metadata.total_latency_ms.
func TestMetadataSummarization_EmptyAgentsMap(t *testing.T) {
t.Parallel()
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
receivedAt := sentAt.Add(5 * time.Second)
metadata := &bot.Metadata{
TotalLatencyMs: 4500,
RetryCount: 2,
Agents: map[string]bot.AgentTelemetry{},
}
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(metadata, sentAt, receivedAt)
assert.Equal(t, 4500, latencyMs)
assert.Equal(t, 0, tokensIn)
assert.Equal(t, 0, tokensOut)
}
// TestStripCachedResults_OnRealAariaSample exercises StripCachedResults
// against the authoritative Aaria session_state shape Q committed to
// plans/chatbot.stuff/sample.response.json. The sample carries the full
// 10-key session-state envelope including a populated cached_results
// map keyed by result hash; stripping must leave the other 9 keys
// intact and the resulting bytes must fit comfortably under
// chatbot.StateMaxBytes.
//
// The other-key preservation check uses json.Unmarshal->reflect.DeepEqual
// rather than byte equality because Go's encoding/json does not
// guarantee map-key emission order: a sound implementation may re-emit
// keys in a different order than the source file, but the parsed value
// must be identical. Plan §7 mapping table, plan §8 64KB cap.
func TestStripCachedResults_OnRealAariaSample(t *testing.T) {
t.Parallel()
const samplePath = "../../plans/chatbot.stuff/sample.response.json"
raw, err := os.ReadFile(samplePath)
require.NoErrorf(t, err, "must read fixture at %s", samplePath)
// Sanity-precondition: the fixture has the 10 documented top-level
// keys including cached_results. If Q reshapes the fixture, this
// assertion fails first with a clear pointer at the source of drift.
var input map[string]interface{}
require.NoError(t, json.Unmarshal(raw, &input),
"fixture must be a JSON object")
require.Len(t, input, 10,
"fixture must have 10 top-level keys before stripping; got %d", len(input))
_, hadCached := input["cached_results"]
require.True(t, hadCached, "fixture must include cached_results")
stripped, err := bot.StripCachedResults(json.RawMessage(raw))
require.NoError(t, err)
require.NotNil(t, stripped, "non-empty input must produce non-nil output")
var out map[string]interface{}
require.NoError(t, json.Unmarshal(stripped, &out),
"stripped output must be a valid JSON object")
require.Len(t, out, 9,
"stripping cached_results must leave exactly 9 keys; got %d", len(out))
_, hasCached := out["cached_results"]
assert.False(t, hasCached, "cached_results must be absent after strip")
// Every other key must round-trip with the original value. Plan §7
// is explicit that ONLY cached_results is touched.
expectedKeys := []string{
"turn",
"resolved_entities",
"last_scope",
"last_intent",
"last_query_sql",
"last_query_plan",
"last_result_hash",
"last_scope_signature",
"validator_retry_count",
}
for _, k := range expectedKeys {
require.Containsf(t, out, k,
"key %q must be preserved after strip", k)
assert.Equalf(t, input[k], out[k],
"key %q value must round-trip unchanged", k)
}
assert.Falsef(t, bot.SessionStateOversize(stripped),
"plan §8: real Aaria sample is well under 64 KiB after strip; got %d bytes", len(stripped))
}