Coaster Searching Prep: basic document embedding & searching w/ ollama bindings

.github/workflows/ci.yml

@@ -30,6 +30,19 @@ jobs:
       # `make build`, not `make check`: linting belongs to the `lint` job.
       - run: make build
 
+  test:
+    name: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+      - uses: actions/setup-go@v6
+        with:
+          go-version-file: backend/go.mod
+          cache-dependency-path: |
+            backend/go.sum
+            cli/go.sum
+      - run: make test
+
   generated:
     name: generated code is current
     runs-on: ubuntu-latest

Makefile

@@ -12,7 +12,7 @@ export GOOSE_DRIVER := postgres
 export GOOSE_DBSTRING := $(DATABASE_URL)
 export GOOSE_MIGRATION_DIR := backend/db/migrations
 
-.PHONY: help api rebuild-api watch-api stop-api logs psql install tools generate build check lint verify-gen migration
+.PHONY: help api rebuild-api watch-api stop-api logs psql install tools generate build check test lint verify-gen migration
 
 help: ## Show available targets
 	@grep -E '^[a-zA-Z_-]+:.*?## ' $(MAKEFILE_LIST) | awk -F':.*?## ' '{printf "  \033[36m%-11s\033[0m %s\n", $$1, $$2}'
@@ -54,10 +54,15 @@ build: ## Build and vet both modules
 	cd backend && go build ./... && go vet ./...
 	cd cli && go build ./... && go vet ./...
 
-check: build ## Build, vet, lint, and verify generated code (mirrors CI)
+check: build ## Build, vet, test, lint, and verify generated code (mirrors CI)
+	$(MAKE) test
 	$(MAKE) lint
 	$(MAKE) verify-gen
 
+test: ## Run the unit tests for both modules
+	cd backend && go test ./...
+	cd cli && go test ./...
+
 lint: ## Run golangci-lint on both modules
 	cd backend && golangci-lint run
 	cd cli && golangci-lint run

backend/cmd/api/main.go

@@ -2,9 +2,10 @@ package main
 
 import (
 	"context"
-	"fmt"
+	"log"
 
 	"github.com/jackjustus/credittracker/backend/internal/data"
+	"github.com/jackjustus/credittracker/backend/internal/embed"
 	api "github.com/jackjustus/credittracker/backend/internal/gen/openapi"
 	"github.com/jackjustus/credittracker/backend/internal/handler"
 	"github.com/labstack/echo/v4"
@@ -14,10 +15,15 @@ func main() {
 	ctx := context.Background()
 	dao, err := data.NewDAO(ctx)
 	if err != nil {
-		fmt.Print(err)
+		log.Fatal(err)
 		return
 	}
-	r := handler.NewServer(dao)
+	embedder, err := embed.NewClient()
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+	r := handler.NewServer(dao, embedder)
 	e := echo.New()
 	e.HTTPErrorHandler = func(err error, c echo.Context) {
 		c.Logger().Errorf("%s %s: %v", c.Request().Method, c.Path(), err)
@@ -30,7 +36,7 @@ func main() {
 	))
 	err = e.Start(":8080")
 	if err != nil {
-		fmt.Print(err)
+		log.Fatal(err)
 		return
 	}
 }

backend/cmd/embedder/main.go

@@ -0,0 +1,64 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"strings"
+
+	"github.com/jackjustus/credittracker/backend/internal/data"
+	"github.com/jackjustus/credittracker/backend/internal/embed"
+	"github.com/jackjustus/credittracker/backend/internal/models"
+	"github.com/life4/genesis/slices"
+)
+
+func main() {
+	ctx := context.Background()
+	dao, err := data.NewDAO(ctx)
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+	coasters, err := dao.ListEmbeddingTargets(ctx, embed.Model, 700)
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+
+	docs := slices.Map(coasters, func(c *models.CoasterWithPark) string {
+		doc, _ := CoasterDoc(*c)
+		return doc
+	})
+
+	client, err := embed.NewClient()
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+
+	vecs, err := client.Embed(ctx, docs)
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+
+	// Embed guarantees one vector per doc, so these indexes stay aligned.
+	for i, c := range coasters {
+		if err := dao.SetCoasterEmbedding(ctx, c.Coaster().ID(), docs[i], embed.Model, vecs[i]); err != nil {
+			log.Fatalf("writing %s: %v", c.Coaster().Name(), err)
+		}
+	}
+
+	fmt.Printf("embedded %d coasters with %s\n", len(coasters), embed.Model)
+}
+
+func CoasterDoc(c models.CoasterWithPark) (string, error) {
+	var b strings.Builder
+
+	_, err := fmt.Fprintf(&b, "%s is a coaster at %s", c.Coaster().Name(), c.Park().Name())
+	if err != nil {
+		return "", err
+	}
+
+	return b.String(), nil
+}

