vectorstore.py

import uuid


class vectorstore:
    def __init__(self, pinecone_client):
        self.pc_client = pinecone_client
        self.__PINECONE_INDEX = "ml-yearning"
        self.embedding = {
            "model": "llama-text-embed-v2",
            "field_map": {"text": "chunk_text"}
        }
        self.create_index(self.embedding)
        self.vector_store = self.pc_client.Index(self.__PINECONE_INDEX)


    def create_index(self, embedding):
        if not self.pc_client.has_index(self.__PINECONE_INDEX):
            self.pc_client.create_index_for_model(
                name=self.__PINECONE_INDEX,
                cloud="aws",
                region="us-east-1",
                embed=embedding
            )

    def upsert_doc(self, all_chunks, batch_size=96):
        """Upsert LangChain documents to Pinecone."""
        if all_chunks is None or len(all_chunks) == 0:
            return
        
        # Convert LangChain documents to Pinecone records
        records = []
        for chunk in all_chunks:
            record = {
                "_id": str(uuid.uuid4()),
                "text": chunk.page_content,
            }
            # Flatten metadata into the record (Pinecone doesn't allow nested objects)
            for key, value in chunk.metadata.items():
                # Skip empty values and ensure valid types
                if value is not None and value != "":
                    record[key] = value
            records.append(record)
        
        # Upsert in batches
        for i in range(0, len(records), batch_size):
            batch = records[i:i + batch_size]
            self.vector_store.upsert_records(namespace="default", records=batch)
        
        print(f"Upserted {len(records)} documents")

    def query_doc(self, query, k=3):
        """Search using Pinecone integrated inference."""
        results = self.vector_store.search(
            namespace="default",
            query={"top_k": k, "inputs": {"text": query}},
            fields=["text", "page", "source"]
        )
        return results