Update rag_retriever.py

rag_retriever.py

@@ -1,67 +1,64 @@
-
-from dotenv import load_dotenv
-import os
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.llms.huggingface import HuggingFaceInferenceAPI
-from llama_index.core.settings import Settings
-from llama_index.core import SimpleDirectoryReader
-from llama_index.core.node_parser import SentenceWindowNodeParser
-from pinecone import Pinecone
-from llama_index.core import VectorStoreIndex, StorageContext
-from llama_index.vector_stores.pinecone import PineconeVectorStore
-from llama_index.core.postprocessor import MetadataReplacementPostProcessor
-from llama_index.core.postprocessor import SentenceTransformerRerank
-
-
-def initialize_llm():
-    load_dotenv()
-    HF_TOKEN = os.getenv("HUGGINGFACE_API_KEY")
-    
-    Settings.llm = HuggingFaceInferenceAPI(model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", model_kwargs={"temperature": 0.1}, token=HF_TOKEN)
-    Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
-
-
-
-
-def initialize_pinecone():
-    load_dotenv()
-    api_key = os.environ["PINECONE_API_KEY"]
-    index_name = "chatbot"
-    pc = Pinecone(api_key=api_key)
-    pinecone_index = pc.Index(index_name)
-    vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
-    index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
-    return index
-
-def create_query_engine(index):
-    postproc = MetadataReplacementPostProcessor(target_metadata_key="window")
-    rerank = SentenceTransformerRerank(top_n=2, model="BAAI/bge-reranker-base")
-    query_engine = index.as_query_engine(similarity_top_k = 6,vector_store_query_mode="hybrid",alpha=0.5,node_postprocessors = [postproc, rerank],
-)
-    return query_engine
-
-def generate_prompt(context, query):
-    prompt_template = """
-    You are a highly informed representative of KYC SUD CONSULTING. 
-    Your role is to assist users by answering questions about our company with accurate and helpful information. 
-    Use the provided context to guide your responses effectively. 
-    If a question falls outside your knowledge or the context provided, simply state that you don't have the information.
-
-    Given the following context:
-    {context_str}
-
-    Answer the following question in different words while keeping the meaning intact:
-    {query_str}
-
-    Response:
-    """
-    return prompt_template.format(context_str=context, query_str=query)
-
-
-
-# Function to process input sentence and return response
-def get_response(query_engine, sentence):
-    retrieved_context = query_engine.query(sentence)
-    prompt = generate_prompt(retrieved_context, sentence)
-    response = Settings.llm.complete(prompt)  # Use the correct method to generate response
+
+from dotenv import load_dotenv
+import os
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.llms.huggingface import HuggingFaceInferenceAPI
+from llama_index.core.settings import Settings
+from llama_index.core import SimpleDirectoryReader
+from llama_index.core.node_parser import SentenceWindowNodeParser
+from pinecone import Pinecone
+from llama_index.core import VectorStoreIndex, StorageContext
+from llama_index.vector_stores.pinecone import PineconeVectorStore
+from llama_index.core.postprocessor import MetadataReplacementPostProcessor
+from llama_index.core.postprocessor import SentenceTransformerRerank
+
+
+def initialize_llm():
+    load_dotenv()
+    HF_TOKEN = os.getenv("HUGGINGFACE_API_KEY")
+    
+    Settings.llm = HuggingFaceInferenceAPI(model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", model_kwargs={"temperature": 0.1}, token=HF_TOKEN)
+    Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
+
+
+
+
+def initialize_pinecone():
+    load_dotenv()
+    api_key = os.environ["PINECONE_API_KEY"]
+    index_name = "chatbot"
+    pc = Pinecone(api_key=api_key)
+    pinecone_index = pc.Index(index_name)
+    vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
+    index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
+    return index
+
+def create_query_engine(index):
+    postproc = MetadataReplacementPostProcessor(target_metadata_key="window")
+    rerank = SentenceTransformerRerank(top_n=2, model="BAAI/bge-reranker-base")
+    query_engine = index.as_query_engine(similarity_top_k = 6,vector_store_query_mode="hybrid",alpha=0.5,node_postprocessors = [postproc, rerank],
+)
+    return query_engine
+
+def generate_prompt(context, query):
+    prompt_template = """
+    You are a highly informed representative of KYC SUD CONSULTING. Your role is to assist users by answering questions about our company with accurate and helpful information. Use the provided context to guide your responses effectively. If a question falls outside your knowledge or the context provided, simply state that you don't have the information.
+
+    Given the following context:
+    {context_str}
+
+    Answer the following question in different words while keeping the meaning intact:
+    {query_str}
+
+    Response:
+    """
+    return prompt_template.format(context_str=context, query_str=query)
+
+
+
+# Function to process input sentence and return response
+def get_response(query_engine, sentence):
+    retrieved_context = query_engine.query(sentence)
+    prompt = generate_prompt(retrieved_context, sentence)
+    response = Settings.llm.complete(prompt)  # Use the correct method to generate response
     return response.text