Aigenthix/Graph_RAG7
0
1import os2import json3import logging4from datetime import datetime5from flask import Flask, render_template, request, jsonify, send_file6from flask_cors import CORS7import threading8from pathlib import Path9import numpy as np10import networkx as nx11import matplotlib12matplotlib.use('Agg')13import matplotlib.pyplot as plt14from io import BytesIO15import base6416 17from sentence_transformers import SentenceTransformer18from groq import Groq19import PyPDF220import pandas as pd21from langchain.text_splitter import RecursiveCharacterTextSplitter22 23logging.basicConfig(level=logging.INFO)24logger = logging.getLogger(__name__)25 26app = Flask(__name__)27CORS(app)28 29UPLOAD_FOLDER = Path('./data/uploads')30GRAPH_DATA_FOLDER = Path('./data/graph_data')31UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)32GRAPH_DATA_FOLDER.mkdir(parents=True, exist_ok=True)33 34app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER35app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 102436 37embedding_model = SentenceTransformer('all-MiniLM-L6-v2')38 39groq_api_key = os.getenv('GROQ_API_KEY', '')40if groq_api_key:41 groq_client = Groq(api_key=groq_api_key)42else:43 groq_client = None44 logger.warning("GROQ_API_KEY not set - chat queries will fail")45 46documents_state = {}47graph_data = {}48 49class DocumentProcessor:50 @staticmethod51 def extract_text(file_path):52 file_ext = Path(file_path).suffix.lower()53 54 if file_ext == '.pdf':55 text = DocumentProcessor._extract_pdf(file_path)56 elif file_ext == '.csv':57 text = DocumentProcessor._extract_csv(file_path)58 elif file_ext == '.txt':59 with open(file_path, 'r', encoding='utf-8') as f:60 text = f.read()61 else:62 raise ValueError(f"Unsupported file type: {file_ext}")63 64 return text65 66 @staticmethod67 def _extract_pdf(file_path):68 text = []69 with open(file_path, 'rb') as f:70 reader = PyPDF2.PdfReader(f)71 for page in reader.pages:72 text.append(page.extract_text())73 return '\n'.join(text)74 75 @staticmethod76 def _extract_csv(file_path):77 df = pd.read_csv(file_path)78 return df.to_string()79 80 @staticmethod81 def chunk_text(text, chunk_size=500, overlap=100):82 splitter = RecursiveCharacterTextSplitter(83 chunk_size=chunk_size,84 chunk_overlap=overlap,85 separators=["\n\n", "\n", " ", ""]86 )87 chunks = splitter.split_text(text)88 return chunks89 90class GraphBuilder:91 @staticmethod92 def build_knowledge_graph(chunks, doc_name):93 graph = nx.DiGraph()94 entities = set()95 96 for chunk in chunks:97 words = chunk.split()[:10]98 main_entity = f"{doc_name}_chunk_{chunks.index(chunk)}"99 graph.add_node(main_entity, type='chunk', content=chunk[:200])100 101 for word in words:102 if len(word) > 3:103 word_node = word.lower()104 graph.add_node(word_node, type='entity')105 graph.add_edge(main_entity, word_node, weight=1.0)106 entities.add(word_node)107 108 return graph, list(entities)109 110 @staticmethod111 def visualize_graph(graph, output_path):112 plt.figure(figsize=(14, 10))113 114 if len(graph.nodes()) == 0:115 plt.text(0.5, 0.5, 'Empty Graph', ha='center', va='center')116 else:117 pos = nx.spring_layout(graph, k=2, iterations=50, seed=42)118 119 node_colors = []120 for node in graph.nodes():121 if graph.nodes[node].get('type') == 'chunk':122 node_colors.append('lightblue')123 else:124 node_colors.append('lightgreen')125 126 nx.draw_networkx_nodes(graph, pos, node_color=node_colors,127 node_size=500, alpha=0.9)128 nx.draw_networkx_edges(graph, pos, edge_color='gray',129 arrows=True, alpha=0.5, width=1.5)130 131 labels = {node: node[:15] for node in graph.nodes()}132 nx.draw_networkx_labels(graph, pos, labels, font_size=8)133 134 plt.title('Knowledge Graph Visualization', fontsize=16, fontweight='bold')135 plt.axis('off')136 plt.tight_layout()137 plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')138 plt.close()139 140 logger.info(f"Graph visualization saved to {output_path}")141 142def process_document_async(filename, file_path):143 try:144 logger.info(f"Processing document: {filename}")145 documents_state[filename] = {146 'status': 'processing',147 'progress': 10,148 'error': None,149 'chunks': 0,150 'entities': 0151 }152 153 text = DocumentProcessor.extract_text(file_path)154 documents_state[filename]['progress'] = 40155 156 chunks = DocumentProcessor.chunk_text(text)157 documents_state[filename]['progress'] = 60158 logger.info(f"Created {len(chunks)} chunks from {filename}")159 160 graph, entities = GraphBuilder.build_knowledge_graph(chunks, filename)161 documents_state[filename]['progress'] = 80162 logger.info(f"Built graph with {len(graph.nodes())} nodes for {filename}")163 164 graph_path = GRAPH_DATA_FOLDER / f"{filename}_graph.png"165 GraphBuilder.visualize_graph(graph, graph_path)166 documents_state[filename]['progress'] = 95167 168 embeddings = embedding_model.encode(chunks, show_progress_bar=False)169 170 graph_data[filename] = {171 'chunks': chunks,172 'embeddings': embeddings.tolist(),173 'graph': nx.node_link_data(graph),174 'entities': entities,175 'graph_image': str(graph_path)176 }177 178 documents_state[filename] = {179 'status': 'ready',180 'progress': 100,181 'error': None,182 'chunks': len(chunks),183 'entities': len(entities),184 'graph_image': f"/graph-image/{filename}",185 'timestamp': datetime.now().isoformat()186 }187 188 logger.info(f"Successfully processed {filename}")189 190 except Exception as e:191 logger.error(f"Error processing {filename}: {str(e)}")192 documents_state[filename] = {193 'status': 'error',194 'progress': 0,195 'error': str(e),196 'chunks': 0,197 'entities': 0198 }199 200@app.route('/')201def index():202 return render_template('index.html')203 204@app.route('/api/documents', methods=['GET'])205def get_documents():206 return jsonify({207 'documents': documents_state,208 'api_key_set': bool(os.getenv('GROQ_API_KEY')),209 'timestamp': datetime.now().isoformat()210 })211 212@app.route('/api/upload', methods=['POST'])213def upload_document():214 if 'files' not in request.files:215 return jsonify({'error': 'No files provided'}), 400216 217 files = request.files.getlist('files')218 results = {'successful': 0, 'failed': 0, 'files': []}219 220 for file in files:221 if not file or file.filename == '':222 results['failed'] += 1223 continue224 225 filename = file.filename226 file_path = UPLOAD_FOLDER / filename227 file.save(file_path)228 229 documents_state[filename] = {230 'status': 'queued',231 'progress': 0,232 'error': None,233 'chunks': 0,234 'entities': 0235 }236 237 thread = threading.Thread(target=process_document_async, args=(filename, file_path))238 thread.daemon = True239 thread.start()240 241 results['successful'] += 1242 results['files'].append(filename)243 244 return jsonify({245 'success': True,246 'message': f"✅ {results['successful']} file(s) queued for processing",247 **results248 })249 250@app.route('/api/query', methods=['POST'])251def query():252 data = request.json253 query_text = data.get('query', '').strip()254 doc_name = data.get('document', '')255 256 if not query_text or not doc_name:257 return jsonify({'error': 'Missing query or document'}), 400258 259 if doc_name not in graph_data:260 return jsonify({'error': 'Document not found or not ready'}), 404261 262 if not groq_client:263 return jsonify({'error': 'GROQ_API_KEY not configured. Chat is unavailable.'}), 500264 265 try:266 doc_info = graph_data[doc_name]267 query_embedding = embedding_model.encode(query_text, show_progress_bar=False)268 269 embeddings = np.array(doc_info['embeddings'])270 similarities = np.dot(embeddings, query_embedding) / (271 np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_embedding) + 1e-10272 )273 274 top_k_indices = np.argsort(similarities)[-3:][::-1]275 relevant_chunks = [doc_info['chunks'][i] for i in top_k_indices]276 277 # Filter by lower threshold but always include at least top 1278 high_sim_indices = [i for i in top_k_indices if similarities[i] > 0.1]279 if high_sim_indices:280 top_k_indices = high_sim_indices281 relevant_chunks = [doc_info['chunks'][i] for i in top_k_indices]282 283 # Always use top chunks even if similarity is low284 if not relevant_chunks:285 return jsonify({286 'answer': 'Unable to process this query. Please try with a different document or question.',287 'sources': [],288 'confidence': 0.0289 })290 291 context = "\n".join(relevant_chunks)292 293 message = groq_client.chat.completions.create(294 model="llama-3.3-70b-versatile",295 max_tokens=500,296 messages=[297 {"role": "user", "content": f"""Based on this context, answer the question concisely.298 299Context: {context}300 301Question: {query_text}302 303Answer:"""}304 ]305 )306 307 answer = message.choices[0].message.content.strip()308 309 return jsonify({310 'answer': answer,311 'sources': [f"Chunk {idx+1}" for idx in top_k_indices],312 'confidence': float(max(similarities[top_k_indices]))313 })314 315 except Exception as e:316 logger.error(f"Query error: {str(e)}")317 return jsonify({'error': str(e)}), 500318 319@app.route('/graph-image/<filename>')320def get_graph_image(filename):321 graph_path = GRAPH_DATA_FOLDER / f"{filename}_graph.png"322 if graph_path.exists():323 return send_file(graph_path, mimetype='image/png')324 return jsonify({'error': 'Graph not found'}), 404325 326@app.route('/api/delete/<filename>', methods=['DELETE'])327def delete_document(filename):328 if filename in documents_state:329 del documents_state[filename]330 if filename in graph_data:331 del graph_data[filename]332 333 file_path = UPLOAD_FOLDER / filename334 if file_path.exists():335 file_path.unlink()336 337 graph_path = GRAPH_DATA_FOLDER / f"{filename}_graph.png"338 if graph_path.exists():339 graph_path.unlink()340 341 return jsonify({'success': True})342 343if __name__ == '__main__':344 port = int(os.getenv('PORT', 7860))345 app.run(host='0.0.0.0', port=port, debug=False, threaded=True)