feat: HNSW indexing; refactor doc embeddings

.github/workflows/build-test.yml

@@ -6,10 +6,7 @@ on:
       - main
   push:
      branches: [ "main" ]
-
-
 jobs:
-
   build:
     runs-on: ubuntu-latest
     container: debian:bullseye-slim
@@ -34,22 +31,31 @@ jobs:
           --health-retries 5
 
     steps:
-    - uses: actions/checkout@v3
+      - uses: actions/checkout@v3
 
-    - name: install certs and build-essential (required by CGO)
-      run: apt-get update && apt-get install -y ca-certificates build-essential
-
-    - name: Set up Go
-      uses: actions/setup-go@v4
-      with:
-        go-version: '^1.20'
-    - name: Build
-      run: go build -v ./...
-
-    - name: Test
-      run: CGO_ENABLED=1 go test -tags=testutils -race -p 1 -v ./...
-      env:
-        ZEP_OPENAI_API_KEY: ${{ secrets.ZEP_OPENAI_API_KEY }}
-        ZEP_ANTHROPIC_API_KEY: ${{ secrets.ZEP_ANTHROPIC_API_KEY }}
-        ZEP_STORE_POSTGRES_DSN: 'postgres://postgres:postgres@postgres:5432/?sslmode=disable'
-        ZEP_NLP_SERVER_URL: 'http://nlp:5557'
+      - name: install certs and build-essential (required by CGO)
+        run: apt-get update && apt-get install -y ca-certificates build-essential
+
+      - name: Set up Go
+        uses: actions/setup-go@v4
+        with:
+          go-version: '^1.20'
+
+      - name: Cache Go modules
+        uses: actions/cache@v3
+        with:
+          path: ~/go/pkg/mod
+          key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
+          restore-keys: |
+            ${{ runner.os }}-go-
+
+      - name: Build
+        run: go build -v ./...
+
+      - name: Test
+        run: CGO_ENABLED=1 go test -tags=testutils -race -p 1 -v ./...
+        env:
+          ZEP_OPENAI_API_KEY: ${{ secrets.ZEP_OPENAI_API_KEY }}
+          ZEP_ANTHROPIC_API_KEY: ${{ secrets.ZEP_ANTHROPIC_API_KEY }}
+          ZEP_STORE_POSTGRES_DSN: 'postgres://postgres:postgres@postgres:5432/?sslmode=disable'
+          ZEP_NLP_SERVER_URL: 'http://nlp:5557'

.gitignore

@@ -8,9 +8,14 @@ out/
 *.dll
 *.so
 *.dylib
+
+# Secrets
 .env
 .env.local
 
+# Test data
+test_data
+
 # Test binary, built with `go test -c`
 *.test
 

cmd/zep/cmd.go

@@ -57,11 +57,13 @@ var loadFixturesCmd = &cobra.Command{
 		appState := &models.AppState{
 			Config: cfg,
 		}
-		db := postgres.NewPostgresConn(appState)
+		db, err := postgres.NewPostgresConn(appState)
+		if err != nil {
+			log.Fatalf("Failed to connect to database: %v\n", err)
+		}
 		err = postgres.LoadFixtures(context.Background(), appState, db, fixturePath)
 		if err != nil {
-			fmt.Printf("Failed to load fixtures: %v\n", err)
-			os.Exit(1)
+			log.Fatalf("Failed to load fixtures: %v\n", err)
 		}
 		fmt.Println("Fixtures loaded successfully.")
 	},

cmd/zep/run.go

@@ -100,7 +100,6 @@ func handleCLIOptions(cfg *config.Config) {
 }
 
 // initializeStores initializes the memory and document stores based on the config file / ENV
