import os from openai import OpenAI import numpy as np import sqlite3 import uuid from datetime import datetime from typing import List, Optional from dotenv import load_dotenv from pypdf import PdfReader from pathlib import Path import re from fastapi import FastAPI, HTTPException, UploadFile, File from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse from pydantic import BaseModel import shutil # ============================== # 0. Sozlamalar # ============================== env_path = Path(__file__).resolve().parent / ".env" load_dotenv(env_path, override=True) OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") if not OPENAI_API_KEY: raise RuntimeError("OPENAI_API_KEY topilmadi") openai_client = OpenAI(api_key=OPENAI_API_KEY) OPENAI_EMBED_MODEL = "text-embedding-3-small" OPENAI_CHAT_MODEL = "gpt-4o-mini" CHROMA_DIR = "./chroma_db" CHAT_DB_PATH = "./chats.db" # Use PersistentClient (Modern API) # Fallback for Python 3.14 compatibility try: import chromadb chroma_client = chromadb.PersistentClient(path=CHROMA_DIR) collection = chroma_client.get_or_create_collection(name="rag_docs") RAG_AVAILABLE = True except Exception as e: print(f"WARNING: ChromaDB failed to initialize (likely Python version mismatch): {e}") chromadb = None chroma_client = None collection = None RAG_AVAILABLE = False print("RAG_AVAILABLE:", RAG_AVAILABLE) print("COLLECTION COUNT:", collection.count() if collection else 0) # ============================== # 1. PDF -> TEXT # ============================== # def load_pdf(path: str) -> str: # reader = PdfReader(path) # text = "" # for page in reader.pages: # page_text = page.extract_text() # if page_text: # text += page_text + "\n" # return text def fix_spaced_text(line: str) -> str: if re.fullmatch(r'(?:[A-Za-z]\s+){3,}[A-Za-z]?', line.strip()): return line.replace(" ", "") return line def load_pdf(path: str) -> str: reader = PdfReader(path) lines = [] for page in reader.pages: page_text = page.extract_text() if not page_text: continue for line in page_text.splitlines(): cleaned = fix_spaced_text(line) if cleaned: lines.append(cleaned) text = "\n".join(lines) text = re.sub(r'(?<=\w)\s*@\s*(?=\w)', '@', text) text = re.sub(r'(?<=\w)\s*\.\s*(?=\w)', '.', text) text = re.sub(r'\n{3,}', '\n\n', text) text = re.sub(r'[ \t]{2,}', ' ', text) return text.strip() def extract_email_from_context(context_list: List[str]) -> Optional[str]: text = "\n".join(context_list) text = re.sub(r'(?<=\w)\s*@\s*(?=\w)', '@', text) text = re.sub(r'(?<=\w)\s*\.\s*(?=\w)', '.', text) text = re.sub(r'(?<=\w)\s+(?=\w@)', '', text) text = re.sub(r'(?<=@)\s+(?=\w)', '', text) match = re.search(r'[\w\.-]+@[\w\.-]+\.\w+', text) return match.group(0) if match else None def extract_project_lines(context_list: List[str]) -> List[str]: text = "\n".join(context_list) lines = [line.strip() for line in text.splitlines() if line.strip()] keywords = ["project", "cs core", "trusty", "mily", "corporate solutions"] result = [] for line in lines: low = line.lower() if any(k in low for k in keywords): result.append(line) return result[:8] # ============================== # 2. TEXT -> CHUNKS # ============================== def chunk_text(text, chunk_size=500, overlap=100): chunks = [] start = 0 while start < len(text): end = start + chunk_size chunks.append(text[start:end]) start = end - overlap return chunks # ============================== # 3. CHUNKS -> EMBEDDINGS # ============================== # def embed_texts(texts): # # Ollama embedding for a list of texts # embeddings = [] # for text in texts: # response = client.models.embed_content( # model=GEMINI_EMBED_MODEL, # contents=text # ) # embeddings.append(response.embeddings[0].values) # return embeddings # qwen # def embed_texts(texts): # embeddings = embedder.encode(texts, convert_to_numpy=True, normalize_embeddings=True) # return embeddings.tolist() OPENAI_EMBED_MODEL = "text-embedding-3-small" def embed_texts(texts: List[str]) -> List[List[float]]: cleaned = [t.strip() for t in texts if t and t.strip()] if not cleaned: return [] response = openai_client.embeddings.create( model=OPENAI_EMBED_MODEL, input=cleaned ) return [item.embedding for item in response.data] # ============================== # 4. VECTOR DB (SAVE / LOAD) # ============================== # ============================== # 4. VECTOR DB (SAVE / LOAD) # ============================== def save_to_chroma(chunks, embeddings, doc_id): if not RAG_AVAILABLE: return """ Revised to include doc_id in metadata for deletion support. """ ids = [f"{doc_id}_chunk_{i}" for i in range(len(chunks))] metadatas = [{"chunk_index": i, "doc_id": doc_id} for i in range(len(chunks))] collection.add( documents=chunks, embeddings=embeddings, ids=ids, metadatas=metadatas ) def delete_from_chroma(doc_id): if not RAG_AVAILABLE: return """Delete all chunks associated with a document ID""" try: collection.delete(where={"doc_id": doc_id}) except Exception as e: print(f"Error deleting from Chroma: {e}") # ============================== # 5. SIMILARITY SEARCH # ============================== # ============================== # Helper for RAG Tool # ============================== def find_context(query, top_k=4): if not RAG_AVAILABLE: return [] try: query_embedding = embed_texts([query])[0] results = collection.query( query_embeddings=[query_embedding], n_results=top_k ) return results["documents"][0] except Exception as e: print(f"RAG Error: {e}") return [] def retrieve_documents(query: str) -> str: """ Retrieve relevant information from the uploaded documents based on the query. Use this tool when the user asks questions about specific documents or content that might be contained in the uploaded PDF files. """ if not RAG_AVAILABLE: return "System Notification: RAG system is currently unavailable." contexts = find_context(query) if not contexts: return "No relevant information found in documents." return "\n\n---\n\n".join(contexts) # ... (init_db, CRUD, etc - skipped for brevity in tool call logic, assuming target content matches) # ============================== # 6. CHAT & DOCUMENT DATABASE SETUP # ============================== def init_db(): """Initialize SQLite database for chat and document persistence""" conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() # Create chats table cursor.execute(""" CREATE TABLE IF NOT EXISTS chats ( id TEXT PRIMARY KEY, title TEXT NOT NULL, created_at TEXT NOT NULL ) """) # Create messages table cursor.execute(""" CREATE TABLE IF NOT EXISTS messages ( id TEXT PRIMARY KEY, chat_id TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL, timestamp TEXT NOT NULL, FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE ) """) # Create documents table cursor.execute(""" CREATE TABLE IF NOT EXISTS documents ( id TEXT PRIMARY KEY, filename TEXT NOT NULL, upload_date TEXT NOT NULL, status TEXT NOT NULL ) """) conn.commit() conn.close() # Initialize database init_db() # ============================== # 7. CHAT CRUD OPERATIONS # ============================== def create_chat(title: str = "New Chat") -> dict: """Create a new chat session""" chat_id = str(uuid.uuid4()) created_at = datetime.now().isoformat() conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() cursor.execute( "INSERT INTO chats (id, title, created_at) VALUES (?, ?, ?)", (chat_id, title, created_at) ) conn.commit() conn.close() return {"id": chat_id, "title": title, "created_at": created_at} def get_all_chats() -> List[dict]: """Get all chats ordered by creation date descending""" conn = sqlite3.connect(CHAT_DB_PATH) conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute("SELECT * FROM chats ORDER BY created_at DESC") rows = cursor.fetchall() conn.close() return [dict(row) for row in rows] def get_chat_by_id(chat_id: str) -> Optional[dict]: """Get a single chat by ID""" conn = sqlite3.connect(CHAT_DB_PATH) conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute("SELECT * FROM chats WHERE id = ?", (chat_id,)) row = cursor.fetchone() conn.close() return dict(row) if row else None def update_chat_title(chat_id: str, title: str) -> bool: """Update chat title""" conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() cursor.execute("UPDATE chats SET title = ? WHERE id = ?", (title, chat_id)) affected = cursor.rowcount conn.commit() conn.close() return affected > 0 def delete_chat(chat_id: str) -> bool: """Delete a chat and all its messages""" conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() cursor.execute("DELETE FROM messages WHERE chat_id = ?", (chat_id,)) cursor.execute("DELETE FROM chats WHERE id = ?", (chat_id,)) affected = cursor.rowcount conn.commit() conn.close() return affected > 0 # ============================== # 7.5. DOCUMENT CRUD OPERATIONS # ============================== def create_document_record(filename: str) -> dict: doc_id = str(uuid.uuid4()) upload_date = datetime.now().isoformat() status = "ready" # We process synchronously for now conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() cursor.execute( "INSERT INTO documents (id, filename, upload_date, status) VALUES (?, ?, ?, ?)", (doc_id, filename, upload_date, status) ) conn.commit() conn.close() return { "id": doc_id, "filename": filename, "upload_date": upload_date, "status": status } def get_all_documents() -> List[dict]: conn = sqlite3.connect(CHAT_DB_PATH) conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute("SELECT * FROM documents ORDER BY upload_date DESC") rows = cursor.fetchall() conn.close() return [dict(row) for row in rows] def delete_document_record(doc_id: str) -> bool: conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() cursor.execute("DELETE FROM documents WHERE id = ?", (doc_id,)) affected = cursor.rowcount conn.commit() conn.close() return affected > 0 # ============================== # 8. MESSAGE CRUD OPERATIONS # ============================== def add_message(chat_id: str, role: str, content: str) -> dict: """Add a message to a chat""" message_id = str(uuid.uuid4()) timestamp = datetime.now().isoformat() conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() cursor.execute( "INSERT INTO messages (id, chat_id, role, content, timestamp) VALUES (?, ?, ?, ?, ?)", (message_id, chat_id, role, content, timestamp) ) conn.commit() conn.close() return { "id": message_id, "chat_id": chat_id, "role": role, "content": content, "timestamp": timestamp } def get_chat_messages(chat_id: str) -> List[dict]: """Get all messages for a chat ordered by timestamp""" conn = sqlite3.connect(CHAT_DB_PATH) conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute( "SELECT * FROM messages WHERE chat_id = ? ORDER BY timestamp ASC", (chat_id,) ) rows = cursor.fetchall() conn.close() return [dict(row) for row in rows] # ============================== # 9. RAG-AWARE GENERATION # ============================== # from tools import calculate_expression, get_current_weather # from google.genai.types import Tool, GenerateContentConfig, FunctionDeclaration # ============================== # 9. RAG-AWARE GENERATION # ============================== SYSTEM_PROMPT = ( "You are a document-grounded assistant. " "Answer the user's question using only the provided document context and relevant chat history. " "Do not guess, do not invent facts, and do not add information that is not supported by the document context. " "If the answer is not clearly available in the provided context, say exactly: " "'The exact answer is not clearly available in the document.' " "When the document contains the answer, provide a complete and accurate response with all relevant details found in the context. " "Preserve important names, numbers, dates, email addresses, links, titles, and technical terms exactly as they appear in the document whenever possible." ) # SYSTEM_PROMPT = """You are a helpful assistant. # - Answer general greetings (like 'hi', 'hello') directly and briefly. # - Use `retrieve_documents` ONLY for questions about uploaded files. # - Use `calculate_expression` ONLY for math. # - Use `get_current_weather` ONLY for weather questions. # DO NOT use tools for simple conversation. # """ OPENAI_CHAT_MODEL = "gpt-4o-mini" def generate_rag_response(question: str, context_list: List[str], chat_history: List[dict]) -> str: # context_text = "\n\n---\n\n".join(context_list[:2]) if context_list else "No relevant context found." context_text = "\n\n---\n\n".join(context_list) if context_list else "No relevant context found." history_text = "" for msg in (chat_history[-3:] if len(chat_history) > 3 else chat_history): history_text += f"{msg['role'].upper()}: {msg['content']}\n" prompt = f"""Document context: {context_text} Chat history: {history_text} User question: {question} """ response = openai_client.responses.create( model=OPENAI_CHAT_MODEL, input=[ { "role": "system", "content": [ { "type": "input_text", "text": ( "You are a document-grounded assistant. " "Answer ONLY from the provided document context. " "Do not guess. Do not rewrite names, emails, project names, companies, locations, or technologies. " "If the exact answer is not clearly available in the document, say exactly: " "'The exact answer is not clearly available in the document.' " "Keep the answer short and factual." ), # "text": ( # "You are a document-grounded assistant. " # "Answer ONLY from the provided document context. " # "Do not guess. Do not rewrite names, emails, project names, companies, locations, or technologies. " # "If the exact answer is not clearly available in the document, say exactly: " # "'The exact answer is not clearly available in the document.' " # "Keep the answer short and factual." # ), } ], }, { "role": "user", "content": [ { "type": "input_text", "text": prompt, } ], }, ], ) return (response.output_text or "").strip() or "The exact answer is not clearly available in the document." # Legacy function - kept for backwards compatibility def ask_gemini(question, context_list): return generate_rag_response(question, context_list, []) # ============================== # 10. MAIN PROCESS (PDF Processing) # ============================== if __name__ == "__main__": # No hardcoded PDF loading anymore. Documents are managed via API. print("RAG system initialized. Use API endpoints to manage documents and chats.") # ============================== # 11. FASTAPI APPLICATION # ============================== from fastapi import FastAPI, HTTPException, UploadFile, File from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel app = FastAPI(title="RAG Chat API", version="2.0.0") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) BASE_DIR = os.path.dirname(os.path.abspath(__file__)) PROJECT_ROOT = os.path.abspath(os.path.join(BASE_DIR, "..")) FRONTEND_DIST = os.path.join(PROJECT_ROOT, "dist") FRONTEND_ASSETS = os.path.join(FRONTEND_DIST, "assets") if os.path.isdir(FRONTEND_ASSETS): app.mount("/assets", StaticFiles(directory=FRONTEND_ASSETS), name="assets") # ============================== # 12. REQUEST/RESPONSE MODELS # ============================== class QuestionRequest(BaseModel): question: str class CreateChatRequest(BaseModel): title: Optional[str] = "New Chat" class SendMessageRequest(BaseModel): content: str class UpdateChatRequest(BaseModel): title: str # ============================== # 13. LEGACY ENDPOINT (backwards compatibility) # ============================== @app.post("/ask") async def ask_question(req: QuestionRequest): question = req.question.strip() if not question: raise HTTPException(status_code=400, detail="Savol bo'sh bo'lishi mumkin emas") if collection.count() == 0: raise HTTPException(status_code=404, detail="No documents loaded into the system. Please upload PDFs first.") relevant_chunks = find_context(question, top_k=4) answer = generate_rag_response(question, relevant_chunks, []) return {"answer": answer} # ============================== # 14. CHAT ENDPOINTS # ============================== @app.post("/chats") async def create_new_chat(req: CreateChatRequest = CreateChatRequest()): """Create a new chat session""" chat = create_chat(req.title) return chat @app.get("/chats") async def list_chats(): """Get all chats""" chats = get_all_chats() return {"chats": chats} @app.get("/chats/{chat_id}") async def get_chat(chat_id: str): """Get a single chat by ID""" chat = get_chat_by_id(chat_id) if not chat: raise HTTPException(status_code=404, detail="Chat not found") return chat @app.patch("/chats/{chat_id}") async def update_chat(chat_id: str, req: UpdateChatRequest): """Update chat title""" success = update_chat_title(chat_id, req.title) if not success: raise HTTPException(status_code=404, detail="Chat not found") return {"message": "Chat updated successfully"} @app.delete("/chats/{chat_id}") async def remove_chat(chat_id: str): """Delete a chat and all its messages""" success = delete_chat(chat_id) if not success: raise HTTPException(status_code=404, detail="Chat not found") return {"message": "Chat deleted successfully"} # ============================== # 14.5 DOCUMENT ENDPOINTS # ============================== @app.get("/documents") async def list_documents(): return {"documents": get_all_documents()} @app.post("/documents") async def upload_document(file: UploadFile = File(...)): if not RAG_AVAILABLE: raise HTTPException(status_code=503, detail="RAG system unavailable.") if not file.filename.endswith(".pdf"): raise HTTPException(status_code=400, detail="Only PDF files are allowed") os.makedirs("data/uploads", exist_ok=True) file_path = f"data/uploads/{uuid.uuid4()}_{file.filename}" try: with open(file_path, "wb") as f: content = await file.read() f.write(content) text = load_pdf(file_path) if not text.strip(): raise HTTPException(status_code=400, detail="Could not extract text from PDF") chunks = chunk_text(text) embeddings = embed_texts(chunks) doc = create_document_record(file.filename) save_to_chroma(chunks, embeddings, doc["id"]) return doc except Exception as e: if os.path.exists(file_path): os.remove(file_path) raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}") @app.delete("/documents/{doc_id}") async def delete_document(doc_id: str): success = delete_document_record(doc_id) if not success: raise HTTPException(status_code=404, detail="Document not found") if RAG_AVAILABLE: delete_from_chroma(doc_id) return {"message": "Document deleted successfully"} # ============================== # 15. MESSAGE ENDPOINTS # ============================== @app.get("/chats/{chat_id}/messages") async def get_messages(chat_id: str): """Get all messages for a chat""" chat = get_chat_by_id(chat_id) if not chat: raise HTTPException(status_code=404, detail="Chat not found") messages = get_chat_messages(chat_id) return {"messages": messages} @app.post("/chats/{chat_id}/messages") async def send_message(chat_id: str, req: SendMessageRequest): chat = get_chat_by_id(chat_id) if not chat: raise HTTPException(status_code=404, detail="Chat not found") content = req.content.strip() if not content: raise HTTPException(status_code=400, detail="Message content cannot be empty") user_message = add_message(chat_id, "user", content) messages = get_chat_messages(chat_id) if len(messages) == 1: title = content[:50] + "..." if len(content) > 50 else content update_chat_title(chat_id, title) relevant_chunks = find_context(content, top_k=4) print("RELEVANT CHUNKS:", relevant_chunks) chat_history = messages[:-1] if len(messages) > 1 else [] lower_content = content.lower() if "email" in lower_content or "e-mail" in lower_content or "gmail" in lower_content: email = extract_email_from_context(relevant_chunks) response_text = email if email else "The exact answer is not clearly available in the document." elif "project" in lower_content: project_lines = extract_project_lines(relevant_chunks) if project_lines: response_text = "\n".join(project_lines) else: response_text = generate_rag_response(content, relevant_chunks, chat_history) else: response_text = generate_rag_response(content, relevant_chunks, chat_history) assistant_message = add_message(chat_id, "assistant", response_text) return { "user_message": user_message, "assistant_message": assistant_message } # ____________________________________________________________________ @app.get("/") async def serve_frontend(): index_file = os.path.join(FRONTEND_DIST, "index.html") if os.path.exists(index_file): return FileResponse(index_file) raise HTTPException(status_code=404, detail="Frontend build not found") @app.get("/{full_path:path}") async def serve_spa(full_path: str): api_prefixes = ( "ask", "chats", "documents", "docs", "redoc", "openapi.json", "assets", ) if full_path.startswith(api_prefixes): raise HTTPException(status_code=404, detail="Not found") index_file = os.path.join(FRONTEND_DIST, "index.html") if os.path.exists(index_file): return FileResponse(index_file) raise HTTPException(status_code=404, detail="Frontend build not found") ###################################################### # RESET CHROMA DB RESET CHATS DB REMOVE UPLOADS FILES @app.post("/admin/reset-rag") async def reset_rag(): base_dir = Path(__file__).resolve().parent chroma_dir = base_dir / "chroma_db" chats_db = base_dir / "chats.db" uploads_dir = base_dir / "data" / "uploads" try: if chroma_dir.exists(): shutil.rmtree(chroma_dir) if chats_db.exists(): chats_db.unlink() if uploads_dir.exists(): for item in uploads_dir.iterdir(): if item.is_file(): item.unlink() elif item.is_dir(): shutil.rmtree(item) return {"ok": True, "message": "RAG data reset successful"} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) # ============================== # 16. SERVER STARTUP # ============================== if __name__ == "__main__": # Auto-load logic for data/data.pdf DATA_PDF_PATH = os.path.abspath("../data/data.pdf") # Create tables if not exist init_db() if RAG_AVAILABLE and os.path.exists(DATA_PDF_PATH): print(f"Found default data file: {DATA_PDF_PATH}") try: # Check if likely already indexed (files with name 'data.pdf') # This is a basic check. conn = sqlite3.connect(CHAT_DB_PATH) cursor = conn.cursor() cursor.execute("SELECT id FROM documents WHERE filename = ?", ("data.pdf",)) existing = cursor.fetchone() conn.close() if not existing: print("Auto-loading data.pdf...") text = load_pdf(DATA_PDF_PATH) if text.strip(): chunks = chunk_text(text) embeddings = embed_texts(chunks) doc = create_document_record("data.pdf") save_to_chroma(chunks, embeddings, doc["id"]) print("Successfully auto-loaded data.pdf") else: print("Warning: data.pdf was empty") else: print("data.pdf already indexed.") except Exception as e: print(f"Failed to auto-load data.pdf: {e}") import uvicorn port = int(os.environ.get("PORT", 7860)) uvicorn.run(app, host="0.0.0.0", port=port) # uvicorn.run("server:app", host="0.0.0.0", port=4000)