Spaces:
Paused
Paused
| import os | |
| import tempfile | |
| import logging | |
| from typing import List, Dict | |
| from langchain.embeddings import OpenAIEmbeddings | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain.vectorstores import Chroma | |
| from langchain.chains import ConversationalRetrievalChain | |
| from langchain.document_loaders import PyPDFLoader | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate | |
| from langchain.schema import Document | |
| from langchain.memory import ConversationBufferMemory | |
| import chainlit as cl | |
| from chainlit.types import AskFileResponse | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Set Chainlit config directory | |
| os.environ["CHAINLIT_CONFIG_DIR"] = tempfile.gettempdir() | |
| # App configuration | |
| APP_CONFIG = { | |
| "model_name": "gpt-4o", | |
| "temperature": 0, | |
| "chunk_size": 2000, # Increased for book-length content | |
| "chunk_overlap": 200, | |
| "max_file_size_mb": 100, # Increased for larger books | |
| "max_output_tokens": 4096, | |
| "embedding_model": "text-embedding-3-large" | |
| } | |
| # Initialize text splitter with more sophisticated settings | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=APP_CONFIG["chunk_size"], | |
| chunk_overlap=APP_CONFIG["chunk_overlap"], | |
| length_function=len, | |
| separators=["\n\n", "\n", " ", ""] # Prioritize splitting on paragraph breaks | |
| ) | |
| # System template | |
| system_template = """You are GPT-4o, an advanced AI model with multimodal capabilities and high intelligence. Use the following pieces of context to answer the user's question. | |
| If you don't know the answer, just say that you don't know, don't try to make up an answer. | |
| ALWAYS return a "SOURCES" part in your answer. | |
| The "SOURCES" part should be a reference to the source of the document from which you got your answer. | |
| If the user greets with phrases like Hi, hello, How are you, etc., reply accordingly as well. | |
| Example of your response should be: | |
| The answer is foo | |
| SOURCES: xyz | |
| Begin! | |
| ---------------- | |
| {context}""" | |
| # Create prompt template | |
| messages = [ | |
| SystemMessagePromptTemplate.from_template(system_template), | |
| HumanMessagePromptTemplate.from_template("{question}") | |
| ] | |
| prompt = ChatPromptTemplate.from_messages(messages) | |
| def process_file(file: AskFileResponse) -> List[Dict[str, str]]: | |
| """Process uploaded PDF file and return list of text chunks with metadata.""" | |
| with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=".pdf") as temp_file: | |
| temp_file.write(file.content) | |
| temp_file_path = temp_file.name | |
| loader = PyPDFLoader(temp_file_path) | |
| pages = loader.load_and_split() | |
| chunks_with_metadata = [] | |
| for i, page in enumerate(pages): | |
| chunks = text_splitter.split_text(page.page_content) | |
| for j, chunk in enumerate(chunks): | |
| chunks_with_metadata.append({ | |
| "content": chunk, | |
| "metadata": { | |
| "page": i + 1, | |
| "chunk": j + 1, | |
| "source": f"page_{i+1}_chunk_{j+1}" | |
| } | |
| }) | |
| return chunks_with_metadata | |
| async def start(): | |
| await cl.Message(f"Welcome! Chat session started using {APP_CONFIG['model_name']} for text generation and {APP_CONFIG['embedding_model']} for embeddings. This combination is optimized for processing and querying book-length content. You can now upload a PDF of a book to begin!").send() | |
| files = await cl.AskFileMessage( | |
| content="Please upload a PDF file of a book to begin!", | |
| accept=["application/pdf"], | |
| max_size_mb=APP_CONFIG["max_file_size_mb"], | |
| timeout=3600, | |
| ).send() | |
| if not files: | |
| await cl.Message("No file uploaded. Please try again.").send() | |
| return | |
| file = files[0] | |
| await cl.Message(f"Starting to process `{file.name}`. This may take a few minutes...").send() | |
| try: | |
| # Process the file | |
| await cl.Message("Reading and splitting the document...").send() | |
| chunks_with_metadata = process_file(file) | |
| await cl.Message("Creating embeddings and setting up the retrieval system...").send() | |
| embeddings = OpenAIEmbeddings(model=APP_CONFIG["embedding_model"]) | |
| docsearch = await cl.make_async(Chroma.from_documents)( | |
| [Document(page_content=chunk["content"], metadata=chunk["metadata"]) for chunk in chunks_with_metadata], | |
| embeddings | |
| ) | |
| memory = ConversationBufferMemory( | |
| memory_key="chat_history", | |
| output_key="answer", | |
| return_messages=True, | |
| ) | |
| await cl.Message("Finalizing the chat model...").send() | |
| chain = ConversationalRetrievalChain.from_llm( | |
| ChatOpenAI( | |
| model_name=APP_CONFIG["model_name"], | |
| temperature=APP_CONFIG["temperature"], | |
| max_tokens=APP_CONFIG["max_output_tokens"], | |
| streaming=True | |
| ), | |
| chain_type="stuff", | |
| retriever=docsearch.as_retriever(search_kwargs={"k": 3}), | |
| memory=memory, | |
| return_source_documents=True, | |
| combine_docs_chain_kwargs={"prompt": prompt} | |
| ) | |
| cl.user_session.set("chain", chain) | |
| await cl.Message(content=f"Processing `{file.name}` complete. You can now ask questions about the book!").send() | |
| except Exception as e: | |
| error_msg = f"An error occurred while processing the file: {str(e)}" | |
| logger.error(error_msg) | |
| await cl.Message(content=error_msg).send() | |
| async def main(message: cl.Message): | |
| """Handle user messages and generate responses.""" | |
| try: | |
| chain = cl.user_session.get("chain") | |
| cb = cl.AsyncLangchainCallbackHandler() | |
| res = await chain.acall(message.content, callbacks=[cb]) | |
| answer = res["answer"] | |
| source_documents = res["source_documents"] | |
| text_elements = [] | |
| for source_doc in source_documents: | |
| source_name = f"Page {source_doc.metadata['page']}, Chunk {source_doc.metadata['chunk']}" | |
| text_elements.append(cl.Text(content=source_doc.page_content, name=source_name)) | |
| source_names = [text_el.name for text_el in text_elements] | |
| if source_names: | |
| answer += f"\nSources: {', '.join(source_names)}" | |
| else: | |
| answer += "\nNo sources found" | |
| await cl.Message(content=answer, elements=text_elements).send() | |
| except Exception as e: | |
| error_msg = f"An error occurred: {str(e)}" | |
| logger.error(error_msg) | |
| await cl.Message(content=error_msg).send() | |
| if __name__ == "__main__": | |
| logger.info(f"Starting Chainlit application with model: {APP_CONFIG['model_name']} and embedding model: {APP_CONFIG['embedding_model']}") | |
| cl.run() |