File size: 27,549 Bytes
38df4e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82de5c7
ebdd994
38df4e4
7979b63
 
 
 
 
 
 
 
 
 
 
 
 
38df4e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a430e55
 
38df4e4
 
 
 
 
 
 
 
 
 
82de5c7
 
38df4e4
 
 
 
 
 
 
 
82de5c7
 
 
7753f0c
82de5c7
8693fa5
 
82de5c7
 
 
ebdd994
82de5c7
 
 
288225c
 
 
 
 
 
 
 
82de5c7
 
38df4e4
 
 
 
7753f0c
38df4e4
 
82de5c7
 
38df4e4
 
 
 
 
 
 
 
 
 
 
82de5c7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38df4e4
 
82de5c7
 
 
 
 
 
 
 
 
 
 
38df4e4
82de5c7
38df4e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82de5c7
 
38df4e4
 
 
 
 
 
 
 
 
 
 
 
 
 
82de5c7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38df4e4
 
82de5c7
 
38df4e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82de5c7
 
 
 
 
 
 
 
 
 
 
38df4e4
82de5c7
38df4e4
82de5c7
38df4e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82de5c7
38df4e4
82de5c7
 
 
38df4e4
 
 
 
 
82de5c7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38df4e4
82de5c7
38df4e4
 
 
82de5c7
38df4e4
 
 
 
 
 
 
 
 
 
82de5c7
38df4e4
 
 
 
 
 
 
82de5c7
38df4e4
 
 
82de5c7
38df4e4
 
 
 
 
 
 
 
 
 
 
82de5c7
38df4e4
82de5c7
38df4e4
 
 
 
 
82de5c7
38df4e4
 
 
 
 
 
 
 
 
82de5c7
38df4e4
 
 
 
 
 
82de5c7
38df4e4
 
82de5c7
38df4e4
82de5c7
 
 
 
 
 
38df4e4
82de5c7
38df4e4
82de5c7
 
 
 
 
 
 
 
 
 
 
38df4e4
 
 
 
82de5c7
38df4e4
 
 
 
82de5c7
38df4e4
 
 
 
 
82de5c7
 
 
 
 
 
 
 
 
 
 
 
38df4e4
 
 
 
 
 
 
 
 
 
 
82de5c7
 
38df4e4
 
82de5c7
 
38df4e4
 
 
82de5c7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38df4e4
82de5c7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38df4e4
 
 
82de5c7
38df4e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82de5c7
38df4e4
 
 
 
 
 
 
 
 
 
 
 
82de5c7
38df4e4
 
 
 
 
 
 
 
 
82de5c7
38df4e4
 
 
82de5c7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38df4e4
82de5c7
 
 
 
 
 
 
38df4e4
 
82de5c7
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
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
from transformers import RagRetriever, RagTokenizer, RagSequenceForGeneration
from transformers import AutoTokenizer, AutoModelWithLMHead


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 = "/tmp"  # /tmp is always writable in Hugging Face Spaces
    os.makedirs(temp_dir, exist_ok=True)
    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: Optional[bytes]
    file_type: Optional[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="microsoft/Phi-3-mini-4k-instruct",
    device_map="auto",
    torch_dtype="auto",
    max_new_tokens=256,
    trust_remote_code=True
)

# Initialize RAG components
tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base", trust_remote_code=True)
retriever = RagRetriever.from_pretrained(
    "facebook/rag-token-base",
    index_name="exact",           # or "legacy" for legacy FAISS index
    use_dummy_dataset=False,       # set to False and download the full index for real Wikipedia retrieval
    trust_remote_code=True,       # Trust remote code for dataset loading
    dataset_revision="main",      # Specify a fixed revision
    dataset="wiki_dpr",           # Explicitly specify dataset name
)
rag_model = RagSequenceForGeneration.from_pretrained(
    "facebook/rag-token-base", 
    retriever=retriever,
    trust_remote_code=True
)
# Speech-to-text pipeline
asr_pipe = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-small",
    device="auto"
)

