internal/firestore: implement VectorDB with Firestore
Change-Id: I7f6edd96770b6d295c89c8ee363247738f2f5c31
Reviewed-on: https://go-review.googlesource.com/c/oscar/+/600955
Reviewed-by: Russ Cox <rsc@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/go.mod b/go.mod
index 5de6a20..25dbbd5 100644
--- a/go.mod
+++ b/go.mod
@@ -3,11 +3,13 @@
go 1.23
require (
+ cloud.google.com/go/firestore v1.16.0
github.com/cockroachdb/pebble v1.1.1
github.com/google/generative-ai-go v0.16.0
github.com/google/go-replayers/grpcreplay v1.3.0
golang.org/x/tools v0.23.0
- google.golang.org/api v0.186.0
+ google.golang.org/api v0.189.0
+ google.golang.org/grpc v1.64.1
rsc.io/markdown v0.0.0-20240617154923-1f2ef1438fed
rsc.io/omap v1.2.1-0.20240709133045-40dad5c0c0fb
rsc.io/ordered v1.1.0
@@ -17,11 +19,10 @@
require (
cloud.google.com/go v0.115.0 // indirect
cloud.google.com/go/ai v0.8.0 // indirect
- cloud.google.com/go/auth v0.6.0 // indirect
- cloud.google.com/go/auth/oauth2adapt v0.2.2 // indirect
- cloud.google.com/go/compute/metadata v0.3.0 // indirect
- cloud.google.com/go/firestore v1.15.0 // indirect
- cloud.google.com/go/longrunning v0.5.7 // indirect
+ cloud.google.com/go/auth v0.7.2 // indirect
+ cloud.google.com/go/auth/oauth2adapt v0.2.3 // indirect
+ cloud.google.com/go/compute/metadata v0.5.0 // indirect
+ cloud.google.com/go/longrunning v0.5.9 // indirect
github.com/DataDog/zstd v1.4.5 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/cespare/xxhash/v2 v2.2.0 // indirect
@@ -32,7 +33,7 @@
github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06 // indirect
github.com/felixge/httpsnoop v1.0.4 // indirect
github.com/getsentry/sentry-go v0.27.0 // indirect
- github.com/go-logr/logr v1.4.1 // indirect
+ github.com/go-logr/logr v1.4.2 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
@@ -41,7 +42,7 @@
github.com/google/s2a-go v0.1.7 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
- github.com/googleapis/gax-go/v2 v2.12.5 // indirect
+ github.com/googleapis/gax-go/v2 v2.13.0 // indirect
github.com/klauspost/compress v1.16.0 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/kr/text v0.2.0 // indirect
@@ -66,9 +67,8 @@
golang.org/x/sys v0.22.0 // indirect
golang.org/x/text v0.16.0 // indirect
golang.org/x/time v0.5.0 // indirect
- google.golang.org/genproto v0.0.0-20240617180043-68d350f18fd4 // indirect
- google.golang.org/genproto/googleapis/api v0.0.0-20240617180043-68d350f18fd4 // indirect
- google.golang.org/genproto/googleapis/rpc v0.0.0-20240617180043-68d350f18fd4 // indirect
- google.golang.org/grpc v1.64.1 // indirect
+ google.golang.org/genproto v0.0.0-20240722135656-d784300faade // indirect
+ google.golang.org/genproto/googleapis/api v0.0.0-20240722135656-d784300faade // indirect
+ google.golang.org/genproto/googleapis/rpc v0.0.0-20240722135656-d784300faade // indirect
google.golang.org/protobuf v1.34.2 // indirect
)
diff --git a/go.sum b/go.sum
index cf23ab1..65ee541 100644
--- a/go.sum
+++ b/go.sum
@@ -17,24 +17,30 @@
cloud.google.com/go v0.115.0/go.mod h1:8jIM5vVgoAEoiVxQ/O4BFTfHqulPZgs/ufEzMcFMdWU=
cloud.google.com/go/ai v0.8.0 h1:rXUEz8Wp2OlrM8r1bfmpF2+VKqc1VJpafE3HgzRnD/w=
cloud.google.com/go/ai v0.8.0/go.mod h1:t3Dfk4cM61sytiggo2UyGsDVW3RF1qGZaUKDrZFyqkE=
-cloud.google.com/go/auth v0.6.0 h1:5x+d6b5zdezZ7gmLWD1m/xNjnaQ2YDhmIz/HH3doy1g=
-cloud.google.com/go/auth v0.6.0/go.mod h1:b4acV+jLQDyjwm4OXHYjNvRi4jvGBzHWJRtJcy+2P4g=
+cloud.google.com/go/auth v0.7.0 h1:kf/x9B3WTbBUHkC+1VS8wwwli9TzhSt0vSTVBmMR8Ts=
+cloud.google.com/go/auth v0.7.0/go.mod h1:D+WqdrpcjmiCgWrXmLLxOVq1GACoE36chW6KXoEvuIw=
+cloud.google.com/go/auth v0.7.2 h1:uiha352VrCDMXg+yoBtaD0tUF4Kv9vrtrWPYXwutnDE=
+cloud.google.com/go/auth v0.7.2/go.mod h1:VEc4p5NNxycWQTMQEDQF0bd6aTMb6VgYDXEwiJJQAbs=
cloud.google.com/go/auth/oauth2adapt v0.2.2 h1:+TTV8aXpjeChS9M+aTtN/TjdQnzJvmzKFt//oWu7HX4=
cloud.google.com/go/auth/oauth2adapt v0.2.2/go.mod h1:wcYjgpZI9+Yu7LyYBg4pqSiaRkfEK3GQcpb7C/uyF1Q=
+cloud.google.com/go/auth/oauth2adapt v0.2.3 h1:MlxF+Pd3OmSudg/b1yZ5lJwoXCEaeedAguodky1PcKI=
+cloud.google.com/go/auth/oauth2adapt v0.2.3/go.mod h1:tMQXOfZzFuNuUxOypHlQEXgdfX5cuhwU+ffUuXRJE8I=
cloud.google.com/go/bigquery v1.0.1/go.mod h1:i/xbL2UlR5RvWAURpBYZTtm/cXjCha9lbfbpx4poX+o=
cloud.google.com/go/bigquery v1.3.0/go.mod h1:PjpwJnslEMmckchkHFfq+HTD2DmtT67aNFKH1/VBDHE=
cloud.google.com/go/bigquery v1.4.0/go.mod h1:S8dzgnTigyfTmLBfrtrhyYhwRxG72rYxvftPBK2Dvzc=
cloud.google.com/go/bigquery v1.5.0/go.mod h1:snEHRnqQbz117VIFhE8bmtwIDY80NLUZUMb4Nv6dBIg=
cloud.google.com/go/bigquery v1.7.0/go.mod h1://okPTzCYNXSlb24MZs83e2Do+h+VXtc4gLoIoXIAPc=
cloud.google.com/go/bigquery v1.8.0/go.mod h1:J5hqkt3O0uAFnINi6JXValWIb1v0goeZM77hZzJN/fQ=
-cloud.google.com/go/compute/metadata v0.3.0 h1:Tz+eQXMEqDIKRsmY3cHTL6FVaynIjX2QxYC4trgAKZc=
-cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k=
+cloud.google.com/go/compute/metadata v0.4.0 h1:vHzJCWaM4g8XIcm8kopr3XmDA4Gy/lblD3EhhSux05c=
+cloud.google.com/go/compute/metadata v0.4.0/go.mod h1:SIQh1Kkb4ZJ8zJ874fqVkslA29PRXuleyj6vOzlbK7M=
+cloud.google.com/go/compute/metadata v0.5.0 h1:Zr0eK8JbFv6+Wi4ilXAR8FJ3wyNdpxHKJNPos6LTZOY=
+cloud.google.com/go/compute/metadata v0.5.0/go.mod h1:aHnloV2TPI38yx4s9+wAZhHykWvVCfu7hQbF+9CWoiY=
cloud.google.com/go/datastore v1.0.0/go.mod h1:LXYbyblFSglQ5pkeyhO+Qmw7ukd3C+pD7TKLgZqpHYE=
cloud.google.com/go/datastore v1.1.0/go.mod h1:umbIZjpQpHh4hmRpGhH4tLFup+FVzqBi1b3c64qFpCk=
-cloud.google.com/go/firestore v1.15.0 h1:/k8ppuWOtNuDHt2tsRV42yI21uaGnKDEQnRFeBpbFF8=
-cloud.google.com/go/firestore v1.15.0/go.mod h1:GWOxFXcv8GZUtYpWHw/w6IuYNux/BtmeVTMmjrm4yhk=
-cloud.google.com/go/longrunning v0.5.7 h1:WLbHekDbjK1fVFD3ibpFFVoyizlLRl73I7YKuAKilhU=
-cloud.google.com/go/longrunning v0.5.7/go.mod h1:8GClkudohy1Fxm3owmBGid8W0pSgodEMwEAztp38Xng=
+cloud.google.com/go/firestore v1.16.0 h1:YwmDHcyrxVRErWcgxunzEaZxtNbc8QoFYA/JOEwDPgc=
+cloud.google.com/go/firestore v1.16.0/go.mod h1:+22v/7p+WNBSQwdSwP57vz47aZiY+HrDkrOsJNhk7rg=
+cloud.google.com/go/longrunning v0.5.9 h1:haH9pAuXdPAMqHvzX0zlWQigXT7B0+CL4/2nXXdBo5k=
+cloud.google.com/go/longrunning v0.5.9/go.mod h1:HD+0l9/OOW0za6UWdKJtXoFAX/BGg/3Wj8p10NeWF7c=
cloud.google.com/go/pubsub v1.0.1/go.mod h1:R0Gpsv3s54REJCy4fxDixWD93lHJMoZTyQ2kNxGRt3I=
cloud.google.com/go/pubsub v1.1.0/go.mod h1:EwwdRX2sKPjnvnqCa270oGRyludottCI76h+R3AArQw=
cloud.google.com/go/pubsub v1.2.0/go.mod h1:jhfEVHT8odbXTkndysNHCcx0awwzvfOlguIAii9o8iA=
@@ -108,6 +114,8 @@
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
github.com/go-logr/logr v1.4.1 h1:pKouT5E8xu9zeFC39JXRDukb6JFQPXM5p5I91188VAQ=
github.com/go-logr/logr v1.4.1/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
+github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
@@ -163,8 +171,6 @@
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
-github.com/google/go-replayers/grpcreplay v1.2.0 h1:JlxRH0a1d8lLTXq6Xl+CD5aqMwHNhMyMxvFjqa3EAvg=
-github.com/google/go-replayers/grpcreplay v1.2.0/go.mod h1:v6NgKtkijC0d3e3RW8il6Sy5sqRVUwoQa4mHOGEy8DI=
github.com/google/go-replayers/grpcreplay v1.3.0 h1:1Keyy0m1sIpqstQmgz307zhiJ1pV4uIlFds5weTmxbo=
github.com/google/go-replayers/grpcreplay v1.3.0/go.mod h1:v6NgKtkijC0d3e3RW8il6Sy5sqRVUwoQa4mHOGEy8DI=
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
@@ -189,6 +195,8 @@
github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
github.com/googleapis/gax-go/v2 v2.12.5 h1:8gw9KZK8TiVKB6q3zHY3SBzLnrGp6HQjyfYBYGmXdxA=
github.com/googleapis/gax-go/v2 v2.12.5/go.mod h1:BUDKcWo+RaKq5SC9vVYL0wLADa3VcfswbOMMRmB9H3E=
+github.com/googleapis/gax-go/v2 v2.13.0 h1:yitjD5f7jQHhyDsnhKEBU52NdvvdSeGzlAnDPT0hH1s=
+github.com/googleapis/gax-go/v2 v2.13.0/go.mod h1:Z/fvTZXF8/uw7Xu5GuslPw+bplx6SS338j1Is2S+B7A=
github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
@@ -296,6 +304,8 @@
go.opentelemetry.io/otel v1.26.0/go.mod h1:UmLkJHUAidDval2EICqBMbnAd0/m2vmpf/dAM+fvFs4=
go.opentelemetry.io/otel/metric v1.26.0 h1:7S39CLuY5Jgg9CrnA9HHiEjGMF/X2VHvoXGgSllRz30=
go.opentelemetry.io/otel/metric v1.26.0/go.mod h1:SY+rHOI4cEawI9a7N1A4nIg/nTQXe1ccCNWYOJUrpX4=
+go.opentelemetry.io/otel/sdk v1.24.0 h1:YMPPDNymmQN3ZgczicBY3B6sf9n62Dlj9pWD3ucgoDw=
+go.opentelemetry.io/otel/sdk v1.24.0/go.mod h1:KVrIYw6tEubO9E96HQpcmpTKDVn9gdv35HoYiQWGDFg=
go.opentelemetry.io/otel/trace v1.26.0 h1:1ieeAUb4y0TE26jUFrCIXKpTuVK7uJGN9/Z/2LP5sQA=
go.opentelemetry.io/otel/trace v1.26.0/go.mod h1:4iDxvGDQuUkHve82hJJ8UqrwswHYsZuWCBllGV2U2y0=
golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
@@ -508,8 +518,10 @@
google.golang.org/api v0.28.0/go.mod h1:lIXQywCXRcnZPGlsd8NbLnOjtAoL6em04bJ9+z0MncE=
google.golang.org/api v0.29.0/go.mod h1:Lcubydp8VUV7KeIHD9z2Bys/sm/vGKnG1UHuDBSrHWM=
google.golang.org/api v0.30.0/go.mod h1:QGmEvQ87FHZNiUVJkT14jQNYJ4ZJjdRF23ZXz5138Fc=
-google.golang.org/api v0.186.0 h1:n2OPp+PPXX0Axh4GuSsL5QL8xQCTb2oDwyzPnQvqUug=
-google.golang.org/api v0.186.0/go.mod h1:hvRbBmgoje49RV3xqVXrmP6w93n6ehGgIVPYrGtBFFc=
+google.golang.org/api v0.188.0 h1:51y8fJ/b1AaaBRJr4yWm96fPcuxSo0JcegXE3DaHQHw=
+google.golang.org/api v0.188.0/go.mod h1:VR0d+2SIiWOYG3r/jdm7adPW9hI2aRv9ETOSCQ9Beag=
+google.golang.org/api v0.189.0 h1:equMo30LypAkdkLMBqfeIqtyAnlyig1JSZArl4XPwdI=
+google.golang.org/api v0.189.0/go.mod h1:FLWGJKb0hb+pU2j+rJqwbnsF+ym+fQs73rbJ+KAUgy8=
google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
google.golang.org/appengine v1.5.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
@@ -545,12 +557,18 @@
google.golang.org/genproto v0.0.0-20200729003335-053ba62fc06f/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
google.golang.org/genproto v0.0.0-20200804131852-c06518451d9c/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
google.golang.org/genproto v0.0.0-20200825200019-8632dd797987/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
-google.golang.org/genproto v0.0.0-20240617180043-68d350f18fd4 h1:CUiCqkPw1nNrNQzCCG4WA65m0nAmQiwXHpub3dNyruU=
-google.golang.org/genproto v0.0.0-20240617180043-68d350f18fd4/go.mod h1:EvuUDCulqGgV80RvP1BHuom+smhX4qtlhnNatHuroGQ=
-google.golang.org/genproto/googleapis/api v0.0.0-20240617180043-68d350f18fd4 h1:MuYw1wJzT+ZkybKfaOXKp5hJiZDn2iHaXRw0mRYdHSc=
-google.golang.org/genproto/googleapis/api v0.0.0-20240617180043-68d350f18fd4/go.mod h1:px9SlOOZBg1wM1zdnr8jEL4CNGUBZ+ZKYtNPApNQc4c=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240617180043-68d350f18fd4 h1:Di6ANFilr+S60a4S61ZM00vLdw0IrQOSMS2/6mrnOU0=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240617180043-68d350f18fd4/go.mod h1:Ue6ibwXGpU+dqIcODieyLOcgj7z8+IcskoNIgZxtrFY=
+google.golang.org/genproto v0.0.0-20240708141625-4ad9e859172b h1:dSTjko30weBaMj3eERKc0ZVXW4GudCswM3m+P++ukU0=
+google.golang.org/genproto v0.0.0-20240708141625-4ad9e859172b/go.mod h1:FfBgJBJg9GcpPvKIuHSZ/aE1g2ecGL74upMzGZjiGEY=
+google.golang.org/genproto v0.0.0-20240722135656-d784300faade h1:lKFsS7wpngDgSCeFn7MoLy+wBDQZ1UQIJD4UNM1Qvkg=
+google.golang.org/genproto v0.0.0-20240722135656-d784300faade/go.mod h1:FfBgJBJg9GcpPvKIuHSZ/aE1g2ecGL74upMzGZjiGEY=
+google.golang.org/genproto/googleapis/api v0.0.0-20240701130421-f6361c86f094 h1:0+ozOGcrp+Y8Aq8TLNN2Aliibms5LEzsq99ZZmAGYm0=
+google.golang.org/genproto/googleapis/api v0.0.0-20240701130421-f6361c86f094/go.mod h1:fJ/e3If/Q67Mj99hin0hMhiNyCRmt6BQ2aWIJshUSJw=
+google.golang.org/genproto/googleapis/api v0.0.0-20240722135656-d784300faade h1:WxZOF2yayUHpHSbUE6NMzumUzBxYc3YGwo0YHnbzsJY=
+google.golang.org/genproto/googleapis/api v0.0.0-20240722135656-d784300faade/go.mod h1:mw8MG/Qz5wfgYr6VqVCiZcHe/GJEfI+oGGDCohaVgB0=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240708141625-4ad9e859172b h1:04+jVzTs2XBnOZcPsLnmrTGqltqJbZQ1Ey26hjYdQQ0=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240708141625-4ad9e859172b/go.mod h1:Ue6ibwXGpU+dqIcODieyLOcgj7z8+IcskoNIgZxtrFY=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240722135656-d784300faade h1:oCRSWfwGXQsqlVdErcyTt4A93Y8fo0/9D4b1gnI++qo=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240722135656-d784300faade/go.mod h1:Ue6ibwXGpU+dqIcODieyLOcgj7z8+IcskoNIgZxtrFY=
google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=
@@ -564,8 +582,6 @@
google.golang.org/grpc v1.30.0/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM8pak=
google.golang.org/grpc v1.31.0/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM8pak=
google.golang.org/grpc v1.33.2/go.mod h1:JMHMWHQWaTccqQQlmk3MJZS+GWXOdAesneDmEnv2fbc=
-google.golang.org/grpc v1.64.0 h1:KH3VH9y/MgNQg1dE7b3XfVK0GsPSIzJwdF617gUSbvY=
-google.golang.org/grpc v1.64.0/go.mod h1:oxjF8E3FBnjp+/gVFYdWacaLDx9na1aqy9oovLpxQYg=
google.golang.org/grpc v1.64.1 h1:LKtvyfbX3UGVPFcGqJ9ItpVWW6oN/2XqTxfAnwRRXiA=
google.golang.org/grpc v1.64.1/go.mod h1:hiQF4LFZelK2WKaP6W0L92zGHtiQdZxk8CrSdvyjeP0=
google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
diff --git a/internal/firestore/db_test.go b/internal/firestore/db_test.go
index 4b1662a..2c2ee5d 100644
--- a/internal/firestore/db_test.go
+++ b/internal/firestore/db_test.go
@@ -9,6 +9,7 @@
"flag"
"fmt"
"math/rand/v2"
+ "path/filepath"
"strings"
"testing"
"time"
@@ -23,34 +24,13 @@
)
func TestDB(t *testing.T) {
- ctx := context.Background()
- rr, err := grpcrr.Open("testdata/db.grpcrr")
- if err != nil {
- t.Fatalf("grpcrr.Open: %v", err)
- }
+ rr, fsProject, fsDatabase := openRR(t, "testdata/db.grpcrr")
defer func() {
if err := rr.Close(); err != nil {
t.Fatal(err)
}
}()
-
- var fsProject, fsDatabase string
- if rr.Recording() {
- if *project == "" {
- t.Fatal("recording requires -project")
- }
- // The -database flag can be omitted. We'll use the default one.
- rr.SetInitial([]byte(*project + "," + *database))
- fsProject = *project
- fsDatabase = *database
- } else {
- // Allow -project and -database on replay because other tests might need them.
- var found bool
- fsProject, fsDatabase, found = strings.Cut(string(rr.Initial()), ",")
- if !found {
- t.Fatal("bad initial state")
- }
- }
+ ctx := context.Background()
db, err := NewDB(ctx, &DBOptions{ProjectID: fsProject, Database: fsDatabase, ClientOptions: rr.ClientOptions()})
if err != nil {
@@ -145,6 +125,28 @@
}
}
+func openRR(t *testing.T, file string) (_ *grpcrr.RecordReplay, fsProject, fsDatabase string) {
+ rr, err := grpcrr.Open(filepath.FromSlash(file))
+ if err != nil {
+ t.Fatalf("grpcrr.Open: %v", err)
+ }
+ if rr.Recording() {
+ if *project == "" {
+ t.Fatal("recording requires -project")
+ }
+ // The -database flag can be omitted. We'll use the default one.
+ rr.SetInitial([]byte(*project + "," + *database))
+ return rr, *project, *database
+ }
+ // Allow -project and -database on replay because other tests might need them.
+ var found bool
+ fsProject, fsDatabase, found = strings.Cut(string(rr.Initial()), ",")
+ if !found {
+ t.Fatalf("%s: bad initial state", file)
+ }
+ return rr, fsProject, fsDatabase
+}
+
func panics(f func()) (b bool) {
defer func() {
if recover() != nil {
diff --git a/internal/firestore/firestore.go b/internal/firestore/firestore.go
index 3427a55..4623f82 100644
--- a/internal/firestore/firestore.go
+++ b/internal/firestore/firestore.go
@@ -115,6 +115,9 @@
// If tx is non-nil, the set happens inside the transaction.
func (f *fstore) set(tx *firestore.Transaction, coll, key string, value any) {
dr := f.client.Collection(coll).Doc(key)
+ if dr == nil {
+ f.Panic("firestore set bad doc ref args", "collection", coll, "key", key)
+ }
var err error
if tx == nil {
_, err = dr.Set(context.TODO(), value)
@@ -248,7 +251,7 @@
// estimate the size of val.
func (b *batch) set(id string, val any, valSize int) {
if val == nil {
- panic("firestore batch set: nil value")
+ b.f.Panic("firestore batch set: nil value")
}
b.ops = append(b.ops, &op{id: id, value: val})
b.size += perWriteSize + len(id) + valSize
diff --git a/internal/firestore/testdata/vectordb.grpcrr b/internal/firestore/testdata/vectordb.grpcrr
new file mode 100644
index 0000000..7ab0005
--- /dev/null
+++ b/internal/firestore/testdata/vectordb.grpcrr
@@ -0,0 +1,1415 @@
+RPCReTxt1
+"go-discovery-exp,"
+79
+kind: CREATE_STREAM
+method: "/google.firestore.v1.Firestore/BatchGetDocuments"
+252
+kind: SEND
+message: {
+ [type.googleapis.com/google.firestore.v1.BatchGetDocumentsRequest]: {
+ database: "projects/go-discovery-exp/databases/(default)"
+ documents: "projects/go-discovery-exp/databases/(default)/documents/c/d"
+ }
+}
+ref_index: 1
+261
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.BatchGetDocumentsResponse]: {
+ missing: "projects/go-discovery-exp/databases/(default)/documents/c/d"
+ read_time: {
+ seconds: 1722510074
+ nanos: 371441000
+ }
+ }
+}
+ref_index: 1
+39
+kind: RECV
+is_error: true
+ref_index: 1
+2361
+kind: REQUEST
+method: "/google.firestore.v1.Firestore/Commit"
+message: {
+ [type.googleapis.com/google.firestore.v1.CommitRequest]: {
+ database: "projects/go-discovery-exp/databases/(default)"
+ writes: {
+ update: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/orange2"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.30135461688041687
+ }
+ values: {
+ double_value: 0.3094993531703949
+ }
+ values: {
+ double_value: 0.2633459270000458
+ }
+ values: {
+ double_value: 0.298639714717865
+ }
+ values: {
+ double_value: 0.27963536977767944
+ }
+ values: {
+ double_value: 0.2742055654525757
+ }
+ values: {
+ double_value: 0.1357453316450119
+ }
+ values: {
+ double_value: -0.695016086101532
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+}
+290
+kind: RESPONSE
+message: {
+ [type.googleapis.com/google.firestore.v1.CommitResponse]: {
+ write_results: {
+ update_time: {
+ seconds: 1721666627
+ nanos: 811432000
+ }
+ }
+ commit_time: {
+ seconds: 1722510074
+ nanos: 381426000
+ }
+ }
+}
+ref_index: 5
+2366
+kind: REQUEST
+method: "/google.firestore.v1.Firestore/Commit"
+message: {
+ [type.googleapis.com/google.firestore.v1.CommitRequest]: {
+ database: "projects/go-discovery-exp/databases/(default)"
+ writes: {
+ update: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/orange1"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.30146464705467224
+ }
+ values: {
+ double_value: 0.30961233377456665
+ }
+ values: {
+ double_value: 0.2634420692920685
+ }
+ values: {
+ double_value: 0.29874876141548157
+ }
+ values: {
+ double_value: 0.2797374725341797
+ }
+ values: {
+ double_value: 0.27430567145347595
+ }
+ values: {
+ double_value: 0.13307899236679077
+ }
+ values: {
+ double_value: -0.6952698230743408
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+}
+290
+kind: RESPONSE
+message: {
+ [type.googleapis.com/google.firestore.v1.CommitResponse]: {
+ write_results: {
+ update_time: {
+ seconds: 1721666627
+ nanos: 948188000
+ }
+ }
+ commit_time: {
+ seconds: 1722510074
+ nanos: 381426000
+ }
+ }
+}
+ref_index: 7
+222
+kind: REQUEST
+method: "/google.firestore.v1.Firestore/BeginTransaction"
+message: {
+ [type.googleapis.com/google.firestore.v1.BeginTransactionRequest]: {
+ database: "projects/go-discovery-exp/databases/(default)"
+ }
+}
+401
+kind: RESPONSE
+message: {
+ [type.googleapis.com/google.firestore.v1.BeginTransactionResponse]: {
+ transaction: "\x11\xde\xc38\xb9/P\xaf!\"Y\x00\xd4u\xc8\xcd\x15}i\xe8Q\x0c\xc7/\xe1SfC\xd3$\xef\xa2\x18O\x9a%\xc4ѻ\x83\xea\x13\xa6\xaf\x824\x06\"\xaf\x1e\x89\xfan\xe2ϓ\x97gC\xebܲ\xc3\xfb;\xe2\x82x\xa3\xddM\xd9q倲\xb4A\xee[\x98\x7fP\xbeQ\xb1q{\x9d\xa7\x98I?\xe4\xf6;\x8b\xc3"
+ }
+}
+ref_index: 9
+5781
+kind: REQUEST
+method: "/google.firestore.v1.Firestore/Commit"
+message: {
+ [type.googleapis.com/google.firestore.v1.CommitRequest]: {
+ database: "projects/go-discovery-exp/databases/(default)"
+ writes: {
+ update: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple3"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.2749089002609253
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.30608412623405457
+ }
+ values: {
+ double_value: 0.28624534606933594
+ }
+ values: {
+ double_value: 0.1445397287607193
+ }
+ values: {
+ double_value: -0.7255327701568604
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ writes: {
+ update: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple4"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.274795264005661
+ }
+ values: {
+ double_value: 0.3172893524169922
+ }
+ values: {
+ double_value: 0.3172893524169922
+ }
+ values: {
+ double_value: 0.3059576153755188
+ }
+ values: {
+ double_value: 0.2861270308494568
+ }
+ values: {
+ double_value: 0.1473129242658615
+ }
+ values: {
+ double_value: -0.7252328395843506
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ writes: {
+ update: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/ignore"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.31871140003204346
+ }
+ values: {
+ double_value: 0.3154592514038086
+ }
+ values: {
+ double_value: 0.32521572709083557
+ }
+ values: {
+ double_value: -0.8325522541999817
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ transaction: "\x11\xde\xc38\xb9/P\xaf!\"Y\x00\xd4u\xc8\xcd\x15}i\xe8Q\x0c\xc7/\xe1SfC\xd3$\xef\xa2\x18O\x9a%\xc4ѻ\x83\xea\x13\xa6\xaf\x824\x06\"\xaf\x1e\x89\xfan\xe2ϓ\x97gC\xebܲ\xc3\xfb;\xe2\x82x\xa3\xddM\xd9q倲\xb4A\xee[\x98\x7fP\xbeQ\xb1q{\x9d\xa7\x98I?\xe4\xf6;\x8b\xc3"
+ }
+}
+509
+kind: RESPONSE
+message: {
+ [type.googleapis.com/google.firestore.v1.CommitResponse]: {
+ write_results: {
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ write_results: {
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ write_results: {
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ commit_time: {
+ seconds: 1722510075
+ nanos: 148135000
+ }
+ }
+}
+ref_index: 11
+79
+kind: CREATE_STREAM
+method: "/google.firestore.v1.Firestore/BatchGetDocuments"
+264
+kind: SEND
+message: {
+ [type.googleapis.com/google.firestore.v1.BatchGetDocumentsRequest]: {
+ database: "projects/go-discovery-exp/databases/(default)"
+ documents: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple3"
+ }
+}
+ref_index: 13
+2331
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.BatchGetDocumentsResponse]: {
+ found: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple3"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.2749089002609253
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.30608412623405457
+ }
+ values: {
+ double_value: 0.28624534606933594
+ }
+ values: {
+ double_value: 0.1445397287607193
+ }
+ values: {
+ double_value: -0.7255327701568604
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721666278
+ nanos: 579066000
+ }
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ read_time: {
+ seconds: 1722510074
+ nanos: 371441000
+ }
+ }
+}
+ref_index: 13
+40
+kind: RECV
+is_error: true
+ref_index: 13
+70
+kind: CREATE_STREAM
+method: "/google.firestore.v1.Firestore/RunQuery"
+2278
+kind: SEND
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryRequest]: {
+ parent: "projects/go-discovery-exp/databases/(default)/documents"
+ structured_query: {
+ from: {
+ collection_id: "vectors"
+ }
+ find_nearest: {
+ vector_field: {
+ field_path: "Embedding"
+ }
+ query_vector: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.2746795415878296
+ }
+ values: {
+ double_value: 0.31715577840805054
+ }
+ values: {
+ double_value: 0.31715577840805054
+ }
+ values: {
+ double_value: 0.3058287799358368
+ }
+ values: {
+ double_value: 0.28600654006004333
+ }
+ values: {
+ double_value: 0.15008264780044556
+ }
+ values: {
+ double_value: -0.7249274849891663
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ distance_measure: DOT_PRODUCT
+ limit: {
+ value: 5
+ }
+ }
+ }
+ }
+}
+ref_index: 17
+2322
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryResponse]: {
+ document: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple4"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.274795264005661
+ }
+ values: {
+ double_value: 0.3172893524169922
+ }
+ values: {
+ double_value: 0.3172893524169922
+ }
+ values: {
+ double_value: 0.3059576153755188
+ }
+ values: {
+ double_value: 0.2861270308494568
+ }
+ values: {
+ double_value: 0.1473129242658615
+ }
+ values: {
+ double_value: -0.7252328395843506
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721666278
+ nanos: 579066000
+ }
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ read_time: {
+ seconds: 1722510075
+ nanos: 455689000
+ }
+ }
+}
+ref_index: 17
+2325
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryResponse]: {
+ document: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple3"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.2749089002609253
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.30608412623405457
+ }
+ values: {
+ double_value: 0.28624534606933594
+ }
+ values: {
+ double_value: 0.1445397287607193
+ }
+ values: {
+ double_value: -0.7255327701568604
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721666278
+ nanos: 579066000
+ }
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ read_time: {
+ seconds: 1722510075
+ nanos: 455689000
+ }
+ }
+}
+ref_index: 17
+2346
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryResponse]: {
+ document: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/orange1"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.30146464705467224
+ }
+ values: {
+ double_value: 0.30961233377456665
+ }
+ values: {
+ double_value: 0.2634420692920685
+ }
+ values: {
+ double_value: 0.29874876141548157
+ }
+ values: {
+ double_value: 0.2797374725341797
+ }
+ values: {
+ double_value: 0.27430567145347595
+ }
+ values: {
+ double_value: 0.13307899236679077
+ }
+ values: {
+ double_value: -0.6952698230743408
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721665713
+ nanos: 754362000
+ }
+ update_time: {
+ seconds: 1721666627
+ nanos: 948188000
+ }
+ }
+ read_time: {
+ seconds: 1722510075
+ nanos: 455689000
+ }
+ }
+}
+ref_index: 17
+2341
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryResponse]: {
+ document: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/orange2"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.30135461688041687
+ }
+ values: {
+ double_value: 0.3094993531703949
+ }
+ values: {
+ double_value: 0.2633459270000458
+ }
+ values: {
+ double_value: 0.298639714717865
+ }
+ values: {
+ double_value: 0.27963536977767944
+ }
+ values: {
+ double_value: 0.2742055654525757
+ }
+ values: {
+ double_value: 0.1357453316450119
+ }
+ values: {
+ double_value: -0.695016086101532
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721665713
+ nanos: 507024000
+ }
+ update_time: {
+ seconds: 1721666627
+ nanos: 811432000
+ }
+ }
+ read_time: {
+ seconds: 1722510075
+ nanos: 455689000
+ }
+ }
+}
+ref_index: 17
+40
+kind: RECV
+is_error: true
+ref_index: 17
+79
+kind: CREATE_STREAM
+method: "/google.firestore.v1.Firestore/BatchGetDocuments"
+253
+kind: SEND
+message: {
+ [type.googleapis.com/google.firestore.v1.BatchGetDocumentsRequest]: {
+ database: "projects/go-discovery-exp/databases/(default)"
+ documents: "projects/go-discovery-exp/databases/(default)/documents/c/d"
+ }
+}
+ref_index: 24
+262
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.BatchGetDocumentsResponse]: {
+ missing: "projects/go-discovery-exp/databases/(default)/documents/c/d"
+ read_time: {
+ seconds: 1722510075
+ nanos: 629932000
+ }
+ }
+}
+ref_index: 24
+40
+kind: RECV
+is_error: true
+ref_index: 24
+70
+kind: CREATE_STREAM
+method: "/google.firestore.v1.Firestore/RunQuery"
+2278
+kind: SEND
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryRequest]: {
+ parent: "projects/go-discovery-exp/databases/(default)/documents"
+ structured_query: {
+ from: {
+ collection_id: "vectors"
+ }
+ find_nearest: {
+ vector_field: {
+ field_path: "Embedding"
+ }
+ query_vector: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.2746795415878296
+ }
+ values: {
+ double_value: 0.31715577840805054
+ }
+ values: {
+ double_value: 0.31715577840805054
+ }
+ values: {
+ double_value: 0.3058287799358368
+ }
+ values: {
+ double_value: 0.28600654006004333
+ }
+ values: {
+ double_value: 0.15008264780044556
+ }
+ values: {
+ double_value: -0.7249274849891663
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ distance_measure: DOT_PRODUCT
+ limit: {
+ value: 3
+ }
+ }
+ }
+ }
+}
+ref_index: 28
+2322
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryResponse]: {
+ document: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple4"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.274795264005661
+ }
+ values: {
+ double_value: 0.3172893524169922
+ }
+ values: {
+ double_value: 0.3172893524169922
+ }
+ values: {
+ double_value: 0.3059576153755188
+ }
+ values: {
+ double_value: 0.2861270308494568
+ }
+ values: {
+ double_value: 0.1473129242658615
+ }
+ values: {
+ double_value: -0.7252328395843506
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721666278
+ nanos: 579066000
+ }
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ read_time: {
+ seconds: 1722510075
+ nanos: 721203000
+ }
+ }
+}
+ref_index: 28
+2325
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryResponse]: {
+ document: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/apple3"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.2749089002609253
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.3174206018447876
+ }
+ values: {
+ double_value: 0.30608412623405457
+ }
+ values: {
+ double_value: 0.28624534606933594
+ }
+ values: {
+ double_value: 0.1445397287607193
+ }
+ values: {
+ double_value: -0.7255327701568604
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721666278
+ nanos: 579066000
+ }
+ update_time: {
+ seconds: 1721666628
+ nanos: 229712000
+ }
+ }
+ read_time: {
+ seconds: 1722510075
+ nanos: 721203000
+ }
+ }
+}
+ref_index: 28
+2346
+kind: RECV
+message: {
+ [type.googleapis.com/google.firestore.v1.RunQueryResponse]: {
+ document: {
+ name: "projects/go-discovery-exp/databases/(default)/documents/vectors/orange1"
+ fields: {
+ key: "Embedding"
+ value: {
+ map_value: {
+ fields: {
+ key: "__type__"
+ value: {
+ string_value: "__vector__"
+ }
+ }
+ fields: {
+ key: "value"
+ value: {
+ array_value: {
+ values: {
+ double_value: 0.30146464705467224
+ }
+ values: {
+ double_value: 0.30961233377456665
+ }
+ values: {
+ double_value: 0.2634420692920685
+ }
+ values: {
+ double_value: 0.29874876141548157
+ }
+ values: {
+ double_value: 0.2797374725341797
+ }
+ values: {
+ double_value: 0.27430567145347595
+ }
+ values: {
+ double_value: 0.13307899236679077
+ }
+ values: {
+ double_value: -0.6952698230743408
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ values: {
+ double_value: 0
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ create_time: {
+ seconds: 1721665713
+ nanos: 754362000
+ }
+ update_time: {
+ seconds: 1721666627
+ nanos: 948188000
+ }
+ }
+ read_time: {
+ seconds: 1722510075
+ nanos: 721203000
+ }
+ }
+}
+ref_index: 28
+40
+kind: RECV
+is_error: true
+ref_index: 28
diff --git a/internal/firestore/vectordb.go b/internal/firestore/vectordb.go
new file mode 100644
index 0000000..8a54038
--- /dev/null
+++ b/internal/firestore/vectordb.go
@@ -0,0 +1,135 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package firestore
+
+import (
+ "context"
+ "net/url"
+
+ "cloud.google.com/go/firestore"
+ "golang.org/x/oscar/internal/llm"
+ "golang.org/x/oscar/internal/storage"
+ "google.golang.org/api/iterator"
+)
+
+// A VectorDB is a [storage.VectorDB] using Firestore.
+type VectorDB struct {
+ fs *fstore
+}
+
+// NewVectorDB creates a [VectorDB] with the given [DBOptions].
+// Vectors are stored in the Firestore collection "vectors".
+func NewVectorDB(ctx context.Context, dbopts *DBOptions) (*VectorDB, error) {
+ fs, err := newFstore(ctx, dbopts)
+ if err != nil {
+ return nil, err
+ }
+ return &VectorDB{fs}, nil
+}
+
+// Close closes the [VectorDB], releasing its resources.
+func (db *VectorDB) Close() {
+ db.fs.Close()
+}
+
+const vectorCollection = "vectors"
+
+// A vectorDoc holds an embedding as the Firestore type for a
+// vector of float32s.
+// All Firestore documents must be sets of key-value pairs (structs or maps in Go),
+// so we cannot represent an embedding as a lone firestore.Vector32.
+type vectorDoc struct {
+ Embedding firestore.Vector32
+}
+
+// Set implements [storage.VectorDB.Set].
+func (db *VectorDB) Set(id string, vec llm.Vector) {
+ doc := vectorDoc{firestore.Vector32(vec)}
+ if _, err := db.docref(id).Set(context.TODO(), doc); err != nil {
+ db.fs.Panic("firestore VectorDB Set", "id", id, "err", err)
+ }
+}
+
+// Get implements [storage.VectorDB.Get].
+func (db *VectorDB) Get(id string) (llm.Vector, bool) {
+ docsnap, err := db.docref(id).Get(context.TODO())
+ if err != nil {
+ if isNotFound(err) {
+ return nil, false
+ }
+ db.fs.Panic("firestore VectorDB Get", "id", id, "err", err)
+ }
+ var doc vectorDoc
+ if err := docsnap.DataTo(&doc); err != nil {
+ db.fs.Panic("firestore VectorDB Get", "id", id, "err", err)
+ }
+ return llm.Vector(doc.Embedding), true
+}
+
+// Search implements [storage.VectorDB.Search].
+func (db *VectorDB) Search(vec llm.Vector, n int) []storage.VectorResult {
+ coll := db.fs.client.Collection(vectorCollection)
+ q := coll.FindNearest("Embedding", firestore.Vector32(vec), n, firestore.DistanceMeasureDotProduct, nil)
+ iter := q.Documents(context.TODO())
+ defer iter.Stop()
+ var res []storage.VectorResult
+ for {
+ docsnap, err := iter.Next()
+ if err == iterator.Done {
+ break
+ }
+ if err != nil {
+ db.fs.Panic("firestore VectorDB Search", "err", err)
+ }
+ var doc vectorDoc
+ if err := docsnap.DataTo(&doc); err != nil {
+ db.fs.Panic("firestore VectorDB Search", "err", err)
+ }
+ id, err := url.PathUnescape(docsnap.Ref.ID)
+ if err != nil {
+ db.fs.Panic("firestore VectorDB Search unescape", "id", docsnap.Ref.ID, "err", err)
+ }
+ res = append(res, storage.VectorResult{
+ ID: id,
+ Score: vec.Dot(llm.Vector(doc.Embedding)),
+ })
+ }
+ return res
+}
+
+// docref returns a DocumentReference for the document with the given ID.
+func (db *VectorDB) docref(id string) *firestore.DocumentRef {
+ // Firestore document IDs cannot contain slashes, so escape the ID.
+ return db.fs.client.Collection(vectorCollection).Doc(url.PathEscape(id))
+}
+
+// Flush implements [storage.VectorDB.Flush]. It is a no-op.
+func (db *VectorDB) Flush() {
+ // Firestore operations do not require flushing.
+}
+
+type vBatch struct {
+ b *batch
+}
+
+// Batch implements [storage.VectorDB.Batch].
+func (db *VectorDB) Batch() storage.VectorBatch {
+ return &vBatch{db.fs.newBatch(vectorCollection)}
+}
+
+// Approximate size of a float64 encoded as a Firestore value.
+// (Firestore encodes a float32 as a float64.)
+const perFloatSize = 11
+
+// Set implements [storage.VectorBatch.Set].
+func (b *vBatch) Set(id string, vec llm.Vector) {
+ b.b.set(id, vectorDoc{firestore.Vector32(vec)}, len(vec)*perFloatSize)
+}
+
+// MaybeApply implements [storage.VectorBatch.MaybeApply].
+func (b *vBatch) MaybeApply() bool { return b.b.maybeApply() }
+
+// Apply implements [storage.VectorBatch.Apply].
+func (b *vBatch) Apply() { b.b.apply() }
diff --git a/internal/firestore/vectordb_test.go b/internal/firestore/vectordb_test.go
new file mode 100644
index 0000000..f64a329
--- /dev/null
+++ b/internal/firestore/vectordb_test.go
@@ -0,0 +1,24 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package firestore
+
+import (
+ "context"
+ "testing"
+
+ "golang.org/x/oscar/internal/storage"
+)
+
+func TestVectorDB(t *testing.T) {
+ rr, project, database := openRR(t, "testdata/vectordb.grpcrr")
+ ctx := context.Background()
+ storage.TestVectorDB(t, func() storage.VectorDB {
+ vdb, err := NewVectorDB(ctx, &DBOptions{ProjectID: project, Database: database, ClientOptions: rr.ClientOptions()})
+ if err != nil {
+ t.Fatal(err)
+ }
+ return vdb
+ })
+}