backend/cmd/embedder/main_test.go

@@ -0,0 +1,49 @@
+package main
+
+import (
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/jackjustus/credittracker/backend/internal/gen/db"
+	"github.com/jackjustus/credittracker/backend/internal/models"
+)
+
+func TestCoasterDoc(t *testing.T) {
+	now := pgtype.Timestamptz{Time: time.Now(), Valid: true}
+	parkID := uuid.New()
+
+	c := models.NewCoaster(db.Coaster{
+		ID:             uuid.New(),
+		ParkID:         parkID,
+		Name:           "Kingda Ka",
+		ManufacturedAt: now,
+		CreatedAt:      now,
+		UpdatedAt:      now,
+	})
+	p := models.NewPark(db.Park{
+		ID:        parkID,
+		Name:      "Six Flags Great Adventure",
+		City:      "Jackson",
+		Country:   "United States",
+		CreatedAt: now,
+		UpdatedAt: now,
+	})
+	cwp, err := models.NewCoasterWithPark(c, p)
+	if err != nil {
+		t.Fatalf("NewCoasterWithPark: %v", err)
+	}
+
+	doc, err := CoasterDoc(*cwp)
+	if err != nil {
+		t.Fatalf("CoasterDoc: %v", err)
+	}
+
+	// Pinned exactly: this string is what gets embedded, so a wording change
+	// invalidates every vector already stored under the old phrasing.
+	want := "Kingda Ka is a coaster at Six Flags Great Adventure"
+	if doc != want {
+		t.Errorf("CoasterDoc = %q, want %q", doc, want)
+	}
+}

backend/compose.yaml

@@ -1,6 +1,6 @@
 services:
   db:
-    image: postgres:18
+    image: pgvector/pgvector:pg18-trixie
     restart: unless-stopped
     environment:
       POSTGRES_PASSWORD: testtest
@@ -56,6 +56,14 @@ services:
           path: ./go.sum
     environment:
       DATABASE_URL: postgres://postgres:testtest@db:5432/postgres?sslmode=disable
+      # Ollama runs natively on the host, not in a container: Docker Desktop on
+      # macOS has no GPU passthrough, so a containerized Ollama would be
+      # CPU-only. The host daemon must bind 0.0.0.0 (OLLAMA_HOST) to accept
+      # this; on loopback it only answers processes on the Mac itself.
+      OLLAMA_URL: http://host.docker.internal:11434
+    # Automatic on Docker Desktop; this makes it work on Linux hosts too.
+    extra_hosts:
+      - "host.docker.internal:host-gateway"
     depends_on:
       migrate:
         condition: service_completed_successfully

backend/db/migrations/20260812232147_add_pgvector.sql

