feat: perfect load(). revert rag() to openai only calls

Author: lpm0073

README.md

@@ -35,8 +35,11 @@ python3 -m models.examples.training_services "Microsoft certified Azure AI engin
 # example 4 - prompted assistant
 python3 -m models.examples.training_services_oracle "Oracle database administrator"
 
-# example 5 - RAG
-python3 -m models.examples.rag "./data/" "What is Accounting Based Valuation?"
+# example 5 - Load PDF documents
+python3 -m models.examples.load "./data/"
+
+# example 6 - Retrieval Augmented Generation
+python3 -m models.examples.rag "What is Accounting Based Valuation?"
 ```
 
 ## Requirements

models/examples/load.py

@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+"""Sales Support Model (SSM) Retrieval Augmented Generation (RAG)"""
+import argparse
+
+from ..ssm import SalesSupportModel
+
+
+ssm = SalesSupportModel()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="RAG example")
+    parser.add_argument("filepath", type=str, help="Location of PDF documents")
+    args = parser.parse_args()
+
+    ssm.load(filepath=args.filepath)

models/examples/rag.py

@@ -9,9 +9,8 @@
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="RAG example")
-    parser.add_argument("filepath", type=str, help="Location of PDF documents")
     parser.add_argument("prompt", type=str, help="A question about the PDF contents")
     args = parser.parse_args()
 
-    result = ssm.rag(filepath=args.filepath, prompt=args.prompt)
+    result = ssm.rag(prompt=args.prompt)
     print(result)

models/ssm.py

@@ -8,10 +8,9 @@
 
 import glob
 import os
-from typing import ClassVar, List
+from typing import List  # ClassVar
 
 import pinecone
-from langchain import hub
 from langchain.cache import InMemoryCache
 
 # prompting and chat
@@ -27,16 +26,17 @@
 from langchain.globals import set_llm_cache
 from langchain.llms.openai import OpenAI
 from langchain.prompts import PromptTemplate
-from langchain.schema import HumanMessage, StrOutputParser, SystemMessage
-from langchain.schema.runnable import RunnablePassthrough
+from langchain.schema import HumanMessage, SystemMessage
 from langchain.text_splitter import Document, RecursiveCharacterTextSplitter
 from langchain.vectorstores.pinecone import Pinecone
-from pydantic import BaseModel, ConfigDict, Field  # ValidationError
 
 # this project
 from models.const import Credentials
 
 
+# from pydantic import BaseModel, ConfigDict, Field
+
+
 ###############################################################################
 # initializations
 ###############################################################################
@@ -45,28 +45,23 @@
 set_llm_cache(InMemoryCache())
 
 
-class SalesSupportModel(BaseModel):
+class SalesSupportModel:
     """Sales Support Model (SSM)."""
 
-    Config: ClassVar = ConfigDict(arbitrary_types_allowed=True)
-
     # prompting wrapper
-    chat: ChatOpenAI = Field(
-        default_factory=lambda: ChatOpenAI(
-            api_key=Credentials.OPENAI_API_KEY,
-            organization=Credentials.OPENAI_API_ORGANIZATION,
-            cache=True,
-            max_retries=3,
-            model="gpt-3.5-turbo",
-            temperature=0.0,
-        )
+    chat = ChatOpenAI(
+        api_key=Credentials.OPENAI_API_KEY,
+        organization=Credentials.OPENAI_API_ORGANIZATION,
+        cache=True,
+        max_retries=3,
+        model="gpt-3.5-turbo",
+        temperature=0.0,
     )
 
     # embeddings
-    texts_splitter_results: List[Document] = Field(None, description="Text splitter results")
-    pinecone_search: Pinecone = Field(None, description="Pinecone search")
-    openai_embedding: OpenAIEmbeddings = Field(OpenAIEmbeddings())
-    query_result: List[float] = Field(None, description="Vector database query result")
+    texts_splitter_results: List[Document]
+    openai_embedding = OpenAIEmbeddings()
+    query_result: List[float]
 
     def cached_chat_request(self, system_message: str, human_message: str) -> SystemMessage:
         """Cached chat request."""
@@ -103,13 +98,13 @@ def embed(self, text: str) -> List[float]:
         # pylint: disable=no-member
         self.openai_embedding.embed_query(embedding)
 
-        self.pinecone_search = Pinecone.from_documents(
-            texts_splitter_results,
+        Pinecone.from_documents(
+            documents=texts_splitter_results,
             embedding=self.openai_embedding,
             index_name=Credentials.PINECONE_INDEX_NAME,
         )
 
-    def rag(self, filepath: str, prompt: str):
+    def load(self, filepath: str):
         """
         Embed PDF.
         1. Load PDF document text data
@@ -118,39 +113,52 @@ def rag(self, filepath: str, prompt: str):
         4. Store in Pinecone
         """
 
-        # pylint: disable=unused-variable
-        def format_docs(docs):
-            """Format docs."""
-            return "\n\n".join(doc.page_content for doc in docs)
-
-        for pdf_file in glob.glob(os.path.join(filepath, "*.pdf")):
+        pdf_files = glob.glob(os.path.join(filepath, "*.pdf"))
+        i = 0
+        for pdf_file in pdf_files:
+            i += 1
+            j = len(pdf_files)
+            print(f"Loading PDF {i} of {j}: ")
             loader = PyPDFLoader(file_path=pdf_file)
             docs = loader.load()
+            k = 0
             for doc in docs:
+                k += 1
+                print(k * "-", end="\r")
                 self.embed(doc.page_content)
+        print("Finished loading PDFs")
 
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        splits = text_splitter.split_documents(docs)
-        vectorstore = Pinecone.from_documents(documents=splits, embedding=self.openai_embedding)
-        retriever = vectorstore.as_retriever()
-        prompt = hub.pull("rlm/rag-prompt")
-
-        rag_chain = (
-            {"context": retriever | self.format_docs, "question": RunnablePassthrough()}
-            | prompt
-            | self.chat
-            | StrOutputParser()
-        )
-
-        return rag_chain.invoke(prompt)
-
-    def embedded_prompt(self, prompt: str) -> List[Document]:
+    def rag(self, prompt: str):
         """
         Embedded prompt.
         1. Retrieve prompt: Given a user input, relevant splits are retrieved
            from storage using a Retriever.
         2. Generate: A ChatModel / LLM produces an answer using a prompt that includes
            the question and the retrieved data
         """
-        result = self.pinecone_search.similarity_search(prompt)
-        return result
+
+        # pylint: disable=unused-variable
+        def format_docs(docs):
+            """Format docs."""
+            return "\n\n".join(doc.page_content for doc in docs)
+
+        pinecone_search = Pinecone.from_existing_index(
+            Credentials.PINECONE_INDEX_NAME,
+            embedding=self.openai_embedding,
+        )
+        retriever = pinecone_search.as_retriever()
+
+        # Use the retriever to get relevant documents
+        documents = retriever.get_relevant_documents(query=prompt)
+        print(f"Retrieved {len(documents)} related documents from Pinecone")
+
+        # Generate a prompt from the retrieved documents
+        prompt += " ".join(doc.page_content for doc in documents)
+        print(f"Prompt contains {len(prompt.split())} words")
+        print("Prompt:", prompt)
+        print(doc for doc in documents)
+
+        # Get a response from the GPT-3.5-turbo model
+        response = self.cached_chat_request(system_message="You are a helpful assistant.", human_message=prompt)
+
+        return response

requirements.txt

@@ -19,10 +19,11 @@ codespell==2.2.6
 
 # production
 # ------------
-python-dotenv==1.0.0
-pydantic==2.5.2
 langchain==0.0.343
+langchainhub==0.1.14
 openai==1.3.5
 pinecone-client==2.2.4
+pydantic==2.5.2
 pypdf==3.17.1
+python-dotenv==1.0.0
 tiktoken==0.5.1