# --- BLIP VQA setup ---
device = "cuda" if torch.cuda.is_available() else "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 functions ---
def ensure_final_answer_format(answer_text: str) -> str:
    """Ensure the answer ends with FINAL ANSWER: format"""
    # Check if the answer already contains a FINAL ANSWER section
    if "FINAL ANSWER:" in answer_text:
        # Extract everything after FINAL ANSWER:
        final_answer_part = answer_text.split("FINAL ANSWER:", 1)[1].strip()
        return f"FINAL ANSWER: {final_answer_part}"
    else:
        # If no FINAL ANSWER section exists, wrap the entire answer
        return f"FINAL ANSWER: {answer_text.strip()}"

def extract_entities(text: str) -> List[str]:
    """Extract key entities from text using spaCy if available, or regex fallback"""
    if nlp:
        # Using spaCy for better entity extraction
        doc = nlp(text)
        entities = [ent.text for ent in doc.ents]
        keywords = [token.text for token in doc if token.pos_ in ("PROPN", "NOUN")]
        return entities if entities else keywords
    else:
        # Simple fallback using regex to extract potential keywords
        words = text.lower().split()
        stopwords = ["what", "who", "when", "where", "why", "how", "is", "are", "the", "a", "an", "of", "in", "on", "at"]
        keywords = [word for word in words if word not in stopwords and len(word) > 2]
        return keywords

def answer_question_on_frame(image_path, question):
    """Answer a question about a single video frame using BLIP"""
    try:
        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
    except Exception as e:
        print(f"Error processing frame {image_path}: {str(e)}")
        return "Error processing this frame"

def answer_video_question(frames_dir, question):
    """Answer a question about a video by analyzing extracted frames"""
    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='/tmp/video/', output_filename='downloaded_video.mp4'):
    """Download a YouTube video using yt-dlp"""
    # 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)

    try:
        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
    except Exception as e:
        print(f"Error downloading YouTube video: {str(e)}")
        return None

def extract_frames(video_path, output_dir, frame_interval_seconds=10):
    """Extract frames from a video file at specified intervals"""
    # 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) -> str:
    """Answer questions about an image using the BLIP model"""
    try:
        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
    except Exception as e:
        print(f"Error in image_qa: {str(e)}")
        return f"Error processing image: {str(e)}"

# --- Node functions ---
def router(state: Dict[str, Any]) -> str:
    """Determine the next node based on question content and file type"""
    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"

def node_decide(state: Dict[str, Any]) -> Dict[str, Any]:
    """Router node that decides which node to go to next"""
    print("Running node_decide")
    # Initialize context list if not present
    if 'context' not in state:
        state['context'] = []
    # Add the next state to the state dict
    state["next"] = router(state)
    print(f"Routing to: {state['next']}")
    return state

def node_image(state: Dict[str, Any]) -> Dict[str, Any]:
    """Process image-based questions"""
    print("Running node_image")
    try:
        # Make sure the image file exists
        if not os.path.exists(state['file_path']):
            state['answer'] = ensure_final_answer_format("Image file not found.")
            return state
            
        # Get answer from image QA model
        answer = image_qa(state['file_path'], state['question'])
        
        # Format the final answer
        state['answer'] = ensure_final_answer_format(answer)
        
        # Add document to state for traceability
        image_doc = Document(
            page_content=f"Image analysis result: {answer}",
            metadata={"source": "image_analysis", "file_path": state['file_path']}
        )
        state['context'].append(image_doc)
        
    except Exception as e:
        error_msg = f"Error processing image: {str(e)}"
        print(error_msg)
        state['answer'] = ensure_final_answer_format(error_msg)
    
    return state

