250 lines
6.6 KiB
Go
250 lines
6.6 KiB
Go
package mongo
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
|
|
|
"github.com/andydixon/relata/internal/domain"
|
|
)
|
|
|
|
// ChunkRepo stores item chunks.
|
|
type ChunkRepo struct{ base }
|
|
|
|
var _ domain.ChunkRepository = (*ChunkRepo)(nil)
|
|
|
|
// ReplaceForItem swaps an item's chunks for a new set.
|
|
//
|
|
// Delete-then-insert is what makes re-chunking idempotent: reprocessing an item
|
|
// leaves exactly one generation of chunks, never two overlapping ones. The
|
|
// caller is responsible for removing the corresponding vectors, which it does
|
|
// by deleting the item's vectors before re-embedding.
|
|
func (r *ChunkRepo) ReplaceForItem(ctx context.Context, workspaceID, itemID string, chunks []*domain.Chunk) error {
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
f = append(f, bson.E{Key: "item_id", Value: itemID})
|
|
|
|
if _, err := r.coll.DeleteMany(ctx, f); err != nil {
|
|
r.observe("replace_for_item", started, err)
|
|
return translate(err)
|
|
}
|
|
if len(chunks) == 0 {
|
|
r.observe("replace_for_item", started, nil)
|
|
return nil
|
|
}
|
|
|
|
docs := make([]any, 0, len(chunks))
|
|
for _, c := range chunks {
|
|
if c.WorkspaceID != workspaceID || c.ItemID != itemID {
|
|
return domain.ErrWorkspaceMiss
|
|
}
|
|
docs = append(docs, c)
|
|
}
|
|
_, err = r.coll.InsertMany(ctx, docs)
|
|
r.observe("replace_for_item", started, err)
|
|
return translate(err)
|
|
}
|
|
|
|
// ListByItem returns an item's chunks in sequence order.
|
|
func (r *ChunkRepo) ListByItem(ctx context.Context, workspaceID, itemID string) ([]*domain.Chunk, error) {
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
f = append(f, bson.E{Key: "item_id", Value: itemID})
|
|
|
|
cur, err := r.coll.Find(ctx, f, options.Find().SetSort(bson.D{{Key: "sequence", Value: 1}}))
|
|
r.observe("list_by_item", started, err)
|
|
if err != nil {
|
|
return nil, translate(err)
|
|
}
|
|
var out []*domain.Chunk
|
|
if err := cur.All(ctx, &out); err != nil {
|
|
return nil, err
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// GetMany fetches chunks by ID within a workspace.
|
|
func (r *ChunkRepo) GetMany(ctx context.Context, workspaceID string, ids []string) ([]*domain.Chunk, error) {
|
|
if len(ids) == 0 {
|
|
return nil, nil
|
|
}
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
f = append(f, bson.E{Key: "_id", Value: bson.D{{Key: "$in", Value: stringsToAny(ids)}}})
|
|
|
|
cur, err := r.coll.Find(ctx, f)
|
|
r.observe("get_many", started, err)
|
|
if err != nil {
|
|
return nil, translate(err)
|
|
}
|
|
var out []*domain.Chunk
|
|
if err := cur.All(ctx, &out); err != nil {
|
|
return nil, err
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// MarkEmbedded records that chunks now have vectors under a given model.
|
|
func (r *ChunkRepo) MarkEmbedded(ctx context.Context, workspaceID string, ids []string, model string) error {
|
|
if len(ids) == 0 {
|
|
return nil
|
|
}
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
f = append(f, bson.E{Key: "_id", Value: bson.D{{Key: "$in", Value: stringsToAny(ids)}}})
|
|
|
|
update := bson.D{{Key: "$set", Value: bson.D{
|
|
{Key: "embedded", Value: true},
|
|
{Key: "embedding_model", Value: model},
|
|
}}}
|
|
_, err = r.coll.UpdateMany(ctx, f, update)
|
|
r.observe("mark_embedded", started, err)
|
|
return translate(err)
|
|
}
|
|
|
|
// ListUnembedded finds chunks lacking a vector for the given model.
|
|
//
|
|
// A chunk embedded under a different model counts as unembedded, which is what
|
|
// makes re-embedding after a model change straightforward.
|
|
func (r *ChunkRepo) ListUnembedded(ctx context.Context, workspaceID, model string, limit int) ([]*domain.Chunk, error) {
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
f = append(f, bson.E{Key: "$or", Value: bson.A{
|
|
bson.D{{Key: "embedded", Value: bson.D{{Key: "$ne", Value: true}}}},
|
|
bson.D{{Key: "embedding_model", Value: bson.D{{Key: "$ne", Value: model}}}},
|
|
}})
|
|
|
|
if limit <= 0 || limit > 1000 {
|
|
limit = 200
|
|
}
|
|
cur, err := r.coll.Find(ctx, f, options.Find().SetLimit(int64(limit)).
|
|
SetSort(bson.D{{Key: "item_id", Value: 1}, {Key: "sequence", Value: 1}}))
|
|
r.observe("list_unembedded", started, err)
|
|
if err != nil {
|
|
return nil, translate(err)
|
|
}
|
|
var out []*domain.Chunk
|
|
if err := cur.All(ctx, &out); err != nil {
|
|
return nil, err
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// DeleteByItem removes an item's chunks.
|
|
func (r *ChunkRepo) DeleteByItem(ctx context.Context, workspaceID, itemID string) error {
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
f = append(f, bson.E{Key: "item_id", Value: itemID})
|
|
|
|
_, err = r.coll.DeleteMany(ctx, f)
|
|
r.observe("delete_by_item", started, err)
|
|
return translate(err)
|
|
}
|
|
|
|
// Count returns the number of chunks in a workspace.
|
|
func (r *ChunkRepo) Count(ctx context.Context, workspaceID string) (int64, error) {
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
n, err := r.coll.CountDocuments(ctx, f)
|
|
r.observe("count", started, err)
|
|
return n, translate(err)
|
|
}
|
|
|
|
// TextSearch finds chunks lexically, used to build result excerpts.
|
|
func (r *ChunkRepo) TextSearch(ctx context.Context, workspaceID, query string, limit int) ([]domain.ChunkHit, error) {
|
|
query = strings.TrimSpace(query)
|
|
if query == "" {
|
|
return nil, nil
|
|
}
|
|
ctx, cancel := r.ctx(ctx)
|
|
defer cancel()
|
|
started := time.Now()
|
|
|
|
f, err := wsFilter(workspaceID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
f = append(f, bson.E{Key: "$text", Value: bson.D{{Key: "$search", Value: query}}})
|
|
|
|
if limit <= 0 || limit > 500 {
|
|
limit = 60
|
|
}
|
|
opts := options.Find().
|
|
SetProjection(bson.D{
|
|
{Key: "score", Value: bson.D{{Key: "$meta", Value: "textScore"}}},
|
|
{Key: "item_id", Value: 1},
|
|
{Key: "text", Value: 1},
|
|
}).
|
|
SetSort(bson.D{{Key: "score", Value: bson.D{{Key: "$meta", Value: "textScore"}}}}).
|
|
SetLimit(int64(limit))
|
|
|
|
cur, err := r.coll.Find(ctx, f, opts)
|
|
r.observe("text_search", started, err)
|
|
if err != nil {
|
|
return nil, translate(err)
|
|
}
|
|
defer cur.Close(ctx)
|
|
|
|
var out []domain.ChunkHit
|
|
for cur.Next(ctx) {
|
|
var row struct {
|
|
ID string `bson:"_id"`
|
|
ItemID string `bson:"item_id"`
|
|
Text string `bson:"text"`
|
|
Score float64 `bson:"score"`
|
|
}
|
|
if err := cur.Decode(&row); err != nil {
|
|
continue
|
|
}
|
|
out = append(out, domain.ChunkHit{
|
|
ChunkID: row.ID, ItemID: row.ItemID, Text: row.Text, Score: row.Score,
|
|
})
|
|
}
|
|
return out, cur.Err()
|
|
}
|