chromago

package module
v1.4.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 28, 2024 License: MIT Imports: 14 Imported by: 1

README

Chroma Go

A simple Chroma Vector Database client written in Go

Works with Chroma Version: v0.4.3 - v0.5.x

We invite users to visit the docs site for the library for more in-depth information: Chroma Go Docs

Feature Parity with ChromaDB API

  • ✅ Create Tenant
  • ✅ Get Tenant
  • ✅ Create Database
  • ✅ Get Database
  • ✅ Reset
  • ✅ Heartbeat
  • ✅ List Collections
  • ✅ Count Collections
  • ✅ Get Version
  • ✅ Create Collection
  • ✅ Delete Collection
  • ✅ Collection Add
  • ✅ Collection Get (partial without additional parameters)
  • ✅ Collection Count
  • ✅ Collection Query
  • ✅ Collection Modify Embeddings
  • ✅ Collection Update
  • ✅ Collection Upsert
  • ✅ Collection Delete - delete documents in collection
  • ✅ Authentication (Basic, Token with Authorization header, Token with X-Chroma-Token header)
  • ✅ Private PKI and self-signed certificate support

Embedding API and Models Support

Reranking Functions

From release 0.2.0 the Chroma Go client also supports Reranking functions. The following are supported:

Installation

go get github.com/szirtesitidom/chroma-go

Import:

import (
chroma "github.com/szirtesitidom/chroma-go"
)

Usage

Ensure you have a running instance of Chroma running. We recommend one of the two following options:

The Setup (Cloud-native):

minikube start --profile chromago
minikube profile chromago
helm repo add chroma https://szirtesitidom.github.io/chromadb-chart/
helm repo update
helm install chroma chroma/chromadb --set chromadb.allowReset=true,chromadb.apiVersion=0.4.5

|Note: To delete the minikube cluster: minikube delete --profile chromago

Getting Started

Consider the following example where:

  • We create a new collection
  • Add documents using OpenAI embedding function
  • Query the collection using the same embedding function
package main

import (
	"context"
	"fmt"
	"log"
	"os"

	chroma "github.com/szirtesitidom/chroma-go"
	"github.com/szirtesitidom/chroma-go/collection"
	openai "github.com/szirtesitidom/chroma-go/pkg/embeddings/openai"
	"github.com/szirtesitidom/chroma-go/types"
)

func main() {
	// Create new OpenAI embedding function

	openaiEf, err := openai.NewOpenAIEmbeddingFunction(os.Getenv("OPENAI_API_KEY"))
	if err != nil {
		log.Fatalf("Error creating OpenAI embedding function: %s \n", err)
	}
	// Create a new Chroma client
	client,err := chroma.NewClient(chroma.WithBasePath("http://localhost:8000"))
	if err != nil {
        log.Fatalf("Error creating client: %s \n", err)
    }

	// Create a new collection with options
	newCollection, err := client.NewCollection(
		context.TODO(),
        "test-collection",
		collection.WithMetadata("key1", "value1"),
		collection.WithEmbeddingFunction(openaiEf),
		collection.WithHNSWDistanceFunction(types.L2),
	)
	if err != nil {
		log.Fatalf("Error creating collection: %s \n", err)
	}

	// Create a new record set with to hold the records to insert
	rs, err := types.NewRecordSet(
		types.WithEmbeddingFunction(openaiEf),
		types.WithIDGenerator(types.NewULIDGenerator()),
	)
	if err != nil {
		log.Fatalf("Error creating record set: %s \n", err)
	}
	// Add a few records to the record set
	rs.WithRecord(types.WithDocument("My name is John. And I have two dogs."), types.WithMetadata("key1", "value1"))
	rs.WithRecord(types.WithDocument("My name is Jane. I am a data scientist."), types.WithMetadata("key2", "value2"))

	// Build and validate the record set (this will create embeddings if not already present)
	_, err = rs.BuildAndValidate(context.TODO())
	if err != nil {
		log.Fatalf("Error validating record set: %s \n", err)
	}

	// Add the records to the collection
	_, err = newCollection.AddRecords(context.Background(), rs)
	if err != nil {
		log.Fatalf("Error adding documents: %s \n", err)
	}

	// Count the number of documents in the collection
	countDocs, qrerr := newCollection.Count(context.TODO())
	if qrerr != nil {
		log.Fatalf("Error counting documents: %s \n", qrerr)
	}

	// Query the collection
	fmt.Printf("countDocs: %v\n", countDocs) //this should result in 2
	qr, qrerr := newCollection.Query(context.TODO(), []string{"I love dogs"}, 5, nil, nil, nil)
	if qrerr != nil {
		log.Fatalf("Error querying documents: %s \n", qrerr)
	}
	fmt.Printf("qr: %v\n", qr.Documents[0][0]) //this should result in the document about dogs
}

