| | import os |
| | import uuid |
| | from flask import Flask, request, render_template, session, jsonify, Response, stream_with_context |
| | from werkzeug.utils import secure_filename |
| | from rag_processor import create_rag_chain |
| | from typing import Sequence, Any, List |
| | import fitz |
| | import re |
| | import io |
| | from gtts import gTTS |
| | from langchain_core.documents import Document |
| | from langchain_community.document_loaders import TextLoader, Docx2txtLoader |
| | from langchain.text_splitter import RecursiveCharacterTextSplitter |
| | from langchain_huggingface import HuggingFaceEmbeddings |
| | from langchain_community.vectorstores import FAISS |
| | from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever |
| | from langchain.retrievers.document_compressors.base import BaseDocumentCompressor |
| | from langchain_community.retrievers import BM25Retriever |
| | from langchain_community.chat_message_histories import ChatMessageHistory |
| | from langchain.storage import InMemoryStore |
| | from sentence_transformers.cross_encoder import CrossEncoder |
| |
|
| | app = Flask(__name__) |
| | app.config['SECRET_KEY'] = os.urandom(24) |
| | TEMPERATURE_LABELS = { |
| | '0.2': 'Precise', |
| | '0.4': 'Confident', |
| | '0.6': 'Balanced', |
| | '0.8': 'Flexible', |
| | '1.0': 'Creative', |
| | } |
| | class LocalReranker(BaseDocumentCompressor): |
| | model: Any |
| | top_n: int = 5 |
| |
|
| | class Config: |
| | arbitrary_types_allowed = True |
| |
|
| | def compress_documents(self, documents: Sequence[Document], query: str, |
| | callbacks=None) -> Sequence[Document]: |
| | if not documents: |
| | return [] |
| | pairs = [[query, doc.page_content] for doc in documents] |
| | scores = self.model.predict(pairs, show_progress_bar=False) |
| | doc_scores = list(zip(documents, scores)) |
| | sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], |
| | reverse=True) |
| | top_docs = [] |
| | for (doc, score) in sorted_doc_scores[:self.top_n]: |
| | doc.metadata['rerank_score'] = float(score) |
| | top_docs.append(doc) |
| | return top_docs |
| |
|
| |
|
| | def create_optimized_parent_child_chunks(all_docs): |
| | if not all_docs: |
| | print ('CHUNKING: No input documents provided!') |
| | return ([], [], []) |
| |
|
| | parent_splitter = RecursiveCharacterTextSplitter(chunk_size=900, |
| | chunk_overlap=200, separators=['\n\n', '\n', '. ', '! ', |
| | '? ', '; ', ', ', ' ', '']) |
| | child_splitter = RecursiveCharacterTextSplitter(chunk_size=350, |
| | chunk_overlap=80, separators=['\n', '. ', '! ', '? ', '; ', |
| | ', ', ' ', '']) |
| | parent_docs = parent_splitter.split_documents(all_docs) |
| | doc_ids = [str(uuid.uuid4()) for _ in parent_docs] |
| | child_docs = [] |
| |
|
| | for (i, parent_doc) in enumerate(parent_docs): |
| | parent_id = doc_ids[i] |
| | children = child_splitter.split_documents([parent_doc]) |
| | for (j, child) in enumerate(children): |
| | child.metadata.update({'doc_id': parent_id, |
| | 'chunk_index': j, |
| | 'total_chunks': len(children), |
| | 'is_first_chunk': j == 0, |
| | 'is_last_chunk': j == len(children) |
| | - 1}) |
| | if len(children) > 1: |
| | if j == 0: |
| | child.page_content = '[Beginning] ' + child.page_content |
| | elif j == len(children) - 1: |
| | child.page_content = '[Continues...] ' + child.page_content |
| | child_docs.append(child) |
| |
|
| | print (f"CHUNKING: Created {len(parent_docs)} parent and {len(child_docs)} child chunks." |
| | ) |
| | return (parent_docs, child_docs, doc_ids) |
| | def get_context_aware_parents(docs: List[Document], store: InMemoryStore) -> List[Document]: |
| | if not docs: |
| | return [] |
| | (parent_scores, child_content_by_parent) = ({}, {}) |
| | for doc in docs: |
| | parent_id = doc.metadata.get('doc_id') |
| | if parent_id: |
| | parent_scores[parent_id] = parent_scores.get(parent_id, 0) \ |
| | + 1 |
| | if parent_id not in child_content_by_parent: |
| | child_content_by_parent[parent_id] = [] |
| | child_content_by_parent[parent_id].append(doc.page_content) |
| |
|
| | parent_ids = list(parent_scores.keys()) |
| | parents = store.mget(parent_ids) |
| | enhanced_parents = [] |
| |
|
| | for (i, parent) in enumerate(parents): |
| | if parent is not None: |
| | parent_id = parent_ids[i] |
| | if parent_id in child_content_by_parent: |
| | child_excerpts = '\n'.join(child_content_by_parent[parent_id][:3]) |
| | enhanced_content = f"{parent.page_content}\n\nRelevant excerpts:\n{child_excerpts}" |
| | enhanced_parent =Document(page_content=enhanced_content, |
| | metadata={**parent.metadata, |
| | 'child_relevance_score': parent_scores[parent_id], |
| | 'matching_children': len(child_content_by_parent[parent_id])}) |
| | enhanced_parents.append(enhanced_parent) |
| | else: |
| | print (f"PARENT_FETCH: Parent {parent_ids[i]} not found in store!") |
| |
|
| | enhanced_parents.sort(key=lambda p: p.metadata.get('child_relevance_score', 0), reverse=True) |
| | return enhanced_parents |
| |
|
| |
|
| | is_hf_spaces = bool(os.getenv('SPACE_ID') or os.getenv('SPACES_ZERO_GPU' |
| | )) |
| | app.config['UPLOAD_FOLDER'] = '/tmp/uploads' if is_hf_spaces else 'uploads' |
| |
|
| | try: |
| | os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True) |
| | print (f"Upload folder ready: {app.config['UPLOAD_FOLDER']}") |
| | except Exception as e: |
| | print (f"Failed to create upload folder, falling back to /tmp: {e}") |
| | app.config['UPLOAD_FOLDER'] = '/tmp/uploads' |
| | os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True) |
| |
|
| | session_data = {} |
| | message_histories = {} |
| |
|
| | print ('Loading embedding model...') |
| | try: |
| | EMBEDDING_MODEL = \ |
| | HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2' |
| | , model_kwargs={'device': 'cpu'}, |
| | encode_kwargs={'normalize_embeddings': True}) |
| | print ('Embedding model loaded.') |
| | except Exception as e: |
| | print (f"FATAL: Could not load embedding model. Error: {e}") |
| | raise e |
| |
|
| | print ('Loading reranker model...') |
| | try: |
| | RERANKER_MODEL = \ |
| | CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', |
| | device='cpu') |
| | print ('Reranker model loaded.') |
| | except Exception as e: |
| | print (f"FATAL: Could not load reranker model. Error: {e}") |
| | raise e |
| |
|
| |
|
| | def load_pdf_with_fallback(filepath): |
| | try: |
| | docs = [] |
| | with fitz.open(filepath) as pdf_doc: |
| | for (page_num, page) in enumerate(pdf_doc): |
| | text = page.get_text() |
| | if text.strip(): |
| | docs.append(Document(page_content=text, |
| | metadata={'source': os.path.basename(filepath), |
| | 'page': page_num + 1})) |
| | if docs: |
| | print (f"Loaded PDF: {os.path.basename(filepath)} - {len(docs)} pages" |
| | ) |
| | return docs |
| | else: |
| | raise ValueError('No text content found in PDF.') |
| | except Exception as e: |
| | print (f"PyMuPDF failed for {filepath}: {e}") |
| | raise |
| |
|
| |
|
| | LOADER_MAPPING = {'.txt': TextLoader, '.pdf': load_pdf_with_fallback, |
| | '.docx': Docx2txtLoader} |
| |
|
| |
|
| | def get_session_history(session_id: str) -> ChatMessageHistory: |
| | if session_id not in message_histories: |
| | message_histories[session_id] = ChatMessageHistory() |
| | return message_histories[session_id] |
| |
|
| |
|
| | @app.route('/health', methods=['GET']) |
| | def health_check(): |
| | return (jsonify({'status': 'healthy'}), 200) |
| |
|
| |
|
| | @app.route('/', methods=['GET']) |
| | def index(): |
| | return render_template('index.html') |
| |
|
| |
|
| | @app.route('/upload', methods=['POST']) |
| | def upload_files(): |
| | files = request.files.getlist('file') |
| |
|
| | temperature_str = request.form.get('temperature', '0.2') |
| | temperature = float(temperature_str) |
| | model_name = request.form.get('model_name', |
| | 'moonshotai/kimi-k2-instruct') |
| | print (f"UPLOAD: Model: {model_name}, Temp: {temperature}") |
| |
|
| | if not files or all(f.filename == '' for f in files): |
| | return (jsonify({'status': 'error', |
| | 'message': 'No selected files.'}), 400) |
| |
|
| | (all_docs, processed_files, failed_files) = ([], [], []) |
| | print (f"Processing {len(files)} file(s)...") |
| | for file in files: |
| | if file and file.filename: |
| | filename = secure_filename(file.filename) |
| | filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) |
| | try: |
| | file.save(filepath) |
| | file_ext = os.path.splitext(filename)[1].lower() |
| | if file_ext not in LOADER_MAPPING: |
| | raise ValueError('Unsupported file format.') |
| | loader_func = LOADER_MAPPING[file_ext] |
| | docs = loader_func(filepath) if file_ext == '.pdf' \ |
| | else loader_func(filepath).load() |
| | if not docs: |
| | raise ValueError('No content extracted.') |
| | all_docs.extend(docs) |
| | processed_files.append(filename) |
| | except Exception as e: |
| | print (f"✗ Error processing {filename}: {e}") |
| | failed_files.append(f"{filename} ({e})") |
| |
|
| | if not all_docs: |
| | return (jsonify({'status': 'error', |
| | 'message': f"Failed to process all files. Reasons: {', '.join(failed_files)}" |
| | }), 400) |
| |
|
| | print (f"UPLOAD: Processed {len(processed_files)} files.") |
| | try: |
| | print ('Starting RAG pipeline setup...') |
| | (parent_docs, child_docs, doc_ids) = \ |
| | create_optimized_parent_child_chunks(all_docs) |
| | if not child_docs: |
| | raise ValueError('No child documents created during chunking.') |
| |
|
| | vectorstore = FAISS.from_documents(child_docs, EMBEDDING_MODEL) |
| | store = InMemoryStore() |
| | store.mset(list(zip(doc_ids, parent_docs))) |
| | print (f"Indexed {len(child_docs)} document chunks.") |
| |
|
| | bm25_retriever = BM25Retriever.from_documents(child_docs) |
| | bm25_retriever.k = 12 |
| | faiss_retriever = vectorstore.as_retriever(search_kwargs={'k': 12}) |
| | ensemble_retriever = \ |
| | EnsembleRetriever(retrievers=[bm25_retriever, |
| | faiss_retriever], weights=[0.6, 0.4]) |
| | reranker = LocalReranker(model=RERANKER_MODEL, top_n=5) |
| |
|
| | def get_parents(docs: List[Document]) -> List[Document]: |
| | return get_context_aware_parents(docs, store) |
| |
|
| | compression_retriever = \ |
| | ContextualCompressionRetriever(base_compressor=reranker, |
| | base_retriever=ensemble_retriever) |
| | final_retriever = compression_retriever | get_parents |
| |
|
| | session_id = str(uuid.uuid4()) |
| | (rag_chain, api_key_manager) = \ |
| | create_rag_chain(retriever=final_retriever, |
| | get_session_history_func=get_session_history, |
| | model_name=model_name, |
| | temperature=temperature) |
| |
|
| | session_data[session_id] = {'chain': rag_chain, |
| | 'model_name': model_name, |
| | 'temperature': temperature, |
| | 'api_key_manager': api_key_manager} |
| |
|
| | success_msg = f"Processed: {', '.join(processed_files)}" |
| | if failed_files: |
| | success_msg += f". Failed: {', '.join(failed_files)}" |
| |
|
| | mode_label = TEMPERATURE_LABELS.get(temperature_str, |
| | temperature_str) |
| |
|
| | print (f"UPLOAD COMPLETE: Session {session_id} is ready.") |
| |
|
| | return jsonify({ |
| | 'status': 'success', |
| | 'filename': success_msg, |
| | 'session_id': session_id, |
| | 'model_name': model_name, |
| | 'mode': mode_label, |
| | }) |
| | except Exception as e: |
| | import traceback |
| | traceback.print_exc() |
| | return (jsonify({'status': 'error', |
| | 'message': f'RAG setup failed: {e}'}), 500) |
| |
|
| | @app.route('/chat', methods=['POST', 'GET']) |
| | def chat(): |
| | if request.method == 'GET': |
| | question = request.args.get('question') |
| | session_id = request.args.get('session_id') |
| | print(f"Received GET request for chat: session={session_id}, question={question[:50]}...") |
| | elif request.method == 'POST': |
| | data = request.get_json() |
| | question = data.get('question') |
| | session_id = data.get('session_id') or session.get('session_id') |
| | print(f"Received POST request for chat: session={session_id}, question={question[:50]}...") |
| | else: |
| | return (jsonify({'status': 'error', 'message': 'Method not allowed'}), 405) |
| |
|
| | if not question: |
| | error_msg = "Error: No question provided." |
| | print(f"CHAT Validation Error: {error_msg}") |
| | if request.method == 'GET': |
| | def error_stream(): |
| | yield f'data: {{"error": "{error_msg}"}}\n\n' |
| | return Response(stream_with_context(error_stream()), mimetype='text/event-stream', status=400) |
| | return jsonify({'status': 'error','message': error_msg}), 400 |
| |
|
| | if not session_id or session_id not in session_data: |
| | error_msg = "Error: Invalid session. Please upload documents first." |
| | print(f"CHAT Validation Error: Invalid session {session_id}.") |
| | if request.method == 'GET': |
| | def error_stream(): |
| | yield f'data: {{"error": "{error_msg}"}}\n\n' |
| | return Response(stream_with_context(error_stream()), mimetype='text/event-stream', status=400) |
| | return jsonify({'status': 'error', 'message': error_msg }), 400 |
| | try: |
| | session_info = session_data[session_id] |
| | rag_chain = session_info['chain'] |
| | model_name = session_info['model_name'] |
| | temperature_float = session_info['temperature'] |
| | temperature_str = str(temperature_float) |
| | mode_label = TEMPERATURE_LABELS.get(temperature_str, temperature_str) |
| |
|
| | print (f"CHAT: Streaming response for session {session_id} (Model: {model_name}, Temp: {temperature_float})...") |
| |
|
| | def generate_chunks(): |
| | full_response = '' |
| | try: |
| | stream_iterator = rag_chain.stream({'question': question}, |
| | config={'configurable': {'session_id': session_id}}) |
| |
|
| | for chunk in stream_iterator: |
| | if isinstance(chunk, str): |
| | full_response += chunk |
| | token_escaped = chunk.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n') |
| | model_name_escaped = model_name.replace('"', '\\"') |
| | mode_label_escaped = mode_label.replace('"', '\\"') |
| | yield f'data: {{"token": "{token_escaped}", "model_name": "{model_name_escaped}", "mode": "{mode_label_escaped}"}}\n\n' |
| | else: |
| | print(f"Received non-string chunk: {type(chunk)}") |
| |
|
| |
|
| | print ('CHAT: Streaming finished successfully.') |
| |
|
| | except Exception as e: |
| | print(f"CHAT Error during streaming generation: {e}") |
| | import traceback |
| | traceback.print_exc() |
| | error_msg = f"Error during response generation: {str(e)}".replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n') |
| | yield f'data: {{"error": "{error_msg}"}}\n\n' |
| | return Response(stream_with_context(generate_chunks()), mimetype='text/event-stream') |
| |
|
| | except Exception as e: |
| | print(f"CHAT Setup Error: {e}") |
| | import traceback |
| | traceback.print_exc() |
| | error_msg = f"Error setting up chat stream: {str(e)}" |
| | if request.method == 'GET': |
| | def error_stream(): |
| | clean_error_msg= error_msg.replace("\"", "\\\"").replace("n", "\\n") |
| | yield f'data: {{"error": "{clean_error_msg}"}}\n\n' |
| | return Response(stream_with_context(error_stream()), mimetype='text/event-stream', status=500) |
| | return (jsonify({'status': 'error', 'message': error_msg}), 500) |
| |
|
| |
|
| | def clean_markdown_for_tts(text: str) -> str: |
| | text = re.sub(r'\[.*?\]\(.*?\)', '', text) |
| | text = re.sub(r'[`*_#]', '', text) |
| | text = re.sub(r'^\s*[\-\*\+]\s+', '', text, flags=re.MULTILINE) |
| | text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE) |
| | text = re.sub(r'^\s*>\s?', '', text, flags=re.MULTILINE) |
| | text = re.sub(r'\n+', ' ', text) |
| | text = re.sub(r'\s{2,}', ' ', text) |
| | return text.strip() |
| |
|
| |
|
| | @app.route('/tts', methods=['POST']) |
| | def text_to_speech(): |
| | data = request.get_json() |
| | text = data.get('text') |
| | if not text: |
| | return (jsonify({'status': 'error', |
| | 'message': 'No text provided.'}), 400) |
| | try: |
| | clean_text = clean_markdown_for_tts(text) |
| | if not clean_text: |
| | return (jsonify({'status': 'error', 'message': 'No speakable text found.'}), 400) |
| |
|
| | tts = gTTS(clean_text, lang='en') |
| | mp3_fp = io.BytesIO() |
| | tts.write_to_fp(mp3_fp) |
| | mp3_fp.seek(0) |
| | return Response(mp3_fp, mimetype='audio/mpeg') |
| | except Exception as e: |
| | print (f"TTS Error: {e}") |
| | return (jsonify({'status': 'error', |
| | 'message': 'Failed to generate audio.'}), 500) |
| |
|
| |
|
| | if __name__ == '__main__': |
| | port = int(os.environ.get('PORT', 7860)) |
| | print (f"Starting Flask app on port {port}") |
| | app.run(host='0.0.0.0', port=port, debug=False, threaded=True) |
