Update app.py

app.py

@@ -13,12 +13,19 @@ from langchain_community.vectorstores import Qdrant
 from langchain_core.runnables import RunnablePassthrough, RunnableParallel
 import chainlit as cl
 from pathlib import Path
-from sentence_transformers import SentenceTransformer  # Ensure this import is correct
+from sentence_transformers import SentenceTransformer
 
 load_dotenv()
 
 os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
 
+class SentenceTransformerEmbedding:
+    def __init__(self, model_name):
+        self.model = SentenceTransformer(model_name)
+
+    def embed_documents(self, texts):
+        return self.model.encode(texts, convert_to_tensor=True).tolist()  # Convert to list for compatibility
+
 @cl.on_chat_start
 async def on_chat_start():
     model = ChatOpenAI(streaming=True)
@@ -53,10 +60,10 @@ async def on_chat_start():
 
     sentence_combined_documents = sentence_framework + sentence_blueprint
 
-    # Initialize the SentenceTransformer model properly
-    embedding_model = SentenceTransformer('Cheselle/finetuned-arctic-sentence')
+    # Initialize the custom embedding class
+    embedding_model = SentenceTransformerEmbedding('Cheselle/finetuned-arctic-sentence')
 
-    # Create the Qdrant vector store using the initialized embedding model
+    # Create the Qdrant vector store using the custom embedding model
     sentence_vectorstore = Qdrant.from_documents(
         documents=sentence_combined_documents,
         embedding=embedding_model,  # Ensure this is an instance
@@ -65,7 +72,11 @@ async def on_chat_start():
     )
 
     sentence_retriever = sentence_vectorstore.as_retriever()
-    
+
+    # Check if retriever is initialized correctly
+    if sentence_retriever is None:
+        raise ValueError("Retriever is not initialized correctly.")
+
     # Set the retriever and prompt into session for reuse
     cl.user_session.set("runnable", model)
     cl.user_session.set("retriever", sentence_retriever)
@@ -83,6 +94,11 @@ async def on_message(message: cl.Message):
     print(f"Received message: {message.content}")
 
     # Retrieve relevant context from documents based on the user's message
+    if retriever is None:
+        print("Retriever is not available.")
+        await cl.Message(content="Sorry, the retriever is not initialized.").send()
+        return
+
     relevant_docs = retriever.get_relevant_documents(message.content)
     print(f"Retrieved {len(relevant_docs)} documents.")