Update app.py

app.py

@@ -1,131 +1,113 @@
-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 typing import cast
+from dotenv import load_dotenv
+import os
 from langchain_community.document_loaders import PyMuPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_openai.embeddings import OpenAIEmbeddings
 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 pathlib import Path
+from sentence_transformers import SentenceTransformer  # Ensure this import is correct
 
 load_dotenv()
 
+os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
 
-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()
+@cl.on_chat_start
+async def on_chat_start():
+    model = ChatOpenAI(streaming=True)
 
+    # 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()
+    
+    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)
 
-def metadata_generator(document, name):
-    fixed_text_splitter = RecursiveCharacterTextSplitter(
+    sentence_text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=500,
         chunk_overlap=100,
         separators=["\n\n", "\n", ".", "!", "?"]
     )
-    collection = fixed_text_splitter.split_documents(document)
-    for doc in collection:
-        doc.metadata["source"] = name
-    return collection
-
-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
-
-
-#from transformers import AutoTokenizer, AutoModel
-#import torch
-#embeddings = AutoModel.from_pretrained("Cheselle/finetuned-arctic-sentence")
-#tokenizer = AutoTokenizer.from_pretrained("Cheselle/finetuned-arctic-sentence")
 
-# Assuming ai_framework_document and ai_blueprint_document are lists of langchain_core.documents.Document
-ai_framework_text = "".join([doc.page_content for doc in ai_framework_document])
+    def metadata_generator(document, name, splitter):
+        collection = splitter.split_documents(document)
+        for doc in collection:
+            doc.metadata["source"] = name
+        return collection
 
-# Similarly for ai_blueprint_document
-ai_blueprint_text = "".join([doc.page_content for doc in ai_blueprint_document])
-
-# Now you can use these text variables
+    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
 
+    # Initialize the SentenceTransformer model properly
+    embedding_model = SentenceTransformer('Cheselle/finetuned-arctic-sentence')
 
+    # Create the Qdrant vector store using the initialized embedding model
+    sentence_vectorstore = Qdrant.from_documents(
+        documents=sentence_combined_documents,
+        embedding=embedding_model,  # Ensure this is an instance
+        location=":memory:",
+        collection_name="AI Policy"
+    )
 
-from sentence_transformers import SentenceTransformer
-embedding_model = SentenceTransformer("Cheselle/finetuned-arctic-sentence")
-embeddings = embedding_model.encode(ai_framework_text + ai_blueprint_text)
-#embeddings = embedding_model.encode(ai_framework_text + ai_blueprint_text)
-#embeddings = embedding_model.encode(ai_framework_document + ai_blueprint_document)
+    sentence_retriever = sentence_vectorstore.as_retriever()
+    
+    # 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)
 
-vectorstore = Qdrant.from_documents(
-    documents=combined_documents,
-    embedding=lambda docs: embedding_model.encode([doc.page_content for doc in docs]), 
-    #embedding=embedding_model,
-    #embedding=embeddings,
-    location=":memory:",
-    collection_name="ai_policy"
-)
 
-retriever = vectorstore.as_retriever()
+@cl.on_message
+async def on_message(message: cl.Message):
+    # Get the stored model, retriever, and prompt
+    model = cast(ChatOpenAI, cl.user_session.get("runnable"))
+    retriever = cl.user_session.get("retriever")
+    prompt_template = cl.user_session.get("prompt_template")
 
-## Generation LLM
-llm = ChatOpenAI(model="gpt-4o-mini")
+    # Log the message content
+    print(f"Received message: {message.content}")
 
-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. 
+    # Retrieve relevant context from documents based on the user's message
+    relevant_docs = retriever.get_relevant_documents(message.content)
+    print(f"Retrieved {len(relevant_docs)} documents.")
 
-Context: {context}
-Question: {question}
-"""
+    if not relevant_docs:
+        print("No relevant documents found.")
+        await cl.Message(content="Sorry, I couldn't find any relevant documents.").send()
+        return
 
-rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
+    context = "\n\n".join([doc.page_content for doc in relevant_docs])
 
-retrieval_augmented_qa_chain = (
-    # INVOKE CHAIN WITH: {"question" : "<<SOME USER QUESTION>>"}
-    # "question" : populated by getting the value of the "question" key
-    # "context"  : populated by getting the value of the "question" key and chaining it into the base_retriever
-    {"context": itemgetter("question") | retriever, "question": itemgetter("question")}
-    # "context"  : is assigned to a RunnablePassthrough object (will not be called or considered in the next step)
-    #              by getting the value of the "context" key from the previous step
-    | RunnablePassthrough.assign(context=itemgetter("context"))
-    # "response" : the "context" and "question" values are used to format our prompt object and then piped
-    #              into the LLM and stored in a key called "response"
-    # "context"  : populated by getting the value of the "context" key from the previous step
-    | {"response": rag_prompt | llm, "context": itemgetter("context")}
-)
+    # Log the context to check
+    print(f"Context: {context}")
 
-#alt_rag_chain.invoke({"question" : "What is the AI framework all about?"})
+    # Construct the final RAG prompt
+    final_prompt = prompt_template.format(context=context, question=message.content)
+    print(f"Final prompt: {final_prompt}")
 
-@cl.on_message
-async def handle_message(message):
-    try:
-        # Process the incoming question using the RAG chain
-        result = retrieval_augmented_qa_chain.invoke({"question": message.content})
+    # Initialize a streaming message
+    msg = cl.Message(content="")
 
-        # Create a new message for the response
-        response_message = cl.Message(content=result["response"].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)
 
-        # 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}")
+    await msg.send()