import os import sys # Add the project root to Python path project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) sys.path.append(project_root) import json import os import re import time from typing import Any, Dict, List, Optional from urllib.parse import urlparse import docx import exceptiongroup import google.generativeai as genai import numpy as np import pandas as pd import PyPDF2 import requests from bs4 import BeautifulSoup from dotenv import load_dotenv from nanoid import generate from pydantic import BaseModel, Field from src.utils.vectorDB import VectorStore load_dotenv() class MetadataPDF(BaseModel): key_concepts: List[str] = Field( ..., description="Key concepts related to the topic" ) page_number: int = Field( ..., alias="page-number", description="The page number where this content is located", ) class SegmentPDF(BaseModel): content: str = Field(..., description="The main text of the segment") metadata: MetadataPDF class AnalyzedContentPDF(BaseModel): segments: List[SegmentPDF] = Field( ..., description="List of meaningful content segments" ) # for text data class MetadataTxt(BaseModel): key_concepts: List[str] = Field( ..., description="Key concepts related to the topic" ) class SegmentTxt(BaseModel): content: str = Field(..., description="The main text of the segment") metadata: MetadataTxt class AnalyzedContentTxt(BaseModel): segments: List[SegmentTxt] = Field( ..., description="List of meaningful content segments" ) class GeminiChunker: def __init__(self): self.api_key = os.getenv("GEMINI_API_KEY") genai.configure(api_key=self.api_key) # self.model = genai.GenerativeModel(model_name="gemini-2.0-flash-exp") self.model = genai.GenerativeModel(model_name="gemini-1.5-flash") # self.model = genai.GenerativeModel(model_name="gemini-1.5-pro") def check_file_ready(self, file): while file.state.name == "PROCESSING": print(".", end="") time.sleep(10) file = genai.get_file(file.name) if file.state.name == "FAILED": raise ValueError(f"File processing failed: {file.state.name}") def chunk_with_gemini( self, content: str, content_type: str ) -> List[Dict[str, Any]]: # cleaned_content = "".join( # char for char in content if ord(char) >= 32 or char in "\n\r\t" # ) # # Encode and decode to handle unsupported characters # safe_string = cleaned_content.encode("utf-8", errors="replace").decode("utf-8") # print(safe_string) if content_type == "pdf" or content_type == "docx": """Use Gemini to intelligently chunk content based on semantic understanding""" prompt = f""" Analyze the following {content_type} content first means read whole content first then after divide it into complete and meaningful segments (chunks). Each chunk size has 512 token should: 1. Be self-contained and end at logical boundaries (e.g., complete sentences or sections). 2. Include all text that belongs to a single segment without truncation. 3. Ensure the last chunk is fully complete and not cut off. Return the response strictly in the specified schema format: {{ "content": "segment text here", "metadata": {{ "key_concepts": ["concept1", "concept2"], "page-number": 64 }} }}, // more segments... Content to analyze: {content} Keep the response as pure JSON without any additional text or explanation. Avoid splitting content mid-sentence or mid-thought. All chunks should be complete. """ schema = AnalyzedContentPDF else: prompt = f""" Analyze the following {content_type} content first means read whole content first then after divide it into complete and meaningful segments (chunks). Each chunk should: 1. Be self-contained and end at logical boundaries (e.g., complete sentences or sections). 2. Include all text that belongs to a single segment without truncation. 3. Ensure the last chunk is fully complete and not cut off. Return the response strictly in the specified schema format: {{ "content": "segment text here", "metadata": {{ "key_concepts": ["concept1", "concept2"], "page-number": NA }} }}, // more segments... Content to analyze: {content} Keep the response as pure JSON without any additional text or explanation. Avoid splitting content mid-sentence or mid-thought. All chunks should be complete. """ schema = AnalyzedContentTxt print(schema) try: response = self.model.generate_content( prompt, generation_config=genai.GenerationConfig( response_mime_type="application/json", response_schema=schema, ), ) # print(response.text) cleaned_text = "".join( char for char in response.text if ord(char) >= 32 or char in "\n\r\t" ) with open("chunking_text.txt", "w", encoding="utf-8") as file_text: print("^^^^^^^^^^^^^^^^^^^^^^^^^^^^^") file_text.write(cleaned_text) print("^^^^^^^^^^^^^^^^^^^^^^^^^^^^^") cleaned_text = cleaned_text.encode("utf-8", errors="replace").decode( "utf-8" ) result = json.loads(cleaned_text) # print(result) print("pdf parsing") chunks = [] for segment in result.get("segments", []): # print("################################") # print(segment.get("content", "")) # print("################################") temp_metadata = segment["metadata"] if content_type == "pdf": chunk = { "content": segment.get("content", ""), "metadata": { "topics": temp_metadata.get("key_concepts", []), "page-number": temp_metadata.get("page-number", ""), "type": "pdf", }, } else: chunk = { "content": segment.get("content", ""), "metadata": { "topics": segment.get("key_concepts", []), "page-number": segment.get("page-number", ""), "type": "text", }, } chunks.append(chunk) # print() # print() # print("chunks:") # print(chunks) # print() # print() return chunks except Exception as e: print(f"Error in Gemini chunking: {e}") return [ { "content": content, "metadata": { "topics": "", "page-number": "0", }, } ] def process_media_file( self, file_path: str, media_type: str ) -> List[Dict[str, Any]]: """Process video or audio file using Gemini's media understanding""" try: print("here-0!!!") media_file = genai.upload_file(path=file_path) self.check_file_ready(media_file) print("here-1!!!") if media_type == "video": schema = { "type": "object", "properties": { "segments": { "type": "array", "items": { "type": "object", "properties": { "timestamp": {"type": "string"}, "description": {"type": "string"}, "topics": { "type": "array", "items": {"type": "string"}, }, }, }, } }, } prompt = "Describe this video in detail, breaking it into timestamped segments. Include key events and actions." else: # audio schema = { "type": "object", "properties": { "segments": { "type": "array", "items": { "type": "object", "properties": { "timestamp": {"type": "string"}, "transcription": {"type": "string"}, "speaker": {"type": "string"}, "topics": { "type": "array", "items": {"type": "string"}, }, }, }, } }, } prompt = "Transcribe this audio, identifying speakers and key topics discussed." print("Here-2!!!") response = self.model.generate_content( [media_file, prompt], generation_config=genai.GenerationConfig( response_schema=schema, response_mime_type="application/json" ), ) print("Here-3!!!") # Convert Gemini's media response to our standard chunk format print(response.text) print("Here-4!!!") cleaned_text = "".join( char for char in response.text if ord(char) >= 32 or char in "\n\r\t" ) result = json.loads(cleaned_text) # result = json.loads(response.text) chunks = [] print("Here-5!!!") for segment in result.get("segments", []): if media_type == "video": chunk = { "content": segment.get("description", ""), "metadata": { "timestamp": segment.get("timestamp", ""), "topics": segment.get("key_events", []), "type": "video", }, } else: chunk = { "content": segment.get("transcription", ""), "metadata": { "timestamp": segment.get("timestamp", ""), "speaker": segment.get("speaker", ""), "topics": segment.get("topics", []), "type": "audio", }, } chunks.append(chunk) print("Here-6!!!") return chunks except Exception as e: print(f"Error processing {media_type} file: {e}") return [ { "content": f"Error processing {media_type} file", "metadata": {"type": media_type, "error": str(e)}, } ] class ContentProcessor: def __init__(self): self.gemini_chunker = GeminiChunker() def process_text(self, text: str, source_type: str) -> List[Dict[str, Any]]: """Process any text content using Gemini chunking""" chunks = self.gemini_chunker.chunk_with_gemini(text, source_type) for chunk in chunks: chunk["metadata"]["source_type"] = source_type return chunks def process_pdf(self, file_path: str) -> List[Dict[str, Any]]: with open(file_path, "rb") as file: pdf_reader = PyPDF2.PdfReader(file) full_text = "" for page in pdf_reader.pages: full_text += page.extract_text() + " " return self.process_text(full_text, "pdf") def process_docx(self, file_path: str) -> List[Dict[str, Any]]: doc = docx.Document(file_path) full_text = " ".join([paragraph.text for paragraph in doc.paragraphs]) return self.process_text(full_text, "docx") def process_csv(self, file_path: str) -> List[Dict[str, Any]]: df = pd.read_csv(file_path) # Convert DataFrame to a more readable format for Gemini text_content = df.to_string() return self.process_text(text_content, "csv") def process_webpage(self, url: str) -> List[Dict[str, Any]]: response = requests.get(url) soup = BeautifulSoup(response.text, "html.parser") for script in soup(["script", "style"]): script.decompose() text = soup.get_text() return self.process_text(text, "webpage") def process_video(self, file_path: str) -> List[Dict[str, Any]]: """Process video using Gemini's video understanding capabilities""" print("in process function of video") video_file = genai.upload_file(path=file_path) self.gemini_chunker.check_file_ready(video_file) chunks = self.gemini_chunker.process_media_file( file_path=file_path, media_type="video" ) return chunks def process_audio(self, file_path: str) -> List[Dict[str, Any]]: """Process audio using Gemini's audio understanding capabilities""" print("in process function of audio") audio_file = genai.upload_file(path=file_path) self.gemini_chunker.check_file_ready(audio_file) chunks = self.gemini_chunker.process_media_file( file_path=file_path, media_type="audio" ) return chunks class AgenticRAG: def __init__(self, query_value=False, is_uploaded=False): self.processor = ContentProcessor() self.vector_store = VectorStore(query=query_value, is_uploaded=is_uploaded) if query_value == False and is_uploaded == True: self.json_file_path = "json_file_record.json" else: self.json_file_path = "utils/json_file_record.json" def process_file(self, file_path: str, file_type: Optional[str] = None): if file_type is None: file_type = self._detect_file_type(file_path) if os.path.exists(self.json_file_path): with open(self.json_file_path, "r") as json_file: json_data = json.load(json_file) for record in json_data: if record["file_path"] == file_path: return True # File path exists try: chunks = [] if file_type == "pdf": chunks = self.processor.process_pdf(file_path) elif file_type == "docx": chunks = self.processor.process_docx(file_path) elif file_type == "csv": chunks = self.processor.process_csv(file_path) elif file_type == "url": chunks = self.processor.process_webpage(file_path) elif file_type == "video": chunks = self.processor.process_video(file_path) elif file_type == "audio": chunks = self.processor.process_audio(file_path) elif file_type == "text": with open(file_path, "r") as file: chunks = self.processor.process_text(file.read(), "text") if chunks: # Add source information to metadata print("in processfile fucntion file.") for chunk in chunks: chunk["metadata"]["source"] = file_path print(chunks) # Add to Vector Database. self.vector_store.add_documents(chunks) print(f"Successfully processed {file_path} with {len(chunks)} chunks") return True return False except Exception as e: print(f"Error processing {file_path}: {e}") def _detect_file_type(self, file_path: str) -> str: if file_path.startswith("http"): return "url" extension = file_path.split(".")[-1].lower() type_mapping = { "pdf": "pdf", "docx": "docx", "doc": "docx", "csv": "csv", "txt": "text", "mp3": "audio", "wav": "audio", "mp4": "video", "mov": "video", } return type_mapping.get(extension, "unknown") def query(self, query_text: str, n_results: int = 5) -> Dict: return self.vector_store.query(query_text, n_results) # Define a function to determine the file type based on the extension def get_file_type(file_name: str) -> str: if file_name.endswith(".mp3"): return "audio" elif file_name.endswith(".mp4"): return "video" elif file_name.endswith(".csv"): return "csv" elif file_name.endswith(".pdf"): return "pdf" elif file_name.endswith(".docx"): return "docx" elif file_name.startswith("http"): return "url" else: return "unknown" def main(): # Initialize the RAG system rag = AgenticRAG(is_uploaded=True) # Automatically read the files in the 'data' directory data_directory = "../data" test_files = [] # Loop through all files in the 'data' directory for filename in os.listdir(data_directory): file_path = os.path.join(data_directory, filename) if os.path.isfile(file_path): # Check if it's a file file_type = get_file_type(filename) test_files.append((file_path, file_type)) # Process each file for file_path, file_type in test_files: print(f"\nProcessing {file_path}...") time.sleep(5) rag.process_file(file_path, file_type) if __name__ == "__main__": main()