@@ -0,0 +1,17 @@
+-- +goose Up
+CREATE EXTENSION IF NOT EXISTS vector;
+
+ALTER TABLE coasters
+    ADD COLUMN search_profile          text,        -- text passed to embedding model
+    ADD COLUMN profile_embedding       vector(768),
+    ADD COLUMN profile_embedding_model text,        -- text literal of embedding model used, ex: nomic-embed-text
+    ADD COLUMN profile_embedded_at     timestamptz;
+
+-- +goose Down
+ALTER TABLE coasters
+    DROP COLUMN search_profile          ,
+    DROP COLUMN profile_embedding       ,
+    DROP COLUMN profile_embedding_model ,
+    DROP COLUMN profile_embedded_at     ;
+
+DROP EXTENSION IF EXISTS vector;

backend/db/queries/coaster.sql

@@ -2,3 +2,38 @@
 SELECT sqlc.embed(coasters)
 FROM coasters
 WHERE coasters.id = @id;
+
+-- Candidates for (re-)embedding.. either never embedded, or embedded under a different
+-- model than the parameterized one.
+--
+-- Park is embedded so that the doc producer as as much metadata abt the coaster as possible to
+-- improve embedding outcomes.
+--
+-- name: ListCoastersNeedingEmbedding :many
+SELECT sqlc.embed(coasters), sqlc.embed(parks)
+FROM coasters
+JOIN parks ON parks.id = coasters.park_id
+WHERE coasters.profile_embedding IS NULL
+   OR coasters.profile_embedding_model IS DISTINCT FROM @model::text
+ORDER BY coasters.id
+LIMIT @row_limit;
+
+-- Nearest neighbours to an embedded query vector, closest first.
+-- <=> is cosine distance. non-embedded coasters are excluded.
+-- name: SearchCoasters :many
+SELECT sqlc.embed(coasters)
+FROM coasters
+WHERE profile_embedding IS NOT NULL
+ORDER BY profile_embedding <=> @query_embedding
+LIMIT @row_limit;
+
+-- Stores the vector alongside the exact text and model that produced it, so a
+-- later run can tell whether a row is stale.
+-- rather than changing updated_at, we write to embedded_at. (core coaster metadata isn't changing)
+-- name: SetCoasterEmbedding :exec
+UPDATE coasters
+SET search_profile          = @search_profile::text,
+    profile_embedding       = @profile_embedding,
+    profile_embedding_model = @profile_embedding_model::text,
+    profile_embedded_at     = now()
+WHERE id = @id;

backend/go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/labstack/echo/v4 v4.15.4
 	github.com/life4/genesis v1.10.3
 	github.com/oapi-codegen/runtime v1.6.0
+	github.com/pgvector/pgvector-go v0.4.1
 )
 
 require (
@@ -20,9 +21,9 @@ require (
 	github.com/mattn/go-isatty v0.0.24 // indirect
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
 	github.com/valyala/fasttemplate v1.2.2 // indirect
-	golang.org/x/crypto v0.54.0 // indirect
-	golang.org/x/net v0.57.0 // indirect
+	golang.org/x/crypto v0.55.0 // indirect
+	golang.org/x/net v0.58.0 // indirect
 	golang.org/x/sync v0.22.0 // indirect
 	golang.org/x/sys v0.47.0 // indirect
-	golang.org/x/text v0.40.0 // indirect
+	golang.org/x/text v0.41.0 // indirect
 )

backend/go.sum

@@ -30,6 +30,8 @@ github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsRe
 github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
 github.com/oapi-codegen/runtime v1.6.0 h1:7Xx+GlueD6nRuyKoCPzL434Jfi3BetbiJOrzCHp/VPU=
 github.com/oapi-codegen/runtime v1.6.0/go.mod h1:GwV7hC2hviaMzj+ITfHVRESK5J2W/GefVwIND/bMGvU=
+github.com/pgvector/pgvector-go v0.4.1 h1:Oaj0mC0Ky8KaTweNHHpLwyFlN6a0nUFoo1vgSFTEhPI=
+github.com/pgvector/pgvector-go v0.4.1/go.mod h1:4fSXyjl1TYAIdByAql6JazKWRr2s7J0g4hcRY5cBFCk=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/spkg/bom v0.0.0-20160624110644-59b7046e48ad/go.mod h1:qLr4V1qq6nMqFKkMo8ZTx3f+BZEkzsRUY10Xsm2mwU0=
@@ -44,14 +46,20 @@ github.com/valyala/fasttemplate v1.2.2 h1:lxLXG0uE3Qnshl9QyaK6XJxMXlQZELvChBOCmQ
 github.com/valyala/fasttemplate v1.2.2/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ=
 golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw=
 golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk=
+golang.org/x/crypto v0.55.0 h1:+KWHjbgOaAQ66dh/YlkZKHlz9ZUlq61AFirAR9ntP8M=
+golang.org/x/crypto v0.55.0/go.mod h1:uq0V9dE/fzQuJtbnL+2EhWOE63vo164FY8xqEnV9xis=
 golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE=
 golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU=
+golang.org/x/net v0.58.0 h1:ynWG7rqYi4ccpTEuPZ2QGWHktVEM9DMCj9yzDE0Q7To=
+golang.org/x/net v0.58.0/go.mod h1:YwCddHnFlT7eLQqVprV19OnhLGtc5xOKgE0RyqgfWAU=
 golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
 golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
 golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
 golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
 golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs=
 golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY=
+golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8=
+golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

backend/internal/data/dao.go

@@ -11,6 +11,8 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/jackjustus/credittracker/backend/internal/gen/db"
 	"github.com/jackjustus/credittracker/backend/internal/models"
+	"github.com/life4/genesis/slices"
+	"github.com/pgvector/pgvector-go"
 )
 
 // SentinelUserID is the fixed user ID used for the single local user until
