diff --git a/.vscode/settings.json b/.vscode/settings.json index 419eb30..eb87dcb 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,12 +1,17 @@ -{ - "python.defaultInterpreterPath": "${workspaceFolder}/web/backend/.venv/Scripts/python.exe", - "python.analysis.extraPaths": [ - "${workspaceFolder}/web/backend", - "${workspaceFolder}/web/backend/routes" - ], - "python.autoComplete.extraPaths": [ - "${workspaceFolder}/web/backend", - "${workspaceFolder}/web/backend/routes" - ], - "python.analysis.autoSearchPaths": true -} + { + "python.defaultInterpreterPath": "${workspaceFolder}/web/backend/.venv/Scripts/python.exe", + + "python.analysis.extraPaths": [ + "${workspaceFolder}/web/backend", + "${workspaceFolder}/web/backend/routes", + "${workspaceFolder}/ai-ml" + ], + + "python.autoComplete.extraPaths": [ + "${workspaceFolder}/web/backend", + "${workspaceFolder}/web/backend/routes", + "${workspaceFolder}/ai-ml" + ], + + "python.analysis.autoSearchPaths": true +} \ No newline at end of file diff --git a/ai-ml/ingestion/common/exceptions.py b/ai-ml/ingestion/common/exceptions.py new file mode 100644 index 0000000..0cd4c9f --- /dev/null +++ b/ai-ml/ingestion/common/exceptions.py @@ -0,0 +1,11 @@ +class IngestionError(Exception): + """Base class for all ingestion errors.""" + pass + +class PDFProcessingError(IngestionError): + """Raised when PDF extraction or OCR fails.""" + pass + +class YouTubeTranscriptError(IngestionError): + """Raised when YouTube transcripts are unavailable.""" + pass \ No newline at end of file diff --git a/ai-ml/ingestion/main.py b/ai-ml/ingestion/main.py index 0ffb87e..6cf672e 100644 --- a/ai-ml/ingestion/main.py +++ b/ai-ml/ingestion/main.py @@ -72,8 +72,9 @@ def _chunk_and_store(result: dict, user_id: str) -> dict: async def ingest_pdf_endpoint( file: UploadFile = File(...), user_id: str = Depends(get_current_user_id), + # user_id="test_user", ): - if not file.filename.lower().endswith(".pdf"): + if not file.filename or not file.filename.lower().endswith(".pdf"): raise HTTPException(status_code=400, detail="File must be a PDF") with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: @@ -81,7 +82,10 @@ async def ingest_pdf_endpoint( tmp_path = tmp.name try: - result = ingest_pdf(file_path=tmp_path, original_filename=file.filename) + result = ingest_pdf( + file_path=tmp_path, + original_filename=file.filename + ) storage_info = _chunk_and_store(result, user_id) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @@ -90,21 +94,24 @@ async def ingest_pdf_endpoint( os.remove(tmp_path) return {**result, **storage_info} - - # ----------------------------- # YOUTUBE INGESTION # ----------------------------- @app.post("/ingest/youtube") async def ingest_youtube_endpoint( payload: URLRequest, - user_id: str = Depends(get_current_user_id), + user_id: str = Depends(get_current_user_id), + # user_id = "test_user", ): try: result = ingest_youtube(payload.url) storage_info = _chunk_and_store(result, user_id) except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + # Task 3: If it's our custom error, give a 400. Otherwise, give a 500. + error_msg = str(e) + status_code = 400 if "YouTube" in error_msg else 500 + raise HTTPException(status_code=status_code, detail=error_msg) + return {**result, **storage_info} diff --git a/ai-ml/ingestion/pdf/extractor.py b/ai-ml/ingestion/pdf/extractor.py index 3674e51..59d6a84 100644 --- a/ai-ml/ingestion/pdf/extractor.py +++ b/ai-ml/ingestion/pdf/extractor.py @@ -1,74 +1,46 @@ import os -import fitz # PyMuPDF - +from docling.datamodel.base_models import InputFormat +from docling.document_converter import DocumentConverter from ingestion.pdf.cleaner import clean_pdf_text from ingestion.common.schema import build_result - +from ingestion.common.exceptions import PDFProcessingError def extract_pdf_text(file_path: str) -> str: """ - Open a PDF and extract raw text from every page. + Advanced extraction using IBM's Docling. + Handles OCR, tables, and layout automatically. """ - if not os.path.exists(file_path): - raise FileNotFoundError(f"PDF not found: {file_path}") - - doc = fitz.open(file_path) - - pages_text = [] - - for page in doc: - pages_text.append(page.get_text()) - - doc.close() - - return "\n".join(pages_text) - + raise PDFProcessingError(f"File not found: {file_path}") + + try: + # Initialize the converter + converter = DocumentConverter() + + # Convert the PDF (Docling automatically detects if it needs OCR) + result = converter.convert(file_path) + + # Export to Markdown (best for RAG) or Plain Text + # .export_to_markdown() is highly recommended for LLMs + extracted_text = result.document.export_to_markdown() + + if not extracted_text.strip(): + raise PDFProcessingError("Extraction resulted in empty text.") + + return extracted_text + + except Exception as e: + raise PDFProcessingError(f"Docling failed to process PDF: {str(e)}") def ingest_pdf(file_path: str, original_filename: str) -> dict: - """ - Full PDF ingestion pipeline: - - PDF file - ↓ - Text extraction - ↓ - Text cleaning - ↓ - Build standard output format - """ - - # Step 1: Extract text + """Standard pipeline logic.""" raw_text = extract_pdf_text(file_path) - - # Step 2: Clean extracted text cleaned_text = clean_pdf_text(raw_text) - - # Step 3: Generate user-friendly title title = os.path.splitext(original_filename)[0] - # Step 4: Return common JSON format return build_result( source_type="pdf", title=title, text=cleaned_text, source=file_path, - ) - - -if __name__ == "__main__": - import sys - import json - - if len(sys.argv) < 2: - print("Usage: python extractor.py ") - - else: - pdf_path = sys.argv[1] - - result = ingest_pdf( - file_path=pdf_path, - original_filename=os.path.basename(pdf_path) - ) - - print(json.dumps(result, indent=2)) \ No newline at end of file + ) \ No newline at end of file diff --git a/ai-ml/ingestion/youtube/exceptions.py b/ai-ml/ingestion/youtube/exceptions.py new file mode 100644 index 0000000..e5369af --- /dev/null +++ b/ai-ml/ingestion/youtube/exceptions.py @@ -0,0 +1,70 @@ +"""Structured errors for the YouTube ingestion pipeline. + +Every error carries a machine-readable `code`, an HTTP `status_code` the +API layer should return, and a human-readable `message`. This lets the +FastAPI layer turn these into clean JSON instead of a raw 500 traceback. +""" + + +class YouTubeIngestError(Exception): + """Base class for all YouTube ingestion errors.""" + + code = "youtube_ingest_error" + status_code = 500 + + def __init__( + self, + message: str, + video_id: str | None = None, + details: dict | None = None, + ): + super().__init__(message) + self.message = message + self.video_id = video_id + self.details = details or {} + + def to_dict(self) -> dict: + payload: dict[str, object] = { + "error": self.code, + "message": self.message, + } + + if self.video_id: + payload["video_id"] = self.video_id + + if self.details: + payload["details"] = self.details + + return payload + +class InvalidYouTubeURLError(YouTubeIngestError): + """The URL is not a recognizable YouTube video URL.""" + + code = "invalid_url" + status_code = 400 + + +class VideoUnavailableError(YouTubeIngestError): + """The video is private, deleted, region-locked, or otherwise unreachable.""" + + code = "video_unavailable" + status_code = 404 + + +class TranscriptNotAvailableError(YouTubeIngestError): + """The video exists but no usable transcript could be produced. + + This covers: transcripts disabled by the uploader, no transcript in any + language, a transcript that fetched but came back empty, and transcript + fetch failures after a listing succeeded. + """ + + code = "transcript_not_available" + status_code = 422 + + +class TranscriptFetchError(YouTubeIngestError): + """Transient/upstream failure (rate limiting, IP block, network, etc.).""" + + code = "transcript_fetch_failed" + status_code = 502 \ No newline at end of file diff --git a/ai-ml/ingestion/youtube/transcript.py b/ai-ml/ingestion/youtube/transcript.py index d3fee0f..1b596b2 100644 --- a/ai-ml/ingestion/youtube/transcript.py +++ b/ai-ml/ingestion/youtube/transcript.py @@ -1,16 +1,25 @@ import re +from typing import Any from urllib.parse import urlparse, parse_qs import yt_dlp +from yt_dlp.utils import DownloadError from youtube_transcript_api import YouTubeTranscriptApi from youtube_transcript_api._errors import ( TranscriptsDisabled, NoTranscriptFound, VideoUnavailable, + CouldNotRetrieveTranscript, ) from ingestion.youtube.cleaner import clean_youtube_text +from ingestion.youtube.exceptions import ( + InvalidYouTubeURLError, + VideoUnavailableError, + TranscriptNotAvailableError, + TranscriptFetchError, +) from ingestion.common.schema import build_result @@ -31,7 +40,7 @@ def extract_video_id(url: str) -> str: if match: return match.group(2) - raise ValueError(f"Could not extract video ID from URL: {url}") + raise InvalidYouTubeURLError(f"Could not extract a video ID from URL: {url}") def fetch_metadata(url: str) -> dict: @@ -41,8 +50,15 @@ def fetch_metadata(url: str) -> dict: "skip_download": True, } - with yt_dlp.YoutubeDL(options) as ydl: - info = ydl.extract_info(url, download=False) + try: + with yt_dlp.YoutubeDL(options) as ydl: # type: ignore[arg-type] + info = ydl.extract_info(url, download=False) + except DownloadError as e: + raise VideoUnavailableError( + f"Could not load video metadata for {url}. It may be private, " + f"deleted, region-locked, or age-restricted.", + details={"reason": str(e)}, + ) from e return { "title": info.get("title", ""), @@ -51,28 +67,91 @@ def fetch_metadata(url: str) -> dict: "date": info.get("upload_date"), } +def fetch_transcript(video_id: str, languages=("en",)) -> dict: + """ + Returns {"text": str, "language_code": str, "is_generated": bool}. -def fetch_transcript(video_id: str, languages=("en",)) -> str: - try: - api = YouTubeTranscriptApi() + Raises: + VideoUnavailableError: video is private/deleted/unreachable. + TranscriptNotAvailableError: video exists but has no usable transcript + (disabled by uploader, none in any language, or empty once fetched). + TranscriptFetchError: transient upstream failure (rate limit, IP block). + """ + api = YouTubeTranscriptApi() - transcript = api.fetch( - video_id, - languages=list(languages) + # Step 1: find out what transcripts actually exist for this video. + try: + transcript_list = api.list(video_id) + except VideoUnavailable as e: + raise VideoUnavailableError( + f"The video {video_id} is unavailable, private, or has been removed.", + video_id=video_id, + ) from e + except TranscriptsDisabled as e: + raise TranscriptNotAvailableError( + f"The uploader has disabled transcripts/captions for video {video_id}.", + video_id=video_id, + details={"reason": "transcripts_disabled"}, + ) from e + except CouldNotRetrieveTranscript as e: + # Covers RequestBlocked/IpBlocked/PoTokenRequired/YouTubeRequestFailed etc. + raise TranscriptFetchError( + f"Could not check transcript availability for video {video_id}: {e}", + video_id=video_id, + ) from e + + available = [ + { + "language": t.language, + "language_code": t.language_code, + "is_generated": t.is_generated, + } + for t in transcript_list + ] + + if not available: + raise TranscriptNotAvailableError( + f"No transcript (manual or auto-generated) exists for video {video_id} " + f"in any language.", + video_id=video_id, + details={"available_languages": []}, ) - except (TranscriptsDisabled, NoTranscriptFound): - api = YouTubeTranscriptApi() - transcript = api.fetch(video_id) - - except VideoUnavailable as e: - raise RuntimeError(f"Video unavailable: {video_id}") from e + # Step 2: prefer a manually created transcript in a requested language, + # then an auto-generated one in a requested language, then just take + # whatever exists (manually created first) rather than failing outright. + transcript = None + try: + transcript = transcript_list.find_transcript(list(languages)) + except NoTranscriptFound: + transcript = sorted(transcript_list, key=lambda t: t.is_generated)[0] - merged = " ".join( - segment.text for segment in transcript - ) + # Step 3: actually fetch it. + try: + fetched = transcript.fetch() + except CouldNotRetrieveTranscript as e: + raise TranscriptFetchError( + f"Found a transcript listing for video {video_id} " + f"(language={transcript.language_code}) but failed to fetch it: {e}", + video_id=video_id, + details={"available_languages": available}, + ) from e + + merged = " ".join(segment.text for segment in fetched).strip() + + if not merged: + raise TranscriptNotAvailableError( + f"Transcript for video {video_id} (language={transcript.language_code}) " + f"fetched successfully but contained no text.", + video_id=video_id, + details={"available_languages": available}, + ) - return merged + return { + "text": merged, + "language_code": transcript.language_code, + "is_generated": transcript.is_generated, + } def ingest_youtube(url: str) -> dict: @@ -80,9 +159,9 @@ def ingest_youtube(url: str) -> dict: metadata = fetch_metadata(url) - raw_text = fetch_transcript(video_id) + transcript_data = fetch_transcript(video_id) - cleaned_text = clean_youtube_text(raw_text) + cleaned_text = clean_youtube_text(transcript_data["text"]) result = build_result( source_type="youtube", @@ -95,6 +174,8 @@ def ingest_youtube(url: str) -> dict: "author": metadata["author"], "duration": metadata["duration"], "date": metadata["date"], + "transcript_language": transcript_data["language_code"], + "transcript_auto_generated": transcript_data["is_generated"], }) return result @@ -104,8 +185,14 @@ def ingest_youtube(url: str) -> dict: import sys import json + from ingestion.youtube.exceptions import YouTubeIngestError + if len(sys.argv) < 2: print("Usage: python transcript.py ") else: - result = ingest_youtube(sys.argv[1]) - print(json.dumps(result, indent=2)) \ No newline at end of file + try: + result = ingest_youtube(sys.argv[1]) + print(json.dumps(result, indent=2)) + except YouTubeIngestError as e: + print(json.dumps(e.to_dict(), indent=2)) + sys.exit(1) \ No newline at end of file diff --git a/ai-ml/requirements.txt b/ai-ml/requirements.txt index bc380bd..37b8356 100644 --- a/ai-ml/requirements.txt +++ b/ai-ml/requirements.txt @@ -18,4 +18,7 @@ python-dotenv groq yake pytest -pyjwt \ No newline at end of file +PyJWT +docling +pytesseract +Pillow \ No newline at end of file diff --git a/ai-ml/tests/test_chunker.py b/ai-ml/tests/test_chunker.py new file mode 100644 index 0000000..01137e4 --- /dev/null +++ b/ai-ml/tests/test_chunker.py @@ -0,0 +1,24 @@ +import pytest +from embedding.chunker import chunk_document + +def test_chunk_document_happy_path(): + """Verify document is split into chunks with correct metadata.""" + doc = { + "source_type": "pdf", + "title": "Lab Report", + "text": "word1 " * 20, # 20 words + "metadata": {"author": "Farwa"} + } + + # Testing with small chunk size to force multiple chunks + chunks = chunk_document(doc, chunk_size=10, overlap=2) + + assert len(chunks) > 1 + assert chunks[0]["title"] == "Lab Report" + assert "chunk_index" in chunks[0] + assert chunks[0]["metadata"]["author"] == "Farwa" + +def test_chunk_document_empty_text(): + """Edge Case: Empty text should return empty list.""" + doc = {"text": "", "title": "Empty"} + assert chunk_document(doc) == [] \ No newline at end of file diff --git a/ai-ml/tests/test_embedder.py b/ai-ml/tests/test_embedder.py new file mode 100644 index 0000000..ca1df30 --- /dev/null +++ b/ai-ml/tests/test_embedder.py @@ -0,0 +1,45 @@ +import pytest +from unittest.mock import MagicMock, patch +from embedding.chroma_store import store_chunks + +def test_store_chunks_happy_path(): + """Task 4: Happy Path - Mock everything to avoid disk and AI model delays.""" + + # 1. Provide a chunk that matches your Chunker.py output exactly + mock_chunks = [{ + "text": "sample text", + "chunk_index": 0 + }] + + # 2. We mock the Client and the Model + # This prevents the code from touching C:\Dev\... and from loading AI weights + with patch("embedding.chroma_store.chromadb.PersistentClient") as mock_client_class, \ + patch("embedding.chroma_store.get_embedding_model") as mock_get_model: + + # Setup the database mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_col = MagicMock() + mock_client.get_or_create_collection.return_value = mock_col + + # Setup the model mock + mock_model = MagicMock() + mock_get_model.return_value = mock_model + mock_model.encode.return_value = [[0.1, 0.2]] + + # 3. Call the function + result = store_chunks(mock_chunks, "user_1", "doc_1", "Test Title") + + # 4. Verify + assert result == 1 + # Your code uses upsert, so we check that + assert mock_col.upsert.called + + # Verify the ID was created correctly: {doc_id}_chunk{index} + args, kwargs = mock_col.upsert.call_args + assert kwargs['ids'][0] == "doc_1_chunk0" + +def test_store_chunks_empty_list(): + """Task 4: Edge Case - Empty list returns 0.""" + from embedding.chroma_store import store_chunks + assert store_chunks([], "user", "doc", "title") == 0 \ No newline at end of file