Development

Build
make build
Test
make test
Generate ChromaDB API Client
make generate
Lint
make lint-fix
Local Server

Note: Docker must be installed

make server

References

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func APIEmbeddingToEmbedding 

func APIEmbeddingToEmbedding(embedding openapiclient.EmbeddingsInner) *types.Embedding

func APIEmbeddingsToEmbeddings 

func APIEmbeddingsToEmbeddings(embeddings []openapiclient.EmbeddingsInner) []*types.Embedding

func GetStringTypeOfEmbeddingFunction 

func GetStringTypeOfEmbeddingFunction(ef types.EmbeddingFunction) string

Types

type Client 

type Client struct {
	ApiClient  *openapiclient.APIClient //nolint
	Tenant     string
	Database   string
	APIVersion semver.Version

	BasePath string
	// contains filtered or unexported fields
}

Client represents the ChromaDB Client

func NewClient 

func NewClient(options ...ClientOption) (*Client, error)

func (*Client) CountCollections 

func (c *Client) CountCollections(ctx context.Context) (int32, error)

func (*Client) CreateCollection 

func (c *Client) CreateCollection(ctx context.Context, collectionName string, metadata map[string]interface{}, createOrGet bool, embeddingFunction types.EmbeddingFunction, distanceFunction types.DistanceFunction) (*Collection, error)

func (*Client) CreateDatabase 

func (c *Client) CreateDatabase(ctx context.Context, databaseName string, tenantName *string) (*openapiclient.Database, error)

func (*Client) CreateTenant 

func (c *Client) CreateTenant(ctx context.Context, tenantName string) (*openapiclient.Tenant, error)

func (*Client) DeleteCollection 

func (c *Client) DeleteCollection(ctx context.Context, collectionName string) (*Collection, error)

func (*Client) GetCollection 

func (c *Client) GetCollection(ctx context.Context, collectionName string, embeddingFunction types.EmbeddingFunction) (*Collection, error)

func (*Client) GetDatabase 

func (c *Client) GetDatabase(ctx context.Context, databaseName string, tenantName *string) (*openapiclient.Database, error)

func (*Client) GetTenant 

func (c *Client) GetTenant(ctx context.Context, tenantName string) (*openapiclient.Tenant, error)

func (*Client) Heartbeat 

func (c *Client) Heartbeat(ctx context.Context) (map[string]float32, error)

func (*Client) ListCollections 

func (c *Client) ListCollections(ctx context.Context) ([]*Collection, error)

func (*Client) NewCollection 

func (c *Client) NewCollection(ctx context.Context, name string, options ...collection.Option) (*Collection, error)

func (*Client) PreflightChecks 

func (c *Client) PreflightChecks(ctx context.Context) (map[string]interface{}, error)

func (*Client) Reset 

func (c *Client) Reset(ctx context.Context) (bool, error)

func (*Client) SetDatabase 

func (c *Client) SetDatabase(database string)

func (*Client) SetTenant 

func (c *Client) SetTenant(tenant string)

func (*Client) Version 

func (c *Client) Version(ctx context.Context) (string, error)

type ClientConfiguration 

type ClientConfiguration struct {
	BasePath          string                  `json:"basePath,omitempty"`
	DefaultHeaders    map[string]string       `json:"defaultHeader,omitempty"`
	EmbeddingFunction types.EmbeddingFunction `json:"embeddingFunction,omitempty"`
}

type ClientOption 

type ClientOption func(p *Client) error

func WithAuth 

func WithAuth(provider types.CredentialsProvider) ClientOption

func WithBasePath 

func WithBasePath(basePath string) ClientOption

WithBasePath sets the base path for the client. The base path must be a valid URL.

func WithDatabase 

func WithDatabase(database string) ClientOption

func WithDebug 

func WithDebug(debug bool) ClientOption

func WithDefaultHeaders 

func WithDefaultHeaders(headers map[string]string) ClientOption

