Update agent.py

agent.py

@@ -24,11 +24,9 @@ import numpy as np
 
 import pandas as pd
 import uuid
-from langchain_community.vectorstores import FAISS
-from langchain.schema import Document
 import requests
 import json
-from langchain.schema import Document
+from langchain_core.documents import Document
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
 
@@ -147,11 +145,6 @@ sys_msg = SystemMessage(content=system_prompt)
 json_url = "https://agents-course-unit4-scoring.hf.space/questions"  # Replace with your actual JSON file URL
 response = requests.get(json_url)
 
-# Ensure the request was successful
-if response.status_code != 200:
-    raise Exception(f"Failed to load JSON from {json_url}. Status code: {response.status_code}")
-
-# Ensure the request was successful
 if response.status_code != 200:
     raise Exception(f"Failed to load JSON from {json_url}. Status code: {response.status_code}")
 
@@ -260,11 +253,18 @@ def build_graph(provider: str = "google"):
     def retriever(state: MessagesState):
         """Retriever node"""
         similar_question = vector_store.similarity_search(state["messages"][0].content)
-        example_msg = HumanMessage(
-            content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
-        )
+    
+        if not similar_question:
+            example_msg = HumanMessage(content="No similar question found.")
+        else:
+            example_msg = HumanMessage(
+                content=f"Here I provide a similar question and answer for reference:\n\n{similar_question[0].page_content}",
+            )
+    
         return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
+
+    
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)