From 894b6712382fa75936f20b391f88de37864d94a4 Mon Sep 17 00:00:00 2001 From: Jen Hamon Date: Fri, 19 Jul 2024 11:06:48 -0400 Subject: [PATCH] Add inference plugin and tests --- .github/workflows/testing-integration.yaml | 68 +++++++--------------- poetry.lock | 17 +++++- pyproject.toml | 1 + tests/integration/inference/__init__.py | 0 tests/integration/inference/conftest.py | 6 ++ tests/integration/inference/test_embed.py | 33 +++++++++++ 6 files changed, 77 insertions(+), 48 deletions(-) create mode 100644 tests/integration/inference/__init__.py create mode 100644 tests/integration/inference/conftest.py create mode 100644 tests/integration/inference/test_embed.py diff --git a/.github/workflows/testing-integration.yaml b/.github/workflows/testing-integration.yaml index c2ee05bf..6a4bbdb1 100644 --- a/.github/workflows/testing-integration.yaml +++ b/.github/workflows/testing-integration.yaml @@ -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 diff --git a/poetry.lock b/poetry.lock index 8513ca3d..2f84b616 100644 --- a/poetry.lock +++ b/poetry.lock @@ -868,6 +868,21 @@ pygments = ">=2.12.0" [package.extras] dev = ["hypothesis", "mypy", "pdoc-pyo3-sample-library (==1.0.11)", "pygments (>=2.14.0)", "pytest", "pytest-cov", "pytest-timeout", "ruff", "tox", "types-pygments"] +[[package]] +name = "pinecone-plugin-inference" +version = "1.0.2" +description = "Embeddings plugin for Pinecone SDK" +optional = false +python-versions = "<4.0,>=3.8" +files = [ + {file = "pinecone_plugin_inference-1.0.2-py3-none-any.whl", hash = "sha256:ef94e7f51554e780408cf1507888b120bb0d185b8485a903cbeb9d0176ee03f1"}, + {file = "pinecone_plugin_inference-1.0.2.tar.gz", hash = "sha256:dda62d9ff44dbbf191b11e6e884235f329cf0ec22edc616fe7efcb1e479e6a9a"}, +] + +[package.dependencies] +pinecone-client = ">=4.1.1,<6.0.0" +pinecone-plugin-interface = ">=0.0.7,<0.0.8" + [[package]] name = "pinecone-plugin-interface" version = "0.0.7" @@ -1369,4 +1384,4 @@ grpc = ["googleapis-common-protos", "grpcio", "grpcio", "lz4", "protobuf", "prot [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "1c61ed3540dc2bf608b4c29a42f1a0a6a0a5cf48faba3c1e0968ae783091a40f" +content-hash = "29a770771daba5a228a5b291c4f6d676a12758fe78e6b0d4eeb3ac16b0446e38" diff --git a/pyproject.toml b/pyproject.toml index d248dea8..cf77a99f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 diff --git a/tests/integration/inference/__init__.py b/tests/integration/inference/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/inference/conftest.py b/tests/integration/inference/conftest.py new file mode 100644 index 00000000..9d967dd3 --- /dev/null +++ b/tests/integration/inference/conftest.py @@ -0,0 +1,6 @@ +import pytest +from ..helpers import get_environment_var + +@pytest.fixture() +def api_key(): + return get_environment_var("PINECONE_API_KEY") diff --git a/tests/integration/inference/test_embed.py b/tests/integration/inference/test_embed.py new file mode 100644 index 00000000..4b599fb2 --- /dev/null +++ b/tests/integration/inference/test_embed.py @@ -0,0 +1,33 @@ +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 \ No newline at end of file