First commit
This commit is contained in:
aliakbar
@@ -0,0 +1,278 @@
|
||||
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"
|
||||
)
|
||||
Reference in New Issue
Block a user