|
|
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) |
