From 2a12ed87d64978165451d693f0a8acbe276b66ab Mon Sep 17 00:00:00 2001 From: Luca Mannini <48441989+italianconcerto@users.noreply.github.com> Date: Mon, 2 Dec 2024 16:20:48 +0100 Subject: [PATCH 1/2] Reranking made async --- examples/example_usage.py | 26 ++++++++++++++++-- src/pinecone_async/client.py | 52 ++++++++++++++++++++++++++++++++++-- src/pinecone_async/schema.py | 33 ++++++++++++++++++++++- 3 files changed, 106 insertions(+), 5 deletions(-) diff --git a/examples/example_usage.py b/examples/example_usage.py index 7b7d283..cac2169 100644 --- a/examples/example_usage.py +++ b/examples/example_usage.py @@ -2,7 +2,7 @@ import os from pinecone_async import PineconeClient -async def main(): +async def list_indexes_example(): api_key = os.getenv("PINECONE_API_KEY") if not api_key: raise ValueError("PINECONE_API_KEY environment variable is not set") @@ -11,6 +11,28 @@ async def main(): indexes = await client.list_indexes() print(f"Indexes: {indexes}") +async def rerank_example(): + api_key = os.getenv("PINECONE_API_KEY") + if not api_key: + raise ValueError("PINECONE_API_KEY environment variable is not set") + client = PineconeClient(api_key=api_key) + + result = await client.rerank( + model="bge-reranker-v2-m3", + query="The tech company Apple is known for its innovative products like the iPhone.", + documents=[ + {"id": "vec1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, + {"id": "vec2", "text": "Many people enjoy eating apples as a healthy snack."}, + {"id": "vec3", "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, + {"id": "vec4", "text": "An apple a day keeps the doctor away, as the saying goes."}, + ], + top_n=4, + return_documents=True, + parameters={ + "truncate": "END" + } + ) + print(result) if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(rerank_example()) diff --git a/src/pinecone_async/client.py b/src/pinecone_async/client.py index 87b638f..2eb0a7e 100644 --- a/src/pinecone_async/client.py +++ b/src/pinecone_async/client.py @@ -1,10 +1,10 @@ import os -from typing import Any, Dict, List, Literal +from typing import Any, Dict, List, Literal, Optional import httpx from pinecone_async.exceptions import IndexNotFoundError -from pinecone_async.schema import IndexResponse, PineconePod, Serverless +from pinecone_async.schema import Document, IndexResponse, PineconePod, RerankParameters, RerankRequest, RerankResponse, Serverless class PineconeClient: @@ -99,3 +99,51 @@ async def __aenter__(self): async def __aexit__(self, exc_type, exc_val, exc_tb): await self.close() + async def rerank( + self, + model: str, + query: str, + documents: list[dict[str, str]], + top_n: Optional[int] = None, + return_documents: Optional[bool] = True, + parameters: Optional[dict] = None, + rank_fields: Optional[list[str]] = None + ) -> RerankResponse: + """ + Rerank documents based on their relevance to a query. + Args: + model: The reranking model to use (e.g., "bge-reranker-v2-m3") + query: The query text to compare documents against + documents: List of documents to rerank + top_n: Number of top results to return + return_documents: Whether to include documents in response + parameters: Additional parameters like truncation + rank_fields: Optional list of custom fields to rank on + """ + headers = { + "Api-Key": self.headers["Api-Key"], + "Content-Type": "application/json", + "Accept": "application/json", + "X-Pinecone-API-Version": "2024-10" + } + + request = RerankRequest( + model=model, + query=query, + documents=[Document(**doc) for doc in documents], + top_n=top_n, + return_documents=return_documents, + parameters=RerankParameters(**(parameters or {})), + rank_fields=rank_fields + ) + + async with httpx.AsyncClient(headers=headers) as client: + response = await client.post( + "https://api.pinecone.io/rerank", + json=request.model_dump(exclude_none=True) + ) + + if response.status_code == 200: + return RerankResponse(**response.json()) + else: + raise Exception(f"Failed to rerank: {response.status_code} : {response.text}") \ No newline at end of file diff --git a/src/pinecone_async/schema.py b/src/pinecone_async/schema.py index 6b70d1e..3e1d2ce 100644 --- a/src/pinecone_async/schema.py +++ b/src/pinecone_async/schema.py @@ -124,4 +124,35 @@ class ListResponse(BaseModel): class FetchResponse(BaseModel): """Response from fetch operation""" vectors: Dict[str, PineconeVector] - namespace: Optional[str] = None \ No newline at end of file + namespace: Optional[str] = None + + + +class Document(BaseModel): + id: str + text: str | None = None + my_field: str | None = None # For custom field support + +class RerankParameters(BaseModel): + truncate: Optional[Literal["START", "END", "NONE"]] = "END" + +class RerankRequest(BaseModel): + model: str + query: str + documents: list[Document] + top_n: Optional[int] = None + return_documents: Optional[bool] = True + parameters: Optional[RerankParameters] = None + rank_fields: Optional[list[str]] = None + +class RerankResult(BaseModel): + index: int + document: Optional[Document] = None + score: float + +class RerankUsage(BaseModel): + rerank_units: int + +class RerankResponse(BaseModel): + data: list[RerankResult] + usage: RerankUsage \ No newline at end of file From 6f838aad40e0deed20f67424f81ad0493c943ebd Mon Sep 17 00:00:00 2001 From: Luca Mannini <48441989+italianconcerto@users.noreply.github.com> Date: Mon, 2 Dec 2024 17:40:44 +0100 Subject: [PATCH 2/2] Deleted comment --- src/pinecone_async/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pinecone_async/__init__.py b/src/pinecone_async/__init__.py index 304e67b..958a7f9 100644 --- a/src/pinecone_async/__init__.py +++ b/src/pinecone_async/__init__.py @@ -1,4 +1,3 @@ -# src/pinecone_async/__init__.py from .client import PineconeClient from .index import PineconeIndex from .schema import (