func WithHTTPClient 

func WithHTTPClient(client *http.Client) ClientOption

WithHTTPClient sets a custom http.Client for the client. The option is mutually exclusive with WithSSLCert and WithIgnoreSSLCert.

func WithInsecure 

func WithInsecure() ClientOption

WithInsecure disables SSL certificate verification. This option is not recommended for production use. The option is mutually exclusive with WithHttpClient.

func WithSSLCert 

func WithSSLCert(certPath string) ClientOption

WithSSLCert adds a custom SSL certificate to the client. The certificate must be in PEM format. The Option can be added multiple times to add multiple certificates. The option is mutually exclusive with WithHttpClient.

func WithTenant 

func WithTenant(tenant string) ClientOption

type Collection 

type Collection struct {
	Name              string
	EmbeddingFunction types.EmbeddingFunction
	ApiClient         *openapiclient.APIClient //nolint
	Metadata          map[string]interface{}
	ID                string
	Tenant            string
	Database          string
}

func NewCollection 

func NewCollection(apiClient *openapiclient.APIClient, id string, name string, metadata *map[string]interface{}, embeddingFunction types.EmbeddingFunction, tenant string, database string) *Collection

func (*Collection) Add 

func (c *Collection) Add(ctx context.Context, embeddings []*types.Embedding, metadatas []map[string]interface{}, documents []string, ids []string, timeout int) (*Collection, error)

func (*Collection) AddRecords 

func (c *Collection) AddRecords(ctx context.Context, recordSet *types.RecordSet) (*Collection, error)

func (*Collection) Count 

func (c *Collection) Count(ctx context.Context) (int32, error)

func (*Collection) Delete 

func (c *Collection) Delete(ctx context.Context, ids []string, where map[string]interface{}, whereDocuments map[string]interface{}) ([]string, error)

func (*Collection) Get 

func (c *Collection) Get(ctx context.Context, where map[string]interface{}, whereDocuments map[string]interface{}, ids []string, include []types.QueryEnum) (*GetResults, error)

func (*Collection) GetWithOptions 

func (c *Collection) GetWithOptions(ctx context.Context, options ...types.CollectionQueryOption) (*GetResults, error)

func (*Collection) Modify 

func (c *Collection) Modify(ctx context.Context, embeddings []*types.Embedding, metadatas []map[string]interface{}, documents []string, ids []string) (*Collection, error)

func (*Collection) Query 

func (c *Collection) Query(ctx context.Context, queryTexts []string, nResults int32, where map[string]interface{}, whereDocuments map[string]interface{}, include []types.QueryEnum) (*QueryResults, error)

func (*Collection) QueryWithOptions 

func (c *Collection) QueryWithOptions(ctx context.Context, queryOptions ...types.CollectionQueryOption) (*QueryResults, error)

func (*Collection) String 

func (c *Collection) String() string

func (*Collection) Update 

func (c *Collection) Update(ctx context.Context, newName string, newMetadata *map[string]interface{}) (*Collection, error)

func (*Collection) Upsert 

func (c *Collection) Upsert(ctx context.Context, embeddings []*types.Embedding, metadatas []map[string]interface{}, documents []string, ids []string) (*Collection, error)

type GetResults 

type GetResults struct {
	Ids        []string
	Documents  []string
	Metadatas  []map[string]interface{}
	Embeddings []*types.Embedding
}

type QueryResults 

type QueryResults struct {
	Documents                     [][]string                 `json:"documents,omitempty"`
	Ids                           [][]string                 `json:"ids,omitempty"`
	Metadatas                     [][]map[string]interface{} `json:"metadatas,omitempty"`
	Distances                     [][]float32                `json:"distances,omitempty"`
	QueryTexts                    []string
	QueryEmbeddings               []*types.Embedding
	QueryTextsGeneratedEmbeddings []*types.Embedding // the generated embeddings from the query texts
}

Source Files

View all Source files

Directories

Path Synopsis
pkg
commons/cohere
commons/http
embeddings/cloudflare
embeddings/cohere
embeddings/default_ef
embeddings/gemini
embeddings/hf
embeddings/jina
embeddings/mistral
embeddings/nomic
embeddings/ollama
embeddings/openai
embeddings/together
embeddings/voyage
rerankings
rerankings/cohere
rerankings/hf
rerankings/jina
tokenizers/libtokenizers

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.