Updated to include llamaindex nvidia integrations

app.py

@@ -1,7 +1,10 @@
 # Import modules and classes
 from llama_index.core import VectorStoreIndex, StorageContext, load_index_from_storage
-from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings, NVIDIARerank
+from langchain_nvidia_ai_endpoints import NVIDIARerank
 from llama_index.core.indices.query.query_transform import HyDEQueryTransform
+from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.llms.nvidia import NVIDIA
+from llama_index.embeddings.nvidia import NVIDIAEmbedding
 from llama_index.core.embeddings.utils import resolve_embed_model
 from llama_index.core.query_engine import TransformQueryEngine
 from langchain_core.documents import Document as LangDocument
@@ -16,7 +19,7 @@ nvidia_api_key = os.getenv("NVIDIA_KEY")
 llamaparse_api_key = os.getenv("PARSE_KEY")
 
 # Initialize ChatNVIDIA, NVIDIARerank, and NVIDIAEmbeddings
-client = ChatNVIDIA(
+client = NVIDIA(
     model="meta/llama-3.1-8b-instruct",
     api_key=nvidia_api_key,
     temperature=0.2,
@@ -24,16 +27,7 @@ client = ChatNVIDIA(
     max_tokens=1024
 )
 
-# Custom resolve_embed_model function without callback_manager for NVIDIAEmbeddings
-def custom_resolve_embed_model(embed_model):
-    if isinstance(embed_model, NVIDIAEmbeddings):
-        return embed_model
-    embed_model = resolve_embed_model(embed_model)
-    if hasattr(embed_model, 'callback_manager'):
-        embed_model.callback_manager = Settings.callback_manager
-    return embed_model
-
-embed_model = NVIDIAEmbeddings(
+embed_model = NVIDIAEmbedding(
     model="nvidia/nv-embedqa-e5-v5", 
     api_key=nvidia_api_key, 
     truncate="NONE"
@@ -45,7 +39,7 @@ reranker = NVIDIARerank(
 )
 
 # Set the NVIDIA models globally
-Settings.embed_model = custom_resolve_embed_model(embed_model)
+Settings.embed_model = embed_model
 Settings.llm = client
 
 # Parse the local PDF document
@@ -55,7 +49,12 @@ parser = LlamaParse(
     verbose=True
 )
 
-documents = parser.load_data("C:\\Users\\user\\Documents\\Jan 2024\\Projects\\RAGs\\Files\\PhilDataset.pdf")
+# Get the absolute path of the script's directory
+script_dir = os.path.dirname(os.path.abspath(__file__))
+data_file = os.path.join(script_dir, "PhilDataset.pdf")
+
+# Load the PDF document using the relative path
+documents = parser.load_data(data_file)
 print("Document Parsed")
 
 # Split parsed text into chunks for embedding model
@@ -87,7 +86,7 @@ all_documents = []
 for doc in documents:
     text_chunks = split_text(doc.text)
     for chunk in text_chunks:
-        embedding = embed_model.embed_query(chunk)
+        embedding = embed_model.get_text_embedding(chunk)
         all_embeddings.append(embedding)
         all_documents.append(LlamaDocument(text=chunk))
 print("Embeddings generated")
@@ -139,18 +138,32 @@ def query_model_with_context(question):
     context = ranked_documents[0].page_content
 
     # Send context and question to the client (NVIDIA Llama 3.1 8B model)
+    # Construct the messages using the ChatMessage class
     messages = [
-        {"role": "system", "content": context},
-        {"role": "user", "content": str(question)}
+        ChatMessage(role=MessageRole.SYSTEM, content=context),
+        ChatMessage(role=MessageRole.USER, content=str(question))
     ]
-    completion = client.stream(messages)
-    
-    # Process response
+
+    # Call the chat method to get the response
+    completion = client.chat(messages)
+
+    print(completion)
+
+    # Process response - assuming completion is a single string or a tuple containing a string
     response_text = ""
-    for chunk in completion:
-        if chunk.content is not None:
-            response_text += chunk.content
+
+    if isinstance(completion, (list, tuple)):
+        # Join elements of tuple/list if it's in such format
+        response_text = ' '.join(completion)
+    elif isinstance(completion, str):
+        # Directly assign if it's a string
+        response_text = completion
+    else:
+        # Fallback for unexpected types, convert to string
+        response_text = str(completion)
     
+    response_text = response_text.replace("assistant:", "Final Response:").strip()
+
     return response_text