99 lines
2.4 KiB
Go
99 lines
2.4 KiB
Go
|
package retrieval
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
type ReloadableProvider struct {
|
||
|
|
mu sync.RWMutex
|
||
|
|
current Provider
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewReloadableProvider(current Provider) *ReloadableProvider {
|
||
|
|
return &ReloadableProvider{current: current}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *ReloadableProvider) Update(next Provider) {
|
||
|
|
if p == nil || next == nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
p.mu.Lock()
|
||
|
|
defer p.mu.Unlock()
|
||
|
|
p.current = next
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *ReloadableProvider) Validate() error {
|
||
|
|
return p.snapshot().Validate()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *ReloadableProvider) Embed(ctx context.Context, inputs []string) ([][]float64, error) {
|
||
|
|
return p.snapshot().Embed(ctx, inputs)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *ReloadableProvider) Rerank(ctx context.Context, query string, documents []string, topN int) ([]RerankResult, error) {
|
||
|
|
return p.snapshot().Rerank(ctx, query, documents, topN)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *ReloadableProvider) snapshot() Provider {
|
||
|
|
if p == nil {
|
||
|
|
return disabledProvider{reason: "retrieval provider is nil"}
|
||
|
|
}
|
||
|
|
p.mu.RLock()
|
||
|
|
defer p.mu.RUnlock()
|
||
|
|
if p.current == nil {
|
||
|
|
return disabledProvider{reason: "retrieval provider is not initialized"}
|
||
|
|
}
|
||
|
|
return p.current
|
||
|
|
}
|
||
|
|
|
||
|
|
type ReloadableVectorStore struct {
|
||
|
|
mu sync.RWMutex
|
||
|
|
current VectorStore
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewReloadableVectorStore(current VectorStore) *ReloadableVectorStore {
|
||
|
|
return &ReloadableVectorStore{current: current}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ReloadableVectorStore) Update(next VectorStore) {
|
||
|
|
if s == nil || next == nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
s.mu.Lock()
|
||
|
|
defer s.mu.Unlock()
|
||
|
|
s.current = next
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ReloadableVectorStore) Validate() error {
|
||
|
|
return s.snapshot().Validate()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ReloadableVectorStore) EnsureCollection(ctx context.Context, vectorSize int) error {
|
||
|
|
return s.snapshot().EnsureCollection(ctx, vectorSize)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ReloadableVectorStore) Upsert(ctx context.Context, points []VectorPoint) error {
|
||
|
|
return s.snapshot().Upsert(ctx, points)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ReloadableVectorStore) Search(ctx context.Context, req SearchRequest) ([]SearchPoint, error) {
|
||
|
|
return s.snapshot().Search(ctx, req)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ReloadableVectorStore) DeletePoints(ctx context.Context, pointIDs []string) error {
|
||
|
|
return s.snapshot().DeletePoints(ctx, pointIDs)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ReloadableVectorStore) snapshot() VectorStore {
|
||
|
|
if s == nil {
|
||
|
|
return disabledVectorStore{reason: "vector store is nil"}
|
||
|
|
}
|
||
|
|
s.mu.RLock()
|
||
|
|
defer s.mu.RUnlock()
|
||
|
|
if s.current == nil {
|
||
|
|
return disabledVectorStore{reason: "vector store is not initialized"}
|
||
|
|
}
|
||
|
|
return s.current
|
||
|
|
}
|