Skip to content

Comments

Migrate create_index_for_model from records plugin#438

Merged
jhamon merged 6 commits intorelease-candidate/2025-01from
jhamon/create-index-for-model
Jan 29, 2025
Merged

Migrate create_index_for_model from records plugin#438
jhamon merged 6 commits intorelease-candidate/2025-01from
jhamon/create-index-for-model

Conversation

@jhamon
Copy link
Collaborator

@jhamon jhamon commented Jan 29, 2025

Problem

We need to migrate create_index_for_model functionality from the records plugin into the core of the SDK to provide improved UX around code completions and error handling

Solution

  • Copy & modify implementation of create_index_for_model and IndexEmbed from records plugin
  • I did some moderately heavy refactoring of this code and also the existing create_index method to reduce the amount of duplication.
  • Add records plugin to list of deprecated plugins that will halt startup with an exception
  • Copy relevant tests from plugin and supplement with additional tests to ensure enum types and edge cases are handled properly even after refactoring
  • Reuse recently-defined enums to make the method signature more self-documenting.

Todo

  • Data plane operations in the records plugin will be handled in a separate PR

Usage

These would all be considered valid usage. Enums are available to help know what values are accepted, but you can type the literal strings if you prefer. This flexibility also keeps compatibility with existing usage of the plugin.

from pinecone import (
    Pinecone, 
    EmbedModel, 
    CloudProvider, 
    AwsRegion, 
    IndexEmbed, 
    Metric
)

pc = Pinecone(api_key='key')

# All hand-crafted literals
pc.create_index_for_model(
    name='index-name',
    cloud='aws',
    region='us-east-1',
    embed={
        "model": "multilingual-e5-large", 
        "field_map": {"text": "my-sample-text"},
        "metric": "cosine"
    },
)

# All enum values
pc.create_index_for_model(
    name='index-name',
    cloud=CloudProvider.AWS,
    region=AwsRegion.US_EAST_1,
    embed=IndexEmbed(
        model=EmbedModel.Multilingual_E5_Large, 
        field_map={"text": "my-sample-text"},
        metric=Metric.COSINE
    ),
)

# Mixed literals and enums
pc.create_index_for_model(
    name='index-name',
    cloud='aws',
    region=AwsRegion.US_EAST_1,
    embed={
        "model": EmbedModel.Multilingual_E5_Large, 
        "field_map": {"text": "my-sample-text"},
        "metric": "cosine"
    },
)

Type of Change

  • New feature (non-breaking change which adds functionality)

@jhamon jhamon marked this pull request as ready for review January 29, 2025 14:27
Copy link
Contributor

@austin-denoble austin-denoble left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work, big thanks again for all the plugin porting. Loving the test expansion, and general UX improvement through enums, interface flexibility, etc. 🎉

