279 lines
9.7 KiB
Python
279 lines
9.7 KiB
Python
|
import os
|
||
|
|
import shutil
|
||
|
|
import logging
|
||
|
|
from typing import List, Dict, Any
|
||
|
|
from contextlib import asynccontextmanager
|
||
|
|
|
||
|
|
from fastapi import FastAPI, UploadFile, File, HTTPException
|
||
|
|
from fastapi.responses import JSONResponse
|
||
|
|
from pydantic import BaseModel
|
||
|
|
import chromadb
|
||
|
|
from chromadb.utils import embedding_functions
|
||
|
|
from langchain_community.document_loaders import PyPDFLoader
|
||
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
|
|
from llama_cpp import Llama
|
||
|
|
import uvicorn
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# Configuration
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
MODEL_PATH = os.getenv("MODEL_PATH", "/home/orangepi/daemon/orangepi_service/models/qwen2.5-0.5b-instruct-q8_0.gguf")
|
||
|
|
N_CTX = int(os.getenv("N_CTX", "1024")) # Smaller context for Orange Pi
|
||
|
|
N_THREADS = int(os.getenv("N_THREADS", "4"))
|
||
|
|
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "rag_docs")
|
||
|
|
CHROMA_DIR = os.getenv("CHROMA_DIR", "/home/orangepi/daemon/orangepi_service/models/chroma_db")
|
||
|
|
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "/home/orangepi/daemon/orangepi_service/models/uploaded_pdfs")
|
||
|
|
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "500"))
|
||
|
|
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "50"))
|
||
|
|
TOP_K = int(os.getenv("TOP_K", "3")) # Number of chunks to retrieve
|
||
|
|
|
||
|
|
# Generation defaults
|
||
|
|
TEMP = float(os.getenv("TEMP", "0.3"))
|
||
|
|
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "256"))
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# Logging
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
logging.basicConfig(level=logging.INFO)
|
||
|
|
logger = logging.getLogger("orangepi-rag")
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# Global objects (initialised at startup)
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
llm = None
|
||
|
|
chroma_client = None
|
||
|
|
collection = None
|
||
|
|
embedding_fn = None
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# Initialisation
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
def init_directories():
|
||
|
|
os.makedirs(UPLOAD_DIR, exist_ok=True)
|
||
|
|
os.makedirs(CHROMA_DIR, exist_ok=True)
|
||
|
|
|
||
|
|
def init_chroma():
|
||
|
|
global chroma_client, collection, embedding_fn
|
||
|
|
chroma_client = chromadb.PersistentClient(path=CHROMA_DIR)
|
||
|
|
# Use sentence-transformers embedding function (runs on CPU)
|
||
|
|
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
|
||
|
|
model_name="all-MiniLM-L6-v2"
|
||
|
|
)
|
||
|
|
# Get or create collection
|
||
|
|
collection = chroma_client.get_or_create_collection(
|
||
|
|
name=COLLECTION_NAME,
|
||
|
|
embedding_function=embedding_fn
|
||
|
|
)
|
||
|
|
logger.info(f"Chroma initialised, collection '{COLLECTION_NAME}' has {collection.count()} documents")
|
||
|
|
|
||
|
|
def init_llm():
|
||
|
|
global llm
|
||
|
|
logger.info(f"Loading GGUF model from {MODEL_PATH} ...")
|
||
|
|
llm = Llama(
|
||
|
|
model_path=MODEL_PATH,
|
||
|
|
n_ctx=N_CTX,
|
||
|
|
n_threads=N_THREADS,
|
||
|
|
verbose=False,
|
||
|
|
)
|
||
|
|
logger.info("LLM loaded.")
|
||
|
|
|
||
|
|
def index_pdf(file_path: str, doc_id: str) -> int:
|
||
|
|
"""Load PDF, split, embed, and add to Chroma. Returns number of chunks."""
|
||
|
|
loader = PyPDFLoader(file_path)
|
||
|
|
documents = loader.load()
|
||
|
|
text_splitter = RecursiveCharacterTextSplitter(
|
||
|
|
chunk_size=CHUNK_SIZE,
|
||
|
|
chunk_overlap=CHUNK_OVERLAP,
|
||
|
|
separators=["\n\n", "\n", " ", ""]
|
||
|
|
)
|
||
|
|
chunks = text_splitter.split_documents(documents)
|
||
|
|
|
||
|
|
if not chunks:
|
||
|
|
logger.warning(f"No text extracted from {file_path}")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
ids = [f"{doc_id}_{i}" for i in range(len(chunks))]
|
||
|
|
texts = [chunk.page_content for chunk in chunks]
|
||
|
|
metadatas = [{"source": doc_id, "page": chunk.metadata.get("page", 0)} for chunk in chunks]
|
||
|
|
|
||
|
|
collection.add(
|
||
|
|
ids=ids,
|
||
|
|
documents=texts,
|
||
|
|
metadatas=metadatas
|
||
|
|
)
|
||
|
|
logger.info(f"Indexed {len(chunks)} chunks from {doc_id}")
|
||
|
|
return len(chunks)
|
||
|
|
|
||
|
|
def retrieve_relevant_chunks(question: str) -> List[str]:
|
||
|
|
"""Return top-K document chunks as strings."""
|
||
|
|
results = collection.query(query_texts=[question], n_results=TOP_K)
|
||
|
|
if results and results['documents']:
|
||
|
|
return results['documents'][0]
|
||
|
|
return []
|
||
|
|
|
||
|
|
def generate_answer(question: str, context_chunks: List[str]) -> str:
|
||
|
|
"""Generate answer using local LLM with retrieved context."""
|
||
|
|
if not context_chunks:
|
||
|
|
prompt = f"""Question: {question}\n\nNo relevant documents found. Answer based on your general knowledge.\nAnswer:"""
|
||
|
|
else:
|
||
|
|
context = "\n\n".join(context_chunks)
|
||
|
|
prompt = f"""Use only the following context to answer the question. If the answer is not in the context, say "I don't know".
|
||
|
|
|
||
|
|
Context:
|
||
|
|
{context}
|
||
|
|
|
||
|
|
Question: {question}
|
||
|
|
|
||
|
|
Answer:"""
|
||
|
|
output = llm.create_completion(
|
||
|
|
prompt="Hello how are you?",
|
||
|
|
max_tokens=MAX_TOKENS,
|
||
|
|
temperature=TEMP,
|
||
|
|
stop=None,
|
||
|
|
echo=False
|
||
|
|
)
|
||
|
|
print(output)
|
||
|
|
return output["choices"][0]["text"].strip()
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# FastAPI lifespan
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
@asynccontextmanager
|
||
|
|
async def lifespan(app: FastAPI):
|
||
|
|
init_directories()
|
||
|
|
init_chroma()
|
||
|
|
init_llm()
|
||
|
|
yield
|
||
|
|
# Optional: cleanup
|
||
|
|
logger.info("Shutting down OrangePi RAG service.")
|
||
|
|
|
||
|
|
app = FastAPI(
|
||
|
|
title="OrangePi RAG Service",
|
||
|
|
description="PDF upload, retrieval, and answer generation using a local GGUF model.",
|
||
|
|
version="1.0",
|
||
|
|
lifespan=lifespan,
|
||
|
|
)
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# Request/Response models
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
class QuestionRequest(BaseModel):
|
||
|
|
question: str
|
||
|
|
|
||
|
|
class AnswerResponse(BaseModel):
|
||
|
|
answer: str
|
||
|
|
retrieved_chunks: List[str] = [] # for debugging
|
||
|
|
|
||
|
|
class DocumentInfo(BaseModel):
|
||
|
|
doc_id: str
|
||
|
|
filename: str
|
||
|
|
chunks: int
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# API Endpoints
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
@app.get("/health")
|
||
|
|
async def health():
|
||
|
|
return {"status": "ok", "model_loaded": llm is not None, "docs_indexed": collection.count() if collection else 0}
|
||
|
|
|
||
|
|
@app.post("/upload")
|
||
|
|
async def upload_pdf(file: UploadFile = File(...)):
|
||
|
|
"""Upload a PDF file, store it, and index its contents."""
|
||
|
|
if not file.filename.endswith('.pdf'):
|
||
|
|
raise HTTPException(status_code=400, detail="Only PDF files are allowed")
|
||
|
|
|
||
|
|
# Generate a unique doc_id from filename + timestamp
|
||
|
|
doc_id = file.filename.replace(' ', '_').replace('/', '_')
|
||
|
|
base, ext = os.path.splitext(doc_id)
|
||
|
|
doc_id = f"{base}_{os.urandom(4).hex()}{ext}"
|
||
|
|
file_path = os.path.join(UPLOAD_DIR, doc_id)
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Save file
|
||
|
|
with open(file_path, "wb") as buffer:
|
||
|
|
shutil.copyfileobj(file.file, buffer)
|
||
|
|
|
||
|
|
# Index the PDF
|
||
|
|
num_chunks = index_pdf(file_path, doc_id)
|
||
|
|
|
||
|
|
return JSONResponse(content={
|
||
|
|
"message": "PDF uploaded and indexed",
|
||
|
|
"doc_id": doc_id,
|
||
|
|
"original_filename": file.filename,
|
||
|
|
"chunks_indexed": num_chunks
|
||
|
|
})
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception("Upload/index failed")
|
||
|
|
raise HTTPException(status_code=500, detail=str(e))
|
||
|
|
|
||
|
|
@app.get("/documents")
|
||
|
|
async def list_documents():
|
||
|
|
"""List all indexed documents with metadata."""
|
||
|
|
# Chroma doesn't store document-level metadata easily; we maintain a separate index.
|
||
|
|
# For simplicity, we'll read from the collection's metadata (if we stored source per chunk)
|
||
|
|
# Actually we stored 'source' in each chunk's metadata. Let's aggregate.
|
||
|
|
all_metas = collection.get() # returns ids, metadatas, documents
|
||
|
|
if not all_metas or not all_metas['metadatas']:
|
||
|
|
return {"documents": []}
|
||
|
|
|
||
|
|
doc_map = {}
|
||
|
|
for meta in all_metas['metadatas']:
|
||
|
|
source = meta.get('source')
|
||
|
|
if source:
|
||
|
|
if source not in doc_map:
|
||
|
|
doc_map[source] = 0
|
||
|
|
doc_map[source] += 1
|
||
|
|
docs = [{"doc_id": doc_id, "chunks": count} for doc_id, count in doc_map.items()]
|
||
|
|
return {"documents": docs}
|
||
|
|
|
||
|
|
@app.post("/answer", response_model=AnswerResponse)
|
||
|
|
async def answer_question(req: QuestionRequest):
|
||
|
|
"""RAG answer: retrieve relevant chunks and generate answer."""
|
||
|
|
if llm is None:
|
||
|
|
raise HTTPException(status_code=503, detail="LLM not loaded")
|
||
|
|
|
||
|
|
# Retrieve
|
||
|
|
chunks = retrieve_relevant_chunks(req.question)
|
||
|
|
# Generate
|
||
|
|
answer = generate_answer(req.question, chunks)
|
||
|
|
|
||
|
|
return AnswerResponse(answer=answer, retrieved_chunks=chunks)
|
||
|
|
|
||
|
|
@app.delete("/documents/{doc_id}")
|
||
|
|
async def delete_document(doc_id: str):
|
||
|
|
"""
|
||
|
|
Delete a document and its chunks from the vector database.
|
||
|
|
Optionally also delete the uploaded PDF file.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
# 1. Retrieve all chunks with metadata source == doc_id
|
||
|
|
# Chroma doesn't support direct delete by metadata, so we get all chunk ids with that source
|
||
|
|
all_chunks = collection.get(where={"source": doc_id})
|
||
|
|
if not all_chunks or not all_chunks['ids']:
|
||
|
|
raise HTTPException(status_code=404, detail=f"Document {doc_id} not found")
|
||
|
|
|
||
|
|
# 2. Delete the chunks
|
||
|
|
collection.delete(ids=all_chunks['ids'])
|
||
|
|
|
||
|
|
# 3. (Optional) Delete the physical PDF file if it exists
|
||
|
|
pdf_path = os.path.join(UPLOAD_DIR, doc_id)
|
||
|
|
if os.path.exists(pdf_path):
|
||
|
|
os.remove(pdf_path)
|
||
|
|
|
||
|
|
return JSONResponse(content={"message": f"Document {doc_id} deleted", "chunks_deleted": len(all_chunks['ids'])})
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception(f"Failed to delete document {doc_id}")
|
||
|
|
raise HTTPException(status_code=500, detail=str(e))
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
# Run
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
if __name__ == "__main__":
|
||
|
|
uvicorn.run(
|
||
|
|
"service:app",
|
||
|
|
host="0.0.0.0",
|
||
|
|
port=8080,
|
||
|
|
reload=False,
|
||
|
|
log_level="info"
|
||
|
|
)
|