|
from pinecone import Pinecone, ServerlessSpec |
|
import os |
|
from dotenv import load_dotenv |
|
|
|
load_dotenv() |
|
|
|
class PineconeStore: |
|
def __init__(self): |
|
self.pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY")) |
|
|
|
if 'blip-soccer-highlights' not in self.pc.list_indexes().names(): |
|
self.pc.create_index( |
|
name='blip-soccer-highlights', |
|
dimension=1024, |
|
metric='cosine', |
|
spec=ServerlessSpec(cloud='aws', region=os.getenv("PINECONE_ENV", "us-east-1")) |
|
) |
|
|
|
self.index = self.pc.Index('blip-soccer-highlights') |
|
|
|
def upsert(self, id, vector, metadata): |
|
self.index.upsert([(id, vector.tolist(), metadata)]) |
|
|
|
def query(self, vector, filter_key, top_k): |
|
results = self.index.query(vector.tolist(), top_k=top_k, include_metadata=True) |
|
return [ |
|
match["metadata"] for match in results["matches"] |
|
if filter_key in match["id"] |
|
] |
|
|