Ayomide
71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
from .config import Config
|
|
from pinecone import Pinecone, ServerlessSpec
|
|
from typing import List, Optional
|
|
import time
|
|
|
|
|
|
class VectorStore:
|
|
def __init__(self):
|
|
if Config.VECTOR_STORE_TYPE == "pinecone":
|
|
self.pc = Pinecone(api_key=Config.PINECONE_API_KEY)
|
|
|
|
# Free tier supported regions
|
|
FREE_TIER_SUPPORTED_REGIONS = {
|
|
'aws': 'us-east-1',
|
|
'gcp': 'us-central1'
|
|
}
|
|
|
|
# Check if index exists
|
|
if Config.PINECONE_INDEX not in self.pc.list_indexes().names():
|
|
print(f"Creating new Pinecone index: {Config.PINECONE_INDEX}")
|
|
try:
|
|
# First try AWS free tier region
|
|
self.pc.create_index(
|
|
name=Config.PINECONE_INDEX,
|
|
dimension=1024, # Cohere embed-english-v3.0 dimension
|
|
metric="cosine",
|
|
spec=ServerlessSpec(
|
|
cloud="aws",
|
|
region=FREE_TIER_SUPPORTED_REGIONS['aws']
|
|
)
|
|
)
|
|
except Exception as e:
|
|
print(f"AWS region failed, trying GCP: {str(e)}")
|
|
# Fallback to GCP if AWS fails
|
|
self.pc.create_index(
|
|
name=Config.PINECONE_INDEX,
|
|
dimension=1024,
|
|
metric="cosine",
|
|
spec=ServerlessSpec(
|
|
cloud="gcp",
|
|
region=FREE_TIER_SUPPORTED_REGIONS['gcp']
|
|
)
|
|
)
|
|
|
|
# Wait for index to initialize
|
|
time.sleep(1)
|
|
|
|
self.index = self.pc.Index(Config.PINECONE_INDEX)
|
|
|
|
def upsert_document(self, doc_id: str, embedding: List[float], metadata: dict):
|
|
self.index.upsert(
|
|
vectors=[{
|
|
"id": doc_id,
|
|
"values": embedding,
|
|
"metadata": metadata
|
|
}]
|
|
)
|
|
|
|
def search_similar(self, embedding: List[float], top_k: int = 3):
|
|
return self.index.query(
|
|
vector=embedding,
|
|
top_k=top_k,
|
|
include_metadata=True
|
|
)
|
|
|
|
def get_document(self, doc_id: str) -> Optional[dict]:
|
|
fetch_response = self.index.fetch(ids=[doc_id])
|
|
if doc_id in fetch_response.vectors:
|
|
return fetch_response.vectors[doc_id]
|
|
return None
|