from flask import Flask, request, jsonify, send_from_directory from werkzeug.utils import secure_filename import os from file_processor.processor import FileProcessor import gradio as gr from database.db_manager import DatabaseManager from embeddings.embedding_manager import EmbeddingManager from retrieval.vector_store import VectorStore from transformers import pipeline import numpy as np from config import Config import requests import wikipedia import textwrap app = Flask(__name__) app.config['UPLOAD_FOLDER'] = 'uploads' app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB max file size ALLOWED_EXTENSIONS = {'txt', 'pdf', 'docx'} db_manager = DatabaseManager() embedding_manager = EmbeddingManager() vector_store = None os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True) def allowed_file(filename): return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS def get_wikipedia_content(): # Get content from multiple Wikipedia pages about AI try: # List of AI-related topics topics = [ 'Artificial intelligence', 'Machine learning', 'Deep learning', 'Natural language processing', 'Computer vision' ] all_content = [] for topic in topics: try: # Get the Wikipedia page content page = wikipedia.page(topic) # Add the content all_content.append(page.content[:2000]) # Get first 2000 chars of each topic except wikipedia.exceptions.DisambiguationError as e: # If disambiguation page, take the first suggestion try: page = wikipedia.page(e.options[0]) all_content.append(page.content[:2000]) except: continue except: continue return '\n\n'.join(all_content) except Exception as e: # Fallback content in case of any issues return """ Artificial Intelligence (AI) is the simulation of human intelligence processes by machines, especially computer systems. These processes include learning (the acquisition of information and rules for using the information), reasoning (using rules to reach approximate or definite conclusions) and self-correction. Machine Learning is a subset of artificial intelligence that provides systems the ability to automatically learn and improve from experience without being explicitly programmed. Machine learning focuses on the development of computer programs that can access data and use it to learn for themselves. Deep Learning is part of a broader family of machine learning methods based on artificial neural networks with representation learning. Learning can be supervised, semi-supervised or unsupervised. Natural Language Processing (NLP) is a branch of artificial intelligence that helps computers understand, interpret and manipulate human language. NLP draws from many disciplines, including computer science and computational linguistics, in its pursuit to fill the gap between human communication and computer understanding. """ def chunk_text(text, chunk_size=300): """Split text into chunks of approximately equal size.""" # Split into sentences first (crude approach) sentences = [s.strip() for s in text.split('.') if s.strip()] chunks = [] current_chunk = [] current_length = 0 for sentence in sentences: sentence_length = len(sentence) if current_length + sentence_length > chunk_size and current_chunk: # Join the current chunk and add to chunks chunks.append('. '.join(current_chunk) + '.') current_chunk = [sentence] current_length = sentence_length else: current_chunk.append(sentence) current_length += sentence_length # Add the last chunk if it exists if current_chunk: chunks.append('. '.join(current_chunk) + '.') return chunks def init_vector_store(): global vector_store # Get content from Wikipedia text = get_wikipedia_content() # Chunk the text chunks = chunk_text(text, Config.CHUNK_SIZE) # Get embeddings for all chunks embeddings = embedding_manager.get_embeddings(chunks) # Initialize and populate vector store vector_store = VectorStore(dimension=embeddings.shape[1]) vector_store.add_texts(chunks, embeddings) return vector_store def generate_answer(query: str, context: str) -> str: try: # Using a small model for generation generator = pipeline('text-generation', model='gpt2') prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:" response = generator(prompt, max_length=150, num_return_sequences=1) return response[0]['generated_text'] except Exception as e: # Fallback to a simple extraction-based approach relevant_sentences = [s for s in context.split('.') if query.lower() in s.lower()] if relevant_sentences: return relevant_sentences[0] + '.' return "I apologize, but I couldn't generate a specific answer based on the available information." @app.route('/chat', methods=['POST']) def chat(): data = request.json query = data.get('query') if not query: return jsonify({'error': 'No query provided'}), 400 # Save user message db_manager.save_message('user', query) # Get query embedding query_embedding = embedding_manager.get_embedding(query) # Retrieve relevant chunks results = vector_store.search(query_embedding, Config.TOP_K_RESULTS) context = ' '.join([text for text, _ in results]) # Generate answer answer = generate_answer(query, context) # Save system response db_manager.save_message('system', answer) return jsonify({ 'answer': answer, 'retrieved_chunks': [{'text': text, 'score': score} for text, score in results] }) @app.route('/history', methods=['GET']) def history(): chat_history = db_manager.get_chat_history() return jsonify(chat_history) @app.route('/upload', methods=['POST']) def upload_file(): if 'file' not in request.files: return jsonify({'error': 'No file part'}), 400 file = request.files['file'] if file.filename == '': return jsonify({'error': 'No selected file'}), 400 if file and allowed_file(file.filename): filename = secure_filename(file.filename) filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) file.save(filepath) # Process the file try: processor = FileProcessor() content = processor.process_file(filepath) # Chunk the content chunks = chunk_text(content, Config.CHUNK_SIZE) # Get embeddings for chunks embeddings = embedding_manager.get_embeddings(chunks) # Add to vector store vector_store.add_texts(chunks, embeddings) return jsonify({ 'message': 'File processed successfully', 'chunks_added': len(chunks) }) except Exception as e: return jsonify({'error': str(e)}), 500 finally: # Clean up uploaded file os.remove(filepath) return jsonify({'error': 'Invalid file type'}), 400 def create_gradio_interface(): def chat_function(message, history): # Process query through RAG system query_embedding = embedding_manager.get_embedding(message) results = vector_store.search(query_embedding, Config.TOP_K_RESULTS) context = ' '.join([text for text, _ in results]) answer = generate_answer(message, context) # Save to database db_manager.save_message('user', message) db_manager.save_message('system', answer) return answer def handle_file_upload(file): if file is None: return "No file uploaded" try: processor = FileProcessor() content = processor.process_file(file.name) chunks = chunk_text(content, Config.CHUNK_SIZE) embeddings = embedding_manager.get_embeddings(chunks) vector_store.add_texts(chunks, embeddings) return f"File processed successfully. Added {len(chunks)} chunks to knowledge base." except Exception as e: return f"Error processing file: {str(e)}" # Create Gradio interface with gr.Blocks() as demo: gr.Markdown("# RAG Chatbot") with gr.Row(): file_input = gr.File(label="Upload Document") upload_button = gr.Button("Process File") upload_output = gr.Textbox(label="Upload Status") chatbot = gr.ChatInterface( chat_function, examples=["What is artificial intelligence?", "Explain machine learning"], title="Chat with your documents" ) upload_button.click( handle_file_upload, inputs=[file_input], outputs=[upload_output] ) return demo if __name__ == '__main__': print("Creating database tables...") from database.models import create_tables create_tables() print("Initializing vector store with Wikipedia content...") init_vector_store() print("Starting Gradio interface...") demo = create_gradio_interface() demo.launch(server_name="0.0.0.0", server_port=7860)