import asyncio
import hashlib
import io
import json
import os
import re
import uuid
from datetime import datetime
from typing import Optional, List
from contextlib import asynccontextmanager
from fastapi import FastAPI, UploadFile, File, Form, WebSocket, WebSocketDisconnect, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from sqlmodel import Field, SQLModel, create_engine, Session, select
import chromadb
from chromadb.config import Settings
import tiktoken
import openai
import pypdf
# Configuration
UPLOAD_DIR = "uploads"
DB_URL = "sqlite:///./database.db"
CHROMA_DIR = "./chroma_db"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "your-api-key-here")
CHUNK_SIZE = 500
CHUNK_OVERLAP = 50
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(CHROMA_DIR, exist_ok=True)
# Database Models
class Document(SQLModel, table=True):
id: Optional[str] = Field(default=None, primary_key=True)
name: str
filename: str
markdown_content: str = Field(default="")
upload_time: datetime
status: str = "pending"
downloads: int = 0
class SearchResult(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
query: str
slug: str = Field(index=True)
results_json: str
created_at: datetime
tokens_used: int = 0
cost_eur: float = 0.0
class PromptResult(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
query: str
slug: str = Field(index=True)
prompt_response: str
search_results_json: str
created_at: datetime
input_tokens: int = 0
output_tokens: int = 0
cost_eur: float = 0.0
# Database Setup
engine = create_engine(DB_URL, connect_args={"check_same_thread": False})
def init_db():
SQLModel.metadata.create_all(engine)
# ChromaDB Setup
chroma_client = chromadb.PersistentClient(path=CHROMA_DIR, settings=Settings(anonymized_telemetry=False))
try:
collection = chroma_client.get_collection("documents")
except:
collection = chroma_client.create_collection("documents")
# OpenAI Setup
openai.api_key = OPENAI_API_KEY
tokenizer = tiktoken.get_encoding("cl100k_base")
# WebSocket Manager
class ConnectionManager:
def __init__(self):
self.active_connections: dict[str, WebSocket] = {}
async def connect(self, document_id: str, websocket: WebSocket):
await websocket.accept()
self.active_connections[document_id] = websocket
def disconnect(self, document_id: str):
if document_id in self.active_connections:
del self.active_connections[document_id]
async def send_message(self, document_id: str, message: dict):
if document_id in self.active_connections:
try:
await self.active_connections[document_id].send_json(message)
except:
self.disconnect(document_id)
manager = ConnectionManager()
# Lifespan
@asynccontextmanager
async def lifespan(app: FastAPI):
init_db()
yield
# FastAPI App
app = FastAPI(lifespan=lifespan)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Helper Functions
def generate_slug(text: str) -> str:
"""Generate URL-friendly slug from text"""
slug = re.sub(r'[^\w\s-]', '', text.lower())
slug = re.sub(r'[-\s]+', '-', slug)
return slug[:50] + "-" + hashlib.md5(text.encode()).hexdigest()[:8]
def convert_to_markdown(filepath: str) -> str:
"""Convert various document formats to markdown"""
ext = os.path.splitext(filepath)[1].lower()
if ext == '.pdf':
markdown_text = ""
with open(filepath, 'rb') as f:
pdf_reader = pypdf.PdfReader(f)
for page_num, page in enumerate(pdf_reader.pages):
text = page.extract_text()
markdown_text += f"\n\n## Page {page_num + 1}\n\n{text}"
return markdown_text
elif ext == '.md':
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
return f.read()
else:
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
content = f.read()
lines = content.split('\n')
markdown_lines = []
for line in lines:
if line.strip():
markdown_lines.append(line)
else:
markdown_lines.append('')
return '\n'.join(markdown_lines)
def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[dict]:
tokens = tokenizer.encode(text)
chunks = []
start = 0
# Extract page numbers from markdown headers
page_pattern = re.compile(r'##\s*Page\s+(\d+)', re.IGNORECASE)
while start < len(tokens):
end = start + chunk_size
chunk_tokens = tokens[start:end]
chunk_text = tokenizer.decode(chunk_tokens)
# Find page number in chunk
page_match = page_pattern.search(chunk_text)
page_num = int(page_match.group(1)) if page_match else 1
chunks.append({
"text": chunk_text,
"page": page_num
})
start = end - overlap
return chunks
async def get_embedding(text: str) -> tuple[List[float], int]:
response = await asyncio.to_thread(
openai.embeddings.create,
input=text,
model="text-embedding-3-small"
)
tokens_used = response.usage.total_tokens
return response.data[0].embedding, tokens_used
async def detect_prompt_intent(query: str) -> bool:
"""Detect if query is a prompt vs simple search"""
prompt_keywords = ['make', 'create', 'list', 'summarize', 'explain', 'compare', 'analyze', 'generate', 'write', 'show me', 'give me', 'find all', 'extract', 'what are', 'how many']
query_lower = query.lower()
# Check for question words or action verbs
if any(keyword in query_lower for keyword in prompt_keywords):
return True
if query_lower.endswith('?') and len(query.split()) > 3:
return True
return False
async def execute_prompt(query: str, search_results: List[dict]) -> tuple[str, int, int, float]:
"""Execute prompt using GPT-4 with search results as context"""
# Prepare context from search results
context = "\n\n".join([
f"Document: {r['name']}\nPage: {r.get('page', 'N/A')}\nContent: {r['snippet']}"
for r in search_results[:10] # Use top 10 results
])
messages = [
{"role": "system", "content": "You are a helpful assistant that answers questions based on the provided document context. Be concise and accurate."},
{"role": "user", "content": f"Context from documents:\n\n{context}\n\nUser query: {query}\n\nPlease answer based on the context provided."}
]
response = await asyncio.to_thread(
openai.chat.completions.create,
model="gpt-4o-mini",
messages=messages,
temperature=0.7,
max_tokens=1000
)
input_tokens = response.usage.prompt_tokens
output_tokens = response.usage.completion_tokens
# Calculate cost: GPT-4o-mini pricing
# $0.150 per 1M input tokens, $0.600 per 1M output tokens
cost_usd = (input_tokens / 1_000_000 * 0.150) + (output_tokens / 1_000_000 * 0.600)
cost_eur = cost_usd * 0.92
return response.choices[0].message.content, input_tokens, output_tokens, cost_eur
async def process_document(document_id: str, filepath: str):
total_tokens = 0
await manager.send_message(document_id, {
"step": "reading",
"progress": 5,
"message": "Reading file...",
"tokens": 0,
"cost_eur": 0.0
})
try:
await manager.send_message(document_id, {
"step": "converting",
"progress": 15,
"message": "Converting to markdown...",
"tokens": 0,
"cost_eur": 0.0
})
content = await asyncio.to_thread(convert_to_markdown, filepath)
with Session(engine) as session:
doc = session.get(Document, document_id)
if doc:
doc.markdown_content = content
session.add(doc)
session.commit()
await manager.send_message(document_id, {
"step": "chunking",
"progress": 25,
"message": "Splitting text into chunks...",
"tokens": 0,
"cost_eur": 0.0
})
chunks = chunk_text(content)
await manager.send_message(document_id, {
"step": "embedding",
"progress": 35,
"message": f"Processing {len(chunks)} chunks...",
"tokens": 0,
"cost_eur": 0.0
})
for i, chunk_data in enumerate(chunks):
embedding, tokens = await get_embedding(chunk_data["text"])
total_tokens += tokens
cost_usd = (total_tokens / 1000) * 0.00002
cost_eur = cost_usd * 0.92
collection.add(
ids=[f"{document_id}_chunk_{i}"],
embeddings=[embedding],
documents=[chunk_data["text"]],
metadatas=[{
"document_id": document_id,
"chunk_index": i,
"total_chunks": len(chunks),
"page": chunk_data["page"]
}]
)
progress = 35 + int((i + 1) / len(chunks) * 55)
await manager.send_message(document_id, {
"step": "embedding",
"progress": progress,
"message": f"Embedded chunk {i + 1}/{len(chunks)}",
"tokens": total_tokens,
"cost_eur": cost_eur
})
await manager.send_message(document_id, {
"step": "indexing",
"progress": 95,
"message": "Finalizing index...",
"tokens": total_tokens,
"cost_eur": (total_tokens / 1000) * 0.00002 * 0.92
})
with Session(engine) as session:
doc = session.get(Document, document_id)
if doc:
doc.status = "completed"
session.add(doc)
session.commit()
await manager.send_message(document_id, {
"step": "completed",
"progress": 100,
"message": "Processing complete!",
"tokens": total_tokens,
"cost_eur": (total_tokens / 1000) * 0.00002 * 0.92
})
except Exception as e:
with Session(engine) as session:
doc = session.get(Document, document_id)
if doc:
doc.status = "failed"
session.add(doc)
session.commit()
await manager.send_message(document_id, {
"step": "error",
"progress": 0,
"message": f"Error: {str(e)}",
"tokens": total_tokens,
"cost_eur": 0.0
})
# API Endpoints
@app.post("/api/upload")
async def upload_documents(
background_tasks: BackgroundTasks,
files: List[UploadFile] = File(...)
):
uploaded_docs = []
for file in files:
doc_id = str(uuid.uuid4())
file_ext = os.path.splitext(file.filename)[1]
filename = f"{doc_id}{file_ext}"
filepath = os.path.join(UPLOAD_DIR, filename)
content = await file.read()
with open(filepath, 'wb') as f:
f.write(content)
document = Document(
id=doc_id,
name=file.filename,
filename=filename,
upload_time=datetime.now(),
status="processing"
)
with Session(engine) as session:
session.add(document)
session.commit()
background_tasks.add_task(process_document, doc_id, filepath)
uploaded_docs.append({
"document_id": doc_id,
"name": file.filename,
"status": "processing"
})
return {"documents": uploaded_docs}
@app.get("/api/search")
async def search_documents(query: str, page: int = 1, page_size: int = 10):
try:
slug = generate_slug(query)
# Check cache first
with Session(engine) as session:
cached = session.exec(select(SearchResult).where(SearchResult.slug == slug)).first()
if cached:
results_data = json.loads(cached.results_json)
return {
"results": results_data["results"],
"total": results_data["total"],
"page": page,
"page_size": page_size,
"slug": slug,
"tokens": cached.tokens_used,
"cost_eur": cached.cost_eur,
"cached": True
}
collection_count = collection.count()
if collection_count == 0:
return {
"results": [],
"total": 0,
"page": page,
"page_size": page_size,
"slug": slug,
"tokens": 0,
"cost_eur": 0.0
}
query_embedding, tokens = await get_embedding(query)
cost_eur = (tokens / 1000) * 0.00002 * 0.92
results = collection.query(
query_embeddings=[query_embedding],
n_results=min(50, collection_count)
)
search_results = []
if results['ids'] and len(results['ids'][0]) > 0:
for i in range(len(results['ids'][0])):
document_id = results['metadatas'][0][i]['document_id']
snippet = results['documents'][0][i]
distance = results['distances'][0][i] if 'distances' in results else 0
page_num = results['metadatas'][0][i].get('page', 1)
with Session(engine) as session:
doc = session.get(Document, document_id)
if doc:
search_results.append({
"document_id": document_id,
"name": doc.name,
"snippet": snippet[:300] + "..." if len(snippet) > 300 else snippet,
"chunk_index": results['metadatas'][0][i]['chunk_index'],
"page": page_num,
"upload_date": doc.upload_time.isoformat(),
"score": 1 - distance
})
# Cache results
with Session(engine) as session:
search_cache = SearchResult(
query=query,
slug=slug,
results_json=json.dumps({"results": search_results, "total": len(search_results)}),
created_at=datetime.now(),
tokens_used=tokens,
cost_eur=cost_eur
)
session.add(search_cache)
session.commit()
start_idx = (page - 1) * page_size
end_idx = page * page_size
return {
"results": search_results[start_idx:end_idx],
"total": len(search_results),
"page": page,
"page_size": page_size,
"slug": slug,
"tokens": tokens,
"cost_eur": cost_eur
}
except Exception as e:
print(f"Search error: {str(e)}")
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/search/{slug}")
async def get_cached_search(slug: str):
with Session(engine) as session:
cached = session.exec(select(SearchResult).where(SearchResult.slug == slug)).first()
if not cached:
raise HTTPException(status_code=404, detail="Search not found")
results_data = json.loads(cached.results_json)
return {
"query": cached.query,
"results": results_data["results"],
"total": results_data["total"],
"slug": slug,
"tokens": cached.tokens_used,
"cost_eur": cached.cost_eur
}
@app.post("/api/prompt")
async def execute_prompt_endpoint(query: str):
try:
slug = generate_slug(query)
# Check cache
with Session(engine) as session:
cached = session.exec(select(PromptResult).where(PromptResult.slug == slug)).first()
if cached:
return {
"response": cached.prompt_response,
"search_results": json.loads(cached.search_results_json),
"slug": slug,
"input_tokens": cached.input_tokens,
"output_tokens": cached.output_tokens,
"cost_eur": cached.cost_eur,
"cached": True
}
# First get search results
search_response = await search_documents(query, page=1, page_size=20)
# Execute prompt with results
response_text, input_tokens, output_tokens, cost_eur = await execute_prompt(
query,
search_response["results"]
)
# Add search cost to total
total_cost = cost_eur + search_response.get("cost_eur", 0)
# Cache prompt result
with Session(engine) as session:
prompt_cache = PromptResult(
query=query,
slug=slug,
prompt_response=response_text,
search_results_json=json.dumps(search_response["results"]),
created_at=datetime.now(),
input_tokens=input_tokens,
output_tokens=output_tokens,
cost_eur=total_cost
)
session.add(prompt_cache)
session.commit()
return {
"response": response_text,
"search_results": search_response["results"],
"slug": slug,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"cost_eur": total_cost
}
except Exception as e:
print(f"Prompt error: {str(e)}")
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/prompt/{slug}")
async def get_cached_prompt(slug: str):
with Session(engine) as session:
cached = session.exec(select(PromptResult).where(PromptResult.slug == slug)).first()
if not cached:
raise HTTPException(status_code=404, detail="Prompt result not found")
return {
"query": cached.query,
"response": cached.prompt_response,
"search_results": json.loads(cached.search_results_json),
"slug": slug,
"input_tokens": cached.input_tokens,
"output_tokens": cached.output_tokens,
"cost_eur": cached.cost_eur
}
@app.get("/api/documents")
async def list_documents():
with Session(engine) as session:
documents = session.exec(select(Document)).all()
return [{
"id": doc.id,
"name": doc.name,
"upload_time": doc.upload_time.isoformat(),
"status": doc.status,
"downloads": doc.downloads
} for doc in documents]
@app.get("/api/document/{document_id}")
async def get_document(document_id: str):
with Session(engine) as session:
doc = session.get(Document, document_id)
if not doc:
raise HTTPException(status_code=404, detail="Document not found")
return {
"id": doc.id,
"name": doc.name,
"markdown_content": doc.markdown_content,
"upload_time": doc.upload_time.isoformat(),
"status": doc.status
}
@app.get("/api/download/{document_id}")
async def download_document(document_id: str):
with Session(engine) as session:
doc = session.get(Document, document_id)
if not doc:
raise HTTPException(status_code=404, detail="Document not found")
doc.downloads += 1
session.add(doc)
session.commit()
filepath = os.path.join(UPLOAD_DIR, doc.filename)
if not os.path.exists(filepath):
raise HTTPException(status_code=404, detail="File not found")
return FileResponse(filepath, filename=doc.name)
@app.websocket("/ws/status/{document_id}")
async def websocket_endpoint(websocket: WebSocket, document_id: str):
await manager.connect(document_id, websocket)
try:
while True:
await websocket.receive_text()
except WebSocketDisconnect:
manager.disconnect(document_id)
@app.get("/")
async def read_root():
return FileResponse("index.html")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=9900)
Reference in New Issue View Git Blame Copy Permalink