Ayomide
33 lines
1011 B
Python
33 lines
1011 B
Python
from .config import Config
|
|
from pinecone import Pinecone
|
|
from typing import List, Optional
|
|
|
|
|
|
class VectorStore:
|
|
def __init__(self):
|
|
if Config.VECTOR_STORE_TYPE == "pinecone":
|
|
self.pc = Pinecone(api_key=Config.PINECONE_API_KEY)
|
|
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
|