weichenandCursor
06b9278996
项目 6579 的正式正文已换成真实解析内容，冲击波 PDF 写入 90 条带向量的切块，小文档仍全文注入。 Co-authored-by: Cursor <cursoragent@cursor.com>
60 lines
1.7 KiB
Go
60 lines
1.7 KiB
Go
package materialindex
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
|
|
"chat-agent-system/internal/config"
|
|
)
|
|
|
|
func TestPipelineSearchShockwave(t *testing.T) {
|
|
if os.Getenv("MATERIAL_SEARCH_LIVE") != "1" {
|
|
t.Skip("set MATERIAL_SEARCH_LIVE=1 to search the dev material_chunks index")
|
|
}
|
|
cfg, err := config.Load("")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if cfg.EmbeddingAPIKey == "" || cfg.RerankerAPIKey == "" {
|
|
t.Fatal("embedding and reranker credentials must be in the environment")
|
|
}
|
|
dsn := fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=disable",
|
|
cfg.ContextDbUser, cfg.ContextDbPassword, cfg.ContextDbHost, cfg.ContextDbPort, cfg.ContextDbName)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
|
|
defer cancel()
|
|
pool, err := pgxpool.New(ctx, dsn)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer pool.Close()
|
|
p := NewPipeline(cfg, pool)
|
|
const blob = "57eb6d860bc8526da6f9e6adb4866e7699edcafec219d93f23758f7cc2f447bb"
|
|
hits, err := p.Search(ctx, 6579, blob, "能量参数是多少", 3)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(hits) == 0 {
|
|
t.Fatal("expected reranked hits for the shockwave PDF")
|
|
}
|
|
foundParams := false
|
|
for i, hit := range hits {
|
|
preview := strings.ReplaceAll(hit.Title+" "+hit.Content, "\n", " ")
|
|
if len(preview) > 160 {
|
|
preview = preview[:160]
|
|
}
|
|
t.Logf("hit%d title=%q pages=%d-%d score=%.4f preview=%s", i, hit.Title, hit.PageStart, hit.PageEnd, hit.Score, preview)
|
|
if strings.Contains(hit.Content, "MPa") || strings.Contains(hit.Content, "6-30") {
|
|
foundParams = true
|
|
}
|
|
}
|
|
if !foundParams {
|
|
t.Fatal("reranked hits did not include the energy-parameter fragment")
|
|
}
|
|
}
|