File size: 2,553 Bytes
dcde7f3 4735086 923418b 4735086 dcde7f3 0bfabd3 4735086 dcde7f3 0bfabd3 dcde7f3 0bfabd3 dcde7f3 0bfabd3 dcde7f3 0bfabd3 4735086 0bfabd3 4735086 dcde7f3 |
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 74 75 76 77 78 79 80 81 82 83 |
from fastapi import APIRouter, Depends, HTTPException, status, File, UploadFile
from sqlalchemy.ext.asyncio import AsyncSession
from app.models import User
from app.api.deps import get_db, get_current_user
from app.schema import ChatMessage, AI_chat_input, pdf_input
from app.llm import stream_chat
import uuid
from fastapi.responses import StreamingResponse
from chromadb.api.models.Collection import Collection
from app.api.deps import get_chroma_collection
from app.api.deps import get_db, get_current_user, get_chroma_client
from pathlib import Path
from llama_index.readers.file import PyMuPDFReader
from llama_index.core.node_parser import SentenceSplitter
from typing import Annotated
import shutil
import os
from .quiz import ingest_logic
router = APIRouter(prefix="/notes")
UPLOAD_DIRECTORY = "uploaded_pdfs"
os.makedirs(UPLOAD_DIRECTORY, exist_ok=True)
@router.post("/stream_chat", response_class=StreamingResponse)
async def ai_chat(
Input_model: AI_chat_input,
# db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user)
):
messages_dict = [msg.model_dump() for msg in Input_model.messages]
return StreamingResponse(
stream_chat(messages_dict, Input_model.context),
media_type="text/plain"
)
@router.post("/upload_notes")
async def upload_notes(
file: Annotated[UploadFile, File(description="A PDF file to upload")],
collection: Collection = Depends(get_chroma_collection),
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user)
):
file_path = Path(UPLOAD_DIRECTORY) / file.filename
try:
chunks = await pdf_process(str(file_path))
if not chunks:
raise ValueError("No chunks availible")
await ingest_logic(chunks, collection)
return {"status": "success"}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error processing PDF: {str(e)}")
finally:
if file_path.exists():
os.remove(file_path)
# #--------Helper Functions--------#
async def pdf_process(pdf_path: str):
loader = PyMuPDFReader()
# 5. Load using the file path string
documents = loader.load_data(file_path=pdf_path)
text_splitter = SentenceSplitter(
chunk_size=1000,
chunk_overlap=20
)
text_chunks = []
for doc_idx, doc in enumerate(documents):
cur_text_chunks = text_splitter.split_text(doc.text)
text_chunks.extend(cur_text_chunks)
return text_chunks |