adds retrieval logic

src/_google/docindex.py

@@ -13,6 +13,12 @@
 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
+
 
 class GooglePineconeIndexer:
     """
@@ -234,3 +240,23 @@ def initialize_vectorstore(self, index_name):
                 )
         vectorstore = PineconeVectorStore(index, embed, "text")
         return vectorstore
+
+
+    def retrieve_and_generate(self, prompt: str ,query: str, top_k: int , index_name: str):
+        """
+        Retrieve documents from the Pinecone index and generate a response.
+        """
+        llm = ChatGoogleGenerativeAI(model = self.model_name, google_api_key=self.google_api_key)
+        rag_prompt = PromptTemplate(template = prompt, 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

@@ -12,6 +12,13 @@
 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
+
+
 
 class OpenaiPineconeIndexer:
     """
@@ -231,3 +238,30 @@ def initialize_vectorstore(self, index_name):
                 )
         vectorstore = PineconeVectorStore(index, embed, "text")
         return vectorstore
+
+
+    def retrieve_and_generate(self, prompt: str ,query: str, top_k: int , index_name: str):
+        """
+        Retrieve documents from the Pinecone index and generate a response.
+        """
+        llm = ChatOpenAI(temperature = 0, model = "gpt-3.5-turbo", openai_api_key = self.openai_api_key)
+        rag_prompt = PromptTemplate(template = prompt, 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})
+
+
+
+
+
+