-// TODO: refactor to include document store switch
 func initializeStores(appState *models.AppState) {
 	if appState.Config.Store.Type == "" {
 		log.Fatal(ErrStoreTypeNotSet)
@@ -111,7 +110,10 @@ func initializeStores(appState *models.AppState) {
 		if appState.Config.Store.Postgres.DSN == "" {
 			log.Fatal(ErrPostgresDSNNotSet)
 		}
-		db := postgres.NewPostgresConn(appState)
+		db, err := postgres.NewPostgresConn(appState)
+		if err != nil {
+			log.Fatalf("Failed to connect to database: %v\n", err)
+		}
 		if appState.Config.Log.Level == "debug" {
 			pgDebugLogging(db)
 		}
@@ -122,13 +124,13 @@ func initializeStores(appState *models.AppState) {
 		log.Debug("memoryStore created")
 
 		// create channels for the document embedding processor
-		// TODO: make this a configurable buffer size
 		embeddingTaskChannel := make(
 			chan []models.DocEmbeddingTask,
-			100,
+			// We use the Pool's buffer, so this doesn't need to be large
+			10,
 		)
-		// TODO: make this a configurable buffer size
-		embeddingUpdateChannel := make(chan []models.DocEmbeddingUpdate, 100)
+		// TODO: Make channel size configurable
+		embeddingUpdateChannel := make(chan []models.DocEmbeddingUpdate, 500)
 		documentStore, err := postgres.NewDocumentStore(
 			appState,
 			db,

config.yaml

@@ -26,6 +26,9 @@ extractors:
   documents:
     embeddings:
       enabled: true
+      max_procs: 2
+      chunk_size: 1000
+      buffer_size: 1000
       dimensions: 384
       service: "local"
 #      dimensions: 1536

config/models.go

@@ -45,7 +45,13 @@ type MemoryConfig struct {
 }
 
 type PostgresConfig struct {
-	DSN string `mapstructure:"dsn"`
+	DSN              string           `mapstructure:"dsn"`
+	AvailableIndexes AvailableIndexes `mapstructure:"available_indexes"`
+}
+
+type AvailableIndexes struct {
+	IVFFLAT bool `mapstructure:"ivfflat"`
+	HSNW    bool `mapstructure:"hsnw"`
 }
 
 type ServerConfig struct {
@@ -94,6 +100,12 @@ type EmbeddingsConfig struct {
 	Enabled    bool   `mapstructure:"enabled"`
 	Dimensions int    `mapstructure:"dimensions"`
 	Service    string `mapstructure:"service"`
+	// MaxProcs is the maximum number of concurrent processes to use for embedding tasks.
+	MaxProcs int `mapstructure:"max_procs"`
+	// ChunkSize is the number of documents to embed in a single task.
+	ChunkSize int `mapstructure:"chunk_size"`
+	// BufferSize is the size of the channel buffer for embedding tasks.
+	BufferSize int `mapstructure:"buffer_size"`
 }
 
 type EntityExtractorConfig struct {

docker-compose.yaml

@@ -4,6 +4,7 @@ services:
     image: ghcr.io/getzep/postgres:latest
     container_name: zep-postgres
     restart: on-failure
+    shm_size: "128mb"  # Increase this if vacuuming fails with a "no space left on device" error
     environment:
       - POSTGRES_USER=postgres
       - POSTGRES_PASSWORD=postgres

go.mod

@@ -23,80 +23,80 @@ require (
 	github.com/spf13/viper v1.16.0
 	github.com/stretchr/testify v1.8.4
 	github.com/swaggo/swag/v2 v2.0.0-rc3
-	github.com/uptrace/bun v1.1.15
-	github.com/uptrace/bun/dialect/pgdialect v1.1.15
-	github.com/uptrace/bun/driver/pgdriver v1.1.15
+	github.com/uptrace/bun v1.1.16
+	github.com/uptrace/bun/dialect/pgdialect v1.1.16
+	github.com/uptrace/bun/driver/pgdriver v1.1.16
 	gonum.org/v1/gonum v0.14.0
 )
 
 require (
-	github.com/Masterminds/sprig/v3 v3.2.3
+	github.com/Masterminds/semver/v3 v3.2.1
 	github.com/alecthomas/chroma v0.10.0
+	github.com/dustin/go-humanize v1.0.1
+	github.com/getzep/sprig/v3 v3.0.0-20230930153539-1d7fce7d845e
 	github.com/hashicorp/go-retryablehttp v0.7.4
-	github.com/tmc/langchaingo v0.0.0-20230910230029-9c8845b2b019
-	github.com/uptrace/bun/dbfixture v1.1.15
-	github.com/uptrace/bun/extra/bundebug v1.1.15
+	github.com/tmc/langchaingo v0.0.0-20230929160525-e16b77704b8d
+	github.com/uptrace/bun/dbfixture v1.1.16
+	github.com/uptrace/bun/extra/bundebug v1.1.16
 	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
 	github.com/KyleBanks/depth v1.2.1 // indirect
 	github.com/Masterminds/goutils v1.1.1 // indirect
-	github.com/Masterminds/semver/v3 v3.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.2.0 // indirect
 	github.com/dlclark/regexp2 v1.10.0 // indirect
 	github.com/fatih/color v1.15.0 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
-	github.com/go-openapi/jsonpointer v0.19.6 // indirect
+	github.com/go-openapi/jsonpointer v0.20.0 // indirect
 	github.com/go-openapi/jsonreference v0.20.2 // indirect
 	github.com/go-openapi/spec v0.20.9 // indirect
-	github.com/go-openapi/swag v0.22.3 // indirect
+	github.com/go-openapi/swag v0.22.4 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/goccy/go-json v0.10.2 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
-	github.com/huandu/xstrings v1.3.3 // indirect
-	github.com/imdario/mergo v0.3.11 // indirect
+	github.com/huandu/xstrings v1.4.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
-	github.com/jackc/pgx/v5 v5.4.2 // indirect
+	github.com/jackc/pgx/v5 v5.4.3 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/leodido/go-urn v1.2.4 // indirect
-	github.com/lestrrat-go/blackmagic v1.0.1 // indirect
+	github.com/lestrrat-go/blackmagic v1.0.2 // indirect
 	github.com/lestrrat-go/httpcc v1.0.1 // indirect
 	github.com/lestrrat-go/httprc v1.0.4 // indirect
 	github.com/lestrrat-go/iter v1.0.2 // indirect
-	github.com/lestrrat-go/jwx/v2 v2.0.11 // indirect
+	github.com/lestrrat-go/jwx/v2 v2.0.13 // indirect
 	github.com/lestrrat-go/option v1.0.1 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
-	github.com/mitchellh/copystructure v1.0.0 // indirect
+	github.com/mitchellh/copystructure v1.2.0 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
-	github.com/mitchellh/reflectwalk v1.0.0 // indirect
-	github.com/pelletier/go-toml/v2 v2.0.8 // indirect
+	github.com/mitchellh/reflectwalk v1.0.2 // indirect
+	github.com/pelletier/go-toml/v2 v2.1.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/segmentio/asm v1.2.0 // indirect
-	github.com/shopspring/decimal v1.2.0 // indirect
-	github.com/spf13/afero v1.9.5 // indirect
+	github.com/shopspring/decimal v1.3.1 // indirect
+	github.com/spf13/afero v1.10.0 // indirect
 	github.com/spf13/cast v1.5.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	github.com/subosito/gotenv v1.4.2 // indirect
-	github.com/sv-tools/openapi v0.2.1 // indirect
+	github.com/subosito/gotenv v1.6.0 // indirect
+	github.com/sv-tools/openapi v0.2.2 // indirect
 	github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc // indirect
 	github.com/vmihailenco/msgpack/v5 v5.3.5 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
 	golang.org/x/crypto v0.13.0 // indirect
-	golang.org/x/net v0.10.0 // indirect
+	golang.org/x/net v0.15.0 // indirect
 	golang.org/x/sys v0.12.0 // indirect
 	golang.org/x/text v0.13.0 // indirect
-	golang.org/x/tools v0.8.0 // indirect
+	golang.org/x/tools v0.13.0 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	mellium.im/sasl v0.3.1 // indirect