client

type AddRecordsRequest struct {
	IDs        []string         `json:"ids"`
	Documents  []string         `json:"documents"`
	Embeddings [][]float32      `json:"embeddings"` // Change: No longer omitempty
	Metadatas  []map[string]any `json:"metadatas,omitempty"`
}

type ChromaClient struct {
	URL, Tenant, Database string

	Embedder onnx.EmbedderInterface
	// contains filtered or unexported fields
}

type ChromaClientInterface interface {
	TestConnection() error
	GetTenant() (bool, error)
	ListDatabases() ([]Database, error)
	CreateDatabase(name string) error
	ListCollections() ([]Collection, error)
	AddBatch(collectionID string, docs []string, ids []string) error
	AddBatchGeneric(collectionID string, documents []string, ids []string, metadatas []map[string]any) error
	UpsertBatchGeneric(collectionID string, documents []string, ids []string, metadatas []map[string]any) error
	QueryBatch(collectionId string, queryTexts []string, nResults int) (*QueryResponse, error)
	GetIDByName(name string) (string, error)
	ResolveCollectionID(input string) (string, error)
	DeleteCollection(name string) error
	DeleteRecords(collectionID string, ids []string) error
}

type Collection struct {
	ID        string         `json:"id"`
	Name      string         `json:"name"`
	Tenant    string         `json:"tenant"`
	Database  string         `json:"database"`
	Metadata  map[string]any `json:"metadata"`
	Dimension *int           `json:"dimension"` // Pointer because it can be null
	Config    map[string]any `json:"configuration_json"`
}

type CreateCollectionRequest struct {
	Name        string         `json:"name"`
	Metadata    map[string]any `json:"metadata"`
	GetOrCreate bool           `json:"get_or_create"`
}

type Database struct {
	Id     string `json:"id"`
	Name   string `json:"name"`
	Tenant string `json:"tenant"`
}

type GetRecordsRequest struct {
	Tenant   string   `json:"tenant"`
	Database string   `json:"database"`
	IDs      []string `json:"ids,omitempty"`
	Include  []string `json:"include"`
	Limit    *int     `json:"limit"`
	Offset   *int     `json:"offset"`
}

type GetRecordsResponse struct {
	IDs       []string         `json:"ids"`
	Documents []string         `json:"documents"`
	Metadatas []map[string]any `json:"metadatas"`
}

type IngestRecord struct {
	ID       string         `json:"id"`
	Text     string         `json:"text"`
	Metadata map[string]any `json:"metadata"`
}

type QueryResponse struct {
	IDs       [][]string         `json:"ids"`
	Documents [][]string         `json:"documents"`
	Metadatas [][]map[string]any `json:"metadatas"`
	Distances [][]float32        `json:"distances"`
}