AI_Chat_With_Tools / back /server.py
muhammad1707's picture
added reset endpoint
e93cc31
Raw
History Blame Contribute Delete
26.7 kB
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)