"""
return self.__dict__

def __init__(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks very much for improving this generally, I don't think I left things in the best state in the plugin itself.

Comment on lines +16 to +20
try:
from pinecone_plugins.records import __installables__ # type: ignore

if __installables__ is not None:
raise DeprecatedPluginError("pinecone-plugin-records")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Really like your solution for this. 👍

@jhamon jhamon merged commit 8f2b09d into release-candidate/2025-01 Jan 29, 2025
70 of 71 checks passed
@jhamon jhamon deleted the jhamon/create-index-for-model branch January 29, 2025 20:23
@jhamon jhamon mentioned this pull request Jan 30, 2025
7 tasks
jhamon added a commit that referenced this pull request Jan 30, 2025
## Problem

Migrating `search_records` (aliased to `search`) and `upsert_records`
from the `pinecone-plugin-records` plugin.

## Solution

Working off the content of the records plugin, I have done the
following:

- Adjusted the codegen script to fix the way openapi generator handles
underscore fields such as `_id` and `_score`
- Adjusted the rest library code in `rest_urllib3.py` and
`rest_aiohttp.py` to handle record uploading with content-type
`application/x-ndjson`
- Copied and modified the integration tests from the plugin
- Extracted a lot of the guts of the `upload_records` and
`search_records` methods into the request factory where they could more
easily be unit tested. The logic around parsing user inputs into the
openapi request objects is surprisingly complicated, so I added quite a
lot of new unit tests checking some of those edge cases.
- Compared to the plugin implementation, the major changes are:
  - Made `search` an alias of `search_records`
- Moved away from usages of `.pop()` which mutates the input objects;
this could be confusing for users if they are using those objects for
anything else
  - Added better typing of dict fields
  - Incorporated optional use of enum values for `RerankModel`
- Added asyncio variants of these methods, although most of the guts are
shared in the request factory.

I already handled disallowing the records plugin in yesterday's PR #438 

## Usage

```python
from pinecone import Pinecone, CloudProvider, AwsRegion, EmbedModel, RerankModel

pc = Pinecone(api_key="key")

# Create an index for your embedding model
index_model = pc.create_index_for_model(
    name="my-model-index",
    cloud=CloudProvider.AWS,
    region=AwsRegion.US_EAST_1,
    embed={
        "model": EmbedModel.Multilingual_E5_Large,
        "field_map": {"text": "my_text_field"}
    }
)

# Create an index client
index = pc.Index(host=index_model.host)

# Upsert records
namespace = "target-namespace"
index.upsert_records(
    namespace=namespace,
    records=[
        {
            "_id": "test1",
            "my_text_field": "Apple is a popular fruit known for its sweetness and crisp texture.",
        },
        {
            "_id": "test2",
            "my_text_field": "The tech company Apple is known for its innovative products like the iPhone.",
        },
        {
            "_id": "test3",
            "my_text_field": "Many people enjoy eating apples as a healthy snack.",
        },
        {
            "_id": "test4",
            "my_text_field": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces.",
        },
        {
            "_id": "test5",
            "my_text_field": "An apple a day keeps the doctor away, as the saying goes.",
        },
        {
            "_id": "test6",
            "my_text_field": "Apple Computer Company was founded on April 1, 1976, by Steve Jobs, Steve Wozniak, and Ronald Wayne as a partnership.",
        },
    ],
)

# Search for similar records
response = index.search(
    namespace=namespace,
    query={
        "inputs":{
            "text": "Apple corporation",
        },
        "top_k":3,
    },
    rerank={
        "model": RerankModel.Bge_Reranker_V2_M3,
        "rank_fields": ["my_text_field"],
        "top_n": 3,
    },
)
```

These methods also have asyncio variants available

```python
import asyncio
from pinecone import Pinecone, RerankModel

async def main():
    # Create an index client

    pc = Pinecone(api_key='key')
    index = pc.AsyncioIndex(host='host')

    # Upsert records
    namespace = "target-namespace"
    records = [
        {
            "_id": "test1",
            "my_text_field": "Apple is a popular fruit known for its sweetness and crisp texture.",
        },
        {
            "_id": "test2",
            "my_text_field": "The tech company Apple is known for its innovative products like the iPhone.",
        },
        {
            "_id": "test3",
            "my_text_field": "Many people enjoy eating apples as a healthy snack.",
        },
        {
            "_id": "test4",
            "my_text_field": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces.",
        },
        {
            "_id": "test5",
            "my_text_field": "An apple a day keeps the doctor away, as the saying goes.",
        },
        {
            "_id": "test6",
            "my_text_field": "Apple Computer Company was founded on April 1, 1976, by Steve Jobs, Steve Wozniak, and Ronald Wayne as a partnership.",
        },
    ]
    await index.upsert_records(
        namespace=namespace,
        records=records,
    )

    # Search for similar records
    response = await index.search(
        namespace=namespace,
        query={
            "inputs":{
                "text": "Apple corporation",
            },
            "top_k":3,
        },
        rerank={
            "model": RerankModel.Bge_Reranker_V2_M3,
            "rank_fields": ["my_text_field"],
            "top_n": 3,
        },
    )
    
asyncio.run(main())
```

## Type of Change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update
- [ ] Infrastructure change (CI configs, etc)
- [ ] Non-code change (docs, etc)
- [ ] None of the above: (explain here)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants