feat: refactor pinecone logic and add pinecone unit tests

Author: lpm0073

models/hybrid_search_retreiver.py

@@ -1,5 +1,4 @@
 # -*- coding: utf-8 -*-
-# pylint: disable=too-few-public-methods
 """
 Hybrid Search Retriever. A class that combines the following:
     - OpenAI prompting and ChatModel
@@ -16,20 +15,15 @@
      https://python.langchain.com/docs/integrations/retrievers/pinecone_hybrid_search
 """
 
-# document loading
-import glob
-
 # general purpose imports
 import logging
-import os
 import textwrap
 from typing import Union
 
 # pinecone integration
 import pinecone
 from langchain.cache import InMemoryCache
 from langchain.chat_models import ChatOpenAI
-from langchain.document_loaders import PyPDFLoader
 
 # embedding
 from langchain.embeddings import OpenAIEmbeddings
@@ -42,12 +36,12 @@
 # hybrid search capability
 from langchain.retrievers import PineconeHybridSearchRetriever
 from langchain.schema import BaseMessage, HumanMessage, SystemMessage
-from langchain.text_splitter import Document
 from langchain.vectorstores.pinecone import Pinecone
 from pinecone_text.sparse import BM25Encoder
 
 # this project
 from models.const import Config, Credentials
+from models.pinecone import PineConeIndex, TextSplitter
 
 
 ###############################################################################
@@ -56,37 +50,29 @@
 logging.basicConfig(level=logging.DEBUG if Config.DEBUG_MODE else logging.INFO)
 
 
-class TextSplitter:
-    """
-    Custom text splitter that adds metadata to the Document object
-    which is required by PineconeHybridSearchRetriever.
-    """
-
-    def create_documents(self, texts):
-        """Create documents"""
-        documents = []
-        for text in texts:
-            # Create a Document object with the text and metadata
-            document = Document(page_content=text, metadata={"context": text})
-            documents.append(document)
-        return documents
-
-
 class HybridSearchRetriever:
     """Hybrid Search Retriever"""
 
     _chat: ChatOpenAI = None
     _openai_embeddings: OpenAIEmbeddings = None
-    _pinecone_index: pinecone.Index = None
     _vector_store: Pinecone = None
     _text_splitter: TextSplitter = None
     _b25_encoder: BM25Encoder = None
+    _pinecone: PineConeIndex = None
+    _retriever: PineconeHybridSearchRetriever = None
 
     def __init__(self):
         """Constructor"""
         pinecone.init(api_key=Credentials.PINECONE_API_KEY, environment=Config.PINECONE_ENVIRONMENT)
         set_llm_cache(InMemoryCache())
 
+    @property
+    def pinecone(self) -> PineConeIndex:
+        """PineConeIndex lazy read-only property."""
+        if self._pinecone is None:
+            self._pinecone = PineConeIndex()
+        return self._pinecone
+
     # prompting wrapper
     @property
     def chat(self) -> ChatOpenAI:
@@ -112,19 +98,12 @@ def openai_embeddings(self) -> OpenAIEmbeddings:
             )
         return self._openai_embeddings
 
-    @property
-    def pinecone_index(self) -> pinecone.Index:
-        """pinecone.Index lazy read-only property."""
-        if self._pinecone_index is None:
-            self._pinecone_index = pinecone.Index(index_name=Config.PINECONE_INDEX_NAME)
-        return self._pinecone_index
-
     @property
     def vector_store(self) -> Pinecone:
         """Pinecone lazy read-only property."""
         if self._vector_store is None:
             self._vector_store = Pinecone(
-                index=self.pinecone_index,
+                index=self.pinecone.index,
                 embedding=self.openai_embeddings,
                 text_key=Config.PINECONE_VECTORSTORE_TEXT_KEY,
             )
@@ -144,6 +123,15 @@ def bm25_encoder(self) -> BM25Encoder:
             self._b25_encoder = BM25Encoder().default()
         return self._b25_encoder
 
+    @property
+    def retriever(self) -> PineconeHybridSearchRetriever:
+        """PineconeHybridSearchRetriever lazy read-only property."""
+        if self._retriever is None:
+            self._retriever = PineconeHybridSearchRetriever(
+                embeddings=self.openai_embeddings, sparse_encoder=self.bm25_encoder, index=self.pinecone.index
+            )
+        return self._retriever
+
     def cached_chat_request(
         self, system_message: Union[str, SystemMessage], human_message: Union[str, HumanMessage]
     ) -> BaseMessage:
@@ -169,54 +157,8 @@ def prompt_with_template(
         return retval
 
     def load(self, filepath: str):
-        """
-        Embed PDF.
-        1. Load PDF document text data
-        2. Split into pages
-        3. Embed each page
-        4. Store in Pinecone
-
-        Note: it's important to make sure that the "context" field that holds the document text
-        in the metadata is not indexed. Currently you need to specify explicitly the fields you
-        do want to index. For more information checkout
-        https://docs.pinecone.io/docs/manage-indexes#selective-metadata-indexing
-        """
-        try:
-            logging.info("Deleting index...")
-            pinecone.delete_index(Config.PINECONE_INDEX_NAME)
-        except pinecone.exceptions.PineconeException:
-            logging.info("Index does not exist. Continuing...")
-
-        metadata_config = {
-            "indexed": [Config.PINECONE_VECTORSTORE_TEXT_KEY, "lc_type"],
-            "context": ["lc_text"],
-        }
-        logging.info("Creating index. This may take a few minutes...")
-        pinecone.create_index(
-            Config.PINECONE_INDEX_NAME,
-            dimension=Config.PINECONE_DIMENSIONS,
-            metric=Config.PINECONE_METRIC,
-            metadata_config=metadata_config,
-        )
-
-        pdf_files = glob.glob(os.path.join(filepath, "*.pdf"))
-        i = 0
-        for pdf_file in pdf_files:
-            i += 1
-            j = len(pdf_files)
-            logging.info("Loading PDF %s of %s: %s", i, j, pdf_file)
-            loader = PyPDFLoader(file_path=pdf_file)
-            docs = loader.load()
-            k = 0
-            for doc in docs:
-                k += 1
-                logging.info(k * "-", end="\r")
-                documents = self.text_splitter.create_documents([doc.page_content])
-                document_texts = [doc.page_content for doc in documents]
-                embeddings = self.openai_embeddings.embed_documents(document_texts)
-                self.vector_store.add_documents(documents=documents, embeddings=embeddings)
-
-        logging.info("Finished loading PDFs")
+        """Pdf loader."""
+        self.pinecone.pdf_loader(filepath=filepath)
 
     def rag(self, human_message: Union[str, HumanMessage]):
         """
@@ -241,10 +183,7 @@ def rag(self, human_message: Union[str, HumanMessage]):
         # ---------------------------------------------------------------------
         # 1.) Retrieve relevant documents from Pinecone vector database
         # ---------------------------------------------------------------------
-        retriever = PineconeHybridSearchRetriever(
-            embeddings=self.openai_embeddings, sparse_encoder=self.bm25_encoder, index=self.pinecone_index
-        )
-        documents = retriever.get_relevant_documents(query=human_message.content)
+        documents = self.retriever.get_relevant_documents(query=human_message.content)
 
         # Extract the text from the documents
         document_texts = [doc.page_content for doc in documents]

models/pinecone.py

@@ -0,0 +1,167 @@
+# -*- coding: utf-8 -*-
+"""Pinecone helper functions."""
+
+# document loading
+import glob
+
+# general purpose imports
+import logging
+import os
+
+# pinecone integration
+import pinecone
+from langchain.document_loaders import PyPDFLoader
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.text_splitter import Document
+from langchain.vectorstores.pinecone import Pinecone as LCPinecone
+
+# this project
+from models.const import Config, Credentials
+
+
+# pylint: disable=too-few-public-methods
+class TextSplitter:
+    """
+    Custom text splitter that adds metadata to the Document object
+    which is required by PineconeHybridSearchRetriever.
+    """
+
+    def create_documents(self, texts):
+        """Create documents"""
+        documents = []
+        for text in texts:
+            # Create a Document object with the text and metadata
+            document = Document(page_content=text, metadata={"context": text})
+            documents.append(document)
+        return documents
+
+
+class PineConeIndex:
+    """Pinecone helper class."""
+
+    _index: pinecone.Index = None
+    _index_name: str = None
+    _text_splitter: TextSplitter = None
+    _openai_embeddings: OpenAIEmbeddings = None
+    _vector_store: LCPinecone = None
+
+    def __init__(self, index_name: str = None):
+        self._index_name = index_name or Config.PINECONE_INDEX_NAME
+        self.init()
+
+    @property
+    def vector_store(self) -> LCPinecone:
+        """Pinecone lazy read-only property."""
+        if self._vector_store is None:
+            self._vector_store = LCPinecone(
+                index=self.index,
+                embedding=self.openai_embeddings,
+                text_key=Config.PINECONE_VECTORSTORE_TEXT_KEY,
+            )
+        return self._vector_store
+
+    @property
+    def openai_embeddings(self) -> OpenAIEmbeddings:
+        """OpenAIEmbeddings lazy read-only property."""
+        if self._openai_embeddings is None:
+            self._openai_embeddings = OpenAIEmbeddings(
+                api_key=Credentials.OPENAI_API_KEY, organization=Credentials.OPENAI_API_ORGANIZATION
+            )
+        return self._openai_embeddings
+
+    @property
+    def text_splitter(self) -> TextSplitter:
+        """TextSplitter lazy read-only property."""
+        if self._text_splitter is None:
+            self._text_splitter = TextSplitter()
+        return self._text_splitter
+
+    @property
+    def index_name(self) -> str:
+        """index name."""
+        return self._index_name
+
+    @index_name.setter
+    def index_name(self, value: str) -> None:
+        """Set index name."""
+        if self._index_name != value:
+            self._index_name = value
+            self.initialize()
+
+    @property
+    def index(self) -> pinecone.Index:
+        """pinecone.Index lazy read-only property."""
+        if self._index is None:
+            try:
+                self._index = pinecone.Index(index_name=self.index_name)
+            except pinecone.exceptions.PineconeException:
+                # index does not exist, so create it.
+                self.create()
+                self._index = pinecone.Index(index_name=self.index_name)
+        return self._index
+
+    def init(self):
+        """Initialize Pinecone."""
+        pinecone.init(api_key=Credentials.PINECONE_API_KEY, environment=Config.PINECONE_ENVIRONMENT)
+
+    def delete(self):
+        """Delete index."""
+        try:
+            logging.info("Deleting index...")
+            pinecone.delete_index(self.index_name)
+        except pinecone.exceptions.PineconeException:
+            logging.info("Index does not exist. Continuing...")
+
+    def create(self):
+        """Create index."""
+        metadata_config = {
+            "indexed": [Config.PINECONE_VECTORSTORE_TEXT_KEY, "lc_type"],
+            "context": ["lc_text"],
+        }
+        logging.info("Creating index. This may take a few minutes...")
+
+        pinecone.create_index(
+            self.index_name,
+            dimension=Config.PINECONE_DIMENSIONS,
+            metric=Config.PINECONE_METRIC,
+            metadata_config=metadata_config,
+        )
+
+    def initialize(self):
+        """Initialize index."""
+        self.delete()
+        self.create()
+
+    def pdf_loader(self, filepath: str):
+        """
+        Embed PDF.
+        1. Load PDF document text data
+        2. Split into pages
+        3. Embed each page
+        4. Store in Pinecone
+
+        Note: it's important to make sure that the "context" field that holds the document text
+        in the metadata is not indexed. Currently you need to specify explicitly the fields you
+        do want to index. For more information checkout
+        https://docs.pinecone.io/docs/manage-indexes#selective-metadata-indexing
+        """
+        self.initialize()
+
+        pdf_files = glob.glob(os.path.join(filepath, "*.pdf"))
+        i = 0
+        for pdf_file in pdf_files:
+            i += 1
+            j = len(pdf_files)
+            logging.info("Loading PDF %s of %s: %s", i, j, pdf_file)
+            loader = PyPDFLoader(file_path=pdf_file)
+            docs = loader.load()
+            k = 0
+            for doc in docs:
+                k += 1
+                logging.info(k * "-", end="\r")
+                documents = self.text_splitter.create_documents([doc.page_content])
+                document_texts = [doc.page_content for doc in documents]
+                embeddings = self.openai_embeddings.embed_documents(document_texts)
+                self.vector_store.add_documents(documents=documents, embeddings=embeddings)
+
+        logging.info("Finished loading PDFs")

models/tests/test_hsr.py

@@ -6,9 +6,9 @@
 import pytest  # pylint: disable=unused-import
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings
-from pinecone import Index
 
 from models.hybrid_search_retreiver import HybridSearchRetriever, TextSplitter
+from models.pinecone import PineConeIndex
 
 
 class TestSalesSupportModel:
@@ -28,6 +28,6 @@ def test_02_class_aatribute_types(self):
 
         hsr = HybridSearchRetriever()
         assert isinstance(hsr.chat, ChatOpenAI)
-        assert isinstance(hsr.pinecone_index, Index)
+        assert isinstance(hsr.pinecone, PineConeIndex)
         assert isinstance(hsr.text_splitter, TextSplitter)
         assert isinstance(hsr.openai_embeddings, OpenAIEmbeddings)