├── .cargo └── config.toml ├── .github └── workflows │ ├── CI.yml │ └── pr.yaml ├── .gitignore ├── CHANGELOG.md ├── CONTRIBUTING.md ├── Cargo.toml ├── LEGACY-README.md ├── LICENSE ├── Makefile ├── README.md ├── client_sdk ├── Cargo.toml ├── build.rs └── src │ ├── client │ ├── control_plane.rs │ ├── grpc.rs │ ├── mod.rs │ └── pinecone_client.rs │ ├── data_types.rs │ ├── index.rs │ ├── lib.rs │ ├── proto │ └── vector_service.proto │ └── utils │ ├── conversions.rs │ ├── errors.rs │ ├── mod.rs │ └── python_conversions.rs ├── config.toml ├── index_service ├── Cargo.toml ├── build.rs └── src │ └── lib.rs ├── openapi └── index_service.json ├── pinecone ├── .gitignore ├── Cargo.toml ├── pyproject.toml └── src │ ├── client.rs │ ├── data_types.rs │ ├── index.rs │ ├── lib.rs │ └── utils │ ├── errors.rs │ └── mod.rs └── tests ├── __init__.py ├── unit ├── __init__.py ├── test_control_plane.py ├── test_data_plane.py ├── test_hybrid_search.py ├── test_metadata.py ├── test_upsert_format.py └── tox.ini └── utils ├── __init__.py ├── remote_index.py └── utils.py /.cargo/config.toml: -------------------------------------------------------------------------------- 1 | [target.x86_64-apple-darwin] 2 | rustflags = [ 3 | "-C", "link-arg=-undefined", 4 | "-C", "link-arg=dynamic_lookup", 5 | ] 6 | 7 | [target.aarch64-apple-darwin] 8 | rustflags = [ 9 | "-C", "link-arg=-undefined", 10 | "-C", "link-arg=dynamic_lookup", 11 | ] -------------------------------------------------------------------------------- /.github/workflows/CI.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | tags: 8 | - 'V*' 9 | workflow_dispatch: 10 | 11 | jobs: 12 | generate_openapi: 13 | runs-on: ubuntu-latest 14 | steps: 15 | - uses: actions/checkout@v4 16 | with: 17 | path: pinecone_client 18 | - run: | 19 | cd pinecone_client 20 | make generate-index-service 21 | - uses: actions/upload-artifact@v3 22 | with: 23 | name: index-service 24 | path: pinecone_client/index_service 25 | 26 | linux: 27 | 28 | runs-on: ubuntu-latest 29 | needs: generate_openapi 30 | steps: 31 | - uses: actions/checkout@v4 32 | - uses: actions/download-artifact@v3 33 | with: 34 | name: index-service 35 | path: ./index_service/ 36 | - uses: PyO3/maturin-action@v1 37 | with: 38 | container: quay.io/pypa/manylinux2014_x86_64:latest 39 | manylinux: auto 40 | command: build 41 | args: --release --sdist -o dist -i python3.7 -i python3.8 -i python3.9 -i python3.10 -i python3.11 42 | working-directory: ./pinecone 43 | before-script-linux: | 44 | yum install -y openssl-devel && 45 | PB_REL="https://github.com/protocolbuffers/protobuf/releases" && 46 | curl -LO $PB_REL/download/v22.2/protoc-22.2-linux-x86_64.zip && 47 | unzip protoc-22.2-linux-x86_64.zip -d $HOME/.local && 48 | export PROTOC=$HOME/.local/bin/protoc 49 | docker-options: | 50 | -e PROTOC=$HOME/.local/bin/protoc 51 | - name: Upload wheels 52 | uses: actions/upload-artifact@v3 53 | with: 54 | name: wheels 55 | path: ./pinecone/dist 56 | 57 | 58 | linux_aa64: 59 | runs-on: ubuntu-latest 60 | needs: generate_openapi 61 | steps: 62 | - uses: actions/checkout@v4 63 | - uses: actions/download-artifact@v3 64 | with: 65 | name: index-service 66 | path: ./index_service/ 67 | - uses: PyO3/maturin-action@v1 68 | with: 69 | container: ghcr.io/rust-cross/manylinux2014-cross:aarch64 70 | target: aarch64-unknown-linux-gnu 71 | manylinux: auto 72 | command: build 73 | args: --release --sdist -o dist -i python3.7 -i python3.8 -i python3.9 -i python3.10 -i python3.11 74 | working-directory: ./pinecone 75 | before-script-linux: | 76 | sudo apt-get update && sudo apt-get install -y pkg-config libssl-dev clang 77 | sudo apt install -y protobuf-compiler libprotobuf-dev 78 | - name: Upload wheels 79 | uses: actions/upload-artifact@v3 80 | with: 81 | name: wheels 82 | path: ./pinecone/dist 83 | 84 | windows: 85 | runs-on: windows-latest 86 | needs: generate_openapi 87 | strategy: 88 | matrix: 89 | python_version: [ "python3.7", "python3.8", "python3.9", "python3.10", "python3.11" ] 90 | steps: 91 | - uses: actions/checkout@v4 92 | - name: install protoc 93 | run: choco install protoc 94 | - uses: actions/download-artifact@v3 95 | with: 96 | name: index-service 97 | path: ./index_service/ 98 | - uses: PyO3/maturin-action@v1 99 | with: 100 | command: build 101 | args: --release -o dist -i ${{ matrix.python_version }} 102 | working-directory: ./pinecone 103 | - name: Upload wheels 104 | uses: actions/upload-artifact@v3 105 | with: 106 | name: wheels 107 | path: ./pinecone/dist 108 | 109 | macos: 110 | runs-on: macos-latest 111 | needs: generate_openapi 112 | strategy: 113 | matrix: 114 | python_version: ["python3.7", "python3.8", "python3.9", "python3.10", "python3.11"] 115 | steps: 116 | - uses: actions/checkout@v4 117 | - uses: actions/download-artifact@v3 118 | with: 119 | name: index-service 120 | path: ./index_service/ 121 | - uses: PyO3/maturin-action@v1 122 | with: 123 | command: build 124 | args: --release -o dist --universal2 -i ${{ matrix.python_version }} 125 | working-directory: ./pinecone 126 | - name: Upload wheels 127 | uses: actions/upload-artifact@v3 128 | with: 129 | name: wheels 130 | path: ./pinecone/dist 131 | 132 | test: 133 | name: Test wheels 134 | strategy: 135 | fail-fast: false 136 | matrix: 137 | os: [ubuntu-latest, macos-latest, windows-latest] 138 | python_version: ["3.7", "3.8", "3.9", "3.10", "3.11"] 139 | runs-on: ${{ matrix.os }} 140 | needs: [ macos, linux, linux_aa64, windows ] 141 | steps: 142 | - uses: actions/checkout@v4 143 | - uses: actions/download-artifact@v3 144 | with: 145 | name: wheels 146 | path: wheels 147 | - uses: actions/setup-python@v4 148 | with: 149 | python-version: ${{ matrix.python_version }} 150 | - name: Install wheel 151 | run: python -m pip install --no-index --find-links=./wheels/ pinecone-client 152 | - name: Try importing 153 | run: python -c "import pinecone" 154 | - name: Test 155 | if: false # TODO: decide if we want to run full integration tests 156 | run: make integration-test 157 | 158 | release: 159 | name: Release 160 | runs-on: ubuntu-latest 161 | if: "startsWith(github.ref, 'refs/tags/')" 162 | needs: [ test ] 163 | steps: 164 | - uses: actions/download-artifact@v3 165 | with: 166 | name: wheels 167 | - name: Publish to PyPI 168 | uses: PyO3/maturin-action@v1.39.0 169 | env: 170 | MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }} 171 | MATURIN_REPOSITORY: testpypi 172 | with: 173 | command: upload 174 | args: --skip-existing * -------------------------------------------------------------------------------- /.github/workflows/pr.yaml: -------------------------------------------------------------------------------- 1 | name: Pull-request CI 2 | 3 | on: pull_request 4 | 5 | env: 6 | PINECONE_REGION: eu-west1-gcp 7 | PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY_EU }} 8 | 9 | jobs: 10 | rust-tests: 11 | name: Rust linter and tests 12 | runs-on: ubuntu-latest 13 | steps: 14 | # TODO: make common install steps a separate action 15 | - uses: actions/checkout@v4 16 | - uses: dtolnay/rust-toolchain@stable 17 | - name: Install Build Dependencies 18 | run: sudo apt-get -y update && sudo apt-get -y install pkg-config libssl-dev clang 19 | - name: Install protoc 20 | run: sudo apt install -y protobuf-compiler libprotobuf-dev 21 | - name: Run rustfmt check 22 | run: cargo fmt -p pinecone -p client_sdk -- --check 23 | - name: Run clippy 24 | run: cargo clippy -p pinecone -p client_sdk --tests --benches --examples --bins -- -D warnings 25 | - name: Run cargo tests 26 | if: false # TODO: Fix tests 27 | run: cargo test -p pinecone -p client_sdk 28 | 29 | python-tests: 30 | name: Integration Tests 31 | runs-on: ubuntu-latest 32 | steps: 33 | - uses: actions/setup-python@v4 34 | with: 35 | python-version: 'pypy3.9' 36 | - uses: actions/checkout@v4 37 | - uses: dtolnay/rust-toolchain@stable 38 | - name: Install Build Dependencies 39 | run: sudo apt-get -y update && sudo apt-get -y install pkg-config libssl-dev clang 40 | - name: Install protoc 41 | run: sudo apt install -y protobuf-compiler libprotobuf-dev 42 | - name: Run integration tests 43 | run: 44 | make integration-test 45 | - name: Upload HTML report as artifact 46 | if: always() 47 | uses: actions/upload-artifact@v2 48 | with: 49 | name: pytest-report 50 | path: tests/unit/report.html 51 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | 3 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 4 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 5 | **/Cargo.lock 6 | 7 | # These are backup files generated by rustfmt 8 | **/*.rs.bk 9 | 10 | 11 | # Added by cargo 12 | 13 | **/target 14 | 15 | # Local scripting dir for testing 16 | my_scripts 17 | 18 | # autogen control plane code 19 | index_service/ -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | ## 3.0.0rc1 - breaking changes!! 4 | **New client architecture based on pre-compiled Rust code and gRPC** 5 | ### Breaking changes and migration guide: 6 | #### Client initialization 7 | - The `Client` object is the main entry point for control operations like creating, deleting and configuring pinecone indexes. 8 | It replaces the global `pinecone.init()` function from previous versions. 9 | For most minimal code change, you can initialize a `Client` instance and use it as a drop-in replacement for the global `pinecone` object (see [README](https://github.com/pinecone-io/pinecone-client#migrating-from-pinecone-client-v2)). 10 | - The `envrionement` parameter for `Client` initialization was renamed `region` instead. Similarly, the `PINECONE_ENVIRONMENT` environment variable was renamed to `PINECONE_REGION` as well. 11 | 12 | #### Control plane operations 13 | - The `create_index()` method is mostly unchanged, but a few deprecated or unused parameters were removed: `index_type`, `index_config` 14 | - The `configure_index()` method was removed. Use `scale_index()` instead. 15 | 16 | #### Upsert operation 17 | - `index.upsert()` now supports mixing vector represenations in the same batch (see [Upserting vectors](https://github.com/pinecone-io/pinecone-client#upserting-vectors) 18 | - When used with `async_req=True`, `index.upsert()` now returns an `asyncio` coroutine instead of a `concurrent.futures.Future` object. See [Performance tuning](https://github.com/pinecone-io/pinecone-client#performance-tuning-for-upsering-large-datasets) for more details. 19 | 20 | #### Query operation 21 | - Querying using an existing vector id was separated into a new method `index.query_by_id()`. The `index.query()` method now only accepts a vector values (and optioanl sparse values). 22 | - The both `index.query()` and `index.query_by_id()` now return a list of `QueryResult` objects. The `QueryResult` object has the following attributes: 23 | - `id` - the vector id 24 | - `score` - the ANN score for the given query result 25 | - `values` - optional vector values if `include_values=True` was passed to the query method 26 | - `sparse_values` - optional vector sparse values if `include_values=True` was passed to the query method 27 | - `metadata` - optional vector metadata if `include_metadata=True` was passed to the query method 28 | - `to_dict()` - a method that returns a dictionary representation of the query result 29 | 30 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Building 2 | 3 | ## Install dependencies 4 | 5 | 1. **Install [Rust compiler](https://www.rust-lang.org/tools/install)** 6 | ```bash 7 | curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh 8 | ``` 9 | 2. **Install [protobuf compiler](https://grpc.io/docs/protoc-installation/)** 10 | ```bash 11 | # Linux 12 | sudo apt-get install protobuf-compiler 13 | 14 | # mac OS 15 | brew install protobuf 16 | ``` 17 | Or alternatively install from source: https://github.com/protocolbuffers/protobuf/releases/tag/v22.2 18 | 19 | **Note:** If you are still getting an error like ``Could not find protoc installation`` - set the `PROTOC` environment variable to the `protoc` binary you just installed. 20 | ```bash 21 | export PROTOC=/path/to/protoc 22 | ``` 23 | 24 | 3. **Install lib-ssl** 25 | 26 | If you are getting an error like `Could not find directory of OpenSSL installation`, you need to install lib-ssl. 27 | #### linux 28 | ```bash 29 | sudo apt-get update && sudo apt-get install -y pkg-config libssl-dev 30 | ``` 31 | #### mac OS 32 | ```bash 33 | brew install openssl 34 | ``` 35 | 36 | 4. **Generate OpenAPI client** (Optional, usually done automatically at build time) 37 | 38 | Pinecone uses an OpenAPI spec for control-plane operations like `create_index()`. The OpenAPI client is automatically generated using [openapi-generator](https://github.com/OpenAPITools/openapi-generator/blob/master/docs/generators/rust-server.md) during project build. 39 | This process uses Docker to `docker run` OpenAPI's generator image. 40 | **If you don't have docker installed, or you don't want to use docker** - you can download the generated code from the [latest release](https://github.com/pinecone-io/pinecone-client/releases). 41 | Simply extract the `index_service.zip` file into the `index_service/` folder at the root of the project. 42 | 43 | ## Building from source 44 | ### Python package 45 | #### Using the pyproject.toml file 46 | ```bash 47 | # Create and activate a virtual environment 48 | python3 -m venv venv 49 | source venv/bin/activate 50 | 51 | # Install as editable package 52 | cd pinecone 53 | pip install -e . 54 | 55 | # optionallly, install test dependencies and run tests: 56 | pip install -e .[test] 57 | pytest ../tests/unit 58 | ``` 59 | #### Using `maturin` 60 | ```bash 61 | # Create and activate a virtual environment 62 | python3 -m venv venv 63 | source venv/bin/activate 64 | 65 | # Install maturin 66 | pip install maturin 67 | 68 | # Install pinecone package for development 69 | cd pinecone 70 | maturin develop 71 | ``` 72 | #### Building a wheel for deployment 73 | ```bash 74 | cd pinecone 75 | maturin build --release 76 | ``` 77 | ### Building rust library for linking with other languages 78 | ```bash 79 | cargo build 80 | ``` 81 | 82 | # Contributing 83 | TBD -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = [ 3 | "client_sdk", 4 | "pinecone", 5 | "index_service" 6 | ] 7 | 8 | [workspace.package] 9 | version = "3.0.0" -------------------------------------------------------------------------------- /LEGACY-README.md: -------------------------------------------------------------------------------- 1 | This client is based on pre-compiled Rust code and gRPC, aimed at improving performance and stability. 2 | 3 | Using native gRPC transport, this client is able to achieve a 2x-3x speedup for vector upsert over the previous RESTful client versions, as well as a 10-20% speedup for vector query latency. 4 | 5 | As the client installation is fully self-contained, it does not require any additional dependencies (e.g. `grpcio`), making it easier to install and use in any Python environment. 6 | 7 | 8 | ## Building from source and contributing 9 | See [CONTRIBUTING.md](CONTRIBUTING.md) 10 | 11 | ## Migrating from Pinecone client V2 12 | If you are migrating from pinecone client V2, here is the most minimal code change required to upgrade to V3: 13 | ```python 14 | # Pinecone client V2 15 | import pinecone 16 | pinecone.init() # One time init 17 | ... 18 | index = pinecone.Index("example-index") 19 | index.upsert(...) 20 | 21 | # Pinecone client V3 22 | from pinecone import Client 23 | pinecone = Client() # This is now a `Client` instance 24 | ... 25 | # Unchanged! 26 | index = pinecone.Index("example-index") 27 | index.upsert(...) 28 | ``` 29 | For more API changes see [CHANGELOG.md](CHANGELOG.md) 30 | 31 | # Usage 32 | 33 | ## Index operations 34 | 35 | ### Creating a Client instance 36 | The `Client` is the main entry point for index operations like creating, deleting and configuring Pinecone indexes. 37 | Initializing a `Client` requires your Pinecone API key and a region, which can be passed as either environment variables or as parameters to the `Client` constructor. 38 | 39 | ```python 40 | import os 41 | from pinecone import Client 42 | 43 | # Initialize a client using environment variables 44 | os.environ['PINECONE_API_KEY'] = 'YOUR_API_KEY' 45 | os.environ['PINECONE_REGION'] = 'us-west1-gcp' 46 | client = Client() 47 | 48 | # Initialize a client using parameters 49 | client = Client(api_key = 'YOUR_API_KEY', region = 'us-west1-gcp') 50 | ``` 51 | 52 | ### Creating an index 53 | 54 | The following example creates an index without a metadata configuration. 55 | 56 | By default, all metadata fields are indexed. 57 | 58 | ```python 59 | 60 | from pinecone import Client 61 | 62 | client = Client(api_key="YOUR_API_KEY", region="us-west1-gcp") 63 | 64 | index = client.create_index("example-index", dimension=1024) 65 | ``` 66 | 67 | If some metadata fields contain data payload such as raw text, indexing these fields would make the Pinecone index less efficient. In such cases, it is recommended to configure the index to only index specific metadata fields which are used for query filtering. 68 | 69 | The following example creates an index that only indexes the `"color"` metadata field. 70 | 71 | ```python 72 | metadata_config = { 73 | "indexed": ["color"] 74 | } 75 | 76 | client.create_index("example-index-2", dimension=1024, 77 | metadata_config=metadata_config) 78 | ``` 79 | 80 | #### Listing all indexes 81 | 82 | The following example returns all indexes in your project. 83 | 84 | ```python 85 | active_indexes = client.list_indexes() 86 | ``` 87 | 88 | #### Getting index configuration 89 | 90 | The following example returns information about the index `example-index`. 91 | 92 | ```python 93 | index_description = client.describe_index("example-index") 94 | ``` 95 | 96 | #### Deleting an index 97 | 98 | The following example deletes `example-index`. 99 | 100 | ```python 101 | client.delete_index("example-index") 102 | ``` 103 | 104 | #### Scaling an existing index number of replicas 105 | 106 | The following example changes the number of replicas for `example-index`. 107 | 108 | ```python 109 | new_number_of_replicas = 4 110 | client.scale_index("example-index", replicas=new_number_of_replicas) 111 | ``` 112 | ## Vector operations 113 | ### Creating an Index instance 114 | The index object is the entry point for vector operations like upserting, querying and deleting vectors to a given Pinecone index. 115 | ```python 116 | from pinecone import Client 117 | client = Client(api_key="YOUR_API_KEY", region="us-west1-gcp") 118 | index = client.get_index("example-index") 119 | 120 | # Backwards compatibility 121 | index = client.Index("example-index") 122 | ``` 123 | 124 | #### Printing index statistics 125 | 126 | The following example returns statistics about the index `example-index`. 127 | 128 | ```python 129 | from pinecone import Client 130 | 131 | client = Client(api_key="YOUR_API_KEY", region="us-west1-gcp") 132 | index = client.Index("example-index") 133 | 134 | print(index.describe_index_stats()) 135 | ``` 136 | 137 | 138 | #### Upserting vectors 139 | 140 | The following example upserts vectors to `example-index`. 141 | 142 | ```python 143 | from pinecone import Client, Vector, SparseValues 144 | client = Client(api_key="YOUR_API_KEY", region="us-west1-gcp") 145 | index = client.get_index("example-index") 146 | 147 | upsert_response = index.upsert( 148 | vectors=[ 149 | ("vec1", [0.1, 0.2, 0.3, 0.4], {"genre": "drama"}), 150 | ("vec2", [0.2, 0.3, 0.4, 0.5], {"genre": "action"}), 151 | ], 152 | namespace="example-namespace" 153 | ) 154 | 155 | # Mixing different vector representations is allowed 156 | upsert_response = index.upsert( 157 | vectors=[ 158 | # Tuples 159 | ("vec1", [0.1, 0.2, 0.3, 0.4]), 160 | ("vec2", [0.2, 0.3, 0.4, 0.5], {"genre": "action"}), 161 | # Vector objects 162 | Vector(id='id1', values=[1.0, 2.0, 3.0], metadata={'key': 'value'}), 163 | Vector(id='id3', values=[1.0, 2.0, 3.0], sparse_values=SparseValues(indices=[1, 2], values=[0.2, 0.4])), 164 | # Dictionaries 165 | {'id': 'id1', 'values': [1.0, 2.0, 3.0], 'metadata': {'key': 'value'}}, 166 | {'id': 'id2', 'values': [1.0, 2.0, 3.0], 'sparse_values': {'indices': [1, 2], 'values': [0.2, 0.4]}}, 167 | ], 168 | namespace="example-namespace" 169 | ) 170 | ``` 171 | 172 | #### Querying an index by a new unseen vector 173 | 174 | The following example queries the index `example-index` with metadata 175 | filtering. 176 | 177 | ```python 178 | query_response = index.query( 179 | values=[1.0, 5.3, 8.9, 0.5], # values of a query vector 180 | sparse_values = None, # optional sparse values of the query vector 181 | top_k=10, 182 | namespace="example-namespace", 183 | include_values=True, 184 | include_metadata=True, 185 | filter={ 186 | "genre": {"$in": ["comedy", "documentary", "drama"]} 187 | } 188 | ) 189 | ``` 190 | 191 | #### Querying an index by an existing vector ID 192 | 193 | The following example queries the index `example-index` for the `top_k=10` nearest neighbors of the vector with ID `vec1`. 194 | 195 | ```python 196 | query_response = index.query( 197 | id="vec1", 198 | top_k=10, 199 | namespace="example-namespace", 200 | include_values=True, 201 | include_metadata=True, 202 | ) 203 | ``` 204 | 205 | #### Deleting vectors 206 | 207 | ```python 208 | # Delete vectors by IDs 209 | index.delete(ids=["vec1", "vec2"], namespace="example-namespace") 210 | 211 | # Delete vectors by metadata filters 212 | index.delete_by_metadata(filter={"genre": {"$in": ["comedy", "documentary", "drama"]}}, namespace="example-namespace") 213 | 214 | # Delete all vectors in a given namespace (use namespace="" to delete all vectors in the DEFAULT namespace) 215 | index.delete_all(namespace="example-namespace") 216 | ``` 217 | 218 | #### Fetching vectors by ids 219 | 220 | The following example fetches vectors by ID without querying for nearest neighbors. 221 | 222 | ```python 223 | fetch_response = index.fetch(ids=["vec1", "vec2"], namespace="example-namespace") 224 | ``` 225 | 226 | 227 | #### Update vectors 228 | 229 | The following example updates vectors by ID. 230 | 231 | ```python 232 | update_response = index.update( 233 | id="vec1", 234 | values=[0.1, 0.2, 0.3, 0.4], 235 | set_metadata={"genre": "drama"}, 236 | namespace="example-namespace" 237 | ) 238 | ``` 239 | # Performance tuning for upserting large datasets 240 | To upsert an entire dataset of vectors, we recommend using concurrent batched upsert requests. The following example shows how to do this using the `asyncio` library: 241 | ```python 242 | import asyncio 243 | from pinecone import Client, Vector 244 | 245 | def chunker(seq, batch_size): 246 | return (seq[pos:pos + batch_size] for pos in range(0, len(seq), batch_size)) 247 | 248 | async def async_upload(index, vectors, batch_size, max_concurrent=50): 249 | sem = asyncio.Semaphore(max_concurrent) 250 | async def send_batch(batch): 251 | async with sem: 252 | return await index.upsert(vectors=batch, async_req=True) 253 | 254 | await asyncio.gather(*[send_batch(chunk) for chunk in chunker(vectors, batch_size=batch_size)]) 255 | 256 | # To use it: 257 | client = Client() 258 | index = client.get_index("example-index") 259 | asyncio.run(async_upload(index, vectors, batch_size=100)) 260 | 261 | # In a jupyter notebook, asyncio.run() is not supported. Instead, use 262 | await async_upload(index, vectors, batch_size=100) 263 | ``` 264 | 265 | # Limitations 266 | 267 | ## Code completion and type hints 268 | Due to limitations with the underlying `pyo3` library, code completion and type hints are not available in some IDEs, or might require additional configuration. 269 | - **Jupyter notebooks**: Should work out of the box. 270 | - **VSCode**: Change the [`languageServer`](https://code.visualstudio.com/docs/python/settings-reference#_intellisense-engine-settings) setting to `jedi`. 271 | - **PyCharm**: For the moment, all function signatures would show `(*args, **kwargs)`. We are working on a solution ASAP. (Function docstrings would still show full arguments and type hints). 272 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | VENV_DIR = .temp_venv 2 | 3 | .PHONY: setup venv develop integration-test finish clean 4 | 5 | venv: $(VENV_DIR)/bin/activate 6 | 7 | $(VENV_DIR)/bin/activate: 8 | test -d $(VENV_DIR) || python3 -m venv $(VENV_DIR) 9 | touch $(VENV_DIR)/bin/activate 10 | 11 | develop: venv 12 | . $(VENV_DIR)/bin/activate; cd pinecone && pip3 install -e .[test] 13 | 14 | integration-test: venv develop 15 | . $(VENV_DIR)/bin/activate; cd tests/unit && pytest --self-contained-html --dist=loadscope --numprocesses 4 --durations=10 --durations-min=1.0 --html=report.html 16 | $(MAKE) finish 17 | 18 | finish: venv 19 | rm -rf $(VENV_DIR) 20 | 21 | clean: 22 | rm -rf $(VENV_DIR) 23 | 24 | generate-index-service: 25 | docker run --rm -v "${CURDIR}:/local" openapitools/openapi-generator-cli:v6.3.0 generate --input-spec /local/openapi/index_service.json --generator-name rust --output /local/index_service --additional-properties packageName=index_service --additional-properties packageVersion=0.1.0 --additional-properties withSerde=true --additional-properties supportMultipleResponses=true 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Experimental client warning 2 | 3 | > **_⚠️ Warning_** 4 | > 5 | > This experimental client is not being actively developed 6 | > or maintained. We are leaving it up for anyone who may 7 | > be depending on it, but currently we encourage python 8 | > users to use the official python client developed at 9 | > [pinecone-io/pinecone-python-client](https://github.com/pinecone-io/pinecone-python-client) 10 | 11 | If you understand the risks and wish to continue, you can find the [old README here](./LEGACY-README.md). -------------------------------------------------------------------------------- /client_sdk/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "client_sdk" 3 | version.workspace = true 4 | edition = "2021" 5 | repository = "https://github.com/" # TODO 6 | description = "SDK of the client" # TODO 7 | 8 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 9 | 10 | [lib] 11 | bench = false 12 | 13 | [dependencies] 14 | prost = "0.11" 15 | prost-types = "0.11.0" 16 | reqwest = { version = "0.11.13", features = ["json"]} 17 | serde = { version = "1.0.152", features = ["derive"]} 18 | serde_json = "1.0.91" 19 | thiserror = "1.0.38" 20 | tokio = { version = "1.16.1", features = ["rt-multi-thread"] } 21 | tonic = { version = "0.8", features = ["tls", "tls-roots"] } 22 | webpki-roots = "0.22.6" 23 | pyo3 = { version = "0.18.0", features = ["extension-module"] } 24 | derivative = "2.2.0" 25 | index_service = { version = "0.1.0", path = "../index_service" } 26 | openssl = { version = "0.10", features = ["vendored"] } 27 | 28 | [build-dependencies] 29 | tonic-build = "0.8" 30 | -------------------------------------------------------------------------------- /client_sdk/build.rs: -------------------------------------------------------------------------------- 1 | use std::error::Error; 2 | use std::process::Command; 3 | 4 | fn main() -> Result<(), Box> { 5 | if cfg!(target_os = "macos") { 6 | build_mac_dependencies()? 7 | } 8 | 9 | tonic_build::configure() 10 | .build_server(false) 11 | .compile(&["src/proto/vector_service.proto"], &["src/proto/"])?; 12 | 13 | Ok(()) 14 | } 15 | 16 | fn build_mac_dependencies() -> Result<(), Box> { 17 | // Install dependencies with homebrew 18 | let res = Command::new("brew").args(["install", "protobuf"]).output(); 19 | 20 | match res { 21 | Ok(output) => { 22 | if output.status.success() { 23 | println!("Installed dependencies with Homebrew"); 24 | } else { 25 | println!( 26 | "Failed to install dependencies with Homebrew: {}", 27 | String::from_utf8_lossy(&output.stderr) 28 | ); 29 | } 30 | } 31 | Err(e) => return Err(format!("Failed to install dependencies with Homebrew: {e}").into()), 32 | } 33 | 34 | Ok(()) 35 | } 36 | -------------------------------------------------------------------------------- /client_sdk/src/client/control_plane.rs: -------------------------------------------------------------------------------- 1 | use crate::data_types::Collection; 2 | use crate::data_types::Db; 3 | use crate::data_types::WhoamiResponse; 4 | use crate::utils::errors::PineconeClientError; 5 | use crate::utils::errors::PineconeResult; 6 | use index_service::apis::configuration; 7 | use index_service::apis::index_operations_api; 8 | use index_service::apis::index_operations_api::{ 9 | DescribeCollectionSuccess, DescribeIndexSuccess, ListCollectionsSuccess, ListIndexesSuccess, 10 | }; 11 | use index_service::models::CreateCollectionRequest; 12 | use index_service::models::PatchRequest; 13 | 14 | #[derive(Debug)] 15 | pub struct ControlPlaneClient { 16 | controller_url: String, 17 | configuration: configuration::Configuration, 18 | } 19 | 20 | impl ControlPlaneClient { 21 | pub fn new(controller_url: &str, api_key: &str) -> ControlPlaneClient { 22 | let mut config = configuration::Configuration::new(); 23 | config.base_path = controller_url.to_string(); 24 | config.api_key = Some(configuration::ApiKey { 25 | prefix: None, 26 | key: api_key.to_string(), 27 | }); 28 | config.user_agent = Some("pinecone-rust-client/0.1".to_string()); 29 | // can pass a custom client here 30 | config.client = reqwest::Client::new(); 31 | ControlPlaneClient { 32 | controller_url: controller_url.to_string(), 33 | configuration: config, 34 | } 35 | } 36 | 37 | pub async fn create_index(&self, index: Db) -> PineconeResult<()> { 38 | index_operations_api::create_index(&self.configuration, Some(index.into())).await?; 39 | Ok(()) 40 | } 41 | 42 | pub async fn delete_index(&self, name: &str) -> PineconeResult<()> { 43 | index_operations_api::delete_index(&self.configuration, name).await?; 44 | Ok(()) 45 | } 46 | 47 | pub async fn describe_index(&self, name: &str) -> PineconeResult { 48 | let response = index_operations_api::describe_index(&self.configuration, name).await?; 49 | match response 50 | .entity 51 | .ok_or(PineconeClientError::ControlPlaneParsingError {})? 52 | { 53 | DescribeIndexSuccess::Status200(entity) => Db::try_from(entity), 54 | DescribeIndexSuccess::UnknownValue(val) => { 55 | Err(PineconeClientError::Other(val.to_string())) 56 | } 57 | } 58 | } 59 | 60 | pub async fn list_indexes(&self) -> PineconeResult> { 61 | let response = index_operations_api::list_indexes(&self.configuration).await?; 62 | match response 63 | .entity 64 | .ok_or(PineconeClientError::ControlPlaneParsingError {})? 65 | { 66 | ListIndexesSuccess::Status200(entity) => Ok(entity), 67 | ListIndexesSuccess::UnknownValue(val) => { 68 | Err(PineconeClientError::Other(val.to_string())) 69 | } 70 | } 71 | } 72 | 73 | pub async fn configure_index( 74 | &self, 75 | name: &str, 76 | pod_type: Option, 77 | replicas: Option, 78 | ) -> PineconeResult<()> { 79 | let patch_request = PatchRequest { pod_type, replicas }; 80 | index_operations_api::configure_index(&self.configuration, name, Some(patch_request)) 81 | .await?; 82 | Ok(()) 83 | } 84 | 85 | pub async fn create_collection(&self, collection: Collection) -> PineconeResult<()> { 86 | let collection_request = CreateCollectionRequest::from(collection); 87 | index_operations_api::create_collection(&self.configuration, Some(collection_request)) 88 | .await?; 89 | Ok(()) 90 | } 91 | 92 | pub async fn describe_collection(&self, collection_name: &str) -> PineconeResult { 93 | let response = 94 | index_operations_api::describe_collection(&self.configuration, collection_name).await?; 95 | match response 96 | .entity 97 | .ok_or(PineconeClientError::ControlPlaneParsingError {})? 98 | { 99 | DescribeCollectionSuccess::Status200(entity) => Ok(Collection::from(entity)), 100 | DescribeCollectionSuccess::UnknownValue(val) => { 101 | Err(PineconeClientError::Other(val.to_string())) 102 | } 103 | } 104 | } 105 | 106 | pub async fn delete_collection(&self, collection_name: &str) -> PineconeResult<()> { 107 | index_operations_api::delete_collection(&self.configuration, collection_name).await?; 108 | Ok(()) 109 | } 110 | 111 | pub async fn list_collections(&self) -> PineconeResult> { 112 | let response = index_operations_api::list_collections(&self.configuration).await?; 113 | match response 114 | .entity 115 | .ok_or(PineconeClientError::ControlPlaneParsingError {})? 116 | { 117 | ListCollectionsSuccess::Status200(entity) => Ok(entity), 118 | ListCollectionsSuccess::UnknownValue(val) => { 119 | Err(PineconeClientError::Other(val.to_string())) 120 | } 121 | } 122 | } 123 | 124 | pub async fn whoami(&self) -> PineconeResult { 125 | let rq_client = self.configuration.client.clone(); 126 | let api_key = self 127 | .configuration 128 | .api_key 129 | .as_ref() 130 | .ok_or_else(|| PineconeClientError::ValueError("Error parsing Api Key".into()))? 131 | .key 132 | .as_str(); 133 | if api_key.is_empty() { 134 | return Err(PineconeClientError::ValueError( 135 | "Api key empty or not provided".into(), 136 | )); 137 | } 138 | let response = rq_client 139 | .get(&format!("{}/actions/whoami", self.controller_url)) 140 | .header("Api-Key", api_key) 141 | .send() 142 | .await 143 | .map_err(|e| PineconeClientError::ControlPlaneConnectionError { 144 | region: " ".to_string(), 145 | err: e.to_string(), 146 | })?; 147 | let json_repsonse = response.json::().await.map_err(|e| { 148 | PineconeClientError::ControlPlaneConnectionError { 149 | region: " ".to_string(), 150 | err: e.to_string(), 151 | } 152 | })?; 153 | Ok(json_repsonse) 154 | } 155 | } 156 | 157 | #[cfg(test)] 158 | mod control_plane_tests { 159 | use std::collections::BTreeMap; 160 | 161 | use super::ControlPlaneClient; 162 | use crate::data_types::Collection; 163 | use crate::data_types::Db; 164 | use std::env; 165 | 166 | struct ClientContext { 167 | client: ControlPlaneClient, 168 | } 169 | impl ClientContext { 170 | fn new() -> Self { 171 | let controller_uri = format!( 172 | "https://controller.{}.pinecone.io", 173 | env::var("PINECONE_REGION").unwrap_or_else(|_| "internal-beta".to_string()) 174 | ); 175 | let api_key = env::var("PINECONE_API_KEY").unwrap_or_else(|_| "".to_string()); 176 | let client = ControlPlaneClient::new(controller_uri.as_str(), api_key.as_str()); 177 | ClientContext { client } 178 | } 179 | } 180 | 181 | #[tokio::test] 182 | async fn test_create() { 183 | let context = ClientContext::new(); 184 | let index = Db { 185 | name: "test-index".to_string(), 186 | dimension: 128, 187 | metadata_config: Some( 188 | [( 189 | "indexed".to_string(), 190 | vec!["value1".to_string(), "value2".to_string()], 191 | )] 192 | .iter() 193 | .cloned() 194 | .collect::>>(), 195 | ), 196 | ..Default::default() 197 | }; 198 | let response = context.client.create_index(index).await; 199 | println!("{:?}", response); 200 | assert!(response.is_ok()); 201 | } 202 | 203 | #[tokio::test] 204 | async fn test_get() { 205 | let context = ClientContext::new(); 206 | let response = context.client.describe_index("test-index").await; 207 | println!("{:?}", response); 208 | assert!(response.is_ok()); 209 | } 210 | 211 | #[tokio::test] 212 | async fn test_list() { 213 | let context = ClientContext::new(); 214 | let response = context.client.list_indexes().await; 215 | println!("{:?}", response); 216 | assert!(response.is_ok()); 217 | } 218 | 219 | #[tokio::test] 220 | async fn test_update() { 221 | let context = ClientContext::new(); 222 | let response = context 223 | .client 224 | .configure_index("test-index", None, Some(2)) 225 | .await; 226 | println!("{:?}", response); 227 | assert!(response.is_ok()); 228 | } 229 | 230 | #[tokio::test] 231 | async fn test_create_collection() { 232 | let context = ClientContext::new(); 233 | let collection: Collection = Collection { 234 | name: "test-collection".to_string(), 235 | source: "test-index".to_string(), 236 | ..Default::default() 237 | }; 238 | let response = context.client.create_collection(collection).await; 239 | println!("{:?}", response); 240 | assert!(response.is_ok()); 241 | } 242 | 243 | #[tokio::test] 244 | async fn test_list_collection() { 245 | let context = ClientContext::new(); 246 | let response = context.client.list_collections().await; 247 | println!("{:?}", response); 248 | assert!(response.is_ok()); 249 | } 250 | 251 | #[tokio::test] 252 | async fn test_describe_collection() { 253 | let context = ClientContext::new(); 254 | let response = context.client.describe_collection("test-collection").await; 255 | println!("{:?}", response); 256 | assert!(response.is_ok()); 257 | } 258 | 259 | #[tokio::test] 260 | async fn test_delete_collection() { 261 | let context = ClientContext::new(); 262 | let response = context.client.delete_collection("test-collection").await; 263 | println!("{:?}", response); 264 | assert!(response.is_ok()); 265 | } 266 | 267 | #[tokio::test] 268 | async fn test_delete_index() { 269 | let context = ClientContext::new(); 270 | let response = context.client.delete_index("test-index").await; 271 | println!("{:?}", response); 272 | assert!(response.is_ok()); 273 | } 274 | 275 | #[tokio::test] 276 | async fn test_delete_invalid_timeout() { 277 | let context = ClientContext::new(); 278 | let response = context.client.delete_index("test-index").await; 279 | println!("{:?}", response); 280 | } 281 | 282 | #[tokio::test] 283 | async fn test_whoami() { 284 | let context = ClientContext::new(); 285 | let response = context.client.whoami().await; 286 | println!("{:?}", response); 287 | assert!(response.is_ok()); 288 | } 289 | } 290 | -------------------------------------------------------------------------------- /client_sdk/src/client/grpc.rs: -------------------------------------------------------------------------------- 1 | use self::dataplane_client::UpdateResponse; 2 | pub use self::dataplane_client::{ 3 | ScoredVector as GrpcScoredVector, SparseValues as GrpcSparseValues, Vector as GrpcVector, 4 | }; 5 | use crate::data_types::{ 6 | IndexStats, MetadataValue, NamespaceStats, QueryResult, SparseValues, Vector, 7 | }; 8 | use crate::utils::conversions; 9 | use crate::utils::errors::PineconeResult; 10 | use dataplane_client::vector_service_client::VectorServiceClient; 11 | use dataplane_client::{DescribeIndexStatsRequest, QueryRequest, UpsertRequest}; 12 | use std::collections::{BTreeMap, HashMap}; 13 | use tonic::metadata::Ascii; 14 | use tonic::{ 15 | metadata::MetadataValue as TonicMetadataVal, service::interceptor::InterceptedService, 16 | service::Interceptor, transport::Channel, Request, Status, 17 | }; 18 | 19 | mod dataplane_client { 20 | tonic::include_proto!("_"); 21 | } 22 | 23 | #[derive(Debug, Clone)] 24 | pub struct DataplaneGrpcClient { 25 | inner: VectorServiceClient>, 26 | } 27 | 28 | impl DataplaneGrpcClient { 29 | // TODO: this method shouldn't be public or exposed to python 30 | pub async fn connect( 31 | index_endpoint_url: String, 32 | api_key: &str, 33 | ) -> Result> { 34 | let channel = Channel::from_shared(index_endpoint_url)?.connect().await?; 35 | let token: TonicMetadataVal<_> = api_key.parse()?; 36 | let add_api_key_interceptor = ApiKeyInterceptor { api_token: token }; 37 | let inner = VectorServiceClient::with_interceptor(channel, add_api_key_interceptor); 38 | 39 | Ok(Self { inner }) 40 | } 41 | 42 | pub async fn upsert( 43 | &mut self, 44 | namespace: &str, 45 | vectors: &[Vector], 46 | ) -> Result { 47 | let grpc_vectors: Vec = vectors.iter().map(|v| v.clone().into()).collect(); 48 | let res = self 49 | .inner 50 | .upsert(UpsertRequest { 51 | namespace: namespace.to_string(), 52 | vectors: grpc_vectors, 53 | }) 54 | .await?; 55 | Ok(res.into_inner().upserted_count) 56 | } 57 | 58 | #[allow(clippy::too_many_arguments)] 59 | pub async fn query( 60 | &mut self, 61 | namespace: &str, 62 | id: Option, 63 | values: Option>, 64 | sparse_values: Option, 65 | top_k: u32, 66 | filter: Option>, 67 | include_values: bool, 68 | include_metadata: bool, 69 | ) -> PineconeResult> { 70 | let sparse_vectors = sparse_values.map(|sparse_vector| sparse_vector.into()); 71 | let res = self 72 | .inner 73 | .query(QueryRequest { 74 | namespace: namespace.to_string(), 75 | id: id.unwrap_or_default(), 76 | vector: values.unwrap_or_default(), 77 | sparse_vector: sparse_vectors, 78 | top_k, 79 | filter: filter.map(conversions::hashmap_to_prost_struct), 80 | include_values, 81 | include_metadata, 82 | queries: Vec::default(), // Deprecated 83 | }) 84 | .await?; 85 | 86 | res.into_inner() 87 | .matches 88 | .into_iter() 89 | .map(|sv| sv.try_into()) 90 | .collect() 91 | } 92 | 93 | pub async fn describe_index_stats( 94 | &mut self, 95 | filter: Option>, 96 | ) -> Result { 97 | let res = self 98 | .inner 99 | .describe_index_stats(DescribeIndexStatsRequest { 100 | filter: filter.map(conversions::hashmap_to_prost_struct), 101 | }) 102 | .await? 103 | .into_inner(); 104 | let ns_summaries = res.namespaces; 105 | let mut ns_map: HashMap = 106 | HashMap::with_capacity(ns_summaries.len()); 107 | for (ns_name, ns_summary) in ns_summaries { 108 | ns_map.insert( 109 | ns_name, 110 | NamespaceStats { 111 | vector_count: ns_summary.vector_count, 112 | }, 113 | ); 114 | } 115 | let stats: IndexStats = IndexStats { 116 | namespaces: ns_map, 117 | total_vector_count: res.total_vector_count, 118 | index_fullness: res.index_fullness, 119 | dimension: res.dimension, 120 | }; 121 | Ok(stats) 122 | } 123 | 124 | pub async fn fetch( 125 | &mut self, 126 | namespace: &str, 127 | ids: &[String], 128 | ) -> PineconeResult> { 129 | let res = self 130 | .inner 131 | .fetch(dataplane_client::FetchRequest { 132 | namespace: namespace.to_string(), 133 | ids: ids.to_owned(), 134 | }) 135 | .await?; 136 | let fetch_response = res.into_inner(); 137 | let vectors = fetch_response.vectors; 138 | let mut fetch_vectors: HashMap = HashMap::with_capacity(vectors.len()); 139 | for (id, vector) in vectors { 140 | fetch_vectors.insert(id, vector.try_into()?); 141 | } 142 | Ok(fetch_vectors) 143 | } 144 | 145 | pub async fn delete( 146 | &mut self, 147 | ids: Option>, 148 | namespace: &str, 149 | filter: Option>, 150 | delete_all: bool, 151 | ) -> Result<(), tonic::Status> { 152 | self.inner 153 | .delete(dataplane_client::DeleteRequest { 154 | namespace: namespace.into(), 155 | ids: ids.unwrap_or_default(), 156 | delete_all, 157 | filter: filter.map(conversions::hashmap_to_prost_struct), 158 | }) 159 | .await?; 160 | Ok(()) 161 | } 162 | 163 | pub async fn update( 164 | &mut self, 165 | id: &str, 166 | vector: Option<&Vec>, 167 | sparse_values: Option, 168 | set_metadata: Option>, 169 | namespace: &str, 170 | ) -> Result { 171 | let res = self 172 | .inner 173 | .update(dataplane_client::UpdateRequest { 174 | id: id.into(), 175 | values: match vector { 176 | Some(vec) => vec.clone(), 177 | None => Vec::new(), 178 | }, 179 | sparse_values: sparse_values.map(|sparse_values| sparse_values.into()), 180 | set_metadata: set_metadata.map(conversions::hashmap_to_prost_struct), 181 | namespace: namespace.into(), 182 | }) 183 | .await?; 184 | Ok(res.into_inner()) 185 | } 186 | } 187 | 188 | #[derive(Debug, Clone)] 189 | pub struct ApiKeyInterceptor { 190 | api_token: TonicMetadataVal, 191 | } 192 | 193 | impl Interceptor for ApiKeyInterceptor { 194 | fn call(&mut self, mut request: Request<()>) -> Result, Status> { 195 | // TODO: replace `api_token` with an `Option`, and do a proper `if_some`. 196 | if !self.api_token.is_empty() { 197 | request 198 | .metadata_mut() 199 | .insert("api-key", self.api_token.clone()); 200 | } 201 | Ok(request) 202 | } 203 | } 204 | 205 | /// Get internal grpc client 206 | /// This client would only work from within a pinecone region to the internal endpoint address/ 207 | /// It is meant to be used by internal services within the region that need to communicate with the Index GRPC API 208 | /// TODO: this function shouldn't be exposed by the python client 209 | pub async fn get_internal_grpc_client( 210 | index_endpoint_url: String, 211 | ) -> Result> { 212 | // TODO: Theoretically this method could have simply been one line: 213 | // VectorServiceClient::connect(index_endpoint_url).await? 214 | // But than the return type would be different, DataplaneGrpcClient would need to be Generic. 215 | // so TODO: Find a better way to expose an inner stateless, authentication-less, gRPC client 216 | 217 | let channel = Channel::from_shared(index_endpoint_url)?.connect().await?; 218 | let token: TonicMetadataVal<_> = "".parse()?; 219 | let add_api_key_interceptor = ApiKeyInterceptor { api_token: token }; 220 | let inner = VectorServiceClient::with_interceptor(channel, add_api_key_interceptor); 221 | Ok(DataplaneGrpcClient { inner }) 222 | } 223 | 224 | // todo: add better tests 225 | #[cfg(test)] 226 | mod tests { 227 | use crate::data_types::SparseValues; 228 | 229 | use super::DataplaneGrpcClient; 230 | const INDEX_ENDPOINT: &str = ""; 231 | const KEY: &str = ""; 232 | 233 | fn gen_random_dense_vectors(count: usize, dimension: i32) -> Vec { 234 | let mut vectors = Vec::new(); 235 | for i in 0..count { 236 | let values = vec![0.1; dimension as usize]; 237 | 238 | vectors.push(super::Vector { 239 | id: i.to_string(), 240 | values, 241 | sparse_values: None, 242 | metadata: None, 243 | }); 244 | } 245 | vectors 246 | } 247 | 248 | fn gen_random_mixed_vectors(count: usize, dimension: i32) -> Vec { 249 | let mut vectors = Vec::new(); 250 | for i in 0..count { 251 | let values = vec![0.1; dimension as usize]; 252 | let sparse_values = SparseValues { 253 | indices: vec![0; dimension as usize], 254 | values: vec![0.1; dimension as usize], 255 | }; 256 | vectors.push(super::Vector { 257 | id: i.to_string(), 258 | values, 259 | sparse_values: Some(sparse_values), 260 | metadata: None, 261 | }); 262 | } 263 | vectors 264 | } 265 | 266 | #[tokio::test] 267 | async fn test_upsert() { 268 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 269 | .await 270 | .unwrap(); 271 | let vectors = gen_random_dense_vectors(10, 1024); 272 | let res = client.upsert("ns", &vectors).await; 273 | assert!(res.unwrap() == 10) 274 | } 275 | 276 | #[tokio::test] 277 | async fn test_stats() { 278 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 279 | .await 280 | .unwrap(); 281 | let res = client.describe_index_stats(None).await; 282 | assert!(res.is_ok()); 283 | } 284 | 285 | #[tokio::test] 286 | async fn test_fetch() { 287 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 288 | .await 289 | .unwrap(); 290 | let res = client.fetch("ns", &["1".to_string()]).await; 291 | assert!(res.is_ok()); 292 | } 293 | 294 | #[tokio::test] 295 | async fn test_delete() { 296 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 297 | .await 298 | .unwrap(); 299 | let res = client 300 | .delete(Some(vec![("2".to_string())]), "ns", None, false) 301 | .await; 302 | assert!(res.is_ok()); 303 | } 304 | 305 | #[tokio::test] 306 | async fn test_delete_all() { 307 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 308 | .await 309 | .unwrap(); 310 | let res = client.delete(None, "ns", None, true).await; 311 | assert!(res.is_ok()); 312 | } 313 | 314 | #[tokio::test] 315 | async fn test_update() { 316 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 317 | .await 318 | .unwrap(); 319 | let res = client 320 | .update("1", Some(&vec![0.4; 128]), None, None, "ns") 321 | .await; 322 | assert!(res.is_ok()); 323 | } 324 | 325 | #[tokio::test] 326 | async fn test_mixed_upsert() { 327 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 328 | .await 329 | .unwrap(); 330 | let vectors = gen_random_mixed_vectors(10, 128); 331 | let res = client.upsert("ns", &vectors).await; 332 | assert!(res.unwrap() == 10) 333 | } 334 | 335 | #[tokio::test] 336 | async fn test_fetch_non_existent() { 337 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 338 | .await 339 | .unwrap(); 340 | let res = client.fetch("ns", &["100".to_string()]).await; 341 | assert!(res.unwrap().is_empty()); 342 | } 343 | 344 | #[tokio::test] 345 | async fn test_delete_non_existent() { 346 | let mut client = DataplaneGrpcClient::connect(INDEX_ENDPOINT.to_string(), KEY) 347 | .await 348 | .unwrap(); 349 | let res = client 350 | .delete(Some(vec!["100".to_string()]), "ns", None, false) 351 | .await; 352 | assert!(res.is_ok()); 353 | } 354 | } 355 | -------------------------------------------------------------------------------- /client_sdk/src/client/mod.rs: -------------------------------------------------------------------------------- 1 | mod control_plane; 2 | pub mod grpc; 3 | pub mod pinecone_client; 4 | -------------------------------------------------------------------------------- /client_sdk/src/client/pinecone_client.rs: -------------------------------------------------------------------------------- 1 | use pyo3::Python; 2 | use std::io::Write; 3 | use std::time::{Duration, Instant}; 4 | use std::{env, io}; 5 | 6 | use super::control_plane::ControlPlaneClient; 7 | use super::grpc::DataplaneGrpcClient; 8 | use crate::data_types::{Collection, Db}; 9 | use crate::index::Index; 10 | use crate::utils::errors::PineconeClientError::IndexConnectionError; 11 | use crate::utils::errors::{PineconeClientError, PineconeResult}; 12 | 13 | const DEAULT_PINECONE_REGION: &str = "us-west1-gcp"; 14 | 15 | #[derive(Debug)] 16 | pub struct PineconeClient { 17 | pub api_key: String, 18 | pub region: String, 19 | pub project_id: String, 20 | control_plane_client: ControlPlaneClient, 21 | } 22 | 23 | impl PineconeClient { 24 | pub async fn new( 25 | api_key: Option<&str>, 26 | region: Option<&str>, 27 | project_id: Option<&str>, 28 | ) -> PineconeResult { 29 | let api_key = match api_key { 30 | Some(s) => Ok(s.to_string()), 31 | None => env::var("PINECONE_API_KEY").map_err(|_| PineconeClientError::ValueError("Please provide a valid API key or set the 'PINECONE_API_KEY' environment variable".to_string())) 32 | }?; 33 | let region = match region { 34 | Some(s) => s.to_string(), 35 | None => { 36 | env::var("PINECONE_REGION").unwrap_or_else(|_| DEAULT_PINECONE_REGION.to_string()) 37 | } 38 | }; 39 | // Check if region is empty. For cases where the user sets the region to an empty string 40 | if region.is_empty() { 41 | return Err(PineconeClientError::ValueError( 42 | "Please provide a valid region or set the 'PINECONE_REGION' environment variable" 43 | .to_string(), 44 | )); 45 | } 46 | let control_plane_client = 47 | ControlPlaneClient::new(&PineconeClient::get_controller_url(®ion), &api_key); 48 | let project_id = match project_id { 49 | Some(id) => id.to_string(), 50 | None => PineconeClient::get_project_id(&control_plane_client) 51 | .await 52 | .map_err(|e| match e { 53 | PineconeClientError::ControlPlaneConnectionError { err, .. } => { 54 | PineconeClientError::ControlPlaneConnectionError { 55 | err, 56 | region: region.clone(), 57 | } 58 | } 59 | _ => e, 60 | })?, 61 | }; 62 | 63 | Ok(PineconeClient { 64 | api_key, 65 | region, 66 | project_id, 67 | control_plane_client, 68 | }) 69 | } 70 | 71 | fn get_index_url(&self, index_name: &str) -> String { 72 | let output = format!( 73 | "https://{index_name}-{project_id}.svc.{region}.pinecone.io:443", 74 | index_name = index_name, 75 | project_id = self.project_id, 76 | region = self.region 77 | ); 78 | output 79 | } 80 | 81 | fn get_controller_url(region: &str) -> String { 82 | let output = format!("https://controller.{}.pinecone.io", region); 83 | output 84 | } 85 | 86 | async fn get_dataplane_grpc_client( 87 | &self, 88 | index_name: &str, 89 | ) -> PineconeResult { 90 | let index_endpoint_url = self.get_index_url(index_name); 91 | let client = DataplaneGrpcClient::connect(index_endpoint_url, &self.api_key) 92 | .await 93 | .map_err(|e| IndexConnectionError { 94 | index: index_name.to_string(), 95 | err: e.to_string(), 96 | })?; 97 | Ok(client) 98 | } 99 | 100 | async fn get_project_id(control_plane_client: &ControlPlaneClient) -> PineconeResult { 101 | let whoami_response = control_plane_client.whoami().await?; 102 | Ok(whoami_response.project_name) 103 | } 104 | 105 | pub async fn create_index( 106 | &self, 107 | db: Db, 108 | timeout: Option, 109 | py: Option>, 110 | ) -> PineconeResult<()> { 111 | // If timeout is -ve and not -1 throw an error 112 | let name = db.name.clone(); 113 | // If timeout is -ve and not -1 throw an error 114 | if timeout.is_some() && timeout.unwrap() < -1 { 115 | return Err(PineconeClientError::ValueError( 116 | "Timeout must be -1 or a positive integer".to_string(), 117 | )); 118 | } 119 | self.control_plane_client.create_index(db).await?; 120 | // If -1 then don't wait for index to be ready 121 | if timeout == Some(-1) { 122 | return Ok(()); 123 | } 124 | // block until index is ready 125 | let mut new_index = self.describe_index(&name).await?; 126 | let start_time = Instant::now(); 127 | let max_timeout = Duration::from_secs(timeout.unwrap_or(300) as u64); 128 | if let Some(py) = py { 129 | py.run( 130 | "print(\"Waiting for index to be ready...\", flush=True)", 131 | None, 132 | None, 133 | ) 134 | .map_err(|_| PineconeClientError::Other("Failed to print to stdout".to_string()))?; 135 | } else { 136 | println!("Waiting for index to be ready..."); 137 | io::stdout().flush()?; 138 | } 139 | while new_index.status != Some("Ready".to_string()) { 140 | if let Some(py) = py { 141 | Python::check_signals(py) 142 | .map_err(|_| { 143 | let msg = "Interrupted. Index status unknown. Please call describe_index() to check status"; 144 | println!("{}", msg); 145 | io::stdout().flush().unwrap(); 146 | PineconeClientError::KeyboardInterrupt( 147 | msg.into(), 148 | ) 149 | })?; 150 | } 151 | if start_time.elapsed() > max_timeout { 152 | return Err(PineconeClientError::Other( 153 | "Index creation timed out. Please call describe_index() to check status." 154 | .to_string(), 155 | )); 156 | } 157 | new_index = self.describe_index(&name).await?; 158 | tokio::time::sleep(std::time::Duration::from_secs(5)).await; 159 | } 160 | Ok(()) 161 | } 162 | 163 | pub async fn get_index(&self, index_name: &str) -> PineconeResult { 164 | Ok(Index::new( 165 | index_name.to_string(), 166 | self.get_dataplane_grpc_client(index_name).await?, 167 | )) 168 | } 169 | 170 | pub async fn describe_index(&self, index_name: &str) -> PineconeResult { 171 | self.control_plane_client.describe_index(index_name).await 172 | } 173 | 174 | pub async fn list_indexes(&self) -> PineconeResult> { 175 | self.control_plane_client.list_indexes().await 176 | } 177 | 178 | pub async fn delete_index(&self, index_name: &str, timeout: Option) -> PineconeResult<()> { 179 | // If timeout is -ve and not -1 throw an error 180 | if timeout.is_some() && timeout.unwrap() < -1 { 181 | return Err(PineconeClientError::ValueError( 182 | "Timeout must be -1 or a positive integer".to_string(), 183 | )); 184 | } 185 | self.control_plane_client.delete_index(index_name).await?; 186 | if timeout == Some(-1) { 187 | return Ok(()); 188 | } 189 | // block until index is deleted 190 | println!("Verifying delete..."); 191 | let start_time = Instant::now(); 192 | let max_timeout = Duration::from_secs(timeout.unwrap_or(300) as u64); 193 | while self.list_indexes().await?.contains(&index_name.to_string()) { 194 | if start_time.elapsed() > max_timeout { 195 | return Err(PineconeClientError::Other( 196 | "Index deletion timed out. Please call describe_index to check status." 197 | .to_string(), 198 | )); 199 | } 200 | tokio::time::sleep(std::time::Duration::from_secs(5)).await; 201 | } 202 | Ok(()) 203 | } 204 | 205 | pub async fn configure_index( 206 | &self, 207 | index_name: &str, 208 | pod_type: Option, 209 | replicas: Option, 210 | ) -> PineconeResult<()> { 211 | self.control_plane_client 212 | .configure_index(index_name, pod_type, replicas) 213 | .await 214 | } 215 | 216 | pub async fn create_collection( 217 | &self, 218 | collection_name: &str, 219 | source_index: &str, 220 | ) -> PineconeResult<()> { 221 | let collection = Collection { 222 | name: collection_name.to_string(), 223 | source: source_index.to_string(), 224 | ..Default::default() 225 | }; 226 | self.control_plane_client 227 | .create_collection(collection) 228 | .await 229 | } 230 | 231 | pub async fn describe_collection(&self, collection_name: &str) -> PineconeResult { 232 | self.control_plane_client 233 | .describe_collection(collection_name) 234 | .await 235 | } 236 | 237 | pub async fn list_collections(&self) -> PineconeResult> { 238 | self.control_plane_client.list_collections().await 239 | } 240 | 241 | pub async fn delete_collection(&self, collection_name: &str) -> PineconeResult<()> { 242 | self.control_plane_client 243 | .delete_collection(collection_name) 244 | .await 245 | } 246 | } 247 | 248 | mod tests { 249 | #[tokio::test] 250 | async fn test_env_vars() { 251 | use super::*; 252 | env::set_var("PINECONE_API_KEY", ""); 253 | env::set_var("PINECONE_REGION", ""); 254 | let client = PineconeClient::new(None, None, None).await.unwrap(); 255 | println!("{:?}", client); 256 | } 257 | } 258 | -------------------------------------------------------------------------------- /client_sdk/src/data_types.rs: -------------------------------------------------------------------------------- 1 | use derivative::Derivative; 2 | 3 | use pyo3::types::{PyDict, PyList}; 4 | use serde::Deserialize; 5 | use std::collections::{BTreeMap, HashMap}; 6 | use std::vec::Vec; 7 | 8 | use pyo3::prelude::*; 9 | use pyo3::types::IntoPyDict; 10 | 11 | const SHORT_PRINT_LEN: usize = 5; 12 | 13 | #[derive(Debug, Default, Clone)] 14 | #[pyclass] 15 | #[pyo3(get_all)] 16 | #[pyo3(text_signature = "(indices, values)")] 17 | pub struct SparseValues { 18 | pub indices: Vec, 19 | pub values: Vec, 20 | } 21 | 22 | #[pymethods] 23 | impl SparseValues { 24 | #[new] 25 | #[pyo3(signature = (indices, values))] 26 | pub fn new(indices: Vec, values: Vec) -> Self { 27 | Self { indices, values } 28 | } 29 | 30 | pub fn __repr__(&self) -> Result { 31 | Ok(format!( 32 | "SparseValue:\n indices: {indices:?}...\n values: {values:?} ...", 33 | indices = &self.indices.chunks(5).next().unwrap_or(&Vec::::new()), 34 | values = &self.values.chunks(5).next().unwrap_or(&Vec::::new()) 35 | )) 36 | } 37 | } 38 | 39 | #[derive(Debug, Default, Clone)] 40 | #[pyclass] 41 | #[pyo3(get_all)] 42 | #[pyo3(text_signature = "(id, values, sparse_values=None, metadata=None)")] 43 | pub struct Vector { 44 | pub id: String, 45 | pub values: Vec, 46 | pub sparse_values: Option, 47 | pub metadata: Option>, 48 | } 49 | 50 | #[pymethods] 51 | impl Vector { 52 | #[new] 53 | #[pyo3(signature = (id, values, sparse_values=None, metadata=None))] 54 | pub fn new( 55 | id: String, 56 | values: Vec, 57 | sparse_values: Option, 58 | metadata: Option>, 59 | ) -> Self { 60 | Self { 61 | id, 62 | values, 63 | sparse_values, 64 | metadata, 65 | } 66 | } 67 | 68 | pub fn __repr__(&self, py: Python) -> Result { 69 | Ok("Vector:\n".to_string() + pretty_print_dict(self.to_dict(py), 2)?.as_str()) 70 | } 71 | 72 | pub fn to_dict<'a>(&self, py: Python<'a>) -> &'a PyDict { 73 | let key_vals: Vec<(&str, PyObject)> = vec![ 74 | ("id", self.id.to_object(py)), 75 | ("values", self.values.to_object(py)), 76 | ("sparse_values", self.sparse_values.to_object(py)), 77 | ("metadata", self.metadata.to_object(py)), 78 | ]; 79 | key_vals.into_py_dict(py) 80 | } 81 | } 82 | 83 | #[derive(Debug)] 84 | #[pyclass] 85 | #[pyo3(get_all)] 86 | pub struct UpsertResponse { 87 | pub upserted_count: u32, 88 | } 89 | 90 | #[pymethods] 91 | impl UpsertResponse { 92 | pub fn __repr__(&self, py: Python) -> Result { 93 | Ok("UpsertResponse:\n".to_string() + pretty_print_dict(self.to_dict(py), 2)?.as_str()) 94 | } 95 | 96 | pub fn to_dict<'a>(&self, py: Python<'a>) -> &'a PyDict { 97 | let key_vals: Vec<(&str, PyObject)> = 98 | vec![("upserted_count", self.upserted_count.to_object(py))]; 99 | key_vals.into_py_dict(py) 100 | } 101 | } 102 | 103 | #[derive(Debug)] 104 | #[pyclass] 105 | #[pyo3(get_all, mapping)] 106 | pub struct QueryResult { 107 | pub id: String, 108 | pub score: f32, 109 | pub values: Option>, 110 | pub sparse_values: Option, 111 | pub metadata: Option>, 112 | } 113 | 114 | #[pymethods] 115 | impl QueryResult { 116 | pub fn __repr__(&self, py: Python) -> Result { 117 | Ok("QueryResult:\n".to_string() + pretty_print_dict(self.to_dict(py), 2)?.as_str()) 118 | } 119 | 120 | pub fn to_dict<'a>(&self, py: Python<'a>) -> &'a PyDict { 121 | let key_vals: Vec<(&str, PyObject)> = vec![ 122 | ("id", self.id.to_object(py)), 123 | ("score", self.score.to_object(py)), 124 | ("values", self.values.to_object(py)), 125 | ("sparse_values", self.sparse_values.to_object(py)), 126 | ("metadata", self.metadata.to_object(py)), 127 | ]; 128 | key_vals.into_py_dict(py) 129 | } 130 | } 131 | 132 | #[derive(Deserialize, Debug)] 133 | pub struct WhoamiResponse { 134 | pub project_name: String, 135 | pub user_label: String, 136 | pub user_name: String, 137 | } 138 | 139 | #[derive(Deserialize, Debug, Clone)] 140 | #[pyclass] 141 | #[pyo3(get_all)] 142 | pub struct NamespaceStats { 143 | pub vector_count: u32, 144 | } 145 | 146 | #[pymethods] 147 | impl NamespaceStats { 148 | pub fn to_dict<'a>(&self, py: Python<'a>) -> &'a PyDict { 149 | let key_vals: Vec<(&str, PyObject)> = 150 | vec![("vector_count", self.vector_count.to_object(py))]; 151 | key_vals.into_py_dict(py) 152 | } 153 | } 154 | 155 | #[derive(Deserialize, Debug)] 156 | #[pyclass] 157 | #[pyo3(get_all)] 158 | pub struct IndexStats { 159 | pub namespaces: HashMap, 160 | pub dimension: u32, 161 | pub index_fullness: f32, 162 | pub total_vector_count: u32, 163 | } 164 | 165 | #[pymethods] 166 | impl IndexStats { 167 | pub fn __repr__(&self, py: Python) -> Result { 168 | Ok("Index statistics:\n".to_string() + pretty_print_dict(self.to_dict(py), 2)?.as_str()) 169 | } 170 | 171 | pub fn to_dict<'a>(&self, py: Python<'a>) -> &'a PyDict { 172 | let key_vals: Vec<(&str, PyObject)> = vec![ 173 | ("namespaces", self.namespaces.to_object(py)), 174 | ("dimension", self.dimension.to_object(py)), 175 | ("index_fullness", self.index_fullness.to_object(py)), 176 | ("total_vector_count", self.total_vector_count.to_object(py)), 177 | ]; 178 | key_vals.into_py_dict(py) 179 | } 180 | } 181 | 182 | #[derive(FromPyObject, Debug, Clone)] 183 | pub enum MetadataValue { 184 | StringVal(String), 185 | BoolVal(bool), 186 | NumberVal(f64), 187 | ListVal(Vec), 188 | DictVal(BTreeMap), 189 | } 190 | 191 | #[derive(Derivative, Default, Debug, Clone)] 192 | #[pyclass] 193 | #[pyo3(get_all, mapping)] 194 | pub struct Db { 195 | pub name: String, 196 | pub dimension: i32, 197 | pub metric: Option, 198 | pub replicas: Option, 199 | pub shards: Option, 200 | pub pods: Option, 201 | pub source_collection: Option, 202 | pub metadata_config: Option>>, 203 | pub pod_type: Option, 204 | pub status: Option, 205 | } 206 | 207 | #[derive(Derivative, Default, Debug, Clone)] 208 | #[pyclass] 209 | #[pyo3(get_all, mapping)] 210 | pub struct Collection { 211 | pub name: String, 212 | pub source: String, 213 | pub vector_count: Option, 214 | pub size: Option, 215 | pub status: Option, 216 | } 217 | 218 | #[pymethods] 219 | impl Db { 220 | pub fn __repr__(&self, py: Python) -> Result { 221 | Ok("Index config:\n".to_string() + pretty_print_dict(self.to_dict(py), 2)?.as_str()) 222 | } 223 | 224 | pub fn to_dict<'a>(&self, py: Python<'a>) -> &'a PyDict { 225 | let key_vals: Vec<(&str, PyObject)> = vec![ 226 | ("name", self.name.to_object(py)), 227 | ("dimension", self.dimension.to_object(py)), 228 | ("replicas", self.replicas.to_object(py)), 229 | ("shards", self.shards.to_object(py)), 230 | ("pod_type", self.pod_type.to_object(py)), 231 | ("metric", self.metric.to_object(py)), 232 | ("pods", self.pods.to_object(py)), 233 | ("source_collection", self.source_collection.to_object(py)), 234 | ("metadata_config", self.metadata_config.to_object(py)), 235 | ("status", self.status.to_object(py)), 236 | ]; 237 | key_vals.into_py_dict(py) 238 | } 239 | } 240 | 241 | #[pymethods] 242 | impl Collection { 243 | pub fn __repr__(&self, py: Python) -> Result { 244 | Ok("Collection:\n".to_string() + pretty_print_dict(self.to_dict(py), 2)?.as_str()) 245 | } 246 | 247 | pub fn to_dict<'a>(&self, py: Python<'a>) -> &'a PyDict { 248 | let key_vals: Vec<(&str, PyObject)> = vec![ 249 | ("name", self.name.to_object(py)), 250 | ("source", self.source.to_object(py)), 251 | ("vector_count", self.vector_count.to_object(py)), 252 | ("size", self.size.to_object(py)), 253 | ("status", self.status.to_object(py)), 254 | ]; 255 | key_vals.into_py_dict(py) 256 | } 257 | } 258 | 259 | fn pretty_print_dict(dict: &PyDict, indent: usize) -> Result { 260 | let mut msg = String::new(); 261 | for (k, v) in dict.into_iter() { 262 | if let Ok(dict_val) = v.downcast::() { 263 | let inner_msg = pretty_print_dict(dict_val, indent + 2)?; 264 | msg += format!( 265 | "{:indent$}{key}:\n{val}", 266 | "", 267 | key = k, 268 | val = inner_msg, 269 | indent = indent 270 | ) 271 | .as_str(); 272 | } else if let Ok(list_val) = v.downcast::() { 273 | let short_list = list_val.as_sequence().get_slice(0, SHORT_PRINT_LEN)?; 274 | let ellipsis = if list_val.len() > SHORT_PRINT_LEN { 275 | "..." 276 | } else { 277 | "" 278 | }; 279 | msg += format!( 280 | "{:indent$}{key}: {short_list:.3}{elipsis}\n", 281 | "", 282 | key = k, 283 | short_list = short_list, 284 | elipsis = ellipsis, 285 | indent = indent 286 | ) 287 | .as_str(); 288 | } else { 289 | msg += format!( 290 | "{:indent$}{key}: {val:.3}\n", 291 | "", 292 | key = k, 293 | val = v, 294 | indent = indent 295 | ) 296 | .as_str(); 297 | } 298 | } 299 | Ok(msg) 300 | } 301 | -------------------------------------------------------------------------------- /client_sdk/src/index.rs: -------------------------------------------------------------------------------- 1 | use crate::client::grpc::DataplaneGrpcClient; 2 | use crate::data_types::MetadataValue; 3 | use crate::data_types::{QueryResult, UpsertResponse, Vector}; 4 | use crate::utils::errors::{PineconeClientError, PineconeResult}; 5 | use std::collections::{BTreeMap, HashMap}; 6 | 7 | use crate::data_types::{IndexStats, SparseValues}; 8 | 9 | #[derive(Clone)] 10 | pub struct Index { 11 | pub name: String, 12 | dataplane_client: DataplaneGrpcClient, 13 | } 14 | 15 | impl Index { 16 | pub fn new(index_name: String, dataplane_client: DataplaneGrpcClient) -> Self { 17 | Index { 18 | name: index_name, 19 | dataplane_client, 20 | } 21 | } 22 | 23 | /// The `Upsert` operation writes vectors into a namespace. 24 | /// If a new value is upserted for an existing vector id, it will overwrite the previous value. 25 | /// 26 | /// # Arguments 27 | /// - `namespace` - the name of the namespace to which data will be upserted 28 | /// - `vectors` - a list of vectors to be upserted to the index. 29 | /// 30 | /// # Returns 31 | /// `Ok(list_ids)` with a list of vector ids that were successfully upserted to the Index, or the underlying gRPC error on failure. 32 | 33 | pub async fn upsert( 34 | &mut self, 35 | namespace: &str, 36 | vectors: &[Vector], 37 | batch_size: Option, 38 | ) -> PineconeResult { 39 | if batch_size.is_some() { 40 | todo!("Add proper upsert batching") 41 | } 42 | 43 | let upserted_count = self.dataplane_client.upsert(namespace, vectors).await?; 44 | 45 | if upserted_count != vectors.len() as u32 { 46 | return Err(PineconeClientError::Other(format!( 47 | "Failed to upsert all vectors. Upserted {} out of {} vectors", 48 | upserted_count, 49 | vectors.len() 50 | ))); 51 | } 52 | 53 | Ok(UpsertResponse { upserted_count }) 54 | } 55 | 56 | /// Query 57 | /// 58 | /// The `Query` operation searches a namespace, using a query vector. 59 | /// It retrieves the ids of the most similar items in a namespace, along with their similarity scores. 60 | /// To query by the id of already upserted vector, use `Index.query_by_id()` 61 | /// 62 | /// # Arguments 63 | /// - `namespace` - the name of the namespace in which vectors will be queried 64 | /// - `values` - The values for a new, unseen query vector. This should be the same length as the dimension of the index being queried. The results will be the `top_k` vectors closest to the given vector. Can not be used together with `id` 65 | /// - `sparse_values` - The query vector's sparse values. 66 | /// - `top_k` - The number of results to return for each query. 67 | /// - `filter` - The filter to apply. You can use vector metadata to limit your search. See 68 | /// - `include_values` - Indicates whether vector values are included in the response. 69 | /// - `include_metadata` - Indicates whether metadata is included in the response as well as the ids. 70 | /// 71 | /// # Returns 72 | /// A list of QueryResults 73 | #[allow(clippy::too_many_arguments)] 74 | pub async fn query( 75 | &mut self, 76 | namespace: &str, 77 | values: Option>, 78 | sparse_values: Option, 79 | top_k: u32, 80 | filter: Option>, 81 | include_values: bool, 82 | include_metadata: bool, 83 | ) -> PineconeResult> { 84 | let res = self 85 | .dataplane_client 86 | .query( 87 | namespace, 88 | None, 89 | values, 90 | sparse_values, 91 | top_k, 92 | filter, 93 | include_values, 94 | include_metadata, 95 | ) 96 | .await?; 97 | 98 | Ok(res) 99 | } 100 | 101 | /// Query by id 102 | /// 103 | /// The `Query by id` operation searches a namespace given the `id` of a vector already residing in the Index. 104 | /// It retrieves the ids of the most similar items in a namespace, along with their similarity scores. 105 | /// To query by new unseen vector use `Index.query()` 106 | /// 107 | /// # Arguments 108 | /// - `namespace` - the name of the namespace in which vectors will be queried 109 | /// - `id` - An id of a vector already upserted to the relevant namespace. The results will be the `top_k` nearest neighbours of the vector with the given id. Can not be used together with `values`. 110 | /// - `top_k` - The number of results to return for each query. 111 | /// - `filter` - The filter to apply. You can use vector metadata to limit your search. See 112 | /// - `include_values` - Indicates whether vector values are included in the response. 113 | /// - `include_metadata` - Indicates whether metadata is included in the response as well as the ids. 114 | /// 115 | /// # Returns 116 | /// A list QueryResults 117 | pub async fn query_by_id( 118 | &mut self, 119 | namespace: &str, 120 | id: &str, 121 | top_k: u32, 122 | filter: Option>, 123 | include_values: bool, 124 | include_metadata: bool, 125 | ) -> PineconeResult> { 126 | let res = self 127 | .dataplane_client 128 | .query( 129 | namespace, 130 | Some(id.into()), 131 | None, 132 | None, 133 | top_k, 134 | filter, 135 | include_values, 136 | include_metadata, 137 | ) 138 | .await?; 139 | 140 | Ok(res) 141 | } 142 | 143 | /// Describe index stats 144 | /// 145 | /// The DescribeIndexStats operation returns the number of vectors present in the index, for all the namespaces 146 | /// and the fullness of the index. Can also accept a filter to count the number of vectors matching the filter. 147 | /// 148 | /// # Arguments 149 | /// - `filter` - Optional filter to apply to the stats call. When applied, the stats only refer to matching vectors. 150 | /// 151 | /// # Returns 152 | /// A map of number of vectors per namespace, total vectors and the index fulness. 153 | pub async fn describe_index_stats( 154 | &mut self, 155 | filter: Option>, 156 | ) -> PineconeResult { 157 | let res = self.dataplane_client.describe_index_stats(filter).await?; 158 | Ok(res) 159 | } 160 | 161 | /// Fetch 162 | /// 163 | /// The Fetch operation retrieves the vectors with the given ids from the index. 164 | /// 165 | /// # Arguments 166 | /// - `namespace` - the name of the namespace in which vectors will be fetched 167 | /// - `ids` - A list of ids of vectors already upserted to the relevant namespace. 168 | /// 169 | pub async fn fetch( 170 | &mut self, 171 | namespace: &str, 172 | ids: &[String], 173 | ) -> PineconeResult> { 174 | let res = self.dataplane_client.fetch(namespace, ids).await?; 175 | Ok(res) 176 | } 177 | 178 | /// Update 179 | /// The update operation updates a single vector in the index. 180 | /// 181 | /// # Arguments 182 | /// - `id` - The id of the vector to be updated 183 | /// - `values` - Optional new values for the vector 184 | /// - `set_metadata` - Optional new metadata keys and values to be updated 185 | /// - `namespace` - The name of the namespace in which vectors will be updated 186 | /// 187 | pub async fn update( 188 | &mut self, 189 | id: &str, 190 | values: Option<&Vec>, 191 | sparse_values: Option, 192 | set_metadata: Option>, 193 | namespace: &str, 194 | ) -> PineconeResult<()> { 195 | self.dataplane_client 196 | .update(id, values, sparse_values, set_metadata, namespace) 197 | .await?; 198 | Ok(()) 199 | } 200 | 201 | /// Delete 202 | /// The delete operation deletes a list of vectors from a given namespace. 203 | /// 204 | /// # Arguments 205 | /// - `ids` - ids of the vectors to be deleted 206 | /// - `namespace` - the name of the namespace in which vectors will be deleted 207 | /// 208 | pub async fn delete(&mut self, ids: Vec, namespace: &str) -> PineconeResult<()> { 209 | self.dataplane_client 210 | .delete(Some(ids), namespace, None, false) 211 | .await?; 212 | Ok(()) 213 | } 214 | 215 | /// Delete by filter 216 | /// The delete by filter operation deletes a list of vectors from a given namespace that match the filter. 217 | /// 218 | /// # Arguments 219 | /// - `filter` - filter to be applied to delete the vectors 220 | /// - `namespace` - the name of the namespace in which vectors will be deleted 221 | /// 222 | pub async fn delete_by_metadata( 223 | &mut self, 224 | filter: Option>, 225 | namespace: &str, 226 | ) -> PineconeResult<()> { 227 | self.dataplane_client 228 | .delete(None, namespace, filter, false) 229 | .await?; 230 | Ok(()) 231 | } 232 | 233 | /// Delete all 234 | /// The delete all operation deletes all the vectors from a given namespace. 235 | /// 236 | /// # Arguments 237 | /// - `namespace` - the name of the namespace in which vectors will be deleted 238 | /// 239 | pub async fn delete_all(&mut self, namespace: &str) -> PineconeResult<()> { 240 | self.dataplane_client 241 | .delete(None, namespace, None, true) 242 | .await?; 243 | Ok(()) 244 | } 245 | } 246 | -------------------------------------------------------------------------------- /client_sdk/src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod client; 2 | pub mod data_types; 3 | pub mod index; 4 | pub mod utils; 5 | -------------------------------------------------------------------------------- /client_sdk/src/proto/vector_service.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | option java_package = "io.pinecone.proto"; 4 | option java_multiple_files = true; 5 | 6 | import "google/protobuf/struct.proto"; 7 | 8 | // Vector sparse data. Represented as a list of indices and a list of corresponded values, which must be with the same length. 9 | message SparseValues { 10 | 11 | // The indices of the sparse data. 12 | repeated uint32 indices = 1; 13 | 14 | // The corresponding values of the sparse data, which must be with the same length as the indices. 15 | repeated float values = 2; 16 | } 17 | 18 | message Vector { 19 | 20 | // This is the vector's unique id. 21 | string id = 1; 22 | 23 | // This is the vector data included in the request. 24 | repeated float values = 2; 25 | 26 | // This is the vector sparse data 27 | SparseValues sparse_values = 4; 28 | 29 | // This is the metadata included in the request. 30 | google.protobuf.Struct metadata = 3; 31 | } 32 | 33 | message ScoredVector { 34 | 35 | // This is the vector's unique id. 36 | string id = 1; 37 | 38 | // This is a measure of similarity between this vector and the query vector. The higher the score, the more they are similar. 39 | float score = 2; 40 | 41 | // This is the vector data, if it is requested. 42 | repeated float values = 3; 43 | 44 | // This is the sparse data, if it is requested. 45 | SparseValues sparse_values = 5; 46 | 47 | // This is the metadata, if it is requested. 48 | google.protobuf.Struct metadata = 4; 49 | } 50 | 51 | // The request for the `Upsert` operation. 52 | message UpsertRequest { 53 | // This is the actual array data. 54 | repeated Vector vectors = 1; 55 | 56 | // This is the namespace name where you upsert vectors. 57 | string namespace = 2; 58 | } 59 | 60 | // The response for the `Upsert` operation. 61 | message UpsertResponse { 62 | // The number of vectors upserted. 63 | uint32 upserted_count = 1; 64 | } 65 | 66 | // The request for the `Delete` operation. 67 | message DeleteRequest { 68 | // Vectors to delete. 69 | repeated string ids = 1; 70 | 71 | // This indicates that all vectors in the index namespace should be deleted. 72 | bool delete_all = 2; 73 | 74 | // The namespace to delete vectors from, if applicable. 75 | string namespace = 3; 76 | 77 | // If specified, the metadata filter here will be used to select the vectors to delete. This is mutually exclusive 78 | // with specifying ids to delete in the ids param or using delete_all=True. 79 | // See https://www.pinecone.io/docs/metadata-filtering/` 80 | google.protobuf.Struct filter = 4; 81 | } 82 | 83 | // The response for the `Delete` operation. 84 | message DeleteResponse {} 85 | 86 | // The request for the `Fetch` operation. 87 | message FetchRequest { 88 | // The vector ids to fetch. 89 | repeated string ids = 1; 90 | 91 | string namespace = 2; 92 | } 93 | 94 | // The response for the `Fetch` operation. 95 | message FetchResponse { 96 | // The fetched vectors, in the form of a map between the fetched ids and the fetched vectors 97 | map vectors = 1; 98 | 99 | // The namespace of the vectors. 100 | string namespace = 2; 101 | } 102 | 103 | // A single query vector within a `QueryRequest`. 104 | message QueryVector { 105 | 106 | // The query vector values. This should be the same length as the dimension of the index being queried. 107 | repeated float values = 1; 108 | 109 | // The query sparse values. 110 | SparseValues sparse_values = 5; 111 | 112 | // An override for the number of results to return for this query vector. 113 | uint32 top_k = 2; 114 | 115 | // An override the namespace to search. 116 | string namespace = 3; 117 | 118 | // An override for the metadata filter to apply. This replaces the request-level filter. 119 | google.protobuf.Struct filter = 4; 120 | } 121 | 122 | // The request for the `Query` operation. 123 | message QueryRequest { 124 | // The namespace to query. 125 | string namespace = 1; 126 | 127 | // The number of results to return for each query. 128 | uint32 top_k = 2; 129 | 130 | // The filter to apply if no filter is specified in `QueryVector. You can use vector metadata to limit your search. See https://www.pinecone.io/docs/metadata-filtering/` 131 | google.protobuf.Struct filter = 3; 132 | 133 | // Indicates whether vector values are included in the response. 134 | bool include_values = 4; 135 | 136 | // Indicates whether metadata is included in the response as well as the ids. 137 | bool include_metadata = 5; 138 | 139 | // The query vectors. 140 | repeated QueryVector queries = 6; 141 | 142 | // The query vector. This should be the same length as the dimension of the index being queried. 143 | repeated float vector = 7; 144 | 145 | // The query sparse values. 146 | SparseValues sparse_vector = 9; 147 | 148 | // The id of the vector 149 | string id = 8; 150 | } 151 | 152 | // The query results for a single `QueryVector` 153 | message SingleQueryResults { 154 | // The matches for the vectors. 155 | repeated ScoredVector matches = 1; 156 | 157 | // The namespace for the vectors. 158 | string namespace = 2; 159 | } 160 | 161 | // The response for the `Query` operation. These are the matches found for a particular query vector. The matches are ordered from most similar to least similar. 162 | message QueryResponse { 163 | // The results of each query. The order is the same as `QueryRequest.queries`. 164 | repeated SingleQueryResults results = 1; 165 | 166 | // The matches for the vectors. 167 | repeated ScoredVector matches = 2; 168 | 169 | // The namespace for the vectors. 170 | string namespace = 3; 171 | } 172 | 173 | // The request for the `Upsert` operation. 174 | message UpdateRequest { 175 | // Vector's unique id. 176 | string id = 1; 177 | 178 | // Vector data. 179 | repeated float values = 2; 180 | 181 | // Vector sparse values 182 | SparseValues sparse_values = 5; 183 | 184 | // Metadata to *set* for the vector. 185 | google.protobuf.Struct set_metadata = 3; 186 | 187 | // Namespace name where to update the vector. 188 | string namespace = 4; 189 | } 190 | 191 | // The response for the `Update` operation. 192 | message UpdateResponse { 193 | } 194 | 195 | // The request for the `DescribeIndexStats` operation. 196 | message DescribeIndexStatsRequest { 197 | // If specified, the metadata filter here will be used to select the vectors to get stats about. 198 | // See https://www.pinecone.io/docs/metadata-filtering/` 199 | google.protobuf.Struct filter = 1; 200 | } 201 | 202 | // A summary of the contents of a namespace. 203 | message NamespaceSummary { 204 | // The number of vectors stored in this namespace. Note that updates to this field may lag behind updates to the 205 | // underlying index and corresponding query results, etc. 206 | uint32 vector_count = 1; 207 | } 208 | 209 | // The response for the `DescribeIndexStats` operation. 210 | message DescribeIndexStatsResponse { 211 | 212 | // A mapping for each namespace in the index from namespace name to a summary of its contents. 213 | map namespaces = 1; 214 | 215 | // The dimension of the indexed vectors. 216 | uint32 dimension = 2 ; 217 | 218 | // The fullness of the index. 219 | float index_fullness = 3 ; 220 | 221 | // The total number of vectors in the index 222 | uint32 total_vector_count = 4 ; 223 | } 224 | 225 | // The `VectorService` interface is exposed by Pinecone's vector index services. 226 | // This service could also be called a `gRPC` service or a `REST`-like api. 227 | service VectorService { 228 | // Upsert 229 | // 230 | // The `Upsert` operation writes vectors into a namespace. 231 | // If a new value is upserted for an existing vector id, it will overwrite the previous value. 232 | rpc Upsert(UpsertRequest) returns (UpsertResponse) { 233 | } 234 | 235 | // Delete 236 | // 237 | // The `Delete` operation deletes vectors, by id, from a single namespace. 238 | // You can delete items by their id, from a single namespace. 239 | rpc Delete(DeleteRequest) returns (DeleteResponse) { 240 | } 241 | 242 | // Fetch 243 | // 244 | // The `Fetch` operation looks up and returns vectors, by id, from a single namespace. 245 | // The returned vectors include the vector data and/or metadata. 246 | rpc Fetch(FetchRequest) returns (FetchResponse) { 247 | } 248 | 249 | // Query 250 | // 251 | // The `Query` operation searches a namespace, using one or more query vectors. 252 | // It retrieves the ids of the most similar items in a namespace, along with their similarity scores. 253 | rpc Query(QueryRequest) returns (QueryResponse) { 254 | } 255 | 256 | // Update 257 | // 258 | // The `Update` operation updates vector in a namespace. 259 | // If a value is included, it will overwrite the previous value. 260 | // If a set_metadata is included, the values of the fields specified in it will be added or overwrite the previous value. 261 | rpc Update(UpdateRequest) returns (UpdateResponse) { 262 | } 263 | 264 | // DescribeIndexStats 265 | // 266 | // The `DescribeIndexStats` operation returns statistics about the index's contents. 267 | // For example: The vector count per namespace and the number of dimensions. 268 | rpc DescribeIndexStats(DescribeIndexStatsRequest) returns (DescribeIndexStatsResponse) { 269 | } 270 | } -------------------------------------------------------------------------------- /client_sdk/src/utils/conversions.rs: -------------------------------------------------------------------------------- 1 | use crate::client::grpc::{GrpcScoredVector, GrpcSparseValues, GrpcVector}; 2 | use crate::data_types::{Collection, Db, MetadataValue, QueryResult, SparseValues, Vector}; 3 | use crate::utils::errors::PineconeClientError::{MetadataError, MetadataValueError}; 4 | use crate::utils::errors::{PineconeClientError, PineconeResult}; 5 | use index_service::models::IndexMetaStatus; 6 | use index_service::models::{ 7 | CollectionMeta, CreateCollectionRequest, CreateRequest, CreateRequestMetadataConfig, IndexMeta, 8 | }; 9 | use prost_types::value::Kind; 10 | use prost_types::{ListValue as ProstListValue, Struct, Value as ProstValue}; 11 | use std::collections::BTreeMap; 12 | 13 | impl From for GrpcSparseValues { 14 | fn from(value: SparseValues) -> Self { 15 | GrpcSparseValues { 16 | indices: value.indices, 17 | values: value.values, 18 | } 19 | } 20 | } 21 | 22 | impl From for SparseValues { 23 | fn from(value: GrpcSparseValues) -> Self { 24 | SparseValues { 25 | indices: value.indices, 26 | values: value.values, 27 | } 28 | } 29 | } 30 | 31 | impl TryFrom for MetadataValue { 32 | type Error = PineconeClientError; 33 | 34 | fn try_from(val: ProstValue) -> Result { 35 | if let Some(kind) = val.kind { 36 | match kind { 37 | Kind::NumberValue(v) => Ok(MetadataValue::NumberVal(v)), 38 | Kind::StringValue(v) => Ok(MetadataValue::StringVal(v)), 39 | Kind::BoolValue(v) => Ok(MetadataValue::BoolVal(v)), 40 | Kind::ListValue(v) => { 41 | let mut inners = Vec::new(); 42 | for item in v.values.into_iter() { 43 | let new_val = item.try_into().map_err(|e| match e { 44 | MetadataValueError { val_type } => MetadataValueError { 45 | val_type: format!("{val_type} value in a list"), 46 | }, 47 | _ => e, 48 | })?; 49 | inners.push(new_val); 50 | } 51 | Ok(MetadataValue::ListVal(inners)) 52 | } 53 | Kind::NullValue(_) => Err(MetadataValueError { 54 | val_type: "None".into(), 55 | }), 56 | Kind::StructValue(s) => { 57 | let mut inners = BTreeMap::new(); 58 | for (k, v) in s.fields { 59 | let new_val = v.try_into().map_err(|e| match e { 60 | MetadataValueError { val_type } => MetadataValueError { 61 | val_type: format!("{val_type} value in a dict"), 62 | }, 63 | MetadataError { key, val_type } => MetadataError { 64 | key: format!("{k}: {key}"), 65 | val_type: format!("{val_type} value in a dict"), 66 | }, 67 | _ => e, 68 | })?; 69 | inners.insert(k, new_val); 70 | } 71 | Ok(MetadataValue::DictVal(inners)) 72 | } 73 | } 74 | } else { 75 | Err(MetadataValueError { 76 | val_type: "empty".into(), 77 | }) 78 | } 79 | } 80 | } 81 | 82 | impl From for ProstValue { 83 | fn from(val: MetadataValue) -> Self { 84 | match val { 85 | MetadataValue::StringVal(v) => ProstValue { 86 | kind: Some(Kind::StringValue(v)), 87 | }, 88 | MetadataValue::NumberVal(v) => ProstValue { 89 | kind: Some(Kind::NumberValue(v)), 90 | }, 91 | MetadataValue::BoolVal(v) => ProstValue { 92 | kind: Some(Kind::BoolValue(v)), 93 | }, 94 | MetadataValue::ListVal(v) => { 95 | let new_list = v.into_iter().map(|x| x.into()).collect(); 96 | ProstValue { 97 | kind: Some(Kind::ListValue(ProstListValue { values: new_list })), 98 | } 99 | } 100 | MetadataValue::DictVal(v) => { 101 | let mut new_dict: BTreeMap = BTreeMap::new(); 102 | for (k, v) in v.into_iter() { 103 | new_dict.insert(k, v.into()); 104 | } 105 | let new_struct = Struct { fields: new_dict }; 106 | ProstValue { 107 | kind: Some(Kind::StructValue(new_struct)), 108 | } 109 | } 110 | } 111 | } 112 | } 113 | 114 | impl From for CreateRequest { 115 | fn from(index: Db) -> Self { 116 | CreateRequest { 117 | name: index.name, 118 | dimension: index.dimension, 119 | replicas: index.replicas, 120 | pod_type: index.pod_type, 121 | metric: index.metric, 122 | pods: index.pods, 123 | shards: index.shards, 124 | source_collection: index.source_collection, 125 | metadata_config: index.metadata_config.map(|config| { 126 | Some(Box::new(CreateRequestMetadataConfig { 127 | indexed: config.get("indexed").map(|v| v.to_vec()), 128 | })) 129 | }), 130 | ..Default::default() 131 | } 132 | } 133 | } 134 | 135 | impl TryFrom for Db { 136 | type Error = PineconeClientError; 137 | fn try_from(index_meta: IndexMeta) -> Result { 138 | let db = index_meta.database; 139 | let status = index_meta.status; 140 | let state = status.and_then(|inner_box| { 141 | let inner_struct: IndexMetaStatus = *inner_box; 142 | inner_struct.state 143 | }); 144 | match db { 145 | Some(db) => { 146 | let name = db.name.ok_or_else(|| { 147 | PineconeClientError::ValueError("Failed to parse db name".to_string()) 148 | })?; 149 | let replicas = db.replicas; 150 | let shards = db.shards; 151 | let pod_type = db.pod_type; 152 | let dimension = db.dimension.ok_or_else(|| { 153 | PineconeClientError::Other("Failed to parse db dimension".to_string()) 154 | })?; 155 | let metric = db.metric; 156 | let pods = db.pods; 157 | let source_collection = db.source_collection; 158 | let metadata_config = db.metadata_config.map(|config| { 159 | let indexed = config.indexed.unwrap_or_default(); 160 | let mut map = BTreeMap::new(); 161 | map.insert("indexed".to_string(), indexed); 162 | map 163 | }); 164 | let status = state; 165 | Ok(Db { 166 | name, 167 | dimension, 168 | replicas, 169 | shards, 170 | pod_type, 171 | metric, 172 | pods, 173 | source_collection, 174 | metadata_config, 175 | status, 176 | }) 177 | } 178 | None => Err(PineconeClientError::Other("Failed to parse db".to_string())), 179 | } 180 | } 181 | } 182 | 183 | impl From for CreateCollectionRequest { 184 | fn from(collection: Collection) -> Self { 185 | CreateCollectionRequest { 186 | name: collection.name, 187 | source: collection.source, 188 | } 189 | } 190 | } 191 | 192 | impl From for Collection { 193 | fn from(collection_meta: CollectionMeta) -> Self { 194 | Collection { 195 | name: collection_meta.name.unwrap(), 196 | source: "".to_string(), 197 | vector_count: None, 198 | size: collection_meta.size, 199 | status: collection_meta.status, 200 | } 201 | } 202 | } 203 | 204 | pub fn hashmap_to_prost_struct(dict: BTreeMap) -> Struct { 205 | let mut fields = BTreeMap::new(); 206 | for (k, v) in dict.into_iter() { 207 | fields.insert(k, v.into()); 208 | } 209 | Struct { fields } 210 | } 211 | 212 | pub fn prost_struct_to_hashmap(dict: Struct) -> PineconeResult> { 213 | let mut fields: BTreeMap = BTreeMap::new(); 214 | for (k, v) in dict.fields.into_iter() { 215 | let new_val = v.try_into().map_err(|e| match e { 216 | PineconeClientError::MetadataValueError { val_type } => { 217 | PineconeClientError::MetadataError { 218 | key: k.clone(), 219 | val_type, 220 | } 221 | } 222 | _ => e, 223 | })?; 224 | fields.insert(k, new_val); 225 | } 226 | Ok(fields) 227 | } 228 | 229 | impl From for GrpcVector { 230 | fn from(grpc_vector: Vector) -> Self { 231 | GrpcVector { 232 | id: grpc_vector.id, 233 | values: grpc_vector.values, 234 | sparse_values: grpc_vector 235 | .sparse_values 236 | .map(|sparse_vector| sparse_vector.into()), 237 | metadata: grpc_vector.metadata.map(hashmap_to_prost_struct), 238 | } 239 | } 240 | } 241 | 242 | impl TryFrom for Vector { 243 | type Error = PineconeClientError; 244 | 245 | fn try_from(grpc_vector: GrpcVector) -> Result { 246 | Ok(Vector { 247 | id: grpc_vector.id, 248 | values: grpc_vector.values, 249 | sparse_values: grpc_vector 250 | .sparse_values 251 | .map(|sparse_vector| sparse_vector.into()), 252 | metadata: grpc_vector 253 | .metadata 254 | .map(prost_struct_to_hashmap) 255 | .transpose()?, 256 | }) 257 | } 258 | } 259 | 260 | impl TryFrom for QueryResult { 261 | type Error = PineconeClientError; 262 | 263 | fn try_from(grpc_vector: GrpcScoredVector) -> Result { 264 | Ok(QueryResult { 265 | id: grpc_vector.id, 266 | score: grpc_vector.score, 267 | values: if grpc_vector.values.is_empty() { 268 | None 269 | } else { 270 | Some(grpc_vector.values) 271 | }, 272 | sparse_values: grpc_vector 273 | .sparse_values 274 | .map(|sparse_vector| sparse_vector.into()), 275 | metadata: grpc_vector 276 | .metadata 277 | .map(prost_struct_to_hashmap) 278 | .transpose()?, 279 | }) 280 | } 281 | } 282 | -------------------------------------------------------------------------------- /client_sdk/src/utils/errors.rs: -------------------------------------------------------------------------------- 1 | use std::fmt::Debug; 2 | use thiserror::Error; 3 | 4 | #[derive(Error, Debug)] 5 | pub enum PineconeClientError { 6 | #[error("Invalid value for argument {name}: {found:?})")] 7 | ArgumentError { name: String, found: String }, 8 | 9 | #[error("`{0}`")] 10 | ValueError(String), 11 | 12 | #[error("Error in vector number {vec_num}: Missing key '{key}'")] 13 | UpsertKeyError { key: String, vec_num: usize }, 14 | 15 | #[error("Error in vector number {vec_num}: Found unexpected value for '{key}'. Expected a {expected_type}, found: {actual}")] 16 | UpsertValueError { 17 | key: String, 18 | vec_num: usize, 19 | expected_type: String, 20 | actual: String, 21 | }, 22 | 23 | #[error("Failed to connect to Pinecone's controller on region {region}. Please verify client configuration: API key, region and project_id. \ 24 | See more info: https://docs.pinecone.io/docs/quickstart#2-get-and-verify-your-pinecone-api-key\n\ 25 | Underlying Error: {err}")] 26 | ControlPlaneConnectionError { region: String, err: String }, 27 | 28 | #[error("Failed to connect to index '{index}'. Please verify that an index with that name exists using `client.list_indexes()`. \n\ 29 | Underlying Error: {err}")] 30 | IndexConnectionError { index: String, err: String }, 31 | 32 | #[error(transparent)] 33 | DataplaneOperationError(#[from] tonic::Status), 34 | 35 | #[error(transparent)] 36 | IoError(#[from] std::io::Error), 37 | 38 | #[error("Unsupported metadata value {val_type}. \ 39 | Please see https://docs.pinecone.io/docs/metadata-filtering#supported-metadata-types for allowed metadata types")] 40 | MetadataValueError { val_type: String }, 41 | 42 | #[error("Unsupported metadata value for key {key}: found value of type {val_type}. \ 43 | Please see https://docs.pinecone.io/docs/metadata-filtering#supported-metadata-types for allowed metadata types")] 44 | MetadataError { key: String, val_type: String }, 45 | 46 | #[error("`{0}`")] 47 | Other(String), 48 | 49 | #[error("Operation failed with error code {status_code }. \nUnderlying Error: {err}")] 50 | ControlPlaneOperationError { err: String, status_code: String }, 51 | 52 | #[error("Failed to parse response contents")] 53 | ControlPlaneParsingError {}, 54 | 55 | #[error(transparent)] 56 | DeserializationError(#[from] serde_json::Error), 57 | 58 | #[error("`{0}`")] 59 | KeyboardInterrupt(String), 60 | } 61 | 62 | // TODO: Decide if we want to print the full formatted error on dubug 63 | // impl std::fmt::Debug for PineconeClientError { 64 | // fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 65 | // write!(f, "{}", self.to_string()) 66 | // } 67 | // } 68 | 69 | pub type PineconeResult = Result; 70 | 71 | impl From> for PineconeClientError { 72 | fn from(err: index_service::apis::Error) -> Self { 73 | match err { 74 | index_service::apis::Error::ResponseError(response_error) => { 75 | PineconeClientError::ControlPlaneOperationError { 76 | err: response_error.content, 77 | status_code: response_error.status.to_string(), 78 | } 79 | } 80 | index_service::apis::Error::Reqwest(reqwest_error) => { 81 | if reqwest_error.is_connect() { 82 | PineconeClientError::ControlPlaneConnectionError { 83 | region: "".into(), 84 | err: reqwest_error.to_string(), 85 | } 86 | } else { 87 | PineconeClientError::ControlPlaneOperationError { 88 | err: reqwest_error.to_string(), 89 | status_code: match reqwest_error.status() { 90 | None => "unknown".into(), 91 | Some(c) => c.to_string(), 92 | }, 93 | } 94 | } 95 | } 96 | err => PineconeClientError::Other(err.to_string()), 97 | } 98 | } 99 | } 100 | -------------------------------------------------------------------------------- /client_sdk/src/utils/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod conversions; 2 | pub mod errors; 3 | pub mod python_conversions; 4 | -------------------------------------------------------------------------------- /client_sdk/src/utils/python_conversions.rs: -------------------------------------------------------------------------------- 1 | use crate::data_types::{MetadataValue, NamespaceStats, SparseValues, Vector}; 2 | use crate::utils::errors::PineconeClientError; 3 | use pyo3::types::{IntoPyDict, PyDict}; 4 | use pyo3::{IntoPy, PyObject, Python, ToPyObject}; 5 | use std::collections::{BTreeMap, HashSet}; 6 | 7 | const SPARSE_KEYS: &[&str] = &["indices", "values"]; 8 | const VECTOR_KEYS: &[&str] = &["id", "values", "sparse_values", "metadata"]; 9 | 10 | impl TryFrom<&PyDict> for SparseValues { 11 | type Error = PineconeClientError; 12 | 13 | fn try_from(dict: &PyDict) -> Result { 14 | let allowed_keys: HashSet = SPARSE_KEYS.iter().map(|x| (*x).into()).collect(); 15 | let actual_keys: HashSet = dict 16 | .keys() 17 | .into_iter() 18 | .map(|x| x.extract::()) 19 | .collect::, _>>() 20 | .map_err(|_| { 21 | PineconeClientError::ValueError("Couldn't retrieve dictionary keys".into()) 22 | })?; 23 | 24 | let excess_keys = actual_keys 25 | .difference(&allowed_keys) 26 | .collect::>(); 27 | if !excess_keys.is_empty() { 28 | return Err(PineconeClientError::ValueError(format!( 29 | "Found unexpected keys: {excess_keys:?}", 30 | excess_keys = excess_keys 31 | ))); 32 | } 33 | 34 | let indices = match dict.get_item("indices") { 35 | None => { 36 | return Err(PineconeClientError::UpsertKeyError { 37 | key: "indices".into(), 38 | vec_num: 0, 39 | }) 40 | } 41 | Some(v) => { 42 | v.extract::>() 43 | .map_err(|_| PineconeClientError::UpsertValueError { 44 | key: "indices".into(), 45 | vec_num: 0, 46 | expected_type: "List[int]".into(), 47 | actual: format!("{:?}", v), 48 | })? 49 | } 50 | }; 51 | 52 | let values = match dict.get_item("values") { 53 | None => { 54 | return Err(PineconeClientError::UpsertKeyError { 55 | key: "values".into(), 56 | vec_num: 0, 57 | }) 58 | } 59 | Some(v) => { 60 | v.extract::>() 61 | .map_err(|_| PineconeClientError::UpsertValueError { 62 | key: "values".into(), 63 | vec_num: 0, 64 | expected_type: "List[float]".into(), 65 | actual: format!("{:?}", v), 66 | })? 67 | } 68 | }; 69 | 70 | Ok(SparseValues { indices, values }) 71 | } 72 | } 73 | 74 | impl ToPyObject for SparseValues { 75 | fn to_object(&self, py: Python) -> PyObject { 76 | let dict = [ 77 | ("indices", self.indices.to_object(py)), 78 | ("values", self.values.to_object(py)), 79 | ] 80 | .into_py_dict(py); 81 | dict.to_object(py) 82 | } 83 | } 84 | 85 | impl TryFrom<&PyDict> for Vector { 86 | type Error = PineconeClientError; 87 | 88 | fn try_from(dict: &PyDict) -> Result { 89 | let allowed_keys: HashSet = VECTOR_KEYS.iter().map(|x| (*x).into()).collect(); 90 | let actual_keys: HashSet = dict 91 | .keys() 92 | .into_iter() 93 | .map(|x| x.extract::()) 94 | .collect::, _>>() 95 | .map_err(|_| { 96 | PineconeClientError::ValueError("Couldn't retrieve dictionary keys".into()) 97 | })?; 98 | 99 | let excess_keys = actual_keys 100 | .difference(&allowed_keys) 101 | .collect::>(); 102 | if !excess_keys.is_empty() { 103 | return Err(PineconeClientError::ValueError(format!( 104 | "Found unexpected keys: {excess_keys:?}", 105 | excess_keys = excess_keys 106 | ))); 107 | } 108 | 109 | Ok(Vector { 110 | id: match dict.get_item("id") { 111 | None => { 112 | return Err(PineconeClientError::UpsertKeyError { 113 | key: "id".into(), 114 | vec_num: 0, 115 | }) 116 | } 117 | Some(id) => { 118 | id.extract::() 119 | .map_err(|_| PineconeClientError::UpsertValueError { 120 | key: "id".into(), 121 | vec_num: 0, 122 | expected_type: "String".into(), 123 | actual: format!("{:?}", id), 124 | }) 125 | } 126 | }?, 127 | values: match dict.get_item("values") { 128 | None => { 129 | return Err(PineconeClientError::UpsertKeyError { 130 | key: "values".into(), 131 | vec_num: 0, 132 | }) 133 | } 134 | Some(values) => values.extract::>().map_err(|_| { 135 | PineconeClientError::UpsertValueError { 136 | key: "values".into(), 137 | vec_num: 0, 138 | expected_type: "List[float]".into(), 139 | actual: format!("{:?}", values), 140 | } 141 | })?, 142 | }, 143 | sparse_values: dict 144 | .get_item("sparse_values") 145 | .map(|val| { 146 | let val = val.extract::<&PyDict>().map_err(|_| { 147 | PineconeClientError::UpsertValueError { 148 | key: "sparse_values".into(), 149 | vec_num: 0, 150 | expected_type: "dict".into(), 151 | actual: format!("{:?}", val), 152 | } 153 | })?; 154 | val.try_into().map_err(|e| match e { 155 | PineconeClientError::UpsertKeyError { key, vec_num } => { 156 | PineconeClientError::UpsertKeyError { 157 | key: format!("sparse_values: {key}", key = key), 158 | vec_num, 159 | } 160 | } 161 | PineconeClientError::UpsertValueError { 162 | key, 163 | vec_num, 164 | actual, 165 | expected_type, 166 | } => PineconeClientError::UpsertValueError { 167 | key: format!("sparse_values: {key}", key = key), 168 | vec_num, 169 | actual, 170 | expected_type, 171 | }, 172 | _ => PineconeClientError::ValueError(format!( 173 | "Error in 'sparse_values: {e}", 174 | e = e 175 | )), 176 | }) 177 | }) 178 | .transpose()?, 179 | metadata: dict 180 | .get_item("metadata") 181 | .map(|val| { 182 | val.extract::>() 183 | .map_err(|_| PineconeClientError::UpsertValueError { 184 | key: "metadata".into(), 185 | vec_num: 0, 186 | expected_type: "dict".into(), 187 | actual: format!("{:?}", val), 188 | }) 189 | }) 190 | .transpose()?, 191 | }) 192 | } 193 | } 194 | 195 | impl ToPyObject for NamespaceStats { 196 | fn to_object(&self, py: Python) -> PyObject { 197 | self.to_dict(py).to_object(py) 198 | } 199 | } 200 | 201 | impl ToPyObject for MetadataValue { 202 | fn to_object(&self, py: Python<'_>) -> PyObject { 203 | match self { 204 | MetadataValue::StringVal(v) => v.to_object(py), 205 | MetadataValue::NumberVal(v) => v.to_object(py), 206 | MetadataValue::BoolVal(v) => v.to_object(py), 207 | MetadataValue::ListVal(v) => v.to_object(py), 208 | MetadataValue::DictVal(v) => v.to_object(py), 209 | } 210 | } 211 | } 212 | 213 | impl IntoPy for MetadataValue { 214 | fn into_py(self, py: Python<'_>) -> PyObject { 215 | match self { 216 | MetadataValue::StringVal(v) => v.to_object(py), 217 | MetadataValue::NumberVal(v) => v.to_object(py), 218 | MetadataValue::ListVal(v) => v.to_object(py), 219 | MetadataValue::BoolVal(v) => v.to_object(py), 220 | MetadataValue::DictVal(v) => v.to_object(py), 221 | } 222 | } 223 | } 224 | -------------------------------------------------------------------------------- /config.toml: -------------------------------------------------------------------------------- 1 | [target.x86_64-apple-darwin] 2 | rustflags = [ 3 | "-C", "link-arg=-undefined", 4 | "-C", "link-arg=dynamic_lookup", 5 | ] 6 | 7 | [target.aarch64-apple-darwin] 8 | rustflags = [ 9 | "-C", "link-arg=-undefined", 10 | "-C", "link-arg=dynamic_lookup", 11 | ] -------------------------------------------------------------------------------- /index_service/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "index_service" 3 | version = "0.1.0" 4 | authors = ["OpenAPI Generator team and contributors"] 5 | description = "defaultDescription" 6 | # Override this license by providing a License Object in the OpenAPI. 7 | license = "Unlicense" 8 | edition = "2018" 9 | 10 | [dependencies] 11 | serde = "^1.0" 12 | serde_derive = "^1.0" 13 | serde_with = "^2.0" 14 | serde_json = "^1.0" 15 | url = "^2.2" 16 | uuid = { version = "^1.0", features = ["serde"] } 17 | [dependencies.reqwest] 18 | version = "^0.11" 19 | features = ["json", "multipart"] 20 | -------------------------------------------------------------------------------- /index_service/build.rs: -------------------------------------------------------------------------------- 1 | use std::error::Error; 2 | use std::path::Path; 3 | use std::process::Command; 4 | 5 | fn main() -> Result<(), Box> { 6 | let path = Command::new("cargo").args(["locate-project", "--workspace", "--message-format", "plain"]).output()?; 7 | let path = String::from_utf8( path.stdout)?; 8 | let path = Path::new(&path).parent().unwrap(); 9 | let output = Command::new("make").current_dir(path).args(["generate-index-service"]).output()?; 10 | 11 | //If a previous auto-generated output already exists, don't error out. 12 | //This is needed for CI build on macOS and windows where `docker` is expected to fail 13 | if !output.status.success() { 14 | if path.join("index_service").join("src").join("apis").join("index_operations_api.rs").exists() { 15 | eprintln!("Warning - failed to generate OpenAPI: {output:?}. Found existing index_operations_api.rs. Continuing anyway..."); 16 | } else { 17 | eprintln!("Failed to generate OpenAPI: {output:?}"); 18 | return Err("process failed".into()); 19 | } 20 | } 21 | Ok(()) 22 | } -------------------------------------------------------------------------------- /index_service/src/lib.rs: -------------------------------------------------------------------------------- 1 | #[macro_use] 2 | extern crate serde_derive; 3 | 4 | extern crate serde; 5 | extern crate serde_json; 6 | extern crate url; 7 | extern crate reqwest; 8 | 9 | pub mod apis; 10 | pub mod models; 11 | -------------------------------------------------------------------------------- /pinecone/.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | .pytest_cache/ 6 | *.py[cod] 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | .venv/ 14 | env/ 15 | bin/ 16 | build/ 17 | develop-eggs/ 18 | dist/ 19 | eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | include/ 26 | man/ 27 | venv/ 28 | *.egg-info/ 29 | .installed.cfg 30 | *.egg 31 | 32 | # Installer logs 33 | pip-log.txt 34 | pip-delete-this-directory.txt 35 | pip-selfcheck.json 36 | 37 | # Unit test / coverage reports 38 | htmlcov/ 39 | .tox/ 40 | .coverage 41 | .cache 42 | nosetests.xml 43 | coverage.xml 44 | 45 | # Translations 46 | *.mo 47 | 48 | # Mr Developer 49 | .mr.developer.cfg 50 | .project 51 | .pydevproject 52 | 53 | # Rope 54 | .ropeproject 55 | 56 | # Django stuff: 57 | *.log 58 | *.pot 59 | 60 | .DS_Store 61 | 62 | # Sphinx documentation 63 | docs/_build/ 64 | 65 | # PyCharm 66 | .idea/ 67 | 68 | # VSCode 69 | .vscode/ 70 | 71 | # Pyenv 72 | .python-version -------------------------------------------------------------------------------- /pinecone/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "pinecone" 3 | version.workspace = true 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | [lib] 8 | name = "pinecone" 9 | crate-type = ["cdylib"] 10 | bench=false 11 | 12 | 13 | 14 | [dependencies] 15 | pyo3 = { version = "0.18.0", features = ["extension-module"] } 16 | client_sdk = {path = "../client_sdk" } 17 | tokio = { version = "1.16.1", features = ["rt-multi-thread"] } 18 | reqwest = { version = "0.11.6", features = ["json"] } 19 | pyo3-asyncio = {version = "0.18.0", features = ["tokio-runtime"]} 20 | -------------------------------------------------------------------------------- /pinecone/pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["maturin>=0.14,<0.15"] 3 | build-backend = "maturin" 4 | 5 | [project] 6 | version = "3.0.0rc2" # TODO: remove this key completely once we have a stable release. Version is taken automatically from Cargo.toml. 7 | name = "pinecone-client" 8 | description="Pinecone client and SDK" 9 | requires-python = ">=3.7" 10 | keywords=["Pinecone", "vector", "database", "cloud"] 11 | classifiers = [ 12 | "Programming Language :: Rust", 13 | "Programming Language :: Python :: Implementation :: CPython", 14 | "Programming Language :: Python :: Implementation :: PyPy", 15 | "Intended Audience :: Developers", 16 | "Intended Audience :: Information Technology", 17 | "Intended Audience :: Science/Research", 18 | "Intended Audience :: System Administrators", 19 | "License :: OSI Approved :: Apache Software License", 20 | "Topic :: Database", 21 | "Topic :: Software Development", 22 | "Topic :: Software Development :: Libraries", 23 | "Topic :: Software Development :: Libraries :: Application Frameworks", 24 | "Topic :: Software Development :: Libraries :: Python Modules", 25 | "Operating System :: OS Independent", 26 | ] 27 | 28 | [project.optional-dependencies] 29 | test = [ 30 | "loguru==0.6.0", 31 | "numpy==1.24.2", 32 | "pytest==6.2.5", 33 | "pytest-asyncio==0.16.0", 34 | "pytest-timeout==1.4.2", 35 | "pytest-html==3.1.1", 36 | "pytest-xdist[psutil]==2.5.0", 37 | "requests==2.25.1", 38 | "testcontainers==3.4.1", 39 | "inflection==0.5.1", 40 | ] 41 | 42 | 43 | -------------------------------------------------------------------------------- /pinecone/src/client.rs: -------------------------------------------------------------------------------- 1 | use std::collections::BTreeMap; 2 | 3 | use client_sdk::data_types::{Collection, Db}; 4 | use pyo3::prelude::*; 5 | use tokio::runtime::Runtime; 6 | 7 | use crate::index::Index; 8 | use crate::utils::errors::{PineconeClientError, PineconeResult}; 9 | use client_sdk::client::pinecone_client as core_client; 10 | use client_sdk::utils::errors::{self as core_errors}; 11 | 12 | #[pyclass] 13 | #[pyo3(text_signature = "(api_key=None, region=None, project_id=None)")] 14 | pub struct Client { 15 | inner: core_client::PineconeClient, 16 | runtime: Runtime, 17 | } 18 | 19 | #[pymethods] 20 | impl Client { 21 | #[new] 22 | #[pyo3(signature = (api_key=None, region=None, project_id=None))] 23 | /// Creates a Pinecone client instance. 24 | /// Configuration parameters are usually set as environment variables. If you want to override the environment variables, you can pass them as arguments to the constructor. 25 | /// 26 | /// Args: 27 | /// api_key (str, optional): The API key to use for authentication. Defaults to the value of the `PINECONE_API_KEY` environment variable. See more info here: https://docs.pinecone.io/docs/quickstart#2-get-and-verify-your-pinecone-api-key 28 | /// region (str, optional): The pinecone region to use. Defaults to the value of the `PINECONE_REGION` environment variable, or to `us-west1-gcp` if the environment variable is not set. 29 | /// project_id (str, optional): By default, the client will use project id associated with the API key. If you want to use a different project id, you can pass it as an argument to the constructor. 30 | /// 31 | /// Returns: 32 | /// Client: A Pinecone client instance. 33 | pub fn new( 34 | api_key: Option<&str>, 35 | region: Option<&str>, 36 | project_id: Option<&str>, 37 | ) -> PineconeResult { 38 | let rt = Runtime::new().map_err(core_errors::PineconeClientError::IoError)?; 39 | let client = rt.block_on(core_client::PineconeClient::new( 40 | api_key, region, project_id, 41 | ))?; 42 | 43 | Ok(Self { 44 | inner: client, 45 | runtime: rt, 46 | }) 47 | } 48 | 49 | pub fn __repr__(&self) -> String { 50 | let api_key = self.inner.api_key.split('-').last().unwrap_or("None"); 51 | format!( 52 | "Client:\n API key: ****************-{api_key}\n region: {region}\n project_id: {project_id}", 53 | api_key = api_key, 54 | region = self.inner.region, 55 | project_id = self.inner.project_id 56 | ) 57 | } 58 | 59 | /// Index 60 | /// 61 | /// The Index is the main object for interacting with a Pinecone index. It is used to insert, update, fetch and query vectors. 62 | /// You can create an Index object by calling the `get_index` method on the Pinecone client. 63 | /// You can also create an Index object by calling the `create_index` method on the Pinecone client. 64 | /// This method is a shortcut for `get_index` for backwards compatibility and will eventually be deprecated. 65 | /// 66 | /// Args: 67 | /// name (str): The name an existing Pinecone index to connect to. 68 | /// 69 | /// Returns: 70 | /// Index: The index object. 71 | #[allow(non_snake_case)] 72 | pub fn Index(&self, name: &str) -> PineconeResult { 73 | self.get_index(name) 74 | } 75 | 76 | /// Get an Index object for interacting with a Pinecone index. 77 | /// 78 | /// The Index is the main object for interacting with a Pinecone index. It is used to insert, update, fetch and query vectors. 79 | /// You can create an Index object by calling the `get_index` method on the Pinecone client. 80 | /// You can also create an Index object by calling the `create_index` method on the Pinecone client. 81 | /// 82 | /// Args: 83 | /// name (str): The name an existing Pinecone index to connect to. 84 | /// 85 | /// Returns: 86 | /// Index: The index object. 87 | pub fn get_index(&self, index_name: &str) -> PineconeResult { 88 | let inner_index = self.runtime.block_on(self.inner.get_index(index_name))?; 89 | Ok(Index::new(inner_index, self.runtime.handle().clone())) 90 | } 91 | 92 | /// Creates a new Pinecone index. 93 | /// 94 | /// Args: 95 | /// name (str): The name of the index to be created. The maximum length is 45 characters. 96 | /// dimension (int): The dimensions of the vectors to be inserted in the index. 97 | /// metric (str, optional): The distance metric to be used for similarity search. You can use 'euclidean', 'cosine', or 'dotproduct'. Defaults to 'cosine'. 98 | /// replicas (int, optional): The number of replicas. Replicas duplicate your index. They provide higher availability and throughput. Defaults to 1. 99 | /// shards (int, optional): The number of shards to be used in the index. Defaults to 1. 100 | /// pods (int, optional): The number of pods for the index to use,including replicas. Defaults to 1. 101 | /// pod_type (str, optional): The type of pod to use. One of `s1`, `p1`, or `p2` appended with `.` and one of `x1`, `x2`, `x4`, or `x8`. Defaults to p1.x1. 102 | /// metadata_config (dict, optional): Configuration for the behavior of Pinecone's internal metadata index. By default, all metadata is indexed; when `metadata_config` is present, only specified metadata fields are indexed. To specify metadata fields to index, provide a JSON object of the following form: {"indexed": ["example_metadata_field"]}. 103 | /// source_collection (str, optional): The name of the collection to create an index from. 104 | /// timeout (int, optional): The number of seconds to wait for the index to be created. Defaults to 300 seconds. Pass -1 to avoid waiting for the index to be created. 105 | /// 106 | /// Returns: 107 | /// Index: The index object, if successfully created. 108 | #[pyo3(signature = (name, dimension, metric=None, replicas=None, shards=None, pods=None, pod_type=None, metadata_config=None, source_collection=None, timeout=None))] 109 | #[pyo3( 110 | text_signature = "($self, name, dimension, metric=None, replicas=None, shards=None, pods=None, pod_type=None, metadata_config=None, source_collection=None)" 111 | )] 112 | #[allow(clippy::too_many_arguments)] 113 | pub fn create_index( 114 | &self, 115 | name: &str, 116 | py: Python<'_>, 117 | dimension: i32, 118 | metric: Option, 119 | replicas: Option, 120 | shards: Option, 121 | pods: Option, 122 | pod_type: Option, 123 | metadata_config: Option>>, 124 | source_collection: Option, 125 | timeout: Option, 126 | ) -> PineconeResult { 127 | let db = Db { 128 | name: name.into(), 129 | dimension, 130 | metric, 131 | replicas, 132 | shards, 133 | pods, 134 | pod_type, 135 | metadata_config, 136 | source_collection, 137 | ..Default::default() 138 | }; 139 | self.runtime 140 | .block_on(self.inner.create_index(db, timeout, Some(py)))?; 141 | // If successful return an Index object 142 | self.get_index(name) 143 | } 144 | 145 | /// Delete an index. 146 | /// 147 | /// Args: 148 | /// name (str): The name of the index to delete. 149 | /// timeout (int, optional): The number of seconds to wait for the index to be deleted. Defaults to 300 seconds. Pass -1 to avoid waiting for the index to be deleted. 150 | /// 151 | /// Returns: 152 | /// None 153 | pub fn delete_index(&self, name: &str, timeout: Option) -> PineconeResult<()> { 154 | self.runtime 155 | .block_on(self.inner.delete_index(name, timeout))?; 156 | Ok(()) 157 | } 158 | 159 | /// List all indexes 160 | /// 161 | /// Returns: 162 | /// List[str]: A list of all indexes in the project 163 | pub fn list_indexes(&self) -> PineconeResult> { 164 | let res = self.runtime.block_on(self.inner.list_indexes())?; 165 | Ok(res) 166 | } 167 | 168 | /// Describe an index. 169 | /// 170 | /// Args: 171 | /// name (str): The name of the index to describe. 172 | /// 173 | /// Returns: 174 | /// DB: An object describing the index configuration. 175 | pub fn describe_index(&self, name: &str) -> PineconeResult { 176 | let res = self.runtime.block_on(self.inner.describe_index(name))?; 177 | Ok(res) 178 | } 179 | 180 | #[pyo3(signature = (name, replicas=None, pod_type=None))] 181 | #[pyo3(text_signature = "($self, name, replicas=None, pod_type=None)")] 182 | /// Configure an index. 183 | /// 184 | /// Args: 185 | /// name (str): The name of the index to rescale or configure. 186 | /// replicas (int): The number of replicas to use for the index. 187 | /// pod_type (str): The type of pod to use for the index. 188 | /// 189 | /// Returns: 190 | /// None 191 | pub fn scale_index( 192 | &self, 193 | name: &str, 194 | replicas: Option, 195 | pod_type: Option, 196 | ) -> PineconeResult<()> { 197 | // at least one of replicas or pod_type must be set 198 | if replicas.is_none() && pod_type.is_none() { 199 | return Err(PineconeClientError::from( 200 | core_errors::PineconeClientError::ValueError( 201 | "At least one of replicas or pod_type must be set".into(), 202 | ), 203 | )); 204 | } 205 | self.runtime 206 | .block_on(self.inner.configure_index(name, pod_type, replicas))?; 207 | Ok(()) 208 | } 209 | 210 | /// Create a new collection. 211 | /// 212 | /// Args: 213 | /// name (str): The name of the collection to create. 214 | /// source_index (str): The name of the index to use as the source for the collection. 215 | /// 216 | /// Returns: 217 | /// None 218 | pub fn create_collection( 219 | &self, 220 | name: &str, 221 | source_index: &str, 222 | ) -> Result<(), PineconeClientError> { 223 | self.runtime 224 | .block_on(self.inner.create_collection(name, source_index))?; 225 | Ok(()) 226 | } 227 | 228 | /// Describe a collection 229 | /// 230 | /// Args: 231 | /// name (str): The name of the collection to describe 232 | /// 233 | /// Returns: 234 | /// Collection: The collection description 235 | pub fn describe_collection(&self, name: &str) -> Result { 236 | let res = self 237 | .runtime 238 | .block_on(self.inner.describe_collection(name))?; 239 | Ok(res) 240 | } 241 | 242 | /// List all collections 243 | /// 244 | /// Returns: 245 | /// List[str] - A list of all collections 246 | pub fn list_collections(&self) -> PineconeResult> { 247 | let res = self.runtime.block_on(self.inner.list_collections())?; 248 | Ok(res) 249 | } 250 | 251 | /// Delete a collection 252 | /// 253 | /// Args: 254 | /// name (str): The name of the collection to delete. 255 | /// 256 | /// Returns: 257 | /// None 258 | pub fn delete_collection(&self, name: &str) -> Result<(), PineconeClientError> { 259 | self.runtime.block_on(self.inner.delete_collection(name))?; 260 | Ok(()) 261 | } 262 | } 263 | -------------------------------------------------------------------------------- /pinecone/src/data_types.rs: -------------------------------------------------------------------------------- 1 | use pyo3::types::PyDict; 2 | use pyo3::{FromPyObject, PyAny}; 3 | use std::collections::BTreeMap; 4 | 5 | use crate::utils::errors::{PineconeClientError, PineconeResult}; 6 | use client_sdk::data_types as core_data_types; 7 | use client_sdk::utils::errors::PineconeClientError as core_error; 8 | 9 | #[derive(FromPyObject, Debug, Clone)] 10 | pub enum UpsertRecord<'a> { 11 | Vector(core_data_types::Vector), 12 | TwoTuple((String, Vec)), 13 | ThreeTuple( 14 | ( 15 | String, 16 | Vec, 17 | BTreeMap, 18 | ), 19 | ), 20 | Dict(&'a PyDict), 21 | #[pyo3(transparent)] 22 | Other(&'a PyAny), // This extraction never fails 23 | } 24 | 25 | pub fn convert_upsert_enum_to_vectors( 26 | vectors: Vec, 27 | ) -> PineconeResult> { 28 | let vectors_to_upsert: Vec = vectors.into_iter().enumerate().map(|(i, vec)| { 29 | let new_vec: PineconeResult = match vec.to_owned() { 30 | UpsertRecord::Vector(v) => Ok(v), 31 | UpsertRecord::TwoTuple(t) => Ok(core_data_types::Vector{ id: t.0, values: t.1 , ..Default::default()}), 32 | UpsertRecord::ThreeTuple(t) => Ok(core_data_types::Vector{ id: t.0, values: t.1 , metadata: Some(t.2), ..Default::default()}), 33 | UpsertRecord::Dict(d) => Ok( 34 | d.try_into() 35 | .map_err(|e| match e{ 36 | core_error::UpsertKeyError { key, vec_num: _ } => 37 | core_error::UpsertKeyError {key, vec_num: i}, 38 | core_error::UpsertValueError { key, vec_num: _, actual, expected_type} => 39 | core_error::UpsertValueError {key, vec_num: i, actual, expected_type}, 40 | _ => core_error::ValueError(format!("Error in vector number {i}: {e}", i=i, e=e)) 41 | })? 42 | ), 43 | // TODO: add a dedicated error type, then format this error message in pinecone (the error message is pythonic) 44 | UpsertRecord::Other(val) => Err(PineconeClientError::from( 45 | core_error::ValueError(format!("Error in vector number {i}: Found unexpected value: {val}.\n\ 46 | Allowed types are: Vector; Tuple[str, List[float]]; Tuple[str, List[float], dict]; Dict[str, Any]", i=i, val=val)) 47 | )) 48 | 49 | }; 50 | new_vec 51 | 52 | }).collect::, _>>()?; 53 | Ok(vectors_to_upsert) 54 | } 55 | -------------------------------------------------------------------------------- /pinecone/src/index.rs: -------------------------------------------------------------------------------- 1 | use crate::data_types::convert_upsert_enum_to_vectors; 2 | use crate::data_types::UpsertRecord; 3 | use crate::utils::errors::{PineconeClientError, PineconeResult}; 4 | use client_sdk::data_types as core_data_types; 5 | use client_sdk::index as core_index; 6 | use client_sdk::utils::errors::PineconeClientError as core_error; 7 | use pyo3::prelude::*; 8 | use std::collections::{BTreeMap, HashMap}; 9 | use tokio::runtime::Handle; 10 | 11 | #[pyclass] 12 | pub struct Index { 13 | inner: core_index::Index, 14 | runtime: Handle, 15 | } 16 | 17 | impl Index { 18 | pub fn new(inner: core_index::Index, runtime: Handle) -> Self { 19 | Self { inner, runtime } 20 | } 21 | } 22 | 23 | #[pymethods] 24 | impl Index { 25 | pub fn __repr__(&self) -> String { 26 | format!("Index: \"{name}\"", name = self.inner.name) 27 | } 28 | 29 | #[pyo3(signature = (vectors, namespace="", async_req=false))] 30 | #[pyo3(text_signature = "(vectors, namespace='', async_req=False)")] 31 | /// The `Upsert` operation writes vectors into a namespace. 32 | /// If a new value is upserted for an existing vector id, it will overwrite the previous value. 33 | /// 34 | /// Args: 35 | /// vectors (Union[List[Tuple[str, List[float]]], List[Dict[str, Any]], List[Vector]]): A list of vectors to upsert. 36 | /// A vector can be represented by: 37 | /// - A `Vector` object. 38 | /// - A tuple of the form (id: str, vector: List[float]) or (id: str, vector: List[float], metadata: Dict[str, Union[str, float, int, bool, List[str]]]]) 39 | /// - A dictionary with the keys 'id' (str), 'values' (List[float]), 'sparse_values' (optional dict in the format {'indices': List[int], 'values': List[float]}), 'metadata' (Optional[Dict[str, Any]]) 40 | /// Note: sparse values are not supported when using a tuple. Please use a dictionary or a `Vector` object instead. 41 | /// 42 | /// namespace (Optional[str]): Optional namespace to which data will be upserted. 43 | /// async_req (bool): When set to True, the upsert request will be performed asynchronously, and a "future" (asyncio coroutine) will be returned. 44 | /// 45 | /// Examples: 46 | /// ```python 47 | /// index.upsert([ Vector(id='id1', values=[1.0, 2.0, 3.0], metadata={'key': 'value'}), 48 | /// Vector(id='id3', values=[1.0, 2.0, 3.0], sparse_values=SparseValues(indices=[1, 2], values=[0.2, 0.4])) ]) 49 | /// 50 | /// index.upsert([ ('id1', [1.0, 2.0, 3.0], {'key': 'value'}), 51 | /// ('id2', [1.0, 2.0, 3.0]) ]) 52 | /// 53 | /// index.upsert([ {'id': 'id1', 'values': [1.0, 2.0, 3.0], 'metadata': {'key': 'value'}}, 54 | /// {'id': 'id2', 'values': [1.0, 2.0, 3.0], 'sparse_values': {'indices': [1, 2], 'values': [0.2, 0.4]}} ]) 55 | /// 56 | /// # Mixing different vector representations is also allowed 57 | /// index.upsert([ {'id': 'id1', 'values': [1.0, 2.0, 3.0], 'metadata': {'key': 'value'}, 'sparse_values': {'indices': [1, 2], 'values': [0.2, 0.4]}}, 58 | /// ('id2', [1.0, 2.0, 3.0]), ]) 59 | /// ``` 60 | /// 61 | /// Returns: 62 | /// - If `async_req=False`: 63 | /// UpsertResponse: An upsert response object. Currently has an 'upserted_count' field with vector count. Might be extended in the future. 64 | /// - If `async_req=True`: 65 | /// An `asyncio` coroutine that can be awaited using `await` or `asyncio.gather()`. 66 | pub fn upsert<'a>( 67 | &mut self, 68 | py: Python<'a>, 69 | vectors: Vec, 70 | namespace: &'a str, 71 | async_req: bool, 72 | ) -> PyResult<&'a PyAny> { 73 | // According to tonic's documentation, cloning the generated client is actually quite cheap, 74 | // and that's the recommended behavior: https://docs.rs/tonic/latest/tonic/transport/struct.Channel.html#multiplexing-requests 75 | let mut inner_index = self.inner.clone(); 76 | 77 | let namespace = namespace.to_owned(); 78 | let vectors_to_upsert = 79 | convert_upsert_enum_to_vectors(vectors).map_err(PineconeClientError::from)?; 80 | 81 | if async_req { 82 | pyo3_asyncio::tokio::future_into_py(py, async move { 83 | let res = inner_index 84 | .upsert(&namespace, &vectors_to_upsert, None) 85 | .await 86 | .map_err(PineconeClientError::from)?; 87 | Ok(res) 88 | }) 89 | } else { 90 | pyo3_asyncio::tokio::get_runtime().block_on(async move { 91 | let res = inner_index 92 | .upsert(&namespace, &vectors_to_upsert, None) 93 | .await 94 | .map_err(PineconeClientError::from)?; 95 | Ok(res.into_py(py).into_ref(py)) 96 | }) 97 | } 98 | } 99 | 100 | #[pyo3(signature = (top_k, values=None, sparse_values=None, namespace="", filter=None, include_values=false, include_metadata=false))] 101 | #[pyo3( 102 | text_signature = "($self, top_k, values=None, sparse_values=None, namespace='', filter=None, include_values=False, include_metadata=False)" 103 | )] 104 | /// Query 105 | /// 106 | /// The `Query` operation searches a namespace, using a query vector. 107 | /// It retrieves the ids of the most similar items in a namespace, along with their similarity scores. 108 | /// To query by the id of already upserted vector, use `Index.query_by_id()` 109 | /// 110 | /// Args: 111 | /// top_k (int): The number of results to return for each query. 112 | /// values (Optional[List[float]]): The values for a new, unseen query vector. This should be the same length as the dimension of the index being queried. The results will be the `top_k` vectors closest to the given vector. Can not be used together with `id`. 113 | /// sparse_values (Optional[SparseValues]): The query vector's sparse values. 114 | /// namespace (Optional[str]): Optional namespace in which vectors will be queried. 115 | /// filter (Optional[dict]): The filter to apply. You can use vector metadata to limit your search. See 116 | /// include_values (bool): Indicates whether vector values are included in the response. 117 | /// include_metadata (bool): Indicates whether metadata is included in the response as well as the ids. 118 | /// 119 | /// Returns: 120 | /// list of QueryResults 121 | #[allow(clippy::too_many_arguments)] 122 | pub fn query( 123 | &mut self, 124 | top_k: i32, 125 | values: Option>, 126 | sparse_values: Option, 127 | namespace: &str, 128 | filter: Option>, 129 | include_values: bool, 130 | include_metadata: bool, 131 | ) -> PineconeResult> { 132 | if top_k < 1 { 133 | return Err(core_error::ValueError("top_k must be greater than 0".to_string()).into()); 134 | } 135 | let res = self.runtime.block_on(self.inner.query( 136 | namespace, 137 | values, 138 | sparse_values, 139 | top_k as u32, 140 | filter, 141 | include_values, 142 | include_metadata, 143 | ))?; 144 | Ok(res) 145 | } 146 | 147 | #[pyo3(signature = (id, top_k, namespace="", filter=None, include_values=false, include_metadata=false))] 148 | #[pyo3( 149 | text_signature = "($self, id, top_k, namespace='', filter=None, include_values=False, include_metadata=False)" 150 | )] 151 | /// Query by id 152 | /// 153 | /// The `Query by id` operation searches a namespace given the `id` of a vector already residing in the Index. 154 | /// It retrieves the ids of the most similar items in a namespace, along with their similarity scores. 155 | /// To query by new unseen vector use `Index.query()` 156 | /// 157 | /// Args: 158 | /// id (str): An id of a vector already upserted to the relevant namespace. The results will be the `top_k` nearest neighbours of the vector with the given id. Cannot be used together with `values`. 159 | /// top_k (int): The number of results to return for each query. 160 | /// namespace (Optional[str]): Optional namespace in which vectors will be queried. 161 | /// filter (Optional[dict]): The filter to apply. You can use vector metadata to limit your search. See 162 | /// include_values (bool): Indicates whether vector values are included in the response. 163 | /// include_metadata (bool): Indicates whether metadata is included in the response as well as the ids. 164 | /// 165 | /// Returns: 166 | /// list of QueryResults 167 | pub fn query_by_id( 168 | &mut self, 169 | id: &str, 170 | top_k: i32, 171 | namespace: &str, 172 | filter: Option>, 173 | include_values: bool, 174 | include_metadata: bool, 175 | ) -> PineconeResult> { 176 | if top_k < 1 { 177 | return Err(core_error::ValueError("top_k must be greater than 0".to_string()).into()); 178 | } 179 | let res = self.runtime.block_on(self.inner.query_by_id( 180 | namespace, 181 | id, 182 | top_k as u32, 183 | filter, 184 | include_values, 185 | include_metadata, 186 | ))?; 187 | Ok(res) 188 | } 189 | 190 | #[pyo3(signature = (filter=None))] 191 | #[pyo3(text_signature = "(filter=None)")] 192 | /// Describe index stats. 193 | /// 194 | /// The `DescribeIndexStats` operation returns the number of vectors present in the index, for all the namespaces 195 | /// and the fullness of the index. Can also accept a filter to count the number of vectors matching the filter. 196 | /// 197 | /// Args: 198 | /// filter (Dict[str, Union[str, float, int, bool, List, dict]]): 199 | /// If this parameter is present, the operation only returns statistics for vectors that satisfy the filter. 200 | /// See https://www.pinecone.io/docs/metadata-filtering/.. [optional] 201 | /// 202 | /// Returns: 203 | /// An `IndexStats` object containing index statistics. 204 | pub fn describe_index_stats( 205 | &mut self, 206 | filter: Option>, 207 | ) -> PineconeResult { 208 | let res = self 209 | .runtime 210 | .block_on(self.inner.describe_index_stats(filter))?; 211 | Ok(res) 212 | } 213 | 214 | #[pyo3(signature = (ids, namespace=""))] 215 | #[pyo3(text_signature = "($self, ids, namespace='')")] 216 | /// Fetch 217 | /// 218 | /// The fetch operation looks up and returns vectors, by ID, from a single namespace. 219 | /// The returned vectors include the vector data and/or metadata. 220 | /// 221 | /// Args: 222 | /// ids (List[str]): The vector IDs to fetch. 223 | /// namespace (str): The namespace to fetch vectors from. 224 | /// If not specified, the default namespace is used. [optional] 225 | /// 226 | /// Examples: 227 | /// >>> index.fetch(ids=['id1', 'id2'], namespace='my_namespace') 228 | /// >>> index.fetch(ids=['id1', 'id2']) 229 | /// 230 | /// Returns: a dictionary of vector IDs to the fetched vectors. 231 | pub fn fetch( 232 | &mut self, 233 | ids: Vec, 234 | namespace: &str, 235 | ) -> PineconeResult> { 236 | let res = self.runtime.block_on(self.inner.fetch(namespace, &ids))?; 237 | Ok(res) 238 | } 239 | 240 | #[pyo3(signature = (id, values=None, sparse_values=None, set_metadata=None, namespace=""))] 241 | #[pyo3( 242 | text_signature = "($self, id, values=None, sparse_values=None, set_metadata=None, namespace='')" 243 | )] 244 | /// Update 245 | /// The Update operation updates vector in a namespace. 246 | /// If a value is included, it will overwrite the previous value. 247 | /// If a set_metadata is included, 248 | /// the values of the fields specified in it will be added or overwrite the previous value. 249 | /// 250 | /// Examples: 251 | /// >>> index.update(id='id1', values=[1, 2, 3], namespace='my_namespace') 252 | /// >>> index.update(id='id1', set_metadata={'key': 'value'}, namespace='my_namespace') 253 | /// >>> index.update(id='id1', values=[1, 2, 3], sparse_values=SparseValues(indices=[1, 2], values=[0.2, 0.4]), 254 | /// namespace='my_namespace') 255 | /// 256 | /// Args: 257 | /// id (str): Vector's unique id. 258 | /// values (List[float]): vector values to set. [optional] 259 | /// sparse_values: (SparseValues): sparse values to update for the vector. 260 | /// set_metadata (Dict[str, Union[str, float, int, bool, List[str]]]]): metadata to set for vector. [optional] 261 | /// namespace (str): Namespace name where to update the vector.. [optional] 262 | /// 263 | pub fn update( 264 | &mut self, 265 | id: &str, 266 | values: Option>, 267 | sparse_values: Option, 268 | set_metadata: Option>, 269 | namespace: &str, 270 | ) -> PineconeResult<()> { 271 | self.runtime.block_on(self.inner.update( 272 | id, 273 | values.as_ref(), 274 | sparse_values, 275 | set_metadata, 276 | namespace, 277 | ))?; 278 | Ok(()) 279 | } 280 | 281 | #[pyo3(signature = (ids, namespace=""))] 282 | #[pyo3(text_signature = "($self, ids, namespace='')")] 283 | /// Delete 284 | /// Delete vectors by ID from a given namespace. 285 | /// 286 | /// Args: 287 | /// ids (List[str]): A list of IDs for vectors to be deleted. 288 | /// namespace (str): The name of the namespace from which vectors will be deleted. If None, the default namespace will be used. 289 | /// 290 | /// Returns: 291 | /// None 292 | pub fn delete(&mut self, ids: Vec, namespace: &str) -> PineconeResult<()> { 293 | self.runtime.block_on(self.inner.delete(ids, namespace))?; 294 | Ok(()) 295 | } 296 | 297 | #[pyo3(signature = (filter, namespace=""))] 298 | #[pyo3(text_signature = "($self, filter, namespace='')")] 299 | /// Delete by filter 300 | /// The delete by filter operation deletes a list of vectors from a given namespace that match the filter. 301 | /// 302 | /// Args: 303 | /// filter (Dict[str, Union[str, float, int, bool, List, dict]]): filter to be applied to delete the vectors. See https://www.pinecone.io/docs/metadata-filtering/ 304 | /// namespace (Optional[str]): The name of the namespace from which vectors will be deleted. If None, the default namespace will be used. 305 | /// 306 | /// Returns: 307 | /// None 308 | pub fn delete_by_metadata( 309 | &mut self, 310 | filter: Option>, 311 | namespace: &str, 312 | ) -> PineconeResult<()> { 313 | self.runtime 314 | .block_on(self.inner.delete_by_metadata(filter, namespace))?; 315 | Ok(()) 316 | } 317 | 318 | #[pyo3(signature = (namespace=""))] 319 | #[pyo3(text_signature = "($self, namespace='')")] 320 | /// Delete all 321 | /// The delete all operation deletes all the vectors from a given namespace. 322 | /// 323 | /// Args: 324 | /// namespace (str): The name of the namespace from which vectors will be deleted. If None, the default namespace will be used. 325 | /// 326 | /// Returns: 327 | /// None 328 | pub fn delete_all(&mut self, namespace: &str) -> PineconeResult<()> { 329 | self.runtime.block_on(self.inner.delete_all(namespace))?; 330 | Ok(()) 331 | } 332 | } 333 | -------------------------------------------------------------------------------- /pinecone/src/lib.rs: -------------------------------------------------------------------------------- 1 | // pub extern crate client_sdk; 2 | // 3 | 4 | use pyo3::prelude::*; 5 | 6 | pub mod client; 7 | pub mod data_types; 8 | pub mod index; 9 | pub mod utils; 10 | 11 | use crate::index::Index; 12 | use client::Client; 13 | use client_sdk::data_types as core_data_types; 14 | use utils::errors; 15 | 16 | #[pymodule] 17 | fn pinecone(_py: Python<'_>, m: &PyModule) -> PyResult<()> { 18 | m.add_class::()?; 19 | m.add_class::()?; 20 | m.add_class::()?; 21 | m.add_class::()?; 22 | m.add_class::()?; 23 | m.add_class::()?; 24 | m.add( 25 | "PineconeOpError", 26 | ::type_object(_py), 27 | )?; 28 | m.add_class::()?; 29 | Ok(()) 30 | } 31 | -------------------------------------------------------------------------------- /pinecone/src/utils/errors.rs: -------------------------------------------------------------------------------- 1 | use client_sdk::utils::errors as core_errors; 2 | use pyo3::create_exception; 3 | use pyo3::exceptions; 4 | use pyo3::prelude::*; 5 | 6 | create_exception!( 7 | pinecone_client, 8 | PineconeOpError, 9 | pyo3::exceptions::PyException 10 | ); 11 | 12 | pub struct PineconeClientError { 13 | inner: core_errors::PineconeClientError, 14 | } 15 | 16 | impl From for PineconeClientError { 17 | fn from(error: core_errors::PineconeClientError) -> PineconeClientError { 18 | PineconeClientError { inner: error } 19 | } 20 | } 21 | 22 | impl From for PyErr { 23 | fn from(err: PineconeClientError) -> PyErr { 24 | match err.inner { 25 | core_errors::PineconeClientError::ArgumentError { .. } => { 26 | exceptions::PyValueError::new_err(err.inner.to_string()) 27 | } 28 | core_errors::PineconeClientError::ControlPlaneConnectionError { .. } => { 29 | exceptions::PyConnectionError::new_err(err.inner.to_string()) 30 | } 31 | core_errors::PineconeClientError::IndexConnectionError { .. } => { 32 | exceptions::PyConnectionError::new_err(err.inner.to_string()) 33 | } 34 | core_errors::PineconeClientError::DataplaneOperationError(_) => { 35 | PineconeOpError::new_err(err.inner.to_string()) 36 | } 37 | core_errors::PineconeClientError::IoError(_) => { 38 | exceptions::PyIOError::new_err(err.inner.to_string()) 39 | } 40 | core_errors::PineconeClientError::MetadataValueError { .. } => { 41 | exceptions::PyValueError::new_err(err.inner.to_string()) 42 | } 43 | core_errors::PineconeClientError::MetadataError { .. } => { 44 | exceptions::PyValueError::new_err(err.inner.to_string()) 45 | } 46 | core_errors::PineconeClientError::Other(_) => { 47 | exceptions::PyRuntimeError::new_err(err.inner.to_string()) 48 | } 49 | core_errors::PineconeClientError::ControlPlaneOperationError { .. } => { 50 | PineconeOpError::new_err(err.inner.to_string()) 51 | } 52 | core_errors::PineconeClientError::ControlPlaneParsingError { .. } => { 53 | PineconeOpError::new_err(err.inner.to_string()) 54 | } 55 | core_errors::PineconeClientError::DeserializationError(_) => { 56 | PineconeOpError::new_err(err.inner.to_string()) 57 | } 58 | core_errors::PineconeClientError::ValueError(_) => { 59 | exceptions::PyValueError::new_err(err.inner.to_string()) 60 | } 61 | core_errors::PineconeClientError::UpsertValueError { .. } => { 62 | exceptions::PyValueError::new_err(err.inner.to_string()) 63 | } 64 | core_errors::PineconeClientError::UpsertKeyError { .. } => { 65 | exceptions::PyValueError::new_err(err.inner.to_string()) 66 | } 67 | core_errors::PineconeClientError::KeyboardInterrupt(_) => { 68 | exceptions::PyKeyboardInterrupt::new_err(err.inner.to_string()) 69 | } 70 | } 71 | } 72 | } 73 | 74 | pub type PineconeResult = Result; 75 | -------------------------------------------------------------------------------- /pinecone/src/utils/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod errors; 2 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pinecone-io/pinecone-client/94133b14f85223272d1cb5cb486460a8647af4d8/tests/__init__.py -------------------------------------------------------------------------------- /tests/unit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pinecone-io/pinecone-client/94133b14f85223272d1cb5cb486460a8647af4d8/tests/unit/__init__.py -------------------------------------------------------------------------------- /tests/unit/test_control_plane.py: -------------------------------------------------------------------------------- 1 | """Tests for control plane api calls""" 2 | import os 3 | from pinecone import Client 4 | import pytest 5 | from time import time 6 | import numpy as np 7 | from ..utils.utils import retry_assert 8 | from ..utils.remote_index import PodType 9 | 10 | env = os.getenv('PINECONE_REGION') 11 | key = os.getenv('PINECONE_API_KEY') 12 | client = Client(key, env) 13 | INDEX_NAME_PREFIX = 'control-plane-' + str(np.random.randint(10000)) 14 | d = 512 15 | 16 | INDEX_NAME_KEY = 'index_name' 17 | POD_TYPE_KEY = 'pod_type' 18 | 19 | 20 | @pytest.fixture(scope="module", 21 | params=[ 22 | {INDEX_NAME_KEY: f'{INDEX_NAME_PREFIX}-{PodType.P1}', 23 | POD_TYPE_KEY: PodType.P1} 24 | ], 25 | ids=lambda param: str(param[POD_TYPE_KEY])) 26 | def index_fixture(testrun_uid, request): 27 | index_name = request.param[INDEX_NAME_KEY] + '-' + testrun_uid[:8] 28 | pod_type = request.param[POD_TYPE_KEY] 29 | # Note: relies on grouping strategy –-dist=loadfile to keep different xdist workers 30 | # from repeating this 31 | index_creation_args = {'name': index_name, 32 | 'dimension': d, 33 | 'pod_type': str(pod_type), 34 | 'pods': 2} 35 | client.create_index(**index_creation_args) 36 | 37 | def remove_index(): 38 | if index_name in client.list_indexes(): 39 | client.delete_index(index_name) 40 | 41 | # attempt to remove index even if creation raises exception 42 | request.addfinalizer(remove_index) 43 | 44 | yield index_name, f"{pod_type}.x1" if pod_type.is_implicitly_x1() else str(pod_type) 45 | 46 | # The client interface 47 | def test_client_valid_params(): 48 | # assert no error is raised 49 | pinecone = Client(api_key='api_key', region='env',project_id='project_id') 50 | 51 | def test_env_vars(): 52 | # assuming tess are run with env vars set 53 | pinecone = Client() 54 | 55 | @pytest.fixture(scope="function") 56 | def set_api_key_env_var(testrun_uid): 57 | old_key = os.environ.get('PINECONE_API_KEY') 58 | os.environ['PINECONE_API_KEY'] = "non-existent-key" 59 | yield old_key 60 | if old_key is not None: 61 | os.environ['PINECONE_API_KEY'] = old_key 62 | 63 | @pytest.fixture(scope="function") 64 | def set_region_env_var(testrun_uid): 65 | old_region = os.environ.get('PINECONE_REGION') 66 | os.environ['PINECONE_REGION'] = "non-existent-region" 67 | yield old_region 68 | if old_region is not None: 69 | os.environ['PINECONE_REGION'] = old_region 70 | 71 | def test_env_vars_missing_api_key(set_api_key_env_var): 72 | with pytest.raises(ConnectionError): 73 | pinecone = Client() 74 | 75 | def test_env_vars_missing_region(set_region_env_var): 76 | with pytest.raises(ConnectionError): 77 | pinecone = Client() 78 | 79 | def test_env_var_override(set_api_key_env_var, set_region_env_var): 80 | assert os.environ['PINECONE_API_KEY'] == "non-existent-key" 81 | assert os.environ['PINECONE_REGION'] == "non-existent-region" 82 | pinecone = Client(api_key=set_api_key_env_var, region=set_region_env_var) 83 | pinecone.list_indexes() 84 | 85 | # def test_env_var_override_region(set_region_env_var): 86 | # assert os.environ['PINECONE_REGION'] == "non-existent-region" 87 | # pinecone = Client(region=set_region_env_var) 88 | # pinecone.list_indexes() 89 | def test_client_invalid_params(): 90 | with pytest.raises(TypeError): 91 | pinecone = Client(random_param='random_param') 92 | 93 | def test_client_invalid_region(): 94 | with pytest.raises(ConnectionError): 95 | pinecone = Client(region='invalid_region') 96 | 97 | def test_missing_dimension(): 98 | # Missing Dimension 99 | name = 'test-missing-dim' 100 | with pytest.raises(TypeError): 101 | client.create_index(name) 102 | 103 | 104 | def test_invalid_name(): 105 | # Missing Dimension 106 | name = 'Test-Bad-Name' 107 | # TODO: raise proper exception 108 | with pytest.raises(Exception) as e: 109 | client.create_index(name, 32) 110 | # assert e.value.status == 400 111 | 112 | def test_invalid_name_rfc_1123(): 113 | name = "bad.name.with.periods" 114 | # TODO: raise proper exception 115 | with pytest.raises(Exception) as e: 116 | client.create_index(name, 32) 117 | # assert e.value.status == 400 118 | 119 | @pytest.fixture(scope="module") 120 | def timeout_index(testrun_uid): 121 | name = f'{INDEX_NAME_PREFIX}-create-timeout' + '-' + testrun_uid[:8] 122 | if name in client.list_indexes(): 123 | client.delete_index(name) 124 | yield name 125 | if name in client.list_indexes(): 126 | client.delete_index(name) 127 | 128 | def test_create_timeout(timeout_index): 129 | timeout = 5 # seconds 130 | TOLERANCE = 0.5 # seconds 131 | before = time() 132 | with pytest.raises(RuntimeError) as e: 133 | client.create_index(timeout_index, 32, timeout=timeout) 134 | eplased = time() - before 135 | assert eplased - timeout < TOLERANCE 136 | assert "timed out" in str(e.value) 137 | 138 | def test_create_timeout_invalid(): 139 | timeout = -5 # seconds 140 | with pytest.raises(ValueError) as e: 141 | client.create_index('test-create-timeout-invalid', 32, timeout=timeout) 142 | assert "-1 or a positive integer" in str(e.value) 143 | assert "test-create-timeout-invalid" not in client.list_indexes() 144 | 145 | def test_create_duplicate(index_fixture): 146 | index_name, _ = index_fixture 147 | # Duplicate index 148 | with pytest.raises(Exception): 149 | client.create_index(index_name, 32) 150 | 151 | def test_get(index_fixture): 152 | index_name, pod_type = index_fixture 153 | # Successful Call 154 | result = client.describe_index(index_name) 155 | # TODO: dict vs obj 156 | return_val = {'name': index_name, 'dimension': 512, 'replicas': 1, 'shards': 2, 'pod_type': pod_type, 'metric': 'cosine', 'pods': 2, 'source_collection': None, 'metadata_config': None, 'status': 'Ready'} 157 | assert result.to_dict() == return_val 158 | # Calling non-existent index 159 | with pytest.raises(Exception): 160 | client.describe_index('non-existent-index') 161 | 162 | # Missing Field 163 | with pytest.raises(TypeError): 164 | client.describe_index() 165 | 166 | 167 | def test_update(index_fixture): 168 | index_name, _ = index_fixture 169 | # Scale Up 170 | num_replicas = 2 171 | client.scale_index(name=index_name, replicas=num_replicas) 172 | 173 | retry_assert(lambda: client.describe_index(index_name).status == 'Ready') 174 | meta_obj = client.describe_index(index_name) 175 | assert meta_obj.replicas == 2 176 | assert meta_obj.pods == 4 177 | 178 | # Missing replicas field 179 | # TODO: raise proper exception 180 | with pytest.raises(ValueError): 181 | client.scale_index(index_name) 182 | 183 | # Calling on non-existent index 184 | with pytest.raises(Exception): 185 | client.scale_index('non-existent-index', 2) 186 | 187 | # Scale to zero 188 | num_replicas = 0 189 | client.scale_index(name=index_name, replicas=num_replicas) 190 | retry_assert(lambda: client.describe_index(index_name).status == 'Ready') 191 | meta_obj = client.describe_index(index_name) 192 | assert meta_obj.replicas == num_replicas 193 | assert meta_obj.pods == num_replicas 194 | 195 | 196 | 197 | 198 | def test_delete(index_fixture): 199 | index_name, _ = index_fixture 200 | # Delete existing index 201 | client.delete_index(index_name) 202 | assert index_name not in client.list_indexes() 203 | 204 | # Delete non existent index 205 | with pytest.raises(Exception): 206 | client.delete_index('non-existent-index') 207 | 208 | # Missing Field 209 | with pytest.raises(TypeError): 210 | client.delete_index() 211 | 212 | 213 | def test_create_collection_from_nonexistent_index(): 214 | nonexistent_index_name = 'nonexistent_index_name' 215 | collection_name = 'collection' 216 | with pytest.raises(Exception) as e: 217 | client.create_collection(collection_name, nonexistent_index_name) 218 | assert f'source database {nonexistent_index_name} does not exist' in str(e.value) 219 | 220 | 221 | def test_create_index_from_nonexistent_collection(): 222 | index_name = 'test-index-collection' 223 | nonexistent_collection_name = 'nonexistent_collection_name' 224 | with pytest.raises(Exception) as e: 225 | client.create_index(index_name, d, pods=2, pod_type='s1', 226 | source_collection=nonexistent_collection_name) 227 | assert f'failed to fetch source collection {nonexistent_collection_name}' in str(e.value) 228 | -------------------------------------------------------------------------------- /tests/unit/test_hybrid_search.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | from loguru import logger 5 | from pinecone import Vector, Client, SparseValues 6 | 7 | from ..utils.remote_index import PodType, RemoteIndex 8 | from ..utils.utils import index_fixture_factory, retry_assert, sparse_values, get_vector_count, approx_sparse_equals 9 | 10 | logger.remove() 11 | logger.add(sys.stdout, level=(os.getenv("PINECONE_LOGGING") or "INFO")) 12 | 13 | vector_dim = 512 14 | env = os.getenv('PINECONE_REGION') 15 | api_key = os.getenv('PINECONE_API_KEY') 16 | 17 | INDEX_NAME = 'test-hybrid-search' 18 | 19 | hybrid_index = index_fixture_factory( 20 | [ 21 | (RemoteIndex(pods=2, index_name=f'{INDEX_NAME}-{PodType.S1}', 22 | dimension=vector_dim, pod_type=PodType.S1, metric='dotproduct'), str(PodType.S1)), 23 | ] 24 | ) 25 | 26 | 27 | def sparse_vector(dimension=32000, nnz=120): 28 | indices, values = sparse_values(dimension, nnz) 29 | return SparseValues(indices=indices, values=values) 30 | 31 | 32 | def get_test_data(vector_count=10, no_meta_vector_count=5, dimension=vector_dim, sparse=True): 33 | """repeatably produces same results for a given vector_count""" 34 | meta_vector_count = vector_count - no_meta_vector_count 35 | metadata_choices = [ 36 | {'genre': 'action', 'year': 2020}, 37 | {'genre': 'documentary', 'year': 2021}, 38 | {'genre': 'documentary', 'year': 2005}, 39 | {'genre': 'drama', 'year': 2011}, 40 | ] 41 | no_meta_vectors: list[Vector] = [ 42 | Vector(id=f'vec{i}', values=[i / 1000] * dimension, sparse_values=sparse_vector() if sparse else None) 43 | for i in range(no_meta_vector_count) 44 | ] 45 | meta_vectors: list[Vector] = [ 46 | Vector(id=f'mvec{i}', values=[i / 1000] * dimension, sparse_values=sparse_vector() if sparse else {}, 47 | metadata=metadata_choices[i % len(metadata_choices)]) 48 | for i in range(meta_vector_count) 49 | ] 50 | 51 | return list(meta_vectors) + list(no_meta_vectors) 52 | 53 | 54 | def write_test_data(index, namespace, vector_count=10, no_meta_vector_count=5, dimension=vector_dim, batch_size=300): 55 | """writes vector_count vectors into index, half with metadata half without.""" 56 | data = get_test_data(vector_count, no_meta_vector_count, dimension) 57 | index.upsert(vectors=data, namespace=namespace) 58 | return {vector.id: vector for vector in data} 59 | 60 | 61 | def test_upsert_vectors(hybrid_index): 62 | index, _ = hybrid_index 63 | namespace = 'test_upsert_vectors' 64 | api_response = index.upsert( 65 | vectors=[ 66 | Vector(id='mvec1', values=[0.1] * vector_dim, sparse_values=sparse_vector(), metadata={'genre': 'action', 'year': 2020}), 67 | Vector(id='mvec2', values=[0.2] * vector_dim, sparse_values=sparse_vector(), metadata={'genre': 'documentary', 'year': 2021}), 68 | ], 69 | namespace=namespace, 70 | ) 71 | assert api_response.upserted_count == 2 72 | logger.debug('got openapi upsert with metadata response: {}', api_response) 73 | 74 | 75 | def test_fetch_vectors_mixed_metadata(hybrid_index): 76 | index, _ = hybrid_index 77 | namespace = 'test_fetch_vectors_mixed_metadata' 78 | vector_count = 10 79 | test_data = write_test_data(index, namespace, vector_count, no_meta_vector_count=5) 80 | api_response = index.fetch(ids=['vec1', 'mvec2'], namespace=namespace) 81 | logger.debug('got openapi fetch response: {}', api_response) 82 | 83 | for vector_id in ['mvec2', 'vec1']: 84 | expected_vector = test_data.get(vector_id) 85 | fetched_vector = api_response[vector_id] 86 | assert fetched_vector 87 | assert fetched_vector.values == expected_vector.values 88 | assert fetched_vector.metadata == expected_vector.metadata 89 | assert approx_sparse_equals(fetched_vector.sparse_values, expected_vector.sparse_values) 90 | 91 | 92 | def test_query_simple(hybrid_index): 93 | index, _ = hybrid_index 94 | namespace = 'test_query_simple' 95 | vector_count = 100 96 | write_test_data(index, namespace, vector_count) 97 | # simple query - no filter, no data, no metadata 98 | dense_query_response = index.query( 99 | values=[0.1] * vector_dim, 100 | namespace=namespace, 101 | top_k=10, 102 | include_values=False, 103 | include_metadata=False 104 | ) 105 | logger.debug('got openapi query (no filter, no data, no metadata) response: {}', dense_query_response) 106 | 107 | dense_query_match = dense_query_response[0] 108 | assert not dense_query_match.values 109 | assert not dense_query_match.sparse_values 110 | assert not dense_query_match.metadata 111 | 112 | hybrid_query_response = index.query( 113 | values=[0.1] * vector_dim, 114 | sparse_values=sparse_vector(), 115 | namespace=namespace, 116 | top_k=10, 117 | include_values=False, 118 | include_metadata=False 119 | ) 120 | 121 | assert dense_query_response != hybrid_query_response 122 | 123 | 124 | def test_query_simple_with_include_values(hybrid_index): 125 | index, _ = hybrid_index 126 | namespace = 'test_query_simple_with_values_metadata' 127 | vector_count = 10 128 | test_data = write_test_data(index, namespace, vector_count) 129 | # simple query - no filter, with data, with metadata 130 | api_response = index.query( 131 | values=[0.1] * vector_dim, 132 | sparse_values=sparse_vector(), 133 | namespace=namespace, 134 | top_k=10, 135 | include_values=True, 136 | include_metadata=True 137 | ) 138 | logger.debug('got openapi query (no filter, with data, with metadata) response: {}', api_response) 139 | 140 | first_match_vector = api_response[0] 141 | expected_vector = test_data.get(first_match_vector.id) 142 | assert first_match_vector.values == expected_vector.values 143 | assert approx_sparse_equals(first_match_vector.sparse_values, expected_vector.sparse_values) 144 | if first_match_vector.id.startswith('mvec'): 145 | assert first_match_vector.metadata == test_data.get(first_match_vector.id).metadata 146 | else: 147 | assert not first_match_vector.metadata 148 | 149 | 150 | def test_delete(hybrid_index): 151 | index, _ = hybrid_index 152 | namespace = 'test_delete' 153 | vector_count = 10 154 | test_data = write_test_data(index, namespace, vector_count) 155 | 156 | expected_mvec1 = test_data.get('mvec1') 157 | api_response = index.fetch(ids=['mvec1', 'mvec2'], namespace=namespace) 158 | logger.debug('got openapi fetch response: {}', api_response) 159 | assert api_response and api_response.get('mvec1').values == expected_mvec1.values 160 | assert api_response and approx_sparse_equals(api_response.get('mvec1').sparse_values, expected_mvec1.sparse_values) 161 | 162 | vector_count = get_vector_count(index, namespace) 163 | api_response = index.delete(ids=['vec1', 'vec2'], namespace=namespace) 164 | logger.debug('got openapi delete response: {}', api_response) 165 | retry_assert(lambda: get_vector_count(index, namespace) == (vector_count - 2)) 166 | 167 | 168 | def test_update(hybrid_index): 169 | index, _ = hybrid_index 170 | namespace = 'test_update' 171 | vector_count = 10 172 | test_data = write_test_data(index, namespace, vector_count) 173 | assert get_vector_count(index, namespace) == vector_count 174 | 175 | api_response = index.update(id='mvec1', namespace=namespace, values=test_data.get('mvec2').values) 176 | logger.debug('got openapi update response: {}', api_response) 177 | retry_assert( 178 | lambda: index.fetch(ids=['mvec1'], namespace=namespace)['mvec1'].values == test_data.get('mvec2').values) 179 | fetched_response = index.fetch(ids=['mvec1'], namespace=namespace)['mvec1'] 180 | assert fetched_response.metadata == test_data.get('mvec1').metadata 181 | assert approx_sparse_equals(fetched_response.sparse_values, test_data.get('mvec1').sparse_values) 182 | 183 | api_response = index.update(id='mvec1', namespace=namespace, sparse_values=test_data.get('mvec2').sparse_values) 184 | logger.debug('got openapi update response: {}', api_response) 185 | retry_assert( 186 | lambda: approx_sparse_equals(index.fetch(ids=['mvec1'], namespace=namespace)['mvec1'].sparse_values, test_data.get('mvec2').sparse_values)) 187 | fetched_response = index.fetch(ids=['mvec1'], namespace=namespace)['mvec1'] 188 | assert fetched_response.values == test_data.get('mvec2').values 189 | assert fetched_response.metadata == test_data.get('mvec1').metadata 190 | 191 | api_response = index.update(id='mvec2', namespace=namespace, set_metadata=test_data.get('mvec1').metadata) 192 | logger.debug('got openapi update response: {}', api_response) 193 | retry_assert( 194 | lambda: index.fetch(ids=['mvec2'], namespace=namespace)['mvec2'].metadata == test_data.get('mvec1').metadata) 195 | fetched_response = index.fetch(ids=['mvec1'], namespace=namespace)['mvec1'] 196 | assert fetched_response.values == test_data.get('mvec2').values 197 | assert approx_sparse_equals(fetched_response.sparse_values, test_data.get('mvec2').sparse_values) 198 | 199 | api_response = index.update(id='mvec3', namespace=namespace, values=test_data.get('mvec1').values, 200 | sparse_values=test_data.get('mvec1').sparse_values, set_metadata=test_data.get('mvec1').metadata) 201 | logger.debug('got openapi update response: {}', api_response) 202 | retry_assert( 203 | lambda: index.fetch(ids=['mvec3'], namespace=namespace)['mvec3'].values == test_data.get('mvec1').values) 204 | fetched_response = index.fetch(ids=['mvec1'], namespace=namespace)['mvec1'] 205 | assert approx_sparse_equals(fetched_response.sparse_values, test_data.get('mvec2').sparse_values) 206 | assert fetched_response.metadata == test_data.get('mvec1').metadata 207 | 208 | 209 | def test_upsert_with_dense_only(hybrid_index): 210 | index, _ = hybrid_index 211 | namespace = 'test_upsert_with_dense_only' 212 | vectors = [Vector('0', [0.1] * vector_dim, metadata={'colors': True, 'country': 'greece'}), 213 | Vector('1', [-0.1] * vector_dim, metadata={'colors': False}), 214 | Vector('2', [0.2] * vector_dim, sparse_values=sparse_vector(),metadata={})] 215 | index.upsert(vectors=vectors, namespace=namespace) 216 | 217 | api_response = index.fetch(ids=['0', '1', '2'], namespace=namespace) 218 | for expected_vector in vectors: 219 | vector = api_response.get(expected_vector.id) 220 | assert vector.values == expected_vector.values 221 | if(vector.metadata is not None and expected_vector.metadata!={}): 222 | assert vector.metadata == expected_vector.metadata 223 | assert approx_sparse_equals(vector.sparse_values,expected_vector.sparse_values) 224 | 225 | 226 | def test_query_by_id(hybrid_index): 227 | index, _ = hybrid_index 228 | namespace = 'test_query_by_id' 229 | vector_count = 10 230 | test_data = write_test_data(index, namespace, vector_count) 231 | api_response = index.query_by_id( 232 | id='vec1', 233 | namespace=namespace, 234 | top_k=10, 235 | include_values=True, 236 | include_metadata=True 237 | ) 238 | logger.debug('got openapi query response: {}', api_response) 239 | 240 | assert len(api_response) == 10 241 | for match_vector in api_response: 242 | expected_vector = test_data.get(match_vector.id) 243 | assert match_vector.values == expected_vector.values 244 | assert approx_sparse_equals(expected_vector.sparse_values, match_vector.sparse_values) 245 | if expected_vector.metadata: 246 | assert match_vector.metadata == expected_vector.metadata 247 | else: 248 | assert not match_vector.metadata -------------------------------------------------------------------------------- /tests/unit/test_metadata.py: -------------------------------------------------------------------------------- 1 | import dataclasses 2 | import os 3 | import sys 4 | import uuid 5 | from itertools import cycle 6 | from typing import Any 7 | 8 | import numpy as np 9 | import pytest 10 | from loguru import logger 11 | 12 | from pinecone import Vector, SparseValues 13 | from ..utils.remote_index import RemoteIndex, PodType 14 | from ..utils.utils import retry_assert, index_fixture_factory 15 | 16 | logger.remove() 17 | logger.add(sys.stdout, level=(os.getenv("PINECONE_LOGGING") or "INFO")) 18 | 19 | d = 128 20 | n = 100 21 | 22 | INDEX_NAME_PREFIX = 'test-metadata' 23 | MAPPING_INDEX_NAME_PREFIX = 'test-mapping' 24 | 25 | test_metadata_index = index_fixture_factory( 26 | [ 27 | (RemoteIndex(pods=1, index_name=f'{INDEX_NAME_PREFIX}-{PodType.P1}', 28 | dimension=d, pod_type=PodType.P1), str(PodType.P1)) 29 | ] 30 | ) 31 | 32 | test_metadata_index_with_mapping = index_fixture_factory( 33 | [ 34 | (RemoteIndex(pods=1, index_name=f'{MAPPING_INDEX_NAME_PREFIX}-{PodType.P1}', 35 | dimension=d, pod_type=PodType.P1, metadata_config={"indexed": ["weather"]}), 36 | str(PodType.P1)) 37 | ] 38 | ) 39 | 40 | 41 | @dataclasses.dataclass 42 | class RunData: 43 | ids: Any 44 | vectors: Any 45 | metadata: Any 46 | query_vector: Any 47 | 48 | 49 | def insert_test_data(index, _n, _d, namespace=''): 50 | ids = [str(i) for i in range(_n)] 51 | vectors = [np.random.rand(_d).astype(np.float32).tolist() for _ in range(_n)] 52 | weather_vocab = ['sunny', 'rain', 'cloudy', 'snowy'] 53 | loop = cycle(weather_vocab) 54 | metadata = [{"value": i, 'weather': next(loop), "bool_field": i % 2 == 0} for i in range(_n)] 55 | index.upsert( 56 | vectors=[ 57 | Vector(id=ids[i], values=vectors[i], metadata=metadata[i]) 58 | for i in range(_n) 59 | ], 60 | namespace=namespace 61 | ) 62 | retry_assert(lambda: index.describe_index_stats().namespaces.get(namespace).vector_count == _n, max_tries=10) 63 | 64 | return ids, vectors, metadata 65 | 66 | 67 | @pytest.fixture(scope="module") 68 | def test_data(test_metadata_index): 69 | # Note: relies on grouping strategy –-dist=loadfile to keep different xdist workers 70 | # from running different tests below with different data/metadata 71 | index, _ = test_metadata_index 72 | query_vector = np.random.rand(d).astype(np.float32).tolist() 73 | ids, vectors, metadata = insert_test_data(index, n, d) 74 | yield RunData(ids=ids, vectors=vectors, metadata=metadata, query_vector=query_vector) 75 | 76 | 77 | @pytest.fixture(scope="module") 78 | def test_data_for_mapping(test_metadata_index_with_mapping): 79 | index = test_metadata_index_with_mapping[0] 80 | query_vector = np.random.rand(d).astype(np.float32).tolist() 81 | ids, vectors, metadata = insert_test_data(index, n, d) 82 | yield RunData(ids=ids, vectors=vectors, metadata=metadata, query_vector=query_vector) 83 | 84 | 85 | def get_query_results(index, vector, filter): 86 | return index.query( 87 | values=vector, 88 | filter=filter, 89 | top_k=10, 90 | include_values=True, 91 | include_metadata=True 92 | ) 93 | 94 | 95 | def test_fetch(test_metadata_index, test_data): 96 | index, _ = test_metadata_index 97 | ids = test_data.ids 98 | metadata = test_data.metadata 99 | fetch_response = index.fetch(ids=[ids[0]]) 100 | fetched_metadata = fetch_response[ids[0]].metadata 101 | assert fetched_metadata == metadata[0] 102 | 103 | 104 | def test_delete_eq(test_metadata_index): 105 | index, _ = test_metadata_index 106 | namespace = 'test_delete_eq' 107 | _n = 10 108 | ids, values, metadata = insert_test_data(index, _n, d, namespace=namespace) 109 | first_md = metadata[0]['weather'] 110 | matches = [i for i in range(_n) if metadata[i]['weather'] == first_md] 111 | 112 | fetch_response = index.fetch(ids=ids, namespace=namespace) 113 | fetched_vecs = list(fetch_response.values()) 114 | assert len(fetched_vecs) == len(ids) 115 | assert all(fetch_response[ids[i]].metadata['weather'] == first_md for i in matches) 116 | 117 | index.delete_by_metadata(namespace=namespace, filter={'weather': first_md}) 118 | retry_assert( 119 | lambda: index.describe_index_stats().namespaces[namespace].vector_count == len(ids) - len(matches)) 120 | 121 | fetch_response = index.fetch(ids=ids, namespace=namespace) 122 | assert len(fetch_response.values()) == len(ids) - len(matches) 123 | assert all(vec.metadata.get('weather') != first_md for vec in fetch_response.values()) 124 | 125 | 126 | def test_gt(test_metadata_index, test_data): 127 | index, _ = test_metadata_index 128 | query_vector = test_data.query_vector 129 | gt_filter = {"value": {"$gt": 10}} 130 | query_response = get_query_results(index, query_vector, gt_filter) 131 | result = query_response 132 | for match in result: 133 | assert 'value' in match.metadata 134 | assert match.metadata['value'] > 10 135 | 136 | 137 | def test_lt(test_metadata_index, test_data): 138 | index, _ = test_metadata_index 139 | query_vector = test_data.query_vector 140 | lt_filter = {"value": {"$lt": 10}} 141 | query_response = get_query_results(index, query_vector, lt_filter) 142 | result = query_response 143 | for match in result: 144 | assert match.metadata['value'] < 10 145 | 146 | 147 | def test_eq(test_metadata_index, test_data): 148 | index, _ = test_metadata_index 149 | query_vector = test_data.query_vector 150 | eq_filter = {"value": {"$eq": 25}} 151 | query_response = get_query_results(index, query_vector, eq_filter) 152 | result = query_response 153 | for match in result: 154 | assert match.metadata['value'] == 25 155 | 156 | 157 | def test_boolean_eq(test_metadata_index, test_data): 158 | index, _ = test_metadata_index 159 | query_vector = test_data.query_vector 160 | eq_filter = {"bool_field": True} 161 | query_response = get_query_results(index, query_vector, eq_filter) 162 | result = query_response 163 | for match in result: 164 | assert match.metadata['bool_field'] 165 | 166 | 167 | def test_boolean_ne(test_metadata_index, test_data): 168 | index, _ = test_metadata_index 169 | query_vector = test_data.query_vector 170 | eq_filter = {"bool_field": False} 171 | query_response = get_query_results(index, query_vector, eq_filter) 172 | for match in query_response: 173 | assert not match.metadata['bool_field'] 174 | 175 | 176 | def test_in(test_metadata_index, test_data): 177 | index, _ = test_metadata_index 178 | query_vector = test_data.query_vector 179 | in_filter = {"weather": {"$in": ['snowy', 'sunny']}} 180 | query_response = get_query_results(index, query_vector, in_filter) 181 | for match in query_response: 182 | metadata_value = match.metadata['weather'] 183 | assert metadata_value in ['snowy', 'sunny'] 184 | 185 | 186 | def test_nin(test_metadata_index, test_data): 187 | index, _ = test_metadata_index 188 | query_vector = test_data.query_vector 189 | nin_vals = ['snowy', 'rainy'] 190 | nin_filter = {"weather": {"$nin": nin_vals}} 191 | query_response = get_query_results(index, query_vector, nin_filter) 192 | for match in query_response: 193 | metadata_value = match.metadata['value'] 194 | assert metadata_value not in nin_vals 195 | 196 | 197 | def test_compound_ne_and_lte(test_metadata_index, test_data): 198 | index, _ = test_metadata_index 199 | query_vector = test_data.query_vector 200 | cmp_filter = {"$and": [{"weather": {"$ne": "sunny"}}, {"value": {"$lte": 2}}]} 201 | query_response = get_query_results(index, query_vector, cmp_filter) 202 | for match in query_response: 203 | assert 'sunny' != match.metadata["weather"] 204 | assert match.metadata["value"] <= 2 205 | 206 | 207 | def test_compound_eq_or_gte(test_metadata_index, test_data): 208 | index, _ = test_metadata_index 209 | query_vector = test_data.query_vector 210 | cmp_filter = {"$or": [{"weather": {"$eq": "snowy"}}, {"value": {"$gte": 5}}]} 211 | # cmp_filter = {"$or": [{"weather": {"$eq": "snowy"}}, {"value": {"gte": 5}}]} 212 | query_response = get_query_results(index, query_vector, cmp_filter) 213 | for match in query_response: 214 | w = match.metadata["weather"] 215 | v = match.metadata["value"] 216 | assert w == "snowy" or v >= 5 217 | 218 | 219 | def test_update_md(test_metadata_index): 220 | index, _ = test_metadata_index 221 | vector_id = 'vec-update' 222 | values = np.random.rand(d).astype(np.float32).tolist() 223 | # first write 224 | old_md = {'value': 11, 'weather': 'chilly'} 225 | index.upsert(vectors=[Vector(id=vector_id, values=values, metadata=old_md)]) 226 | retry_assert(lambda: list(index.fetch(ids=[vector_id]).values())[0].metadata == old_md) 227 | 228 | # second write 229 | new_md = {'value': 12, 'weather': 'sunny'} 230 | index.upsert(vectors=[Vector(vector_id, values=values, metadata=new_md)]) 231 | retry_assert(lambda: list(index.fetch(ids=[vector_id]).values())[0].metadata == new_md) 232 | 233 | 234 | def test_multiple_values(test_metadata_index, test_data): 235 | index, _ = test_metadata_index 236 | ids = test_data.ids 237 | vectors = test_data.vectors 238 | query_vector = test_data.query_vector 239 | # multiple values 240 | in_vals = ['snowy', 'rainy', 'chilly'] 241 | unique_value = 812312 242 | md = {'value': unique_value, 'weather': in_vals} 243 | index.upsert(vectors=[Vector(id=ids[0], values=vectors[0], metadata=md)]) 244 | retry_assert(lambda: list(index.fetch(ids=[ids[0]]).values())[0].metadata == md) 245 | 246 | for val in in_vals: 247 | eq_filter = {'value': unique_value, 'weather': val} 248 | query_response = get_query_results(index, query_vector, eq_filter) 249 | matches = query_response 250 | assert len(matches) == 1 251 | assert matches[0].metadata['value'] == unique_value 252 | assert matches[0].metadata['weather'] == in_vals 253 | 254 | 255 | # TODO: Fix this test after metadata config is finalized 256 | # def test_metadata_mapping(test_metadata_index_with_mapping, test_data_for_mapping): 257 | # index = test_metadata_index_with_mapping[0] 258 | # query_vector = test_data_for_mapping.query_vector 259 | 260 | # # Check for a non-indexd field 261 | # eq_filter = {"value": {"$eq": 25}} 262 | # query_response = get_query_results(index, query_vector, eq_filter) 263 | # result = query_response[0] 264 | 265 | # assert len(result) == 0 266 | 267 | # # Check for an indexed field 268 | # eq_filter = {"weather": {"$eq": "sunny"}} 269 | # query_response = get_query_results(index, query_vector, eq_filter) 270 | # result = query_response[0] 271 | # print(result) 272 | # assert len(result) != 0 -------------------------------------------------------------------------------- /tests/unit/test_upsert_format.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | 3 | import numpy as np 4 | import pytest 5 | from pinecone import Client, Vector, SparseValues 6 | import os 7 | from ..utils.remote_index import RemoteIndex, PodType 8 | from ..utils.utils import index_fixture_factory, retry_assert 9 | 10 | vector_dim = 512 11 | env = os.getenv('PINECONE_REGION') 12 | api_key = os.getenv('PINECONE_API_KEY') 13 | client = Client(api_key,env) 14 | INDEX_NAME_PREFIX = 'upsert-format' 15 | 16 | test_data_plane_index = index_fixture_factory( 17 | [ 18 | (RemoteIndex(pods=2, index_name=f'{INDEX_NAME_PREFIX}-{PodType.P1}', 19 | dimension=vector_dim, pod_type=PodType.P1), str(PodType.P1)) 20 | ] 21 | ) 22 | 23 | def assert_sparse_vectors(sparse_vec_1:SparseValues,sparse_vec_2:dict): 24 | assert sparse_vec_1.indices == sparse_vec_2['indices'] 25 | assert sparse_vec_1.values == sparse_vec_2['values'] 26 | 27 | 28 | def get_random_metadata(): 29 | return { 30 | 'some_string': np.random.choice(['action', 'documentary', 'drama']), 31 | 'some_int': np.random.randint(2000, 2021), 32 | 'some_bool' : np.random.choice([True, False]), 33 | 'some_float' : np.random.rand(), 34 | 'some_list' : [np.random.choice(['action', 'documentary', 'drama']) for _ in range(5)] 35 | } 36 | 37 | def get_random_vector(): 38 | return np.random.rand(vector_dim).astype(np.float32).tolist() 39 | 40 | def get_random_sparse_vector(): 41 | indices = np.random.choice(vector_dim, 10, replace=False).astype(np.int32).tolist() 42 | values = np.random.rand(10).astype(np.float32).tolist() 43 | return SparseValues(indices=indices, values=values) 44 | 45 | def get_random_sparse_dict(): 46 | indices = np.random.choice(vector_dim, 10, replace=False).astype(np.int32).tolist() 47 | values = np.random.rand(10).astype(np.float32).tolist() 48 | return {'indices': indices, 'values': values} 49 | 50 | def sparese_dict_to_vec(sparse_vec): 51 | return SparseValues(indices=sparse_vec['indices'], values=sparse_vec['values']) 52 | 53 | def test_upsert_tuplesOfIdVec_UpserWithoutMD(test_data_plane_index): 54 | index, _ = test_data_plane_index 55 | index.upsert([('vec1', get_random_vector()), ('vec2', get_random_vector())], namespace='ns') 56 | 57 | 58 | def test_upsert_tuplesOfIdVecMD_UpsertVectorsWithMD(test_data_plane_index): 59 | index, _ = test_data_plane_index 60 | index.upsert([('vec1', get_random_vector(), get_random_metadata()), ('vec2', get_random_vector(), get_random_metadata())], namespace='ns') 61 | 62 | def test_upsert_vectors_upsertInputVectorsSparse(test_data_plane_index): 63 | index, _ = test_data_plane_index 64 | index.upsert([Vector(id='vec1', values=get_random_vector(), metadata=get_random_metadata(), 65 | sparse_values=get_random_sparse_vector()), 66 | Vector(id='vec2', values=get_random_vector(), metadata=get_random_metadata(), 67 | sparse_values=get_random_sparse_vector())], 68 | namespace='ns') 69 | 70 | def test_upsert_dict(test_data_plane_index): 71 | index, _ = test_data_plane_index 72 | dict1 = {'id': 'vec1', 'values': get_random_vector()} 73 | dict2 = {'id': 'vec2', 'values': get_random_vector()} 74 | index.upsert([dict1, dict2], namespace='ns') 75 | fetched_vectors = index.fetch(['vec1', 'vec2'], namespace='ns') 76 | assert len(fetched_vectors) == 2 77 | assert fetched_vectors.get('vec1').id == 'vec1' 78 | assert fetched_vectors.get('vec2').id == 'vec2' 79 | assert fetched_vectors.get('vec1').values == dict1['values'] 80 | assert fetched_vectors.get('vec2').values == dict2['values'] 81 | 82 | 83 | 84 | def test_upsert_dict_md(test_data_plane_index): 85 | index, _ = test_data_plane_index 86 | dict1 = {'id': 'vec1', 'values': get_random_vector(), 'metadata': get_random_metadata()} 87 | dict2 = {'id': 'vec2', 'values': get_random_vector(), 'metadata': get_random_metadata()} 88 | index.upsert([dict1, dict2], namespace='ns') 89 | fetched_vectors = index.fetch(['vec1', 'vec2'], namespace='ns') 90 | assert len(fetched_vectors) == 2 91 | assert fetched_vectors.get('vec1').id == 'vec1' 92 | assert fetched_vectors.get('vec2').id == 'vec2' 93 | assert fetched_vectors.get('vec1').values == dict1['values'] 94 | assert fetched_vectors.get('vec2').values == dict2['values'] 95 | assert fetched_vectors.get('vec1').metadata == dict1['metadata'] 96 | assert fetched_vectors.get('vec2').metadata == dict2['metadata'] 97 | 98 | def test_upsert_dict_sparse(test_data_plane_index): 99 | index, _ = test_data_plane_index 100 | dict1 = {'id': 'vec1', 'values': get_random_vector(), 101 | 'sparse_values': get_random_sparse_dict()} 102 | dict2 = {'id': 'vec2', 'values': get_random_vector(), 103 | 'sparse_values': get_random_sparse_dict()} 104 | index.upsert([dict1, dict2], namespace='ns') 105 | fetched_vectors = index.fetch(['vec1', 'vec2'], namespace='ns') 106 | assert len(fetched_vectors) == 2 107 | assert fetched_vectors.get('vec1').id == 'vec1' 108 | assert fetched_vectors.get('vec2').id == 'vec2' 109 | assert fetched_vectors.get('vec1').values == dict1['values'] 110 | assert fetched_vectors.get('vec2').values == dict2['values'] 111 | assert_sparse_vectors(fetched_vectors.get('vec1').sparse_values, dict1['sparse_values']) 112 | assert_sparse_vectors(fetched_vectors.get('vec2').sparse_values, dict2['sparse_values']) 113 | 114 | def test_upsert_dict_sparse_md(test_data_plane_index): 115 | index, _ = test_data_plane_index 116 | dict1 = {'id': 'vec1', 'values': get_random_vector(), 117 | 'sparse_values': get_random_sparse_dict(), 'metadata': get_random_metadata()} 118 | dict2 = {'id': 'vec2', 'values': get_random_vector(), 119 | 'sparse_values': get_random_sparse_dict(), 'metadata': get_random_metadata()} 120 | index.upsert([dict1, dict2], namespace='ns') 121 | fetched_vectors = index.fetch(['vec1', 'vec2'], namespace='ns') 122 | assert len(fetched_vectors) == 2 123 | assert fetched_vectors.get('vec1').id == 'vec1' 124 | assert fetched_vectors.get('vec2').id == 'vec2' 125 | assert fetched_vectors.get('vec1').values == dict1['values'] 126 | assert fetched_vectors.get('vec2').values == dict2['values'] 127 | assert_sparse_vectors(fetched_vectors.get('vec1').sparse_values, dict1['sparse_values']) 128 | assert_sparse_vectors(fetched_vectors.get('vec2').sparse_values, dict2['sparse_values']) 129 | assert fetched_vectors.get('vec1').metadata == dict1['metadata'] 130 | assert fetched_vectors.get('vec2').metadata == dict2['metadata'] 131 | 132 | 133 | def test_upsert_dict_negative(test_data_plane_index): 134 | index, _ = test_data_plane_index 135 | 136 | # Missing required keys 137 | dict1 = {'values': get_random_vector()} 138 | dict2 = {'id': 'vec2'} 139 | with pytest.raises(ValueError): 140 | index.upsert([dict1, dict2]) 141 | with pytest.raises(ValueError): 142 | index.upsert([dict1]) 143 | with pytest.raises(ValueError): 144 | index.upsert([dict2]) 145 | 146 | # Excess keys 147 | dict1 = {'id': 'vec1', 'values': get_random_vector()} 148 | dict2 = {'id': 'vec2', 'values': get_random_vector(), 'animal': 'dog'} 149 | with pytest.raises(ValueError) as e: 150 | index.upsert([dict1, dict2]) 151 | assert 'animal' in str(e.value) 152 | 153 | dict1 = {'id': 'vec1', 'values': get_random_vector(), 'metadatta': get_random_metadata()} 154 | dict2 = {'id': 'vec2', 'values': get_random_vector()} 155 | with pytest.raises(ValueError) as e: 156 | index.upsert([dict1, dict2]) 157 | assert 'metadatta' in str(e.value) 158 | 159 | @pytest.mark.parametrize("key,new_val", [ 160 | ("id", 4.2), 161 | ("id", ['vec1']), 162 | ("values", ['the', 'lazy', 'fox']), 163 | ("values", 'the lazy fox'), 164 | ("values", 0.5), 165 | ("metadata", np.nan), 166 | ("metadata", ['key1', 'key2']), 167 | ("sparse_values", 'cat'), 168 | ("sparse_values", []), 169 | ]) 170 | def test_upsert_dict_negative_types(test_data_plane_index, key, new_val): 171 | index, _ = test_data_plane_index 172 | full_dict1 = {'id': 'vec1', 'values': get_random_vector(), 173 | 'sparse_values': get_random_sparse_dict(), 174 | 'metadata': get_random_metadata()} 175 | 176 | dict1 = deepcopy(full_dict1) 177 | dict1[key] = new_val 178 | with pytest.raises(ValueError) as e: 179 | index.upsert([dict1]) 180 | assert key in str(e.value) 181 | 182 | @pytest.mark.parametrize("key,new_val", [ 183 | ("indices", 3), 184 | ("indices", [1.2, 0.5]), 185 | ("values", ['1', '4.4']), 186 | ("values", 0.5), 187 | ]) 188 | def test_upsert_dict_negative_types_sparse(test_data_plane_index, key, new_val): 189 | index, _ = test_data_plane_index 190 | 191 | full_dict1 = {'id': 'vec1', 'values': get_random_vector(), 192 | 'sparse_values': get_random_sparse_dict(), 193 | 'metadata': get_random_metadata()} 194 | 195 | dict1 = deepcopy(full_dict1) 196 | dict1['sparse_values'][key] = new_val 197 | # TODO: Lenght mismatch between indices and values should be done on client or server? 198 | with pytest.raises((Exception,ValueError)) as e: 199 | index.upsert([dict1]) 200 | assert key in str(e.value) 201 | -------------------------------------------------------------------------------- /tests/unit/tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = py37 , py38, py39, py310 3 | skipsdist = True 4 | 5 | [testenv] 6 | deps= 7 | -r {toxinidir}/requirements.txt 8 | 9 | [testenv:integration] 10 | setenv: 11 | PINECONE_API_KEY = {env:PINECONE_API_KEY} 12 | PINECONE_ENVIRONMENT = {env:PINECONE_ENVIRONMENT} 13 | PINECONE_LOGGING = {env:PINECONE_LOGGING:} 14 | PINECONE_INDEX_NAME = {env:PINECONE_INDEX_NAME:} 15 | commands = 16 | pytest --self-contained-html --dist=loadscope --numprocesses 4 --durations=10 --durations-min=1.0 {posargs} 17 | -------------------------------------------------------------------------------- /tests/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pinecone-io/pinecone-client/94133b14f85223272d1cb5cb486460a8647af4d8/tests/utils/__init__.py -------------------------------------------------------------------------------- /tests/utils/remote_index.py: -------------------------------------------------------------------------------- 1 | import time 2 | from loguru import logger 3 | import os 4 | from enum import Enum 5 | import pinecone as pinecone 6 | 7 | from urllib3.exceptions import MaxRetryError 8 | 9 | QUOTA = 2 10 | 11 | class PodType(Enum): 12 | """ 13 | Enum for pod types 14 | """ 15 | P1 = 'p1' 16 | S1 = 's1' 17 | P2 = 'p2' 18 | P1_X2 = 'p1.x2' 19 | S1_X2 = 's1.x2' 20 | P2_X2 = 'p2.x2' 21 | S1H = 's1h' 22 | 23 | def __str__(self): 24 | return self.value 25 | 26 | def as_name(self): 27 | return self.value.replace('.', '-') 28 | 29 | def is_implicitly_x1(self): 30 | return '.' not in self.value 31 | 32 | class RemoteIndex: 33 | index = None 34 | 35 | def __init__(self, pods=1, index_name=None, dimension=512, pod_type="p1", metadata_config=None, 36 | _openapi_client_config=None, source_collection='',metric='cosine'): 37 | self.env = os.getenv('PINECONE_ENVIRONMENT') 38 | self.key = os.getenv('PINECONE_API_KEY') 39 | self.pod_type = pod_type 40 | self.pods = pods 41 | self.index_name = index_name if index_name else 'sdk-citest-{0}'.format(pod_type) 42 | self.dimension = dimension 43 | self.metadata_config = metadata_config 44 | self.source_collection = source_collection 45 | self.metric = metric 46 | self.client = pinecone.Client(self.key, self.env) 47 | 48 | def __enter__(self): 49 | if self.index_name not in self.client.list_indexes(): 50 | index_creation_args = {'name': self.index_name, 51 | 'dimension': self.dimension, 52 | 'pod_type': str(self.pod_type), 53 | 'pods': self.pods, 54 | 'metadata_config': self.metadata_config, 55 | 'source_collection': self.source_collection, 56 | 'metric': self.metric} 57 | self.client.create_index(**index_creation_args) 58 | 59 | self.index = self.client.get_index(self.index_name) 60 | return self.index 61 | 62 | @staticmethod 63 | def wait_for_ready(index): 64 | logger.info('waiting until index gets ready...') 65 | max_attempts = 30 66 | for i in range(max_attempts): 67 | try: 68 | time.sleep(1) 69 | index.describe_index_stats() 70 | break 71 | except (Exception, MaxRetryError): 72 | if i + 1 == max_attempts: 73 | logger.info("Index didn't get ready in time. Raising error...") 74 | raise 75 | 76 | def __exit__(self, exc_type, exc_val, exc_tb): 77 | print('deleting index') 78 | self.client.delete_index(self.index_name) 79 | -------------------------------------------------------------------------------- /tests/utils/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | 4 | from loguru import logger 5 | import pinecone 6 | from pinecone import Client 7 | import pytest 8 | import random 9 | from .remote_index import RemoteIndex 10 | from _pytest.python_api import approx 11 | import numpy as np 12 | 13 | 14 | def index_fixture_factory(remote_indices: [(RemoteIndex, str)]): 15 | """ 16 | Creates and returns a pytest fixture for creating/tearing down indexes to test against. 17 | - adds the xdist testrun_uid to the index name unless include_random_suffix is set False 18 | - fixture yields a pinecone.Index object 19 | """ 20 | 21 | @pytest.fixture(scope="module", params=[remote_index[0] for remote_index in remote_indices], 22 | ids=[remote_index[1] for remote_index in remote_indices]) 23 | def index_fixture(testrun_uid, request): 24 | request.param.index_name = request.param.index_name + '-' + testrun_uid[:8] 25 | 26 | def remove_index(): 27 | env = os.getenv('PINECONE_ENVIRONMENT') 28 | api_key = os.getenv('PINECONE_API_KEY') 29 | client = Client(api_key, env) 30 | if request.param.index_name in client.list_indexes(): 31 | client.delete_index(request.param.index_name) 32 | 33 | # attempt to remove index even if creation raises exception 34 | request.addfinalizer(remove_index) 35 | 36 | logger.info('Proceeding with index_name {}', request.param.index_name) 37 | with request.param as index: 38 | yield index, request.param.index_name 39 | 40 | return index_fixture 41 | 42 | 43 | def retry_assert(fun, max_tries=5): 44 | wait_time = 0.5 45 | for i in range(max_tries): 46 | try: 47 | assert fun() 48 | return 49 | except Exception as e: 50 | if i == max_tries - 1: 51 | raise 52 | time.sleep(wait_time) 53 | wait_time = wait_time * 2 54 | 55 | def sparse_values(dimension=32000, nnz=120): 56 | indices = [] 57 | values = [] 58 | threshold = nnz / dimension 59 | for i in range(dimension): 60 | key = random.uniform(0, 1) 61 | if key < threshold: 62 | indices.append(i) 63 | values.append(random.uniform(0, 1)) 64 | return indices, values 65 | 66 | 67 | def get_vector_count(index, namespace): 68 | stats = index.describe_index_stats().namespaces 69 | if namespace not in stats: 70 | return 0 71 | return stats[namespace].vector_count 72 | 73 | 74 | def approx_sparse_equals(sv1, sv2): 75 | if sv1 is None and sv2 is None: 76 | return True 77 | return sv1.indices == sv2.indices and sv1.values == approx(sv2.values) --------------------------------------------------------------------------------