Spaces:
Sleeping
Sleeping
| 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) | |
| # ============================== | |
| 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 | |
| # ============================== | |
| async def create_new_chat(req: CreateChatRequest = CreateChatRequest()): | |
| """Create a new chat session""" | |
| chat = create_chat(req.title) | |
| return chat | |
| async def list_chats(): | |
| """Get all chats""" | |
| chats = get_all_chats() | |
| return {"chats": chats} | |
| 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 | |
| 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"} | |
| 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 | |
| # ============================== | |
| async def list_documents(): | |
| return {"documents": get_all_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)}") | |
| 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 | |
| # ============================== | |
| 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} | |
| 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 | |
| } | |
| # ____________________________________________________________________ | |
| 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") | |
| 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 | |
| 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) | |