Skip to content

Commit

Permalink
Add inference plugin and tests (#370)
Browse files Browse the repository at this point in the history
## Problem

We want to ship inference as part of the core SDK

## Solution

Add a dependency on v1.0.2 of the inference plugin, which includes a
recent fix to plugin installation for use with the GRPC client.

## Usage

The plugin should automatically be installed and available for usage.

```python
from pinecone import Pinecone

pc = Pinecone(api_key='key')

embedding_model = "multilingual-e5-large"
embeddings = pc.inference.embed(
    model=embedding_model,
    inputs=["The quick brown fox jumps over the lazy dog.", "lorem ipsum"],
    parameters={"input_type": "query", "truncate": "END"},
)
```

## Type of Change

- [x] New feature (non-breaking change which adds functionality)

## Test Plan

See tests passing.
  • Loading branch information
jhamon authored Jul 19, 2024
1 parent 8e55a96 commit ddfa538
Show file tree
Hide file tree
Showing 6 changed files with 79 additions and 48 deletions.
68 changes: 21 additions & 47 deletions .github/workflows/testing-integration.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -3,53 +3,27 @@ name: "Integration Tests"
workflow_call: {}

jobs:
# setup-index:
# name: Setup proxyconfig test index
# runs-on: ubuntu-latest
# outputs:
# index_name: ${{ steps.create-index.outputs.index_name }}
# steps:
# - uses: actions/checkout@v4
# - name: Create index
# id: create-index
# uses: ./.github/actions/create-index
# with:
# PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }}
# NAME_PREFIX: 'proxyconfig-'
# REGION: 'us-west-2'
# CLOUD: 'aws'
# DIMENSION: 1536
# METRIC: 'cosine'

# proxy-config:
# name: Proxy config tests
# needs: [setup-index]
# runs-on: ubuntu-latest
# steps:
# - uses: actions/checkout@v4
# - uses: actions/setup-python@v5
# with:
# python-version: 3.9
# - name: Setup Poetry
# uses: ./.github/actions/setup-poetry
# - name: 'Run integration tests (proxy config)'
# run: |
# poetry run pytest tests/integration/proxy_config -s -v
# env:
# PINECONE_API_KEY: '${{ secrets.PINECONE_API_KEY }}'
# PINECONE_INDEX_NAME: ${{ needs.setup-index.outputs.index_name }}
# - name: Upload logs
# if: always()
# uses: actions/upload-artifact@v4
# with:
# name: proxy_config_test_logs
# path: tests/integration/proxy_config/logs
# - name: Cleanup index
# if: always()
# uses: ./.github/actions/delete-index
# with:
# PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }}
# INDEX_NAME: ${{ needs.setup-index.outputs.index_name }}
plugin-inference:
name: Test inference plugin
runs-on: ubuntu-latest
strategy:
matrix:
python_version: [3.8, 3.12]
steps:
- uses: actions/checkout@v4
- name: 'Set up Python ${{ matrix.python_version }}'
uses: actions/setup-python@v5
with:
python-version: '${{ matrix.python_version }}'
- name: Setup Poetry
uses: ./.github/actions/setup-poetry
with:
include_grpc: 'true'
- name: 'Run integration tests'
run: poetry run pytest tests/integration/inference -s -vv
env:
PINECONE_DEBUG_CURL: 'true'
PINECONE_API_KEY: '${{ secrets.PINECONE_API_KEY }}'

data-plane-serverless:
name: Data plane serverless integration tests
Expand Down
17 changes: 16 additions & 1 deletion poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ lz4 = { version = ">=3.1.3", optional = true }
protobuf = { version = "^4.25", optional = true }
protoc-gen-openapiv2 = {version = "^0.0.1", optional = true }
pinecone-plugin-interface = "^0.0.7"
pinecone-plugin-inference = "1.0.2"

[tool.poetry.group.types]
optional = true
Expand Down
Empty file.
7 changes: 7 additions & 0 deletions tests/integration/inference/conftest.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
import pytest
from ..helpers import get_environment_var


@pytest.fixture()
def api_key():
return get_environment_var("PINECONE_API_KEY")
34 changes: 34 additions & 0 deletions tests/integration/inference/test_embed.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
from pinecone import Pinecone
from pinecone.grpc import PineconeGRPC


class TestInferencePlugin:
def test_embed(self, api_key):
pc = Pinecone(api_key=api_key)

embedding_model = "multilingual-e5-large"
embeddings = pc.inference.embed(
model=embedding_model,
inputs=["The quick brown fox jumps over the lazy dog.", "lorem ipsum"],
parameters={"input_type": "query", "truncate": "END"},
)

assert len(embeddings.get("data")) == 2
assert len(embeddings.get("data")[0]["values"]) == 1024
assert len(embeddings.get("data")[1]["values"]) == 1024
assert embeddings.get("model") == embedding_model

def test_embed_grpc(self, api_key):
pc = PineconeGRPC(api_key=api_key)

embedding_model = "multilingual-e5-large"
embeddings = pc.inference.embed(
model=embedding_model,
inputs=["The quick brown fox jumps over the lazy dog.", "lorem ipsum"],
parameters={"input_type": "query", "truncate": "END"},
)

assert len(embeddings.get("data")) == 2
assert len(embeddings.get("data")[0]["values"]) == 1024
assert len(embeddings.get("data")[1]["values"]) == 1024
assert embeddings.get("model") == embedding_model

0 comments on commit ddfa538

Please sign in to comment.