Spaces:
Sleeping
Sleeping
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)}" |