| """ |
| Pinecone vector store wrapper. |
| Reconstructed to match the interface app.py expects: |
| - create_vector_store_from_env() |
| - .add_documents(documents, chunk_size, chunk_overlap) |
| - .get_index_stats() |
| - .similarity_search(...) / .similarity_search_with_score(...) |
| - .delete_index() |
| """ |
|
|
| import os |
| import time |
| from typing import List, Tuple |
|
|
| from langchain.schema import Document |
| from langchain.text_splitter import RecursiveCharacterTextSplitter |
| from langchain_openai import OpenAIEmbeddings |
| from langchain_pinecone import PineconeVectorStore |
| from pinecone import Pinecone, ServerlessSpec |
|
|
| |
| EMBEDDING_DIMENSION = 1536 |
|
|
|
|
| class VectorStore: |
| """Thin wrapper around a Pinecone serverless index + LangChain embeddings.""" |
|
|
| def __init__( |
| self, |
| api_key: str, |
| index_name: str, |
| embedding_model: str = "text-embedding-3-small", |
| cloud: str = "aws", |
| region: str = "us-east-1", |
| ): |
| self.index_name = index_name |
| self.pc = Pinecone(api_key=api_key) |
| self.embeddings = OpenAIEmbeddings(model=embedding_model) |
|
|
| existing = [idx["name"] for idx in self.pc.list_indexes()] |
| if index_name not in existing: |
| self.pc.create_index( |
| name=index_name, |
| dimension=EMBEDDING_DIMENSION, |
| metric="cosine", |
| spec=ServerlessSpec(cloud=cloud, region=region), |
| ) |
| |
| while not self.pc.describe_index(index_name).status["ready"]: |
| time.sleep(1) |
|
|
| self.index = self.pc.Index(index_name) |
| self._store = PineconeVectorStore( |
| index=self.index, embedding=self.embeddings |
| ) |
|
|
| def add_documents( |
| self, |
| documents: List[Document], |
| chunk_size: int = 1000, |
| chunk_overlap: int = 200, |
| ) -> int: |
| """Split documents into chunks and upsert them. Returns chunk count.""" |
| splitter = RecursiveCharacterTextSplitter( |
| chunk_size=chunk_size, |
| chunk_overlap=chunk_overlap, |
| separators=["\n\n", "\n", ". ", " ", ""], |
| ) |
| chunks = splitter.split_documents(documents) |
| if chunks: |
| self._store.add_documents(chunks) |
| return len(chunks) |
|
|
| def similarity_search(self, query: str, k: int = 5) -> List[Document]: |
| return self._store.similarity_search(query, k=k) |
|
|
| def similarity_search_with_score( |
| self, query: str, k: int = 5 |
| ) -> List[Tuple[Document, float]]: |
| return self._store.similarity_search_with_score(query, k=k) |
|
|
| def as_retriever(self, k: int = 5): |
| return self._store.as_retriever(search_kwargs={"k": k}) |
|
|
| def get_index_stats(self) -> dict: |
| stats = self.index.describe_index_stats() |
| return { |
| "total_vector_count": stats.get("total_vector_count", 0), |
| "dimension": stats.get("dimension", EMBEDDING_DIMENSION), |
| } |
|
|
| def delete_index(self) -> None: |
| """Destroy the entire index. Irreversible.""" |
| self.pc.delete_index(self.index_name) |
|
|
|
|
| def create_vector_store_from_env() -> VectorStore: |
| """Build a VectorStore from environment variables.""" |
| api_key = os.getenv("PINECONE_API_KEY") |
| index_name = os.getenv("PINECONE_INDEX_NAME", "rag-demo") |
| region = os.getenv("PINECONE_ENVIRONMENT", "us-east-1") |
| embedding_model = os.getenv( |
| "OPENAI_EMBEDDING_MODEL", "text-embedding-3-small" |
| ) |
|
|
| if not api_key: |
| raise ValueError("PINECONE_API_KEY is not set") |
|
|
| |
| region = region.replace("-aws", "").replace("-gcp", "") |
|
|
| return VectorStore( |
| api_key=api_key, |
| index_name=index_name, |
| embedding_model=embedding_model, |
| region=region, |
| ) |
|
|
