Skip to content

Conversation

jhamon
Copy link
Collaborator

@jhamon jhamon commented Jan 31, 2025

Problem

Previous work implemented asyncio for the db data plane, and now we want to roll out a similar approach for the db control plane and inference as well.

Solution

  • Extract request construction logic out of Pinecone and move it to a request factory
  • Implement PineconeAsyncio using the request factory to keep most of the method-specific logic the same.
  • Add new integration tests using the asyncio code path. These are mostly modified from the existing serverless integration tests.
  • Update tests for the asyncio index client to reflect new setup steps
  • Some refactorings around async context management to address log warnings being shown from aiohttp

Usage

The async version of the client has some async setup/teardown related to the underlying aiohttp library being used. You can either use the async with syntax to have the async context automatically managed for you. Or, if you prefer, you can take the responsibility to close the async context yourself by using close().

Context management option 1: Using async with

import asyncio
from pinecone import (
    PineconeAsyncio,
    ServerlessSpec,
    CloudProvider,
    AwsRegion,
)

async def main():
    async with PineconeAsyncio(api_key="key") as pc:
        await pc.create_index(
            name="my-index",
            metric="cosine",
            spec=ServerlessSpec(
                cloud=CloudProvider.AWS, 
                region=AwsRegion.US_EAST_1
            ),
        )
        
asyncio.run(main())

Context management option 2: Manually close()

import asyncio
from pinecone import (
    PineconeAsyncio,
    ServerlessSpec,
    CloudProvider,
    AwsRegion,
)

async def main():
    pc = PineconeAsyncio(api_key="key")
    await pc.create_index(
        name="my-index",
        metric="cosine",
        spec=ServerlessSpec(
            cloud=CloudProvider.AWS, 
            region=AwsRegion.US_EAST_1
        ),
    )
    await pc.close() # <-- Don't forget to close the client when you are done mking network calls
        
asyncio.run(main())

Sparse index example

import asyncio
import random
from pinecone import (
    PineconeAsyncio,
    ServerlessSpec,
    CloudProvider,
    AwsRegion,
    Metric,
    VectorType,
    Vector,
    SparseValues,
)

async def main():
    async with PineconeAsyncio() as pc:
        # Create a sparse index
        index_name = "my-index2"
        
        if not await pc.has_index(index_name):
            await pc.create_index(
                name=index_name,
                metric=Metric.DOTPRODUCT,
                spec=ServerlessSpec(
                    cloud=CloudProvider.AWS, 
                    region=AwsRegion.US_EAST_1
                ),
                vector_type=VectorType.SPARSE,
                tags={
                    "env": "testing",
                }
            )
        
        # Get the index host
        description = await pc.describe_index(name=index_name)
        
        # Make an index client
        async with pc.Index(host=description.host) as idx:

            # Upsert some sparse vectors
            await idx.upsert(
                vectors=[
                    Vector(
                        id=str(i), 
                        sparse_values=SparseValues(
                            indices=[j for j in range(100)], 
                            values=[random.random() for _ in range(100)]
                        )
                    ) for i in range(50)
                ]
            )
            
            # Query the index
            query_results = await idx.query(
                top_k=5,
                sparse_vector=SparseValues(
                    indices=[5, 10, 20], 
                    values=[0.5, 0.5, 0.5]
                ),
            )
            print(query_results)

asyncio.run(main())

Type of Change

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

@jhamon jhamon changed the base branch from main to release-candidate/2025-01 January 31, 2025 10:57
@jhamon jhamon marked this pull request as ready for review January 31, 2025 14:54
@jhamon jhamon merged commit 6e81b12 into release-candidate/2025-01 Jan 31, 2025
75 checks passed
@jhamon jhamon deleted the jhamon/asyncio-control-plane branch January 31, 2025 14:55
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.

1 participant