Update app.py

app.py

@@ -11,7 +11,7 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from docarray import DocumentArray
+from docarray import Document, DocumentArray
 from sentence_transformers import SentenceTransformer
 
 # StreamHandler to intercept streaming output from the LLM.
@@ -26,6 +26,19 @@ class StreamHandler(BaseCallbackHandler):
         self.text += token
         self.container.markdown(self.text)
 
+from langchain.retrievers import BaseRetriever
+
+class SimpleEmbeddingRetriever(BaseRetriever):
+    def __init__(self, documents):
+        self.documents = documents
+
+    def _get_relevant_documents(self, query: str, num_documents: int = 5):
+        query_doc = Document(text=query)
+        query_embedding = self.documents.embeddings.model.encode([query_doc.text])[0]
+        query_doc.embedding = query_embedding
+        scores = self.documents.match(query_doc, limit=num_documents, metric='cosine', use_scipy=True)
+        return [(doc.text, score) for doc, score in scores]
+
 
 @st.cache_data
 def get_page_urls(url):
@@ -54,11 +67,13 @@ def get_url_content(url):
             text = ' '.join([c.get_text().strip() for c in content if c.get_text().strip() != ''])
 
         # Create a single document with metadata
-        return DocumentArray([{'text': text, 'tags': {'url': url}}])
+        document = Document(text=text, tags={'url': url})
+        return DocumentArray([document])
     except Exception as e:
         st.error(f"Failed to process URL content: {e}")
         return DocumentArray()
 
+
 @st.cache_resource
 def get_retriever(urls):
     documents = DocumentArray()
@@ -72,7 +87,7 @@ def get_retriever(urls):
     for doc, emb in zip(documents, embeddings):
         doc.embedding = emb
     
-    return documents
+    return SimpleEmbeddingRetriever(documents)
 
 
 @st.cache_resource