@@ -41,6 +43,55 @@ func (dao *DAO) HydrateCoaster(ctx context.Context, coasterID uuid.UUID) (*model
 	return models.NewCoaster(row.Coaster), nil
 }
 
+func (dao *DAO) ListEmbeddingTargets(ctx context.Context, model string, limit int32) ([]*models.CoasterWithPark, error) {
+	rows, err := dao.dbc.ListCoastersNeedingEmbedding(ctx, db.ListCoastersNeedingEmbeddingParams{
+		Model:    model,
+		RowLimit: limit,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	coasters := make([]*models.CoasterWithPark, 0, len(rows))
+	for _, row := range rows {
+		cwp, err := models.NewCoasterWithPark(models.NewCoaster(row.Coaster), models.NewPark(row.Park))
+		if err != nil {
+			return nil, err
+		}
+		coasters = append(coasters, cwp)
+	}
+	return coasters, nil
+}
+
+// SearchCoasters returns the coasters whose embeddings sit closest to the given query vector.
+// The caller supplies an already embedded query.
+func (dao *DAO) SearchCoasters(ctx context.Context, query []float32, limit int32) (models.Coasters, error) {
+	embedding := pgvector.NewVector(query)
+	rows, err := dao.dbc.SearchCoasters(ctx, db.SearchCoastersParams{
+		QueryEmbedding: &embedding,
+		RowLimit:       limit,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return slices.Map(rows, func(row db.SearchCoastersRow) *models.Coaster {
+		return models.NewCoaster(row.Coaster)
+	}), nil
+}
+
+// SetCoasterEmbedding writes a vector along with the document text and model
+// that produced it.
+func (dao *DAO) SetCoasterEmbedding(ctx context.Context, coasterID uuid.UUID, doc, model string, vec []float32) error {
+	embedding := pgvector.NewVector(vec)
+	return dao.dbc.SetCoasterEmbedding(ctx, db.SetCoasterEmbeddingParams{
+		ID:                    coasterID,
+		SearchProfile:         doc,
+		ProfileEmbedding:      &embedding,
+		ProfileEmbeddingModel: model,
+	})
+}
+
 func (dao *DAO) GetCredits(ctx context.Context, userID uuid.UUID) (*models.Credits, error) {
 	events, err := dao.dbc.GetRideEvents(ctx, userID)
 	if err != nil {