import numpy as np import spacy import tempfile import glob import yt_dlp import shutil import cv2 import librosa import wikipedia from typing import TypedDict, List, Optional, Dict, Any from langchain.docstore.document import Document from langchain.prompts import PromptTemplate from langchain_community.document_loaders import WikipediaLoader from langgraph.graph import START, END, StateGraph from langchain_core.messages import AnyMessage, HumanMessage, AIMessage # If you are using it from langchain_community.retrievers import BM25Retriever # If you are using it from langgraph.prebuilt import ToolNode, tools_condition # If you are using it from langchain.vectorstores import FAISS from langchain.embeddings import HuggingFaceEmbeddings from langchain.schema import Document from transformers import BlipProcessor, BlipForQuestionAnswering, pipeline from io import BytesIO from sentence_transformers import SentenceTransformer import os import re from PIL import Image # This is correctly imported, but was being used incorrectly import numpy as np from collections import Counter import torch from transformers import BlipProcessor, BlipForQuestionAnswering, pipeline from typing import TypedDict, List, Optional, Dict, Any, Literal, Tuple from langgraph.graph import StateGraph, START, END from langchain.docstore.document import Document nlp = spacy.load("en_core_web_sm") # Define file extension sets for each category PICTURE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'} AUDIO_EXTENSIONS = {'.mp3', '.wav', '.aac', '.flac', '.ogg', '.m4a', '.wma'} CODE_EXTENSIONS = {'.py', '.js', '.java', '.cpp', '.c', '.cs', '.rb', '.go', '.php', '.html', '.css', '.ts'} SPREADSHEET_EXTENSIONS = { '.xls', '.xlsx', '.xlsm', '.xlsb', '.xlt', '.xltx', '.xltm', '.ods', '.ots', '.csv', '.tsv', '.sxc', '.stc', '.dif', '.gsheet', '.numbers', '.numbers-tef', '.nmbtemplate', '.fods', '.123', '.wk1', '.wk2', '.wks', '.wku', '.wr1', '.gnumeric', '.gnm', '.xml', '.pmvx', '.pmdx', '.pmv', '.uos', '.txt' } def get_file_type(filename: str) -> str: if not filename or '.' not in filename or filename == '': return '' ext = filename.lower().rsplit('.', 1)[-1] dot_ext = f'.{ext}' if dot_ext in PICTURE_EXTENSIONS: return 'picture' elif dot_ext in AUDIO_EXTENSIONS: return 'audio' elif dot_ext in CODE_EXTENSIONS: return 'code' elif dot_ext in SPREADSHEET_EXTENSIONS: return 'spreadsheet' else: return 'unknown' def write_bytes_to_temp_dir(file_bytes: bytes, file_name: str) -> str: """ Writes bytes to a file in the system temporary directory using the provided file_name. Returns the full path to the saved file. The file will persist until manually deleted or the OS cleans the temp directory. """ temp_dir = tempfile.gettempdir() file_path = os.path.join(temp_dir, file_name) with open(file_path, 'wb') as f: f.write(file_bytes) print(f"File written to: {file_path}") return file_path # 1. Define the State type class State(TypedDict, total=False): question: str task_id: str input_file: bytes file_type: str context: List[Document] # Using LangChain's Document class file_path: Optional[str] youtube_url: Optional[str] answer: Optional[str] frame_answers: Optional[list] next: Optional[str] # Added to track the next node # --- LLM pipeline for general questions --- llm_pipe = pipeline("text-generation", #model="meta-llama/Llama-3.3-70B-Instruct", #model="meta-llama/Meta-Llama-3-8B-Instruct", #model="Qwen/Qwen2-7B-Instruct", #model="microsoft/Phi-4-reasoning", model="microsoft/Phi-3-mini-4k-instruct", device_map="auto", #device_map={ "": 0 }, # "" means the whole model #max_memory={0: "10GiB"}, torch_dtype="auto", max_new_tokens=256) # Speech-to-text pipeline asr_pipe = pipeline( "automatic-speech-recognition", model="openai/whisper-small", device=-1 #device_map={"", 0}, #max_memory = {0: "4.5GiB"}, #device_map="auto" ) # --- Your BLIP VQA setup --- #device = "cuda" if torch.cuda.is_available() else "cpu" device = "cpu" vqa_model_name = "Salesforce/blip-vqa-base" processor_vqa = BlipProcessor.from_pretrained(vqa_model_name) # Attempt to load model to GPU; fall back to CPU if OOM try: model_vqa = BlipForQuestionAnswering.from_pretrained(vqa_model_name).to(device) except torch.cuda.OutOfMemoryError: print("WARNING: Loading model to CPU due to insufficient GPU memory.") device = "cpu" # Switch device to CPU model_vqa = BlipForQuestionAnswering.from_pretrained(vqa_model_name).to(device) # --- Helper: Answer question on a single frame --- def answer_question_on_frame(image_path, question): # Fixed: Properly use the PIL Image module image = Image.open(image_path).convert('RGB') inputs = processor_vqa(image, question, return_tensors="pt").to(device) out = model_vqa.generate(**inputs) answer = processor_vqa.decode(out[0], skip_special_tokens=True) return answer # --- Helper: Answer question about the whole video --- def answer_video_question(frames_dir, question): valid_exts = ('.jpg', '.jpeg', '.png') # Check if directory exists if not os.path.exists(frames_dir): return { "most_common_answer": "No frames found to analyze.", "all_answers": [], "answer_counts": Counter() } frame_files = [os.path.join(frames_dir, f) for f in os.listdir(frames_dir) if f.lower().endswith(valid_exts)] # Sort frames properly by number def get_frame_number(filename): match = re.search(r'(\d+)', os.path.basename(filename)) return int(match.group(1)) if match else 0 frame_files = sorted(frame_files, key=get_frame_number) if not frame_files: return { "most_common_answer": "No valid image frames found.", "all_answers": [], "answer_counts": Counter() } answers = [] for frame_path in frame_files: try: ans = answer_question_on_frame(frame_path, question) answers.append(ans) print(f"Processed frame: {os.path.basename(frame_path)}, Answer: {ans}") except Exception as e: print(f"Error processing frame {frame_path}: {str(e)}") if not answers: return { "most_common_answer": "Could not analyze any frames successfully.", "all_answers": [], "answer_counts": Counter() } counted = Counter(answers) most_common_answer, freq = counted.most_common(1)[0] return { "most_common_answer": most_common_answer, "all_answers": answers, "answer_counts": counted } def download_youtube_video(url, output_dir='/content/video/', output_filename='downloaded_video.mp4'): # Ensure the output directory exists os.makedirs(output_dir, exist_ok=True) # Delete all files in the output directory files = glob.glob(os.path.join(output_dir, '*')) for f in files: try: os.remove(f) except Exception as e: print(f"Error deleting {f}: {str(e)}") # Set output path for yt-dlp output_path = os.path.join(output_dir, output_filename) ydl_opts = { 'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best', 'outtmpl': output_path, 'quiet': True, 'merge_output_format': 'mp4', # Ensures merged output is mp4 'postprocessors': [{ 'key': 'FFmpegVideoConvertor', 'preferedformat': 'mp4', # Recode if needed }] } with yt_dlp.YoutubeDL(ydl_opts) as ydl: ydl.download([url]) return output_path # --- Helper: Extract frames from video --- def extract_frames(video_path, output_dir, frame_interval_seconds=10): # --- Clean output directory before extracting new frames --- if os.path.exists(output_dir): for filename in os.listdir(output_dir): file_path = os.path.join(output_dir, filename) try: if os.path.isfile(file_path) or os.path.islink(file_path): os.unlink(file_path) elif os.path.isdir(file_path): shutil.rmtree(file_path) except Exception as e: print(f'Failed to delete {file_path}. Reason: {e}') else: os.makedirs(output_dir, exist_ok=True) try: cap = cv2.VideoCapture(video_path) if not cap.isOpened(): print("Error: Could not open video.") return False fps = cap.get(cv2.CAP_PROP_FPS) frame_interval = int(fps * frame_interval_seconds) count = 0 saved = 0 while True: ret, frame = cap.read() if not ret: break if count % frame_interval == 0: frame_filename = os.path.join(output_dir, f"frame_{count:06d}.jpg") cv2.imwrite(frame_filename, frame) saved += 1 count += 1 cap.release() print(f"Extracted {saved} frames.") return saved > 0 except Exception as e: print(f"Exception during frame extraction: {e}") return False def image_qa(image_path: str, question: str, model_name: str = vqa_model_name) -> str: """ Answers questions about images using Hugging Face's VQA pipeline. Args: image_path: Path to local image file or URL question: Natural language question about the image model_name: Pretrained VQA model (default: good general-purpose model) Returns: str: The model's best answer """ # Create VQA pipeline with specified model vqa_pipeline = pipeline("visual-question-answering", model=model_name) # Get predictions (automatically handles local files/URLs) results = vqa_pipeline(image=image_path, question=question, top_k=1) # Return top answer return results[0]['answer'] def router(state: Dict[str, Any]) -> str: """Determine the next node based on whether the question contains a YouTube URL or references Wikipedia.""" question = state.get('question', '') # Pattern for Wikipedia and similar sources wiki_pattern = r"(wikipedia\.org|wiki|encyclopedia|britannica\.com|encyclop[a|æ]dia)" has_wiki = re.search(wiki_pattern, question, re.IGNORECASE) is not None # Pattern for YouTube yt_pattern = r"(https?://)?(www\.)?(youtube\.com|youtu\.be)/[^\s]+" has_youtube = re.search(yt_pattern, question) is not None # Check for image has_image = state.get('file_type') == 'picture' # Check for audio has_audio = state.get('file_type') == 'audio' print(f"Has Wikipedia reference: {has_wiki}") print(f"Has YouTube link: {has_youtube}") print(f"Has picture file: {has_image}") print(f"Has audio file: {has_audio}") if has_wiki: return "retrieve" elif has_youtube: # Store the extracted YouTube URL in the state url_match = re.search(r"(https?://[^\s]+)", question) if url_match: state['youtube_url'] = url_match.group(0) return "video" elif has_image: return "image" elif has_audio: return "audio" else: return "llm" # --- Node Implementation --- def node_image(state: Dict[str, Any]) -> Dict[str, Any]: """Router node that decides which node to go to next.""" print("Running node_image") # Add the next state to the state dict img = Image.open(state['file_path']) state['answer'] = image_qa(state['file_path'], state['question']) return state def node_decide(state: Dict[str, Any]) -> Dict[str, Any]: """Router node that decides which node to go to next.""" print("Running node_decide") # Add the next state to the state dict state["next"] = router(state) print(f"Routing to: {state['next']}") return state def node_video(state: Dict[str, Any]) -> Dict[str, Any]: print("Running node_video") youtube_url = state.get('youtube_url') if not youtube_url: state['answer'] = "No YouTube URL found in the question." return state question = state['question'] # Extract the actual question part (remove the URL) question_text = re.sub(r'https?://[^\s]+', '', question).strip() if not question_text.endswith('?'): question_text += '?' video_file = download_youtube_video(youtube_url) if not video_file or not os.path.exists(video_file): state['answer'] = "Failed to download the video." return state frames_dir = "/tmp/frames" os.makedirs(frames_dir, exist_ok=True) success = extract_frames(video_path=video_file, output_dir=frames_dir, frame_interval_seconds=10) if not success: state['answer'] = "Failed to extract frames from the video." return state result = answer_video_question(frames_dir, question_text) state['answer'] = result['most_common_answer'] state['frame_answers'] = result['all_answers'] # Create Document objects for each frame analysis frame_documents = [] for i, ans in enumerate(result['all_answers']): doc = Document( page_content=f"Frame {i}: {ans}", metadata={"frame_number": i, "source": "video_analysis"} ) frame_documents.append(doc) # Add documents to state if not already present if 'context' not in state: state['context'] = [] state['context'].extend(frame_documents) print(f"Video answer: {state['answer']}") return state def node_audio_rag(state: Dict[str, Any]) -> Dict[str, Any]: print(f"Processing audio file: {state['file_path']}") try: # Step 1: Transcribe audio audio, sr = librosa.load(state['file_path'], sr=16000) asr_result = asr_pipe({"raw": audio, "sampling_rate": sr}) audio_transcript = asr_result['text'] print(f"Audio transcript: {audio_transcript}") # Step 2: Store ONLY the transcript in the vector store transcript_doc = [Document(page_content=audio_transcript)] embeddings = HuggingFaceEmbeddings(model_name='BAAI/bge-large-en-v1.5') vector_db = FAISS.from_documents(transcript_doc, embedding=embeddings) # Step 3: Retrieve relevant docs for the user's question question = state['question'] similar_docs = vector_db.similarity_search(question, k=1) # Only one doc in store retrieved_context = "\n".join([doc.page_content for doc in similar_docs]) # Step 4: Augment prompt and generate answer prompt = ( f"Use the following context to answer the question.\n" f"Context:\n{retrieved_context}\n\n" f"Question: {question}\nAnswer:" ) llm_response = llm_pipe(prompt) state['answer'] = llm_response[0]['generated_text'] except Exception as e: error_msg = f"Audio processing error: {str(e)}" print(error_msg) state['answer'] = error_msg return state def node_llm(state: Dict[str, Any]) -> Dict[str, Any]: print("Running node_llm") question = state['question'] # Optionally add context from state (e.g., Wikipedia/Wikidata content) context_text = "" if 'article_content' in state and state['article_content']: context_text = f"\n\nBackground Information:\n{state['article_content']}\n" elif 'context' in state and state['context']: context_text = "\n\n".join([doc.page_content for doc in state['context']]) # Compose a detailed prompt prompt = ( "You are an expert researcher. Answer the user's question as accurately as possible. " "If the text appears to be scrambled, try to unscramble the text for the user" "If the information is incomplete or ambiguous, provide your best estimate based on the available evidence, and clearly explain any assumptions or reasoning you use. " "If the answer requires multiple steps or deeper analysis, break down the question into sub-questions and answer them step by step, citing the relevant context for each step.\n\n" f"Question: {question}" f"{context_text}\n" "Answer:" ) # Add document to state for traceability query_doc = Document( page_content=prompt, metadata={"source": "llm_prompt"} ) if 'context' not in state: state['context'] = [] state['context'].append(query_doc) try: result = llm_pipe(prompt) state['answer'] = result[0]['generated_text'] except Exception as e: print(f"Error in LLM processing: {str(e)}") state['answer'] = f"An error occurred while processing your question: {str(e)}" print(f"LLM answer: {state['answer']}") return state # --- Define the edge condition function --- def get_next_node(state: Dict[str, Any]) -> str: """Get the next node from the state.""" return state["next"] # 2. Improved Wikipedia Retrieval Node def extract_keywords(question: str) -> List[str]: doc = nlp(question) keywords = [token.text for token in doc if token.pos_ in ("PROPN", "NOUN")] # Extract proper nouns and nouns return keywords def extract_entities(question: str) -> List[str]: doc = nlp(question) entities = [ent.text for ent in doc.ents] return entities if entities else [token.text for token in doc if token.pos_ in ("PROPN", "NOUN")] def retrieve(state: State) -> dict: keywords = extract_entities(state["question"]) query = " ".join(keywords) search_results = wikipedia.search(query) selected_page = search_results[0] if search_results else None if selected_page: loader = WikipediaLoader( query=selected_page, lang="en", load_max_docs=1, doc_content_chars_max=100000, load_all_available_meta=True ) docs = loader.load() # Chunk the article for finer retrieval from langchain.text_splitter import RecursiveCharacterTextSplitter splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200) all_chunks = [] for doc in docs: chunks = splitter.split_text(doc.page_content) all_chunks.extend([Document(page_content=chunk) for chunk in chunks]) # Optionally: re-rank or filter chunks here return {"context": all_chunks} else: return {"context": []} # 3. Prompt Template for General QA prompt = PromptTemplate( input_variables=["question", "context"], template=( "You are an expert researcher. Given the following context from Wikipedia, answer the user's question as accurately as possible. " "If the text appears to be scrambled, try to unscramble the text for the user" "If the information is incomplete or ambiguous, provide your best estimate based on the available evidence, and clearly explain any assumptions or reasoning you use. " "If the answer requires multiple steps or deeper analysis, break down the question into sub-questions and answer them step by step, citing the relevant context for each step." "Context:\n{context}\n\n" "Question: {question}\n\n" "Best Estimate Answer:" ) ) """ def generate(state: State) -> dict: # Concatenate all context documents into a single string docs_content = "\n\n".join(doc.page_content for doc in state["context"]) # Format the prompt for the LLM prompt_str = prompt.format(question=state["question"], context=docs_content) # Generate answer response = llm.invoke(prompt_str) return {"answer": response} """ def generate(state: dict) -> dict: # Concatenate all context documents into a single string docs_content = "\n\n".join(doc.page_content for doc in state["context"]) # Format the prompt for the LLM prompt_str = prompt.format(question=state["question"], context=docs_content) # Generate answer using Hugging Face pipeline response = llm_pipe(prompt_str) # Extract generated text answer = response[0]["generated_text"] return {"answer": answer} # Create the StateGraph graph = StateGraph(State) # Add nodes graph.add_node("decide", node_decide) graph.add_node("video", node_video) graph.add_node("llm", node_llm) graph.add_node("retrieve", retrieve) graph.add_node("generate", generate) graph.add_node("image", node_image) graph.add_node("audio", node_audio_rag) # Add edge from START to decide graph.add_edge(START, "decide") graph.add_edge("retrieve", "generate") # Add conditional edges from decide to video or llm based on question graph.add_conditional_edges( "decide", get_next_node, { "video": "video", "llm": "llm", "retrieve": "retrieve", "image": "image", "audio": "audio" } ) # Add edges from video and llm to END to terminate the graph graph.add_edge("video", END) graph.add_edge("llm", END) graph.add_edge("generate", END) graph.add_edge("image", END) graph.add_edge("audio", END) # Compile the graph agent = graph.compile() # --- Usage Example --- def intelligent_agent(state: State) -> str: """Process a question using the appropriate pipeline based on content.""" #state = State(question= question) try: final_state = agent.invoke(state) return final_state.get('answer', "No answer found.") except Exception as e: print(f"Error in agent execution: {str(e)}") return f"An error occurred: {str(e)}"