├── README.md ├── go.mod ├── go.sum └── pgvector ├── pgvector.go └── pgvector_test.go /README.md: -------------------------------------------------------------------------------- 1 | # PostgreSQL Vector Database implementation for LangChain Go 2 | 3 | [langchaingo](https://github.com/tmc/langchaingo) extension to use [pgvector](https://github.com/pgvector/pgvector) as a vector database for your Go applications. It uses the [pgvector-go](https://github.com/pgvector/pgvector-go) library along with [pgx](https://github.com/jackc/pgx) driver. 4 | 5 | You can use this in your LangChain applications as a standalone vector database or more likely, as part of a chain. For example, in a RAG implementation: 6 | 7 | ```go 8 | import( 9 | "github.com/abhirockzz/langchain-go-postgresql-vectorstore/pgvector" 10 | //... 11 | ) 12 | func ragToRiches(){ 13 | 14 | bedrockClaudeLLM, err := claude.New("us-east-1") 15 | 16 | tableName := "test_table" 17 | textColumnName := "text_data" 18 | embeddingStoreColumnName := "embedding_data" 19 | 20 | amazonTitanEmbedder, err := titan_embedding.New("us-east-1") 21 | 22 | pgVectorStore, err := pgvector.New(pgConnString, 23 | tableName, 24 | embeddingStoreColumnName, 25 | textColumnName, 26 | false, 27 | amazonTitanEmbedder) 28 | 29 | result, err := chains.Run( 30 | context.Background(), 31 | chains.NewRetrievalQAFromLLM( 32 | bedrockClaudeLLM, 33 | vectorstores.ToRetriever(pgVectorStore, numOfResults), 34 | ), 35 | question, 36 | chains.WithMaxTokens(8091), 37 | ) 38 | } 39 | ``` -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/abhirockzz/langchain-go-postgresql-vectorstore 2 | 3 | go 1.19 4 | 5 | require ( 6 | github.com/jackc/pgx/v5 v5.4.3 7 | github.com/pgvector/pgvector-go v0.1.1 8 | github.com/stretchr/testify v1.8.4 9 | github.com/testcontainers/testcontainers-go v0.25.0 10 | github.com/tmc/langchaingo v0.0.0-20230929160525-e16b77704b8d 11 | ) 12 | 13 | require ( 14 | dario.cat/mergo v1.0.0 // indirect 15 | github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect 16 | github.com/Microsoft/go-winio v0.6.1 // indirect 17 | github.com/Microsoft/hcsshim v0.11.0 // indirect 18 | github.com/cenkalti/backoff/v4 v4.2.1 // indirect 19 | github.com/containerd/containerd v1.7.6 // indirect 20 | github.com/cpuguy83/dockercfg v0.3.1 // indirect 21 | github.com/davecgh/go-spew v1.1.1 // indirect 22 | github.com/dlclark/regexp2 v1.8.1 // indirect 23 | github.com/docker/distribution v2.8.2+incompatible // indirect 24 | github.com/docker/docker v24.0.6+incompatible // indirect 25 | github.com/docker/go-connections v0.4.0 // indirect 26 | github.com/docker/go-units v0.5.0 // indirect 27 | github.com/go-ole/go-ole v1.2.6 // indirect 28 | github.com/gogo/protobuf v1.3.2 // indirect 29 | github.com/golang/protobuf v1.5.3 // indirect 30 | github.com/google/uuid v1.3.1 // indirect 31 | github.com/jackc/pgpassfile v1.0.0 // indirect 32 | github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect 33 | github.com/jackc/puddle/v2 v2.2.1 // indirect 34 | github.com/klauspost/compress v1.16.0 // indirect 35 | github.com/kr/text v0.2.0 // indirect 36 | github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect 37 | github.com/magiconair/properties v1.8.7 // indirect 38 | github.com/moby/patternmatcher v0.5.0 // indirect 39 | github.com/moby/sys/sequential v0.5.0 // indirect 40 | github.com/moby/term v0.5.0 // indirect 41 | github.com/morikuni/aec v1.0.0 // indirect 42 | github.com/opencontainers/go-digest v1.0.0 // indirect 43 | github.com/opencontainers/image-spec v1.1.0-rc4 // indirect 44 | github.com/opencontainers/runc v1.1.5 // indirect 45 | github.com/pkg/errors v0.9.1 // indirect 46 | github.com/pkoukk/tiktoken-go v0.1.2 // indirect 47 | github.com/pmezard/go-difflib v1.0.0 // indirect 48 | github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect 49 | github.com/shirou/gopsutil/v3 v3.23.8 // indirect 50 | github.com/shoenig/go-m1cpu v0.1.6 // indirect 51 | github.com/sirupsen/logrus v1.9.0 // indirect 52 | github.com/tklauser/go-sysconf v0.3.12 // indirect 53 | github.com/tklauser/numcpus v0.6.1 // indirect 54 | github.com/yusufpapurcu/wmi v1.2.3 // indirect 55 | golang.org/x/crypto v0.13.0 // indirect 56 | golang.org/x/exp v0.0.0-20230510235704-dd950f8aeaea // indirect 57 | golang.org/x/mod v0.9.0 // indirect 58 | golang.org/x/net v0.15.0 // indirect 59 | golang.org/x/sync v0.3.0 // indirect 60 | golang.org/x/sys v0.12.0 // indirect 61 | golang.org/x/text v0.13.0 // indirect 62 | golang.org/x/tools v0.7.0 // indirect 63 | google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc // indirect 64 | google.golang.org/grpc v1.57.0 // indirect 65 | google.golang.org/protobuf v1.31.0 // indirect 66 | gopkg.in/yaml.v3 v3.0.1 // indirect 67 | ) 68 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | dario.cat/mergo v1.0.0 h1:AGCNq9Evsj31mOgNPcLyXc+4PNABt905YmuqPYYpBWk= 2 | dario.cat/mergo v1.0.0/go.mod h1:uNxQE+84aUszobStD9th8a29P2fMDhsBdgRYvZOxGmk= 3 | github.com/AdaLogics/go-fuzz-headers v0.0.0-20230811130428-ced1acdcaa24 h1:bvDV9vkmnHYOMsOr4WLk+Vo07yKIzd94sVoIqshQ4bU= 4 | github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 h1:UQHMgLO+TxOElx5B5HZ4hJQsoJ/PvUvKRhJHDQXO8P8= 5 | github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= 6 | github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= 7 | github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow= 8 | github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM= 9 | github.com/Microsoft/hcsshim v0.11.0 h1:7EFNIY4igHEXUdj1zXgAyU3fLc7QfOKHbkldRVTBdiM= 10 | github.com/Microsoft/hcsshim v0.11.0/go.mod h1:OEthFdQv/AD2RAdzR6Mm1N1KPCztGKDurW1Z8b8VGMM= 11 | github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM= 12 | github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= 13 | github.com/checkpoint-restore/go-criu/v5 v5.3.0/go.mod h1:E/eQpaFtUKGOOSEBZgmKAcn+zUUwWxqcaKZlF54wK8E= 14 | github.com/cilium/ebpf v0.7.0/go.mod h1:/oI2+1shJiTGAMgl6/RgJr36Eo1jzrRcAWbcXO2usCA= 15 | github.com/containerd/console v1.0.3/go.mod h1:7LqA/THxQ86k76b8c/EMSiaJ3h1eZkMkXar0TQ1gf3U= 16 | github.com/containerd/containerd v1.7.6 h1:oNAVsnhPoy4BTPQivLgTzI9Oleml9l/+eYIDYXRCYo8= 17 | github.com/containerd/containerd v1.7.6/go.mod h1:SY6lrkkuJT40BVNO37tlYTSnKJnP5AXBc0fhx0q+TJ4= 18 | github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= 19 | github.com/cpuguy83/dockercfg v0.3.1 h1:/FpZ+JaygUR/lZP2NlFI2DVfrOEMAIKP5wWEJdoYe9E= 20 | github.com/cpuguy83/dockercfg v0.3.1/go.mod h1:sugsbF4//dDlL/i+S+rtpIWp+5h0BHJHfjj5/jFyUJc= 21 | github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= 22 | github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= 23 | github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY= 24 | github.com/cyphar/filepath-securejoin v0.2.3/go.mod h1:aPGpWjXOXUn2NCNjFvBE6aRxGGx79pTxQpKOJNYHHl4= 25 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 26 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 27 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 28 | github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0= 29 | github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= 30 | github.com/docker/distribution v2.8.2+incompatible h1:T3de5rq0dB1j30rp0sA2rER+m322EBzniBPB6ZIzuh8= 31 | github.com/docker/distribution v2.8.2+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w= 32 | github.com/docker/docker v24.0.6+incompatible h1:hceabKCtUgDqPu+qm0NgsaXf28Ljf4/pWFL7xjWWDgE= 33 | github.com/docker/docker v24.0.6+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= 34 | github.com/docker/go-connections v0.4.0 h1:El9xVISelRB7BuFusrZozjnkIM5YnzCViNKohAFqRJQ= 35 | github.com/docker/go-connections v0.4.0/go.mod h1:Gbd7IOopHjR8Iph03tsViu4nIes5XhDvyHbTtUxmeec= 36 | github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= 37 | github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= 38 | github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= 39 | github.com/frankban/quicktest v1.11.3/go.mod h1:wRf/ReqHper53s+kmmSZizM8NamnL3IM0I9ntUbOk+k= 40 | github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= 41 | github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= 42 | github.com/go-pg/pg/v10 v10.11.0 h1:CMKJqLgTrfpE/aOVeLdybezR2om071Vh38OLZjsyMI0= 43 | github.com/go-pg/zerochecker v0.2.0 h1:pp7f72c3DobMWOb2ErtZsnrPaSvHd2W4o9//8HtF4mU= 44 | github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= 45 | github.com/godbus/dbus/v5 v5.0.6/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= 46 | github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= 47 | github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= 48 | github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= 49 | github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg= 50 | github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= 51 | github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 52 | github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 53 | github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 54 | github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= 55 | github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= 56 | github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4= 57 | github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 58 | github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= 59 | github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= 60 | github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk= 61 | github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= 62 | github.com/jackc/pgx/v5 v5.4.3 h1:cxFyXhxlvAifxnkKKdlxv8XqUf59tDlYjnV5YYfsJJY= 63 | github.com/jackc/pgx/v5 v5.4.3/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA= 64 | github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk= 65 | github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= 66 | github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= 67 | github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= 68 | github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= 69 | github.com/klauspost/compress v1.16.0 h1:iULayQNOReoYUe+1qtKOqw9CwJv3aNQu8ivo7lw1HU4= 70 | github.com/klauspost/compress v1.16.0/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE= 71 | github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= 72 | github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= 73 | github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= 74 | github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= 75 | github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= 76 | github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= 77 | github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= 78 | github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= 79 | github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY= 80 | github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= 81 | github.com/moby/patternmatcher v0.5.0 h1:YCZgJOeULcxLw1Q+sVR636pmS7sPEn1Qo2iAN6M7DBo= 82 | github.com/moby/patternmatcher v0.5.0/go.mod h1:hDPoyOpDY7OrrMDLaYoY3hf52gNCR/YOUYxkhApJIxc= 83 | github.com/moby/sys/mountinfo v0.5.0/go.mod h1:3bMD3Rg+zkqx8MRYPi7Pyb0Ie97QEBmdxbhnCLlSvSU= 84 | github.com/moby/sys/sequential v0.5.0 h1:OPvI35Lzn9K04PBbCLW0g4LcFAJgHsvXsRyewg5lXtc= 85 | github.com/moby/sys/sequential v0.5.0/go.mod h1:tH2cOOs5V9MlPiXcQzRC+eEyab644PWKGRYaaV5ZZlo= 86 | github.com/moby/term v0.5.0 h1:xt8Q1nalod/v7BqbG21f8mQPqH+xAaC9C3N3wfWbVP0= 87 | github.com/moby/term v0.5.0/go.mod h1:8FzsFHVUBGZdbDsJw/ot+X+d5HLUbvklYLJ9uGfcI3Y= 88 | github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= 89 | github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= 90 | github.com/mrunalp/fileutils v0.5.0/go.mod h1:M1WthSahJixYnrXQl/DFQuteStB1weuxD2QJNHXfbSQ= 91 | github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= 92 | github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= 93 | github.com/opencontainers/image-spec v1.1.0-rc4 h1:oOxKUJWnFC4YGHCCMNql1x4YaDfYBTS5Y4x/Cgeo1E0= 94 | github.com/opencontainers/image-spec v1.1.0-rc4/go.mod h1:X4pATf0uXsnn3g5aiGIsVnJBR4mxhKzfwmvK/B2NTm8= 95 | github.com/opencontainers/runc v1.1.5 h1:L44KXEpKmfWDcS02aeGm8QNTFXTo2D+8MYGDIJ/GDEs= 96 | github.com/opencontainers/runc v1.1.5/go.mod h1:1J5XiS+vdZ3wCyZybsuxXZWGrgSr8fFJHLXuG2PsnNg= 97 | github.com/opencontainers/runtime-spec v1.0.3-0.20210326190908-1c3f411f0417/go.mod h1:jwyrGlmzljRJv/Fgzds9SsS/C5hL+LL3ko9hs6T5lQ0= 98 | github.com/opencontainers/selinux v1.10.0/go.mod h1:2i0OySw99QjzBBQByd1Gr9gSjvuho1lHsJxIJ3gGbJI= 99 | github.com/pgvector/pgvector-go v0.1.1 h1:kqJigGctFnlWvskUiYIvJRNwUtQl/aMSUZVs0YWQe+g= 100 | github.com/pgvector/pgvector-go v0.1.1/go.mod h1:wLJgD/ODkdtd2LJK4l6evHXTuG+8PxymYAVomKHOWac= 101 | github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= 102 | github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 103 | github.com/pkoukk/tiktoken-go v0.1.2 h1:u7PCSBiWJ3nJYoTGShyM9iHXz4dNyYkurwwp+GHtyHY= 104 | github.com/pkoukk/tiktoken-go v0.1.2/go.mod h1:boMWvk9pQCOTx11pgu0DrIdrAKgQzzJKUP6vLXaz7Rw= 105 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 106 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 107 | github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw= 108 | github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= 109 | github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M= 110 | github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= 111 | github.com/seccomp/libseccomp-golang v0.9.2-0.20220502022130-f33da4d89646/go.mod h1:JA8cRccbGaA1s33RQf7Y1+q9gHmZX1yB/z9WDN1C6fg= 112 | github.com/shirou/gopsutil/v3 v3.23.8 h1:xnATPiybo6GgdRoC4YoGnxXZFRc3dqQTGi73oLvvBrE= 113 | github.com/shirou/gopsutil/v3 v3.23.8/go.mod h1:7hmCaBn+2ZwaZOr6jmPBZDfawwMGuo1id3C6aM8EDqQ= 114 | github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= 115 | github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= 116 | github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= 117 | github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k= 118 | github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc= 119 | github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0= 120 | github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0= 121 | github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= 122 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 123 | github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= 124 | github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= 125 | github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= 126 | github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= 127 | github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 128 | github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 129 | github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= 130 | github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= 131 | github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= 132 | github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635/go.mod h1:hkRG7XYTFWNJGYcbNJQlaLq0fg1yr4J4t/NcTQtrfww= 133 | github.com/testcontainers/testcontainers-go v0.25.0 h1:erH6cQjsaJrH+rJDU9qIf89KFdhK0Bft0aEZHlYC3Vs= 134 | github.com/testcontainers/testcontainers-go v0.25.0/go.mod h1:4sC9SiJyzD1XFi59q8umTQYWxnkweEc5OjVtTUlJzqQ= 135 | github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= 136 | github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= 137 | github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= 138 | github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= 139 | github.com/tmc/langchaingo v0.0.0-20230929160525-e16b77704b8d h1:i4+wYULVM2/3Yb/aDE7Z4s2v7vqtQERQWh5lopBEuig= 140 | github.com/tmc/langchaingo v0.0.0-20230929160525-e16b77704b8d/go.mod h1:R+a8fqt6nmKyYj7KSpr/m9oxqE6OJLbLyO9pxeHpjLU= 141 | github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc h1:9lRDQMhESg+zvGYmW5DyG0UqvY96Bu5QYsTLvCHdrgo= 142 | github.com/uptrace/bun v1.1.12 h1:sOjDVHxNTuM6dNGaba0wUuz7KvDE1BmNu9Gqs2gJSXQ= 143 | github.com/uptrace/bun/dialect/pgdialect v1.1.12 h1:m/CM1UfOkoBTglGO5CUTKnIKKOApOYxkcP2qn0F9tJk= 144 | github.com/uptrace/bun/driver/pgdriver v1.1.12 h1:3rRWB1GK0psTJrHwxzNfEij2MLibggiLdTqjTtfHc1w= 145 | github.com/urfave/cli v1.22.1/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0= 146 | github.com/vishvananda/netlink v1.1.0/go.mod h1:cTgwzPIzzgDAYoQrMm0EdrjRUBkTqKYppBueQtXaqoE= 147 | github.com/vishvananda/netns v0.0.0-20191106174202-0a2b9b5464df/go.mod h1:JP3t17pCcGlemwknint6hfoeCVQrEMVwxRLRjXpq+BU= 148 | github.com/vmihailenco/bufpool v0.1.11 h1:gOq2WmBrq0i2yW5QJ16ykccQ4wH9UyEsgLm6czKAd94= 149 | github.com/vmihailenco/msgpack/v5 v5.3.5 h1:5gO0H1iULLWGhs2H5tbAHIZTV8/cYafcFOr9znI5mJU= 150 | github.com/vmihailenco/tagparser v0.1.2 h1:gnjoVuB/kljJ5wICEEOpx98oXMWPLj22G67Vbd1qPqc= 151 | github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g= 152 | github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= 153 | github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= 154 | github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= 155 | github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= 156 | golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= 157 | golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= 158 | golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= 159 | golang.org/x/crypto v0.13.0 h1:mvySKfSWJ+UKUii46M40LOvyWfN0s2U+46/jDd0e6Ck= 160 | golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc= 161 | golang.org/x/exp v0.0.0-20230510235704-dd950f8aeaea h1:vLCWI/yYrdEHyN2JzIzPO3aaQJHQdp89IZBA/+azVC4= 162 | golang.org/x/exp v0.0.0-20230510235704-dd950f8aeaea/go.mod h1:V1LtkGg67GoY2N1AnLN78QLrzxkLyJw7RJb1gzOOz9w= 163 | golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= 164 | golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= 165 | golang.org/x/mod v0.9.0 h1:KENHtAZL2y3NLMYZeHY9DW8HW8V+kQyJsY/V9JlKvCs= 166 | golang.org/x/mod v0.9.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= 167 | golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= 168 | golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= 169 | golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= 170 | golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= 171 | golang.org/x/net v0.0.0-20201224014010-6772e930b67b/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= 172 | golang.org/x/net v0.15.0 h1:ugBLEUaxABaB5AJqW9enI0ACdci2RUd4eP51NTBvuJ8= 173 | golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk= 174 | golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 175 | golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 176 | golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 177 | golang.org/x/sync v0.3.0 h1:ftCYgMx6zT/asHUrPw8BLLscYtGznsLAnjq5RH9P66E= 178 | golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= 179 | golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 180 | golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 181 | golang.org/x/sys v0.0.0-20190606203320-7fc4e5ec1444/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 182 | golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 183 | golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 184 | golang.org/x/sys v0.0.0-20191115151921-52ab43148777/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 185 | golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 186 | golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 187 | golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 188 | golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 189 | golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 190 | golang.org/x/sys v0.0.0-20210906170528-6f6e22806c34/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 191 | golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 192 | golang.org/x/sys v0.0.0-20211116061358-0a5406a5449c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 193 | golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 194 | golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 195 | golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 196 | golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o= 197 | golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 198 | golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= 199 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 200 | golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= 201 | golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k= 202 | golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= 203 | golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4= 204 | golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= 205 | golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= 206 | golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= 207 | golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= 208 | golang.org/x/tools v0.7.0 h1:W4OVu8VVOaIO0yzWMNdepAulS7YfoS3Zabrm8DOXXU4= 209 | golang.org/x/tools v0.7.0/go.mod h1:4pg6aUX35JBAogB10C9AtvVL+qowtN4pT3CGSQex14s= 210 | golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 211 | golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 212 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 213 | golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 214 | google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc h1:XSJ8Vk1SWuNr8S18z1NZSziL0CPIXLCCMDOEFtHBOFc= 215 | google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc/go.mod h1:66JfowdXAEgad5O9NnYcsNPLCPZJD++2L9X0PCMODrA= 216 | google.golang.org/grpc v1.57.0 h1:kfzNeI/klCGD2YPMUlaGNT3pxvYfga7smW3Vth8Zsiw= 217 | google.golang.org/grpc v1.57.0/go.mod h1:Sd+9RMTACXwmub0zcNY2c4arhtrbBYD1AUHI/dt16Mo= 218 | google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= 219 | google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= 220 | google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= 221 | google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8= 222 | google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= 223 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 224 | gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= 225 | gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 226 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 227 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 228 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 229 | gotest.tools/v3 v3.5.0 h1:Ljk6PdHdOhAb5aDMWXjDLMMhph+BpztA4v1QdqEW2eY= 230 | mellium.im/sasl v0.3.1 h1:wE0LW6g7U83vhvxjC1IY8DnXM+EU095yeo8XClvCdfo= 231 | -------------------------------------------------------------------------------- /pgvector/pgvector.go: -------------------------------------------------------------------------------- 1 | package pgvector 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "strings" 8 | 9 | "github.com/jackc/pgx/v5/pgxpool" 10 | pgv "github.com/pgvector/pgvector-go" 11 | "github.com/tmc/langchaingo/embeddings" 12 | "github.com/tmc/langchaingo/schema" 13 | "github.com/tmc/langchaingo/vectorstores" 14 | ) 15 | 16 | type Store struct { 17 | embedder embeddings.Embedder 18 | pool *pgxpool.Pool 19 | 20 | // name of the table 21 | tableName string 22 | 23 | //name of the column in which text data will be stored. this will come from the "PageContent" of the langchain doc 24 | textColumnName string 25 | 26 | //name of the column in which embedded vector data will be stored. the data from "PageContent" of the langchain doc will go through the embedding (vector creation) process 27 | embeddingStoreColumnName string 28 | 29 | //if true, the langchain doc Metadata will be saved to postgresql as well. in that case the, column(s) needs to exist in advance 30 | saveMetadata bool 31 | 32 | // attributes for similarity search 33 | //searchKey string // name of the column whose value needs to returned by search 34 | 35 | //optional - data for these columns will be added to resulting langchain doc Metadata 36 | QueryAttributes []string 37 | } 38 | 39 | func New(pgConnectionString, tableName, embeddingStoreColumnName, textColumnName string, saveMetadata bool, embedder embeddings.Embedder) (Store, error) { 40 | //connection string example - postgres://postgres:postgres@localhost/postgres 41 | pool, err := pgxpool.New(context.Background(), pgConnectionString) 42 | 43 | if err != nil { 44 | return Store{}, err 45 | } 46 | 47 | return Store{embedder: embedder, 48 | tableName: tableName, 49 | embeddingStoreColumnName: embeddingStoreColumnName, 50 | textColumnName: textColumnName, 51 | pool: pool, 52 | saveMetadata: saveMetadata}, nil 53 | } 54 | 55 | var ErrEmbedderWrongNumberVectors = errors.New( 56 | "number of vectors from embedder does not match number of documents", 57 | ) 58 | 59 | func (store Store) AddDocuments(ctx context.Context, docs []schema.Document, options ...vectorstores.Option) error { 60 | 61 | texts := make([]string, 0, len(docs)) 62 | for _, doc := range docs { 63 | texts = append(texts, doc.PageContent) 64 | } 65 | 66 | vectors, err := store.embedder.EmbedDocuments(ctx, texts) 67 | if err != nil { 68 | return err 69 | } 70 | 71 | if len(vectors) != len(docs) { 72 | return ErrEmbedderWrongNumberVectors 73 | } 74 | 75 | metadatas := make([]map[string]any, 0, len(docs)) 76 | 77 | for i := 0; i < len(docs); i++ { 78 | metadata := make(map[string]any, len(docs[i].Metadata)) 79 | for key, value := range docs[i].Metadata { 80 | metadata[key] = value 81 | } 82 | 83 | metadatas = append(metadatas, metadata) 84 | } 85 | 86 | for i, doc := range docs { 87 | 88 | data := map[string]any{} 89 | data[store.embeddingStoreColumnName] = pgv.NewVector(vectors[i]) 90 | data[store.textColumnName] = doc.PageContent 91 | 92 | metadata := metadatas[i] 93 | 94 | query, values := store.generateInsertQueryWithValues(data, metadata) 95 | //fmt.Println("generated query:", query) 96 | 97 | _, err := store.pool.Exec(context.Background(), query, values...) 98 | if err != nil { 99 | return err 100 | } 101 | } 102 | 103 | return nil 104 | } 105 | 106 | func (store Store) generateInsertQueryWithValues(data, metadata map[string]any) (string, []any) { 107 | 108 | //INSERT INTO test_table (data, embedding) VALUES ($1, $2) 109 | //INSERT INTO test_table (data, embedding, other_data) VALUES ($1, $2, $3) 110 | 111 | // generate column names and placeholders dynamically 112 | var columns []string 113 | var placeholders []string 114 | var values []any 115 | 116 | for column, value := range data { 117 | columns = append(columns, column) 118 | placeholders = append(placeholders, fmt.Sprintf("$%d", len(placeholders)+1)) 119 | values = append(values, value) 120 | } 121 | 122 | if store.saveMetadata { 123 | for column, value := range metadata { 124 | columns = append(columns, column) 125 | placeholders = append(placeholders, fmt.Sprintf("$%d", len(placeholders)+1)) 126 | values = append(values, value) 127 | } 128 | } 129 | 130 | sqlQuery := fmt.Sprintf( 131 | "INSERT INTO %s (%s) VALUES (%s)", 132 | store.tableName, 133 | strings.Join(columns, ", "), 134 | strings.Join(placeholders, ", "), 135 | ) 136 | 137 | return sqlQuery, values 138 | } 139 | 140 | func (store Store) SimilaritySearch(ctx context.Context, searchString string, numDocuments int, options ...vectorstores.Option) ([]schema.Document, error) { 141 | 142 | //fmt.Println("similarity search for", searchString, "with max docs", numDocuments) 143 | 144 | vector, err := store.embedder.EmbedQuery(ctx, searchString) 145 | if err != nil { 146 | return nil, err 147 | } 148 | 149 | opts := vectorstores.Options{} 150 | for _, opt := range options { 151 | opt(&opts) 152 | } 153 | 154 | query := store.generateSelectQuery(numDocuments, opts.ScoreThreshold) 155 | 156 | rows, err := store.pool.Query(context.Background(), query, pgv.NewVector(vector)) 157 | 158 | if err != nil { 159 | return nil, err 160 | } 161 | 162 | defer rows.Close() 163 | 164 | docs := []schema.Document{} 165 | doc := schema.Document{} 166 | 167 | for rows.Next() { 168 | vals, err := rows.Values() 169 | 170 | if err != nil { 171 | return nil, err 172 | } 173 | 174 | doc.PageContent = vals[0].(string) 175 | 176 | score := vals[1].(float64) 177 | doc.Score = float32(score) 178 | 179 | metadata := make(map[string]any) 180 | for i := 2; i <= len(vals)-1; i++ { 181 | metadata[store.QueryAttributes[i-2]] = vals[i] 182 | } 183 | 184 | doc.Metadata = metadata 185 | 186 | docs = append(docs, doc) 187 | } 188 | 189 | return docs, nil 190 | } 191 | 192 | const queryFormatWithQueryAttributes = "SELECT %s, 1 - (%s <=> $1) as similarity_score, %s FROM %s WHERE 1 - (embedding <=> $1) > %v ORDER BY similarity_score DESC LIMIT %d" 193 | 194 | const queryFormat = "SELECT %s, 1 - (%s <=> $1) as similarity_score FROM %s WHERE 1 - (embedding <=> $1) > %v ORDER BY similarity_score DESC LIMIT %d" 195 | 196 | func (store Store) generateSelectQuery(numDocuments int, threshold float32) string { 197 | 198 | //SELECT data, 1 - (embedding <=> $1) as similarity_score FROM test_table WHERE 1 - (embedding <=> $1) > 0.5 ORDER BY similarity_score DESC LIMIT 5 199 | //SELECT data, 1 - (embedding <=> $1) as similarity_score, other_data FROM test_table WHERE 1 - (embedding <=> $1) > 0.5 ORDER BY similarity_score DESC LIMIT 5 200 | 201 | var sqlQuery string 202 | 203 | if len(store.QueryAttributes) > 0 { 204 | 205 | sqlQuery = fmt.Sprintf(queryFormatWithQueryAttributes, store.textColumnName, store.embeddingStoreColumnName, strings.Join(store.QueryAttributes, ","), store.tableName, threshold, numDocuments) 206 | 207 | } else { 208 | sqlQuery = fmt.Sprintf(queryFormat, store.textColumnName, store.embeddingStoreColumnName, store.tableName, threshold, numDocuments) 209 | } 210 | 211 | //fmt.Println("search query -", sqlQuery) 212 | 213 | return sqlQuery 214 | } 215 | -------------------------------------------------------------------------------- /pgvector/pgvector_test.go: -------------------------------------------------------------------------------- 1 | package pgvector 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | pgv "github.com/pgvector/pgvector-go" 8 | "github.com/stretchr/testify/assert" 9 | testcontainers "github.com/testcontainers/testcontainers-go" 10 | "github.com/testcontainers/testcontainers-go/wait" 11 | 12 | "github.com/tmc/langchaingo/schema" 13 | ) 14 | 15 | func TestAddDocuments(t *testing.T) { 16 | 17 | tableName := "test_table" 18 | textColumnName := "data" 19 | embeddingStoreColumnName := "embedding" 20 | 21 | postgresContainer, err := startPostgresContainer() 22 | assert.Nil(t, err) 23 | 24 | ctx := context.Background() 25 | 26 | defer postgresContainer.Terminate(ctx) 27 | 28 | ip, err := postgresContainer.Host(ctx) 29 | assert.Nil(t, err) 30 | 31 | port, err := postgresContainer.MappedPort(ctx, "5432") 32 | assert.Nil(t, err) 33 | 34 | pgConnString := "postgres://postgres:postgres@" + ip + ":" + port.Port() + "/postgres" 35 | 36 | saveDocMetadtaToTable := false 37 | 38 | pgStore, err := New(pgConnString, tableName, embeddingStoreColumnName, textColumnName, saveDocMetadtaToTable, MockEmbedder{}) 39 | assert.Nil(t, err) 40 | 41 | //create extension and table 42 | _, err = pgStore.pool.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector;") 43 | assert.Nil(t, err) 44 | 45 | _, err = pgStore.pool.Exec(ctx, "CREATE TABLE IF NOT EXISTS test_table (id bigserial primary key, data text, embedding vector(1));") 46 | assert.Nil(t, err) 47 | 48 | docs := []schema.Document{{PageContent: "foo"}} 49 | 50 | err = pgStore.AddDocuments(context.Background(), docs) 51 | assert.Nil(t, err) 52 | 53 | rows, err := pgStore.pool.Query(context.Background(), "select count(*) from test_table;") 54 | assert.Nil(t, err) 55 | 56 | for rows.Next() { 57 | var count int 58 | err = rows.Scan(&count) 59 | 60 | assert.Nil(t, err) 61 | assert.Equal(t, 1, count) 62 | } 63 | 64 | rows, err = pgStore.pool.Query(context.Background(), "select data, embedding from test_table;") 65 | assert.Nil(t, err) 66 | 67 | for rows.Next() { 68 | 69 | var data string 70 | var embedding pgv.Vector 71 | 72 | err = rows.Scan(&data, &embedding) 73 | 74 | assert.Nil(t, err) 75 | assert.Equal(t, "foo", data) 76 | assert.Equal(t, []float32{42.42}, embedding.Slice()) 77 | } 78 | } 79 | 80 | func TestAddDocumentsWithMetadata(t *testing.T) { 81 | 82 | tableName := "test_table" 83 | textColumnName := "data" 84 | embeddingStoreColumnName := "embedding" 85 | 86 | postgresContainer, err := startPostgresContainer() 87 | assert.Nil(t, err) 88 | 89 | ctx := context.Background() 90 | 91 | defer postgresContainer.Terminate(ctx) 92 | 93 | ip, err := postgresContainer.Host(ctx) 94 | assert.Nil(t, err) 95 | 96 | port, err := postgresContainer.MappedPort(ctx, "5432") 97 | assert.Nil(t, err) 98 | 99 | pgConnString := "postgres://postgres:postgres@" + ip + ":" + port.Port() + "/postgres" 100 | 101 | saveDocMetadtaToTable := true 102 | 103 | pgStore, err := New(pgConnString, tableName, embeddingStoreColumnName, textColumnName, saveDocMetadtaToTable, MockEmbedder{}) 104 | assert.Nil(t, err) 105 | 106 | //create extension and table 107 | _, err = pgStore.pool.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector;") 108 | assert.Nil(t, err) 109 | 110 | _, err = pgStore.pool.Exec(ctx, "CREATE TABLE IF NOT EXISTS test_table (id bigserial primary key, data text, other_data text, embedding vector(1));") 111 | 112 | assert.Nil(t, err) 113 | 114 | docs := []schema.Document{{PageContent: "foo", Metadata: map[string]any{"other_data": "bar"}}} 115 | 116 | err = pgStore.AddDocuments(context.Background(), docs) 117 | assert.Nil(t, err) 118 | 119 | rows, err := pgStore.pool.Query(context.Background(), "select count(*) from test_table;") 120 | assert.Nil(t, err) 121 | 122 | for rows.Next() { 123 | var count int 124 | err = rows.Scan(&count) 125 | 126 | assert.Nil(t, err) 127 | assert.Equal(t, 1, count) 128 | } 129 | 130 | rows, err = pgStore.pool.Query(context.Background(), "select data, embedding, other_data from test_table;") 131 | assert.Nil(t, err) 132 | 133 | for rows.Next() { 134 | 135 | var data string 136 | var embedding pgv.Vector 137 | var metadata string 138 | 139 | err = rows.Scan(&data, &embedding, &metadata) 140 | 141 | assert.Nil(t, err) 142 | assert.Equal(t, "foo", data) 143 | assert.Equal(t, []float32{42.42}, embedding.Slice()) 144 | assert.Equal(t, "bar", metadata) 145 | 146 | } 147 | } 148 | 149 | func TestSimilaritySearch(t *testing.T) { 150 | 151 | tableName := "test_table" 152 | textColumnName := "data" 153 | embeddingStoreColumnName := "embedding" 154 | 155 | postgresContainer, err := startPostgresContainer() 156 | assert.Nil(t, err) 157 | 158 | ctx := context.Background() 159 | 160 | defer postgresContainer.Terminate(ctx) 161 | 162 | ip, err := postgresContainer.Host(ctx) 163 | assert.Nil(t, err) 164 | 165 | port, err := postgresContainer.MappedPort(ctx, "5432") 166 | assert.Nil(t, err) 167 | 168 | pgConnString := "postgres://postgres:postgres@" + ip + ":" + port.Port() + "/postgres" 169 | 170 | saveDocMetadtaToTable := false 171 | 172 | pgStore, err := New(pgConnString, tableName, embeddingStoreColumnName, textColumnName, saveDocMetadtaToTable, MockEmbedder{}) 173 | assert.Nil(t, err) 174 | 175 | //create extension and table 176 | _, err = pgStore.pool.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector;") 177 | assert.Nil(t, err) 178 | 179 | _, err = pgStore.pool.Exec(ctx, "CREATE TABLE IF NOT EXISTS test_table (id bigserial primary key, data text, embedding vector(1));") 180 | assert.Nil(t, err) 181 | 182 | docs := []schema.Document{{PageContent: "foo"}} 183 | 184 | err = pgStore.AddDocuments(context.Background(), docs) 185 | assert.Nil(t, err) 186 | 187 | searchResults, err := pgStore.SimilaritySearch(ctx, "doesn't really matter", 1) 188 | assert.Nil(t, err) 189 | 190 | assert.Equal(t, 1, len(searchResults)) 191 | assert.Equal(t, "foo", searchResults[0].PageContent) 192 | assert.Equal(t, float32(1), searchResults[0].Score) 193 | 194 | } 195 | 196 | func TestGenerateSelectQuery(t *testing.T) { 197 | 198 | tableName := "test_table" 199 | textColumnName := "data" 200 | embeddingStoreColumnName := "embedding" 201 | saveDocMetadtaToTable := false 202 | 203 | pgStore, err := New("postgres://postgres:postgres@localhost/postgres", tableName, embeddingStoreColumnName, textColumnName, saveDocMetadtaToTable, MockEmbedder{}) 204 | assert.Nil(t, err) 205 | 206 | //pgStore.QueryAttributes = []string{} 207 | 208 | query := pgStore.generateSelectQuery(5, 0.5) 209 | expectedQuery := "SELECT data, 1 - (embedding <=> $1) as similarity_score FROM test_table WHERE 1 - (embedding <=> $1) > 0.5 ORDER BY similarity_score DESC LIMIT 5" 210 | 211 | assert.Equal(t, expectedQuery, query) 212 | } 213 | 214 | func TestGenerateSelectQueryWithQueryAttributes(t *testing.T) { 215 | 216 | tableName := "test_table" 217 | textColumnName := "data" 218 | embeddingStoreColumnName := "embedding" 219 | saveDocMetadtaToTable := false 220 | 221 | pgStore, err := New("postgres://postgres:postgres@localhost/postgres", tableName, embeddingStoreColumnName, textColumnName, saveDocMetadtaToTable, MockEmbedder{}) 222 | assert.Nil(t, err) 223 | 224 | pgStore.QueryAttributes = []string{"other_data"} 225 | 226 | query := pgStore.generateSelectQuery(5, 0.5) 227 | expectedQuery := "SELECT data, 1 - (embedding <=> $1) as similarity_score, other_data FROM test_table WHERE 1 - (embedding <=> $1) > 0.5 ORDER BY similarity_score DESC LIMIT 5" 228 | 229 | assert.Equal(t, expectedQuery, query) 230 | } 231 | 232 | func TestGenerateInsertQueryWithValues(t *testing.T) { 233 | 234 | tableName := "test_table" 235 | textColumnName := "data" 236 | embeddingStoreColumnName := "embedding" 237 | saveDocMetadtaToTable := false 238 | 239 | pgStore, err := New("postgres://postgres:postgres@localhost/postgres", tableName, embeddingStoreColumnName, textColumnName, saveDocMetadtaToTable, MockEmbedder{}) 240 | assert.Nil(t, err) 241 | 242 | pgStore.saveMetadata = false 243 | 244 | //pgStore.QueryAttributes = []string{"other_data"} 245 | 246 | query, values := pgStore.generateInsertQueryWithValues(map[string]any{"data": "foo", "embedding": pgv.NewVector([]float32{42})}, nil) 247 | 248 | expectedQuery := "INSERT INTO test_table (data, embedding) VALUES ($1, $2)" 249 | assert.Equal(t, expectedQuery, query) 250 | 251 | expectedValues := []any{"foo", pgv.NewVector([]float32{42})} 252 | assert.Equal(t, expectedValues, values) 253 | } 254 | 255 | func TestGenerateInsertQueryWithValuesAndMetadata(t *testing.T) { 256 | 257 | tableName := "test_table" 258 | textColumnName := "data" 259 | embeddingStoreColumnName := "embedding" 260 | saveDocMetadtaToTable := false 261 | 262 | pgStore, err := New("postgres://postgres:postgres@localhost/postgres", tableName, embeddingStoreColumnName, textColumnName, saveDocMetadtaToTable, MockEmbedder{}) 263 | assert.Nil(t, err) 264 | 265 | pgStore.saveMetadata = true 266 | 267 | //pgStore.QueryAttributes = []string{"other_data"} 268 | 269 | query, values := pgStore.generateInsertQueryWithValues(map[string]any{"data": "foo", "embedding": pgv.NewVector([]float32{42})}, map[string]any{"other_data": "bar"}) 270 | 271 | expectedQuery := "INSERT INTO test_table (data, embedding, other_data) VALUES ($1, $2, $3)" 272 | assert.Equal(t, expectedQuery, query) 273 | 274 | expectedValues := []any{"foo", pgv.NewVector([]float32{42}), "bar"} 275 | assert.Equal(t, expectedValues, values) 276 | 277 | } 278 | 279 | type MockEmbedder struct{} 280 | 281 | func (m MockEmbedder) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error) { 282 | return [][]float32{{42.42}}, nil 283 | } 284 | 285 | func (m MockEmbedder) EmbedQuery(ctx context.Context, text string) ([]float32, error) { 286 | return []float32{42.42}, nil 287 | } 288 | 289 | func startPostgresContainer() (testcontainers.Container, error) { 290 | ctx := context.Background() 291 | 292 | req := testcontainers.ContainerRequest{ 293 | Image: "ankane/pgvector:latest", 294 | ExposedPorts: []string{"5432/tcp"}, 295 | WaitingFor: wait.ForListeningPort("5432/tcp"), 296 | Env: map[string]string{ 297 | "POSTGRES_PASSWORD": "postgres", 298 | "POSTGRES_USER": "postgres", 299 | }, 300 | } 301 | postgresContainer, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ 302 | ContainerRequest: req, 303 | Started: true, 304 | }) 305 | if err != nil { 306 | return nil, err 307 | } 308 | 309 | return postgresContainer, nil 310 | } 311 | --------------------------------------------------------------------------------