Update app.py

app.py

@@ -1,138 +1,123 @@
+import re
+from langchain_openai import OpenAIEmbeddings
 from langchain_openai import ChatOpenAI
+from langchain_openai.embeddings import OpenAIEmbeddings
+
 from langchain.prompts import ChatPromptTemplate
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.schema import StrOutputParser
-from langchain.schema.runnable import Runnable
-from langchain.schema.runnable.config import RunnableConfig
-from dotenv import load_dotenv
-import os
+
 from langchain_community.document_loaders import PyMuPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Qdrant
+
+from langchain_core.runnables import RunnablePassthrough, RunnableParallel
+from langchain_core.documents import Document
+
+from operator import itemgetter
+import os
+from dotenv import load_dotenv
 import chainlit as cl
+from langchain.embeddings.base import Embeddings
 from sentence_transformers import SentenceTransformer
 
 # Load environment variables
 load_dotenv()
-os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
 
-class SentenceTransformerEmbedding:
-    def __init__(self, model_name):
+# Custom wrapper for SentenceTransformer to work with Langchain
+class LangchainSentenceTransformerEmbeddings(Embeddings):
+    def __init__(self, model_name: str):
         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
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        # Encode the documents using SentenceTransformer's encode method
+        return self.model.encode(texts)
 
-    def __call__(self, texts):
-        return self.embed_documents(texts)  # Make it callable
+    def embed_query(self, text: str) -> list[float]:
+        # Encode a single query using SentenceTransformer's encode method
+        return self.model.encode([text])[0]
 
-@cl.on_chat_start
-async def on_chat_start():
-    model = ChatOpenAI(streaming=True)
+# Initialize the custom embedding model
+embedding_model = LangchainSentenceTransformerEmbeddings("Cheselle/finetuned-arctic-sentence")
 
-    # Load documents
-    ai_framework_document = PyMuPDFLoader(file_path="https://nvlpubs.nist.gov/nistpubs/ai/NIST.AI.600-1.pdf").load()
-    ai_blueprint_document = PyMuPDFLoader(file_path="https://www.whitehouse.gov/wp-content/uploads/2022/10/Blueprint-for-an-AI-Bill-of-Rights.pdf").load()
+# Load the documents using PyMuPDFLoader
+ai_framework_document = PyMuPDFLoader(file_path="https://nvlpubs.nist.gov/nistpubs/ai/NIST.AI.600-1.pdf").load()
+ai_blueprint_document = PyMuPDFLoader(file_path="https://www.whitehouse.gov/wp-content/uploads/2022/10/Blueprint-for-an-AI-Bill-of-Rights.pdf").load()
 
-    print("Documents loaded.")
-
-    RAG_PROMPT = """\
-    Given a provided context and question, you must answer the question based only on context.
-
-    Context: {context}
-    Question: {question}
-    """
-
-    rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
-
-    sentence_text_splitter = RecursiveCharacterTextSplitter(
+# Metadata generator function to add metadata to documents
+def metadata_generator(document, name):
+    fixed_text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=500,
         chunk_overlap=100,
         separators=["\n\n", "\n", ".", "!", "?"]
     )
