chromago

package module
v0.1.7 Latest Latest
Warning

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

 Go to latest Published: Oct 12, 2024 License: MIT Imports: 8 Imported by: 0

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/hxllyl/chroma-go

Import:

import (
chroma "github.com/hxllyl/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://hxllyl.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/hxllyl/chroma-go"
	"github.com/hxllyl/chroma-go/collection"
	openai "github.com/hxllyl/chroma-go/openai"
	"github.com/hxllyl/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 

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

Client represents the ChromaDB Client

func NewClient 

func NewClient(basePath string, 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, 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 WithDatabase 

func WithDatabase(database string) ClientOption

func WithDebug 

func WithDebug(debug bool) ClientOption

func WithDefaultHeaders 

func WithDefaultHeaders(headers map[string]string) ClientOption

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) (*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"`
}

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.