def node_video(state: Dict[str, Any]) -> Dict[str, Any]:
    """Process video-based questions"""
    print("Running node_video")
    youtube_url = state.get('youtube_url')
    if not youtube_url:
        state['answer'] = ensure_final_answer_format("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'] = ensure_final_answer_format("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'] = ensure_final_answer_format("Failed to extract frames from the video.")
        return state

    result = answer_video_question(frames_dir, question_text)
    final_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
    state['context'].extend(frame_documents)
    state['answer'] = ensure_final_answer_format(final_answer)

    print(f"Video answer: {state['answer']}")
    return state

def node_audio_rag(state: Dict[str, Any]) -> Dict[str, Any]:
    """Process audio-based questions"""
    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 transcript in 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)
        retrieved_context = "\n".join([doc.page_content for doc in similar_docs])

        # Step 4: Generate answer
        prompt = (
            f"You are an AI assistant that answers questions about audio content.\n\n"
            f"Audio transcript: {retrieved_context}\n\n"
            f"Question: {question}\n\n"
            f"Based only on the provided audio transcript, answer the question. "
            f"If the transcript does not contain relevant information, state that clearly.\n\n"
            f"End your response with 'FINAL ANSWER: ' followed by a concise answer."
        )
        
        llm_response = llm_pipe(prompt)
        answer_text = llm_response[0]['generated_text']
        
        # Add documents to state
        state['context'].extend(transcript_doc)
        state['context'].append(Document(
            page_content=prompt,
            metadata={"source": "audio_analysis_prompt"}
        ))
        
        # Ensure final answer format
        state['answer'] = ensure_final_answer_format(answer_text)

    except Exception as e:
        error_msg = f"Audio processing error: {str(e)}"
        print(error_msg)
        state['answer'] = ensure_final_answer_format(error_msg)

    return state

def node_llm(state: Dict[str, Any]) -> Dict[str, Any]:
    """Process general knowledge questions with LLM"""
    print("Running node_llm")
    question = state['question']

    # Compose a detailed prompt
    prompt = (
        "You are an AI assistant that answers questions using your general knowledge. "
        "Follow these steps:\n\n"
        "1. If the question appears to be scrambled or jumbled, first try to unscramble or reconstruct the intended meaning.\n"
        "2. Analyze the question (unscrambled if needed) and use your own knowledge to answer it.\n"
        "3. If the question can't be answered with certainty, provide your best estimate and clearly explain any assumptions.\n"
        "4. Format your answer using these rules:\n"
        "   - Numbers: Plain digits without commas/units (e.g. 1234567)\n"
        "   - Strings: Minimal words, no articles/abbreviations\n"
        "   - Lists: comma-separated values without extra formatting\n\n"
        "5. Always conclude with:\n"
        "FINAL ANSWER: [your answer] (replace bracketed text)\n\n"
        f"Current question: {question}"
    )

    # Add document to state for traceability
    query_doc = Document(
        page_content=prompt,
        metadata={"source": "llm_prompt"}
    )
    state['context'].append(query_doc)

    try:
        result = llm_pipe(prompt)
        answer_text = result[0]['generated_text']
        state['answer'] = ensure_final_answer_format(answer_text)
    except Exception as e:
        print(f"Error in LLM processing: {str(e)}")
        error_msg = f"An error occurred while processing your question: {str(e)}"
        state['answer'] = ensure_final_answer_format(error_msg)

    print(f"LLM answer: {state['answer']}")
    return state