-
-    def metadata_generator(document, name, splitter):
-        collection = splitter.split_documents(document)
-        for doc in collection:
-            doc.metadata["source"] = name
-        return collection
-
-    sentence_framework = metadata_generator(ai_framework_document, "AI Framework", sentence_text_splitter)
-    sentence_blueprint = metadata_generator(ai_blueprint_document, "AI Blueprint", sentence_text_splitter)
-
-    sentence_combined_documents = sentence_framework + sentence_blueprint
-
-    print(f"Total documents to embed: {len(sentence_combined_documents)}")
-
-    # Limit the number of documents processed for debugging
-    max_documents = 10
-    sentence_combined_documents = sentence_combined_documents[:max_documents]
-
-    # Initialize the embedding model instance
-    embedding_model = SentenceTransformerEmbedding('Cheselle/finetuned-arctic-sentence')
-
-    # Create the Qdrant vector store using the embedding instance
-    sentence_vectorstore = Qdrant.from_documents(
-        documents=sentence_combined_documents,
-        embedding=embedding_model,
-        location=":memory:",
-        collection_name="AI Policy"
-    )
-
-    print("Vector store created.")
-
-    # Create retriever from the vector store
-    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)
-    cl.user_session.set("prompt_template", rag_prompt)
-
+    collection = fixed_text_splitter.split_documents(document)
+    for doc in collection:
+        doc.metadata["source"] = name
+    return collection
+
+# Generate metadata for the loaded documents
+recursive_framework_document = metadata_generator(ai_framework_document, "AI Framework")
+recursive_blueprint_document = metadata_generator(ai_blueprint_document, "AI Blueprint")
+combined_documents = recursive_framework_document + recursive_blueprint_document
+
+# Combine the content of the documents
+ai_framework_text = "".join([doc.page_content for doc in ai_framework_document])
+ai_blueprint_text = "".join([doc.page_content for doc in ai_blueprint_document])
+
+# Pass the custom embedding model to Qdrant to create a vectorstore
+vectorstore = Qdrant.from_documents(
+    documents=combined_documents,  # List of documents
+    embedding=embedding_model,     # Custom Langchain wrapper for SentenceTransformer
+    location=":memory:",
+    collection_name="ai_policy"
+)
+
+# Set up the retriever
+retriever = vectorstore.as_retriever()
+
+# LLM configuration
+llm = ChatOpenAI(model="gpt-4o-mini")
+
+# Define the RAG (Retrieval-Augmented Generation) prompt template
+RAG_PROMPT = """\
+You are an AI Policy Expert. 
+Given a provided context and question, you must answer the question based only on context. 
+Think through your answer carefully and step by step. 
+Context: {context}
+Question: {question}
+"""
+
+rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
+
+# Define the retrieval-augmented QA chain
+retrieval_augmented_qa_chain = (
+    {"context": itemgetter("question") | retriever, "question": itemgetter("question")}
+    | RunnablePassthrough.assign(context=itemgetter("context"))
+    | {"response": rag_prompt | llm, "context": itemgetter("context")}
+)
+
+# Chainlit event handler for receiving messages
 @cl.on_message
-async def on_message(message: cl.Message):
-    # Get the stored model, retriever, and prompt
-    model = cl.user_session.get("runnable")
-    retriever = cl.user_session.get("retriever")
-    prompt_template = cl.user_session.get("prompt_template")
-
-    # Log the message content
-    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.")
-
-    if not relevant_docs:
-        print("No relevant documents found.")
-        await cl.Message(content="Sorry, I couldn't find any relevant documents.").send()
-        return
-
-    context = "\n\n".join([doc.page_content for doc in relevant_docs])
-
-    # Log the context to check
-    print(f"Context: {context}")
-
-    # Construct the final RAG prompt
-    final_prompt = prompt_template.format(context=context, question=message.content)
-    print(f"Final prompt: {final_prompt}")
-
-    # Initialize a streaming message
-    msg = cl.Message(content="")
-
-    # Stream the response from the model
-    async for chunk in model.astream(
-        final_prompt,
-        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
-    ):
-        await msg.stream_token(chunk.content)
-
-    await msg.send()
+async def handle_message(message):
+    try:
+        # Process the incoming question using the RAG chain
+        result = retrieval_augmented_qa_chain.invoke({"question": message.content})
+
+        # Create a new message for the response
+        response_message = cl.Message(content=result["response"].content)
+
+        # Send the response back to the user
+        await response_message.send()
+    
+    except Exception as e:
+        # Handle any exception and log it or send a response back to the user
+        error_message = cl.Message(content=f"An error occurred: {str(e)}")
+        await error_message.send()
+        print(f"Error occurred: {e}")
+
+# Run the Chainlit server
+if __name__ == "__main__":
+    try:
+        cl.run()
+    except Exception as e:
+        print(f"Server error occurred: {e}")