diff --git a/rag_tutorials/llm_app_hybrid_RAG_claude/local_main.py b/rag_tutorials/llm_app_hybrid_RAG_claude/local_main.py new file mode 100644 index 0000000..94c93e1 --- /dev/null +++ b/rag_tutorials/llm_app_hybrid_RAG_claude/local_main.py @@ -0,0 +1,279 @@ +import os +import logging +import chainlit as cl +from raglite import RAGLiteConfig, insert_document, hybrid_search, retrieve_chunks, rerank_chunks, rag +from rerankers import Reranker +from typing import List, Tuple +from pathlib import Path +from chainlit.input_widget import TextInput +import time + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +my_config = None + +# System prompt for local LLM +RAG_SYSTEM_PROMPT = """ +You are a friendly and knowledgeable assistant that provides complete and insightful answers. +Answer the user's question using only the context below. +When responding, you MUST NOT reference the existence of the context, directly or indirectly. +Instead, you MUST treat the context as if its contents are entirely part of your working memory. +""".strip() + +def initialize_config(settings: dict) -> RAGLiteConfig: + """Initialize RAGLite configuration with local models.""" + try: + return RAGLiteConfig( + db_url=settings["DBUrl"], + llm=f"llama-cpp-python/{settings['LLMPath']}", + embedder=f"llama-cpp-python/{settings['EmbedderPath']}", + embedder_normalize=True, + chunk_max_size=2000, + embedder_sentence_window_size=2, + reranker=( + ("en", Reranker("ms-marco-MiniLM-L-12-v2", model_type="flashrank")), + ("other", Reranker("ms-marco-MultiBERT-L-12", model_type="flashrank")), + ) + ) + except KeyError as e: + raise ValueError(f"Missing required setting: {e}") + +def process_document(file_path: str) -> None: + """Process and embed a document into the database.""" + logger.info(f"Starting to process document: {file_path}") + try: + import time + start_time = time.time() + + # Inserting document into PostgreSQL database + insert_document(Path(file_path), config=my_config) + + processing_time = time.time() - start_time + logger.info(f"Document processed and embedded in {processing_time:.2f} seconds") + + except Exception as e: + logger.error(f"Error processing document: {str(e)}") + raise + +def perform_search(query: str) -> List[dict]: + """Perform hybrid search and reranking on the query.""" + logger.info(f"Performing hybrid search for: {query}") + try: + chunk_ids, scores = hybrid_search(query, num_results=10, config=my_config) + logger.debug(f"Found {len(chunk_ids)} chunks with scores: {scores}") + + if not chunk_ids: + logger.info("No relevant chunks found in database") + return [] + + chunks = retrieve_chunks(chunk_ids, config=my_config) + reranked_chunks = rerank_chunks(query, chunks, config=my_config) + return reranked_chunks + except Exception as e: + logger.error(f"Search error: {str(e)}") + return [] + +@cl.on_settings_update +async def handle_settings_update(settings: dict): + """Handle settings updates when user submits the form.""" + try: + def validate_path(path: str, path_type: str) -> bool: + if not path: + raise ValueError(f"{path_type} path is required") + return True + + def validate_db_url(url: str) -> bool: + if not url: + raise ValueError("Database URL is required") + if not url.startswith(('sqlite://', 'postgresql://')): + raise ValueError("Invalid database URL format") + return True + + # Validate settings + validate_path(settings["LLMPath"], "LLM") + validate_path(settings["EmbedderPath"], "Embedder") + validate_db_url(settings["DBUrl"]) + + global my_config + my_config = initialize_config(settings) + cl.user_session.set("settings", settings) + + await cl.Message(content="✅ Successfully configured with local models!").send() + + # Ask for file upload + files = await cl.AskFileMessage( + content="Please upload one or more PDF documents to begin!", + accept=["application/pdf"], + max_size_mb=20, + timeout=300, + max_files=5 + ).send() + + if files: + success = False + + # Process uploaded files + for file in files: + logger.info(f"Starting to process file: {file.name}") + + # Create new message for each file + await cl.Message(f"Processing {file.name}...").send() + + try: + logger.info(f"Embedding document: {file.path}") + process_document(file_path=file.path) + + success = True + await cl.Message(f"✅ Successfully processed: {file.name}").send() + logger.info(f"Successfully processed and embedded: {file.name}") + + except Exception as proc_error: + error_msg = f"Failed to process {file.name}: {str(proc_error)}" + logger.error(error_msg) + await cl.Message(f"❌ {error_msg}").send() + continue + + if success: + # Send completion message + await cl.Message( + content="✅ Documents are ready! You can now ask questions about them." + ).send() + + # Store session state + cl.user_session.set("documents_loaded", True) + + # Reset the chat interface + await cl.Message(content="Ask your first question:").send() + + # Clear any existing message elements + cl.user_session.set("message_elements", []) + + else: + await cl.Message( + content="❌ No documents were successfully processed. Please try uploading again." + ).send() + + except Exception as e: + error_msg = f"❌ Error with provided settings: {str(e)}" + logger.error(error_msg) + await cl.Message(content=error_msg).send() + + +@cl.on_chat_start +async def start() -> None: + """Initialize chat and request model paths.""" + try: + logger.info("Chat session started") + cl.user_session.set("chat_history", []) + + await cl.Message( + content="""# 🤖 Local LLM-Powered Hybrid Search-RAG Assistant + + Welcome! To get started: + Enter your model paths and database URL below on the ChatSettings widget + """ + ).send() + + await cl.ChatSettings([ + TextInput( + id="LLMPath", + label="LLM Path", + initial="", + placeholder="e.g., bartowski/Llama-3.2-3B-Instruct-GGUF/*Q4_K_M.gguf@4096" + ), + TextInput( + id="EmbedderPath", + label="Embedder Path", + initial="", + placeholder="e.g., lm-kit/bge-m3-gguf/*F16.gguf@1024" + ), + TextInput( + id="DBUrl", + label="Database URL", + initial="", + placeholder="sqlite:///raglite.sqlite or postgresql://user:pass@host:port/db" + ), + ]).send() + + except Exception as e: + logger.error(f"Error in chat start: {str(e)}") + await cl.Message(content=f"Error initializing chat: {str(e)}").send() + +@cl.on_message +async def message_handler(message: cl.Message) -> None: + """Handle user queries using local RAG.""" + try: + if not cl.user_session.get("documents_loaded"): + await cl.Message(content="❌ Please upload and process documents first!").send() + return + + if not my_config: + await cl.Message(content="❌ Please configure your model paths first!").send() + return + + msg = cl.Message(content="") + await msg.send() + + query = message.content.strip() + if not query: + await cl.Message(content="❌ Please enter a valid question.").send() + return + + logger.info(f"Processing query: {query}") + + try: + reranked_chunks = perform_search(query) + + if not reranked_chunks: + await cl.Message(content="No relevant information found in the documents.").send() + return + + chat_history = cl.user_session.get("chat_history", []) + formatted_messages = [] + for user_msg, assistant_msg in chat_history: + formatted_messages.append({"role": "user", "content": user_msg}) + formatted_messages.append({"role": "assistant", "content": assistant_msg}) + + response_stream = rag( + prompt=query, + system_prompt=RAG_SYSTEM_PROMPT, + search=hybrid_search, + messages=formatted_messages, + max_contexts=5, + config=my_config + ) + + full_response = "" + for chunk in response_stream: + if chunk: + full_response += chunk + await msg.stream_token(chunk) + + if not full_response: + await cl.Message(content="❌ No response generated. Please try rephrasing your question.").send() + return + + chat_history.append((query, full_response)) + cl.user_session.set("chat_history", chat_history) + + except Exception as e: + logger.error(f"RAG error: {str(e)}") + await cl.Message(content=f"❌ Error generating response: {str(e)}").send() + + except Exception as e: + error_msg = f"Error processing your question: {str(e)}" + logger.error(error_msg) + await cl.Message(content=f"❌ {error_msg}").send() + +@cl.on_chat_end +async def on_chat_end(): + """Clean up session data when chat ends.""" + try: + cl.user_session.clear() + logger.info("Chat session ended, cleared session data") + except Exception as e: + logger.error(f"Error clearing session data: {e}") + +if __name__ == "__main__": + cl.run()