rag

type Client struct {
	// contains filtered or unexported fields
}

type ClientConfig struct {
	// GCP project ID.
	Project string

	// GCP region (e.g., "us-east5").
	Location string

	// Embedding model name (e.g., "gemini-embedding-001").
	EmbeddingModel string

	// Vertex AI Matching Engine index resource name.
	// Format: projects/{project}/locations/{location}/indexes/{index}
	IndexName string

	// Vertex AI index endpoint resource ID.
	IndexEndpointID string

	// Deployed index ID within the endpoint.
	DeployedIndexID string

	// PublicDomainName is the public endpoint domain for the index endpoint
	// (e.g., "1234.us-central1-5678.vdb.vertexai.goog").
	// Required for public endpoints. Leave empty for private (VPC) endpoints.
	PublicDomainName string

	// GCSBucket enables durable persistence of embeddings to GCS.
	// When set, stores write to both GCS and Matching Engine.
	// Records in GCS can be re-embedded when upgrading models.
	// Optional — leave empty for Matching Engine only.
	GCSBucket string

	// GCSPrefix is an optional path prefix within the GCS bucket.
	GCSPrefix string

	// Dimensions specifies the output dimensionality for embeddings.
	// If 0, defaults to DefaultDimensions (768).
	// Note: different models support different ranges (e.g., text-embedding-005: 1-768).
	Dimensions int
}

type Embedder struct {
	// contains filtered or unexported fields
}

type GCSStore struct {
	// contains filtered or unexported fields
}

type MatchingEngineRetriever struct {
	// contains filtered or unexported fields
}

type MatchingEngineStore struct {
	// contains filtered or unexported fields
}

type MultiStore struct {
	// contains filtered or unexported fields
}

type Result struct {
	// ID is the unique identifier of the matched datapoint.
	ID string

	// Distance is the cosine distance from the query vector.
	// Lower values indicate higher similarity (0 = identical, 2 = opposite).
	Distance float64

	// Metadata contains the key-value pairs stored with the datapoint.
	Metadata map[string]string
}

type Retriever interface {
	// Search finds the nearest neighbors to the query vector.
	// Results are ordered by distance (most similar first).
	// Returns an empty (non-nil) slice when no results match.
	Search(ctx context.Context, query []float32, opts SearchOptions) ([]Result, error)

	// Close releases resources held by the retriever.
	Close() error
}

type SearchOptions struct {
	// TopK is the maximum number of results to return.
	// Defaults to DefaultTopK (5) when zero.
	TopK int

	// DistanceThreshold is the maximum cosine distance for results.
	// Lower values mean stricter matching (higher similarity required).
	//
	// When zero (the default), no threshold filtering is applied and all
	// TopK results are returned. This lets you examine raw distances to
	// calibrate the right threshold for your corpus.
	//
	// Set to a positive value (e.g., 0.4) to filter out results with
	// distance greater than that value.
	//
	// Examples:
	//   SearchOptions{TopK: 5}                       // no filtering, return all 5
	//   SearchOptions{TopK: 10, DistanceThreshold: 0.3} // strict: only very similar
	//   SearchOptions{TopK: 10, DistanceThreshold: 0.6} // moderate: related content
	DistanceThreshold float64

	// Restricts narrows results to datapoints whose stored restricts overlap
	// the supplied allow lists. The map is keyed by namespace; each value
	// is the set of allowed values for that namespace. Restricts are
	// AND-ed across namespaces (a result must match every namespace) and
	// OR-ed within a namespace (any value matches).
	//
	// Datapoints carry restricts via WithRestricts at write time. A
	// datapoint with no value for a queried namespace is excluded — there
	// is no implicit "untagged" group.
	//
	// Use restricts to partition a single index into logical sub-corpora
	// (e.g., by tenant, language, or domain) without standing up a
	// separate index per partition. Leave nil/empty to search across the
	// whole corpus.
	//
	// Example — return only fixes from the "package-build" domain:
	//
	//   SearchOptions{
	//     TopK: 5,
	//     Restricts: map[string][]string{"domain": {"package-build"}},
	//   }
	Restricts map[string][]string
}

type Store interface {
	// Upsert inserts or updates a datapoint in the vector index.
	// The id must be non-empty and unique within the index. Metadata values
	// are stored alongside the vector for retrieval; pass UpsertOption
	// values (e.g. WithRestricts) to attach index-level restrict tags
	// usable by SearchOptions.Restricts.
	Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string, opts ...UpsertOption) error

	// Close releases resources held by the store.
	Close() error
}

type TaskType string

type UpsertOption func(*upsertConfig)