blip / src /pinecone_store.py
amezi's picture
more minor bux fixedd
2e8783e
raw
history blame
1.02 kB
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, # This must match InternVL2_5-8B-MPO
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"]
]