chromago

package module
v0.1.3 Latest Latest
Warning

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

Go to latest
Published: Mar 27, 2024 License: MIT Imports: 8 Imported by: 10

README

Chroma Go

A simple Chroma Vector Database client written in Go

Works with Chroma Version: v0.4.3 - v0.4.24

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
  • ✅ 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 API
  • ✅ Cohere API (including Multi-language support)
  • ✅ Sentence Transformers (HuggingFace Inference API)
  • 🚫 PaLM API
  • 🚫 Custom Embedding Function
  • ✅ HuggingFace Embedding Inference Server Function
  • ✅ Ollama Embedding Function 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"`
}

Directories

Path Synopsis

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL