File size: 2,643 Bytes
fd1c9c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
from fastapi import FastAPI, File, UploadFile, HTTPException
from pydantic import BaseModel
import fitz # PyMuPDF
import faiss
from sentence_transformers import SentenceTransformer
import numpy as np
from phi.agent import Agent
from phi.model.groq import Groq
app = FastAPI()
# Load embedding model
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
# Global storage
pdf_text_chunks = []
index = None
def agent_response(question, retrieved_text):
"""Generate response using AI model based on retrieved text."""
agent = Agent(
model=Groq(id="llama-3.3-70b-versatile"),
markdown=True,
description="You are an AI assistant that provides the answer based on the provided document.",
instructions=[
f"First read the question carefully. The question is: **{question}**",
f"Then read the document provided to you as a text. The document is: \n**{retrieved_text}**\n",
"Finally answer the question based on the provided document only. Don't try to give random responses."
]
)
response = agent.run(question + '\n' + retrieved_text).content
return response
@app.post("/upload-pdf/")
async def upload_pdf(file: UploadFile = File(...)):
"""Extract text from PDF, create FAISS index."""
global pdf_text_chunks, index
pdf_text_chunks = []
# Read the uploaded file into memory
pdf_data = await file.read()
with fitz.open("pdf", pdf_data) as doc:
for page in doc:
pdf_text_chunks.append(page.get_text("text"))
# Embed the chunks
embeddings = embedding_model.encode(pdf_text_chunks, convert_to_numpy=True)
# Create FAISS index
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(embeddings)
return {"message": "PDF processed successfully!"}
class QueryRequest(BaseModel):
question: str
@app.post("/chat/")
async def chat(request: QueryRequest):
"""Retrieve the most relevant chunk and generate a response."""
global index, pdf_text_chunks
if index is None:
raise HTTPException(status_code=400, detail="No PDF uploaded yet.")
# Search for relevant text
query_embedding = embedding_model.encode([request.question], convert_to_numpy=True)
_, indices = index.search(query_embedding, 5) # Get top 5 matches
retrieved_texts = [pdf_text_chunks[idx] for idx in indices[0]]
retrieved_text_combined = "\n\n".join(retrieved_texts)
response = agent_response(request.question, retrieved_text_combined)
return {"user": request.question, "response": response}
|