chromago

package module
v0.1.4 Latest Latest
Warning

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

 Go to latest Published: Jun 26, 2024 License: MIT Imports: 8 Imported by: 32

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)

Embedding Functions Support

  • ✅ OpenAI Embedding Support
  • ✅ Cohere API (including Multi-language support)
  • ✅ Sentence Transformers (HuggingFace Inference API)
  • ✅ Google Gemini Embedding Support
  • 🚫 Custom Embedding Function
  • ✅ HuggingFace Embedding Inference Server Function
  • ✅ Ollama Embedding Support
  • ✅ Cloudflare Workers AI Embedding Support
  • ✅ TogetherAI Embedding Support
  • ✅ VoyageAI Embedding Support

Installation

go get github.com/amikos-tech/chroma-go

Import:

import (
chroma "github.com/amikos-tech/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://amikos-tech.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/amikos-tech/chroma-go"
	"github.com/amikos-tech/chroma-go/collection"
	openai "github.com/amikos-tech/chroma-go/openai"
	"github.com/amikos-tech/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 := chroma.NewClient("localhost:8000")

	// Create a new collection with options
	newCollection, err := client.NewCollection(
		context.TODO(),
		collection.WithName("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 added in v0.1.0 

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

func APIEmbeddingsToEmbeddings added in v0.1.0 

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
	// contains filtered or unexported fields
}

Client represents the ChromaDB Client

func NewClient 

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

func (*Client) CountCollections added in v0.1.0 

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 added in v0.1.0 

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

func (*Client) CreateTenant added in v0.1.0 

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 added in v0.1.0 

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

func (*Client) GetTenant added in v0.1.0 

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 added in v0.1.0 

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

func (*Client) PreflightChecks added in v0.1.0 

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 added in v0.1.0 

func (c *Client) SetDatabase(database string)

func (*Client) SetTenant added in v0.1.0 

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 added in v0.1.0 

type ClientOption func(p *Client) error

func WithAuth added in v0.1.1 

func WithAuth(provider types.CredentialsProvider) ClientOption

func WithDatabase added in v0.1.0 

func WithDatabase(database string) ClientOption

func WithDebug added in v0.1.0 

func WithDebug(debug bool) ClientOption

func WithDefaultHeaders added in v0.1.0 

func WithDefaultHeaders(headers map[string]string) ClientOption

func WithTenant added in v0.1.0 

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) (*Collection, error)

func (*Collection) AddRecords added in v0.1.0 

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 added in v0.1.0 

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 added in v0.1.0 

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

func (*Collection) String added in v0.1.0 

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 added in v0.1.0 

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"`
}

Source Files

View all Source files

Directories

Path Synopsis
pkg
embeddings/cloudflare
embeddings/gemini
embeddings/mistral
embeddings/ollama
embeddings/together
embeddings/voyage

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.