messages update

app.py

@@ -2,15 +2,12 @@ import gradio as gr
 from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
 from dotenv import load_dotenv
 import os
+from custom_retriever import query_engine
 
 load_dotenv()
 
-documents = SimpleDirectoryReader("data").load_data()
-index = VectorStoreIndex.from_documents(documents)
-query_engine = index.as_query_engine()
-
-def update(query):
-    response = query_engine.query(query)
+def get_response(query):
+    response = query_engine.custom_query(query)
     return str(response)
 
 with gr.Blocks() as demo:
@@ -19,6 +16,6 @@ with gr.Blocks() as demo:
         inp = gr.Textbox(placeholder="Please enter your query", label="User Query")
         out = gr.Textbox()
     btn = gr.Button("Run")
-    btn.click(fn=update, inputs=inp, outputs=out)
+    btn.click(fn=get_response, inputs=inp, outputs=out)
 
 demo.launch()

custom_retriever.py

@@ -0,0 +1,98 @@
+from llama_index.vector_stores.pinecone import PineconeVectorStore
+from llama_index.core.schema import NodeWithScore
+from typing import List, Any, Optional
+from llama_index.core import PromptTemplate, get_response_synthesizer, VectorStoreIndex, ChatPromptTemplate, QueryBundle
+from llama_index.core.vector_stores import VectorStoreQuery
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever
+from llama_index.core.query_engine import CustomQueryEngine
+from llama_index.core.response_synthesizers import BaseSynthesizer
+from llama_index.llms.openai.base import OpenAI
+from pinecone import Pinecone, PodSpec
+from dotenv import load_dotenv
+
+load_dotenv()
+
+llm = OpenAI(model="gpt-4-0125-preview", temperature=0.7)
+pc = Pinecone(api_key="cf85ee13-6be8-4b4b-b512-fbad9f4ce781")
+
+embed_model = OpenAIEmbedding(model="text-embedding-3-small")
+synthesizer = get_response_synthesizer(response_mode="compact")
+
+index = pc.Index("test-index")
+
+qa_prompt = PromptTemplate(
+    "Context \n"
+    "---------------------\n"
+    "{context_str}\n"
+    "---------------------\n"
+    "Given the context and without prior knowledge answer the user query.\n"
+    "User Query: {query_str}\n"
+    "Answer: ",
+)
+
+class PineconeRetriever(BaseRetriever):
+    def __init__(self, vector_store: PineconeVectorStore, embed_model: Any, query_mode: str = "default", similarity_top_k: int = 10) -> None:
+        self._vector_store = vector_store
+        self._embed_model = embed_model
+        self._query_mode = query_mode
+        self._similarity_top_k = similarity_top_k
+        super().__init__()
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        print("Yup we are using the custom retriever")
+        query_embedding = embed_model.get_query_embedding(query_bundle)
+        vector_store_query = VectorStoreQuery(
+            query_embedding=query_embedding,
+            similarity_top_k=self._similarity_top_k,
+            mode=self._query_mode,
+        )
+        query_result = self._vector_store.query(vector_store_query)
+        nodes_with_scores = []
+        for index, node in enumerate(query_result.nodes):
+            score: Optional[float] = None
+            if query_result.similarities is not None:
+                score = query_result.similarities[index]
+            nodes_with_scores.append(NodeWithScore(node=node, score=score))
+
+        return nodes_with_scores
+    
+class RAGStringQueryEngine(CustomQueryEngine):
+    """RAG String Query Engine."""
+
+    retriever: PineconeRetriever
+    response_synthesizer: BaseSynthesizer
+    llm: OpenAI
+    qa_prompt: PromptTemplate
+
+    def custom_query(self, query_str: str):
+        nodes = self.retriever._retrieve(query_str)
+        context_str = "\n\n".join([n.node.get_content() for n in nodes])
+        qa_prompt_str = qa_prompt.format(
+            context_str=context_str, 
+            query_str=query_str,
+        )
+        print("QA Prompt STR: ", qa_prompt_str)
+        response = self.llm.complete(qa_prompt_str)
+        return str(response)
+    
+    def print_retrieved_nodes(self, query_str: str):
+        nodes = self.retriever._retrieve(query_str)
+        # print(nodes[0].score)
+        for n in nodes:
+            print("Node Score: ", n.score)
+            print(n.node.get_content())
+
+
+vector_store = PineconeVectorStore(pinecone_index=index)
+vector_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
+# retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=2)
+retriever = PineconeRetriever(vector_store=vector_store, embed_model=embed_model, query_mode="default", similarity_top_k=4)
+
+
+query_engine = RAGStringQueryEngine(
+        retriever=retriever,
+        response_synthesizer=synthesizer,
+        llm=llm,
+        qa_prompt=qa_prompt,
+    )

requirements.txt

@@ -1,3 +1,4 @@
 llama-index==0.10.11
 gradio==4.28.3
-python-dotenv==1.0.1
+python-dotenv==1.0.1
+pinecone-client==3.0.3