def retrieve(state: State) -> State:
    """Retrieve relevant documents using RAG"""
    print("Running retrieve")
    question = state["question"]
    
    try:
        # Tokenize the question
        inputs = tokenizer(question, return_tensors="pt")
        
        # Get doc_ids by using the retriever directly
        question_hidden_states = rag_model.question_encoder(inputs["input_ids"])[0]
        docs_dict = retriever(
            inputs["input_ids"].numpy(),
            question_hidden_states.detach().numpy(), 
            return_tensors="pt"
        )
        
        # Extract the retrieved passages
        all_chunks = []
        
        # Debug print to see what's in docs_dict
        print(f"docs_dict keys: {docs_dict.keys()}")
        
        # Check for different possible keys that might contain the documents
        doc_text_key = None
        for possible_key in ['retrieved_doc_text', 'doc_text', 'texts', 'documents']:
            if possible_key in docs_dict:
                doc_text_key = possible_key
                break
        
        if doc_text_key:
            # Access the retrieved document texts from the docs_dict
            for i in range(len(docs_dict["doc_ids"][0])):
                doc_text = docs_dict[doc_text_key][0][i]
                all_chunks.append(Document(page_content=doc_text))
                
            print(f"Retrieved {len(all_chunks)} documents")
        else:
            # Fallback: Try to extract document text from doc_ids
            doc_ids = docs_dict.get("doc_ids", [[]])[0]
            print(f"Retrieved doc_ids: {doc_ids}")
            
            # Create minimal document stubs from IDs
            for doc_id in doc_ids:
                stub_text = f"Information related to document ID: {doc_id}"
                all_chunks.append(Document(page_content=stub_text))
            
            print(f"Created {len(all_chunks)} document stubs from IDs")
        
        # Add documents to state context
        if not state.get('context'):
            state['context'] = []
        state['context'].extend(all_chunks)
        
    except Exception as e:
        print(f"Error in retrieval: {str(e)}")
        # Create an error document
        error_doc = Document(
            page_content=f"Error during retrieval: {str(e)}",
            metadata={"source": "retrieval_error"}
        )
        if not state.get('context'):
            state['context'] = []
        state['context'].append(error_doc)
    
    return state

def generate(state: State) -> State:
    """Generate an answer based on retrieved documents"""
    print("Running generate")
    
    try:
        # Check if context exists
        if not state.get('context') or len(state['context']) == 0:
            state['answer'] = ensure_final_answer_format("No relevant information found to answer your question.")
            return state
            
        # 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 = PromptTemplate(
            input_variables=["question", "context"],
            template=(
                "You are an AI assistant that answers questions using retrieved context. "
                "Follow these steps:\n\n"
                "1. Analyze the provided context:\n{context}\n\n"
                "2. If the context contains scrambled text, first attempt to reconstruct meaningful information\n"
                "3. If the question can't be answered from context alone, combine context with general knowledge "
                "but clearly state this limitation\n"
                "4. Format your answer using these rules:\n"
                "   - Numbers: Plain digits without commas/units (e.g. 1234567)\n"
                "   - Strings: Minimal words, no articles/abbreviations\n"
                "   - Lists: comma-separated values without extra formatting\n\n"
                "5. Always conclude with:\n"
                "FINAL ANSWER: [your answer] (replace bracketed text)\n\n"
                "Current question: {question}"
            )
        ).format(question=state["question"], context=docs_content)
        
        # Generate answer using the LLM pipeline
        response = llm_pipe(prompt_str)
        answer_text = response[0]["generated_text"]
        
        # Ensure answer has the FINAL ANSWER format
        state['answer'] = ensure_final_answer_format(answer_text)
        
    except Exception as e:
        print(f"Error in generate node: {str(e)}")
        error_msg = f"Error generating answer: {str(e)}"
        state['answer'] = ensure_final_answer_format(error_msg)
        
    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"]

# 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 other nodes based on question
graph.add_conditional_edges(
    "decide",
    get_next_node,
    {
        "video": "video",
        "llm": "llm",
        "retrieve": "retrieve",
        "image": "image",
        "audio": "audio"
    }
)

# Add edges from all terminal nodes to END
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()

# --- Intelligent Agent Function ---
def intelligent_agent(state: State) -> str:
    """Process a question using the appropriate pipeline based on content."""
    try:
        # Ensure state has proper structure
        if not isinstance(state, dict):
            return "FINAL ANSWER: Error - input must be a valid State dictionary"
            
        # Make sure question exists
        if 'question' not in state:
            return "FINAL ANSWER: Error - question is required"
            
        # Initialize context if not present
        if 'context' not in state:
            state['context'] = []
            
        print(f"Processing question: {state['question']}")
        
        # Invoke the agent with the state
        final_state = agent.invoke(state)
        
        # Ensure answer has FINAL ANSWER format
        answer = final_state.get('answer', "No answer found.")
        formatted_answer = ensure_final_answer_format(answer)
        
        return formatted_answer
        
    except Exception as e:
        print(f"Error in agent execution: {str(e)}")
        return f"FINAL ANSWER: An error occurred - {str(e)}"