├── .coveragerc
├── .github
├── ISSUE_TEMPLATE
│ ├── bug_report.md
│ └── rfc.md
└── workflows
│ ├── pre-commit.yaml
│ └── tests.yml
├── .gitignore
├── .pre-commit-config.yaml
├── LICENSE
├── README.md
├── docs
├── api.md
├── assets
│ ├── hf_configure_endpoint.png
│ ├── hf_configure_endpoint_2.png
│ ├── hf_endpoint_url.png
│ ├── hf_new_endpoint.png
│ ├── supabase_connection_info.png
│ ├── supabase_new_project.png
│ ├── supabase_new_project_prompt.png
│ └── supabase_sign_up.png
├── concepts_adapters.md
├── concepts_collections.md
├── concepts_indexes.md
├── concepts_metadata.md
├── hosting.md
├── index.md
├── integrations_amazon_bedrock.md
├── integrations_huggingface_inference_endpoints.md
├── integrations_openai.md
└── support_changelog.md
├── mkdocs.yml
├── pyproject.toml
├── pytest.ini
├── setup.py
└── src
├── py.typed
├── tests
├── conftest.py
├── test_adapters.py
├── test_client.py
├── test_collection.py
└── test_issue_90.py
└── vecs
├── __init__.py
├── adapter
├── __init__.py
├── base.py
├── markdown.py
├── noop.py
└── text.py
├── client.py
├── collection.py
└── exc.py
/.coveragerc:
--------------------------------------------------------------------------------
1 | [report]
2 | exclude_lines =
3 | pragma: no cover
4 | raise AssertionError
5 | raise NotImplementedError
6 | Unreachable
7 | if __name__ == .__main__.:
8 | if TYPE_CHECKING:
9 | pass
10 | except ImportError
11 |
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/bug_report.md:
--------------------------------------------------------------------------------
1 | ---
2 | name: Bug report
3 | about: Create a report to help us improve
4 | title: ''
5 | labels: triage-required
6 | assignees: olirice
7 |
8 | ---
9 |
10 | **Describe the bug**
11 | A clear and concise description of what the bug is.
12 |
13 | **To Reproduce**
14 | Steps to reproduce the behavior:
15 | 1.
16 | 2.
17 | 3.
18 | 4.
19 |
20 | **Expected behavior**
21 | A clear and concise description of what you expected to happen.
22 |
23 | **Screenshots**
24 | If applicable, add screenshots to help explain your problem.
25 |
26 | **Versions:**
27 | - PostgreSQL: [e.g. 14.1]
28 | - vecs version: e.g. 0.2.6
29 |
30 | **Additional context**
31 | Add any other context about the problem here.
32 |
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/rfc.md:
--------------------------------------------------------------------------------
1 | ---
2 | name: RFC
3 | about: Request for Comment
4 | title: ''
5 | labels: ''
6 | assignees: ''
7 |
8 | ---
9 |
10 | # Summary
11 | [summary]: #summary
12 |
13 | Short explanation of the feature.
14 |
15 | # Rationale
16 | [rationale]: #rationale
17 |
18 | Why should we do this?
19 |
20 | # Design
21 | [design]: #design
22 |
23 | An dense explanation in sufficient detail that someone familiar with the
24 | project could implement the feature. Specifics and corner cases should be covered.
25 |
26 | # Examples
27 | [examples]: #examples
28 |
29 | Illustrations and examples to clarify descriptions from previous sections.
30 |
31 | # Drawbacks
32 | [drawbacks]: #drawbacks
33 |
34 | What are the negative trade-offs?
35 |
36 | # Alternatives
37 | [alternatives]: #alternatives
38 |
39 | What other solutions have been considered?
40 |
41 | # Unresolved Questions
42 | [unresolved]: #unresolved-questions
43 |
44 | What parts of problem space or proposed designs are unknown or TBD?
45 |
--------------------------------------------------------------------------------
/.github/workflows/pre-commit.yaml:
--------------------------------------------------------------------------------
1 | name: pre-commit
2 |
3 | on: [push, pull_request]
4 |
5 | jobs:
6 | build:
7 | runs-on: ubuntu-latest
8 |
9 | steps:
10 | - name: checkout
11 | uses: actions/checkout@v3
12 |
13 | - name: set up python 3.9
14 | uses: actions/setup-python@v4
15 | with:
16 | python-version: 3.9
17 |
18 | - name: install pre-commit
19 | run: |
20 | python -m pip install --upgrade pip
21 | pip install pre-commit
22 |
23 | - name: run pre-commit hooks
24 | run: |
25 | pre-commit run --all-files
26 |
--------------------------------------------------------------------------------
/.github/workflows/tests.yml:
--------------------------------------------------------------------------------
1 | name: tests
2 |
3 | on: [push, pull_request]
4 |
5 | jobs:
6 | build:
7 | runs-on: ubuntu-latest
8 | strategy:
9 | matrix:
10 | python-version: ['3.8']
11 | postgres-version: ['15.1.1.78']
12 |
13 | services:
14 | postgres:
15 | image: supabase/postgres:${{ matrix.postgres-version }}
16 | env:
17 | POSTGRES_DB: vecs_db
18 | POSTGRES_USER: postgres
19 | POSTGRES_PASSWORD: password
20 | ports:
21 | - 5611:5432
22 | options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
23 |
24 | steps:
25 | - uses: actions/checkout@v3
26 |
27 | - name: set up python ${{ matrix.python-version }}
28 | uses: actions/setup-python@v4
29 | with:
30 | python-version: ${{ matrix.python-version }}
31 |
32 | - name: install
33 | run: |
34 | pip install --upgrade pip
35 | pip install wheel
36 | pip install -e ".[dev, text_embedding]"
37 |
38 | - name: test with coverage
39 | run: |
40 | pip install coverage coveralls
41 | coverage run --source=vecs -m pytest
42 | coverage report
43 |
44 | - name: upload coverage to coveralls
45 | run: coveralls
46 | env:
47 | coveralls_repo_token: ${{ secrets.COVERALLS_REPO_TOKEN }}
48 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
49 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Temporary Python files
2 | *.pyc
3 | *.egg-info
4 | __pycache__
5 | .ipynb_checkpoints
6 | .DS_STORE
7 |
8 | # pyenv
9 | .python-version
10 |
11 | .benchmarks
12 |
13 | pip-wheel-metadata/
14 |
15 | .vscode
16 |
17 | # Temporary OS files
18 | Icon*
19 |
20 | # Pytest cache
21 | .pytest_cache/*
22 |
23 | # Virtual environment
24 | venv/*
25 |
26 | # Temporary virtual environment files
27 | /.cache/
28 | /.venv/
29 |
30 | # Temporary server files
31 | .env
32 | *.pid
33 | *.swp
34 |
35 | # Generated documentation
36 | /site/
37 | /*.html
38 | /*.rst
39 |
40 | # Google Drive
41 | *.gdoc
42 | *.gsheet
43 | *.gslides
44 | *.gdraw
45 |
46 | # Testing and coverage results
47 | /.pytest/
48 | /.coverage
49 | /.coverage.*
50 | /htmlcov/
51 | /xmlreport/
52 | /pyunit.xml
53 | /tmp/
54 | *.tmp
55 |
56 | # Build and release directories
57 | /build/
58 | /dist/
59 | *.spec
60 |
61 | # Sublime Text
62 | *.sublime-workspace
63 |
64 | # Eclipse
65 | .settings
66 |
--------------------------------------------------------------------------------
/.pre-commit-config.yaml:
--------------------------------------------------------------------------------
1 | repos:
2 | - repo: https://github.com/pre-commit/mirrors-isort
3 | rev: v5.6.4
4 | hooks:
5 | - id: isort
6 | args: ['--multi-line=3', '--trailing-comma', '--force-grid-wrap=0', '--use-parentheses', '--line-width=88']
7 |
8 |
9 | - repo: https://github.com/pre-commit/pre-commit-hooks
10 | rev: v3.3.0
11 | hooks:
12 | - id: trailing-whitespace
13 | - id: check-added-large-files
14 | - id: check-yaml
15 | - id: mixed-line-ending
16 | args: ['--fix=lf']
17 |
18 | - repo: https://github.com/humitos/mirrors-autoflake.git
19 | rev: v1.1
20 | hooks:
21 | - id: autoflake
22 | args: ['--in-place', '--remove-all-unused-imports']
23 |
24 | - repo: https://github.com/ambv/black
25 | rev: 22.10.0
26 | hooks:
27 | - id: black
28 | language_version: python3.9
29 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Apache License
2 | Version 2.0, January 2004
3 | http://www.apache.org/licenses/
4 |
5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6 |
7 | 1. Definitions.
8 |
9 | "License" shall mean the terms and conditions for use, reproduction,
10 | and distribution as defined by Sections 1 through 9 of this document.
11 |
12 | "Licensor" shall mean the copyright owner or entity authorized by
13 | the copyright owner that is granting the License.
14 |
15 | "Legal Entity" shall mean the union of the acting entity and all
16 | other entities that control, are controlled by, or are under common
17 | control with that entity. For the purposes of this definition,
18 | "control" means (i) the power, direct or indirect, to cause the
19 | direction or management of such entity, whether by contract or
20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the
21 | outstanding shares, or (iii) beneficial ownership of such entity.
22 |
23 | "You" (or "Your") shall mean an individual or Legal Entity
24 | exercising permissions granted by this License.
25 |
26 | "Source" form shall mean the preferred form for making modifications,
27 | including but not limited to software source code, documentation
28 | source, and configuration files.
29 |
30 | "Object" form shall mean any form resulting from mechanical
31 | transformation or translation of a Source form, including but
32 | not limited to compiled object code, generated documentation,
33 | and conversions to other media types.
34 |
35 | "Work" shall mean the work of authorship, whether in Source or
36 | Object form, made available under the License, as indicated by a
37 | copyright notice that is included in or attached to the work
38 | (an example is provided in the Appendix below).
39 |
40 | "Derivative Works" shall mean any work, whether in Source or Object
41 | form, that is based on (or derived from) the Work and for which the
42 | editorial revisions, annotations, elaborations, or other modifications
43 | represent, as a whole, an original work of authorship. For the purposes
44 | of this License, Derivative Works shall not include works that remain
45 | separable from, or merely link (or bind by name) to the interfaces of,
46 | the Work and Derivative Works thereof.
47 |
48 | "Contribution" shall mean any work of authorship, including
49 | the original version of the Work and any modifications or additions
50 | to that Work or Derivative Works thereof, that is intentionally
51 | submitted to Licensor for inclusion in the Work by the copyright owner
52 | or by an individual or Legal Entity authorized to submit on behalf of
53 | the copyright owner. For the purposes of this definition, "submitted"
54 | means any form of electronic, verbal, or written communication sent
55 | to the Licensor or its representatives, including but not limited to
56 | communication on electronic mailing lists, source code control systems,
57 | and issue tracking systems that are managed by, or on behalf of, the
58 | Licensor for the purpose of discussing and improving the Work, but
59 | excluding communication that is conspicuously marked or otherwise
60 | designated in writing by the copyright owner as "Not a Contribution."
61 |
62 | "Contributor" shall mean Licensor and any individual or Legal Entity
63 | on behalf of whom a Contribution has been received by Licensor and
64 | subsequently incorporated within the Work.
65 |
66 | 2. Grant of Copyright License. Subject to the terms and conditions of
67 | this License, each Contributor hereby grants to You a perpetual,
68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69 | copyright license to reproduce, prepare Derivative Works of,
70 | publicly display, publicly perform, sublicense, and distribute the
71 | Work and such Derivative Works in Source or Object form.
72 |
73 | 3. Grant of Patent License. Subject to the terms and conditions of
74 | this License, each Contributor hereby grants to You a perpetual,
75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76 | (except as stated in this section) patent license to make, have made,
77 | use, offer to sell, sell, import, and otherwise transfer the Work,
78 | where such license applies only to those patent claims licensable
79 | by such Contributor that are necessarily infringed by their
80 | Contribution(s) alone or by combination of their Contribution(s)
81 | with the Work to which such Contribution(s) was submitted. If You
82 | institute patent litigation against any entity (including a
83 | cross-claim or counterclaim in a lawsuit) alleging that the Work
84 | or a Contribution incorporated within the Work constitutes direct
85 | or contributory patent infringement, then any patent licenses
86 | granted to You under this License for that Work shall terminate
87 | as of the date such litigation is filed.
88 |
89 | 4. Redistribution. You may reproduce and distribute copies of the
90 | Work or Derivative Works thereof in any medium, with or without
91 | modifications, and in Source or Object form, provided that You
92 | meet the following conditions:
93 |
94 | (a) You must give any other recipients of the Work or
95 | Derivative Works a copy of this License; and
96 |
97 | (b) You must cause any modified files to carry prominent notices
98 | stating that You changed the files; and
99 |
100 | (c) You must retain, in the Source form of any Derivative Works
101 | that You distribute, all copyright, patent, trademark, and
102 | attribution notices from the Source form of the Work,
103 | excluding those notices that do not pertain to any part of
104 | the Derivative Works; and
105 |
106 | (d) If the Work includes a "NOTICE" text file as part of its
107 | distribution, then any Derivative Works that You distribute must
108 | include a readable copy of the attribution notices contained
109 | within such NOTICE file, excluding those notices that do not
110 | pertain to any part of the Derivative Works, in at least one
111 | of the following places: within a NOTICE text file distributed
112 | as part of the Derivative Works; within the Source form or
113 | documentation, if provided along with the Derivative Works; or,
114 | within a display generated by the Derivative Works, if and
115 | wherever such third-party notices normally appear. The contents
116 | of the NOTICE file are for informational purposes only and
117 | do not modify the License. You may add Your own attribution
118 | notices within Derivative Works that You distribute, alongside
119 | or as an addendum to the NOTICE text from the Work, provided
120 | that such additional attribution notices cannot be construed
121 | as modifying the License.
122 |
123 | You may add Your own copyright statement to Your modifications and
124 | may provide additional or different license terms and conditions
125 | for use, reproduction, or distribution of Your modifications, or
126 | for any such Derivative Works as a whole, provided Your use,
127 | reproduction, and distribution of the Work otherwise complies with
128 | the conditions stated in this License.
129 |
130 | 5. Submission of Contributions. Unless You explicitly state otherwise,
131 | any Contribution intentionally submitted for inclusion in the Work
132 | by You to the Licensor shall be under the terms and conditions of
133 | this License, without any additional terms or conditions.
134 | Notwithstanding the above, nothing herein shall supersede or modify
135 | the terms of any separate license agreement you may have executed
136 | with Licensor regarding such Contributions.
137 |
138 | 6. Trademarks. This License does not grant permission to use the trade
139 | names, trademarks, service marks, or product names of the Licensor,
140 | except as required for reasonable and customary use in describing the
141 | origin of the Work and reproducing the content of the NOTICE file.
142 |
143 | 7. Disclaimer of Warranty. Unless required by applicable law or
144 | agreed to in writing, Licensor provides the Work (and each
145 | Contributor provides its Contributions) on an "AS IS" BASIS,
146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 | implied, including, without limitation, any warranties or conditions
148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 | PARTICULAR PURPOSE. You are solely responsible for determining the
150 | appropriateness of using or redistributing the Work and assume any
151 | risks associated with Your exercise of permissions under this License.
152 |
153 | 8. Limitation of Liability. In no event and under no legal theory,
154 | whether in tort (including negligence), contract, or otherwise,
155 | unless required by applicable law (such as deliberate and grossly
156 | negligent acts) or agreed to in writing, shall any Contributor be
157 | liable to You for damages, including any direct, indirect, special,
158 | incidental, or consequential damages of any character arising as a
159 | result of this License or out of the use or inability to use the
160 | Work (including but not limited to damages for loss of goodwill,
161 | work stoppage, computer failure or malfunction, or any and all
162 | other commercial damages or losses), even if such Contributor
163 | has been advised of the possibility of such damages.
164 |
165 | 9. Accepting Warranty or Additional Liability. While redistributing
166 | the Work or Derivative Works thereof, You may choose to offer,
167 | and charge a fee for, acceptance of support, warranty, indemnity,
168 | or other liability obligations and/or rights consistent with this
169 | License. However, in accepting such obligations, You may act only
170 | on Your own behalf and on Your sole responsibility, not on behalf
171 | of any other Contributor, and only if You agree to indemnify,
172 | defend, and hold each Contributor harmless for any liability
173 | incurred by, or claims asserted against, such Contributor by reason
174 | of your accepting any such warranty or additional liability.
175 |
176 | END OF TERMS AND CONDITIONS
177 |
178 | APPENDIX: How to apply the Apache License to your work.
179 |
180 | To apply the Apache License to your work, attach the following
181 | boilerplate notice, with the fields enclosed by brackets "[]"
182 | replaced with your own identifying information. (Don't include
183 | the brackets!) The text should be enclosed in the appropriate
184 | comment syntax for the file format. We also recommend that a
185 | file or class name and description of purpose be included on the
186 | same "printed page" as the copyright notice for easier
187 | identification within third-party archives.
188 |
189 | Copyright [yyyy] [name of copyright owner]
190 |
191 | Licensed under the Apache License, Version 2.0 (the "License");
192 | you may not use this file except in compliance with the License.
193 | You may obtain a copy of the License at
194 |
195 | http://www.apache.org/licenses/LICENSE-2.0
196 |
197 | Unless required by applicable law or agreed to in writing, software
198 | distributed under the License is distributed on an "AS IS" BASIS,
199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 | See the License for the specific language governing permissions and
201 | limitations under the License.
202 |
203 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # vecs
2 |
3 |
4 | 
5 |
6 | 
7 |
8 |
9 | 
10 |
11 |
12 |
13 |
14 | 
15 | 
16 | 
17 |
18 |
19 | ---
20 |
21 | **Documentation**: https://supabase.github.io/vecs/
22 |
23 | **Source Code**: https://github.com/supabase/vecs
24 |
25 | ---
26 |
27 | `vecs` is a python client for managing and querying vector stores in PostgreSQL with the [pgvector extension](https://github.com/pgvector/pgvector). This guide will help you get started with using vecs.
28 |
29 | If you don't have a Postgres database with the pgvector ready, see [hosting](https://supabase.github.io/vecs/hosting/) for easy options.
30 |
31 | ## Installation
32 |
33 | Requires:
34 |
35 | - Python 3.7+
36 |
37 | You can install vecs using pip:
38 |
39 | ```sh
40 | pip install vecs
41 | ```
42 |
43 | ## Usage
44 |
45 | Visit the [quickstart guide](https://supabase.github.io/vecs/api) for more complete info.
46 |
47 | ```python
48 | import vecs
49 |
50 | DB_CONNECTION = "postgresql://:@:/"
51 |
52 | # create vector store client
53 | vx = vecs.create_client(DB_CONNECTION)
54 |
55 | # create a collection of vectors with 3 dimensions
56 | docs = vx.get_or_create_collection(name="docs", dimension=3)
57 |
58 | # add records to the *docs* collection
59 | docs.upsert(
60 | records=[
61 | (
62 | "vec0", # the vector's identifier
63 | [0.1, 0.2, 0.3], # the vector. list or np.array
64 | {"year": 1973} # associated metadata
65 | ),
66 | (
67 | "vec1",
68 | [0.7, 0.8, 0.9],
69 | {"year": 2012}
70 | )
71 | ]
72 | )
73 |
74 | # index the collection for fast search performance
75 | docs.create_index()
76 |
77 | # query the collection filtering metadata for "year" = 2012
78 | docs.query(
79 | data=[0.4,0.5,0.6], # required
80 | limit=1, # number of records to return
81 | filters={"year": {"$eq": 2012}}, # metadata filters
82 | )
83 |
84 | # Returns: ["vec1"]
85 | ```
86 |
--------------------------------------------------------------------------------
/docs/api.md:
--------------------------------------------------------------------------------
1 | # API
2 |
3 | `vecs` is a python client for managing and querying vector stores in PostgreSQL with the [pgvector extension](https://github.com/pgvector/pgvector). This guide will help you get started with using vecs.
4 |
5 | If you don't have a Postgres database with the pgvector ready, see [hosting](hosting.md) for easy options.
6 |
7 | ## Installation
8 |
9 | Requires:
10 |
11 | - Python 3.7+
12 |
13 | You can install vecs using pip:
14 |
15 | ```bash
16 | pip install vecs
17 | ```
18 |
19 | ## Usage
20 |
21 | ## Connecting
22 |
23 | Before you can interact with vecs, create the client to communicate with Postgres. If you haven't started a Postgres instance yet, see [hosting](hosting.md).
24 | ```python
25 | import vecs
26 |
27 | DB_CONNECTION = "postgresql://:@:/"
28 |
29 | # create vector store client
30 | vx = vecs.create_client(DB_CONNECTION)
31 | ```
32 |
33 | ## Get or Create a Collection
34 |
35 | You can get a collection (or create if it doesn't exist), specifying the collection's name and the number of dimensions for the vectors you intend to store.
36 |
37 | ```python
38 | docs = vx.get_or_create_collection(name="docs", dimension=3)
39 | ```
40 |
41 | ## Upserting vectors
42 |
43 | `vecs` combines the concepts of "insert" and "update" into "upsert". Upserting records adds them to the collection if the `id` is not present, or updates the existing record if the `id` does exist.
44 |
45 | ```python
46 | # add records to the collection
47 | docs.upsert(
48 | records=[
49 | (
50 | "vec0", # the vector's identifier
51 | [0.1, 0.2, 0.3], # the vector. list or np.array
52 | {"year": 1973} # associated metadata
53 | ),
54 | (
55 | "vec1",
56 | [0.7, 0.8, 0.9],
57 | {"year": 2012}
58 | )
59 | ]
60 | )
61 | ```
62 |
63 | ## Deleting vectors
64 |
65 | Deleting records removes them from the collection. To delete records, specify a list of `ids` or metadata filters to the `delete` method. The ids of the sucessfully deleted records are returned from the method. Note that attempting to delete non-existent records does not raise an error.
66 |
67 | ```python
68 | docs.delete(ids=["vec0", "vec1"])
69 | # or delete by a metadata filter
70 | docs.delete(filters={"year": {"$eq": 2012}})
71 | ```
72 |
73 | ## Create an index
74 |
75 | Collections can be queried immediately after being created.
76 | However, for good throughput, the collection should be indexed after records have been upserted.
77 |
78 | Only one index may exist per-collection. By default, creating an index will replace any existing index.
79 |
80 | To create an index:
81 |
82 | ```python
83 | docs.create_index()
84 | ```
85 |
86 | You may optionally provide a distance measure and index method.
87 |
88 | Available options for distance `measure` are:
89 |
90 | - `vecs.IndexMeasure.cosine_distance`
91 | - `vecs.IndexMeasure.l2_distance`
92 | - `vecs.IndexMeasure.l1_distance`
93 | - `vecs.IndexMeasure.max_inner_product`
94 |
95 | which correspond to different methods for comparing query vectors to the vectors in the database.
96 |
97 | If you aren't sure which to use, the default of cosine_distance is the most widely compatible with off-the-shelf embedding methods.
98 |
99 | Available options for index `method` are:
100 |
101 | - `vecs.IndexMethod.auto`
102 | - `vecs.IndexMethod.hnsw`
103 | - `vecs.IndexMethod.ivfflat`
104 |
105 | Where `auto` selects the best available index method, `hnsw` uses the [HNSW](https://github.com/pgvector/pgvector#hnsw) method and `ivfflat` uses [IVFFlat](https://github.com/pgvector/pgvector#ivfflat).
106 |
107 | HNSW and IVFFlat indexes both allow for parameterization to control the speed/accuracy tradeoff. vecs provides sane defaults for these parameters. For a greater level of control you can optionally pass an instance of `vecs.IndexArgsIVFFlat` or `vecs.IndexArgsHNSW` to `create_index`'s `index_arguments` argument. Descriptions of the impact for each parameter are available in the [pgvector docs](https://github.com/pgvector/pgvector).
108 |
109 | When using IVFFlat indexes, the index must be created __after__ the collection has been populated with records. Building an IVFFlat index on an empty collection will result in significantly reduced recall. You can continue upserting new documents after the index has been created, but should rebuild the index if the size of the collection more than doubles since the last index operation.
110 |
111 | HNSW indexes can be created immediately after the collection without populating records.
112 |
113 | To manually specify `method`, `measure`, and `index_arguments` add them as arguments to `create_index` for example:
114 |
115 | ```python
116 | docs.create_index(
117 | method=IndexMethod.hnsw,
118 | measure=IndexMeasure.cosine_distance,
119 | index_arguments=IndexArgsHNSW(m=8),
120 | )
121 | ```
122 |
123 | !!! note
124 | The time required to create an index grows with the number of records and size of vectors.
125 | For a few thousand records expect sub-minute a response in under a minute. It may take a few
126 | minutes for larger collections.
127 |
128 | ## Query
129 |
130 | Given a collection `docs` with several records:
131 |
132 | ### Basic
133 |
134 | The simplest form of search is to provide a query vector.
135 |
136 | !!! note
137 | Indexes are essential for good performance. See [creating an index](#create-an-index) for more info.
138 |
139 | If you do not create an index, every query will return a warning
140 | ```text
141 | query does not have a covering index for cosine_similarity. See Collection.create_index
142 | ```
143 | that incldues the `IndexMeasure` you should index.
144 |
145 |
146 |
147 | ```python
148 | docs.query(
149 | data=[0.4,0.5,0.6], # required
150 | limit=5, # number of records to return
151 | filters={}, # metadata filters
152 | measure="cosine_distance", # distance measure to use
153 | include_value=False, # should distance measure values be returned?
154 | include_metadata=False, # should record metadata be returned?
155 | include_vector=False, # should vectors be returned?
156 | )
157 | ```
158 |
159 | Which returns a list of vector record `ids`.
160 |
161 |
162 | ### Metadata Filtering
163 |
164 | The metadata that is associated with each record can also be filtered during a query.
165 |
166 | As an example, `{"year": {"$eq": 2005}}` filters a `year` metadata key to be equal to 2005
167 |
168 | In context:
169 |
170 | ```python
171 | docs.query(
172 | data=[0.4,0.5,0.6],
173 | filters={"year": {"$eq": 2012}}, # metadata filters
174 | )
175 | ```
176 |
177 | For a complete reference, see the [metadata guide](concepts_metadata.md).
178 |
179 |
180 | ### Disconnect
181 |
182 | When you're done with a collection, be sure to disconnect the client from the database.
183 |
184 | ```python
185 | vx.disconnect()
186 | ```
187 |
188 | alternatively, use the client as a context manager and it will automatically close the connection on exit.
189 |
190 |
191 | ```python
192 | import vecs
193 |
194 | DB_CONNECTION = "postgresql://:@:/"
195 |
196 | # create vector store client
197 | with vecs.create_client(DB_CONNECTION) as vx:
198 | # do some work here
199 | pass
200 |
201 | # connections are now closed
202 | ```
203 |
204 |
205 | ## Adapters
206 |
207 | Adapters are an optional feature to transform data before adding to or querying from a collection. Adapters make it possible to interact with a collection using only your project's native data type (eg. just raw text), rather than manually handling vectors.
208 |
209 | For a complete list of available adapters, see [built-in adapters](concepts_adapters.md#built-in-adapters).
210 |
211 | As an example, we'll create a collection with an adapter that chunks text into paragraphs and converts each chunk into an embedding vector using the `all-MiniLM-L6-v2` model.
212 |
213 | First, install `vecs` with optional dependencies for text embeddings:
214 | ```sh
215 | pip install "vecs[text_embedding]"
216 | ```
217 |
218 | Then create a collection with an adapter to chunk text into paragraphs and embed each paragraph using the `all-MiniLM-L6-v2` 384 dimensional text embedding model.
219 |
220 | ```python
221 | import vecs
222 | from vecs.adapter import Adapter, ParagraphChunker, TextEmbedding
223 |
224 | # create vector store client
225 | vx = vecs.Client("postgresql://:@:/")
226 |
227 | # create a collection with an adapter
228 | docs = vx.get_or_create_collection(
229 | name="docs",
230 | adapter=Adapter(
231 | [
232 | ParagraphChunker(skip_during_query=True),
233 | TextEmbedding(model='all-MiniLM-L6-v2'),
234 | ]
235 | )
236 | )
237 |
238 | ```
239 |
240 | With the adapter registered against the collection, we can upsert records into the collection passing in text rather than vectors.
241 |
242 | ```python
243 | # add records to the collection using text as the media type
244 | docs.upsert(
245 | records=[
246 | (
247 | "vec0",
248 | "four score and ....", # <- note that we can now pass text here
249 | {"year": 1973}
250 | ),
251 | (
252 | "vec1",
253 | "hello, world!",
254 | {"year": "2012"}
255 | )
256 | ]
257 | )
258 | ```
259 |
260 | Similarly, we can query the collection using text.
261 | ```python
262 |
263 | # search by text
264 | docs.query(data="foo bar")
265 | ```
266 |
267 |
268 |
269 | ---------
270 | ## Deprecated
271 |
272 | ### Create collection
273 |
274 | !!! note
275 | Deprecated: use [get_or_create_collection](#get-or-create-a-collection)
276 |
277 | You can create a collection to store vectors specifying the collections name and the number of dimensions in the vectors you intend to store.
278 |
279 | ```python
280 | docs = vx.create_collection(name="docs", dimension=3)
281 | ```
282 |
283 |
284 | ### Get an existing collection
285 |
286 | !!! note
287 | Deprecated: use [get_or_create_collection](#get-or-create-a-collection)
288 |
289 | To access a previously created collection, use `get_collection` to retrieve it by name
290 |
291 | ```python
292 | docs = vx.get_collection(name="docs")
293 | ```
294 |
295 |
296 |
--------------------------------------------------------------------------------
/docs/assets/hf_configure_endpoint.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/hf_configure_endpoint.png
--------------------------------------------------------------------------------
/docs/assets/hf_configure_endpoint_2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/hf_configure_endpoint_2.png
--------------------------------------------------------------------------------
/docs/assets/hf_endpoint_url.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/hf_endpoint_url.png
--------------------------------------------------------------------------------
/docs/assets/hf_new_endpoint.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/hf_new_endpoint.png
--------------------------------------------------------------------------------
/docs/assets/supabase_connection_info.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/supabase_connection_info.png
--------------------------------------------------------------------------------
/docs/assets/supabase_new_project.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/supabase_new_project.png
--------------------------------------------------------------------------------
/docs/assets/supabase_new_project_prompt.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/supabase_new_project_prompt.png
--------------------------------------------------------------------------------
/docs/assets/supabase_sign_up.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/supabase/vecs/HEAD/docs/assets/supabase_sign_up.png
--------------------------------------------------------------------------------
/docs/concepts_adapters.md:
--------------------------------------------------------------------------------
1 | # Adapters
2 |
3 | Adapters are an optional feature to transform data before adding to or querying from a collection. Adapters provide a customizable and modular way to express data transformations and make interacting with collections more ergonomic.
4 |
5 | Additionally, adapter transformations are applied lazily and can internally batch operations which can make them more memory and CPU efficient compared to manually executing transforms.
6 |
7 | ## Example:
8 | As an example, we'll create a collection with an adapter that chunks text into paragraphs and converts each chunk into an embedding vector using the `all-MiniLM-L6-v2` model.
9 |
10 | First, install `vecs` with optional dependencies for text embeddings:
11 | ```sh
12 | pip install "vecs[text_embedding]"
13 | ```
14 |
15 | Then create a collection with an adapter to chunk text into paragraphs and embed each paragraph using the `all-MiniLM-L6-v2` 384 dimensional text embedding model.
16 |
17 | ```python
18 | import vecs
19 | from vecs.adapter import Adapter, ParagraphChunker, TextEmbedding
20 |
21 | # create vector store client
22 | vx = vecs.Client("postgresql://:@:/")
23 |
24 | # create a collection with an adapter
25 | docs = vx.get_or_create_collection(
26 | name="docs",
27 | adapter=Adapter(
28 | [
29 | ParagraphChunker(skip_during_query=True),
30 | TextEmbedding(model='all-MiniLM-L6-v2'),
31 | ]
32 | )
33 | )
34 |
35 | ```
36 |
37 | With the adapter registered against the collection, we can upsert records into the collection passing in text rather than vectors.
38 |
39 | ```python
40 | # add records to the collection using text as the media type
41 | docs.upsert(
42 | records=[
43 | (
44 | "vec0",
45 | "four score and ....", # <- note that we can now pass text here
46 | {"year": 1973}
47 | ),
48 | (
49 | "vec1",
50 | "hello, world!",
51 | {"year": "2012"}
52 | )
53 | ]
54 | )
55 | ```
56 |
57 | Similarly, we can query the collection using text.
58 | ```python
59 |
60 | # search by text
61 | docs.query(data="foo bar")
62 | ```
63 |
64 | In summary, `Adapter`s allow you to work with a collection as though they store your prefered data type natively.
65 |
66 |
67 | ## Built-in Adapters
68 |
69 | vecs provides several built-in Adapters.
70 |
71 | - [ParagraphChunker](#paragraphchunker)
72 | - [TextEmbedding](#textembedding)
73 |
74 | Have an idea for a useful adapter? [Open an issue](https://github.com/supabase/vecs/issues/new/choose) requesting it.
75 |
76 | ### ParagraphChunker
77 |
78 | The `ParagraphChunker` `AdapterStep` splits text media into paragraphs and yields each paragraph as a separate record. That can be a useful preprocessing step when upserting large documents that contain multiple paragraphs. The `ParagraphChunker` delimits paragraphs by two consecutive line breaks `\n\n`.
79 |
80 | `ParagrphChunker` is a pre-preocessing step and must be used in combination with another adapter step like `TextEmbedding` to transform the chunked text into a vector.
81 |
82 |
83 | ```python
84 | from vecs.adapter import Adapter, ParagraphChunker
85 |
86 | ...
87 |
88 | vx.get_or_create_collection(
89 | name="docs",
90 | adapter=Adapter(
91 | [
92 | ParagraphChunker(skip_during_query=True),
93 | ...
94 | ]
95 | )
96 | )
97 | ```
98 |
99 | When querying the collection, you probably do not want to chunk the text. To skip text chunking during queries, set the `skip_during_query` argument to `True`. Setting `skip_during_query` to `False` will raise an exception if the input text contains more than one paragraph.
100 |
101 |
102 | ### TextEmbedding
103 |
104 | The `TextEmbedding` `AdapterStep` accepts text and converts it into a vector that can be consumed by the `Collection`. `TextEmbedding` supports all models available in the [`sentence_transformers`](https://www.sbert.net) package. A complete list of supported models is available in `vecs.adapter.TextEmbeddingModel`.
105 |
106 | ```python
107 | from vecs.adapter import Adapter, TextEmbedding
108 | ...
109 |
110 | vx.get_or_create_collection(
111 | name="docs",
112 | adapter=Adapter(
113 | [
114 | TextEmbedding(model='all-MiniLM-L6-v2')
115 | ]
116 | )
117 | )
118 |
119 | # search by text
120 | docs.query(data="foo bar")
121 | ```
122 |
123 | ## Interface
124 |
125 | Adapters are objects that take in data in the form of `Iterable[Tuple[str, Any, Optional[Dict]]]` where `Tuple[str, Any, Optional[Dict]]]` represents records of `(id, media, metadata)`.
126 |
127 | The main use of Adapters is to transform the media part of the records into a form that is ready to be ingested into the collection (like converting text into embeddings). However, Adapters can also modify the `id` or `metadata` if required.
128 |
129 | Due to the common interface, adapters may be comprised of multiple adapter steps to create multi-stage preprocessing pipelines. For example, a multi-step adapter might first convert text into chunks and then convert each text chunk into an embedding vector.
130 |
131 |
132 |
133 |
--------------------------------------------------------------------------------
/docs/concepts_collections.md:
--------------------------------------------------------------------------------
1 | # Collections
2 |
3 | A collection is an group of vector records.
4 | Records can be [added to or updated in](api.md/#upserting-vectors) a collection.
5 | Collections can be [queried](api.md/#query) at any time, but should be [indexed](api.md/#create-an-index) for scalable query performance.
6 |
7 | Each vector record has the form:
8 |
9 | ```python
10 | Record (
11 | id: String
12 | vec: Numeric[]
13 | metadata: JSON
14 | )
15 | ```
16 |
17 | For example:
18 |
19 | ```python
20 | ("vec1", [0.1, 0.2, 0.3], {"year": 1990})
21 | ```
22 |
23 | Underneath every `vecs` collection is a Postgres table
24 |
25 | ```sql
26 | create table (
27 | id string primary key,
28 | vec vector(),
29 | metadata jsonb
30 | )
31 | ```
32 |
33 | where rows in the table map 1:1 with vecs vector records.
34 |
35 | It is safe to select collection tables from outside the vecs client but issuing DDL is not recommended.
36 |
--------------------------------------------------------------------------------
/docs/concepts_indexes.md:
--------------------------------------------------------------------------------
1 | # Indexes
2 |
3 | Indexes are tools for optimizing query performance of a [collection](concepts_collections.md).
4 |
5 | Collections can be [queried](api.md/#query) without an index, but that will emit a python warning and should never be done in production.
6 |
7 | ```text
8 | query does not have a covering index for cosine_similarity. See Collection.create_index
9 | ```
10 |
11 | As each query vector must be checked against every record in the collection. When the number of dimensions and/or number of records becomes large, that becomes extremely slow and computationally expensive.
12 |
13 | An index is a heuristic data structure that pre-computes distances between key points in the vector space. It is smaller and can be traversed more quickly than the whole collection enabling much more performant searching.
14 |
15 | Only one index may exist per-collection. An index optimizes a collection for searching according to a selected distance measure.
16 |
17 | To create an index:
18 |
19 | ```python
20 | docs.create_index()
21 | ```
22 |
23 | You may optionally provide a distance measure and index method.
24 |
25 | Available options for distance `measure` are:
26 |
27 | - `vecs.IndexMeasure.cosine_distance`
28 | - `vecs.IndexMeasure.l2_distance`
29 | - `vecs.IndexMeasure.l1_distance`
30 | - `vecs.IndexMeasure.max_inner_product`
31 |
32 | which correspond to different methods for comparing query vectors to the vectors in the database.
33 |
34 | If you aren't sure which to use, the default of cosine_distance is the most widely compatible with off-the-shelf embedding methods.
35 |
36 | Available options for index `method` are:
37 |
38 | - `vecs.IndexMethod.auto`
39 | - `vecs.IndexMethod.hnsw`
40 | - `vecs.IndexMethod.ivfflat`
41 |
42 | Where `auto` selects the best available index method, `hnsw` uses the [HNSW](https://github.com/pgvector/pgvector#hnsw) method and `ivfflat` uses [IVFFlat](https://github.com/pgvector/pgvector#ivfflat).
43 |
44 | HNSW and IVFFlat indexes both allow for parameterization to control the speed/accuracy tradeoff. vecs provides sane defaults for these parameters. For a greater level of control you can optionally pass an instance of `vecs.IndexArgsIVFFlat` or `vecs.IndexArgsHNSW` to `create_index`'s `index_arguments` argument. Descriptions of the impact for each parameter are available in the [pgvector docs](https://github.com/pgvector/pgvector).
45 |
46 | When using IVFFlat indexes, the index must be created __after__ the collection has been populated with records. Building an IVFFlat index on an empty collection will result in significantly reduced recall. You can continue upserting new documents after the index has been created, but should rebuild the index if the size of the collection more than doubles since the last index operation.
47 |
48 | HNSW indexes can be created immediately after the collection without populating records.
49 |
50 | To manually specify `method`, `measure`, and `index_arguments` add them as arguments to `create_index` for example:
51 |
52 | ```python
53 | docs.create_index(
54 | method=IndexMethod.hnsw,
55 | measure=IndexMeasure.cosine_distance,
56 | index_arguments=IndexArgsHNSW(m=8),
57 | )
58 | ```
59 |
60 | !!! note
61 | The time required to create an index grows with the number of records and size of vectors.
62 | For a few thousand records expect sub-minute a response in under a minute. It may take a few
63 | minutes for larger collections.
64 |
--------------------------------------------------------------------------------
/docs/concepts_metadata.md:
--------------------------------------------------------------------------------
1 | # Metadata
2 |
3 | vecs allows you to associate key-value pairs of metadata with indexes and ids in your collections.
4 | You can then add filters to queries that reference the metadata metadata.
5 |
6 | ## Types
7 | Metadata is stored as binary JSON. As a result, allowed metadata types are drawn from JSON primitive types.
8 |
9 | - Boolean
10 | - String
11 | - Number
12 |
13 | The technical limit of a metadata field associated with a vector is 1GB.
14 | In practice you should keep metadata fields as small as possible to maximize performance.
15 |
16 | ## Metadata Query Language
17 |
18 | The metadata query language is based loosely on [mongodb's selectors](https://www.mongodb.com/docs/manual/reference/operator/query/).
19 |
20 | `vecs` currently supports a subset of those operators.
21 |
22 |
23 | ### Comparison Operators
24 |
25 | Comparison operators compare a provided value with a value stored in metadata field of the vector store.
26 |
27 | | Operator | Description |
28 | | --------- | ----------- |
29 | | $eq | Matches values that are equal to a specified value |
30 | | $ne | Matches values that are not equal to a specified value |
31 | | $gt | Matches values that are greater than a specified value |
32 | | $gte | Matches values that are greater than or equal to a specified value |
33 | | $lt | Matches values that are less than a specified value |
34 | | $lte | Matches values that are less than or equal to a specified value |
35 | | $in | Matches values that are contained by scalar list of specified values |
36 | | $contains | Matches values where a scalar is contained within an array metadata field |
37 |
38 |
39 | ### Logical Operators
40 |
41 | Logical operators compose other operators, and can be nested.
42 |
43 | | Operator | Description |
44 | | --------- | ----------- |
45 | | $and | Joins query clauses with a logical AND returns all documents that match the conditions of both clauses. |
46 | | $or | Joins query clauses with a logical OR returns all documents that match the conditions of either clause. |
47 |
48 |
49 | ### Performance
50 |
51 | For best performance, use scalar key-value pairs for metadata and prefer `$eq`, `$and` and `$or` filters where possible.
52 | Those variants are most consistently able to make use of indexes.
53 |
54 |
55 |
56 | ### Examples
57 |
58 | ---
59 |
60 | `year` equals 2020
61 |
62 | ```json
63 | {"year": {"$eq": 2020}}
64 | ```
65 |
66 | ---
67 |
68 | `year` equals 2020 or `gross` greater than or equal to 5000.0
69 |
70 | ```json
71 | {
72 | "$or": [
73 | {"year": {"$eq": 2020}},
74 | {"gross": {"$gte": 5000.0}}
75 | ]
76 | }
77 | ```
78 |
79 | ---
80 |
81 | `last_name` is less than "Brown" and `is_priority_customer` is true
82 |
83 | ```json
84 | {
85 | "$and": [
86 | {"last_name": {"$lt": "Brown"}},
87 | {"is_priority_customer": {"$gte": 5000.00}}
88 | ]
89 | }
90 | ```
91 |
92 | ---
93 |
94 | `priority` contained by ["enterprise", "pro"]
95 |
96 | ```json
97 | {
98 | "priority": {"$in": ["enterprise", "pro"]}
99 | }
100 | ```
101 |
102 | `tags`, an array, contains the string "important"
103 |
104 | ```json
105 | {
106 | "tags": {"$contains": "important"}
107 | }
108 | ```
109 |
--------------------------------------------------------------------------------
/docs/hosting.md:
--------------------------------------------------------------------------------
1 | # Deployment
2 |
3 |
4 | `vecs` is comatible with any Postgres 13+ with the [pgvector](https://github.com/pgvector/pgvector) extension installed.
5 |
6 | In the following we show we show instructions for hosting a database on Supabase and locally in docker since both are fast and free.
7 |
8 |
9 | ## Supabase
10 |
11 | ### Cloud Hosted
12 |
13 | #### Create an account
14 |
15 | Create a supabase account at [https://app.supabase.com/sign-up](https://app.supabase.com/sign-up).
16 |
17 | 
18 |
19 | #### Create a new project
20 |
21 | Select `New Project`
22 |
23 | 
24 |
25 | Complete the prompts. Be sure to remember or write down your password as we'll need that when connecting with vecs.
26 |
27 | 
28 |
29 | #### Connection Info
30 |
31 | On the project page, navigate to `Settings` > `Database` > `Database Settings`
32 |
33 | 
34 |
35 | and substitue those fields into the conenction string
36 |
37 | ```text
38 | postgresql://:@:/
39 | ```
40 | i.e.
41 | ```text
42 | postgres://postgres:[YOUR PASSWORD]@db.cvykdyhlwwwojivopztl.supabase.co:5432/postgres
43 | ```
44 |
45 | Keep that connection string secret and safe. Its your `DB_CONNECTION` in the [quickstart guide](api.md),
46 |
47 |
48 | ### Local
49 |
50 | You can also use Supabase locally on your machine. Doing so will keep your project setup consistent when deploying to hosted Supabase.
51 |
52 | ### Install the CLI
53 |
54 | To install the CLI, use the relevant system instructions below
55 |
56 | === "macOS"
57 |
58 | ```sh
59 | brew install supabase/tap/supabase
60 | ```
61 |
62 | === "Windows"
63 |
64 | ```sh
65 | scoop bucket add supabase https://github.com/supabase/scoop-bucket.git
66 | scoop install supabase
67 | ```
68 |
69 | === "Linux"
70 |
71 | Linux packages are provided in Releases. To install, download the .apk/.deb/.rpm file depending on your package manager and run one of the following:
72 |
73 | ```sh
74 | sudo apk add --allow-untrusted <...>.apk
75 | ```
76 | or
77 | ```sh
78 | sudo dpkg -i <...>.deb
79 | ```
80 | or
81 | ```sh
82 | sudo rpm -i <...>.rpm
83 | ```
84 |
85 | === "npm"
86 |
87 | ```shell
88 | npm install supabase --save-dev
89 | ```
90 |
91 | ### Start the Project
92 |
93 | From your project directory, create the `supabase/` sub-directory required for supabase projects by running:
94 |
95 | ```sh
96 | supabase init
97 | ```
98 |
99 | next start the application using:
100 |
101 | ```sh
102 | supabase start
103 | ```
104 |
105 | which will download the latest Supabase containers and provide a URL to each service:
106 |
107 | ```text
108 | Seeding data supabase/seed.sql...me...
109 | Started supabase local development setup.
110 |
111 | API URL: http://localhost:54321
112 | GraphQL URL: http://localhost:54321/graphql/v1
113 | DB URL: postgresql://postgres:postgres@localhost:54322/postgres
114 | Studio URL: http://localhost:54323
115 | Inbucket URL: http://localhost:54324
116 | JWT secret: super-secret-jwt-token-with-at-least-32-characters-long
117 | anon key: eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFz
118 | service_role key: eyJhbGciOiJIUzI1NiIsInR5cClJzdJsyH-qQwv8Hdp7fsn3W0YpN81IU
119 | ```
120 |
121 | The service we need for `vecs` is `DB URL`. Note it down for use as our `DB_CONNECTION`
122 |
123 | ```text
124 | postgresql://:@:/
125 | ```
126 |
127 | For more info on running a local Supabase project, checkout the [Supabase CLI guide](https://supabase.com/docs/guides/cli)
128 |
129 | ## Docker
130 |
131 | Install docker if you don't have it already at [Get Docker](https://docs.docker.com/get-docker/)
132 |
133 |
134 | #### Start the Postgres Container
135 |
136 | Next, run
137 | ```sh
138 | docker run --rm -d \
139 | --name vecs_hosting_guide \
140 | -p 5019:5432 \
141 | -e POSTGRES_DB=vecs_db \
142 | -e POSTGRES_PASSWORD=password \
143 | -e POSTGRES_USER=postgres \
144 | supabase/postgres:15.1.0.74
145 | ```
146 |
147 | #### Connection Info
148 | Substitue the values from the previous section into the postgres conenction string
149 |
150 | ```text
151 | postgresql://:@:/
152 | ```
153 | i.e.
154 | ```text
155 | postgresql://postgres:password@localhost:5019/vecs_db
156 | ```
157 |
158 | Keep that connection string secret and safe. Its your `DB_CONNECTION` in the [quickstart guide](api.md)
159 |
--------------------------------------------------------------------------------
/docs/index.md:
--------------------------------------------------------------------------------
1 | # vecs
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 | ---
20 |
21 | **Documentation**: https://supabase.github.io/vecs/
22 |
23 | **Source Code**: https://github.com/supabase/vecs
24 |
25 | ---
26 |
27 |
28 | Vecs is a Python client library for managing and querying vector stores in PostgreSQL, leveraging the capabilities of the [pgvector extension](https://github.com/pgvector/pgvector).
29 |
30 | ## Overview
31 |
32 | - Vector Management: create collections to persist and update vectors in a PostgreSQL database.
33 | - Querying: Query vectors efficiently using measures such as cosine distance, l2 distance, l1 distance, or max inner product.
34 | - Metadata: Each vector can have associated metadata, which can also be used as filters during queries.
35 | - Hybrid Data: vecs creates its own schema and can coexist with your existing relational data
36 |
37 |
38 | Visit the [quickstart guide](api.md) for how to get started.
39 |
40 | ## TL;DR
41 |
42 | ### Install
43 |
44 | ```sh
45 | pip install vecs
46 | ```
47 |
48 | ### Usage
49 |
50 |
51 | ```python
52 | import vecs
53 |
54 | DB_CONNECTION = "postgresql://:@:/"
55 |
56 | # create vector store client
57 | vx = vecs.create_client(DB_CONNECTION)
58 |
59 | # create a collection of vectors with 3 dimensions
60 | docs = vx.get_or_create_collection(name="docs", dimension=3)
61 |
62 | # add records to the *docs* collection
63 | docs.upsert(
64 | records=[
65 | (
66 | "vec0", # the vector's identifier
67 | [0.1, 0.2, 0.3], # the vector. list or np.array
68 | {"year": 1973} # associated metadata
69 | ),
70 | (
71 | "vec1",
72 | [0.7, 0.8, 0.9],
73 | {"year": 2012}
74 | )
75 | ]
76 | )
77 |
78 | # index the collection for fast search performance
79 | docs.create_index()
80 |
81 | # query the collection filtering metadata for "year" = 2012
82 | docs.query(
83 | data=[0.4,0.5,0.6], # required
84 | limit=1, # number of records to return
85 | filters={"year": {"$eq": 2012}}, # metadata filters
86 | )
87 |
88 | # Returns: ["vec1"]
89 |
90 | # Disconnect from the database
91 | vx.disconnect()
92 | ```
93 |
--------------------------------------------------------------------------------
/docs/integrations_amazon_bedrock.md:
--------------------------------------------------------------------------------
1 | # Integration: Amazon Bedrock
2 |
3 | This guide will walk you through an example using Amazon Bedrock SDK with `vecs`. We will create embeddings using the Amazon Titan Embeddings G1 – Text v1.2 (amazon.titan-embed-text-v1) model, insert these embeddings into a PostgreSQL database using vecs, and then query the collection to find the most similar sentences to a given query sentence.
4 |
5 | ## Create an Environment
6 |
7 | First, you need to set up your environment. You will need Python 3.7+ with the `vecs` and `boto3` libraries installed.
8 |
9 | You can install the necessary Python libraries using pip:
10 |
11 | ```sh
12 | pip install vecs boto3
13 | ```
14 |
15 | You'll also need:
16 |
17 | - [Credentials to your AWS account](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html)
18 | - [A Postgres Database with the pgvector extension](hosting.md)
19 |
20 | ## Create Embeddings
21 |
22 | Next, we will use Amazon’s Titan Embedding G1 - Text v1.2 model to create embeddings for a set of sentences.
23 |
24 | ```python
25 | import boto3
26 | import vecs
27 | import json
28 |
29 | client = boto3.client(
30 | 'bedrock-runtime',
31 | region_name='us-east-1',
32 | # Credentials from your AWS account
33 | aws_access_key_id='',
34 | aws_secret_access_key='',
35 | aws_session_token='',
36 | )
37 |
38 | dataset = [
39 | "The cat sat on the mat.",
40 | "The quick brown fox jumps over the lazy dog.",
41 | "Friends, Romans, countrymen, lend me your ears",
42 | "To be or not to be, that is the question.",
43 | ]
44 |
45 | embeddings = []
46 |
47 | for sentence in dataset:
48 | # invoke the embeddings model for each sentence
49 | response = client.invoke_model(
50 | body= json.dumps({"inputText": sentence}),
51 | modelId= "amazon.titan-embed-text-v1",
52 | accept = "application/json",
53 | contentType = "application/json"
54 | )
55 | # collect the embedding from the response
56 | response_body = json.loads(response["body"].read())
57 | # add the embedding to the embedding list
58 | embeddings.append((sentence, response_body.get("embedding"), {}))
59 |
60 | ```
61 |
62 | ### Store the Embeddings with vecs
63 |
64 | Now that we have our embeddings, we can insert them into a PostgreSQL database using vecs.
65 |
66 | ```python
67 | import vecs
68 |
69 | DB_CONNECTION = "postgresql://:@:/"
70 |
71 | # create vector store client
72 | vx = vecs.Client(DB_CONNECTION)
73 |
74 | # create a collection named 'sentences' with 1536 dimensional vectors
75 | # to match the default dimension of the Titan Embeddings G1 - Text model
76 | sentences = vx.get_or_create_collection(name="sentences", dimension=1536)
77 |
78 | # upsert the embeddings into the 'sentences' collection
79 | sentences.upsert(records=embeddings)
80 |
81 | # create an index for the 'sentences' collection
82 | sentences.create_index()
83 | ```
84 |
85 | ### Querying for Most Similar Sentences
86 |
87 | Now, we query the `sentences` collection to find the most similar sentences to a sample query sentence. First need to create an embedding for the query sentence. Next, we query the collection we created earlier to find the most similar sentences.
88 |
89 | ```python
90 | query_sentence = "A quick animal jumps over a lazy one."
91 |
92 | # create vector store client
93 | vx = vecs.Client(DB_CONNECTION)
94 |
95 | # create an embedding for the query sentence
96 | response = client.invoke_model(
97 | body= json.dumps({"inputText": query_sentence}),
98 | modelId= "amazon.titan-embed-text-v1",
99 | accept = "application/json",
100 | contentType = "application/json"
101 | )
102 |
103 | response_body = json.loads(response["body"].read())
104 |
105 | query_embedding = response_body.get("embedding")
106 |
107 | # query the 'sentences' collection for the most similar sentences
108 | results = sentences.query(
109 | data=query_embedding,
110 | limit=3,
111 | include_value = True
112 | )
113 |
114 | # print the results
115 | for result in results:
116 | print(result)
117 | ```
118 |
119 | This returns the most similar 3 records and their distance to the query vector.
120 | ```
121 | ('The quick brown fox jumps over the lazy dog.', 0.27600620558852)
122 | ('The cat sat on the mat.', 0.609986272479202)
123 | ('To be or not to be, that is the question.', 0.744849503688346)
124 | ```
125 |
--------------------------------------------------------------------------------
/docs/integrations_huggingface_inference_endpoints.md:
--------------------------------------------------------------------------------
1 | # Integration: Hugging Face Inference Endpoints
2 |
3 | This guide will walk you through an example integration of the Hugging Face Inference API with vecs. We will create embeddings using Hugging Face's `sentence-transformers/all-MiniLM-L6-v2` model, insert these embeddings into a PostgreSQL database using vecs, and then query vecs to find the most similar sentences to a given query sentence.
4 |
5 | ## Create a Hugging Face Inference Endpoint
6 |
7 | Head over to [Hugging Face's inference endpoints](https://ui.endpoints.huggingface.co/endpoints) and select `New Endpoint`.
8 |
9 | 
10 |
11 | Configure your endpoint with your model and provider of choice. In this example we'll use `sentence-transformers/all-MiniLM-L6-v2` and `AWS`.
12 |
13 | 
14 |
15 | Under "Advanced Configuration" select "Sentence Embeddings" as the "Task". Then click "Create Endpoint"
16 |
17 | 
18 |
19 |
20 |
21 | Once the endpoint starts up, take note of the `Endpoint URL`
22 | 
23 |
24 | !!! tip
25 |
26 | Don't forget to pause or delete your Hugging Face Inference Endpoint when you're not using it
27 |
28 | Finally, [create and copy an API key](https://huggingface.co/settings/tokens) we can use to authenticate with the inference endpoint.
29 |
30 | ## Create an Environment
31 |
32 | Next, you need to set up your environment. You will need Python 3.7+ with the vecs and requests installed.
33 |
34 | ```bash
35 | pip install vecs requests
36 | ```
37 |
38 | You'll also need [a Postgres Database with the pgvector extension](hosting.md)
39 |
40 | ## Create Embeddings
41 |
42 | We can use the Hugging Face endpoint to create embeddings for a set of sentences.
43 |
44 | ```python
45 | import requests
46 | import json
47 |
48 | huggingface_endpoint_url = ''
49 | huggingface_api_key = ''
50 |
51 | dataset = [
52 | "The cat sat on the mat.",
53 | "The quick brown fox jumps over the lazy dog.",
54 | "Friends, Romans, countrymen, lend me your ears",
55 | "To be or not to be, that is the question.",
56 | ]
57 |
58 | records = []
59 |
60 | for sentence in dataset:
61 | response = requests.post(
62 | huggingface_endpoint_url,
63 | headers={
64 | "Authorization": f"Bearer {huggingface_api_key}",
65 | "Content-Type": "application/json"
66 | },
67 | json={"inputs": sentence}
68 | )
69 | embedding = response.json()["embeddings"]
70 | records.append((sentence, embedding, {}))
71 | ```
72 |
73 | ## Store the Embeddings with vecs
74 |
75 | Now that we have our embeddings, we can insert them into a PostgreSQL database using vecs subbing in your DB_CONNECTION string.
76 |
77 | ```python
78 | import vecs
79 |
80 | DB_CONNECTION = "postgresql://:@:/"
81 |
82 | # create vector store client
83 | vx = vecs.Client(DB_CONNECTION)
84 |
85 | # create a collection named 'sentences' with 384 dimensional vectors (default dimension for paraphrase-MiniLM-L6-v2)
86 | sentences = vx.get_or_create_collection(name="sentences", dimension=384)
87 |
88 | # upsert the embeddings into the 'sentences' collection
89 | sentences.upsert(records=records)
90 |
91 | # create an index for the 'sentences' collection
92 | sentences.create_index()
93 | ```
94 |
95 |
96 | ## Querying for Most Similar Sentences
97 |
98 | Finally, we can query vecs to find the most similar sentences to a given query sentence. The query sentence is embedded using the same method as the sentences in the dataset, then we query the `sentences` collection with vecs.
99 |
100 | ```python
101 | query_sentence = "A quick animal jumps over a lazy one."
102 |
103 | # create an embedding for the query sentence
104 | response = requests.post(
105 | huggingface_endpoint_url,
106 | headers={
107 | "Authorization": f"Bearer {huggingface_api_key}",
108 | "Content-Type": "application/json"
109 | },
110 | json={"inputs": query_sentence}
111 | )
112 | query_embedding = response.json()["embeddings"]
113 |
114 | # query the 'sentences' collection for the most similar sentences
115 | results = sentences.query(
116 | data=query_embedding,
117 | limit=3,
118 | include_value = True
119 | )
120 |
121 | # print the results
122 | for result in results:
123 | print(result)
124 | ```
125 |
126 | Returns the most similar 3 records and theirdistance to the query vector.
127 |
128 | ```text
129 | ('The quick brown fox jumps over the lazy dog.', 0.256648302882697)
130 | ('The cat sat on the mat.', 0.78635900041167)
131 | ('To be or not to be, that is the question.', 1.04114070479544)
132 | ```
133 |
--------------------------------------------------------------------------------
/docs/integrations_openai.md:
--------------------------------------------------------------------------------
1 | # Integration: Open AI
2 |
3 | This guide will walk you through an example integration of the OpenAI API with the vecs Python library. We will create embeddings using OpenAI's `text-embedding-ada-002` model, insert these embeddings into a PostgreSQL database using vecs, and then query vecs to find the most similar sentences to a given query sentence.
4 |
5 | ## Create an Environment
6 |
7 | First, you need to set up your environment. You will need Python 3.7+ with the `vecs` and `openai` libraries installed.
8 |
9 | You can install the necessary Python libraries using pip:
10 |
11 | ```sh
12 | pip install vecs openai
13 | ```
14 |
15 | You'll also need:
16 |
17 | - [An OpenAI API Key](https://platform.openai.com/account/api-keys)
18 | - [A Postgres Database with the pgvector extension](hosting.md)
19 |
20 | ## Create Embeddings
21 |
22 | Next, we will use OpenAI's `text-embedding-ada-002` model to create embeddings for a set of sentences.
23 |
24 | ```python
25 | import openai
26 |
27 | openai.api_key = ''
28 |
29 | dataset = [
30 | "The cat sat on the mat.",
31 | "The quick brown fox jumps over the lazy dog.",
32 | "Friends, Romans, countrymen, lend me your ears",
33 | "To be or not to be, that is the question.",
34 | ]
35 |
36 | embeddings = []
37 |
38 | for sentence in dataset:
39 | response = openai.Embedding.create(
40 | model="text-embedding-ada-002",
41 | input=[sentence]
42 | )
43 | embeddings.append((sentence, response["data"][0]["embedding"], {}))
44 | ```
45 |
46 | ### Store the Embeddings with vecs
47 |
48 | Now that we have our embeddings, we can insert them into a PostgreSQL database using vecs.
49 |
50 | ```python
51 | import vecs
52 |
53 | DB_CONNECTION = "postgresql://:@:/"
54 |
55 | # create vector store client
56 | vx = vecs.Client(DB_CONNECTION)
57 |
58 | # create a collection named 'sentences' with 1536 dimensional vectors (default dimension for text-embedding-ada-002)
59 | sentences = vx.get_or_create_collection(name="sentences", dimension=1536)
60 |
61 | # upsert the embeddings into the 'sentences' collection
62 | sentences.upsert(records=embeddings)
63 |
64 | # create an index for the 'sentences' collection
65 | sentences.create_index()
66 | ```
67 |
68 | ### Querying for Most Similar Sentences
69 |
70 | Finally, we can query vecs to find the most similar sentences to a given query sentence. We will first need to create an embedding for the query sentence using the `text-embedding-ada-002` model.
71 |
72 | ```python
73 | query_sentence = "A quick animal jumps over a lazy one."
74 |
75 | # create an embedding for the query sentence
76 | response = openai.Embedding.create(
77 | model="text-embedding-ada-002",
78 | input=[query_sentence]
79 | )
80 | query_embedding = response["data"][0]["embedding"]
81 |
82 | # query the 'sentences' collection for the most similar sentences
83 | results = sentences.query(
84 | data=query_embedding,
85 | limit=3,
86 | include_value = True
87 | )
88 |
89 | # print the results
90 | for result in results:
91 | print(result)
92 | ```
93 |
94 | Returns the most similar 3 records and their distance to the query vector.
95 | ```
96 | ('The quick brown fox jumps over the lazy dog.', 0.0633971456300456)
97 | ('The cat sat on the mat.', 0.16474785399561)
98 | ('To be or not to be, that is the question.', 0.24531234467506)
99 | ```
100 |
--------------------------------------------------------------------------------
/docs/support_changelog.md:
--------------------------------------------------------------------------------
1 | # Changelog
2 |
3 |
4 | ## 0.1.0
5 |
6 | - Initial release
7 |
8 | ## 0.2.7
9 |
10 | - Feature: Added `vecs.Collection.disconnect()` to drop database connection
11 | - Feature: `vecs.Client` can be used as a context maanger to auto-close connections
12 | - Feature: Uses (indexed) containment operator `@>` for metadata equality filters where possible
13 | - Docs: Added docstrings to all methods, functions and modules
14 |
15 | ## 0.3.0
16 |
17 | - Feature: Collections can have `adapters` allowing upserting/querying by native media t types
18 | - Breaking Change: Renamed argument `Collection.upsert(vectors, ...)` to `Collection.upsert(records, ...)` in support of adapters
19 | - Breaking Change: Renamed argument `Collection.query(query_vector, ...)` to `Collection.query(data, ...)` in support of adapters
20 |
21 | ## 0.3.1
22 |
23 | - Feature: Metadata filtering with `$in`
24 |
25 | ## 0.4.0
26 |
27 | - Feature: pgvector 0.5.0
28 | - Feature: HNSW index support
29 |
30 | ## 0.4.1
31 |
32 | - Bugfix: removed errant print statement
33 |
34 | ## 0.4.2
35 |
36 | - Feature: Parameterized IVFFlat and HNSW indexes
37 | - Feature: Delete using metadata filter
38 |
39 | ## 0.4.3
40 |
41 | - Feature: Metadata filtering with `$contains`
42 |
43 | ## main
44 |
--------------------------------------------------------------------------------
/mkdocs.yml:
--------------------------------------------------------------------------------
1 | site_name: vecs
2 | site_url: https://supabase.github.io/vecs
3 | site_description: Vector Database Client for pgvector
4 |
5 | repo_name: supabase/vecs
6 | repo_url: https://github.com/supabase/vecs
7 |
8 | theme:
9 | name: 'readthedocs'
10 | features:
11 | - navigation.expand
12 | palette:
13 | primary: grey
14 | accent: red
15 | nav:
16 | - Introduction: index.md
17 | - Guides:
18 | - API: api.md
19 | - Hosting: hosting.md
20 | - Concepts:
21 | - Collections: concepts_collections.md
22 | - Adapters: concepts_adapters.md
23 | - Indexes: concepts_indexes.md
24 | - Metadata: concepts_metadata.md
25 | - Integrations:
26 | - OpenAI: integrations_openai.md
27 | - Bedrock: integrations_amazon_bedrock.md
28 | - HuggingFace Inference Endpoints: integrations_huggingface_inference_endpoints.md
29 | - Support:
30 | - Changelog: support_changelog.md
31 |
32 | markdown_extensions:
33 | - admonition
34 | - pymdownx.highlight:
35 | anchor_linenums: true
36 | line_spans: __span
37 | pygments_lang_class: true
38 | - pymdownx.inlinehilite
39 | - pymdownx.snippets
40 | - pymdownx.superfences
41 | - pymdownx.tabbed:
42 | alternate_style: true
43 |
--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
1 | [tool.black]
2 | line-length = 88
3 | exclude = '''
4 | /(
5 | \.git
6 | | \.hg
7 | | \.mypy_cache
8 | | \.tox
9 | | \.venv
10 | | _build
11 | | buck-out
12 | | build
13 | | dist
14 | )/
15 | '''
16 |
--------------------------------------------------------------------------------
/pytest.ini:
--------------------------------------------------------------------------------
1 | [pytest]
2 | addopts = src/tests
3 |
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 |
3 | """Setup script for the package."""
4 |
5 | import logging
6 | import os
7 | import sys
8 | from pathlib import Path
9 |
10 | import setuptools
11 |
12 | PACKAGE_NAME = "vecs"
13 | MINIMUM_PYTHON_VERSION = (3, 8, 0, "", 0)
14 |
15 |
16 | def check_python_version():
17 | """Exit when the Python version is too low."""
18 | if sys.version_info < MINIMUM_PYTHON_VERSION:
19 | sys.exit(
20 | "At least Python {0}.{1}.{2} is required.".format(
21 | *MINIMUM_PYTHON_VERSION[:3]
22 | )
23 | )
24 |
25 |
26 | def read_package_variable(key, filename="__init__.py"):
27 | """Read the value of a variable from the package without importing."""
28 | module_path = os.path.join("src", PACKAGE_NAME, filename)
29 | with open(module_path) as module:
30 | for line in module:
31 | parts = line.strip().split(" ", 2)
32 | if parts[:-1] == [key, "="]:
33 | return parts[-1].strip("'").strip('"')
34 | logging.warning("'%s' not found in '%s'", key, module_path)
35 | raise KeyError(key)
36 |
37 |
38 | check_python_version()
39 |
40 | long_description = (Path(__file__).parent / "README.md").read_text()
41 |
42 | REQUIRES = [
43 | "pgvector==0.3.*",
44 | "sqlalchemy==2.*",
45 | "psycopg2-binary==2.9.*",
46 | "flupy==1.*",
47 | "deprecated==1.2.*",
48 | ]
49 |
50 |
51 | setuptools.setup(
52 | name=read_package_variable("__project__"),
53 | version=read_package_variable("__version__"),
54 | description="pgvector client",
55 | long_description=long_description,
56 | long_description_content_type="text/markdown",
57 | url="https://github.com/supabase/vecs",
58 | author="Oliver Rice",
59 | packages=setuptools.find_packages("src", exclude="tests"),
60 | package_dir={"": "src"},
61 | package_data={"": ["py.typed"]},
62 | tests_require=["pytest"],
63 | license="MIT",
64 | classifiers=[
65 | "Development Status :: 4 - Beta",
66 | "Natural Language :: English",
67 | "Operating System :: OS Independent",
68 | "Programming Language :: Python",
69 | "Programming Language :: Python :: 3",
70 | "Programming Language :: Python :: 3.8",
71 | "Programming Language :: Python :: 3.9",
72 | "Programming Language :: Python :: 3.10",
73 | "Programming Language :: Python :: 3.11",
74 | "Programming Language :: Python :: 3.12",
75 | ],
76 | install_requires=REQUIRES,
77 | extras_require={
78 | "dev": ["pytest", "parse", "numpy", "pytest-cov"],
79 | "docs": [
80 | "mkdocs",
81 | "pygments",
82 | "pymdown-extensions",
83 | "pymarkdown",
84 | "mike",
85 | ],
86 | "text_embedding": ["sentence-transformers==2.*"],
87 | },
88 | )
89 |
--------------------------------------------------------------------------------
/src/py.typed:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/src/tests/conftest.py:
--------------------------------------------------------------------------------
1 | # pylint: disable=redefined-outer-name,no-member
2 |
3 | import json
4 | import os
5 | import subprocess
6 | import time
7 | from typing import Generator
8 |
9 | import pytest
10 | from parse import parse
11 | from sqlalchemy import create_engine, text
12 |
13 | import vecs
14 |
15 | PYTEST_DB = "postgresql://postgres:password@localhost:5611/vecs_db"
16 |
17 |
18 | @pytest.fixture(scope="session")
19 | def maybe_start_pg() -> Generator[None, None, None]:
20 | """Creates a postgres 15 docker container that can be connected
21 | to using the PYTEST_DB connection string"""
22 |
23 | container_name = "vecs_pg"
24 | image = "supabase/postgres:15.1.1.78"
25 |
26 | connection_template = "postgresql://{user}:{pw}@{host}:{port:d}/{db}"
27 | conn_args = parse(connection_template, PYTEST_DB)
28 |
29 | if "GITHUB_SHA" in os.environ:
30 | yield
31 | return
32 |
33 | try:
34 | is_running = (
35 | subprocess.check_output(
36 | ["docker", "inspect", "-f", "{{.State.Running}}", container_name]
37 | )
38 | .decode()
39 | .strip()
40 | == "true"
41 | )
42 | except subprocess.CalledProcessError:
43 | # Can't inspect container if it isn't running
44 | is_running = False
45 |
46 | if is_running:
47 | yield
48 | return
49 |
50 | out = subprocess.check_output(
51 | [
52 | "docker",
53 | "run",
54 | "--rm",
55 | "--name",
56 | container_name,
57 | "-p",
58 | f"{conn_args['port']}:5432",
59 | "-d",
60 | "-e",
61 | f"POSTGRES_DB={conn_args['db']}",
62 | "-e",
63 | f"POSTGRES_PASSWORD={conn_args['pw']}",
64 | "-e",
65 | f"POSTGRES_USER={conn_args['user']}",
66 | "--health-cmd",
67 | "pg_isready",
68 | "--health-interval",
69 | "3s",
70 | "--health-timeout",
71 | "3s",
72 | "--health-retries",
73 | "15",
74 | image,
75 | ]
76 | )
77 | # Wait for postgres to become healthy
78 | for _ in range(10):
79 | out = subprocess.check_output(["docker", "inspect", container_name])
80 | inspect_info = json.loads(out)[0]
81 | health_status = inspect_info["State"]["Health"]["Status"]
82 | if health_status == "healthy":
83 | break
84 | else:
85 | time.sleep(1)
86 | else:
87 | raise Exception("Could not reach postgres comtainer. Check docker installation")
88 | yield
89 | # subprocess.call(["docker", "stop", container_name])
90 | return
91 |
92 |
93 | @pytest.fixture(scope="function")
94 | def clean_db(maybe_start_pg: None) -> Generator[str, None, None]:
95 | eng = create_engine(PYTEST_DB)
96 | with eng.begin() as connection:
97 | connection.execute(text("drop schema if exists vecs cascade;"))
98 | yield PYTEST_DB
99 | eng.dispose()
100 |
101 |
102 | @pytest.fixture(scope="function")
103 | def client(clean_db: str) -> Generator[vecs.Client, None, None]:
104 | client_ = vecs.create_client(clean_db)
105 | yield client_
106 |
--------------------------------------------------------------------------------
/src/tests/test_adapters.py:
--------------------------------------------------------------------------------
1 | from itertools import zip_longest
2 |
3 | import numpy as np
4 | import pytest
5 |
6 | import vecs
7 | from vecs.adapter import Adapter, AdapterContext, AdapterStep
8 | from vecs.adapter.markdown import MarkdownChunker
9 | from vecs.adapter.noop import NoOp
10 | from vecs.adapter.text import ParagraphChunker, TextEmbedding
11 | from vecs.exc import ArgError, MismatchedDimension
12 |
13 |
14 | def test_create_collection_with_adapter(client: vecs.Client) -> None:
15 | client.get_or_create_collection(
16 | name="ping",
17 | adapter=Adapter([TextEmbedding(model="all-MiniLM-L6-v2")]),
18 | )
19 | # Mismatched between existing collection dim (384) and provided dimension (1)
20 | with pytest.raises(MismatchedDimension):
21 | client.get_or_create_collection(name="ping", dimension=1)
22 |
23 | # Mismatched between dim arg and exported dim from adapter
24 | with pytest.raises(MismatchedDimension):
25 | client.get_or_create_collection(
26 | name="pong",
27 | dimension=9,
28 | adapter=Adapter([TextEmbedding(model="all-MiniLM-L6-v2")]),
29 | )
30 |
31 | client.get_or_create_collection(
32 | name="foo",
33 | dimension=16,
34 | )
35 | # Mismatched between exported dim from adapter and existing collection
36 | with pytest.raises(MismatchedDimension):
37 | client.get_or_create_collection(
38 | name="foo",
39 | adapter=Adapter([TextEmbedding(model="all-MiniLM-L6-v2")]),
40 | )
41 |
42 | # No detectable dimension
43 | with pytest.raises(ArgError):
44 | client.get_or_create_collection(name="foo")
45 |
46 |
47 | def test_adapter_with_no_steps_error() -> None:
48 | with pytest.raises(ArgError):
49 | adapter = Adapter(steps=[])
50 |
51 |
52 | def test_adapter_step_has_default_exported_dim() -> None:
53 | assert ParagraphChunker(skip_during_query=True).exported_dimension is None
54 |
55 |
56 | def test_adapter_does_not_export_dimension() -> None:
57 | class Dummy(AdapterStep):
58 | def __call__(
59 | self,
60 | records,
61 | adapter_context,
62 | ):
63 | raise Exception("not relevant to the test")
64 |
65 | adapter_step = Dummy()
66 | assert adapter_step.exported_dimension is None
67 | assert Adapter([adapter_step]).exported_dimension is None
68 |
69 |
70 | def test_noop_adapter_dimension() -> None:
71 | noop = NoOp(dimension=9)
72 | assert noop.exported_dimension == 9
73 |
74 |
75 | def test_paragraph_chunker_adapter() -> None:
76 | chunker = ParagraphChunker(skip_during_query=True)
77 | res = [
78 | x
79 | for x in chunker(
80 | [("1", "first para\n\nnext para", {})], AdapterContext("upsert")
81 | )
82 | ]
83 | assert res == [("1_para_000", "first para", {}), ("1_para_001", "next para", {})]
84 |
85 | res = [
86 | x
87 | for x in chunker([("", "first para\n\nnext para", {})], AdapterContext("query"))
88 | ]
89 | assert res == [("", "first para\n\nnext para", {})]
90 |
91 |
92 | def test_text_embedding_adapter() -> None:
93 | emb = TextEmbedding(model="all-MiniLM-L6-v2")
94 |
95 | records = [
96 | ("1", "first one", {"a": 1}),
97 | ("2", "next one", {"a": 2}),
98 | ("3", "last one", {"a": 3}),
99 | ]
100 | res = [x for x in emb(records, AdapterContext("upsert"))]
101 | assert len(res) == 3
102 | assert res[0][0] == "1"
103 | assert res[0][2] == {"a": 1}
104 | assert len(res[0][1]) == 384
105 |
106 | # test larger batch size does not impact result
107 | ada = TextEmbedding(model="all-MiniLM-L6-v2", batch_size=2)
108 | res_batch = [x for x in ada(records, AdapterContext("upsert"))]
109 | for (l_id, l_vec, l_meta), (r_id, r_vec, r_meta) in zip_longest(res, res_batch): # type: ignore
110 | assert l_id == r_id
111 | assert np.allclose(l_vec, r_vec, rtol=0.003)
112 | assert l_meta == r_meta
113 |
114 |
115 | def test_text_integration_adapter(client: vecs.Client) -> None:
116 | docs = client.get_or_create_collection(
117 | name="docs",
118 | adapter=Adapter(
119 | [
120 | ParagraphChunker(skip_during_query=False),
121 | TextEmbedding(model="all-MiniLM-L6-v2"),
122 | ]
123 | ),
124 | )
125 |
126 | docs.upsert(
127 | [
128 | ("1", "world hello", {"a": 1}),
129 | ("2", "foo bar", {}),
130 | ("3", "bar baz", {}),
131 | ("4", "long text\n\nshould split", {}),
132 | ]
133 | )
134 |
135 | docs.create_index()
136 |
137 | res = docs.query(data="hi", limit=1)
138 | assert res == ["1_para_000"]
139 |
140 | # the last record in the dataset should be split by the paragraph chunker
141 | res = docs.query(data="hi", limit=10)
142 | assert len(res) == 5
143 |
144 | # providing a vector works if you skip the adapter
145 |
146 | docs.upsert(
147 | [
148 | ("6", np.zeros(384), {}),
149 | ("7", np.zeros(384), {}),
150 | ],
151 | skip_adapter=True,
152 | )
153 |
154 | res = docs.query(data=np.zeros(384), limit=2, skip_adapter=True)
155 | assert len(res) == 2
156 |
157 | with pytest.raises(ArgError):
158 | # if the cardinality changes due to adapter pre-processing, raise an error
159 | docs.query(data="I split \n\n into multiple records \n\n not good")
160 |
161 |
162 | def test_markdown_chunker_normal_headings() -> None:
163 | chunker = MarkdownChunker(skip_during_query=True)
164 | res = [
165 | x
166 | for x in chunker(
167 | [
168 | (
169 | "1",
170 | "# heading 1\n## heading 2\n### heading 3\n#### heading 4\nwith some text\n##### heading 5\n###### heading 6",
171 | {"some": 1},
172 | )
173 | ],
174 | AdapterContext("upsert"),
175 | max_tokens=30,
176 | )
177 | ]
178 | assert res == [
179 | ("1_head_000", "# heading 1\n", {"some": 1}),
180 | ("1_head_001", "## heading 2\n", {"some": 1}),
181 | ("1_head_002", "### heading 3\n", {"some": 1}),
182 | ("1_head_003", "#### heading 4\nwith some text\n", {"some": 1}),
183 | ("1_head_004", "##### heading 5\n", {"some": 1}),
184 | ("1_head_005", "###### heading 6", {"some": 1}),
185 | ]
186 |
187 | res = [
188 | x
189 | for x in chunker([("", "# heading1\n# heading2", {})], AdapterContext("query"))
190 | ]
191 | assert res == [("", "# heading1\n# heading2", {})]
192 |
193 |
194 | def test_invalid_headings() -> None:
195 | # these should all clump into one chunk since none of them are valid headings
196 |
197 | chunker = MarkdownChunker(skip_during_query=True)
198 | res = [
199 | x
200 | for x in chunker(
201 | [
202 | (
203 | "1",
204 | "#heading 1\n####### heading 2\nheading ???\n-=-===\n#!## heading !3",
205 | {},
206 | )
207 | ],
208 | AdapterContext("upsert"),
209 | max_tokens=30,
210 | )
211 | ]
212 | assert res == [
213 | (
214 | "1_head_000",
215 | "#heading 1\n####### heading 2\nheading ???\n-=-===\n#!## heading !3",
216 | {},
217 | ),
218 | ]
219 |
220 |
221 | def test_max_tokens() -> None:
222 | chunker = MarkdownChunker(skip_during_query=True)
223 |
224 | res = [
225 | x
226 | for x in chunker(
227 | [
228 | (
229 | "2",
230 | "this is quite a long sentence which will have to be split into two chunks",
231 | {},
232 | )
233 | ],
234 | AdapterContext("upsert"),
235 | max_tokens=10,
236 | )
237 | ]
238 | assert res == [
239 | ("2_head_000", "this is quite a long sentence which will have to", {}),
240 | ("2_head_001", "be split into two chunks", {}),
241 | ]
242 |
243 | res = [
244 | x
245 | for x in chunker(
246 | [
247 | (
248 | "3",
249 | "this sentence is so long that it won't even fit in two chunks, it'll have to go into three chunks which is exciting",
250 | {},
251 | )
252 | ],
253 | AdapterContext("upsert"),
254 | max_tokens=10,
255 | )
256 | ]
257 | assert res == [
258 | ("3_head_000", "this sentence is so long that it won't even fit", {}),
259 | ("3_head_001", "in two chunks, it'll have to go into three chunks", {}),
260 | ("3_head_002", "which is exciting", {}),
261 | ]
262 |
263 | with pytest.raises(ValueError):
264 | res = [
265 | x
266 | for x in chunker(
267 | [
268 | (
269 | "4",
270 | "this doesn't really matter since it will throw an error anyway",
271 | {},
272 | )
273 | ],
274 | AdapterContext("upsert"),
275 | max_tokens=-5,
276 | )
277 | ]
278 |
--------------------------------------------------------------------------------
/src/tests/test_client.py:
--------------------------------------------------------------------------------
1 | import pytest
2 |
3 | import vecs
4 |
5 |
6 | def test_extracts_vector_version(client: vecs.Client) -> None:
7 | # pgvector version is sucessfully extracted
8 | assert client.vector_version != ""
9 | assert client.vector_version.count(".") >= 2
10 |
11 |
12 | def test_create_collection(client: vecs.Client) -> None:
13 | with pytest.warns(DeprecationWarning):
14 | client.create_collection(name="docs", dimension=384)
15 |
16 | with pytest.raises(vecs.exc.CollectionAlreadyExists):
17 | client.create_collection(name="docs", dimension=384)
18 |
19 |
20 | def test_get_or_create_collection(client: vecs.Client) -> None:
21 | client.get_or_create_collection(name="resumes", dimension=1536)
22 | # no error is raised
23 | client.get_or_create_collection(name="resumes", dimension=1536)
24 |
25 |
26 | def test_get_or_create_collection_dim_change(client: vecs.Client) -> None:
27 | client.get_or_create_collection(name="resumes", dimension=1536)
28 | with pytest.raises(vecs.exc.MismatchedDimension):
29 | client.get_or_create_collection(name="resumes", dimension=1)
30 |
31 |
32 | def test_get_collection(client: vecs.Client) -> None:
33 | with pytest.warns(DeprecationWarning):
34 | with pytest.raises(vecs.exc.CollectionNotFound):
35 | client.get_collection(name="foo")
36 |
37 | client.create_collection(name="foo", dimension=384)
38 |
39 | foo = client.get_collection(name="foo")
40 | assert foo.name == "foo"
41 |
42 |
43 | def test_list_collections(client: vecs.Client) -> None:
44 | assert len(client.list_collections()) == 0
45 | client.get_or_create_collection(name="docs", dimension=384)
46 | client.get_or_create_collection(name="books", dimension=1586)
47 | collections = client.list_collections()
48 | assert len(collections) == 2
49 |
50 |
51 | def test_delete_collection(client: vecs.Client) -> None:
52 | client.get_or_create_collection(name="books", dimension=1586)
53 | collections = client.list_collections()
54 | assert len(collections) == 1
55 |
56 | client.delete_collection("books")
57 |
58 | collections = client.list_collections()
59 | assert len(collections) == 0
60 |
61 | # does not raise when does not exist
62 | client.delete_collection("books")
63 |
64 |
65 | def test_dispose(client: vecs.Client) -> None:
66 | # Connect and disconnect in context manager
67 | with client:
68 | client.get_or_create_collection(name="books", dimension=1)
69 | collections = client.list_collections()
70 | assert len(collections) == 1
71 |
72 | # engine.dispose re-creates the connection pool so
73 | # confirm that the client can still re-connect transparently
74 | assert len(client.list_collections()) == 1
75 |
--------------------------------------------------------------------------------
/src/tests/test_collection.py:
--------------------------------------------------------------------------------
1 | import itertools
2 | import random
3 |
4 | import numpy as np
5 | import pytest
6 |
7 | import vecs
8 | from vecs import IndexArgsHNSW, IndexArgsIVFFlat, IndexMethod
9 | from vecs.exc import ArgError
10 |
11 |
12 | def test_upsert(client: vecs.Client) -> None:
13 | n_records = 100
14 | dim = 384
15 |
16 | movies = client.get_or_create_collection(name="ping", dimension=dim)
17 |
18 | # collection initially empty
19 | assert len(movies) == 0
20 |
21 | records = [
22 | (
23 | f"vec{ix}",
24 | vec,
25 | {
26 | "genre": random.choice(["action", "rom-com", "drama"]),
27 | "year": int(50 * random.random()) + 1970,
28 | },
29 | )
30 | for ix, vec in enumerate(np.random.random((n_records, dim)))
31 | ]
32 |
33 | # insert works
34 | movies.upsert(records)
35 | assert len(movies) == n_records
36 |
37 | # upserting overwrites
38 | new_record = ("vec0", np.zeros(384), {})
39 | movies.upsert([new_record])
40 | db_record = movies["vec0"]
41 | db_record[0] == new_record[0]
42 | db_record[1] == new_record[1]
43 | db_record[2] == new_record[2]
44 |
45 |
46 | def test_fetch(client: vecs.Client) -> None:
47 | n_records = 100
48 | dim = 384
49 |
50 | movies = client.get_or_create_collection(name="ping", dimension=dim)
51 |
52 | records = [
53 | (
54 | f"vec{ix}",
55 | vec,
56 | {
57 | "genre": random.choice(["action", "rom-com", "drama"]),
58 | "year": int(50 * random.random()) + 1970,
59 | },
60 | )
61 | for ix, vec in enumerate(np.random.random((n_records, dim)))
62 | ]
63 |
64 | # insert works
65 | movies.upsert(records)
66 |
67 | # test basic usage
68 | fetch_ids = ["vec0", "vec15", "vec99"]
69 | res = movies.fetch(ids=fetch_ids)
70 | assert len(res) == 3
71 | ids = set([x[0] for x in res])
72 | assert all([x in ids for x in fetch_ids])
73 |
74 | # test one of the keys does not exist not an error
75 | fetch_ids = ["vec0", "vec15", "does not exist"]
76 | res = movies.fetch(ids=fetch_ids)
77 | assert len(res) == 2
78 |
79 | # bad input
80 | with pytest.raises(vecs.exc.ArgError):
81 | movies.fetch(ids="should_be_a_list")
82 |
83 |
84 | def test_delete(client: vecs.Client) -> None:
85 | n_records = 100
86 | dim = 384
87 |
88 | movies = client.get_or_create_collection(name="ping", dimension=dim)
89 |
90 | records = [
91 | (
92 | f"vec{ix}",
93 | vec,
94 | {
95 | "genre": genre,
96 | "year": int(50 * random.random()) + 1970,
97 | },
98 | )
99 | for (ix, vec), genre in zip(
100 | enumerate(np.random.random((n_records, dim))),
101 | itertools.cycle(["action", "rom-com", "drama"]),
102 | )
103 | ]
104 |
105 | # insert works
106 | movies.upsert(records)
107 |
108 | # delete by IDs.
109 | delete_ids = ["vec0", "vec15", "vec99"]
110 | movies.delete(ids=delete_ids)
111 | assert len(movies) == n_records - len(delete_ids)
112 |
113 | # insert works
114 | movies.upsert(records)
115 |
116 | # delete with filters
117 | genre_to_delete = "action"
118 | deleted_ids_by_genre = movies.delete(filters={"genre": {"$eq": genre_to_delete}})
119 | assert len(deleted_ids_by_genre) == 34
120 |
121 | # bad input
122 | with pytest.raises(vecs.exc.ArgError):
123 | movies.delete(ids="should_be_a_list")
124 |
125 | # bad input: neither ids nor filters provided.
126 | with pytest.raises(vecs.exc.ArgError):
127 | movies.delete()
128 |
129 | # bad input: should only provide either ids or filters, not both
130 | with pytest.raises(vecs.exc.ArgError):
131 | movies.delete(ids=["vec0"], filters={"genre": {"$eq": genre_to_delete}})
132 |
133 |
134 | def test_repr(client: vecs.Client) -> None:
135 | movies = client.get_or_create_collection(name="movies", dimension=99)
136 | assert repr(movies) == 'vecs.Collection(name="movies", dimension=99)'
137 |
138 |
139 | def test_getitem(client: vecs.Client) -> None:
140 | movies = client.get_or_create_collection(name="movies", dimension=3)
141 | movies.upsert(records=[("1", [1, 2, 3], {})])
142 |
143 | assert movies["1"] is not None
144 | assert len(movies["1"]) == 3
145 |
146 | with pytest.raises(KeyError):
147 | assert movies["2"] is not None
148 |
149 | with pytest.raises(vecs.exc.ArgError):
150 | movies[["only strings work not lists"]]
151 |
152 |
153 | @pytest.mark.filterwarnings("ignore:Query does")
154 | def test_query(client: vecs.Client) -> None:
155 | n_records = 100
156 | dim = 64
157 |
158 | bar = client.get_or_create_collection(name="bar", dimension=dim)
159 |
160 | records = [
161 | (
162 | f"vec{ix}",
163 | vec,
164 | {
165 | "genre": random.choice(["action", "rom-com", "drama"]),
166 | "year": int(50 * random.random()) + 1970,
167 | },
168 | )
169 | for ix, vec in enumerate(np.random.random((n_records, dim)))
170 | ]
171 |
172 | bar.upsert(records)
173 |
174 | _, query_vec, query_meta = bar["vec5"]
175 |
176 | top_k = 7
177 |
178 | res = bar.query(
179 | data=query_vec,
180 | limit=top_k,
181 | filters=None,
182 | measure="cosine_distance",
183 | include_value=False,
184 | include_metadata=False,
185 | )
186 |
187 | # correct number of results
188 | assert len(res) == top_k
189 | # most similar to self
190 | assert res[0] == "vec5"
191 |
192 | with pytest.raises(vecs.exc.ArgError):
193 | res = bar.query(
194 | data=query_vec,
195 | limit=1001,
196 | )
197 |
198 | with pytest.raises(vecs.exc.ArgError):
199 | res = bar.query(
200 | data=query_vec,
201 | probes=0,
202 | )
203 |
204 | with pytest.raises(vecs.exc.ArgError):
205 | res = bar.query(
206 | data=query_vec,
207 | probes=-1,
208 | )
209 |
210 | with pytest.raises(vecs.exc.ArgError):
211 | res = bar.query(
212 | data=query_vec,
213 | probes="a", # type: ignore
214 | )
215 |
216 | with pytest.raises(vecs.exc.ArgError):
217 | res = bar.query(data=query_vec, limit=top_k, measure="invalid")
218 |
219 | # skip_adapter has no effect (no adapter present)
220 | res = bar.query(data=query_vec, limit=top_k, skip_adapter=True)
221 | assert len(res) == top_k
222 |
223 | # include_value
224 | res = bar.query(
225 | data=query_vec,
226 | limit=top_k,
227 | filters=None,
228 | measure="cosine_distance",
229 | include_value=True,
230 | )
231 | assert len(res[0]) == 2
232 | assert res[0][0] == "vec5"
233 | assert pytest.approx(res[0][1]) == 0
234 |
235 | # include_metadata
236 | res = bar.query(
237 | data=query_vec,
238 | limit=top_k,
239 | filters=None,
240 | measure="cosine_distance",
241 | include_metadata=True,
242 | )
243 | assert len(res[0]) == 2
244 | assert res[0][0] == "vec5"
245 | assert res[0][1] == query_meta
246 |
247 | # include_vector
248 | res = bar.query(
249 | data=query_vec,
250 | limit=top_k,
251 | filters=None,
252 | measure="cosine_distance",
253 | include_vector=True,
254 | )
255 | assert len(res[0]) == 2
256 | assert res[0][0] == "vec5"
257 | assert all(res[0][1] == query_vec)
258 |
259 | # test for different numbers of probes
260 | assert len(bar.query(data=query_vec, limit=top_k, probes=10)) == top_k
261 |
262 | assert len(bar.query(data=query_vec, limit=top_k, probes=5)) == top_k
263 |
264 | assert len(bar.query(data=query_vec, limit=top_k, probes=1)) == top_k
265 |
266 | assert len(bar.query(data=query_vec, limit=top_k, probes=999)) == top_k
267 |
268 |
269 | @pytest.mark.filterwarnings("ignore:Query does")
270 | def test_query_filters(client: vecs.Client) -> None:
271 | n_records = 100
272 | dim = 4
273 |
274 | bar = client.get_or_create_collection(name="bar", dimension=dim)
275 |
276 | records = [
277 | (f"0", [0, 0, 0, 0], {"year": 1990}),
278 | (f"1", [1, 0, 0, 0], {"year": 1995}),
279 | (f"2", [1, 1, 0, 0], {"year": 2005}),
280 | (f"3", [1, 1, 1, 0], {"year": 2001}),
281 | (f"4", [1, 1, 1, 1], {"year": 1985}),
282 | (f"5", [2, 1, 1, 1], {"year": 1863}),
283 | (f"6", [2, 2, 1, 1], {"year": 2021}),
284 | (f"7", [2, 2, 2, 1], {"year": 2019}),
285 | (f"8", [2, 2, 2, 2], {"year": 2003}),
286 | (f"9", [3, 2, 2, 2], {"year": 1997}),
287 | ]
288 |
289 | bar.upsert(records)
290 |
291 | query_rec = records[0]
292 |
293 | res = bar.query(
294 | data=query_rec[1],
295 | limit=3,
296 | filters={"year": {"$lt": 1990}},
297 | measure="cosine_distance",
298 | include_value=False,
299 | include_metadata=False,
300 | )
301 |
302 | assert res
303 |
304 | with pytest.raises(vecs.exc.FilterError):
305 | bar.query(
306 | data=query_rec[1],
307 | limit=3,
308 | filters=["wrong type"],
309 | measure="cosine_distance",
310 | )
311 |
312 | with pytest.raises(vecs.exc.FilterError):
313 | bar.query(
314 | data=query_rec[1],
315 | limit=3,
316 | # multiple keys
317 | filters={"key1": {"$eq": "v"}, "key2": {"$eq": "v"}},
318 | measure="cosine_distance",
319 | )
320 |
321 | with pytest.raises(vecs.exc.FilterError):
322 | bar.query(
323 | data=query_rec[1],
324 | limit=3,
325 | # bad key
326 | filters={1: {"$eq": "v"}},
327 | measure="cosine_distance",
328 | )
329 |
330 | with pytest.raises(vecs.exc.FilterError):
331 | bar.query(
332 | data=query_rec[1],
333 | limit=3,
334 | # and requires a list
335 | filters={"$and": {"year": {"$eq": 1997}}},
336 | measure="cosine_distance",
337 | )
338 |
339 | # AND
340 | assert (
341 | len(
342 | bar.query(
343 | data=query_rec[1],
344 | limit=3,
345 | # and requires a list
346 | filters={
347 | "$and": [
348 | {"year": {"$eq": 1997}},
349 | {"year": {"$eq": 1997}},
350 | ]
351 | },
352 | measure="cosine_distance",
353 | )
354 | )
355 | == 1
356 | )
357 |
358 | # OR
359 | assert (
360 | len(
361 | bar.query(
362 | data=query_rec[1],
363 | limit=3,
364 | # and requires a list
365 | filters={
366 | "$or": [
367 | {"year": {"$eq": 1997}},
368 | {"year": {"$eq": 2001}},
369 | ]
370 | },
371 | measure="cosine_distance",
372 | )
373 | )
374 | == 2
375 | )
376 |
377 | with pytest.raises(vecs.exc.FilterError):
378 | bar.query(
379 | data=query_rec[1],
380 | limit=3,
381 | # bad value, too many conditions
382 | filters={"year": {"$eq": 1997, "$ne": 1998}},
383 | measure="cosine_distance",
384 | )
385 |
386 | with pytest.raises(vecs.exc.FilterError):
387 | bar.query(
388 | data=query_rec[1],
389 | limit=3,
390 | # bad value, unknown operator
391 | filters={"year": {"$no_op": 1997}},
392 | measure="cosine_distance",
393 | )
394 |
395 | # ne
396 | assert (
397 | len(
398 | bar.query(
399 | data=query_rec[1],
400 | limit=3,
401 | # and requires a list
402 | filters={"year": {"$ne": 2000}},
403 | measure="cosine_distance",
404 | )
405 | )
406 | == 3
407 | )
408 |
409 | # lte
410 | assert (
411 | len(
412 | bar.query(
413 | data=query_rec[1],
414 | limit=3,
415 | # and requires a list
416 | filters={"year": {"$lte": 1989}},
417 | measure="cosine_distance",
418 | )
419 | )
420 | == 2
421 | )
422 |
423 | # gt
424 | assert (
425 | len(
426 | bar.query(
427 | data=query_rec[1],
428 | limit=3,
429 | # and requires a list
430 | filters={"year": {"$gt": 2019}},
431 | measure="cosine_distance",
432 | )
433 | )
434 | == 1
435 | )
436 |
437 | # gte
438 | assert (
439 | len(
440 | bar.query(
441 | data=query_rec[1],
442 | limit=3,
443 | # and requires a list
444 | filters={"year": {"$gte": 2019}},
445 | measure="cosine_distance",
446 | )
447 | )
448 | == 2
449 | )
450 |
451 |
452 | def test_filters_eq(client: vecs.Client) -> None:
453 | bar = client.get_or_create_collection(name="bar", dimension=4)
454 |
455 | records = [
456 | ("0", [0, 0, 0, 0], {"a": 1}),
457 | ("1", [1, 0, 0, 0], {"a": 2}),
458 | ("2", [1, 1, 0, 0], {"a": 3}),
459 | ("3", [1, 1, 1, 0], {"b": [1, 2]}),
460 | ("4", [1, 1, 1, 1], {"b": [1, 3]}),
461 | ("5", [1, 1, 1, 1], {"b": 1}),
462 | ("6", [1, 1, 1, 1], {"c": {"d": "hi"}}),
463 | ]
464 |
465 | bar.upsert(records)
466 | bar.create_index()
467 |
468 | # Simple equality of number: has match
469 | assert bar.query(
470 | data=[0, 0, 0, 0],
471 | limit=3,
472 | filters={"a": {"$eq": 1}},
473 | ) == ["0"]
474 |
475 | # Simple equality of number: no match
476 | assert (
477 | bar.query(
478 | data=[0, 0, 0, 0],
479 | limit=3,
480 | filters={"a": {"$eq": 5}},
481 | )
482 | == []
483 | )
484 |
485 | # Equality of array to value: no match
486 | assert (
487 | bar.query(
488 | data=[0, 0, 0, 0],
489 | limit=3,
490 | filters={"a": {"$eq": [1]}},
491 | )
492 | == []
493 | )
494 |
495 | # Equality of value to array: no match
496 | assert (
497 | bar.query(
498 | data=[0, 0, 0, 0],
499 | limit=3,
500 | filters={"b": {"$eq": 2}},
501 | )
502 | == []
503 | )
504 |
505 | # Equality of sub-array to array: no match
506 | assert (
507 | bar.query(
508 | data=[0, 0, 0, 0],
509 | limit=3,
510 | filters={"b": {"$eq": [1]}},
511 | )
512 | == []
513 | )
514 |
515 | # Equality of array to array: match
516 | assert bar.query(
517 | data=[0, 0, 0, 0],
518 | limit=3,
519 | filters={"b": {"$eq": [1, 2]}},
520 | ) == ["3"]
521 |
522 | # Equality of scalar to dict (key matches): no match
523 | assert (
524 | bar.query(
525 | data=[0, 0, 0, 0],
526 | limit=3,
527 | filters={"c": {"$eq": "d"}},
528 | )
529 | == []
530 | )
531 |
532 | # Equality of scalar to dict (value matches): no match
533 | assert (
534 | bar.query(
535 | data=[0, 0, 0, 0],
536 | limit=3,
537 | filters={"c": {"$eq": "hi"}},
538 | )
539 | == []
540 | )
541 |
542 | # Equality of dict to dict: match
543 | assert bar.query(
544 | data=[0, 0, 0, 0],
545 | limit=3,
546 | filters={"c": {"$eq": {"d": "hi"}}},
547 | ) == ["6"]
548 |
549 |
550 | def test_filters_in(client: vecs.Client) -> None:
551 | bar = client.get_or_create_collection(name="bar", dimension=4)
552 |
553 | records = [
554 | ("0", [0, 0, 0, 0], {"a": 1, "b": 2}),
555 | ("1", [1, 0, 0, 0], {"a": [1, 2, 3]}),
556 | ("2", [1, 1, 0, 0], {"a": {"1": "2"}}),
557 | ("3", [0, 0, 0, 0], {"a": "1"}),
558 | ]
559 |
560 | bar.upsert(records)
561 | bar.create_index()
562 |
563 | # int value of "a" is contained by [1, 2]
564 | assert bar.query(
565 | data=[0, 0, 0, 0],
566 | limit=3,
567 | filters={"a": {"$in": [1, 2]}},
568 | ) == ["0"]
569 |
570 | # str value of "a" is contained by ["1", "2"]
571 | assert bar.query(
572 | data=[0, 0, 0, 0],
573 | limit=3,
574 | filters={"a": {"$in": ["1", "2"]}},
575 | ) == ["3"]
576 |
577 | with pytest.raises(vecs.exc.FilterError):
578 | bar.query(
579 | data=[0, 0, 0, 0],
580 | limit=3,
581 | filters={"a": {"$in": 1}}, # error, value should be a list
582 | )
583 |
584 | with pytest.raises(vecs.exc.FilterError):
585 | bar.query(
586 | data=[0, 0, 0, 0],
587 | limit=3,
588 | filters={"a": {"$in": [1, [2]]}}, # error, element must be scalar
589 | )
590 |
591 |
592 | def test_filters_contains(client: vecs.Client) -> None:
593 | bar = client.get_or_create_collection(name="bar", dimension=4)
594 |
595 | records = [
596 | ("0", [0, 0, 0, 0], {"a": 1, "b": 2}),
597 | ("1", [1, 0, 0, 0], {"a": [1, 2, 3]}),
598 | ("2", [1, 1, 0, 0], {"a": {"1": "2", "x": "y"}}),
599 | ("3", [0, 0, 0, 0], {"a": ["1"]}),
600 | ("4", [1, 0, 0, 0], {"a": [4, 3, 2, 1]}),
601 | ("5", [1, 0, 0, 0], {"a": [2]}),
602 | ]
603 |
604 | bar.upsert(records)
605 | bar.create_index()
606 |
607 | # Test $contains operator for int value
608 | assert bar.query(
609 | data=[0, 0, 0, 0],
610 | limit=3,
611 | filters={"a": {"$contains": 1}},
612 | ) == ["1", "4"]
613 |
614 | # Test $contains operator for string value. Strings treated differently than ints
615 | assert bar.query(
616 | data=[0, 0, 0, 0],
617 | limit=3,
618 | filters={"a": {"$contains": "1"}},
619 | ) == ["3"]
620 |
621 | # Test $contains operator for non-existent value
622 | assert (
623 | bar.query(
624 | data=[0, 0, 0, 0],
625 | limit=3,
626 | filters={"a": {"$contains": 5}},
627 | )
628 | == []
629 | )
630 |
631 | # Test $contains requires a scalar value
632 | with pytest.raises(vecs.exc.FilterError):
633 | bar.query(
634 | data=[1, 0, 0, 0],
635 | limit=3,
636 | filters={"a": {"$contains": [1, 2, 3]}},
637 | )
638 |
639 | with pytest.raises(vecs.exc.FilterError):
640 | bar.query(
641 | data=[1, 0, 0, 0],
642 | limit=3,
643 | filters={"a": {"$contains": {"a": 1}}},
644 | )
645 |
646 |
647 | def test_access_index(client: vecs.Client) -> None:
648 | dim = 4
649 | bar = client.get_or_create_collection(name="bar", dimension=dim)
650 | assert bar.index is None
651 |
652 |
653 | def test_create_index(client: vecs.Client) -> None:
654 | dim = 4
655 | bar = client.get_or_create_collection(name="bar", dimension=dim)
656 |
657 | bar.create_index()
658 |
659 | assert bar.index is not None
660 |
661 | with pytest.raises(vecs.exc.ArgError):
662 | bar.create_index(replace=False)
663 |
664 | with pytest.raises(vecs.exc.ArgError):
665 | bar.create_index(method="does not exist")
666 |
667 | with pytest.raises(vecs.exc.ArgError):
668 | bar.create_index(measure="does not exist")
669 |
670 | bar.query(
671 | data=[1, 2, 3, 4],
672 | limit=1,
673 | measure="cosine_distance",
674 | )
675 |
676 |
677 | def test_ivfflat(client: vecs.Client) -> None:
678 | dim = 4
679 | bar = client.get_or_create_collection(name="bar", dimension=dim)
680 | bar.upsert([("a", [1, 2, 3, 4], {})])
681 |
682 | bar.create_index(method="ivfflat") # type: ignore
683 | results = bar.query(data=[1, 2, 3, 4], limit=1, probes=50)
684 | assert len(results) == 1
685 |
686 | bar.create_index(method=IndexMethod.ivfflat, replace=True) # type: ignore
687 | results = bar.query(
688 | data=[1, 2, 3, 4],
689 | limit=1,
690 | )
691 | assert len(results) == 1
692 |
693 |
694 | def test_hnsw(client: vecs.Client) -> None:
695 | dim = 4
696 | bar = client.get_or_create_collection(name="bar", dimension=dim)
697 | bar.upsert([("a", [1, 2, 3, 4], {})])
698 |
699 | bar.create_index(method="hnsw") # type: ignore
700 | results = bar.query(
701 | data=[1, 2, 3, 4],
702 | limit=1,
703 | )
704 | assert len(results) == 1
705 |
706 | bar.create_index(method=IndexMethod.hnsw, replace=True) # type: ignore
707 | results = bar.query(data=[1, 2, 3, 4], limit=1, ef_search=50)
708 | assert len(results) == 1
709 |
710 |
711 | def test_index_build_args(client: vecs.Client) -> None:
712 | dim = 4
713 | bar = client.get_or_create_collection(name="bar", dimension=dim)
714 | bar.upsert([("a", [1, 2, 3, 4], {})])
715 |
716 | # Test that default value for nlists is used in absence of index build args
717 | bar.create_index(method="ivfflat")
718 | [nlists] = [i for i in bar.index.split("_") if i.startswith("nl")]
719 | assert int(nlists.strip("nl")) == 30
720 |
721 | # Test nlists is honored when supplied
722 | bar.create_index(
723 | method=IndexMethod.ivfflat,
724 | index_arguments=IndexArgsIVFFlat(n_lists=123),
725 | replace=True,
726 | )
727 | [nlists] = [i for i in bar.index.split("_") if i.startswith("nl")]
728 | assert int(nlists.strip("nl")) == 123
729 |
730 | # Test that default values for m and ef_construction are used in absence of
731 | # index build args
732 | bar.create_index(method="hnsw", replace=True)
733 | [m] = [i for i in bar.index.split("_") if i.startswith("m")]
734 | [ef_construction] = [i for i in bar.index.split("_") if i.startswith("efc")]
735 | assert int(m.strip("m")) == 16
736 | assert int(ef_construction.strip("efc")) == 64
737 |
738 | # Test m and ef_construction is honored when supplied
739 | bar.create_index(
740 | method="hnsw",
741 | index_arguments=IndexArgsHNSW(m=8, ef_construction=123),
742 | replace=True,
743 | )
744 | [m] = [i for i in bar.index.split("_") if i.startswith("m")]
745 | [ef_construction] = [i for i in bar.index.split("_") if i.startswith("efc")]
746 | assert int(m.strip("m")) == 8
747 | assert int(ef_construction.strip("efc")) == 123
748 |
749 | # Test m is honored and ef_construction is default when _only_ m is supplied
750 | bar.create_index(method="hnsw", index_arguments=IndexArgsHNSW(m=8), replace=True)
751 | [m] = [i for i in bar.index.split("_") if i.startswith("m")]
752 | [ef_construction] = [i for i in bar.index.split("_") if i.startswith("efc")]
753 | assert int(m.strip("m")) == 8
754 | assert int(ef_construction.strip("efc")) == 64
755 |
756 | # Test m is default and ef_construction is honoured when _only_
757 | # ef_construction is supplied
758 | bar.create_index(
759 | method="hnsw", index_arguments=IndexArgsHNSW(ef_construction=123), replace=True
760 | )
761 | [m] = [i for i in bar.index.split("_") if i.startswith("m")]
762 | [ef_construction] = [i for i in bar.index.split("_") if i.startswith("efc")]
763 | assert int(m.strip("m")) == 16
764 | assert int(ef_construction.strip("efc")) == 123
765 |
766 | # Test that exception is raised when index build args don't match
767 | # the requested index type
768 | with pytest.raises(vecs.exc.ArgError):
769 | bar.create_index(
770 | method=IndexMethod.ivfflat, index_arguments=IndexArgsHNSW(), replace=True
771 | )
772 | with pytest.raises(vecs.exc.ArgError):
773 | bar.create_index(
774 | method=IndexMethod.hnsw,
775 | index_arguments=IndexArgsIVFFlat(n_lists=123),
776 | replace=True,
777 | )
778 |
779 | # Test that excpetion is raised index build args are supplied by the
780 | # IndexMethod.auto index is specified
781 | with pytest.raises(vecs.exc.ArgError):
782 | bar.create_index(
783 | method=IndexMethod.auto,
784 | index_arguments=IndexArgsIVFFlat(n_lists=123),
785 | replace=True,
786 | )
787 | with pytest.raises(vecs.exc.ArgError):
788 | bar.create_index(
789 | method=IndexMethod.auto,
790 | index_arguments=IndexArgsHNSW(),
791 | replace=True,
792 | )
793 |
794 |
795 | def test_cosine_index_query(client: vecs.Client) -> None:
796 | dim = 4
797 | bar = client.get_or_create_collection(name="bar", dimension=dim)
798 | bar.upsert([("a", [1, 2, 3, 4], {})])
799 | bar.create_index(measure=vecs.IndexMeasure.cosine_distance)
800 | results = bar.query(
801 | data=[1, 2, 3, 4],
802 | limit=1,
803 | measure="cosine_distance",
804 | )
805 | assert len(results) == 1
806 |
807 |
808 | def test_l2_index_query(client: vecs.Client) -> None:
809 | dim = 4
810 | bar = client.get_or_create_collection(name="bar", dimension=dim)
811 | bar.upsert([("a", [1, 2, 3, 4], {})])
812 | bar.create_index(measure=vecs.IndexMeasure.l2_distance)
813 | results = bar.query(
814 | data=[1, 2, 3, 4],
815 | limit=1,
816 | measure="l2_distance",
817 | )
818 | assert len(results) == 1
819 |
820 |
821 | def test_l1_index_query(client: vecs.Client) -> None:
822 | dim = 4
823 | bar = client.get_or_create_collection(name="bar", dimension=dim)
824 | bar.upsert([("a", [1, 2, 3, 4], {})])
825 | bar.create_index(measure=vecs.IndexMeasure.l1_distance)
826 | results = bar.query(
827 | data=[1, 2, 3, 4],
828 | limit=1,
829 | measure="l1_distance",
830 | )
831 | assert len(results) == 1
832 |
833 |
834 | def test_max_inner_product_index_query(client: vecs.Client) -> None:
835 | dim = 4
836 | bar = client.get_or_create_collection(name="bar", dimension=dim)
837 | bar.upsert([("a", [1, 2, 3, 4], {})])
838 | bar.create_index(measure=vecs.IndexMeasure.max_inner_product)
839 | results = bar.query(
840 | data=[1, 2, 3, 4],
841 | limit=1,
842 | measure="max_inner_product",
843 | )
844 | assert len(results) == 1
845 |
846 |
847 | def test_mismatch_measure(client: vecs.Client) -> None:
848 | dim = 4
849 | bar = client.get_or_create_collection(name="bar", dimension=dim)
850 | bar.upsert([("a", [1, 2, 3, 4], {})])
851 | bar.create_index(measure=vecs.IndexMeasure.max_inner_product)
852 | with pytest.warns(UserWarning):
853 | results = bar.query(
854 | data=[1, 2, 3, 4],
855 | limit=1,
856 | # wrong measure
857 | measure="cosine_distance",
858 | )
859 | assert len(results) == 1
860 |
861 |
862 | def test_is_indexed_for_measure(client: vecs.Client) -> None:
863 | bar = client.get_or_create_collection(name="bar", dimension=4)
864 |
865 | bar.create_index(measure=vecs.IndexMeasure.max_inner_product)
866 | assert not bar.is_indexed_for_measure("invalid") # type: ignore
867 | assert bar.is_indexed_for_measure(vecs.IndexMeasure.max_inner_product)
868 | assert not bar.is_indexed_for_measure(vecs.IndexMeasure.cosine_distance)
869 |
870 | bar.create_index(measure=vecs.IndexMeasure.cosine_distance, replace=True)
871 | assert bar.is_indexed_for_measure(vecs.IndexMeasure.cosine_distance)
872 |
873 |
874 | def test_failover_ivfflat(client: vecs.Client) -> None:
875 | """Test that index fails over to ivfflat on 0.4.0
876 | This is already covered by CI's test matrix but it is convenient for faster feedback
877 | to include it when running on the latest version of pgvector
878 | """
879 | client.vector_version = "0.4.1"
880 | dim = 4
881 | bar = client.get_or_create_collection(name="bar", dimension=dim)
882 | bar.upsert([("a", [1, 2, 3, 4], {})])
883 | # this executes an otherwise uncovered line of code that selects ivfflat when mode is 'auto'
884 | # and hnsw is unavailable
885 | bar.create_index(method=IndexMethod.auto)
886 |
887 |
888 | def test_hnsw_unavailable_error(client: vecs.Client) -> None:
889 | """Test that index fails over to ivfflat on 0.4.0
890 | This is already covered by CI's test matrix but it is convenient for faster feedback
891 | to include it when running on the latest version of pgvector
892 | """
893 | client.vector_version = "0.4.1"
894 | dim = 4
895 | bar = client.get_or_create_collection(name="bar", dimension=dim)
896 | with pytest.raises(ArgError):
897 | bar.create_index(method=IndexMethod.hnsw)
898 |
--------------------------------------------------------------------------------
/src/tests/test_issue_90.py:
--------------------------------------------------------------------------------
1 | import vecs
2 |
3 |
4 | def test_issue_90_multiple_index_support(client: vecs.Client) -> None:
5 | # Create a collection
6 | col1 = client.get_or_create_collection(name="col1", dimension=3)
7 |
8 | # Upsert some records
9 | col1.upsert(
10 | records=[
11 | (
12 | "vec0", # the vector's identifier
13 | [0.1, 0.2, 0.3], # the vector. list or np.array
14 | {"year": 1973}, # associated metadata
15 | ),
16 | ("vec1", [0.7, 0.8, 0.9], {"year": 2012}),
17 | ]
18 | )
19 |
20 | # Creat an index on the first collection
21 | col1.create_index()
22 |
23 | # Create a second collection
24 | col2 = client.get_or_create_collection(name="col2", dimension=3)
25 |
26 | # Create an index on the second collection
27 | col2.create_index()
28 |
29 | assert col1.index is not None
30 | assert col2.index is not None
31 |
32 | assert col1.index != col2.index
33 |
--------------------------------------------------------------------------------
/src/vecs/__init__.py:
--------------------------------------------------------------------------------
1 | from vecs import exc
2 | from vecs.client import Client
3 | from vecs.collection import (
4 | Collection,
5 | IndexArgsHNSW,
6 | IndexArgsIVFFlat,
7 | IndexMeasure,
8 | IndexMethod,
9 | )
10 |
11 | __project__ = "vecs"
12 | __version__ = "0.4.5"
13 |
14 |
15 | __all__ = [
16 | "IndexArgsIVFFlat",
17 | "IndexArgsHNSW",
18 | "IndexMethod",
19 | "IndexMeasure",
20 | "Collection",
21 | "Client",
22 | "exc",
23 | ]
24 |
25 |
26 | def create_client(connection_string: str) -> Client:
27 | """Creates a client from a Postgres connection string"""
28 | return Client(connection_string)
29 |
--------------------------------------------------------------------------------
/src/vecs/adapter/__init__.py:
--------------------------------------------------------------------------------
1 | from .base import Adapter, AdapterContext, AdapterStep
2 | from .markdown import MarkdownChunker
3 | from .noop import NoOp
4 | from .text import ParagraphChunker, TextEmbedding, TextEmbeddingModel
5 |
6 | __all__ = [
7 | "Adapter",
8 | "AdapterContext",
9 | "AdapterStep",
10 | "NoOp",
11 | "ParagraphChunker",
12 | "TextEmbedding",
13 | "TextEmbeddingModel",
14 | "MarkdownChunker",
15 | ]
16 |
--------------------------------------------------------------------------------
/src/vecs/adapter/base.py:
--------------------------------------------------------------------------------
1 | """
2 | The `vecs.experimental.adapter.base` module provides abstract classes and utilities
3 | for creating and handling adapters in vecs. Adapters allow users to interact with
4 | a collection using media types other than vectors.
5 |
6 | All public classes, enums, and functions are re-exported by `vecs.adapters` module.
7 | """
8 |
9 | from abc import ABC, abstractmethod
10 | from enum import Enum
11 | from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple
12 |
13 | from vecs.exc import ArgError
14 |
15 |
16 | class AdapterContext(str, Enum):
17 | """
18 | An enum representing the different contexts in which a Pipeline
19 | will be invoked.
20 |
21 | Attributes:
22 | upsert (str): The Collection.upsert method
23 | query (str): The Collection.query method
24 | """
25 |
26 | upsert = "upsert"
27 | query = "query"
28 |
29 |
30 | class AdapterStep(ABC):
31 | """
32 | Abstract class representing a step in the adapter pipeline.
33 |
34 | Each adapter step should adapt a user media into a tuple of:
35 | - id (str)
36 | - media (unknown type)
37 | - metadata (dict)
38 |
39 | If the user provides id or metadata, default production is overridden.
40 | """
41 |
42 | @property
43 | def exported_dimension(self) -> Optional[int]:
44 | """
45 | Property that should be overridden by subclasses to provide the output dimension
46 | of the adapter step.
47 | """
48 | return None
49 |
50 | @abstractmethod
51 | def __call__(
52 | self,
53 | records: Iterable[Tuple[str, Any, Optional[Dict]]],
54 | adapter_context: AdapterContext,
55 | ) -> Generator[Tuple[str, Any, Dict], None, None]:
56 | """
57 | Abstract method that should be overridden by subclasses to handle each record.
58 | """
59 |
60 |
61 | class Adapter:
62 | """
63 | Class representing a sequence of AdapterStep instances forming a pipeline.
64 | """
65 |
66 | def __init__(self, steps: List[AdapterStep]):
67 | """
68 | Initialize an Adapter instance with a list of AdapterStep instances.
69 |
70 | Args:
71 | steps: List of AdapterStep instances.
72 |
73 | Raises:
74 | ArgError: Raised if the steps list is empty.
75 | """
76 | self.steps = steps
77 | if len(steps) < 1:
78 | raise ArgError("Adapter must contain at least 1 step")
79 |
80 | @property
81 | def exported_dimension(self) -> Optional[int]:
82 | """
83 | The output dimension of the adapter. Returns the exported dimension of the last
84 | AdapterStep that provides one (from end to start of the steps list).
85 | """
86 | for step in reversed(self.steps):
87 | step_dim = step.exported_dimension
88 | if step_dim is not None:
89 | return step_dim
90 | return None
91 |
92 | def __call__(
93 | self,
94 | records: Iterable[Tuple[str, Any, Optional[Dict]]],
95 | adapter_context: AdapterContext,
96 | ) -> Generator[Tuple[str, Any, Dict], None, None]:
97 | """
98 | Invokes the adapter pipeline on an iterable of records.
99 |
100 | Args:
101 | records: Iterable of tuples each containing an id, a media and an optional dict.
102 | adapter_context: Context of the adapter.
103 |
104 | Yields:
105 | Tuples each containing an id, a media and a dict.
106 | """
107 | pipeline = records
108 | for step in self.steps:
109 | pipeline = step(pipeline, adapter_context)
110 |
111 | yield from pipeline # type: ignore
112 |
--------------------------------------------------------------------------------
/src/vecs/adapter/markdown.py:
--------------------------------------------------------------------------------
1 | import re
2 | from typing import Any, Dict, Generator, Iterable, Optional, Tuple
3 |
4 | from flupy import flu
5 |
6 | from .base import AdapterContext, AdapterStep
7 |
8 |
9 | class MarkdownChunker(AdapterStep):
10 | """
11 | MarkdownChunker is an AdapterStep that splits a markdown string into chunks where a heading signifies the start of a chunk, and yields each chunk as a separate record.
12 | """
13 |
14 | def __init__(self, *, skip_during_query: bool):
15 | """
16 | Initializes the MarkdownChunker adapter.
17 |
18 | Args:
19 | skip_during_query (bool): Whether to skip chunking during querying.
20 | """
21 | self.skip_during_query = skip_during_query
22 |
23 | @staticmethod
24 | def split_by_heading(md: str, max_tokens: int) -> Generator[str, None, None]:
25 | regex_split = r"^(#{1,6}\s+.+)$"
26 | headings = [
27 | match.span()[0]
28 | for match in re.finditer(regex_split, md, flags=re.MULTILINE)
29 | ]
30 |
31 | if headings == [] or headings[0] != 0:
32 | headings.insert(0, 0)
33 |
34 | sections = [md[i:j] for i, j in zip(headings, headings[1:] + [None])]
35 |
36 | for section in sections:
37 | chunks = flu(section.split(" ")).chunk(max_tokens)
38 |
39 | is_not_useless_chunk = lambda i: not i in ["", "\n", []]
40 |
41 | joined_chunks = filter(
42 | is_not_useless_chunk, [" ".join(chunk) for chunk in chunks]
43 | )
44 |
45 | for joined_chunk in joined_chunks:
46 | yield joined_chunk
47 |
48 | def __call__(
49 | self,
50 | records: Iterable[Tuple[str, Any, Optional[Dict]]],
51 | adapter_context: AdapterContext,
52 | max_tokens: int = 99999999,
53 | ) -> Generator[Tuple[str, Any, Dict], None, None]:
54 | """
55 | Splits each markdown string in the records into chunks where each heading starts a new chunk, and yields each chunk
56 | as a separate record. If the `skip_during_query` attribute is set to True,
57 | this step is skipped during querying.
58 |
59 | Args:
60 | records (Iterable[Tuple[str, Any, Optional[Dict]]]): Iterable of tuples each containing an id, a markdown string and an optional dict.
61 | adapter_context (AdapterContext): Context of the adapter.
62 | max_tokens (int): The maximum number of tokens per chunk
63 |
64 | Yields:
65 | Tuple[str, Any, Dict]: The id appended with chunk index, the chunk, and the metadata.
66 | """
67 | if max_tokens and max_tokens < 1:
68 | raise ValueError("max_tokens must be a nonzero positive integer")
69 |
70 | if adapter_context == AdapterContext("query") and self.skip_during_query:
71 | for id, markdown, metadata in records:
72 | yield (id, markdown, metadata or {})
73 | else:
74 | for id, markdown, metadata in records:
75 | headings = MarkdownChunker.split_by_heading(markdown, max_tokens)
76 | for heading_ix, heading in enumerate(headings):
77 | yield (
78 | f"{id}_head_{str(heading_ix).zfill(3)}",
79 | heading,
80 | metadata or {},
81 | )
82 |
--------------------------------------------------------------------------------
/src/vecs/adapter/noop.py:
--------------------------------------------------------------------------------
1 | """
2 | The `vecs.experimental.adapter.noop` module provides a default no-op (no operation) adapter
3 | that passes the inputs through without any modification. This can be useful when no specific
4 | adapter processing is required.
5 |
6 | All public classes, enums, and functions are re-exported by `vecs.adapters` module.
7 | """
8 |
9 | from typing import Any, Dict, Generator, Iterable, Optional, Tuple
10 |
11 | from .base import AdapterContext, AdapterStep
12 |
13 |
14 | class NoOp(AdapterStep):
15 | """
16 | NoOp is a no-operation AdapterStep. It is a default adapter that passes through
17 | the input records without any modifications.
18 | """
19 |
20 | def __init__(self, dimension: int):
21 | """
22 | Initializes the NoOp adapter with a dimension.
23 |
24 | Args:
25 | dimension (int): The dimension of the input vectors.
26 | """
27 | self._dimension = dimension
28 |
29 | @property
30 | def exported_dimension(self) -> Optional[int]:
31 | """
32 | Returns the dimension of the adapter.
33 |
34 | Returns:
35 | int: The dimension of the input vectors.
36 | """
37 | return self._dimension
38 |
39 | def __call__(
40 | self,
41 | records: Iterable[Tuple[str, Any, Optional[Dict]]],
42 | adapter_context: AdapterContext,
43 | ) -> Generator[Tuple[str, Any, Dict], None, None]:
44 | """
45 | Yields the input records without any modification.
46 |
47 | Args:
48 | records: Iterable of tuples each containing an id, a media and an optional dict.
49 | adapter_context: Context of the adapter.
50 |
51 | Yields:
52 | Tuple[str, Any, Dict]: The input record.
53 | """
54 | for id, media, metadata in records:
55 | yield (id, media, metadata or {})
56 |
--------------------------------------------------------------------------------
/src/vecs/adapter/text.py:
--------------------------------------------------------------------------------
1 | """
2 | The `vecs.experimental.adapter.text` module provides adapter steps specifically designed for
3 | handling text data. It provides two main classes, `TextEmbedding` and `ParagraphChunker`.
4 |
5 | All public classes, enums, and functions are re-exported by `vecs.adapters` module.
6 | """
7 | from typing import Any, Dict, Generator, Iterable, Literal, Optional, Tuple
8 |
9 | from flupy import flu
10 |
11 | from vecs.exc import MissingDependency
12 |
13 | from .base import AdapterContext, AdapterStep
14 |
15 | TextEmbeddingModel = Literal[
16 | "all-mpnet-base-v2",
17 | "multi-qa-mpnet-base-dot-v1",
18 | "all-distilroberta-v1",
19 | "all-MiniLM-L12-v2",
20 | "multi-qa-distilbert-cos-v1",
21 | "all-MiniLM-L6-v2",
22 | "multi-qa-MiniLM-L6-cos-v1",
23 | "paraphrase-multilingual-mpnet-base-v2",
24 | "paraphrase-albert-small-v2",
25 | "paraphrase-multilingual-MiniLM-L12-v2",
26 | "paraphrase-MiniLM-L3-v2",
27 | "distiluse-base-multilingual-cased-v1",
28 | "distiluse-base-multilingual-cased-v2",
29 | ]
30 |
31 |
32 | class TextEmbedding(AdapterStep):
33 | """
34 | TextEmbedding is an AdapterStep that converts text media into
35 | embeddings using a specified sentence transformers model.
36 | """
37 |
38 | def __init__(
39 | self,
40 | *,
41 | model: TextEmbeddingModel,
42 | batch_size: int = 8,
43 | use_auth_token: str = None,
44 | ):
45 | """
46 | Initializes the TextEmbedding adapter with a sentence transformers model.
47 |
48 | Args:
49 | model (TextEmbeddingModel): The sentence transformers model to use for embeddings.
50 | batch_size (int): The number of records to encode simultaneously.
51 | use_auth_token (str): The HuggingFace Hub auth token to use for private models.
52 |
53 | Raises:
54 | MissingDependency: If the sentence_transformers library is not installed.
55 | """
56 | try:
57 | from sentence_transformers import SentenceTransformer as ST
58 | except ImportError:
59 | raise MissingDependency(
60 | "Missing feature vecs[text_embedding]. Hint: `pip install 'vecs[text_embedding]'`"
61 | )
62 |
63 | self.model = ST(model, use_auth_token=use_auth_token)
64 | self._exported_dimension = self.model.get_sentence_embedding_dimension()
65 | self.batch_size = batch_size
66 |
67 | @property
68 | def exported_dimension(self) -> Optional[int]:
69 | """
70 | Returns the dimension of the embeddings produced by the sentence transformers model.
71 |
72 | Returns:
73 | int: The dimension of the embeddings.
74 | """
75 | return self._exported_dimension
76 |
77 | def __call__(
78 | self,
79 | records: Iterable[Tuple[str, Any, Optional[Dict]]],
80 | adapter_context: AdapterContext, # pyright: ignore
81 | ) -> Generator[Tuple[str, Any, Dict], None, None]:
82 | """
83 | Converts each media in the records to an embedding and yields the result.
84 |
85 | Args:
86 | records: Iterable of tuples each containing an id, a media and an optional dict.
87 | adapter_context: Context of the adapter.
88 |
89 | Yields:
90 | Tuple[str, Any, Dict]: The id, the embedding, and the metadata.
91 | """
92 | for batch in flu(records).chunk(self.batch_size):
93 | batch_records = [x for x in batch]
94 | media = [text for _, text, _ in batch_records]
95 |
96 | embeddings = self.model.encode(media, normalize_embeddings=True)
97 |
98 | for (id, _, metadata), embedding in zip(batch_records, embeddings): # type: ignore
99 | yield (id, embedding, metadata or {})
100 |
101 |
102 | class ParagraphChunker(AdapterStep):
103 | """
104 | ParagraphChunker is an AdapterStep that splits text media into
105 | paragraphs and yields each paragraph as a separate record.
106 | """
107 |
108 | def __init__(self, *, skip_during_query: bool):
109 | """
110 | Initializes the ParagraphChunker adapter.
111 |
112 | Args:
113 | skip_during_query (bool): Whether to skip chunking during querying.
114 | """
115 | self.skip_during_query = skip_during_query
116 |
117 | def __call__(
118 | self,
119 | records: Iterable[Tuple[str, Any, Optional[Dict]]],
120 | adapter_context: AdapterContext,
121 | ) -> Generator[Tuple[str, Any, Dict], None, None]:
122 | """
123 | Splits each media in the records into paragraphs and yields each paragraph
124 | as a separate record. If the `skip_during_query` attribute is set to True,
125 | this step is skipped during querying.
126 |
127 | Args:
128 | records (Iterable[Tuple[str, Any, Optional[Dict]]]): Iterable of tuples each containing an id, a media and an optional dict.
129 | adapter_context (AdapterContext): Context of the adapter.
130 |
131 | Yields:
132 | Tuple[str, Any, Dict]: The id appended with paragraph index, the paragraph, and the metadata.
133 | """
134 | if adapter_context == AdapterContext("query") and self.skip_during_query:
135 | for id, media, metadata in records:
136 | yield (id, media, metadata or {})
137 | else:
138 | for id, media, metadata in records:
139 | paragraphs = media.split("\n\n")
140 |
141 | for paragraph_ix, paragraph in enumerate(paragraphs):
142 | yield (
143 | f"{id}_para_{str(paragraph_ix).zfill(3)}",
144 | paragraph,
145 | metadata or {},
146 | )
147 |
--------------------------------------------------------------------------------
/src/vecs/client.py:
--------------------------------------------------------------------------------
1 | """
2 | Defines the 'Client' class
3 |
4 | Importing from the `vecs.client` directly is not supported.
5 | All public classes, enums, and functions are re-exported by the top level `vecs` module.
6 | """
7 |
8 | from __future__ import annotations
9 |
10 | from typing import TYPE_CHECKING, List, Optional
11 |
12 | from deprecated import deprecated
13 | from sqlalchemy import MetaData, create_engine, text
14 | from sqlalchemy.orm import sessionmaker
15 |
16 | from vecs.adapter import Adapter
17 | from vecs.exc import CollectionNotFound
18 |
19 | if TYPE_CHECKING:
20 | from vecs.collection import Collection
21 |
22 |
23 | class Client:
24 | """
25 | The `vecs.Client` class serves as an interface to a PostgreSQL database with pgvector support. It facilitates
26 | the creation, retrieval, listing and deletion of vector collections, while managing connections to the
27 | database.
28 |
29 | A `Client` instance represents a connection to a PostgreSQL database. This connection can be used to create
30 | and manipulate vector collections, where each collection is a group of vector records in a PostgreSQL table.
31 |
32 | The `vecs.Client` class can be also supports usage as a context manager to ensure the connection to the database
33 | is properly closed after operations, or used directly.
34 |
35 | Example usage:
36 |
37 | DB_CONNECTION = "postgresql://:@:/"
38 |
39 | with vecs.create_client(DB_CONNECTION) as vx:
40 | # do some work
41 | pass
42 |
43 | # OR
44 |
45 | vx = vecs.create_client(DB_CONNECTION)
46 | # do some work
47 | vx.disconnect()
48 | """
49 |
50 | def __init__(self, connection_string: str):
51 | """
52 | Initialize a Client instance.
53 |
54 | Args:
55 | connection_string (str): A string representing the database connection information.
56 |
57 | Returns:
58 | None
59 | """
60 | self.engine = create_engine(connection_string)
61 | self.meta = MetaData(schema="vecs")
62 | self.Session = sessionmaker(self.engine)
63 |
64 | with self.Session() as sess:
65 | with sess.begin():
66 | sess.execute(text("create schema if not exists vecs;"))
67 | sess.execute(text("create extension if not exists vector;"))
68 | self.vector_version: str = sess.execute(
69 | text(
70 | "select installed_version from pg_available_extensions where name = 'vector' limit 1;"
71 | )
72 | ).scalar_one()
73 |
74 | def _supports_hnsw(self):
75 | return (
76 | not self.vector_version.startswith("0.4")
77 | and not self.vector_version.startswith("0.3")
78 | and not self.vector_version.startswith("0.2")
79 | and not self.vector_version.startswith("0.1")
80 | and not self.vector_version.startswith("0.0")
81 | )
82 |
83 | def get_or_create_collection(
84 | self,
85 | name: str,
86 | *,
87 | dimension: Optional[int] = None,
88 | adapter: Optional[Adapter] = None,
89 | ) -> Collection:
90 | """
91 | Get a vector collection by name, or create it if no collection with
92 | *name* exists.
93 |
94 | Args:
95 | name (str): The name of the collection.
96 |
97 | Keyword Args:
98 | dimension (int): The dimensionality of the vectors in the collection.
99 | pipeline (int): The dimensionality of the vectors in the collection.
100 |
101 | Returns:
102 | Collection: The created collection.
103 |
104 | Raises:
105 | CollectionAlreadyExists: If a collection with the same name already exists
106 | """
107 | from vecs.collection import Collection
108 |
109 | adapter_dimension = adapter.exported_dimension if adapter else None
110 |
111 | collection = Collection(
112 | name=name,
113 | dimension=dimension or adapter_dimension, # type: ignore
114 | client=self,
115 | adapter=adapter,
116 | )
117 |
118 | return collection._create_if_not_exists()
119 |
120 | @deprecated("use Client.get_or_create_collection")
121 | def create_collection(self, name: str, dimension: int) -> Collection:
122 | """
123 | Create a new vector collection.
124 |
125 | Args:
126 | name (str): The name of the collection.
127 | dimension (int): The dimensionality of the vectors in the collection.
128 |
129 | Returns:
130 | Collection: The created collection.
131 |
132 | Raises:
133 | CollectionAlreadyExists: If a collection with the same name already exists
134 | """
135 | from vecs.collection import Collection
136 |
137 | return Collection(name, dimension, self)._create()
138 |
139 | @deprecated("use Client.get_or_create_collection")
140 | def get_collection(self, name: str) -> Collection:
141 | """
142 | Retrieve an existing vector collection.
143 |
144 | Args:
145 | name (str): The name of the collection.
146 |
147 | Returns:
148 | Collection: The retrieved collection.
149 |
150 | Raises:
151 | CollectionNotFound: If no collection with the given name exists.
152 | """
153 | from vecs.collection import Collection
154 |
155 | query = text(
156 | f"""
157 | select
158 | relname as table_name,
159 | atttypmod as embedding_dim
160 | from
161 | pg_class pc
162 | join pg_attribute pa
163 | on pc.oid = pa.attrelid
164 | where
165 | pc.relnamespace = 'vecs'::regnamespace
166 | and pc.relkind = 'r'
167 | and pa.attname = 'vec'
168 | and not pc.relname ^@ '_'
169 | and pc.relname = :name
170 | """
171 | ).bindparams(name=name)
172 | with self.Session() as sess:
173 | query_result = sess.execute(query).fetchone()
174 |
175 | if query_result is None:
176 | raise CollectionNotFound("No collection found with requested name")
177 |
178 | name, dimension = query_result
179 | return Collection(
180 | name,
181 | dimension,
182 | self,
183 | )
184 |
185 | def list_collections(self) -> List["Collection"]:
186 | """
187 | List all vector collections.
188 |
189 | Returns:
190 | list[Collection]: A list of all collections.
191 | """
192 | from vecs.collection import Collection
193 |
194 | return Collection._list_collections(self)
195 |
196 | def delete_collection(self, name: str) -> None:
197 | """
198 | Delete a vector collection.
199 |
200 | If no collection with requested name exists, does nothing.
201 |
202 | Args:
203 | name (str): The name of the collection.
204 |
205 | Returns:
206 | None
207 | """
208 | from vecs.collection import Collection
209 |
210 | Collection(name, -1, self)._drop()
211 | return
212 |
213 | def disconnect(self) -> None:
214 | """
215 | Disconnect the client from the database.
216 |
217 | Returns:
218 | None
219 | """
220 | self.engine.dispose()
221 | return
222 |
223 | def __enter__(self) -> "Client":
224 | """
225 | Enable use of the 'with' statement.
226 |
227 | Returns:
228 | Client: The current instance of the Client.
229 | """
230 |
231 | return self
232 |
233 | def __exit__(self, exc_type, exc_val, exc_tb):
234 | """
235 | Disconnect the client on exiting the 'with' statement context.
236 |
237 | Args:
238 | exc_type: The exception type, if any.
239 | exc_val: The exception value, if any.
240 | exc_tb: The traceback, if any.
241 |
242 | Returns:
243 | None
244 | """
245 | self.disconnect()
246 | return
247 |
--------------------------------------------------------------------------------
/src/vecs/collection.py:
--------------------------------------------------------------------------------
1 | """
2 | Defines the 'Collection' class
3 |
4 | Importing from the `vecs.collection` directly is not supported.
5 | All public classes, enums, and functions are re-exported by the top level `vecs` module.
6 | """
7 |
8 | from __future__ import annotations
9 |
10 | import math
11 | import uuid
12 | import warnings
13 | from dataclasses import dataclass
14 | from enum import Enum
15 | from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
16 |
17 | from flupy import flu
18 | from pgvector.sqlalchemy import Vector
19 | from sqlalchemy import (
20 | Column,
21 | MetaData,
22 | String,
23 | Table,
24 | and_,
25 | cast,
26 | delete,
27 | func,
28 | or_,
29 | select,
30 | text,
31 | )
32 | from sqlalchemy.dialects import postgresql
33 |
34 | from vecs.adapter import Adapter, AdapterContext, NoOp
35 | from vecs.exc import (
36 | ArgError,
37 | CollectionAlreadyExists,
38 | CollectionNotFound,
39 | FilterError,
40 | MismatchedDimension,
41 | Unreachable,
42 | )
43 |
44 | if TYPE_CHECKING:
45 | from vecs.client import Client
46 |
47 |
48 | MetadataValues = Union[str, int, float, bool, List[str]]
49 | Metadata = Dict[str, MetadataValues]
50 | Numeric = Union[int, float, complex]
51 | Record = Tuple[str, Iterable[Numeric], Metadata]
52 |
53 |
54 | class IndexMethod(str, Enum):
55 | """
56 | An enum representing the index methods available.
57 |
58 | This class currently only supports the 'ivfflat' method but may
59 | expand in the future.
60 |
61 | Attributes:
62 | auto (str): Automatically choose the best available index method.
63 | ivfflat (str): The ivfflat index method.
64 | hnsw (str): The hnsw index method.
65 | """
66 |
67 | auto = "auto"
68 | ivfflat = "ivfflat"
69 | hnsw = "hnsw"
70 |
71 |
72 | class IndexMeasure(str, Enum):
73 | """
74 | An enum representing the types of distance measures available for indexing.
75 |
76 | Attributes:
77 | cosine_distance (str): The cosine distance measure for indexing.
78 | l2_distance (str): The Euclidean (L2) distance measure for indexing.
79 | max_inner_product (str): The maximum inner product measure for indexing.
80 | """
81 |
82 | cosine_distance = "cosine_distance"
83 | l2_distance = "l2_distance"
84 | max_inner_product = "max_inner_product"
85 | l1_distance = "l1_distance"
86 |
87 |
88 | @dataclass
89 | class IndexArgsIVFFlat:
90 | """
91 | A class for arguments that can optionally be supplied to the index creation
92 | method when building an IVFFlat type index.
93 |
94 | Attributes:
95 | nlist (int): The number of IVF centroids that the index should use
96 | """
97 |
98 | n_lists: int
99 |
100 |
101 | @dataclass
102 | class IndexArgsHNSW:
103 | """
104 | A class for arguments that can optionally be supplied to the index creation
105 | method when building an HNSW type index.
106 |
107 | Ref: https://github.com/pgvector/pgvector#index-options
108 |
109 | Both attributes are Optional in case the user only wants to specify one and
110 | leave the other as default
111 |
112 | Attributes:
113 | m (int): Maximum number of connections per node per layer (default: 16)
114 | ef_construction (int): Size of the dynamic candidate list for
115 | constructing the graph (default: 64)
116 | """
117 |
118 | m: Optional[int] = 16
119 | ef_construction: Optional[int] = 64
120 |
121 |
122 | INDEX_MEASURE_TO_OPS = {
123 | # Maps the IndexMeasure enum options to the SQL ops string required by
124 | # the pgvector `create index` statement
125 | IndexMeasure.cosine_distance: "vector_cosine_ops",
126 | IndexMeasure.l2_distance: "vector_l2_ops",
127 | IndexMeasure.max_inner_product: "vector_ip_ops",
128 | IndexMeasure.l1_distance: "vector_l1_ops",
129 | }
130 |
131 | INDEX_MEASURE_TO_SQLA_ACC = {
132 | IndexMeasure.cosine_distance: lambda x: x.cosine_distance,
133 | IndexMeasure.l2_distance: lambda x: x.l2_distance,
134 | IndexMeasure.max_inner_product: lambda x: x.max_inner_product,
135 | IndexMeasure.l1_distance: lambda x: x.l1_distance,
136 | }
137 |
138 |
139 | class Collection:
140 | """
141 | The `vecs.Collection` class represents a collection of vectors within a PostgreSQL database with pgvector support.
142 | It provides methods to manage (create, delete, fetch, upsert), index, and perform similarity searches on these vector collections.
143 |
144 | The collections are stored in separate tables in the database, with each vector associated with an identifier and optional metadata.
145 |
146 | Example usage:
147 |
148 | with vecs.create_client(DB_CONNECTION) as vx:
149 | collection = vx.create_collection(name="docs", dimension=3)
150 | collection.upsert([("id1", [1, 1, 1], {"key": "value"})])
151 | # Further operations on 'collection'
152 |
153 | Public Attributes:
154 | name: The name of the vector collection.
155 | dimension: The dimension of vectors in the collection.
156 |
157 | Note: Some methods of this class can raise exceptions from the `vecs.exc` module if errors occur.
158 | """
159 |
160 | def __init__(
161 | self,
162 | name: str,
163 | dimension: int,
164 | client: Client,
165 | adapter: Optional[Adapter] = None,
166 | ):
167 | """
168 | Initializes a new instance of the `Collection` class.
169 |
170 | During expected use, developers initialize instances of `Collection` using the
171 | `vecs.Client` with `vecs.Client.create_collection(...)` rather than directly.
172 |
173 | Args:
174 | name (str): The name of the collection.
175 | dimension (int): The dimension of the vectors in the collection.
176 | client (Client): The client to use for interacting with the database.
177 | """
178 | self.client = client
179 | self.name = name
180 | self.dimension = dimension
181 | self.table = build_table(name, client.meta, dimension)
182 | self._index: Optional[str] = None
183 | self.adapter = adapter or Adapter(steps=[NoOp(dimension=dimension)])
184 |
185 | reported_dimensions = set(
186 | [
187 | x
188 | for x in [
189 | dimension,
190 | adapter.exported_dimension if adapter else None,
191 | ]
192 | if x is not None
193 | ]
194 | )
195 | if len(reported_dimensions) == 0:
196 | raise ArgError("One of dimension or adapter must provide a dimension")
197 | elif len(reported_dimensions) > 1:
198 | raise MismatchedDimension(
199 | "Dimensions reported by adapter, dimension, and collection do not match"
200 | )
201 |
202 | def __repr__(self):
203 | """
204 | Returns a string representation of the `Collection` instance.
205 |
206 | Returns:
207 | str: A string representation of the `Collection` instance.
208 | """
209 | return f'vecs.Collection(name="{self.name}", dimension={self.dimension})'
210 |
211 | def __len__(self) -> int:
212 | """
213 | Returns the number of vectors in the collection.
214 |
215 | Returns:
216 | int: The number of vectors in the collection.
217 | """
218 | with self.client.Session() as sess:
219 | with sess.begin():
220 | stmt = select(func.count()).select_from(self.table)
221 | return sess.execute(stmt).scalar() or 0
222 |
223 | def _create_if_not_exists(self):
224 | """
225 | PRIVATE
226 |
227 | Creates a new collection in the database if it doesn't already exist
228 |
229 | Returns:
230 | Collection: The found or created collection.
231 | """
232 | query = text(
233 | f"""
234 | select
235 | relname as table_name,
236 | atttypmod as embedding_dim
237 | from
238 | pg_class pc
239 | join pg_attribute pa
240 | on pc.oid = pa.attrelid
241 | where
242 | pc.relnamespace = 'vecs'::regnamespace
243 | and pc.relkind = 'r'
244 | and pa.attname = 'vec'
245 | and not pc.relname ^@ '_'
246 | and pc.relname = :name
247 | """
248 | ).bindparams(name=self.name)
249 | with self.client.Session() as sess:
250 | query_result = sess.execute(query).fetchone()
251 |
252 | if query_result:
253 | _, collection_dimension = query_result
254 | else:
255 | collection_dimension = None
256 |
257 | reported_dimensions = set(
258 | [x for x in [self.dimension, collection_dimension] if x is not None]
259 | )
260 | if len(reported_dimensions) > 1:
261 | raise MismatchedDimension(
262 | "Dimensions reported by adapter, dimension, and existing collection do not match"
263 | )
264 |
265 | if not collection_dimension:
266 | self.table.create(self.client.engine)
267 |
268 | return self
269 |
270 | def _create(self):
271 | """
272 | PRIVATE
273 |
274 | Creates a new collection in the database. Raises a `vecs.exc.CollectionAlreadyExists`
275 | exception if a collection with the specified name already exists.
276 |
277 | Returns:
278 | Collection: The newly created collection.
279 | """
280 |
281 | collection_exists = self.__class__._does_collection_exist(
282 | self.client, self.name
283 | )
284 | if collection_exists:
285 | raise CollectionAlreadyExists(
286 | "Collection with requested name already exists"
287 | )
288 | self.table.create(self.client.engine)
289 |
290 | unique_string = str(uuid.uuid4()).replace("-", "_")[0:7]
291 | with self.client.Session() as sess:
292 | sess.execute(
293 | text(
294 | f"""
295 | create index ix_meta_{unique_string}
296 | on vecs."{self.table.name}"
297 | using gin ( metadata jsonb_path_ops )
298 | """
299 | )
300 | )
301 | return self
302 |
303 | def _drop(self):
304 | """
305 | PRIVATE
306 |
307 | Deletes the collection from the database. Raises a `vecs.exc.CollectionNotFound`
308 | exception if no collection with the specified name exists.
309 |
310 | Returns:
311 | Collection: The deleted collection.
312 | """
313 | from sqlalchemy.schema import DropTable
314 |
315 | with self.client.Session() as sess:
316 | sess.execute(DropTable(self.table, if_exists=True))
317 | sess.commit()
318 |
319 | return self
320 |
321 | def upsert(
322 | self, records: Iterable[Tuple[str, Any, Metadata]], skip_adapter: bool = False
323 | ) -> None:
324 | """
325 | Inserts or updates *vectors* records in the collection.
326 |
327 | Args:
328 | records (Iterable[Tuple[str, Any, Metadata]]): An iterable of content to upsert.
329 | Each record is a tuple where:
330 | - the first element is a unique string identifier
331 | - the second element is an iterable of numeric values or relevant input type for the
332 | adapter assigned to the collection
333 | - the third element is metadata associated with the vector
334 |
335 | skip_adapter (bool): Should the adapter be skipped while upserting. i.e. if vectors are being
336 | provided, rather than a media type that needs to be transformed
337 | """
338 |
339 | chunk_size = 500
340 |
341 | if skip_adapter:
342 | pipeline = flu(records).chunk(chunk_size)
343 | else:
344 | # Construct a lazy pipeline of steps to transform and chunk user input
345 | pipeline = flu(self.adapter(records, AdapterContext("upsert"))).chunk(
346 | chunk_size
347 | )
348 |
349 | with self.client.Session() as sess:
350 | with sess.begin():
351 | for chunk in pipeline:
352 | stmt = postgresql.insert(self.table).values(chunk)
353 | stmt = stmt.on_conflict_do_update(
354 | index_elements=[self.table.c.id],
355 | set_=dict(
356 | vec=stmt.excluded.vec, metadata=stmt.excluded.metadata
357 | ),
358 | )
359 | sess.execute(stmt)
360 | return None
361 |
362 | def fetch(self, ids: Iterable[str]) -> List[Record]:
363 | """
364 | Fetches vectors from the collection by their identifiers.
365 |
366 | Args:
367 | ids (Iterable[str]): An iterable of vector identifiers.
368 |
369 | Returns:
370 | List[Record]: A list of the fetched vectors.
371 | """
372 | if isinstance(ids, str):
373 | raise ArgError("ids must be a list of strings")
374 |
375 | chunk_size = 12
376 | records = []
377 | with self.client.Session() as sess:
378 | with sess.begin():
379 | for id_chunk in flu(ids).chunk(chunk_size):
380 | stmt = select(self.table).where(self.table.c.id.in_(id_chunk))
381 | chunk_records = sess.execute(stmt)
382 | records.extend(chunk_records)
383 | return records
384 |
385 | def delete(
386 | self, ids: Optional[Iterable[str]] = None, filters: Optional[Metadata] = None
387 | ) -> List[str]:
388 | """
389 | Deletes vectors from the collection by matching filters or ids.
390 |
391 | Args:
392 | ids (Iterable[str], optional): An iterable of vector identifiers.
393 | filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
394 |
395 | Returns:
396 | List[str]: A list of the identifiers of the deleted vectors.
397 | """
398 | if ids is None and filters is None:
399 | raise ArgError("Either ids or filters must be provided.")
400 |
401 | if ids is not None and filters is not None:
402 | raise ArgError("Either ids or filters must be provided, not both.")
403 |
404 | if isinstance(ids, str):
405 | raise ArgError("ids must be a list of strings")
406 |
407 | ids = ids or []
408 | filters = filters or {}
409 | del_ids = []
410 |
411 | with self.client.Session() as sess:
412 | with sess.begin():
413 | if ids:
414 | for id_chunk in flu(ids).chunk(12):
415 | stmt = (
416 | delete(self.table)
417 | .where(self.table.c.id.in_(id_chunk))
418 | .returning(self.table.c.id)
419 | )
420 | del_ids.extend(sess.execute(stmt).scalars() or [])
421 |
422 | if filters:
423 | meta_filter = build_filters(self.table.c.metadata, filters)
424 | stmt = (
425 | delete(self.table).where(meta_filter).returning(self.table.c.id) # type: ignore
426 | )
427 | result = sess.execute(stmt).scalars()
428 | del_ids.extend(result.fetchall())
429 |
430 | return del_ids
431 |
432 | def __getitem__(self, items):
433 | """
434 | Fetches a vector from the collection by its identifier.
435 |
436 | Args:
437 | items (str): The identifier of the vector.
438 |
439 | Returns:
440 | Record: The fetched vector.
441 | """
442 | if not isinstance(items, str):
443 | raise ArgError("items must be a string id")
444 |
445 | row = self.fetch([items])
446 |
447 | if row == []:
448 | raise KeyError("no item found with requested id")
449 | return row[0]
450 |
451 | def query(
452 | self,
453 | data: Union[Iterable[Numeric], Any],
454 | limit: int = 10,
455 | filters: Optional[Dict] = None,
456 | measure: Union[IndexMeasure, str] = IndexMeasure.cosine_distance,
457 | include_value: bool = False,
458 | include_metadata: bool = False,
459 | include_vector: bool = False,
460 | *,
461 | probes: Optional[int] = None,
462 | ef_search: Optional[int] = None,
463 | skip_adapter: bool = False,
464 | ) -> Union[List[Record], List[str]]:
465 | """
466 | Executes a similarity search in the collection.
467 |
468 | The return type is dependent on arguments *include_value* and *include_metadata*
469 |
470 | Args:
471 | data (Any): The vector to use as the query.
472 | limit (int, optional): The maximum number of results to return. Defaults to 10.
473 | filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
474 | measure (Union[IndexMeasure, str], optional): The distance measure to use for the search. Defaults to 'cosine_distance'.
475 | include_value (bool, optional): Whether to include the distance value in the results. Defaults to False.
476 | include_metadata (bool, optional): Whether to include the metadata in the results. Defaults to False.
477 | probes (Optional[Int], optional): Number of ivfflat index lists to query. Higher increases accuracy but decreases speed
478 | ef_search (Optional[Int], optional): Size of the dynamic candidate list for HNSW index search. Higher increases accuracy but decreases speed
479 | skip_adapter (bool, optional): When True, skips any associated adapter and queries using a literal vector provided to *data*
480 |
481 | Returns:
482 | Union[List[Record], List[str]]: The result of the similarity search.
483 | """
484 |
485 | if probes is None:
486 | probes = 10
487 |
488 | if ef_search is None:
489 | ef_search = 40
490 |
491 | if not isinstance(probes, int):
492 | raise ArgError("probes must be an integer")
493 |
494 | if probes < 1:
495 | raise ArgError("probes must be >= 1")
496 |
497 | if limit > 1000:
498 | raise ArgError("limit must be <= 1000")
499 |
500 | # ValueError on bad input
501 | try:
502 | imeasure = IndexMeasure(measure)
503 | except ValueError:
504 | raise ArgError("Invalid index measure")
505 |
506 | if not self.is_indexed_for_measure(imeasure):
507 | warnings.warn(
508 | UserWarning(
509 | f"Query does not have a covering index for {imeasure}. See Collection.create_index"
510 | )
511 | )
512 |
513 | if skip_adapter:
514 | adapted_query = [("", data, {})]
515 | else:
516 | # Adapt the query using the pipeline
517 | adapted_query = [
518 | x
519 | for x in self.adapter(
520 | records=[("", data, {})], adapter_context=AdapterContext("query")
521 | )
522 | ]
523 |
524 | if len(adapted_query) != 1:
525 | raise ArgError("Failed to produce exactly one query vector from input")
526 |
527 | _, vec, _ = adapted_query[0]
528 |
529 | distance_lambda = INDEX_MEASURE_TO_SQLA_ACC.get(imeasure)
530 | if distance_lambda is None:
531 | # unreachable
532 | raise ArgError("invalid distance_measure") # pragma: no cover
533 |
534 | distance_clause = distance_lambda(self.table.c.vec)(vec)
535 |
536 | cols = [self.table.c.id]
537 |
538 | if include_value:
539 | cols.append(distance_clause)
540 |
541 | if include_vector:
542 | cols.append(self.table.c.vec)
543 |
544 | if include_metadata:
545 | cols.append(self.table.c.metadata)
546 |
547 | stmt = select(*cols)
548 | if filters:
549 | stmt = stmt.filter(
550 | build_filters(self.table.c.metadata, filters) # type: ignore
551 | )
552 |
553 | stmt = stmt.order_by(distance_clause)
554 | stmt = stmt.limit(limit)
555 |
556 | with self.client.Session() as sess:
557 | with sess.begin():
558 | # index ignored if greater than n_lists
559 | sess.execute(
560 | text("set local ivfflat.probes = :probes").bindparams(probes=probes)
561 | )
562 | if self.client._supports_hnsw():
563 | sess.execute(
564 | text("set local hnsw.ef_search = :ef_search").bindparams(
565 | ef_search=ef_search
566 | )
567 | )
568 | if len(cols) == 1:
569 | return [str(x) for x in sess.scalars(stmt).fetchall()]
570 | return sess.execute(stmt).fetchall() or []
571 |
572 | @classmethod
573 | def _list_collections(cls, client: "Client") -> List["Collection"]:
574 | """
575 | PRIVATE
576 |
577 | Retrieves all collections from the database.
578 |
579 | Args:
580 | client (Client): The database client.
581 |
582 | Returns:
583 | List[Collection]: A list of all existing collections.
584 | """
585 |
586 | query = text(
587 | """
588 | select
589 | relname as table_name,
590 | atttypmod as embedding_dim
591 | from
592 | pg_class pc
593 | join pg_attribute pa
594 | on pc.oid = pa.attrelid
595 | where
596 | pc.relnamespace = 'vecs'::regnamespace
597 | and pc.relkind = 'r'
598 | and pa.attname = 'vec'
599 | and not pc.relname ^@ '_'
600 | """
601 | )
602 | xc = []
603 | with client.Session() as sess:
604 | for name, dimension in sess.execute(query):
605 | existing_collection = cls(name, dimension, client)
606 | xc.append(existing_collection)
607 | return xc
608 |
609 | @classmethod
610 | def _does_collection_exist(cls, client: "Client", name: str) -> bool:
611 | """
612 | PRIVATE
613 |
614 | Checks if a collection with a given name exists within the database
615 |
616 | Args:
617 | client (Client): The database client.
618 | name (str): The name of the collection
619 |
620 | Returns:
621 | Exists: Whether the collection exists or not
622 | """
623 |
624 | try:
625 | client.get_collection(name)
626 | return True
627 | except CollectionNotFound:
628 | return False
629 |
630 | @property
631 | def index(self) -> Optional[str]:
632 | """
633 | PRIVATE
634 |
635 | Note:
636 | The `index` property is private and expected to undergo refactoring.
637 | Do not rely on it's output.
638 |
639 | Retrieves the SQL name of the collection's vector index, if it exists.
640 |
641 | Returns:
642 | Optional[str]: The name of the index, or None if no index exists.
643 | """
644 |
645 | if self._index is None:
646 | query = text(
647 | """
648 | select
649 | pi.relname as index_name
650 | from
651 | pg_class pi -- index info
652 | join pg_index i -- extend index info
653 | on pi.oid = i.indexrelid
654 | join pg_class pt -- owning table info
655 | on pt.oid = i.indrelid
656 | where
657 | pi.relnamespace = 'vecs'::regnamespace
658 | and pi.relname ilike 'ix_vector%'
659 | and pi.relkind = 'i'
660 | and pt.relname = :table_name
661 | """
662 | )
663 | with self.client.Session() as sess:
664 | ix_name = sess.execute(query, {"table_name": self.name}).scalar()
665 | self._index = ix_name
666 | return self._index
667 |
668 | def is_indexed_for_measure(self, measure: IndexMeasure):
669 | """
670 | Checks if the collection is indexed for a specific measure.
671 |
672 | Args:
673 | measure (IndexMeasure): The measure to check for.
674 |
675 | Returns:
676 | bool: True if the collection is indexed for the measure, False otherwise.
677 | """
678 |
679 | index_name = self.index
680 | if index_name is None:
681 | return False
682 |
683 | ops = INDEX_MEASURE_TO_OPS.get(measure)
684 | if ops is None:
685 | return False
686 |
687 | if ops in index_name:
688 | return True
689 |
690 | return False
691 |
692 | def create_index(
693 | self,
694 | measure: IndexMeasure = IndexMeasure.cosine_distance,
695 | method: IndexMethod = IndexMethod.auto,
696 | index_arguments: Optional[Union[IndexArgsIVFFlat, IndexArgsHNSW]] = None,
697 | replace=True,
698 | ) -> None:
699 | """
700 | Creates an index for the collection.
701 |
702 | Note:
703 | When `vecs` creates an index on a pgvector column in PostgreSQL, it uses a multi-step
704 | process that enables performant indexes to be built for large collections with low end
705 | database hardware.
706 |
707 | Those steps are:
708 |
709 | - Creates a new table with a different name
710 | - Randomly selects records from the existing table
711 | - Inserts the random records from the existing table into the new table
712 | - Creates the requested vector index on the new table
713 | - Upserts all data from the existing table into the new table
714 | - Drops the existing table
715 | - Renames the new table to the existing tables name
716 |
717 | If you create dependencies (like views) on the table that underpins
718 | a `vecs.Collection` the `create_index` step may require you to drop those dependencies before
719 | it will succeed.
720 |
721 | Args:
722 | measure (IndexMeasure, optional): The measure to index for. Defaults to 'cosine_distance'.
723 | method (IndexMethod, optional): The indexing method to use. Defaults to 'auto'.
724 | index_arguments: (IndexArgsIVFFlat | IndexArgsHNSW, optional): Index type specific arguments
725 | replace (bool, optional): Whether to replace the existing index. Defaults to True.
726 |
727 | Raises:
728 | ArgError: If an invalid index method is used, or if *replace* is False and an index already exists.
729 | """
730 |
731 | if method not in (IndexMethod.ivfflat, IndexMethod.hnsw, IndexMethod.auto):
732 | raise ArgError("invalid index method")
733 |
734 | if index_arguments:
735 | # Disallow case where user submits index arguments but uses the
736 | # IndexMethod.auto index (index build arguments should only be
737 | # used with a specific index)
738 | if method == IndexMethod.auto:
739 | raise ArgError(
740 | "Index build parameters are not allowed when using the IndexMethod.auto index."
741 | )
742 | # Disallow case where user specifies one index type but submits
743 | # index build arguments for the other index type
744 | if (
745 | isinstance(index_arguments, IndexArgsHNSW)
746 | and method != IndexMethod.hnsw
747 | ) or (
748 | isinstance(index_arguments, IndexArgsIVFFlat)
749 | and method != IndexMethod.ivfflat
750 | ):
751 | raise ArgError(
752 | f"{index_arguments.__class__.__name__} build parameters were supplied but {method} index was specified."
753 | )
754 |
755 | if method == IndexMethod.auto:
756 | if self.client._supports_hnsw():
757 | method = IndexMethod.hnsw
758 | else:
759 | method = IndexMethod.ivfflat
760 |
761 | if method == IndexMethod.hnsw and not self.client._supports_hnsw():
762 | raise ArgError(
763 | "HNSW Unavailable. Upgrade your pgvector installation to > 0.5.0 to enable HNSW support"
764 | )
765 |
766 | ops = INDEX_MEASURE_TO_OPS.get(measure)
767 | if ops is None:
768 | raise ArgError("Unknown index measure")
769 |
770 | unique_string = str(uuid.uuid4()).replace("-", "_")[0:7]
771 |
772 | with self.client.Session() as sess:
773 | with sess.begin():
774 | if self.index is not None:
775 | if replace:
776 | sess.execute(text(f'drop index vecs."{self.index}";'))
777 | self._index = None
778 | else:
779 | raise ArgError("replace is set to False but an index exists")
780 |
781 | if method == IndexMethod.ivfflat:
782 | if not index_arguments:
783 | n_records: int = sess.execute(func.count(self.table.c.id)).scalar() # type: ignore
784 |
785 | n_lists = (
786 | int(max(n_records / 1000, 30))
787 | if n_records < 1_000_000
788 | else int(math.sqrt(n_records))
789 | )
790 | else:
791 | # The following mypy error is ignored because mypy
792 | # complains that `index_arguments` is typed as a union
793 | # of IndexArgsIVFFlat and IndexArgsHNSW types,
794 | # which both don't necessarily contain the `n_lists`
795 | # parameter, however we have validated that the
796 | # correct type is being used above.
797 | n_lists = index_arguments.n_lists # type: ignore
798 |
799 | sess.execute(
800 | text(
801 | f"""
802 | create index ix_{ops}_ivfflat_nl{n_lists}_{unique_string}
803 | on vecs."{self.table.name}"
804 | using ivfflat (vec {ops}) with (lists={n_lists})
805 | """
806 | )
807 | )
808 |
809 | if method == IndexMethod.hnsw:
810 | if not index_arguments:
811 | index_arguments = IndexArgsHNSW()
812 |
813 | # See above for explanation of why the following lines
814 | # are ignored
815 | m = index_arguments.m # type: ignore
816 | ef_construction = index_arguments.ef_construction # type: ignore
817 |
818 | sess.execute(
819 | text(
820 | f"""
821 | create index ix_{ops}_hnsw_m{m}_efc{ef_construction}_{unique_string}
822 | on vecs."{self.table.name}"
823 | using hnsw (vec {ops}) WITH (m={m}, ef_construction={ef_construction});
824 | """
825 | )
826 | )
827 |
828 | return None
829 |
830 |
831 | def build_filters(json_col: Column, filters: Dict):
832 | """
833 | PRIVATE
834 |
835 | Builds filters for SQL query based on provided dictionary.
836 |
837 | Args:
838 | json_col (Column): The column in the database table.
839 | filters (Dict): The dictionary specifying filter conditions.
840 |
841 | Raises:
842 | FilterError: If filter conditions are not correctly formatted.
843 |
844 | Returns:
845 | The filter clause for the SQL query.
846 | """
847 |
848 | if not isinstance(filters, dict):
849 | raise FilterError("filters must be a dict")
850 |
851 | if len(filters) > 1:
852 | raise FilterError("max 1 entry per filter")
853 |
854 | for key, value in filters.items():
855 | if not isinstance(key, str):
856 | raise FilterError("*filters* keys must be strings")
857 |
858 | if key in ("$and", "$or"):
859 | if not isinstance(value, list):
860 | raise FilterError(
861 | "$and/$or filters must have associated list of conditions"
862 | )
863 |
864 | if key == "$and":
865 | return and_(*[build_filters(json_col, subcond) for subcond in value])
866 |
867 | if key == "$or":
868 | return or_(*[build_filters(json_col, subcond) for subcond in value])
869 |
870 | raise Unreachable()
871 |
872 | if isinstance(value, dict):
873 | if len(value) > 1:
874 | raise FilterError("only one operator permitted")
875 | for operator, clause in value.items():
876 | if operator not in (
877 | "$eq",
878 | "$ne",
879 | "$lt",
880 | "$lte",
881 | "$gt",
882 | "$gte",
883 | "$in",
884 | "$contains",
885 | ):
886 | raise FilterError("unknown operator")
887 |
888 | # equality of singular values can take advantage of the metadata index
889 | # using containment operator. Containment can not be used to test equality
890 | # of lists or dicts so we restrict to single values with a __len__ check.
891 | if operator == "$eq" and not hasattr(clause, "__len__"):
892 | contains_value = cast({key: clause}, postgresql.JSONB)
893 | return json_col.op("@>")(contains_value)
894 |
895 | if operator == "$in":
896 | if not isinstance(clause, list):
897 | raise FilterError("argument to $in filter must be a list")
898 |
899 | for elem in clause:
900 | if not isinstance(elem, (int, str, float)):
901 | raise FilterError(
902 | "argument to $in filter must be a list of scalars"
903 | )
904 |
905 | # cast the array of scalars to a postgres array of jsonb so we can
906 | # directly compare json types in the query
907 | contains_value = [cast(elem, postgresql.JSONB) for elem in clause]
908 | return json_col.op("->")(key).in_(contains_value)
909 |
910 | matches_value = cast(clause, postgresql.JSONB)
911 |
912 | # @> in Postgres is heavily overloaded.
913 | # By default, it will return True for
914 | #
915 | # scalar in array
916 | # '[1, 2, 3]'::jsonb @> '1'::jsonb -- true#
917 | # equality:
918 | # '1'::jsonb @> '1'::jsonb -- true
919 | # key value pair in object
920 | # '{"a": 1, "b": 2}'::jsonb @> '{"a": 1}'::jsonb -- true
921 | #
922 | # At this time we only want to allow "scalar in array" so
923 | # we assert that the clause is a scalar and the target metadata
924 | # is an array
925 | if operator == "$contains":
926 | if not isinstance(clause, (int, str, float)):
927 | raise FilterError(
928 | "argument to $contains filter must be a scalar"
929 | )
930 |
931 | return and_(
932 | json_col.op("->")(key).contains(matches_value),
933 | func.jsonb_typeof(json_col.op("->")(key)) == "array",
934 | )
935 |
936 | # handles non-singular values
937 | if operator == "$eq":
938 | return json_col.op("->")(key) == matches_value
939 |
940 | elif operator == "$ne":
941 | return json_col.op("->")(key) != matches_value
942 |
943 | elif operator == "$lt":
944 | return json_col.op("->")(key) < matches_value
945 |
946 | elif operator == "$lte":
947 | return json_col.op("->")(key) <= matches_value
948 |
949 | elif operator == "$gt":
950 | return json_col.op("->")(key) > matches_value
951 |
952 | elif operator == "$gte":
953 | return json_col.op("->")(key) >= matches_value
954 |
955 | else:
956 | raise Unreachable()
957 |
958 |
959 | def build_table(name: str, meta: MetaData, dimension: int) -> Table:
960 | """
961 | PRIVATE
962 |
963 | Builds a SQLAlchemy model underpinning a `vecs.Collection`.
964 |
965 | Args:
966 | name (str): The name of the table.
967 | meta (MetaData): MetaData instance associated with the SQL database.
968 | dimension: The dimension of the vectors in the collection.
969 |
970 | Returns:
971 | Table: The constructed SQL table.
972 | """
973 | return Table(
974 | name,
975 | meta,
976 | Column("id", String, primary_key=True),
977 | Column("vec", Vector(dimension), nullable=False),
978 | Column(
979 | "metadata",
980 | postgresql.JSONB,
981 | server_default=text("'{}'::jsonb"),
982 | nullable=False,
983 | ),
984 | extend_existing=True,
985 | )
986 |
--------------------------------------------------------------------------------
/src/vecs/exc.py:
--------------------------------------------------------------------------------
1 | __all__ = [
2 | "VecsException",
3 | "CollectionAlreadyExists",
4 | "CollectionNotFound",
5 | "ArgError",
6 | "FilterError",
7 | "IndexNotFound",
8 | "Unreachable",
9 | ]
10 |
11 |
12 | class VecsException(Exception):
13 | """
14 | Base exception class for the 'vecs' package.
15 | All custom exceptions in the 'vecs' package should derive from this class.
16 | """
17 |
18 | ...
19 |
20 |
21 | class CollectionAlreadyExists(VecsException):
22 | """
23 | Exception raised when attempting to create a collection that already exists.
24 | """
25 |
26 | ...
27 |
28 |
29 | class CollectionNotFound(VecsException):
30 | """
31 | Exception raised when attempting to access or manipulate a collection that does not exist.
32 | """
33 |
34 | ...
35 |
36 |
37 | class ArgError(VecsException):
38 | """
39 | Exception raised for invalid arguments when calling a method.
40 | """
41 |
42 | ...
43 |
44 |
45 | class MismatchedDimension(ArgError):
46 | """
47 | Exception raised when multiple sources of truth for a collection's embedding dimension do not match.
48 | """
49 |
50 | ...
51 |
52 |
53 | class FilterError(VecsException):
54 | """
55 | Exception raised when there's an error related to filter usage in a query.
56 | """
57 |
58 | ...
59 |
60 |
61 | class IndexNotFound(VecsException):
62 | """
63 | Exception raised when attempting to access an index that does not exist.
64 | """
65 |
66 | ...
67 |
68 |
69 | class Unreachable(VecsException):
70 | """
71 | Exception raised when an unreachable part of the code is executed.
72 | This is typically used for error handling in cases that should be logically impossible.
73 | """
74 |
75 | ...
76 |
77 |
78 | class MissingDependency(VecsException, ImportError):
79 | """
80 | Exception raised when attempting to access a feature that requires an optional dependency when the optional dependency is not present.
81 | """
82 |
83 | ...
84 |
--------------------------------------------------------------------------------