File size: 2,102 Bytes
16d5a75
 
 
55e58da
 
 
 
 
16d5a75
 
 
 
 
 
 
744b763
55e58da
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
from .llm import embeddings
import os
from langchain_pinecone import PineconeVectorStore
from pydantic import BaseModel
from langchain_google_genai.embeddings import GoogleGenerativeAIEmbeddings
from typing import List, Dict, Any, Optional
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStore

API_PINCONE_KEY = os.getenv("PINECONE_API_KEY")

test_rag_vector_store = PineconeVectorStore(
    index_name="rag-vector-store",
    embedding=embeddings,
    pinecone_api_key=API_PINCONE_KEY,
)


class PineconeVectorStoreCRUD(BaseModel):
    index_name: str
    embedding: GoogleGenerativeAIEmbeddings
    pinecone_api_key: str

    def __init__(
        self,
        index_name: str,
        embedding: GoogleGenerativeAIEmbeddings,
        pinecone_api_key: str,
        k: int = 5,
        score_threshold: float = 0.3,
    ) -> VectorStore:
        self.vector_store = PineconeVectorStore(
            index_name=index_name,
            embedding=embedding,
            pinecone_api_key=pinecone_api_key,
        )
        self.retriever = self.vector_store.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={"k": k, "score_threshold": score_threshold},
        )

    async def search(self, query: str, filter: Optional[Dict[str, Any]] = None):
        return await self.retriever.ainvoke(query, filter=filter)

    async def add_documents(self, documents: List[Document], ids: List[str]):
        await self.vector_store.aadd_documents(documents, ids=ids)

    async def get_documents(self, filter: Optional[Dict[str, Any]] = None):
        return await self.vector_store.asimilarity_search("", filter=filter)

    async def delete_documents(self, ids: List[str]):
        await self.vector_store.adelete(ids=ids)

    async def update_documents(self, documents: List[Document], ids: List[str]):
        await self.vector_store.aadd_documents(documents, ids=ids)


test_rag_vector_store = PineconeVectorStore(
    index_name="rag-vector-store",
    embedding=embeddings,
    pinecone_api_key=API_PINCONE_KEY,
)