adds retrieval logic

requirements.txt

@@ -11,4 +11,5 @@ google.generativeai==0.4.1
 python-dotenv==1.0.1
 docx2txt==0.8
 python-docx==1.1.0
-markdown==3.6
+markdown==3.6
+langchain-core==0.1.46

src/_google/docindex.py

@@ -7,12 +7,18 @@
 import tiktoken
 from typing import List
 from _openai.doc_model import Page
-from langchain_pinecone import PineconeVectorStore 
 import google.generativeai as genai
 from pathlib import Path
 from langchain_community.document_loaders import UnstructuredWordDocumentLoader
 from langchain_community.document_loaders import UnstructuredMarkdownLoader
 from langchain_community.document_loaders import UnstructuredHTMLLoader
+from langchain_pinecone import PineconeVectorStore
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from operator import itemgetter
+from langchain_google_genai import ChatGoogleGenerativeAI
+from src.config import Config
+
 
 class GooglePineconeIndexer:
     """
@@ -234,3 +240,28 @@ def initialize_vectorstore(self, index_name):
                 )
         vectorstore = PineconeVectorStore(index, embed, "text")
         return vectorstore
+
+
+    def retrieve_and_generate(self,query: str, index_name: str, model_name: str = 'gemini-pro', top_k: int =5):
+        """
+        Retrieve documents from the Pinecone index and generate a response.
+        Args:
+            query: The qury from the user
+            index_name: The name of the Pinecone index
+            model_name: The name of the model to use : defaults to 'gemini-pro'
+            top_k: The number of documents to retrieve from the index : defaults to 5
+        """
+        llm = ChatGoogleGenerativeAI(model = Config.default_google_model, google_api_key=self.google_api_key)
+        rag_prompt = PromptTemplate(template = Config.template_str, input_variables = ["query", "context"])
+        vector_store = self.initialize_vectorstore(index_name)
+        retriever = vector_store.as_retriver(search_kwargs = {"k": top_k})
+        rag_chain = (
+            {"context": itemgetter("query")| retriever,
+             "query": itemgetter("query"),
+             }
+             | rag_prompt
+             | llm
+             | StrOutputParser()
+        )
+
+        return rag_chain.invoke({"query": query})

src/_openai/docindex.py

@@ -7,11 +7,18 @@
 import tiktoken
 from typing import List
 from .doc_model import Page
-from langchain_pinecone import PineconeVectorStore 
 from pathlib import Path
 from langchain_community.document_loaders import UnstructuredWordDocumentLoader
 from langchain_community.document_loaders import UnstructuredMarkdownLoader
 from langchain_community.document_loaders import UnstructuredHTMLLoader
+from langchain_pinecone import PineconeVectorStore
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from operator import itemgetter
+from langchain_openai import ChatOpenAI
+from src.config import Config
+
+
 
 class OpenaiPineconeIndexer:
     """
@@ -231,3 +238,35 @@ def initialize_vectorstore(self, index_name):
                 )
         vectorstore = PineconeVectorStore(index, embed, "text")
         return vectorstore
+
+
+    def retrieve_and_generate(self,query: str, index_name: str, model_name: str = 'gpt-3.5-turbo-1106', top_k: int =5):
+        """
+        Retrieve documents from the Pinecone index and generate a response.
+        Args:
+            query: The query from the user
+            index_name: The name of the Pinecone index
+            model_name: The name of the model to use : defaults to 'gpt-3.5-turbo-1106'
+            top_k: The number of documents to retrieve from the index : defaults to 5
+        """
+        llm = ChatOpenAI(model = Config.default_openai_model, openai_api_key = self.openai_api_key)
+        rag_prompt = PromptTemplate(template = Config.template_str, input_variables = ["query", "context"])
+
+        vector_store = self.initialize_vectorstore(index_name)
+        retriever = vector_store.as_retriver(search_kwargs = {"k": top_k})
+        rag_chain = (
+            {"context": itemgetter("query")| retriever,
+             "query": itemgetter("query"),
+             }
+             | rag_prompt
+             | llm
+             | StrOutputParser()
+        )
+
+        return rag_chain.invoke({"query": query})
+
+
+
+
+
+

src/config.py

@@ -0,0 +1,11 @@
+class Config:
+    template_str = """
+    You are very helpful assistant for question answering tasks. Use the pieces of retrieved context to answer question given. If you do not know 
+    the answer, Just say that you do not know the answer instead of making up an answer.
+
+    Retrieved context: {context}
+    Query: {query}
+    """
+
+    default_google_model = "gemini-pro"
+    default_openai_model = "gpt-3.5-turbo-0125"