mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 09:01:31 +00:00
refactor: move Ochre vector store to the Ochre domain
handlers/ochrevector becomes domains/ochre/vector. Callers import it as ochrevector because vector is a common local name in the gateway.
This commit is contained in:
@@ -9,7 +9,7 @@ import (
|
||||
"os"
|
||||
"time"
|
||||
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
@@ -61,8 +61,8 @@ var applianceTeardownFn = func(ctx context.Context, purgeMetadata bool) error {
|
||||
return err
|
||||
}
|
||||
defer nc.Close()
|
||||
_, err = natsmsg.NATSRequest[handlers_ochrevector.TeardownApplianceResponse](ctx, nc,
|
||||
handlers_ochrevector.SubjectTeardownAppliance, &handlers_ochrevector.TeardownApplianceRequest{PurgeMetadata: purgeMetadata},
|
||||
_, err = natsmsg.NATSRequest[ochrevector.TeardownApplianceResponse](ctx, nc,
|
||||
ochrevector.SubjectTeardownAppliance, &ochrevector.TeardownApplianceRequest{PurgeMetadata: purgeMetadata},
|
||||
applianceTeardownTimeout, awsidentifiers.GlobalAccountID)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -9,8 +9,8 @@ import (
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
awsidentifiers "github.com/mulgadc/spinifex/spinifex/foundation/aws/identifiers"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
"github.com/pterm/pterm"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
@@ -165,12 +165,12 @@ func init() {
|
||||
// vectorServiceFn indirects the NATS-backed client so the Run functions'
|
||||
// connect/exit control flow can be tested without a live daemon, mirroring
|
||||
// endpointServiceFn.
|
||||
var vectorServiceFn = func() (handlers_ochrevector.VectorService, func(), error) {
|
||||
var vectorServiceFn = func() (ochrevector.VectorService, func(), error) {
|
||||
_, nc, err := loadConfigAndConnectFn()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return handlers_ochrevector.NewNATSVectorService(nc), nc.Close, nil
|
||||
return ochrevector.NewNATSVectorService(nc), nc.Close, nil
|
||||
}
|
||||
|
||||
// vectorBackupRestoreTimeout bounds the backup/restore NATS round trip: it
|
||||
@@ -184,14 +184,14 @@ const vectorBackupRestoreTimeout = 10 * time.Minute
|
||||
// target account is named in the request payload, not derived from the
|
||||
// caller's own header identity, since the operator CLI always connects as
|
||||
// the global account but must be able to target any tenant's schema.
|
||||
var backupAccountFn = func(ctx context.Context, accountID string) (*handlers_ochrevector.BackupAccountResponse, error) {
|
||||
var backupAccountFn = func(ctx context.Context, accountID string) (*ochrevector.BackupAccountResponse, error) {
|
||||
_, nc, err := loadConfigAndConnectFn()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer nc.Close()
|
||||
return natsmsg.NATSRequest[handlers_ochrevector.BackupAccountResponse](ctx, nc,
|
||||
handlers_ochrevector.SubjectBackupAccount, &handlers_ochrevector.BackupAccountRequest{AccountID: accountID},
|
||||
return natsmsg.NATSRequest[ochrevector.BackupAccountResponse](ctx, nc,
|
||||
ochrevector.SubjectBackupAccount, &ochrevector.BackupAccountRequest{AccountID: accountID},
|
||||
vectorBackupRestoreTimeout, awsidentifiers.GlobalAccountID)
|
||||
}
|
||||
|
||||
@@ -203,14 +203,14 @@ var restoreAccountFn = func(ctx context.Context, accountID, objectKey string) er
|
||||
return err
|
||||
}
|
||||
defer nc.Close()
|
||||
_, err = natsmsg.NATSRequest[handlers_ochrevector.RestoreAccountResponse](ctx, nc,
|
||||
handlers_ochrevector.SubjectRestoreAccount, &handlers_ochrevector.RestoreAccountRequest{AccountID: accountID, ObjectKey: objectKey},
|
||||
_, err = natsmsg.NATSRequest[ochrevector.RestoreAccountResponse](ctx, nc,
|
||||
ochrevector.SubjectRestoreAccount, &ochrevector.RestoreAccountRequest{AccountID: accountID, ObjectKey: objectKey},
|
||||
vectorBackupRestoreTimeout, awsidentifiers.GlobalAccountID)
|
||||
return err
|
||||
}
|
||||
|
||||
// runBackupAccount is the testable core of 'ochre vector backup'.
|
||||
func runBackupAccount(ctx context.Context, accountID string, backup func(context.Context, string) (*handlers_ochrevector.BackupAccountResponse, error)) (string, error) {
|
||||
func runBackupAccount(ctx context.Context, accountID string, backup func(context.Context, string) (*ochrevector.BackupAccountResponse, error)) (string, error) {
|
||||
out, err := backup(ctx, accountID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -227,7 +227,7 @@ func runRestoreAccount(ctx context.Context, accountID, objectKey string, restore
|
||||
}
|
||||
|
||||
// formatIndexRecord renders one index record as aligned key/value lines.
|
||||
func formatIndexRecord(rec handlers_ochrevector.Record) string {
|
||||
func formatIndexRecord(rec ochrevector.Record) string {
|
||||
rows := [][2]string{
|
||||
{"Index ID", rec.ID},
|
||||
{"Name", rec.Name},
|
||||
@@ -246,7 +246,7 @@ func formatIndexRecord(rec handlers_ochrevector.Record) string {
|
||||
}
|
||||
|
||||
// formatJobRecord renders one ingestion job record as aligned key/value lines.
|
||||
func formatJobRecord(job handlers_ochrevector.JobRecord) string {
|
||||
func formatJobRecord(job ochrevector.JobRecord) string {
|
||||
rows := [][2]string{
|
||||
{"Job ID", job.ID},
|
||||
{"Index ID", job.IndexID},
|
||||
@@ -285,7 +285,7 @@ func truncateForTable(s string, maxRunes int) string {
|
||||
|
||||
// formatQueryResults renders 'ochre vector query' output: a truncated table
|
||||
// by default, or the full response as indented JSON when asJSON is set.
|
||||
func formatQueryResults(results []handlers_ochrevector.QueryResult, asJSON bool) (string, error) {
|
||||
func formatQueryResults(results []ochrevector.QueryResult, asJSON bool) (string, error) {
|
||||
if asJSON {
|
||||
data, err := json.MarshalIndent(results, "", " ")
|
||||
if err != nil {
|
||||
@@ -311,11 +311,11 @@ func formatQueryResults(results []handlers_ochrevector.QueryResult, asJSON bool)
|
||||
// parseFilterFlag decodes --filter's JSON-encoded Filter AST. An empty raw
|
||||
// string is not an error: it means no filter, matching QueryRequest.Filter's
|
||||
// nil-is-unfiltered contract.
|
||||
func parseFilterFlag(raw string) (*handlers_ochrevector.Filter, error) {
|
||||
func parseFilterFlag(raw string) (*ochrevector.Filter, error) {
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var f handlers_ochrevector.Filter
|
||||
var f ochrevector.Filter
|
||||
if err := json.Unmarshal([]byte(raw), &f); err != nil {
|
||||
return nil, fmt.Errorf("decode --filter: %w", err)
|
||||
}
|
||||
@@ -325,8 +325,8 @@ func parseFilterFlag(raw string) (*handlers_ochrevector.Filter, error) {
|
||||
// runIndexCreate is the testable core of 'ochre vector index create': it
|
||||
// mints a fresh index ID (indexes are named, not caller-numbered) and
|
||||
// creates it. Returns the message to print.
|
||||
func runIndexCreate(ctx context.Context, svc handlers_ochrevector.VectorService, name string, dimension int, model string) (string, error) {
|
||||
out, err := svc.CreateIndex(ctx, &handlers_ochrevector.CreateIndexRequest{
|
||||
func runIndexCreate(ctx context.Context, svc ochrevector.VectorService, name string, dimension int, model string) (string, error) {
|
||||
out, err := svc.CreateIndex(ctx, &ochrevector.CreateIndexRequest{
|
||||
IndexID: awsidentifiers.GenerateResourceID("idx"),
|
||||
Name: name,
|
||||
Dimension: dimension,
|
||||
@@ -339,8 +339,8 @@ func runIndexCreate(ctx context.Context, svc handlers_ochrevector.VectorService,
|
||||
}
|
||||
|
||||
// runIndexDelete is the testable core of 'ochre vector index delete'.
|
||||
func runIndexDelete(ctx context.Context, svc handlers_ochrevector.VectorService, indexID string) (string, error) {
|
||||
if _, err := svc.DeleteIndex(ctx, &handlers_ochrevector.DeleteIndexRequest{IndexID: indexID}, awsidentifiers.GlobalAccountID); err != nil {
|
||||
func runIndexDelete(ctx context.Context, svc ochrevector.VectorService, indexID string) (string, error) {
|
||||
if _, err := svc.DeleteIndex(ctx, &ochrevector.DeleteIndexRequest{IndexID: indexID}, awsidentifiers.GlobalAccountID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return fmt.Sprintf("✅ Index %s deleted.", indexID), nil
|
||||
@@ -348,8 +348,8 @@ func runIndexDelete(ctx context.Context, svc handlers_ochrevector.VectorService,
|
||||
|
||||
// listIndexesOutput renders 'ochre vector index list'. Split from its Run
|
||||
// function so it is testable against a fake service with no NATS connection.
|
||||
func listIndexesOutput(ctx context.Context, svc handlers_ochrevector.VectorService) (string, error) {
|
||||
out, err := svc.ListIndexes(ctx, &handlers_ochrevector.ListIndexesRequest{}, awsidentifiers.GlobalAccountID)
|
||||
func listIndexesOutput(ctx context.Context, svc ochrevector.VectorService) (string, error) {
|
||||
out, err := svc.ListIndexes(ctx, &ochrevector.ListIndexesRequest{}, awsidentifiers.GlobalAccountID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -365,11 +365,11 @@ func listIndexesOutput(ctx context.Context, svc handlers_ochrevector.VectorServi
|
||||
|
||||
// runIngest is the testable core of 'ochre vector ingest'. The daemon stamps
|
||||
// the index's own embedding model/dimension onto the source spec server-side
|
||||
// (see handlers_ochrevector.vectorService.Ingest), so none is sent here.
|
||||
func runIngest(ctx context.Context, svc handlers_ochrevector.VectorService, indexID, bucket, prefix string, chunkSize, chunkOverlap int, meta map[string]string) (string, error) {
|
||||
out, err := svc.Ingest(ctx, &handlers_ochrevector.IngestRequest{
|
||||
// (see ochrevector.vectorService.Ingest), so none is sent here.
|
||||
func runIngest(ctx context.Context, svc ochrevector.VectorService, indexID, bucket, prefix string, chunkSize, chunkOverlap int, meta map[string]string) (string, error) {
|
||||
out, err := svc.Ingest(ctx, &ochrevector.IngestRequest{
|
||||
IndexID: indexID,
|
||||
Source: handlers_ochrevector.SourceSpec{
|
||||
Source: ochrevector.SourceSpec{
|
||||
Bucket: bucket,
|
||||
Prefix: prefix,
|
||||
ChunkSize: chunkSize,
|
||||
@@ -384,8 +384,8 @@ func runIngest(ctx context.Context, svc handlers_ochrevector.VectorService, inde
|
||||
}
|
||||
|
||||
// runJobDescribe is the testable core of 'ochre vector job describe'.
|
||||
func runJobDescribe(ctx context.Context, svc handlers_ochrevector.VectorService, jobID string) (string, error) {
|
||||
out, err := svc.DescribeJob(ctx, &handlers_ochrevector.DescribeJobRequest{JobID: jobID}, awsidentifiers.GlobalAccountID)
|
||||
func runJobDescribe(ctx context.Context, svc ochrevector.VectorService, jobID string) (string, error) {
|
||||
out, err := svc.DescribeJob(ctx, &ochrevector.DescribeJobRequest{JobID: jobID}, awsidentifiers.GlobalAccountID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -393,12 +393,12 @@ func runJobDescribe(ctx context.Context, svc handlers_ochrevector.VectorService,
|
||||
}
|
||||
|
||||
// runQuery is the testable core of 'ochre vector query'.
|
||||
func runQuery(ctx context.Context, svc handlers_ochrevector.VectorService, indexID, text string, k int, filterJSON string, asJSON bool) (string, error) {
|
||||
func runQuery(ctx context.Context, svc ochrevector.VectorService, indexID, text string, k int, filterJSON string, asJSON bool) (string, error) {
|
||||
filter, err := parseFilterFlag(filterJSON)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
out, err := svc.Query(ctx, &handlers_ochrevector.QueryRequest{IndexID: indexID, Text: text, K: k, Filter: filter}, awsidentifiers.GlobalAccountID)
|
||||
out, err := svc.Query(ctx, &ochrevector.QueryRequest{IndexID: indexID, Text: text, K: k, Filter: filter}, awsidentifiers.GlobalAccountID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
@@ -10,36 +10,36 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// fakeVectorService serves scripted responses so the CLI's testable cores
|
||||
// can be exercised without a daemon or a live NATS connection.
|
||||
type fakeVectorService struct {
|
||||
createIndex handlers_ochrevector.CreateIndexResponse
|
||||
createIndex ochrevector.CreateIndexResponse
|
||||
createIndexErr error
|
||||
deleteIndexErr error
|
||||
list []handlers_ochrevector.Record
|
||||
list []ochrevector.Record
|
||||
listErr error
|
||||
ingest handlers_ochrevector.IngestResponse
|
||||
ingest ochrevector.IngestResponse
|
||||
ingestErr error
|
||||
describeJob handlers_ochrevector.JobRecord
|
||||
describeJob ochrevector.JobRecord
|
||||
describeErr error
|
||||
query []handlers_ochrevector.QueryResult
|
||||
query []ochrevector.QueryResult
|
||||
queryErr error
|
||||
listJobs []handlers_ochrevector.JobRecord
|
||||
listJobs []ochrevector.JobRecord
|
||||
listJobsErr error
|
||||
stopJob handlers_ochrevector.JobRecord
|
||||
stopJob ochrevector.JobRecord
|
||||
stopJobErr error
|
||||
|
||||
createIndexInputs []*handlers_ochrevector.CreateIndexRequest
|
||||
deleteIndexInputs []*handlers_ochrevector.DeleteIndexRequest
|
||||
ingestInputs []*handlers_ochrevector.IngestRequest
|
||||
queryInputs []*handlers_ochrevector.QueryRequest
|
||||
createIndexInputs []*ochrevector.CreateIndexRequest
|
||||
deleteIndexInputs []*ochrevector.DeleteIndexRequest
|
||||
ingestInputs []*ochrevector.IngestRequest
|
||||
queryInputs []*ochrevector.QueryRequest
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) CreateIndex(_ context.Context, in *handlers_ochrevector.CreateIndexRequest, _ string) (*handlers_ochrevector.CreateIndexResponse, error) {
|
||||
func (f *fakeVectorService) CreateIndex(_ context.Context, in *ochrevector.CreateIndexRequest, _ string) (*ochrevector.CreateIndexResponse, error) {
|
||||
f.createIndexInputs = append(f.createIndexInputs, in)
|
||||
if f.createIndexErr != nil {
|
||||
return nil, f.createIndexErr
|
||||
@@ -51,22 +51,22 @@ func (f *fakeVectorService) CreateIndex(_ context.Context, in *handlers_ochrevec
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) DeleteIndex(_ context.Context, in *handlers_ochrevector.DeleteIndexRequest, _ string) (*handlers_ochrevector.DeleteIndexResponse, error) {
|
||||
func (f *fakeVectorService) DeleteIndex(_ context.Context, in *ochrevector.DeleteIndexRequest, _ string) (*ochrevector.DeleteIndexResponse, error) {
|
||||
f.deleteIndexInputs = append(f.deleteIndexInputs, in)
|
||||
if f.deleteIndexErr != nil {
|
||||
return nil, f.deleteIndexErr
|
||||
}
|
||||
return &handlers_ochrevector.DeleteIndexResponse{}, nil
|
||||
return &ochrevector.DeleteIndexResponse{}, nil
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) ListIndexes(_ context.Context, _ *handlers_ochrevector.ListIndexesRequest, _ string) (*handlers_ochrevector.ListIndexesResponse, error) {
|
||||
func (f *fakeVectorService) ListIndexes(_ context.Context, _ *ochrevector.ListIndexesRequest, _ string) (*ochrevector.ListIndexesResponse, error) {
|
||||
if f.listErr != nil {
|
||||
return nil, f.listErr
|
||||
}
|
||||
return &handlers_ochrevector.ListIndexesResponse{Indexes: f.list}, nil
|
||||
return &ochrevector.ListIndexesResponse{Indexes: f.list}, nil
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) Ingest(_ context.Context, in *handlers_ochrevector.IngestRequest, _ string) (*handlers_ochrevector.IngestResponse, error) {
|
||||
func (f *fakeVectorService) Ingest(_ context.Context, in *ochrevector.IngestRequest, _ string) (*ochrevector.IngestResponse, error) {
|
||||
f.ingestInputs = append(f.ingestInputs, in)
|
||||
if f.ingestErr != nil {
|
||||
return nil, f.ingestErr
|
||||
@@ -74,29 +74,29 @@ func (f *fakeVectorService) Ingest(_ context.Context, in *handlers_ochrevector.I
|
||||
return &f.ingest, nil
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) DescribeJob(_ context.Context, _ *handlers_ochrevector.DescribeJobRequest, _ string) (*handlers_ochrevector.DescribeJobResponse, error) {
|
||||
func (f *fakeVectorService) DescribeJob(_ context.Context, _ *ochrevector.DescribeJobRequest, _ string) (*ochrevector.DescribeJobResponse, error) {
|
||||
if f.describeErr != nil {
|
||||
return nil, f.describeErr
|
||||
}
|
||||
return &handlers_ochrevector.DescribeJobResponse{Job: f.describeJob}, nil
|
||||
return &ochrevector.DescribeJobResponse{Job: f.describeJob}, nil
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) Query(_ context.Context, in *handlers_ochrevector.QueryRequest, _ string) (*handlers_ochrevector.QueryResponse, error) {
|
||||
func (f *fakeVectorService) Query(_ context.Context, in *ochrevector.QueryRequest, _ string) (*ochrevector.QueryResponse, error) {
|
||||
f.queryInputs = append(f.queryInputs, in)
|
||||
if f.queryErr != nil {
|
||||
return nil, f.queryErr
|
||||
}
|
||||
return &handlers_ochrevector.QueryResponse{Results: f.query}, nil
|
||||
return &ochrevector.QueryResponse{Results: f.query}, nil
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) ListJobs(_ context.Context, _ *handlers_ochrevector.ListJobsRequest, _ string) (*handlers_ochrevector.ListJobsResponse, error) {
|
||||
func (f *fakeVectorService) ListJobs(_ context.Context, _ *ochrevector.ListJobsRequest, _ string) (*ochrevector.ListJobsResponse, error) {
|
||||
if f.listJobsErr != nil {
|
||||
return nil, f.listJobsErr
|
||||
}
|
||||
return &handlers_ochrevector.ListJobsResponse{Jobs: f.listJobs}, nil
|
||||
return &ochrevector.ListJobsResponse{Jobs: f.listJobs}, nil
|
||||
}
|
||||
|
||||
func (f *fakeVectorService) StopJob(_ context.Context, req *handlers_ochrevector.StopJobRequest, _ string) (*handlers_ochrevector.StopJobResponse, error) {
|
||||
func (f *fakeVectorService) StopJob(_ context.Context, req *ochrevector.StopJobRequest, _ string) (*ochrevector.StopJobResponse, error) {
|
||||
if f.stopJobErr != nil {
|
||||
return nil, f.stopJobErr
|
||||
}
|
||||
@@ -104,16 +104,16 @@ func (f *fakeVectorService) StopJob(_ context.Context, req *handlers_ochrevector
|
||||
if job.ID == "" {
|
||||
job.ID = req.JobID
|
||||
}
|
||||
return &handlers_ochrevector.StopJobResponse{Job: job}, nil
|
||||
return &ochrevector.StopJobResponse{Job: job}, nil
|
||||
}
|
||||
|
||||
var _ handlers_ochrevector.VectorService = (*fakeVectorService)(nil)
|
||||
var _ ochrevector.VectorService = (*fakeVectorService)(nil)
|
||||
|
||||
const testIndexID = "idx-0123456789abcdef0"
|
||||
|
||||
func TestRunIndexCreate_MintsIDAndPrintsRecord(t *testing.T) {
|
||||
svc := &fakeVectorService{createIndex: handlers_ochrevector.CreateIndexResponse{
|
||||
Index: handlers_ochrevector.Record{Name: "kb1", State: handlers_ochrevector.StateReady, Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5"},
|
||||
svc := &fakeVectorService{createIndex: ochrevector.CreateIndexResponse{
|
||||
Index: ochrevector.Record{Name: "kb1", State: ochrevector.StateReady, Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5"},
|
||||
}}
|
||||
|
||||
msg, err := runIndexCreate(context.Background(), svc, "kb1", 768, "nomic-embed-text-v1.5")
|
||||
@@ -154,8 +154,8 @@ func TestListIndexesOutput_NoIndexes(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestListIndexesOutput_ListsIndexes(t *testing.T) {
|
||||
svc := &fakeVectorService{list: []handlers_ochrevector.Record{
|
||||
{ID: testIndexID, Name: "kb1", State: handlers_ochrevector.StateReady, Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5"},
|
||||
svc := &fakeVectorService{list: []ochrevector.Record{
|
||||
{ID: testIndexID, Name: "kb1", State: ochrevector.StateReady, Dimension: 768, EmbeddingModel: "nomic-embed-text-v1.5"},
|
||||
}}
|
||||
msg, err := listIndexesOutput(context.Background(), svc)
|
||||
require.NoError(t, err)
|
||||
@@ -170,8 +170,8 @@ func TestListIndexesOutput_ErrorSurfaces(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunIngest_NoModelFlagSentByCLI(t *testing.T) {
|
||||
svc := &fakeVectorService{ingest: handlers_ochrevector.IngestResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStatePending},
|
||||
svc := &fakeVectorService{ingest: ochrevector.IngestResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-abc", IndexID: testIndexID, State: ochrevector.JobStatePending},
|
||||
}}
|
||||
|
||||
msg, err := runIngest(context.Background(), svc, testIndexID, "docs", "kb/", 100, 10, map[string]string{"team": "eng"})
|
||||
@@ -198,8 +198,8 @@ func TestRunIngest_ErrorSurfaces(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunJobDescribe_PrintsRecord(t *testing.T) {
|
||||
svc := &fakeVectorService{describeJob: handlers_ochrevector.JobRecord{
|
||||
ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStateReady, DocumentsDone: 3, DocumentsTotal: 3,
|
||||
svc := &fakeVectorService{describeJob: ochrevector.JobRecord{
|
||||
ID: "job-abc", IndexID: testIndexID, State: ochrevector.JobStateReady, DocumentsDone: 3, DocumentsTotal: 3,
|
||||
}}
|
||||
out, err := runJobDescribe(context.Background(), svc, "job-abc")
|
||||
require.NoError(t, err)
|
||||
@@ -215,7 +215,7 @@ func TestRunJobDescribe_ErrorSurfaces(t *testing.T) {
|
||||
|
||||
func TestRunQuery_TableTruncatesLongChunk(t *testing.T) {
|
||||
longChunk := strings.Repeat("x", cliChunkPreviewChars+50)
|
||||
svc := &fakeVectorService{query: []handlers_ochrevector.QueryResult{
|
||||
svc := &fakeVectorService{query: []ochrevector.QueryResult{
|
||||
{Chunk: longChunk, SourceKey: "docs/a.txt", SourceOffset: 0, Score: 0.9},
|
||||
}}
|
||||
|
||||
@@ -232,7 +232,7 @@ func TestRunQuery_TableTruncatesLongChunk(t *testing.T) {
|
||||
|
||||
func TestRunQuery_JSONFlagPrintsFullChunk(t *testing.T) {
|
||||
longChunk := strings.Repeat("x", cliChunkPreviewChars+50)
|
||||
svc := &fakeVectorService{query: []handlers_ochrevector.QueryResult{
|
||||
svc := &fakeVectorService{query: []ochrevector.QueryResult{
|
||||
{Chunk: longChunk, SourceKey: "docs/a.txt", Score: 0.9},
|
||||
}}
|
||||
|
||||
@@ -255,7 +255,7 @@ func TestRunQuery_FilterFlagReachesRequest(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Len(t, svc.queryInputs, 1)
|
||||
require.NotNil(t, svc.queryInputs[0].Filter)
|
||||
require.Equal(t, handlers_ochrevector.FilterEquals, svc.queryInputs[0].Filter.Op)
|
||||
require.Equal(t, ochrevector.FilterEquals, svc.queryInputs[0].Filter.Op)
|
||||
require.Equal(t, "category", svc.queryInputs[0].Filter.Key)
|
||||
}
|
||||
|
||||
@@ -272,11 +272,11 @@ func TestRunQuery_ErrorSurfaces(t *testing.T) {
|
||||
|
||||
// withVectorService swaps vectorServiceFn for one returning svc, so a Run
|
||||
// wrapper's real connect/validate/exit control flow runs without a daemon.
|
||||
func withVectorService(t *testing.T, svc handlers_ochrevector.VectorService, connErr error) {
|
||||
func withVectorService(t *testing.T, svc ochrevector.VectorService, connErr error) {
|
||||
t.Helper()
|
||||
orig := vectorServiceFn
|
||||
t.Cleanup(func() { vectorServiceFn = orig })
|
||||
vectorServiceFn = func() (handlers_ochrevector.VectorService, func(), error) {
|
||||
vectorServiceFn = func() (ochrevector.VectorService, func(), error) {
|
||||
if connErr != nil {
|
||||
return nil, nil, connErr
|
||||
}
|
||||
@@ -285,8 +285,8 @@ func withVectorService(t *testing.T, svc handlers_ochrevector.VectorService, con
|
||||
}
|
||||
|
||||
func TestRunOchreVectorIndexCreate_PrintsRecord(t *testing.T) {
|
||||
withVectorService(t, &fakeVectorService{createIndex: handlers_ochrevector.CreateIndexResponse{
|
||||
Index: handlers_ochrevector.Record{ID: testIndexID, State: handlers_ochrevector.StateReady},
|
||||
withVectorService(t, &fakeVectorService{createIndex: ochrevector.CreateIndexResponse{
|
||||
Index: ochrevector.Record{ID: testIndexID, State: ochrevector.StateReady},
|
||||
}}, nil)
|
||||
|
||||
cmd := *ochreVectorIndexCreateCmd
|
||||
@@ -326,8 +326,8 @@ func TestRunOchreVectorIndexDelete_ReportsDeletion(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunOchreVectorIndexList_PrintsTable(t *testing.T) {
|
||||
withVectorService(t, &fakeVectorService{list: []handlers_ochrevector.Record{
|
||||
{ID: testIndexID, Name: "kb1", State: handlers_ochrevector.StateReady},
|
||||
withVectorService(t, &fakeVectorService{list: []ochrevector.Record{
|
||||
{ID: testIndexID, Name: "kb1", State: ochrevector.StateReady},
|
||||
}}, nil)
|
||||
|
||||
var out string
|
||||
@@ -339,8 +339,8 @@ func TestRunOchreVectorIndexList_PrintsTable(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunOchreVectorIngest_PrintsJobID(t *testing.T) {
|
||||
withVectorService(t, &fakeVectorService{ingest: handlers_ochrevector.IngestResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStatePending},
|
||||
withVectorService(t, &fakeVectorService{ingest: ochrevector.IngestResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-abc", IndexID: testIndexID, State: ochrevector.JobStatePending},
|
||||
}}, nil)
|
||||
|
||||
cmd := *ochreVectorIngestCmd
|
||||
@@ -356,8 +356,8 @@ func TestRunOchreVectorIngest_PrintsJobID(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunOchreVectorJobDescribe_PrintsRecord(t *testing.T) {
|
||||
withVectorService(t, &fakeVectorService{describeJob: handlers_ochrevector.JobRecord{
|
||||
ID: "job-abc", IndexID: testIndexID, State: handlers_ochrevector.JobStateReady,
|
||||
withVectorService(t, &fakeVectorService{describeJob: ochrevector.JobRecord{
|
||||
ID: "job-abc", IndexID: testIndexID, State: ochrevector.JobStateReady,
|
||||
}}, nil)
|
||||
|
||||
var out string
|
||||
@@ -369,7 +369,7 @@ func TestRunOchreVectorJobDescribe_PrintsRecord(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunOchreVectorQuery_PrintsTable(t *testing.T) {
|
||||
withVectorService(t, &fakeVectorService{query: []handlers_ochrevector.QueryResult{
|
||||
withVectorService(t, &fakeVectorService{query: []ochrevector.QueryResult{
|
||||
{Chunk: "hello world", SourceKey: "docs/a.txt", Score: 0.9},
|
||||
}}, nil)
|
||||
|
||||
@@ -399,9 +399,9 @@ func TestRunOchreVectorQuery_ConnectFailureExits1(t *testing.T) {
|
||||
const testBackupAccountID = "123456789012"
|
||||
|
||||
func TestRunBackupAccount_PrintsObjectKeyAndSize(t *testing.T) {
|
||||
backup := func(_ context.Context, accountID string) (*handlers_ochrevector.BackupAccountResponse, error) {
|
||||
backup := func(_ context.Context, accountID string) (*ochrevector.BackupAccountResponse, error) {
|
||||
require.Equal(t, testBackupAccountID, accountID)
|
||||
return &handlers_ochrevector.BackupAccountResponse{ObjectKey: testBackupAccountID + "/backup.sql.gz", SizeBytes: 4096}, nil
|
||||
return &ochrevector.BackupAccountResponse{ObjectKey: testBackupAccountID + "/backup.sql.gz", SizeBytes: 4096}, nil
|
||||
}
|
||||
|
||||
msg, err := runBackupAccount(context.Background(), testBackupAccountID, backup)
|
||||
@@ -411,7 +411,7 @@ func TestRunBackupAccount_PrintsObjectKeyAndSize(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunBackupAccount_ErrorSurfaces(t *testing.T) {
|
||||
backup := func(context.Context, string) (*handlers_ochrevector.BackupAccountResponse, error) {
|
||||
backup := func(context.Context, string) (*ochrevector.BackupAccountResponse, error) {
|
||||
return nil, errors.New("ochrevector: appliance not available")
|
||||
}
|
||||
|
||||
@@ -445,7 +445,7 @@ func TestRunRestoreAccount_ErrorSurfaces(t *testing.T) {
|
||||
// withBackupAccount swaps backupAccountFn for one that never dials NATS, so
|
||||
// the Run wrapper's flag/exit control flow is exercised without a live
|
||||
// daemon, mirroring withApplianceTeardown.
|
||||
func withBackupAccount(t *testing.T, fn func(context.Context, string) (*handlers_ochrevector.BackupAccountResponse, error)) {
|
||||
func withBackupAccount(t *testing.T, fn func(context.Context, string) (*ochrevector.BackupAccountResponse, error)) {
|
||||
t.Helper()
|
||||
orig := backupAccountFn
|
||||
t.Cleanup(func() { backupAccountFn = orig })
|
||||
@@ -461,8 +461,8 @@ func withRestoreAccount(t *testing.T, fn func(context.Context, string, string) e
|
||||
}
|
||||
|
||||
func TestRunOchreVectorBackup_PrintsSuccess(t *testing.T) {
|
||||
withBackupAccount(t, func(context.Context, string) (*handlers_ochrevector.BackupAccountResponse, error) {
|
||||
return &handlers_ochrevector.BackupAccountResponse{ObjectKey: testBackupAccountID + "/backup.sql.gz", SizeBytes: 1024}, nil
|
||||
withBackupAccount(t, func(context.Context, string) (*ochrevector.BackupAccountResponse, error) {
|
||||
return &ochrevector.BackupAccountResponse{ObjectKey: testBackupAccountID + "/backup.sql.gz", SizeBytes: 1024}, nil
|
||||
})
|
||||
|
||||
cmd := *ochreVectorBackupCmd
|
||||
@@ -477,7 +477,7 @@ func TestRunOchreVectorBackup_PrintsSuccess(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunOchreVectorBackup_ErrorExits1(t *testing.T) {
|
||||
withBackupAccount(t, func(context.Context, string) (*handlers_ochrevector.BackupAccountResponse, error) {
|
||||
withBackupAccount(t, func(context.Context, string) (*ochrevector.BackupAccountResponse, error) {
|
||||
return nil, errors.New("ochrevector: appliance not available")
|
||||
})
|
||||
|
||||
|
||||
@@ -59,7 +59,7 @@ services:
|
||||
- spinifex/domains/acm
|
||||
- spinifex/handlers/iam
|
||||
- spinifex/handlers/rds
|
||||
- spinifex/handlers/ochrevector
|
||||
- spinifex/domains/ochre/vector
|
||||
subscribes:
|
||||
# EC2 instance lifecycle.
|
||||
- "ec2.RunInstances"
|
||||
@@ -136,7 +136,7 @@ services:
|
||||
# {account}.{db-instance}, so the subject is built at runtime and only the
|
||||
# prefix is declarable.
|
||||
- "rds.bus.*"
|
||||
# Ochre vector store tenant surface (handlers/ochrevector/api.go),
|
||||
# Ochre vector store tenant surface (domains/ochre/vector/api.go),
|
||||
# gated behind config.OchreVectorConfig.Enabled (default off).
|
||||
- "ochre.vector.createIndex"
|
||||
- "ochre.vector.deleteIndex"
|
||||
|
||||
@@ -45,6 +45,7 @@ import (
|
||||
"github.com/mulgadc/spinifex/spinifex/domains/ec2/instancetypes"
|
||||
handlers_ecr "github.com/mulgadc/spinifex/spinifex/domains/ecr"
|
||||
"github.com/mulgadc/spinifex/spinifex/domains/ochre"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/state/clustersize"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/state/kvutil"
|
||||
@@ -69,7 +70,6 @@ import (
|
||||
handlers_eks "github.com/mulgadc/spinifex/spinifex/handlers/eks"
|
||||
handlers_elbv2 "github.com/mulgadc/spinifex/spinifex/handlers/elbv2"
|
||||
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
handlers_rds "github.com/mulgadc/spinifex/spinifex/handlers/rds"
|
||||
"github.com/mulgadc/spinifex/spinifex/network/external"
|
||||
"github.com/mulgadc/spinifex/spinifex/network/external/dhcp"
|
||||
@@ -173,9 +173,9 @@ type Daemon struct {
|
||||
bedrockReaper *ochre.Reaper
|
||||
acmService *acmdomain.ACMServiceImpl
|
||||
acmRenewalWorker *acmdomain.Worker
|
||||
ochreVectorService handlers_ochrevector.VectorService
|
||||
ochreAppliance *handlers_ochrevector.Appliance
|
||||
ochreBackupService *handlers_ochrevector.BackupService
|
||||
ochreVectorService ochrevector.VectorService
|
||||
ochreAppliance *ochrevector.Appliance
|
||||
ochreBackupService *ochrevector.BackupService
|
||||
ecrMetaService *handlers_ecr.MetaServiceImpl
|
||||
routeTableService *handlers_ec2_routetable.RouteTableServiceImpl
|
||||
natGatewayService *handlers_ec2_natgw.NatGatewayServiceImpl
|
||||
|
||||
@@ -11,9 +11,9 @@ import (
|
||||
|
||||
"github.com/mulgadc/bluebottle/pkg/masterkey"
|
||||
"github.com/mulgadc/spinifex/spinifex/domains/ochre"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/telemetry"
|
||||
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
"github.com/mulgadc/spinifex/spinifex/network/host"
|
||||
"github.com/mulgadc/spinifex/spinifex/providers/objectstore"
|
||||
"github.com/nats-io/nats.go/jetstream"
|
||||
@@ -66,8 +66,8 @@ func (d *Daemon) startOchreVector() {
|
||||
return
|
||||
}
|
||||
|
||||
launcher := handlers_ochrevector.NewRDSApplianceLauncher(d.natsConn, ochreApplianceLaunchTimeout)
|
||||
appliance, err := handlers_ochrevector.NewAppliance(js, masterKey, launcher)
|
||||
launcher := ochrevector.NewRDSApplianceLauncher(d.natsConn, ochreApplianceLaunchTimeout)
|
||||
appliance, err := ochrevector.NewAppliance(js, masterKey, launcher)
|
||||
if err != nil {
|
||||
slog.Warn("Ochre vector store disabled: appliance construction failed", "err", err)
|
||||
return
|
||||
@@ -75,7 +75,7 @@ func (d *Daemon) startOchreVector() {
|
||||
// Give the daemon a routed presence in the appliance's system-VPC subnet
|
||||
// before Connect dials it: a tag-filtered ENI describe resolves the real
|
||||
// dial IP and subnet, never the unroutable vanity endpoint hostname.
|
||||
appliance.WithHostPort(handlers_ochrevector.HostPortDeps{
|
||||
appliance.WithHostPort(ochrevector.HostPortDeps{
|
||||
VPC: d.vpcService,
|
||||
HostPort: host.NewOVSPlumber(),
|
||||
NodeID: d.node,
|
||||
@@ -88,10 +88,10 @@ func (d *Daemon) startOchreVector() {
|
||||
// reconstructed; kb/ds are gateway-owned metadata this daemon never reads
|
||||
// or writes itself -- they exist here solely so Teardown can purge them
|
||||
// alongside registry/jobs.
|
||||
registry := handlers_ochrevector.NewRegistry(js)
|
||||
jobs := handlers_ochrevector.NewJobStore(js)
|
||||
registry := ochrevector.NewRegistry(js)
|
||||
jobs := ochrevector.NewJobStore(js)
|
||||
appliance.WithStores(registry, jobs)
|
||||
appliance.WithKBStores(handlers_ochrevector.NewKBStore(js), handlers_ochrevector.NewDataSourceStore(js))
|
||||
appliance.WithKBStores(ochrevector.NewKBStore(js), ochrevector.NewDataSourceStore(js))
|
||||
|
||||
// Publish the appliance and register the operator teardown subject BEFORE
|
||||
// Connect: recovering a broken singleton is teardown's whole purpose, so it
|
||||
@@ -101,7 +101,7 @@ func (d *Daemon) startOchreVector() {
|
||||
d.ochreAppliance = appliance
|
||||
d.mu.Unlock()
|
||||
if err := d.registerNatsSubs([]natsSub{
|
||||
{handlers_ochrevector.SubjectTeardownAppliance, handleNATSRequest(d.node, d.handleOchreApplianceTeardown), "spinifex-workers"},
|
||||
{ochrevector.SubjectTeardownAppliance, handleNATSRequest(d.node, d.handleOchreApplianceTeardown), "spinifex-workers"},
|
||||
}); err != nil {
|
||||
slog.Error("Ochre vector store: failed to register appliance teardown subject", "err", err)
|
||||
}
|
||||
@@ -122,7 +122,7 @@ func (d *Daemon) startOchreVector() {
|
||||
|
||||
// RerankModel unset leaves reranker nil: Query degrades to plain KNN
|
||||
// rather than resolving a default rerank model nobody asked for.
|
||||
var reranker handlers_ochrevector.Reranker
|
||||
var reranker ochrevector.Reranker
|
||||
if cfg.RerankModel != "" {
|
||||
reranker = gateway_bedrock.NewReranker(endpointResolver, cfg.RerankModel)
|
||||
}
|
||||
@@ -132,10 +132,10 @@ func (d *Daemon) startOchreVector() {
|
||||
|
||||
// registry and jobs were already built above (and attached to appliance
|
||||
// via WithStores) before Connect was attempted; reused here, not rebuilt.
|
||||
service := handlers_ochrevector.NewService(registry, backend)
|
||||
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, store, embedder)
|
||||
service := ochrevector.NewService(registry, backend)
|
||||
ingest := ochrevector.NewIngestService(jobs, registry, backend, store, embedder)
|
||||
|
||||
vectorService := handlers_ochrevector.NewVectorService(service, ingest, jobs, registry, backend, embedder, reranker)
|
||||
vectorService := ochrevector.NewVectorService(service, ingest, jobs, registry, backend, embedder, reranker)
|
||||
|
||||
// A shutdown that lands in the gap between Connect succeeding above and
|
||||
// this check does not leave anything to unwind: nothing has been
|
||||
@@ -154,28 +154,28 @@ func (d *Daemon) startOchreVector() {
|
||||
// table-driven mechanism subscribeAll itself uses, so a queue-group
|
||||
// registration here is indistinguishable from one made at boot.
|
||||
subs := []natsSub{
|
||||
{handlers_ochrevector.SubjectCreateIndex, handleNATSRequest(d.node, vectorService.CreateIndex), "spinifex-workers"},
|
||||
{handlers_ochrevector.SubjectDeleteIndex, handleNATSRequest(d.node, vectorService.DeleteIndex), "spinifex-workers"},
|
||||
{handlers_ochrevector.SubjectListIndexes, handleNATSRequest(d.node, vectorService.ListIndexes), "spinifex-workers"},
|
||||
{handlers_ochrevector.SubjectIngest, handleNATSRequest(d.node, vectorService.Ingest), "spinifex-workers"},
|
||||
{handlers_ochrevector.SubjectDescribeJob, handleNATSRequest(d.node, vectorService.DescribeJob), "spinifex-workers"},
|
||||
{handlers_ochrevector.SubjectQuery, handleNATSRequest(d.node, vectorService.Query), "spinifex-workers"},
|
||||
{handlers_ochrevector.SubjectListJobs, handleNATSRequest(d.node, vectorService.ListJobs), "spinifex-workers"},
|
||||
{handlers_ochrevector.SubjectStopJob, handleNATSRequest(d.node, vectorService.StopJob), "spinifex-workers"},
|
||||
{ochrevector.SubjectCreateIndex, handleNATSRequest(d.node, vectorService.CreateIndex), "spinifex-workers"},
|
||||
{ochrevector.SubjectDeleteIndex, handleNATSRequest(d.node, vectorService.DeleteIndex), "spinifex-workers"},
|
||||
{ochrevector.SubjectListIndexes, handleNATSRequest(d.node, vectorService.ListIndexes), "spinifex-workers"},
|
||||
{ochrevector.SubjectIngest, handleNATSRequest(d.node, vectorService.Ingest), "spinifex-workers"},
|
||||
{ochrevector.SubjectDescribeJob, handleNATSRequest(d.node, vectorService.DescribeJob), "spinifex-workers"},
|
||||
{ochrevector.SubjectQuery, handleNATSRequest(d.node, vectorService.Query), "spinifex-workers"},
|
||||
{ochrevector.SubjectListJobs, handleNATSRequest(d.node, vectorService.ListJobs), "spinifex-workers"},
|
||||
{ochrevector.SubjectStopJob, handleNATSRequest(d.node, vectorService.StopJob), "spinifex-workers"},
|
||||
}
|
||||
|
||||
// backup/restore need RegrantAccount alongside EnsureAccount, which is
|
||||
// not part of VectorBackend itself; every real backend (pgxBackend)
|
||||
// implements it, so this only ever skips registration for a future
|
||||
// VectorBackend that does not.
|
||||
if granter, ok := backend.(handlers_ochrevector.AccountGranter); ok {
|
||||
backupSvc := handlers_ochrevector.NewBackupService(appliance, granter, store, handlers_ochrevector.ExecPgDumper{})
|
||||
if granter, ok := backend.(ochrevector.AccountGranter); ok {
|
||||
backupSvc := ochrevector.NewBackupService(appliance, granter, store, ochrevector.ExecPgDumper{})
|
||||
d.mu.Lock()
|
||||
d.ochreBackupService = backupSvc
|
||||
d.mu.Unlock()
|
||||
subs = append(subs,
|
||||
natsSub{handlers_ochrevector.SubjectBackupAccount, handleNATSRequest(d.node, backupSvc.Backup), "spinifex-workers"},
|
||||
natsSub{handlers_ochrevector.SubjectRestoreAccount, handleNATSRequest(d.node, backupSvc.Restore), "spinifex-workers"},
|
||||
natsSub{ochrevector.SubjectBackupAccount, handleNATSRequest(d.node, backupSvc.Backup), "spinifex-workers"},
|
||||
natsSub{ochrevector.SubjectRestoreAccount, handleNATSRequest(d.node, backupSvc.Restore), "spinifex-workers"},
|
||||
)
|
||||
} else {
|
||||
slog.Warn("Ochre vector store: backend does not support account backup/restore; subjects not registered")
|
||||
@@ -208,7 +208,7 @@ func (d *Daemon) startOchreVector() {
|
||||
// EnsureAccount/CreateIndex call and no ingest job for it. Best-effort: a
|
||||
// failure on one record is logged and the pass continues to the rest rather
|
||||
// than failing the appliance connect that already succeeded.
|
||||
func (d *Daemon) reconcileOchreIndexes(ctx context.Context, registry *handlers_ochrevector.Registry, backend handlers_ochrevector.VectorBackend, ingest *handlers_ochrevector.IngestService) {
|
||||
func (d *Daemon) reconcileOchreIndexes(ctx context.Context, registry *ochrevector.Registry, backend ochrevector.VectorBackend, ingest *ochrevector.IngestService) {
|
||||
recs, err := registry.ListAll(ctx)
|
||||
if err != nil {
|
||||
slog.Warn("Ochre vector store: reconcile indexes: list registry failed", "err", err)
|
||||
@@ -229,7 +229,7 @@ func (d *Daemon) reconcileOchreIndexes(ctx context.Context, registry *handlers_o
|
||||
"account", rec.AccountID, "index", rec.ID, "err", err)
|
||||
continue
|
||||
}
|
||||
spec := handlers_ochrevector.IndexSpec{ID: rec.ID, Dimension: rec.Dimension}
|
||||
spec := ochrevector.IndexSpec{ID: rec.ID, Dimension: rec.Dimension}
|
||||
if err := backend.CreateIndex(ctx, rec.AccountID, spec); err != nil {
|
||||
slog.Warn("Ochre vector store: reconcile indexes: create index failed",
|
||||
"account", rec.AccountID, "index", rec.ID, "err", err)
|
||||
@@ -253,7 +253,7 @@ const ochreIngestSweepInterval = 15 * time.Second
|
||||
// runOchreIngestScheduler drives the ingestion sweep on a timer until the
|
||||
// daemon context is cancelled. An initial sweep keeps submit latency low; the
|
||||
// ticker then re-drives PENDING and crash-abandoned RUNNING jobs.
|
||||
func (d *Daemon) runOchreIngestScheduler(ingest *handlers_ochrevector.IngestService) {
|
||||
func (d *Daemon) runOchreIngestScheduler(ingest *ochrevector.IngestService) {
|
||||
if err := ingest.Sweep(d.ctx); err != nil {
|
||||
slog.Warn("Ochre vector store: initial ingest sweep", "err", err)
|
||||
}
|
||||
@@ -278,7 +278,7 @@ func (d *Daemon) runOchreIngestScheduler(ingest *handlers_ochrevector.IngestServ
|
||||
// an earlier call) has nothing to act on, which is reported as an error
|
||||
// rather than silently accepted -- an operator asking to tear down expects
|
||||
// one to have existed.
|
||||
func (d *Daemon) handleOchreApplianceTeardown(ctx context.Context, req *handlers_ochrevector.TeardownApplianceRequest, _ string) (*handlers_ochrevector.TeardownApplianceResponse, error) {
|
||||
func (d *Daemon) handleOchreApplianceTeardown(ctx context.Context, req *ochrevector.TeardownApplianceRequest, _ string) (*ochrevector.TeardownApplianceResponse, error) {
|
||||
d.mu.Lock()
|
||||
appliance := d.ochreAppliance
|
||||
d.mu.Unlock()
|
||||
@@ -297,13 +297,13 @@ func (d *Daemon) handleOchreApplianceTeardown(ctx context.Context, req *handlers
|
||||
d.ochreBackupService = nil
|
||||
d.mu.Unlock()
|
||||
|
||||
return &handlers_ochrevector.TeardownApplianceResponse{}, nil
|
||||
return &ochrevector.TeardownApplianceResponse{}, nil
|
||||
}
|
||||
|
||||
// connectOchreAppliance drives Ensure-then-Connect until it succeeds or the
|
||||
// daemon shuts down, so a torn-down-then-readopted appliance heals on its own
|
||||
// rather than disabling the feature permanently after a few early failures.
|
||||
func (d *Daemon) connectOchreAppliance(appliance *handlers_ochrevector.Appliance) (handlers_ochrevector.VectorBackend, error) {
|
||||
func (d *Daemon) connectOchreAppliance(appliance *ochrevector.Appliance) (ochrevector.VectorBackend, error) {
|
||||
return retryUntilContext(d.ctx, ochreStartupInitialBackoff, ochreStartupMaxBackoff,
|
||||
func(attempt int, backoff time.Duration, err error) {
|
||||
if attempt <= ochreStartupLogAttempts || attempt%ochreStartupLogEvery == 0 {
|
||||
@@ -311,7 +311,7 @@ func (d *Daemon) connectOchreAppliance(appliance *handlers_ochrevector.Appliance
|
||||
"attempt", attempt, "backoff_ms", otelsetup.Millis(backoff), "err", err)
|
||||
}
|
||||
},
|
||||
func() (handlers_ochrevector.VectorBackend, error) {
|
||||
func() (ochrevector.VectorBackend, error) {
|
||||
return d.ensureAndConnectOchreApplianceOnce(appliance)
|
||||
})
|
||||
}
|
||||
@@ -349,7 +349,7 @@ func retryUntilContext[T any](ctx context.Context, initialBackoff, maxBackoff ti
|
||||
// attempt, bounding Ensure by ochreApplianceLaunchTimeout the same way the
|
||||
// original one-shot call did; Connect is bounded by d.ctx alone, since it
|
||||
// does its own network dial rather than a long poll loop.
|
||||
func (d *Daemon) ensureAndConnectOchreApplianceOnce(appliance *handlers_ochrevector.Appliance) (handlers_ochrevector.VectorBackend, error) {
|
||||
func (d *Daemon) ensureAndConnectOchreApplianceOnce(appliance *ochrevector.Appliance) (ochrevector.VectorBackend, error) {
|
||||
ensureCtx, cancel := context.WithTimeout(d.ctx, ochreApplianceLaunchTimeout)
|
||||
defer cancel()
|
||||
if _, err := appliance.Ensure(ensureCtx); err != nil {
|
||||
|
||||
@@ -13,12 +13,12 @@ import (
|
||||
|
||||
"github.com/mulgadc/spinifex/internal/testkit"
|
||||
"github.com/mulgadc/spinifex/spinifex/bootstrap/config"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// reconcileFakeBackend is a minimal handlers_ochrevector.VectorBackend
|
||||
// reconcileFakeBackend is a minimal ochrevector.VectorBackend
|
||||
// double for reconcileOchreIndexes tests: only IndexExists/EnsureAccount/
|
||||
// CreateIndex are exercised, so every other method is an unused stub.
|
||||
type reconcileFakeBackend struct {
|
||||
@@ -26,7 +26,7 @@ type reconcileFakeBackend struct {
|
||||
indexes map[string]bool // key: accountID+"/"+indexID
|
||||
|
||||
ensureAccountCalls []string
|
||||
createIndexCalls []handlers_ochrevector.IndexSpec
|
||||
createIndexCalls []ochrevector.IndexSpec
|
||||
|
||||
// Error-injection fields for the reconcile error-branch tests below: each
|
||||
// makes the matching backend method fail on every call, so a test that
|
||||
@@ -38,10 +38,10 @@ type reconcileFakeBackend struct {
|
||||
// onCreateIndex runs after a successful CreateIndex, letting a test mutate
|
||||
// shared state (e.g. delete the registry record) without any production
|
||||
// code change, so a later ReconcileFromSource call in the same pass fails.
|
||||
onCreateIndex func(ctx context.Context, accountID string, spec handlers_ochrevector.IndexSpec)
|
||||
onCreateIndex func(ctx context.Context, accountID string, spec ochrevector.IndexSpec)
|
||||
}
|
||||
|
||||
var _ handlers_ochrevector.VectorBackend = (*reconcileFakeBackend)(nil)
|
||||
var _ ochrevector.VectorBackend = (*reconcileFakeBackend)(nil)
|
||||
|
||||
func newReconcileFakeBackend() *reconcileFakeBackend {
|
||||
return &reconcileFakeBackend{indexes: map[string]bool{}}
|
||||
@@ -57,7 +57,7 @@ func (f *reconcileFakeBackend) EnsureAccount(_ context.Context, accountID string
|
||||
return err
|
||||
}
|
||||
|
||||
func (f *reconcileFakeBackend) CreateIndex(ctx context.Context, accountID string, spec handlers_ochrevector.IndexSpec) error {
|
||||
func (f *reconcileFakeBackend) CreateIndex(ctx context.Context, accountID string, spec ochrevector.IndexSpec) error {
|
||||
f.mu.Lock()
|
||||
f.createIndexCalls = append(f.createIndexCalls, spec)
|
||||
err := f.createIndexErr
|
||||
@@ -91,11 +91,11 @@ func (f *reconcileFakeBackend) DropIndex(_ context.Context, accountID, indexID s
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *reconcileFakeBackend) ReplaceDocument(_ context.Context, _, _, _ string, _ []handlers_ochrevector.VectorRow) error {
|
||||
func (f *reconcileFakeBackend) ReplaceDocument(_ context.Context, _, _, _ string, _ []ochrevector.VectorRow) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *reconcileFakeBackend) Query(_ context.Context, _, _ string, _ []float32, _ int, _ *handlers_ochrevector.Filter) ([]handlers_ochrevector.QueryResult, error) {
|
||||
func (f *reconcileFakeBackend) Query(_ context.Context, _, _ string, _ []float32, _ int, _ *ochrevector.Filter) ([]ochrevector.QueryResult, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -173,7 +173,7 @@ func TestRetryUntilContext_StopsOnContextCancel(t *testing.T) {
|
||||
func TestHandleOchreApplianceTeardown_NilApplianceRefuses(t *testing.T) {
|
||||
d := &Daemon{}
|
||||
|
||||
out, err := d.handleOchreApplianceTeardown(context.Background(), &handlers_ochrevector.TeardownApplianceRequest{}, "")
|
||||
out, err := d.handleOchreApplianceTeardown(context.Background(), &ochrevector.TeardownApplianceRequest{}, "")
|
||||
|
||||
require.Error(t, err)
|
||||
assert.Nil(t, out)
|
||||
@@ -188,17 +188,17 @@ const reconcileTestAccount = "111111111111"
|
||||
// one PENDING ingest job per persisted SourceSpec.
|
||||
func TestReconcileOchreIndexes_RecreatesMissingIndexAndEnqueuesJob(t *testing.T) {
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
registry := handlers_ochrevector.NewRegistry(js)
|
||||
jobs := handlers_ochrevector.NewJobStore(js)
|
||||
registry := ochrevector.NewRegistry(js)
|
||||
jobs := ochrevector.NewJobStore(js)
|
||||
backend := newReconcileFakeBackend()
|
||||
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
ingest := ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now().UTC()
|
||||
source := handlers_ochrevector.SourceSpec{Bucket: "docs", Prefix: "kb/", ChunkSize: 100, ChunkOverlap: 10, EmbeddingModel: "stub-embed", Dimension: 2}
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, handlers_ochrevector.Record{
|
||||
ID: "idx-one", Dimension: 2, State: handlers_ochrevector.StateReady,
|
||||
SourceSpecs: []handlers_ochrevector.SourceSpec{source}, CreatedAt: now, UpdatedAt: now,
|
||||
source := ochrevector.SourceSpec{Bucket: "docs", Prefix: "kb/", ChunkSize: 100, ChunkOverlap: 10, EmbeddingModel: "stub-embed", Dimension: 2}
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, ochrevector.Record{
|
||||
ID: "idx-one", Dimension: 2, State: ochrevector.StateReady,
|
||||
SourceSpecs: []ochrevector.SourceSpec{source}, CreatedAt: now, UpdatedAt: now,
|
||||
}))
|
||||
|
||||
d := &Daemon{}
|
||||
@@ -211,7 +211,7 @@ func TestReconcileOchreIndexes_RecreatesMissingIndexAndEnqueuesJob(t *testing.T)
|
||||
jobRecs, err := jobs.ListAll(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, jobRecs, 1, "must enqueue exactly one job for the one persisted SourceSpec")
|
||||
assert.Equal(t, handlers_ochrevector.JobStatePending, jobRecs[0].State)
|
||||
assert.Equal(t, ochrevector.JobStatePending, jobRecs[0].State)
|
||||
assert.Equal(t, source, jobRecs[0].Source)
|
||||
}
|
||||
|
||||
@@ -220,20 +220,20 @@ func TestReconcileOchreIndexes_RecreatesMissingIndexAndEnqueuesJob(t *testing.T)
|
||||
// EnsureAccount/CreateIndex call and no ingest job enqueued for it.
|
||||
func TestReconcileOchreIndexes_ExistingIndexIsNoOp(t *testing.T) {
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
registry := handlers_ochrevector.NewRegistry(js)
|
||||
jobs := handlers_ochrevector.NewJobStore(js)
|
||||
registry := ochrevector.NewRegistry(js)
|
||||
jobs := ochrevector.NewJobStore(js)
|
||||
backend := newReconcileFakeBackend()
|
||||
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
ingest := ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now().UTC()
|
||||
source := handlers_ochrevector.SourceSpec{Bucket: "docs", Prefix: "kb/", Dimension: 2}
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, handlers_ochrevector.Record{
|
||||
ID: "idx-one", Dimension: 2, State: handlers_ochrevector.StateReady,
|
||||
SourceSpecs: []handlers_ochrevector.SourceSpec{source}, CreatedAt: now, UpdatedAt: now,
|
||||
source := ochrevector.SourceSpec{Bucket: "docs", Prefix: "kb/", Dimension: 2}
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, ochrevector.Record{
|
||||
ID: "idx-one", Dimension: 2, State: ochrevector.StateReady,
|
||||
SourceSpecs: []ochrevector.SourceSpec{source}, CreatedAt: now, UpdatedAt: now,
|
||||
}))
|
||||
// Pre-seed the backend as if the table already survived a restart.
|
||||
require.NoError(t, backend.CreateIndex(ctx, reconcileTestAccount, handlers_ochrevector.IndexSpec{ID: "idx-one", Dimension: 2}))
|
||||
require.NoError(t, backend.CreateIndex(ctx, reconcileTestAccount, ochrevector.IndexSpec{ID: "idx-one", Dimension: 2}))
|
||||
backend.mu.Lock()
|
||||
backend.createIndexCalls = nil // reset so the assertion below is reconcile's own calls, not the seed
|
||||
backend.mu.Unlock()
|
||||
@@ -254,22 +254,22 @@ func TestReconcileOchreIndexes_ExistingIndexIsNoOp(t *testing.T) {
|
||||
// logs and skips its own record without panicking, rather than failing the
|
||||
// whole reconcile pass.
|
||||
func TestReconcileOchreIndexes_ErrorsAreLoggedAndSkipped(t *testing.T) {
|
||||
newRecord := func(id string) handlers_ochrevector.Record {
|
||||
newRecord := func(id string) ochrevector.Record {
|
||||
now := time.Now().UTC()
|
||||
return handlers_ochrevector.Record{
|
||||
ID: id, Dimension: 2, State: handlers_ochrevector.StateReady,
|
||||
SourceSpecs: []handlers_ochrevector.SourceSpec{{Bucket: "docs", Dimension: 2}},
|
||||
return ochrevector.Record{
|
||||
ID: id, Dimension: 2, State: ochrevector.StateReady,
|
||||
SourceSpecs: []ochrevector.SourceSpec{{Bucket: "docs", Dimension: 2}},
|
||||
CreatedAt: now, UpdatedAt: now,
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("IndexExists error skips the record", func(t *testing.T) {
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
registry := handlers_ochrevector.NewRegistry(js)
|
||||
jobs := handlers_ochrevector.NewJobStore(js)
|
||||
registry := ochrevector.NewRegistry(js)
|
||||
jobs := ochrevector.NewJobStore(js)
|
||||
backend := newReconcileFakeBackend()
|
||||
backend.indexExistsErr = errors.New("index exists boom")
|
||||
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
ingest := ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
|
||||
ctx := context.Background()
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, newRecord("idx-exists-err")))
|
||||
@@ -286,11 +286,11 @@ func TestReconcileOchreIndexes_ErrorsAreLoggedAndSkipped(t *testing.T) {
|
||||
|
||||
t.Run("EnsureAccount error skips the record", func(t *testing.T) {
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
registry := handlers_ochrevector.NewRegistry(js)
|
||||
jobs := handlers_ochrevector.NewJobStore(js)
|
||||
registry := ochrevector.NewRegistry(js)
|
||||
jobs := ochrevector.NewJobStore(js)
|
||||
backend := newReconcileFakeBackend()
|
||||
backend.ensureAccountErr = errors.New("ensure account boom")
|
||||
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
ingest := ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
|
||||
ctx := context.Background()
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, newRecord("idx-ensure-err")))
|
||||
@@ -307,11 +307,11 @@ func TestReconcileOchreIndexes_ErrorsAreLoggedAndSkipped(t *testing.T) {
|
||||
|
||||
t.Run("CreateIndex error skips the record", func(t *testing.T) {
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
registry := handlers_ochrevector.NewRegistry(js)
|
||||
jobs := handlers_ochrevector.NewJobStore(js)
|
||||
registry := ochrevector.NewRegistry(js)
|
||||
jobs := ochrevector.NewJobStore(js)
|
||||
backend := newReconcileFakeBackend()
|
||||
backend.createIndexErr = errors.New("create index boom")
|
||||
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
ingest := ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
|
||||
ctx := context.Background()
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, newRecord("idx-create-err")))
|
||||
@@ -327,18 +327,18 @@ func TestReconcileOchreIndexes_ErrorsAreLoggedAndSkipped(t *testing.T) {
|
||||
|
||||
t.Run("ReconcileFromSource error is logged and does not panic", func(t *testing.T) {
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
registry := handlers_ochrevector.NewRegistry(js)
|
||||
jobs := handlers_ochrevector.NewJobStore(js)
|
||||
registry := ochrevector.NewRegistry(js)
|
||||
jobs := ochrevector.NewJobStore(js)
|
||||
backend := newReconcileFakeBackend()
|
||||
ctx := context.Background()
|
||||
// Delete the just-created record as a side effect of CreateIndex
|
||||
// succeeding, so the ReconcileFromSource call right after it misses
|
||||
// the registry and returns ErrIndexNotFound -- no production code
|
||||
// change needed to force that failure.
|
||||
backend.onCreateIndex = func(ctx context.Context, accountID string, spec handlers_ochrevector.IndexSpec) {
|
||||
backend.onCreateIndex = func(ctx context.Context, accountID string, spec ochrevector.IndexSpec) {
|
||||
require.NoError(t, registry.Delete(ctx, accountID, spec.ID))
|
||||
}
|
||||
ingest := handlers_ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
ingest := ochrevector.NewIngestService(jobs, registry, backend, nil, nil)
|
||||
|
||||
require.NoError(t, registry.Reserve(ctx, reconcileTestAccount, newRecord("idx-ingest-err")))
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// exported surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -1,7 +1,7 @@
|
||||
// Package handlers_ochrevector holds Ochre's tenant vector store: the
|
||||
// Package vector holds Ochre's tenant vector store: the
|
||||
// pgvector backend, the KV index registry, and the create/delete/list/
|
||||
// reconcile orchestration over both.
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import "context"
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// exported surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// orchestration tests.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"strings"
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
awsidentifiers "github.com/mulgadc/spinifex/spinifex/foundation/aws/identifiers"
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// jobs_test.go) and exercise CAS/prefix-scan behavior directly.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"testing"
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -1,4 +1,4 @@
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
// exported surface to drive them through.
|
||||
//
|
||||
//test:in-package
|
||||
package handlers_ochrevector
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -15,9 +15,9 @@ import (
|
||||
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockagent"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
"github.com/nats-io/nats.go"
|
||||
)
|
||||
|
||||
@@ -29,7 +29,7 @@ type bedrockAgentRoute = restRoute[bedrockAgentRouteHandler]
|
||||
// kb/ds are gw.BedrockAgentKB / gw.BedrockAgentDataSources; vector is
|
||||
// gw.BedrockAgentVector, the NATSVectorService forwarding client to .9's
|
||||
// daemon-side VectorService.
|
||||
type bedrockAgentRouteHandler func(ctx context.Context, accountID, region string, params []string, body []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error)
|
||||
type bedrockAgentRouteHandler func(ctx context.Context, accountID, region string, params []string, body []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService) (any, error)
|
||||
|
||||
// bedrockAgentRoutes is the dispatch table. Real AWS HTTP paths/methods
|
||||
// (verified against the vendored aws-sdk-go bedrockagent request
|
||||
@@ -37,7 +37,7 @@ type bedrockAgentRouteHandler func(ctx context.Context, accountID, region string
|
||||
// chi matches exactly.
|
||||
var bedrockAgentRoutes = []bedrockAgentRoute{
|
||||
{"PUT", "/knowledgebases/", "CreateKnowledgeBase",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, _ *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, _ *ochrevector.DataSourceStore, vector ochrevector.VectorService) (any, error) {
|
||||
input := new(bedrockagent.CreateKnowledgeBaseInput)
|
||||
if len(b) > 0 {
|
||||
if err := json.Unmarshal(b, input); err != nil {
|
||||
@@ -47,19 +47,19 @@ var bedrockAgentRoutes = []bedrockAgentRoute{
|
||||
return CreateKnowledgeBase(ctx, acct, region, kb, vector, input)
|
||||
}},
|
||||
{"POST", "/knowledgebases/", "ListKnowledgeBases",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, _ *handlers_ochrevector.DataSourceStore, _ handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, _ *ochrevector.DataSourceStore, _ ochrevector.VectorService) (any, error) {
|
||||
return ListKnowledgeBases(ctx, acct, kb, new(bedrockagent.ListKnowledgeBasesInput))
|
||||
}},
|
||||
{"GET", "/knowledgebases/{knowledgeBaseId}", "GetKnowledgeBase",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, _ *handlers_ochrevector.DataSourceStore, _ handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, _ *ochrevector.DataSourceStore, _ ochrevector.VectorService) (any, error) {
|
||||
return GetKnowledgeBase(ctx, acct, region, kb, &bedrockagent.GetKnowledgeBaseInput{KnowledgeBaseId: aws.String(p[0])})
|
||||
}},
|
||||
{"DELETE", "/knowledgebases/{knowledgeBaseId}", "DeleteKnowledgeBase",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService) (any, error) {
|
||||
return DeleteKnowledgeBase(ctx, acct, kb, ds, vector, &bedrockagent.DeleteKnowledgeBaseInput{KnowledgeBaseId: aws.String(p[0])})
|
||||
}},
|
||||
{"PUT", "/knowledgebases/{knowledgeBaseId}/datasources/", "CreateDataSource",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, _ handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, _ ochrevector.VectorService) (any, error) {
|
||||
input := new(bedrockagent.CreateDataSourceInput)
|
||||
if len(b) > 0 {
|
||||
if err := json.Unmarshal(b, input); err != nil {
|
||||
@@ -70,31 +70,31 @@ var bedrockAgentRoutes = []bedrockAgentRoute{
|
||||
return CreateDataSource(ctx, acct, region, kb, ds, input)
|
||||
}},
|
||||
{"POST", "/knowledgebases/{knowledgeBaseId}/datasources/", "ListDataSources",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, _ handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, _ ochrevector.VectorService) (any, error) {
|
||||
return ListDataSources(ctx, acct, kb, ds, &bedrockagent.ListDataSourcesInput{KnowledgeBaseId: aws.String(p[0])})
|
||||
}},
|
||||
{"GET", "/knowledgebases/{knowledgeBaseId}/datasources/{dataSourceId}", "GetDataSource",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, _ *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, _ handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, _ *ochrevector.KBStore, ds *ochrevector.DataSourceStore, _ ochrevector.VectorService) (any, error) {
|
||||
return GetDataSource(ctx, acct, region, ds, &bedrockagent.GetDataSourceInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1])})
|
||||
}},
|
||||
{"DELETE", "/knowledgebases/{knowledgeBaseId}/datasources/{dataSourceId}", "DeleteDataSource",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, _ *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, _ handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, _ *ochrevector.KBStore, ds *ochrevector.DataSourceStore, _ ochrevector.VectorService) (any, error) {
|
||||
return DeleteDataSource(ctx, acct, ds, &bedrockagent.DeleteDataSourceInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1])})
|
||||
}},
|
||||
{"PUT", "/knowledgebases/{knowledgeBaseId}/datasources/{dataSourceId}/ingestionjobs/", "StartIngestionJob",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService) (any, error) {
|
||||
return StartIngestionJob(ctx, acct, kb, ds, vector, &bedrockagent.StartIngestionJobInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1])})
|
||||
}},
|
||||
{"POST", "/knowledgebases/{knowledgeBaseId}/datasources/{dataSourceId}/ingestionjobs/", "ListIngestionJobs",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService) (any, error) {
|
||||
return ListIngestionJobs(ctx, acct, kb, ds, vector, &bedrockagent.ListIngestionJobsInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1])})
|
||||
}},
|
||||
{"GET", "/knowledgebases/{knowledgeBaseId}/datasources/{dataSourceId}/ingestionjobs/{ingestionJobId}", "GetIngestionJob",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService) (any, error) {
|
||||
return GetIngestionJob(ctx, acct, kb, ds, vector, &bedrockagent.GetIngestionJobInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1]), IngestionJobId: aws.String(p[2])})
|
||||
}},
|
||||
{"PUT", "/knowledgebases/{knowledgeBaseId}/datasources/{dataSourceId}/ingestionjobs/{ingestionJobId}/stop", "StopIngestionJob",
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) {
|
||||
func(ctx context.Context, acct, region string, p []string, b []byte, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService) (any, error) {
|
||||
return StopIngestionJob(ctx, acct, kb, ds, vector, &StopIngestionJobInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1]), IngestionJobId: aws.String(p[2])})
|
||||
}},
|
||||
}
|
||||
@@ -196,11 +196,11 @@ func formatS3BucketARN(bucket string) string {
|
||||
// so only StateReady is ever actually persisted here in Pass 1.
|
||||
func kbStatusToAWS(status string) string {
|
||||
switch status {
|
||||
case handlers_ochrevector.StateReady:
|
||||
case ochrevector.StateReady:
|
||||
return bedrockagent.KnowledgeBaseStatusActive
|
||||
case handlers_ochrevector.StateCreating:
|
||||
case ochrevector.StateCreating:
|
||||
return bedrockagent.KnowledgeBaseStatusCreating
|
||||
case handlers_ochrevector.StateDeleting:
|
||||
case ochrevector.StateDeleting:
|
||||
return bedrockagent.KnowledgeBaseStatusDeleting
|
||||
default:
|
||||
return bedrockagent.KnowledgeBaseStatusFailed
|
||||
@@ -210,7 +210,7 @@ func kbStatusToAWS(status string) string {
|
||||
// dataSourceStatusToAWS is kbStatusToAWS's sibling for DataSourceStatus.
|
||||
func dataSourceStatusToAWS(status string) string {
|
||||
switch status {
|
||||
case handlers_ochrevector.StateDeleting:
|
||||
case ochrevector.StateDeleting:
|
||||
return bedrockagent.DataSourceStatusDeleting
|
||||
default:
|
||||
return bedrockagent.DataSourceStatusAvailable
|
||||
@@ -228,13 +228,13 @@ const ingestionJobStatusStopped = "STOPPED"
|
||||
// wire values.
|
||||
func jobStateToAWS(state string) string {
|
||||
switch state {
|
||||
case handlers_ochrevector.JobStatePending:
|
||||
case ochrevector.JobStatePending:
|
||||
return bedrockagent.IngestionJobStatusStarting
|
||||
case handlers_ochrevector.JobStateRunning:
|
||||
case ochrevector.JobStateRunning:
|
||||
return bedrockagent.IngestionJobStatusInProgress
|
||||
case handlers_ochrevector.JobStateReady:
|
||||
case ochrevector.JobStateReady:
|
||||
return bedrockagent.IngestionJobStatusComplete
|
||||
case handlers_ochrevector.JobStateStopped:
|
||||
case ochrevector.JobStateStopped:
|
||||
return ingestionJobStatusStopped
|
||||
default:
|
||||
return bedrockagent.IngestionJobStatusFailed
|
||||
@@ -264,9 +264,9 @@ func translateVectorErr(err error) error {
|
||||
switch {
|
||||
case err == nil:
|
||||
return nil
|
||||
case errors.Is(err, handlers_ochrevector.ErrIndexNotFound), errors.Is(err, handlers_ochrevector.ErrJobNotFound):
|
||||
case errors.Is(err, ochrevector.ErrIndexNotFound), errors.Is(err, ochrevector.ErrJobNotFound):
|
||||
return errors.New(awserrors.ErrorResourceNotFoundException)
|
||||
case errors.Is(err, handlers_ochrevector.ErrIndexExists):
|
||||
case errors.Is(err, ochrevector.ErrIndexExists):
|
||||
return errors.New(awserrors.ErrorConflictException)
|
||||
case errors.Is(err, nats.ErrNoResponders), errors.Is(err, nats.ErrTimeout), errors.Is(err, context.DeadlineExceeded):
|
||||
return errors.New(awserrors.ErrorServiceUnavailableException)
|
||||
@@ -279,7 +279,7 @@ func translateVectorErr(err error) error {
|
||||
// StorageConfigJSON blob is decoded back into the typed AWS struct so
|
||||
// Get/List echo exactly what CreateKnowledgeBase accepted (D5: accepted and
|
||||
// stubbed, never acted on).
|
||||
func kbRecordToOutput(region, accountID string, rec handlers_ochrevector.KBRecord) (*bedrockagent.KnowledgeBase, error) {
|
||||
func kbRecordToOutput(region, accountID string, rec ochrevector.KBRecord) (*bedrockagent.KnowledgeBase, error) {
|
||||
storageConfig := new(bedrockagent.StorageConfiguration)
|
||||
if len(rec.StorageConfigJSON) > 0 {
|
||||
if err := json.Unmarshal(rec.StorageConfigJSON, storageConfig); err != nil {
|
||||
@@ -321,7 +321,7 @@ func kbRecordToOutput(region, accountID string, rec handlers_ochrevector.KBRecor
|
||||
// round-trip echo but never honored (D5). A KBStore.Create collision (id
|
||||
// reuse) rolls the just-created index back, so no orphan index survives a
|
||||
// failed claim.
|
||||
func CreateKnowledgeBase(ctx context.Context, accountID, region string, kb *handlers_ochrevector.KBStore, vector handlers_ochrevector.VectorService, input *bedrockagent.CreateKnowledgeBaseInput) (*bedrockagent.CreateKnowledgeBaseOutput, error) {
|
||||
func CreateKnowledgeBase(ctx context.Context, accountID, region string, kb *ochrevector.KBStore, vector ochrevector.VectorService, input *bedrockagent.CreateKnowledgeBaseInput) (*bedrockagent.CreateKnowledgeBaseOutput, error) {
|
||||
if input == nil || aws.StringValue(input.Name) == "" || aws.StringValue(input.RoleArn) == "" ||
|
||||
input.StorageConfiguration == nil || input.KnowledgeBaseConfiguration == nil {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
@@ -347,7 +347,7 @@ func CreateKnowledgeBase(ctx context.Context, accountID, region string, kb *hand
|
||||
}
|
||||
|
||||
id := uuid.NewV4().String()
|
||||
indexResp, err := vector.CreateIndex(ctx, &handlers_ochrevector.CreateIndexRequest{
|
||||
indexResp, err := vector.CreateIndex(ctx, &ochrevector.CreateIndexRequest{
|
||||
IndexID: id,
|
||||
Name: aws.StringValue(input.Name),
|
||||
Dimension: dimension,
|
||||
@@ -363,11 +363,11 @@ func CreateKnowledgeBase(ctx context.Context, accountID, region string, kb *hand
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
rec := handlers_ochrevector.KBRecord{
|
||||
rec := ochrevector.KBRecord{
|
||||
ID: id,
|
||||
Name: aws.StringValue(input.Name),
|
||||
Description: aws.StringValue(input.Description),
|
||||
Status: handlers_ochrevector.StateReady,
|
||||
Status: ochrevector.StateReady,
|
||||
EmbeddingModel: embeddingModel,
|
||||
EmbeddingModelArn: embeddingModelArn,
|
||||
Dimension: dimension,
|
||||
@@ -383,10 +383,10 @@ func CreateKnowledgeBase(ctx context.Context, accountID, region string, kb *hand
|
||||
// below, so a rollback failure here must not be swallowed -- log it
|
||||
// at Error with the orphaned index id so it stays observable even
|
||||
// though the caller never sees it.
|
||||
if _, delErr := vector.DeleteIndex(ctx, &handlers_ochrevector.DeleteIndexRequest{IndexID: id}, accountID); delErr != nil {
|
||||
if _, delErr := vector.DeleteIndex(ctx, &ochrevector.DeleteIndexRequest{IndexID: id}, accountID); delErr != nil {
|
||||
slog.ErrorContext(ctx, "bedrock-agent: rollback delete index after knowledge base claim failure left an orphaned index", "index", id, "err", delErr)
|
||||
}
|
||||
if errors.Is(err, handlers_ochrevector.ErrKBExists) {
|
||||
if errors.Is(err, ochrevector.ErrKBExists) {
|
||||
return nil, errors.New(awserrors.ErrorConflictException)
|
||||
}
|
||||
return nil, err
|
||||
@@ -402,7 +402,7 @@ func CreateKnowledgeBase(ctx context.Context, accountID, region string, kb *hand
|
||||
// GetKnowledgeBase looks up id in kb, returning ResourceNotFoundException for
|
||||
// a foreign account or an unknown id alike so a caller cannot distinguish
|
||||
// "not yours" from "does not exist".
|
||||
func GetKnowledgeBase(ctx context.Context, accountID, region string, kb *handlers_ochrevector.KBStore, input *bedrockagent.GetKnowledgeBaseInput) (*bedrockagent.GetKnowledgeBaseOutput, error) {
|
||||
func GetKnowledgeBase(ctx context.Context, accountID, region string, kb *ochrevector.KBStore, input *bedrockagent.GetKnowledgeBaseInput) (*bedrockagent.GetKnowledgeBaseOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -424,12 +424,12 @@ func GetKnowledgeBase(ctx context.Context, accountID, region string, kb *handler
|
||||
// ListKnowledgeBases returns the caller account's own knowledge bases,
|
||||
// sorted by creation time (id as a deterministic tie-breaker), mirroring
|
||||
// gateway_bedrock.ListGuardrails.
|
||||
func ListKnowledgeBases(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, _ *bedrockagent.ListKnowledgeBasesInput) (*bedrockagent.ListKnowledgeBasesOutput, error) {
|
||||
func ListKnowledgeBases(ctx context.Context, accountID string, kb *ochrevector.KBStore, _ *bedrockagent.ListKnowledgeBasesInput) (*bedrockagent.ListKnowledgeBasesOutput, error) {
|
||||
recs, err := kb.List(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
slices.SortFunc(recs, func(a, b handlers_ochrevector.KBRecord) int {
|
||||
slices.SortFunc(recs, func(a, b ochrevector.KBRecord) int {
|
||||
if c := a.CreatedAt.Compare(b.CreatedAt); c != 0 {
|
||||
return c
|
||||
}
|
||||
@@ -455,7 +455,7 @@ func ListKnowledgeBases(ctx context.Context, accountID string, kb *handlers_ochr
|
||||
// ResourceNotFoundException), unlike gateway_bedrock.DeleteGuardrail's
|
||||
// idempotent-delete contract, because AWS's own DeleteKnowledgeBase does the
|
||||
// same.
|
||||
func DeleteKnowledgeBase(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService, input *bedrockagent.DeleteKnowledgeBaseInput) (*bedrockagent.DeleteKnowledgeBaseOutput, error) {
|
||||
func DeleteKnowledgeBase(ctx context.Context, accountID string, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService, input *bedrockagent.DeleteKnowledgeBaseInput) (*bedrockagent.DeleteKnowledgeBaseOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -483,7 +483,7 @@ func DeleteKnowledgeBase(ctx context.Context, accountID string, kb *handlers_och
|
||||
// delete just because its index is already gone (Delete is idempotent
|
||||
// w.r.t. the index). Any other error still aborts before the record
|
||||
// itself is removed.
|
||||
if _, err := vector.DeleteIndex(ctx, &handlers_ochrevector.DeleteIndexRequest{IndexID: rec.IndexID}, accountID); err != nil && !errors.Is(err, handlers_ochrevector.ErrIndexNotFound) {
|
||||
if _, err := vector.DeleteIndex(ctx, &ochrevector.DeleteIndexRequest{IndexID: rec.IndexID}, accountID); err != nil && !errors.Is(err, ochrevector.ErrIndexNotFound) {
|
||||
return nil, translateVectorErr(err)
|
||||
}
|
||||
if err := kb.Delete(ctx, accountID, id); err != nil {
|
||||
@@ -498,7 +498,7 @@ func DeleteKnowledgeBase(ctx context.Context, accountID string, kb *handlers_och
|
||||
// dataSourceRecordToOutput builds the AWS DataSource shape from rec. Only S3
|
||||
// data sources are supported (CreateDataSource rejects any other type), so
|
||||
// this always renders an S3Configuration.
|
||||
func dataSourceRecordToOutput(rec handlers_ochrevector.DataSourceRecord) *bedrockagent.DataSource {
|
||||
func dataSourceRecordToOutput(rec ochrevector.DataSourceRecord) *bedrockagent.DataSource {
|
||||
var inclusionPrefixes []*string
|
||||
if rec.Source.Prefix != "" {
|
||||
inclusionPrefixes = []*string{aws.String(rec.Source.Prefix)}
|
||||
@@ -532,7 +532,7 @@ func dataSourceRecordToOutput(rec handlers_ochrevector.DataSourceRecord) *bedroc
|
||||
// CreateDataSourceInput carries no arbitrary metadata map to source it from
|
||||
// (real per-document metadata comes from S3 .metadata.json sidecar files at
|
||||
// ingest time, out of scope for Pass 1).
|
||||
func CreateDataSource(ctx context.Context, accountID, region string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, input *bedrockagent.CreateDataSourceInput) (*bedrockagent.CreateDataSourceOutput, error) {
|
||||
func CreateDataSource(ctx context.Context, accountID, region string, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, input *bedrockagent.CreateDataSourceInput) (*bedrockagent.CreateDataSourceOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.Name) == "" || input.DataSourceConfiguration == nil {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -573,13 +573,13 @@ func CreateDataSource(ctx context.Context, accountID, region string, kb *handler
|
||||
|
||||
id := uuid.NewV4().String()
|
||||
now := time.Now().UTC()
|
||||
rec := handlers_ochrevector.DataSourceRecord{
|
||||
rec := ochrevector.DataSourceRecord{
|
||||
ID: id,
|
||||
KnowledgeBaseID: kbID,
|
||||
Name: aws.StringValue(input.Name),
|
||||
Description: aws.StringValue(input.Description),
|
||||
Status: handlers_ochrevector.StateReady,
|
||||
Source: handlers_ochrevector.SourceSpec{
|
||||
Status: ochrevector.StateReady,
|
||||
Source: ochrevector.SourceSpec{
|
||||
Bucket: bucketNameFromS3ARN(bucketArn),
|
||||
Prefix: prefix,
|
||||
ChunkSize: chunkSize,
|
||||
@@ -591,7 +591,7 @@ func CreateDataSource(ctx context.Context, accountID, region string, kb *handler
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if err := ds.Create(ctx, accountID, rec); err != nil {
|
||||
if errors.Is(err, handlers_ochrevector.ErrDataSourceExists) {
|
||||
if errors.Is(err, ochrevector.ErrDataSourceExists) {
|
||||
return nil, errors.New(awserrors.ErrorConflictException)
|
||||
}
|
||||
return nil, err
|
||||
@@ -602,7 +602,7 @@ func CreateDataSource(ctx context.Context, accountID, region string, kb *handler
|
||||
// GetDataSource looks up dataSourceId scoped to knowledgeBaseId: a data
|
||||
// source that exists but belongs to a different knowledge base reports
|
||||
// ResourceNotFoundException, the same as one that does not exist at all.
|
||||
func GetDataSource(ctx context.Context, accountID, region string, ds *handlers_ochrevector.DataSourceStore, input *bedrockagent.GetDataSourceInput) (*bedrockagent.GetDataSourceOutput, error) {
|
||||
func GetDataSource(ctx context.Context, accountID, region string, ds *ochrevector.DataSourceStore, input *bedrockagent.GetDataSourceInput) (*bedrockagent.GetDataSourceOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.DataSourceId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -619,7 +619,7 @@ func GetDataSource(ctx context.Context, accountID, region string, ds *handlers_o
|
||||
|
||||
// ListDataSources returns knowledgeBaseId's data sources, sorted by creation
|
||||
// time (id as a deterministic tie-breaker), mirroring ListKnowledgeBases.
|
||||
func ListDataSources(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, input *bedrockagent.ListDataSourcesInput) (*bedrockagent.ListDataSourcesOutput, error) {
|
||||
func ListDataSources(ctx context.Context, accountID string, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, input *bedrockagent.ListDataSourcesInput) (*bedrockagent.ListDataSourcesOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -636,7 +636,7 @@ func ListDataSources(ctx context.Context, accountID string, kb *handlers_ochreve
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
slices.SortFunc(recs, func(a, b handlers_ochrevector.DataSourceRecord) int {
|
||||
slices.SortFunc(recs, func(a, b ochrevector.DataSourceRecord) int {
|
||||
if c := a.CreatedAt.Compare(b.CreatedAt); c != 0 {
|
||||
return c
|
||||
}
|
||||
@@ -659,7 +659,7 @@ func ListDataSources(ctx context.Context, accountID string, kb *handlers_ochreve
|
||||
// DeleteDataSource removes dataSourceId scoped to knowledgeBaseId. Not
|
||||
// idempotent on a foreign/unknown id (reports ResourceNotFoundException), the
|
||||
// same contract as DeleteKnowledgeBase and real AWS's own DeleteDataSource.
|
||||
func DeleteDataSource(ctx context.Context, accountID string, ds *handlers_ochrevector.DataSourceStore, input *bedrockagent.DeleteDataSourceInput) (*bedrockagent.DeleteDataSourceOutput, error) {
|
||||
func DeleteDataSource(ctx context.Context, accountID string, ds *ochrevector.DataSourceStore, input *bedrockagent.DeleteDataSourceInput) (*bedrockagent.DeleteDataSourceOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.DataSourceId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -685,7 +685,7 @@ func DeleteDataSource(ctx context.Context, accountID string, ds *handlers_ochrev
|
||||
// addressed kbID/dsID (the caller's own resolved ids, not necessarily rebuilt
|
||||
// from job.IndexID/job.DataSourceID, since AWS's IngestionJob always renders
|
||||
// under the path it was requested through).
|
||||
func jobRecordToOutput(kbID, dsID string, job handlers_ochrevector.JobRecord) *bedrockagent.IngestionJob {
|
||||
func jobRecordToOutput(kbID, dsID string, job ochrevector.JobRecord) *bedrockagent.IngestionJob {
|
||||
var failureReasons []*string
|
||||
for _, fd := range job.FailedDocuments {
|
||||
failureReasons = append(failureReasons, aws.String(fmt.Sprintf("%s: %s", fd.SourceKey, fd.Reason)))
|
||||
@@ -716,7 +716,7 @@ func jobRecordToOutput(kbID, dsID string, job handlers_ochrevector.JobRecord) *b
|
||||
// VectorService.Ingest. dsRec.ID is stamped onto the request as DataSourceID,
|
||||
// so the resulting job carries an exact link back to the data source that
|
||||
// started it.
|
||||
func StartIngestionJob(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService, input *bedrockagent.StartIngestionJobInput) (*bedrockagent.StartIngestionJobOutput, error) {
|
||||
func StartIngestionJob(ctx context.Context, accountID string, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService, input *bedrockagent.StartIngestionJobInput) (*bedrockagent.StartIngestionJobOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.DataSourceId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -739,7 +739,7 @@ func StartIngestionJob(ctx context.Context, accountID string, kb *handlers_ochre
|
||||
return nil, errDataSourceNotFound(dsID)
|
||||
}
|
||||
|
||||
resp, err := vector.Ingest(ctx, &handlers_ochrevector.IngestRequest{IndexID: kbRec.IndexID, Source: dsRec.Source, DataSourceID: dsRec.ID}, accountID)
|
||||
resp, err := vector.Ingest(ctx, &ochrevector.IngestRequest{IndexID: kbRec.IndexID, Source: dsRec.Source, DataSourceID: dsRec.ID}, accountID)
|
||||
if err != nil {
|
||||
return nil, translateVectorErr(err)
|
||||
}
|
||||
@@ -755,7 +755,7 @@ func StartIngestionJob(ctx context.Context, accountID string, kb *handlers_ochre
|
||||
// KB/data source's job by guessing its id. A job with an empty DataSourceID
|
||||
// (started directly against ochre.vector.ingest, with no bedrock-agent data
|
||||
// source involved) never matches any dataSourceId here.
|
||||
func GetIngestionJob(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService, input *bedrockagent.GetIngestionJobInput) (*bedrockagent.GetIngestionJobOutput, error) {
|
||||
func GetIngestionJob(ctx context.Context, accountID string, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService, input *bedrockagent.GetIngestionJobInput) (*bedrockagent.GetIngestionJobOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.DataSourceId) == "" || aws.StringValue(input.IngestionJobId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -779,7 +779,7 @@ func GetIngestionJob(ctx context.Context, accountID string, kb *handlers_ochreve
|
||||
return nil, errDataSourceNotFound(dsID)
|
||||
}
|
||||
|
||||
resp, err := vector.DescribeJob(ctx, &handlers_ochrevector.DescribeJobRequest{JobID: jobID}, accountID)
|
||||
resp, err := vector.DescribeJob(ctx, &ochrevector.DescribeJobRequest{JobID: jobID}, accountID)
|
||||
if err != nil {
|
||||
return nil, translateVectorErr(err)
|
||||
}
|
||||
@@ -810,7 +810,7 @@ type StopIngestionJobOutput struct {
|
||||
// must match) before cancelling it via VectorService.StopJob, so a
|
||||
// foreign/mismatched knowledgeBaseId or dataSourceId in the path cannot be
|
||||
// used to stop another KB/data source's job by guessing its id.
|
||||
func StopIngestionJob(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService, input *StopIngestionJobInput) (*StopIngestionJobOutput, error) {
|
||||
func StopIngestionJob(ctx context.Context, accountID string, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService, input *StopIngestionJobInput) (*StopIngestionJobOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.DataSourceId) == "" || aws.StringValue(input.IngestionJobId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -838,7 +838,7 @@ func StopIngestionJob(ctx context.Context, accountID string, kb *handlers_ochrev
|
||||
// enforces: StopJob alone has no knowledge base/data source to compare
|
||||
// against, so a foreign job id must be rejected before anything is
|
||||
// cancelled.
|
||||
describeResp, err := vector.DescribeJob(ctx, &handlers_ochrevector.DescribeJobRequest{JobID: jobID}, accountID)
|
||||
describeResp, err := vector.DescribeJob(ctx, &ochrevector.DescribeJobRequest{JobID: jobID}, accountID)
|
||||
if err != nil {
|
||||
return nil, translateVectorErr(err)
|
||||
}
|
||||
@@ -846,7 +846,7 @@ func StopIngestionJob(ctx context.Context, accountID string, kb *handlers_ochrev
|
||||
return nil, errIngestionJobNotFound(jobID)
|
||||
}
|
||||
|
||||
resp, err := vector.StopJob(ctx, &handlers_ochrevector.StopJobRequest{JobID: jobID}, accountID)
|
||||
resp, err := vector.StopJob(ctx, &ochrevector.StopJobRequest{JobID: jobID}, accountID)
|
||||
if err != nil {
|
||||
return nil, translateVectorErr(err)
|
||||
}
|
||||
@@ -857,7 +857,7 @@ func StopIngestionJob(ctx context.Context, accountID string, kb *handlers_ochrev
|
||||
// ListIngestionJobs lists knowledgeBaseId's jobs via the new
|
||||
// VectorService.ListJobs, filtered to dataSourceId's own bound index and
|
||||
// exact DataSourceID, and sorted by start time.
|
||||
func ListIngestionJobs(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService, input *bedrockagent.ListIngestionJobsInput) (*bedrockagent.ListIngestionJobsOutput, error) {
|
||||
func ListIngestionJobs(ctx context.Context, accountID string, kb *ochrevector.KBStore, ds *ochrevector.DataSourceStore, vector ochrevector.VectorService, input *bedrockagent.ListIngestionJobsInput) (*bedrockagent.ListIngestionJobsOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.DataSourceId) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -880,7 +880,7 @@ func ListIngestionJobs(ctx context.Context, accountID string, kb *handlers_ochre
|
||||
return nil, errDataSourceNotFound(dsID)
|
||||
}
|
||||
|
||||
resp, err := vector.ListJobs(ctx, &handlers_ochrevector.ListJobsRequest{}, accountID)
|
||||
resp, err := vector.ListJobs(ctx, &ochrevector.ListJobsRequest{}, accountID)
|
||||
if err != nil {
|
||||
return nil, translateVectorErr(err)
|
||||
}
|
||||
|
||||
@@ -16,8 +16,8 @@ import (
|
||||
|
||||
"github.com/aws/aws-sdk-go/service/bedrockagent"
|
||||
"github.com/mulgadc/spinifex/internal/testkit"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -26,13 +26,13 @@ import (
|
||||
// gateway-owned KB/DataSource stores over a real (embedded) JetStream and an
|
||||
// allow-everything IAMService, so the policy gate passes on a real evaluation
|
||||
// rather than a bypass, mirroring newBedrockRequestGateway.
|
||||
func newBedrockAgentRequestGateway(t *testing.T, vector handlers_ochrevector.VectorService) *GatewayConfig {
|
||||
func newBedrockAgentRequestGateway(t *testing.T, vector ochrevector.VectorService) *GatewayConfig {
|
||||
t.Helper()
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
return &GatewayConfig{
|
||||
IAMService: allowAllIAMService(),
|
||||
BedrockAgentKB: handlers_ochrevector.NewKBStore(js),
|
||||
BedrockAgentDataSources: handlers_ochrevector.NewDataSourceStore(js),
|
||||
BedrockAgentKB: ochrevector.NewKBStore(js),
|
||||
BedrockAgentDataSources: ochrevector.NewDataSourceStore(js),
|
||||
BedrockAgentVector: vector,
|
||||
}
|
||||
}
|
||||
@@ -256,8 +256,8 @@ func TestBedrockAgentRequest_DeleteKnowledgeBase(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestBedrockAgentRequest_StartAndListAndGetIngestionJob(t *testing.T) {
|
||||
vector := &fakeBedrockAgentVectorService{ingestResp: handlers_ochrevector.IngestResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", State: handlers_ochrevector.JobStatePending},
|
||||
vector := &fakeBedrockAgentVectorService{ingestResp: ochrevector.IngestResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", State: ochrevector.JobStatePending},
|
||||
}}
|
||||
gw := newBedrockAgentRequestGateway(t, vector)
|
||||
kbID := createTestKnowledgeBase(t, gw)
|
||||
@@ -275,8 +275,8 @@ func TestBedrockAgentRequest_StartAndListAndGetIngestionJob(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
vector.ingestResp.Job.IndexID = kbRec.IndexID
|
||||
vector.ingestResp.Job.DataSourceID = dsRec.ID
|
||||
vector.describeJobResp = handlers_ochrevector.DescribeJobResponse{Job: vector.ingestResp.Job}
|
||||
vector.listJobsResp = handlers_ochrevector.ListJobsResponse{Jobs: []handlers_ochrevector.JobRecord{vector.ingestResp.Job}}
|
||||
vector.describeJobResp = ochrevector.DescribeJobResponse{Job: vector.ingestResp.Job}
|
||||
vector.listJobsResp = ochrevector.ListJobsResponse{Jobs: []ochrevector.JobRecord{vector.ingestResp.Job}}
|
||||
|
||||
startReq := bedrockAgentRequestWithAccount(http.MethodPut, "/knowledgebases/"+kbID+"/datasources/"+dsID+"/ingestionjobs/", "")
|
||||
w := httptest.NewRecorder()
|
||||
@@ -309,8 +309,8 @@ func TestBedrockAgentRequest_StartAndListAndGetIngestionJob(t *testing.T) {
|
||||
// ErrInvalidAction for it) and round-trips a STOPPED job through the full
|
||||
// HTTP dispatch path.
|
||||
func TestBedrockAgentRequest_StopIngestionJob(t *testing.T) {
|
||||
vector := &fakeBedrockAgentVectorService{ingestResp: handlers_ochrevector.IngestResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", State: handlers_ochrevector.JobStatePending},
|
||||
vector := &fakeBedrockAgentVectorService{ingestResp: ochrevector.IngestResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", State: ochrevector.JobStatePending},
|
||||
}}
|
||||
gw := newBedrockAgentRequestGateway(t, vector)
|
||||
kbID := createTestKnowledgeBase(t, gw)
|
||||
@@ -322,10 +322,10 @@ func TestBedrockAgentRequest_StopIngestionJob(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
vector.ingestResp.Job.IndexID = kbRec.IndexID
|
||||
vector.ingestResp.Job.DataSourceID = dsRec.ID
|
||||
vector.describeJobResp = handlers_ochrevector.DescribeJobResponse{Job: vector.ingestResp.Job}
|
||||
vector.describeJobResp = ochrevector.DescribeJobResponse{Job: vector.ingestResp.Job}
|
||||
stopped := vector.ingestResp.Job
|
||||
stopped.State = handlers_ochrevector.JobStateStopped
|
||||
vector.stopJobResp = handlers_ochrevector.StopJobResponse{Job: stopped}
|
||||
stopped.State = ochrevector.JobStateStopped
|
||||
vector.stopJobResp = ochrevector.StopJobResponse{Job: stopped}
|
||||
|
||||
startReq := bedrockAgentRequestWithAccount(http.MethodPut, "/knowledgebases/"+kbID+"/datasources/"+dsID+"/ingestionjobs/", "")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
@@ -18,8 +18,8 @@ import (
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockagent"
|
||||
"github.com/mulgadc/spinifex/internal/testkit"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -31,55 +31,55 @@ const bedrockAgentTestAccount = "111111111111"
|
||||
// mapping logic can be exercised without a daemon or a live NATS connection,
|
||||
// mirroring cmd/spinifex/cmd's fakeVectorService for this package.
|
||||
type fakeBedrockAgentVectorService struct {
|
||||
createIndexReq *handlers_ochrevector.CreateIndexRequest
|
||||
createIndexReq *ochrevector.CreateIndexRequest
|
||||
createIndexErr error
|
||||
|
||||
deletedIndexIDs []string
|
||||
deleteIndexErr error
|
||||
|
||||
ingestReq *handlers_ochrevector.IngestRequest
|
||||
ingestResp handlers_ochrevector.IngestResponse
|
||||
ingestReq *ochrevector.IngestRequest
|
||||
ingestResp ochrevector.IngestResponse
|
||||
ingestErr error
|
||||
|
||||
describeJobResp handlers_ochrevector.DescribeJobResponse
|
||||
describeJobResp ochrevector.DescribeJobResponse
|
||||
describeJobErr error
|
||||
|
||||
listJobsResp handlers_ochrevector.ListJobsResponse
|
||||
listJobsResp ochrevector.ListJobsResponse
|
||||
listJobsErr error
|
||||
|
||||
queryReq *handlers_ochrevector.QueryRequest
|
||||
queryResp handlers_ochrevector.QueryResponse
|
||||
queryReq *ochrevector.QueryRequest
|
||||
queryResp ochrevector.QueryResponse
|
||||
queryErr error
|
||||
|
||||
stopJobReq *handlers_ochrevector.StopJobRequest
|
||||
stopJobResp handlers_ochrevector.StopJobResponse
|
||||
stopJobReq *ochrevector.StopJobRequest
|
||||
stopJobResp ochrevector.StopJobResponse
|
||||
stopJobErr error
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) CreateIndex(_ context.Context, req *handlers_ochrevector.CreateIndexRequest, _ string) (*handlers_ochrevector.CreateIndexResponse, error) {
|
||||
func (f *fakeBedrockAgentVectorService) CreateIndex(_ context.Context, req *ochrevector.CreateIndexRequest, _ string) (*ochrevector.CreateIndexResponse, error) {
|
||||
f.createIndexReq = req
|
||||
if f.createIndexErr != nil {
|
||||
return nil, f.createIndexErr
|
||||
}
|
||||
return &handlers_ochrevector.CreateIndexResponse{Index: handlers_ochrevector.Record{
|
||||
return &ochrevector.CreateIndexResponse{Index: ochrevector.Record{
|
||||
ID: req.IndexID, Name: req.Name, Dimension: req.Dimension,
|
||||
EmbeddingModel: req.EmbeddingModel, State: handlers_ochrevector.StateReady,
|
||||
EmbeddingModel: req.EmbeddingModel, State: ochrevector.StateReady,
|
||||
}}, nil
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) DeleteIndex(_ context.Context, req *handlers_ochrevector.DeleteIndexRequest, _ string) (*handlers_ochrevector.DeleteIndexResponse, error) {
|
||||
func (f *fakeBedrockAgentVectorService) DeleteIndex(_ context.Context, req *ochrevector.DeleteIndexRequest, _ string) (*ochrevector.DeleteIndexResponse, error) {
|
||||
f.deletedIndexIDs = append(f.deletedIndexIDs, req.IndexID)
|
||||
if f.deleteIndexErr != nil {
|
||||
return nil, f.deleteIndexErr
|
||||
}
|
||||
return &handlers_ochrevector.DeleteIndexResponse{}, nil
|
||||
return &ochrevector.DeleteIndexResponse{}, nil
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) ListIndexes(_ context.Context, _ *handlers_ochrevector.ListIndexesRequest, _ string) (*handlers_ochrevector.ListIndexesResponse, error) {
|
||||
return &handlers_ochrevector.ListIndexesResponse{}, nil
|
||||
func (f *fakeBedrockAgentVectorService) ListIndexes(_ context.Context, _ *ochrevector.ListIndexesRequest, _ string) (*ochrevector.ListIndexesResponse, error) {
|
||||
return &ochrevector.ListIndexesResponse{}, nil
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) Ingest(_ context.Context, req *handlers_ochrevector.IngestRequest, _ string) (*handlers_ochrevector.IngestResponse, error) {
|
||||
func (f *fakeBedrockAgentVectorService) Ingest(_ context.Context, req *ochrevector.IngestRequest, _ string) (*ochrevector.IngestResponse, error) {
|
||||
f.ingestReq = req
|
||||
if f.ingestErr != nil {
|
||||
return nil, f.ingestErr
|
||||
@@ -88,7 +88,7 @@ func (f *fakeBedrockAgentVectorService) Ingest(_ context.Context, req *handlers_
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) DescribeJob(_ context.Context, _ *handlers_ochrevector.DescribeJobRequest, _ string) (*handlers_ochrevector.DescribeJobResponse, error) {
|
||||
func (f *fakeBedrockAgentVectorService) DescribeJob(_ context.Context, _ *ochrevector.DescribeJobRequest, _ string) (*ochrevector.DescribeJobResponse, error) {
|
||||
if f.describeJobErr != nil {
|
||||
return nil, f.describeJobErr
|
||||
}
|
||||
@@ -96,7 +96,7 @@ func (f *fakeBedrockAgentVectorService) DescribeJob(_ context.Context, _ *handle
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) ListJobs(_ context.Context, _ *handlers_ochrevector.ListJobsRequest, _ string) (*handlers_ochrevector.ListJobsResponse, error) {
|
||||
func (f *fakeBedrockAgentVectorService) ListJobs(_ context.Context, _ *ochrevector.ListJobsRequest, _ string) (*ochrevector.ListJobsResponse, error) {
|
||||
if f.listJobsErr != nil {
|
||||
return nil, f.listJobsErr
|
||||
}
|
||||
@@ -104,7 +104,7 @@ func (f *fakeBedrockAgentVectorService) ListJobs(_ context.Context, _ *handlers_
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) Query(_ context.Context, req *handlers_ochrevector.QueryRequest, _ string) (*handlers_ochrevector.QueryResponse, error) {
|
||||
func (f *fakeBedrockAgentVectorService) Query(_ context.Context, req *ochrevector.QueryRequest, _ string) (*ochrevector.QueryResponse, error) {
|
||||
f.queryReq = req
|
||||
if f.queryErr != nil {
|
||||
return nil, f.queryErr
|
||||
@@ -113,7 +113,7 @@ func (f *fakeBedrockAgentVectorService) Query(_ context.Context, req *handlers_o
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
func (f *fakeBedrockAgentVectorService) StopJob(_ context.Context, req *handlers_ochrevector.StopJobRequest, _ string) (*handlers_ochrevector.StopJobResponse, error) {
|
||||
func (f *fakeBedrockAgentVectorService) StopJob(_ context.Context, req *ochrevector.StopJobRequest, _ string) (*ochrevector.StopJobResponse, error) {
|
||||
f.stopJobReq = req
|
||||
if f.stopJobErr != nil {
|
||||
return nil, f.stopJobErr
|
||||
@@ -122,12 +122,12 @@ func (f *fakeBedrockAgentVectorService) StopJob(_ context.Context, req *handlers
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
var _ handlers_ochrevector.VectorService = (*fakeBedrockAgentVectorService)(nil)
|
||||
var _ ochrevector.VectorService = (*fakeBedrockAgentVectorService)(nil)
|
||||
|
||||
func newBedrockAgentTestStores(t *testing.T) (*handlers_ochrevector.KBStore, *handlers_ochrevector.DataSourceStore) {
|
||||
func newBedrockAgentTestStores(t *testing.T) (*ochrevector.KBStore, *ochrevector.DataSourceStore) {
|
||||
t.Helper()
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
return handlers_ochrevector.NewKBStore(js), handlers_ochrevector.NewDataSourceStore(js)
|
||||
return ochrevector.NewKBStore(js), ochrevector.NewDataSourceStore(js)
|
||||
}
|
||||
|
||||
func TestEmbeddingModelIDFromARN(t *testing.T) {
|
||||
@@ -150,10 +150,10 @@ func TestBucketNameFromS3ARN_And_FormatS3BucketARN_RoundTrip(t *testing.T) {
|
||||
|
||||
func TestKBStatusToAWS(t *testing.T) {
|
||||
cases := []struct{ status, want string }{
|
||||
{handlers_ochrevector.StateReady, bedrockagent.KnowledgeBaseStatusActive},
|
||||
{handlers_ochrevector.StateCreating, bedrockagent.KnowledgeBaseStatusCreating},
|
||||
{handlers_ochrevector.StateDeleting, bedrockagent.KnowledgeBaseStatusDeleting},
|
||||
{handlers_ochrevector.StateStale, bedrockagent.KnowledgeBaseStatusFailed},
|
||||
{ochrevector.StateReady, bedrockagent.KnowledgeBaseStatusActive},
|
||||
{ochrevector.StateCreating, bedrockagent.KnowledgeBaseStatusCreating},
|
||||
{ochrevector.StateDeleting, bedrockagent.KnowledgeBaseStatusDeleting},
|
||||
{ochrevector.StateStale, bedrockagent.KnowledgeBaseStatusFailed},
|
||||
{"unknown", bedrockagent.KnowledgeBaseStatusFailed},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
@@ -162,17 +162,17 @@ func TestKBStatusToAWS(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestDataSourceStatusToAWS(t *testing.T) {
|
||||
assert.Equal(t, bedrockagent.DataSourceStatusDeleting, dataSourceStatusToAWS(handlers_ochrevector.StateDeleting))
|
||||
assert.Equal(t, bedrockagent.DataSourceStatusAvailable, dataSourceStatusToAWS(handlers_ochrevector.StateReady))
|
||||
assert.Equal(t, bedrockagent.DataSourceStatusDeleting, dataSourceStatusToAWS(ochrevector.StateDeleting))
|
||||
assert.Equal(t, bedrockagent.DataSourceStatusAvailable, dataSourceStatusToAWS(ochrevector.StateReady))
|
||||
}
|
||||
|
||||
func TestJobStateToAWS(t *testing.T) {
|
||||
cases := []struct{ state, want string }{
|
||||
{handlers_ochrevector.JobStatePending, bedrockagent.IngestionJobStatusStarting},
|
||||
{handlers_ochrevector.JobStateRunning, bedrockagent.IngestionJobStatusInProgress},
|
||||
{handlers_ochrevector.JobStateReady, bedrockagent.IngestionJobStatusComplete},
|
||||
{handlers_ochrevector.JobStateFailed, bedrockagent.IngestionJobStatusFailed},
|
||||
{handlers_ochrevector.JobStateStopped, ingestionJobStatusStopped},
|
||||
{ochrevector.JobStatePending, bedrockagent.IngestionJobStatusStarting},
|
||||
{ochrevector.JobStateRunning, bedrockagent.IngestionJobStatusInProgress},
|
||||
{ochrevector.JobStateReady, bedrockagent.IngestionJobStatusComplete},
|
||||
{ochrevector.JobStateFailed, bedrockagent.IngestionJobStatusFailed},
|
||||
{ochrevector.JobStateStopped, ingestionJobStatusStopped},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
assert.Equal(t, tc.want, jobStateToAWS(tc.state))
|
||||
@@ -181,9 +181,9 @@ func TestJobStateToAWS(t *testing.T) {
|
||||
|
||||
func TestTranslateVectorErr(t *testing.T) {
|
||||
assert.NoError(t, translateVectorErr(nil))
|
||||
assert.True(t, awserrors.IsErrorCode(translateVectorErr(handlers_ochrevector.ErrIndexNotFound), awserrors.ErrorResourceNotFoundException))
|
||||
assert.True(t, awserrors.IsErrorCode(translateVectorErr(handlers_ochrevector.ErrJobNotFound), awserrors.ErrorResourceNotFoundException))
|
||||
assert.True(t, awserrors.IsErrorCode(translateVectorErr(handlers_ochrevector.ErrIndexExists), awserrors.ErrorConflictException))
|
||||
assert.True(t, awserrors.IsErrorCode(translateVectorErr(ochrevector.ErrIndexNotFound), awserrors.ErrorResourceNotFoundException))
|
||||
assert.True(t, awserrors.IsErrorCode(translateVectorErr(ochrevector.ErrJobNotFound), awserrors.ErrorResourceNotFoundException))
|
||||
assert.True(t, awserrors.IsErrorCode(translateVectorErr(ochrevector.ErrIndexExists), awserrors.ErrorConflictException))
|
||||
}
|
||||
|
||||
// TestTranslateVectorErr_BackendDownMapsToServiceUnavailable proves a
|
||||
@@ -278,7 +278,7 @@ func TestCreateKnowledgeBase_RollbackFailureIsLoggedNotSwallowed(t *testing.T) {
|
||||
// A KBStore with no JetStream client makes kb.Create fail deterministically
|
||||
// (independent of the randomly generated knowledge base id), so the
|
||||
// rollback path is reached without needing to force a real id collision.
|
||||
kb := handlers_ochrevector.NewKBStore(nil)
|
||||
kb := ochrevector.NewKBStore(nil)
|
||||
vector := &fakeBedrockAgentVectorService{deleteIndexErr: errors.New("rollback backend unavailable")}
|
||||
|
||||
_, err := CreateKnowledgeBase(context.Background(), bedrockAgentTestAccount, "us-east-1", kb, vector, validCreateKBInput())
|
||||
@@ -301,9 +301,9 @@ func TestDeleteKnowledgeBase_CascadesDataSourcesAndDeletesBoundIndex(t *testing.
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: handlers_ochrevector.StateReady}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-2", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: ochrevector.StateReady}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-2", KnowledgeBaseID: "kb-1"}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{}
|
||||
_, err := DeleteKnowledgeBase(ctx, bedrockAgentTestAccount, kb, ds, vector, &bedrockagent.DeleteKnowledgeBaseInput{KnowledgeBaseId: aws.String("kb-1")})
|
||||
@@ -338,9 +338,9 @@ func TestDeleteKnowledgeBase_UnknownIDReturnsNotFound(t *testing.T) {
|
||||
func TestDeleteKnowledgeBase_TreatsMissingIndexAsAlreadyDeleted(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-gone", Status: handlers_ochrevector.StateReady}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-gone", Status: ochrevector.StateReady}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{deleteIndexErr: handlers_ochrevector.ErrIndexNotFound}
|
||||
vector := &fakeBedrockAgentVectorService{deleteIndexErr: ochrevector.ErrIndexNotFound}
|
||||
_, err := DeleteKnowledgeBase(ctx, bedrockAgentTestAccount, kb, ds, vector, &bedrockagent.DeleteKnowledgeBaseInput{KnowledgeBaseId: aws.String("kb-1")})
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -356,7 +356,7 @@ func TestDeleteKnowledgeBase_TreatsMissingIndexAsAlreadyDeleted(t *testing.T) {
|
||||
func TestDeleteKnowledgeBase_GenuineIndexDeleteErrorStillAborts(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: handlers_ochrevector.StateReady}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: ochrevector.StateReady}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{deleteIndexErr: errors.New("backend unavailable")}
|
||||
_, err := DeleteKnowledgeBase(ctx, bedrockAgentTestAccount, kb, ds, vector, &bedrockagent.DeleteKnowledgeBaseInput{KnowledgeBaseId: aws.String("kb-1")})
|
||||
@@ -386,8 +386,8 @@ func createDataSourceInput(chunking *bedrockagent.VectorIngestionConfiguration)
|
||||
func TestCreateDataSource_ChunkingUnitConversion(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{
|
||||
ID: "kb-1", IndexID: "idx-1", Status: handlers_ochrevector.StateReady,
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{
|
||||
ID: "kb-1", IndexID: "idx-1", Status: ochrevector.StateReady,
|
||||
EmbeddingModel: "titan", Dimension: 1024,
|
||||
}))
|
||||
|
||||
@@ -423,7 +423,7 @@ func TestCreateDataSource_ChunkingUnitConversion(t *testing.T) {
|
||||
func TestCreateDataSource_OmittedChunkingLeavesZeroValues(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: handlers_ochrevector.StateReady}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: ochrevector.StateReady}))
|
||||
|
||||
out, err := CreateDataSource(ctx, bedrockAgentTestAccount, "us-east-1", kb, ds, createDataSourceInput(nil))
|
||||
require.NoError(t, err)
|
||||
@@ -438,7 +438,7 @@ func TestCreateDataSource_OmittedChunkingLeavesZeroValues(t *testing.T) {
|
||||
func TestCreateDataSource_RejectsNonS3Type(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
|
||||
input := createDataSourceInput(nil)
|
||||
input.DataSourceConfiguration.Type = aws.String("WEB")
|
||||
@@ -459,14 +459,14 @@ func TestCreateDataSource_UnknownKnowledgeBaseIsNotFound(t *testing.T) {
|
||||
func TestStartIngestionJob_BuildsIngestRequestFromDataSourceAgainstBoundIndex(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: handlers_ochrevector.StateReady}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1", Status: ochrevector.StateReady}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{
|
||||
ID: "ds-1", KnowledgeBaseID: "kb-1",
|
||||
Source: handlers_ochrevector.SourceSpec{Bucket: "b1", Prefix: "p1"},
|
||||
Source: ochrevector.SourceSpec{Bucket: "b1", Prefix: "p1"},
|
||||
}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{ingestResp: handlers_ochrevector.IngestResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: handlers_ochrevector.JobStatePending},
|
||||
vector := &fakeBedrockAgentVectorService{ingestResp: ochrevector.IngestResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: ochrevector.JobStatePending},
|
||||
}}
|
||||
out, err := StartIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &bedrockagent.StartIngestionJobInput{
|
||||
KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-1"),
|
||||
@@ -484,14 +484,14 @@ func TestStartIngestionJob_BuildsIngestRequestFromDataSourceAgainstBoundIndex(t
|
||||
func TestGetIngestionJob_RejectsJobFromForeignIndex(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{
|
||||
ID: "ds-1", KnowledgeBaseID: "kb-1",
|
||||
Source: handlers_ochrevector.SourceSpec{Bucket: "b1", Prefix: "p1"},
|
||||
Source: ochrevector.SourceSpec{Bucket: "b1", Prefix: "p1"},
|
||||
}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: handlers_ochrevector.DescribeJobResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-other", DataSourceID: "ds-1"},
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: ochrevector.DescribeJobResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", IndexID: "idx-other", DataSourceID: "ds-1"},
|
||||
}}
|
||||
_, err := GetIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &bedrockagent.GetIngestionJobInput{
|
||||
KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-1"), IngestionJobId: aws.String("job-1"),
|
||||
@@ -510,14 +510,14 @@ func TestGetIngestionJob_RejectsJobFromForeignIndex(t *testing.T) {
|
||||
func TestGetIngestionJob_RejectsWrongDataSource(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-a", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-b", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-a", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-b", KnowledgeBaseID: "kb-1"}))
|
||||
|
||||
// job-1 really belongs to ds-a (its exact DataSourceID), but the request
|
||||
// below addresses it through ds-b's path.
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: handlers_ochrevector.DescribeJobResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-a"},
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: ochrevector.DescribeJobResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-a"},
|
||||
}}
|
||||
_, err := GetIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &bedrockagent.GetIngestionJobInput{
|
||||
KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-b"), IngestionJobId: aws.String("job-1"),
|
||||
@@ -540,11 +540,11 @@ func TestGetIngestionJob_RejectsWrongDataSource(t *testing.T) {
|
||||
func TestGetIngestionJob_RejectsEmptyDataSourceID(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: handlers_ochrevector.DescribeJobResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: ""},
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: ochrevector.DescribeJobResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: ""},
|
||||
}}
|
||||
_, err := GetIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &bedrockagent.GetIngestionJobInput{
|
||||
KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-1"), IngestionJobId: aws.String("job-1"),
|
||||
@@ -556,12 +556,12 @@ func TestGetIngestionJob_RejectsEmptyDataSourceID(t *testing.T) {
|
||||
func TestListIngestionJobs_FiltersToBoundIndexAndExactDataSourceID(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
|
||||
now := time.Now().UTC()
|
||||
vector := &fakeBedrockAgentVectorService{listJobsResp: handlers_ochrevector.ListJobsResponse{Jobs: []handlers_ochrevector.JobRecord{
|
||||
{ID: "job-match", IndexID: "idx-1", DataSourceID: "ds-1", State: handlers_ochrevector.JobStateReady, CreatedAt: now, UpdatedAt: now},
|
||||
vector := &fakeBedrockAgentVectorService{listJobsResp: ochrevector.ListJobsResponse{Jobs: []ochrevector.JobRecord{
|
||||
{ID: "job-match", IndexID: "idx-1", DataSourceID: "ds-1", State: ochrevector.JobStateReady, CreatedAt: now, UpdatedAt: now},
|
||||
{ID: "job-wrong-index", IndexID: "idx-2", DataSourceID: "ds-1", CreatedAt: now, UpdatedAt: now},
|
||||
{ID: "job-wrong-datasource", IndexID: "idx-1", DataSourceID: "ds-other", CreatedAt: now, UpdatedAt: now},
|
||||
{ID: "job-no-datasource", IndexID: "idx-1", DataSourceID: "", CreatedAt: now, UpdatedAt: now},
|
||||
@@ -582,15 +582,15 @@ func TestListIngestionJobs_FiltersToBoundIndexAndExactDataSourceID(t *testing.T)
|
||||
func TestStopIngestionJob_CancelsBoundJobAndReturnsStoppedStatus(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{
|
||||
describeJobResp: handlers_ochrevector.DescribeJobResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: handlers_ochrevector.JobStateRunning},
|
||||
describeJobResp: ochrevector.DescribeJobResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: ochrevector.JobStateRunning},
|
||||
},
|
||||
stopJobResp: handlers_ochrevector.StopJobResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: handlers_ochrevector.JobStateStopped},
|
||||
stopJobResp: ochrevector.StopJobResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: ochrevector.JobStateStopped},
|
||||
},
|
||||
}
|
||||
out, err := StopIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &StopIngestionJobInput{
|
||||
@@ -609,11 +609,11 @@ func TestStopIngestionJob_CancelsBoundJobAndReturnsStoppedStatus(t *testing.T) {
|
||||
func TestStopIngestionJob_RejectsJobFromForeignIndex(t *testing.T) {
|
||||
kb, ds := newBedrockAgentTestStores(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"}))
|
||||
require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"}))
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: handlers_ochrevector.DescribeJobResponse{
|
||||
Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-other", DataSourceID: "ds-1"},
|
||||
vector := &fakeBedrockAgentVectorService{describeJobResp: ochrevector.DescribeJobResponse{
|
||||
Job: ochrevector.JobRecord{ID: "job-1", IndexID: "idx-other", DataSourceID: "ds-1"},
|
||||
}}
|
||||
_, err := StopIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &StopIngestionJobInput{
|
||||
KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-1"), IngestionJobId: aws.String("job-1"),
|
||||
|
||||
@@ -13,9 +13,9 @@ import (
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockagentruntime"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockruntime"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
)
|
||||
|
||||
// converseFn is gateway_bedrock.Converse partially applied over the request's
|
||||
@@ -37,14 +37,14 @@ type bedrockAgentRuntimeRoute = restRoute[bedrockAgentRuntimeRouteHandler]
|
||||
// gw.BedrockAgentVector, the same stores bedrock-agent's control plane uses.
|
||||
// converse reaches gateway_bedrock.Converse for RetrieveAndGenerate's
|
||||
// generation step.
|
||||
type bedrockAgentRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, kb *handlers_ochrevector.KBStore, vector handlers_ochrevector.VectorService, converse converseFn) (any, error)
|
||||
type bedrockAgentRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, kb *ochrevector.KBStore, vector ochrevector.VectorService, converse converseFn) (any, error)
|
||||
|
||||
// bedrockAgentRuntimeRoutes is the dispatch table. Real AWS HTTP paths/methods,
|
||||
// verified against the vendored aws-sdk-go bedrockagentruntime request
|
||||
// definitions (opRetrieve/opRetrieveAndGenerate).
|
||||
var bedrockAgentRuntimeRoutes = []bedrockAgentRuntimeRoute{
|
||||
{"POST", "/knowledgebases/{knowledgeBaseId}/retrieve", "Retrieve",
|
||||
func(ctx context.Context, acct string, p []string, b []byte, kb *handlers_ochrevector.KBStore, vector handlers_ochrevector.VectorService, _ converseFn) (any, error) {
|
||||
func(ctx context.Context, acct string, p []string, b []byte, kb *ochrevector.KBStore, vector ochrevector.VectorService, _ converseFn) (any, error) {
|
||||
input := new(bedrockagentruntime.RetrieveInput)
|
||||
if len(b) > 0 {
|
||||
if err := json.Unmarshal(b, input); err != nil {
|
||||
@@ -55,7 +55,7 @@ var bedrockAgentRuntimeRoutes = []bedrockAgentRuntimeRoute{
|
||||
return Retrieve(ctx, acct, kb, vector, b, input)
|
||||
}},
|
||||
{"POST", "/retrieveAndGenerate", "RetrieveAndGenerate",
|
||||
func(ctx context.Context, acct string, _ []string, b []byte, kb *handlers_ochrevector.KBStore, vector handlers_ochrevector.VectorService, converse converseFn) (any, error) {
|
||||
func(ctx context.Context, acct string, _ []string, b []byte, kb *ochrevector.KBStore, vector ochrevector.VectorService, converse converseFn) (any, error) {
|
||||
input := new(bedrockagentruntime.RetrieveAndGenerateInput)
|
||||
if len(b) > 0 {
|
||||
if err := json.Unmarshal(b, input); err != nil {
|
||||
@@ -187,73 +187,73 @@ func (f *wireFilter) isZero() bool {
|
||||
// one operator per node); an over-specified node resolves the first non-nil
|
||||
// field in the fixed order below rather than erroring, consistent with how
|
||||
// the rest of this gateway treats an over-specified request as best-effort.
|
||||
func (f *wireFilter) toFilter() (*handlers_ochrevector.Filter, error) {
|
||||
func (f *wireFilter) toFilter() (*ochrevector.Filter, error) {
|
||||
if f.isZero() {
|
||||
return nil, nil
|
||||
}
|
||||
switch {
|
||||
case f.Equals != nil:
|
||||
return handlers_ochrevector.Equals(f.Equals.Key, f.Equals.Value), nil
|
||||
return ochrevector.Equals(f.Equals.Key, f.Equals.Value), nil
|
||||
case f.NotEquals != nil:
|
||||
return handlers_ochrevector.NotEquals(f.NotEquals.Key, f.NotEquals.Value), nil
|
||||
return ochrevector.NotEquals(f.NotEquals.Key, f.NotEquals.Value), nil
|
||||
case f.GreaterThan != nil:
|
||||
return handlers_ochrevector.GreaterThan(f.GreaterThan.Key, f.GreaterThan.Value), nil
|
||||
return ochrevector.GreaterThan(f.GreaterThan.Key, f.GreaterThan.Value), nil
|
||||
case f.GreaterThanOrEquals != nil:
|
||||
return handlers_ochrevector.GreaterThanOrEquals(f.GreaterThanOrEquals.Key, f.GreaterThanOrEquals.Value), nil
|
||||
return ochrevector.GreaterThanOrEquals(f.GreaterThanOrEquals.Key, f.GreaterThanOrEquals.Value), nil
|
||||
case f.LessThan != nil:
|
||||
return handlers_ochrevector.LessThan(f.LessThan.Key, f.LessThan.Value), nil
|
||||
return ochrevector.LessThan(f.LessThan.Key, f.LessThan.Value), nil
|
||||
case f.LessThanOrEquals != nil:
|
||||
return handlers_ochrevector.LessThanOrEquals(f.LessThanOrEquals.Key, f.LessThanOrEquals.Value), nil
|
||||
return ochrevector.LessThanOrEquals(f.LessThanOrEquals.Key, f.LessThanOrEquals.Value), nil
|
||||
case f.In != nil:
|
||||
values, err := wireFilterStringSlice(f.In.Value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("bedrock-agent-runtime: filter \"in\" on %q: %w", f.In.Key, err)
|
||||
}
|
||||
return handlers_ochrevector.In(f.In.Key, values), nil
|
||||
return ochrevector.In(f.In.Key, values), nil
|
||||
case f.NotIn != nil:
|
||||
values, err := wireFilterStringSlice(f.NotIn.Value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("bedrock-agent-runtime: filter \"notIn\" on %q: %w", f.NotIn.Key, err)
|
||||
}
|
||||
return handlers_ochrevector.NotIn(f.NotIn.Key, values), nil
|
||||
return ochrevector.NotIn(f.NotIn.Key, values), nil
|
||||
case f.StartsWith != nil:
|
||||
prefix, ok := f.StartsWith.Value.(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("bedrock-agent-runtime: filter \"startsWith\" on %q requires a string value", f.StartsWith.Key)
|
||||
}
|
||||
return handlers_ochrevector.StartsWith(f.StartsWith.Key, prefix), nil
|
||||
return ochrevector.StartsWith(f.StartsWith.Key, prefix), nil
|
||||
case f.StringContains != nil:
|
||||
substr, ok := f.StringContains.Value.(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("bedrock-agent-runtime: filter \"stringContains\" on %q requires a string value", f.StringContains.Key)
|
||||
}
|
||||
return handlers_ochrevector.StringContains(f.StringContains.Key, substr), nil
|
||||
return ochrevector.StringContains(f.StringContains.Key, substr), nil
|
||||
case f.ListContains != nil:
|
||||
value, ok := f.ListContains.Value.(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("bedrock-agent-runtime: filter \"listContains\" on %q requires a string value", f.ListContains.Key)
|
||||
}
|
||||
return handlers_ochrevector.ListContains(f.ListContains.Key, value), nil
|
||||
return ochrevector.ListContains(f.ListContains.Key, value), nil
|
||||
case len(f.AndAll) > 0:
|
||||
children, err := wireFiltersToChildren(f.AndAll)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return handlers_ochrevector.AndAll(children...), nil
|
||||
return ochrevector.AndAll(children...), nil
|
||||
default: // len(f.OrAll) > 0, the only remaining non-zero case.
|
||||
children, err := wireFiltersToChildren(f.OrAll)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return handlers_ochrevector.OrAll(children...), nil
|
||||
return ochrevector.OrAll(children...), nil
|
||||
}
|
||||
}
|
||||
|
||||
// wireFiltersToChildren translates a combinator's child list, rejecting any
|
||||
// child that itself carries no operator (a malformed request, not a filter
|
||||
// that matches nothing).
|
||||
func wireFiltersToChildren(fs []wireFilter) ([]*handlers_ochrevector.Filter, error) {
|
||||
children := make([]*handlers_ochrevector.Filter, 0, len(fs))
|
||||
func wireFiltersToChildren(fs []wireFilter) ([]*ochrevector.Filter, error) {
|
||||
children := make([]*ochrevector.Filter, 0, len(fs))
|
||||
for i := range fs {
|
||||
child, err := fs[i].toFilter()
|
||||
if err != nil {
|
||||
@@ -299,7 +299,7 @@ type retrieveFilterEnvelope struct {
|
||||
|
||||
// decodeRetrieveFilter extracts and translates Retrieve's filter from the raw
|
||||
// request body, returning (nil, nil) when no filter was sent.
|
||||
func decodeRetrieveFilter(body []byte) (*handlers_ochrevector.Filter, error) {
|
||||
func decodeRetrieveFilter(body []byte) (*ochrevector.Filter, error) {
|
||||
if len(body) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -331,7 +331,7 @@ type retrieveAndGenerateFilterEnvelope struct {
|
||||
|
||||
// decodeRetrieveAndGenerateFilter is decodeRetrieveFilter's sibling for
|
||||
// RetrieveAndGenerate.
|
||||
func decodeRetrieveAndGenerateFilter(body []byte) (*handlers_ochrevector.Filter, error) {
|
||||
func decodeRetrieveAndGenerateFilter(body []byte) (*ochrevector.Filter, error) {
|
||||
if len(body) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -355,7 +355,7 @@ func decodeRetrieveAndGenerateFilter(body []byte) (*handlers_ochrevector.Filter,
|
||||
// to put it in (same SDK-vintage gap as wireFilter, but here it only drops an
|
||||
// echo of already-visible data rather than silently changing behaviour, so it
|
||||
// is left unset rather than worked around).
|
||||
func queryResultToRetrievalResult(r handlers_ochrevector.QueryResult) *bedrockagentruntime.KnowledgeBaseRetrievalResult {
|
||||
func queryResultToRetrievalResult(r ochrevector.QueryResult) *bedrockagentruntime.KnowledgeBaseRetrievalResult {
|
||||
return &bedrockagentruntime.KnowledgeBaseRetrievalResult{
|
||||
Content: &bedrockagentruntime.RetrievalResultContent{Text: aws.String(r.Chunk)},
|
||||
Location: &bedrockagentruntime.RetrievalResultLocation{
|
||||
@@ -377,7 +377,7 @@ const defaultRetrieveResults = 5
|
||||
// response, and an incoming NextToken is accepted and ignored rather than
|
||||
// rejected, matching how storageConfiguration/roleArn are accepted-and-
|
||||
// stubbed elsewhere in this gateway.
|
||||
func Retrieve(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, vector handlers_ochrevector.VectorService, body []byte, input *bedrockagentruntime.RetrieveInput) (*bedrockagentruntime.RetrieveOutput, error) {
|
||||
func Retrieve(ctx context.Context, accountID string, kb *ochrevector.KBStore, vector ochrevector.VectorService, body []byte, input *bedrockagentruntime.RetrieveInput) (*bedrockagentruntime.RetrieveOutput, error) {
|
||||
if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || input.RetrievalQuery == nil || aws.StringValue(input.RetrievalQuery.Text) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -403,7 +403,7 @@ func Retrieve(ctx context.Context, accountID string, kb *handlers_ochrevector.KB
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
|
||||
resp, err := vector.Query(ctx, &handlers_ochrevector.QueryRequest{
|
||||
resp, err := vector.Query(ctx, &ochrevector.QueryRequest{
|
||||
IndexID: kbRec.IndexID,
|
||||
Text: aws.StringValue(input.RetrievalQuery.Text),
|
||||
K: numResults,
|
||||
@@ -468,7 +468,7 @@ func converseOutputText(out *bedrockruntime.ConverseOutput) string {
|
||||
// the retrieved chunks. D6: sessionId is server-generated when the caller
|
||||
// omits it, and echoed either way -- no conversational memory is persisted,
|
||||
// so a reused sessionId has no effect on this or any later call.
|
||||
func RetrieveAndGenerate(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, vector handlers_ochrevector.VectorService, converse converseFn, body []byte, input *bedrockagentruntime.RetrieveAndGenerateInput) (*bedrockagentruntime.RetrieveAndGenerateOutput, error) {
|
||||
func RetrieveAndGenerate(ctx context.Context, accountID string, kb *ochrevector.KBStore, vector ochrevector.VectorService, converse converseFn, body []byte, input *bedrockagentruntime.RetrieveAndGenerateInput) (*bedrockagentruntime.RetrieveAndGenerateOutput, error) {
|
||||
if input == nil || input.Input == nil || aws.StringValue(input.Input.Text) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
@@ -505,7 +505,7 @@ func RetrieveAndGenerate(ctx context.Context, accountID string, kb *handlers_och
|
||||
}
|
||||
|
||||
queryText := aws.StringValue(input.Input.Text)
|
||||
queryResp, err := vector.Query(ctx, &handlers_ochrevector.QueryRequest{
|
||||
queryResp, err := vector.Query(ctx, &ochrevector.QueryRequest{
|
||||
IndexID: kbRec.IndexID,
|
||||
Text: queryText,
|
||||
K: numResults,
|
||||
|
||||
@@ -16,8 +16,8 @@ import (
|
||||
|
||||
"github.com/aws/aws-sdk-go/service/bedrockagentruntime"
|
||||
"github.com/mulgadc/spinifex/internal/testkit"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -27,13 +27,13 @@ import (
|
||||
// stores, allow-everything IAMService) plus newBedrockRequestGateway's vLLM
|
||||
// routing (BedrockEndpoints/BedrockAccess), since RetrieveAndGenerate's
|
||||
// generation step reaches gateway_bedrock.Converse in-process.
|
||||
func newBedrockAgentRuntimeRequestGateway(t *testing.T, vector handlers_ochrevector.VectorService, vllmURL string) *GatewayConfig {
|
||||
func newBedrockAgentRuntimeRequestGateway(t *testing.T, vector ochrevector.VectorService, vllmURL string) *GatewayConfig {
|
||||
t.Helper()
|
||||
_, _, js := testutil.StartTestJetStream(t)
|
||||
return &GatewayConfig{
|
||||
IAMService: allowAllIAMService(),
|
||||
BedrockAgentKB: handlers_ochrevector.NewKBStore(js),
|
||||
BedrockAgentDataSources: handlers_ochrevector.NewDataSourceStore(js),
|
||||
BedrockAgentKB: ochrevector.NewKBStore(js),
|
||||
BedrockAgentDataSources: ochrevector.NewDataSourceStore(js),
|
||||
BedrockAgentVector: vector,
|
||||
BedrockEndpoints: map[string]string{
|
||||
bedrockTestLlamaModelID: vllmURL,
|
||||
@@ -52,16 +52,16 @@ func bedrockAgentRuntimeRequestWithAccount(method, path, body string) *http.Requ
|
||||
|
||||
func newTestKnowledgeBaseRecord(t *testing.T, gw *GatewayConfig, kbID, indexID string) {
|
||||
t.Helper()
|
||||
require.NoError(t, gw.BedrockAgentKB.Create(context.Background(), bedrockAgentTestAccount, handlers_ochrevector.KBRecord{
|
||||
ID: kbID, Name: "docs", Status: handlers_ochrevector.StateReady,
|
||||
require.NoError(t, gw.BedrockAgentKB.Create(context.Background(), bedrockAgentTestAccount, ochrevector.KBRecord{
|
||||
ID: kbID, Name: "docs", Status: ochrevector.StateReady,
|
||||
EmbeddingModel: "amazon.titan-embed-text-v2:0", Dimension: 1024, IndexID: indexID,
|
||||
}))
|
||||
}
|
||||
|
||||
func TestBedrockAgentRuntimeRequest_Retrieve_HappyPath(t *testing.T) {
|
||||
ts := newVLLMStub(t)
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: handlers_ochrevector.QueryResponse{
|
||||
Results: []handlers_ochrevector.QueryResult{
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: ochrevector.QueryResponse{
|
||||
Results: []ochrevector.QueryResult{
|
||||
{Chunk: "the dragon lives in a cave", SourceKey: "docs/a.txt", Score: 0.9},
|
||||
},
|
||||
}}
|
||||
@@ -102,7 +102,7 @@ func TestBedrockAgentRuntimeRequest_Retrieve_WithFilter(t *testing.T) {
|
||||
|
||||
require.NotNil(t, vector.queryReq)
|
||||
require.NotNil(t, vector.queryReq.Filter)
|
||||
assert.Equal(t, handlers_ochrevector.FilterEquals, vector.queryReq.Filter.Op)
|
||||
assert.Equal(t, ochrevector.FilterEquals, vector.queryReq.Filter.Op)
|
||||
assert.Equal(t, "genre", vector.queryReq.Filter.Key)
|
||||
}
|
||||
|
||||
@@ -131,8 +131,8 @@ func TestBedrockAgentRuntimeRequest_Retrieve_MalformedBodyReturnsValidationExcep
|
||||
|
||||
func TestBedrockAgentRuntimeRequest_RetrieveAndGenerate_HappyPath(t *testing.T) {
|
||||
ts := newVLLMStub(t)
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: handlers_ochrevector.QueryResponse{
|
||||
Results: []handlers_ochrevector.QueryResult{
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: ochrevector.QueryResponse{
|
||||
Results: []ochrevector.QueryResult{
|
||||
{Chunk: "the dragon lives in a cave", SourceKey: "docs/a.txt", Score: 0.9},
|
||||
},
|
||||
}}
|
||||
|
||||
@@ -14,8 +14,8 @@ import (
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockagentruntime"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockruntime"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -24,19 +24,19 @@ func TestWireFilter_ToFilter_Leaves(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
filter wireFilter
|
||||
wantOp handlers_ochrevector.FilterOp
|
||||
wantOp ochrevector.FilterOp
|
||||
wantK string
|
||||
wantV any
|
||||
}{
|
||||
{"equals", wireFilter{Equals: &wireFilterAttr{Key: "genre", Value: "fiction"}}, handlers_ochrevector.FilterEquals, "genre", "fiction"},
|
||||
{"notEquals", wireFilter{NotEquals: &wireFilterAttr{Key: "genre", Value: "fiction"}}, handlers_ochrevector.FilterNotEquals, "genre", "fiction"},
|
||||
{"greaterThan", wireFilter{GreaterThan: &wireFilterAttr{Key: "year", Value: 2000.0}}, handlers_ochrevector.FilterGreaterThan, "year", 2000.0},
|
||||
{"greaterThanOrEquals", wireFilter{GreaterThanOrEquals: &wireFilterAttr{Key: "year", Value: 2000.0}}, handlers_ochrevector.FilterGreaterThanOrEqual, "year", 2000.0},
|
||||
{"lessThan", wireFilter{LessThan: &wireFilterAttr{Key: "year", Value: 2000.0}}, handlers_ochrevector.FilterLessThan, "year", 2000.0},
|
||||
{"lessThanOrEquals", wireFilter{LessThanOrEquals: &wireFilterAttr{Key: "year", Value: 2000.0}}, handlers_ochrevector.FilterLessThanOrEqual, "year", 2000.0},
|
||||
{"startsWith", wireFilter{StartsWith: &wireFilterAttr{Key: "title", Value: "The "}}, handlers_ochrevector.FilterStartsWith, "title", "The "},
|
||||
{"stringContains", wireFilter{StringContains: &wireFilterAttr{Key: "title", Value: "dragon"}}, handlers_ochrevector.FilterStringContains, "title", "dragon"},
|
||||
{"listContains", wireFilter{ListContains: &wireFilterAttr{Key: "tags", Value: "scifi"}}, handlers_ochrevector.FilterListContains, "tags", "scifi"},
|
||||
{"equals", wireFilter{Equals: &wireFilterAttr{Key: "genre", Value: "fiction"}}, ochrevector.FilterEquals, "genre", "fiction"},
|
||||
{"notEquals", wireFilter{NotEquals: &wireFilterAttr{Key: "genre", Value: "fiction"}}, ochrevector.FilterNotEquals, "genre", "fiction"},
|
||||
{"greaterThan", wireFilter{GreaterThan: &wireFilterAttr{Key: "year", Value: 2000.0}}, ochrevector.FilterGreaterThan, "year", 2000.0},
|
||||
{"greaterThanOrEquals", wireFilter{GreaterThanOrEquals: &wireFilterAttr{Key: "year", Value: 2000.0}}, ochrevector.FilterGreaterThanOrEqual, "year", 2000.0},
|
||||
{"lessThan", wireFilter{LessThan: &wireFilterAttr{Key: "year", Value: 2000.0}}, ochrevector.FilterLessThan, "year", 2000.0},
|
||||
{"lessThanOrEquals", wireFilter{LessThanOrEquals: &wireFilterAttr{Key: "year", Value: 2000.0}}, ochrevector.FilterLessThanOrEqual, "year", 2000.0},
|
||||
{"startsWith", wireFilter{StartsWith: &wireFilterAttr{Key: "title", Value: "The "}}, ochrevector.FilterStartsWith, "title", "The "},
|
||||
{"stringContains", wireFilter{StringContains: &wireFilterAttr{Key: "title", Value: "dragon"}}, ochrevector.FilterStringContains, "title", "dragon"},
|
||||
{"listContains", wireFilter{ListContains: &wireFilterAttr{Key: "tags", Value: "scifi"}}, ochrevector.FilterListContains, "tags", "scifi"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
@@ -54,13 +54,13 @@ func TestWireFilter_ToFilter_InNotIn(t *testing.T) {
|
||||
f := wireFilter{In: &wireFilterAttr{Key: "genre", Value: []any{"fiction", "scifi"}}}
|
||||
got, err := f.toFilter()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, handlers_ochrevector.FilterIn, got.Op)
|
||||
assert.Equal(t, ochrevector.FilterIn, got.Op)
|
||||
assert.Equal(t, []string{"fiction", "scifi"}, got.Value)
|
||||
|
||||
f = wireFilter{NotIn: &wireFilterAttr{Key: "genre", Value: []any{"horror"}}}
|
||||
got, err = f.toFilter()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, handlers_ochrevector.FilterNotIn, got.Op)
|
||||
assert.Equal(t, ochrevector.FilterNotIn, got.Op)
|
||||
assert.Equal(t, []string{"horror"}, got.Value)
|
||||
}
|
||||
|
||||
@@ -77,10 +77,10 @@ func TestWireFilter_ToFilter_AndOrAll(t *testing.T) {
|
||||
}}
|
||||
got, err := f.toFilter()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, handlers_ochrevector.FilterAndAll, got.Op)
|
||||
assert.Equal(t, ochrevector.FilterAndAll, got.Op)
|
||||
require.Len(t, got.Children, 2)
|
||||
assert.Equal(t, handlers_ochrevector.FilterEquals, got.Children[0].Op)
|
||||
assert.Equal(t, handlers_ochrevector.FilterGreaterThan, got.Children[1].Op)
|
||||
assert.Equal(t, ochrevector.FilterEquals, got.Children[0].Op)
|
||||
assert.Equal(t, ochrevector.FilterGreaterThan, got.Children[1].Op)
|
||||
|
||||
f = wireFilter{OrAll: []wireFilter{
|
||||
{Equals: &wireFilterAttr{Key: "genre", Value: "fiction"}},
|
||||
@@ -88,7 +88,7 @@ func TestWireFilter_ToFilter_AndOrAll(t *testing.T) {
|
||||
}}
|
||||
got, err = f.toFilter()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, handlers_ochrevector.FilterOrAll, got.Op)
|
||||
assert.Equal(t, ochrevector.FilterOrAll, got.Op)
|
||||
require.Len(t, got.Children, 2)
|
||||
}
|
||||
|
||||
@@ -103,7 +103,7 @@ func TestWireFilter_ToFilter_NestedCombinator(t *testing.T) {
|
||||
got, err := f.toFilter()
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got.Children, 2)
|
||||
assert.Equal(t, handlers_ochrevector.FilterOrAll, got.Children[1].Op)
|
||||
assert.Equal(t, ochrevector.FilterOrAll, got.Children[1].Op)
|
||||
require.Len(t, got.Children[1].Children, 2)
|
||||
}
|
||||
|
||||
@@ -139,7 +139,7 @@ func TestDecodeRetrieveFilter(t *testing.T) {
|
||||
f, err := decodeRetrieveFilter(body)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, f)
|
||||
assert.Equal(t, handlers_ochrevector.FilterEquals, f.Op)
|
||||
assert.Equal(t, ochrevector.FilterEquals, f.Op)
|
||||
assert.Equal(t, "genre", f.Key)
|
||||
}
|
||||
|
||||
@@ -175,7 +175,7 @@ func TestDecodeRetrieveAndGenerateFilter(t *testing.T) {
|
||||
f, err := decodeRetrieveAndGenerateFilter(body)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, f)
|
||||
assert.Equal(t, handlers_ochrevector.FilterAndAll, f.Op)
|
||||
assert.Equal(t, ochrevector.FilterAndAll, f.Op)
|
||||
require.Len(t, f.Children, 2)
|
||||
}
|
||||
|
||||
@@ -203,7 +203,7 @@ func TestConverseOutputText(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestQueryResultToRetrievalResult(t *testing.T) {
|
||||
r := handlers_ochrevector.QueryResult{Chunk: "some text", SourceKey: "docs/a.txt", Score: 0.87}
|
||||
r := ochrevector.QueryResult{Chunk: "some text", SourceKey: "docs/a.txt", Score: 0.87}
|
||||
got := queryResultToRetrievalResult(r)
|
||||
assert.Equal(t, "some text", *got.Content.Text)
|
||||
assert.Equal(t, bedrockagentruntime.RetrievalResultLocationTypeS3, *got.Location.Type)
|
||||
@@ -211,10 +211,10 @@ func TestQueryResultToRetrievalResult(t *testing.T) {
|
||||
assert.InDelta(t, 0.87, *got.Score, 0.0001)
|
||||
}
|
||||
|
||||
func newTestKBRecord(t *testing.T, kb *handlers_ochrevector.KBStore, accountID, kbID, indexID string) {
|
||||
func newTestKBRecord(t *testing.T, kb *ochrevector.KBStore, accountID, kbID, indexID string) {
|
||||
t.Helper()
|
||||
require.NoError(t, kb.Create(context.Background(), accountID, handlers_ochrevector.KBRecord{
|
||||
ID: kbID, Name: "docs", Status: handlers_ochrevector.StateReady,
|
||||
require.NoError(t, kb.Create(context.Background(), accountID, ochrevector.KBRecord{
|
||||
ID: kbID, Name: "docs", Status: ochrevector.StateReady,
|
||||
EmbeddingModel: "amazon.titan-embed-text-v2:0", Dimension: 1024, IndexID: indexID,
|
||||
}))
|
||||
}
|
||||
@@ -223,8 +223,8 @@ func TestRetrieve_MapsRequestAndResponse(t *testing.T) {
|
||||
kb, _ := newBedrockAgentTestStores(t)
|
||||
newTestKBRecord(t, kb, bedrockAgentTestAccount, "kb-1", "idx-1")
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: handlers_ochrevector.QueryResponse{
|
||||
Results: []handlers_ochrevector.QueryResult{
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: ochrevector.QueryResponse{
|
||||
Results: []ochrevector.QueryResult{
|
||||
{Chunk: "chunk a", SourceKey: "docs/a.txt", Score: 0.9},
|
||||
{Chunk: "chunk b", SourceKey: "docs/b.txt", Score: 0.5},
|
||||
},
|
||||
@@ -254,7 +254,7 @@ func TestRetrieve_MapsRequestAndResponse(t *testing.T) {
|
||||
assert.Equal(t, "dragons", vector.queryReq.Text)
|
||||
assert.Equal(t, 2, vector.queryReq.K)
|
||||
require.NotNil(t, vector.queryReq.Filter)
|
||||
assert.Equal(t, handlers_ochrevector.FilterEquals, vector.queryReq.Filter.Op)
|
||||
assert.Equal(t, ochrevector.FilterEquals, vector.queryReq.Filter.Op)
|
||||
}
|
||||
|
||||
func TestRetrieve_UnknownKnowledgeBase(t *testing.T) {
|
||||
@@ -310,8 +310,8 @@ func TestRetrieveAndGenerate_HappyPath(t *testing.T) {
|
||||
kb, _ := newBedrockAgentTestStores(t)
|
||||
newTestKBRecord(t, kb, bedrockAgentTestAccount, "kb-1", "idx-1")
|
||||
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: handlers_ochrevector.QueryResponse{
|
||||
Results: []handlers_ochrevector.QueryResult{{Chunk: "the dragon lives in a cave", SourceKey: "docs/a.txt", Score: 0.9}},
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: ochrevector.QueryResponse{
|
||||
Results: []ochrevector.QueryResult{{Chunk: "the dragon lives in a cave", SourceKey: "docs/a.txt", Score: 0.9}},
|
||||
}}
|
||||
fc := &fakeConverse{resp: fakeConverseOutput("The dragon lives in a cave.")}
|
||||
|
||||
@@ -366,8 +366,8 @@ func TestRetrieveAndGenerate_EchoesCallerSessionID(t *testing.T) {
|
||||
func TestRetrieveAndGenerate_UsesCustomPromptTemplate(t *testing.T) {
|
||||
kb, _ := newBedrockAgentTestStores(t)
|
||||
newTestKBRecord(t, kb, bedrockAgentTestAccount, "kb-1", "idx-1")
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: handlers_ochrevector.QueryResponse{
|
||||
Results: []handlers_ochrevector.QueryResult{{Chunk: "ctx chunk", SourceKey: "k", Score: 0.5}},
|
||||
vector := &fakeBedrockAgentVectorService{queryResp: ochrevector.QueryResponse{
|
||||
Results: []ochrevector.QueryResult{{Chunk: "ctx chunk", SourceKey: "k", Score: 0.5}},
|
||||
}}
|
||||
fc := &fakeConverse{resp: fakeConverseOutput("answer")}
|
||||
|
||||
|
||||
@@ -32,6 +32,7 @@ import (
|
||||
ecrauth "github.com/mulgadc/spinifex/spinifex/domains/ecr/auth"
|
||||
awsapi "github.com/mulgadc/spinifex/spinifex/domains/ecr/awsapi"
|
||||
ecrregistry "github.com/mulgadc/spinifex/spinifex/domains/ecr/registry"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/aws/policy"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/telemetry"
|
||||
@@ -40,7 +41,6 @@ import (
|
||||
gateway_eks "github.com/mulgadc/spinifex/spinifex/gateway/eks"
|
||||
gateway_sts "github.com/mulgadc/spinifex/spinifex/gateway/sts"
|
||||
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota"
|
||||
handlers_sts "github.com/mulgadc/spinifex/spinifex/handlers/sts"
|
||||
ingresshttp "github.com/mulgadc/spinifex/spinifex/ingress/http"
|
||||
@@ -234,14 +234,14 @@ type GatewayConfig struct {
|
||||
// not the daemon-owned vector engine). Nil for either fails
|
||||
// BedrockAgent_Request with ServerInternal rather than panicking, the same
|
||||
// as an unconfigured gw.NATSConn does for every other service.
|
||||
BedrockAgentKB *handlers_ochrevector.KBStore
|
||||
BedrockAgentDataSources *handlers_ochrevector.DataSourceStore
|
||||
BedrockAgentKB *ochrevector.KBStore
|
||||
BedrockAgentDataSources *ochrevector.DataSourceStore
|
||||
// BedrockAgentVector forwards CreateIndex/DeleteIndex/Ingest/DescribeJob/
|
||||
// ListJobs calls to .9's daemon-side VectorService over NATS
|
||||
// (handlers_ochrevector.NewNATSVectorService). It is the interface, not
|
||||
// (ochrevector.NewNATSVectorService). It is the interface, not
|
||||
// the concrete client, so a test can inject a fake without a live NATS
|
||||
// connection.
|
||||
BedrockAgentVector handlers_ochrevector.VectorService
|
||||
BedrockAgentVector ochrevector.VectorService
|
||||
}
|
||||
|
||||
var supportedServices = map[string]bool{
|
||||
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
awsapi "github.com/mulgadc/spinifex/spinifex/domains/ecr/awsapi"
|
||||
ecrregistry "github.com/mulgadc/spinifex/spinifex/domains/ecr/registry"
|
||||
"github.com/mulgadc/spinifex/spinifex/domains/ochre"
|
||||
ochrevector "github.com/mulgadc/spinifex/spinifex/domains/ochre/vector"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/lifecycle/reconciler"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/state/kvstore"
|
||||
"github.com/mulgadc/spinifex/spinifex/foundation/state/kvutil"
|
||||
@@ -37,7 +38,6 @@ import (
|
||||
gateway_ec2_instance "github.com/mulgadc/spinifex/spinifex/gateway/ec2/instance"
|
||||
handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance"
|
||||
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
|
||||
handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector"
|
||||
handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota"
|
||||
handlers_sts "github.com/mulgadc/spinifex/spinifex/handlers/sts"
|
||||
"github.com/mulgadc/spinifex/spinifex/network/reconcile"
|
||||
@@ -405,9 +405,9 @@ func launchService(config *config.ClusterConfig) error {
|
||||
// bedrock stores above are, rather than a second NATS hop through the
|
||||
// daemon. bedrockAgentVector forwards CreateIndex/Ingest/DescribeJob/
|
||||
// ListJobs/etc to .9's daemon-side VectorService over NATS.
|
||||
bedrockAgentKB := handlers_ochrevector.NewKBStore(js)
|
||||
bedrockAgentDataSources := handlers_ochrevector.NewDataSourceStore(js)
|
||||
bedrockAgentVector := handlers_ochrevector.NewNATSVectorService(natsConn)
|
||||
bedrockAgentKB := ochrevector.NewKBStore(js)
|
||||
bedrockAgentDataSources := ochrevector.NewDataSourceStore(js)
|
||||
bedrockAgentVector := ochrevector.NewNATSVectorService(natsConn)
|
||||
|
||||
// Bedrock invocation records: every Converse/InvokeModel call (streaming
|
||||
// or not) is published to the invocation stream, then fanned out by
|
||||
|
||||
Reference in New Issue
Block a user