chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,335 @@
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import List, Dict, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
import assemblyai as aai
|
||||
from src.document_processing.doc_processor import DocumentChunk
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpeakerSegment:
|
||||
"""Represents a speaker segment with timing and content"""
|
||||
speaker: str
|
||||
start_time: float
|
||||
end_time: float
|
||||
text: str
|
||||
confidence: float
|
||||
|
||||
def get_timestamp_str(self) -> str:
|
||||
def format_time(seconds):
|
||||
minutes = int(seconds // 60)
|
||||
seconds = int(seconds % 60)
|
||||
return f"{minutes:02d}:{seconds:02d}"
|
||||
|
||||
return f"[{format_time(self.start_time)} - {format_time(self.end_time)}]"
|
||||
|
||||
|
||||
class AudioTranscriber:
|
||||
def __init__(self, api_key: str):
|
||||
self.api_key = api_key
|
||||
aai.settings.api_key = api_key
|
||||
|
||||
self.supported_formats = {
|
||||
'.mp3', '.wav', '.m4a', '.aac', '.ogg',
|
||||
'.flac', '.wma', '.opus', '.mp4', '.mov', '.avi'
|
||||
}
|
||||
|
||||
logger.info("AudioTranscriber initialized with AssemblyAI")
|
||||
|
||||
def transcribe_audio(
|
||||
self,
|
||||
audio_path: str,
|
||||
enable_speaker_diarization: bool = True,
|
||||
enable_auto_punctuation: bool = True,
|
||||
audio_language: str = "en",
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 100
|
||||
) -> List[DocumentChunk]:
|
||||
|
||||
audio_path = Path(audio_path)
|
||||
|
||||
if not audio_path.exists():
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
if audio_path.suffix.lower() not in self.supported_formats:
|
||||
raise ValueError(f"Unsupported audio format: {audio_path.suffix}")
|
||||
|
||||
logger.info(f"Starting transcription for: {audio_path.name}")
|
||||
|
||||
try:
|
||||
config = aai.TranscriptionConfig(
|
||||
speaker_labels=enable_speaker_diarization,
|
||||
punctuate=enable_auto_punctuation,
|
||||
language_code=audio_language,
|
||||
)
|
||||
|
||||
transcriber = aai.Transcriber(config=config)
|
||||
transcript = transcriber.transcribe(str(audio_path))
|
||||
|
||||
if transcript.status == aai.TranscriptStatus.error:
|
||||
raise Exception(f"Transcription failed: {transcript.error}")
|
||||
|
||||
logger.info(f"Transcription completed for: {audio_path.name}")
|
||||
|
||||
return self._process_transcript_to_chunks(
|
||||
transcript,
|
||||
audio_path.name,
|
||||
chunk_size,
|
||||
chunk_overlap
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error transcribing audio {audio_path.name}: {str(e)}")
|
||||
raise
|
||||
|
||||
def _process_transcript_to_chunks(
|
||||
self,
|
||||
transcript: aai.Transcript,
|
||||
source_file: str,
|
||||
chunk_size: int,
|
||||
chunk_overlap: int
|
||||
) -> List[DocumentChunk]:
|
||||
|
||||
chunks = []
|
||||
transcript_metadata = {
|
||||
'duration_seconds': transcript.audio_duration,
|
||||
'confidence': transcript.confidence,
|
||||
'audio_url': transcript.audio_url,
|
||||
'transcription_id': transcript.id
|
||||
}
|
||||
|
||||
if hasattr(transcript, 'utterances') and transcript.utterances:
|
||||
chunks = self._create_chunks_with_speakers(
|
||||
transcript.utterances,
|
||||
source_file,
|
||||
chunk_size,
|
||||
chunk_overlap,
|
||||
transcript_metadata
|
||||
)
|
||||
else:
|
||||
chunks = self._create_chunks_without_speakers(
|
||||
transcript.text,
|
||||
source_file,
|
||||
chunk_size,
|
||||
chunk_overlap,
|
||||
transcript_metadata
|
||||
)
|
||||
|
||||
logger.info(f"Created {len(chunks)} chunks from transcript")
|
||||
return chunks
|
||||
|
||||
def _create_chunks_with_speakers(
|
||||
self,
|
||||
utterances: List[aai.Utterance],
|
||||
source_file: str,
|
||||
chunk_size: int,
|
||||
chunk_overlap: int,
|
||||
base_metadata: Dict[str, Any]
|
||||
) -> List[DocumentChunk]:
|
||||
|
||||
chunks = []
|
||||
current_text = ""
|
||||
current_speakers = []
|
||||
current_timestamps = []
|
||||
chunk_index = 0
|
||||
start_char = 0
|
||||
|
||||
for utterance in utterances:
|
||||
speaker_label = f"Speaker {utterance.speaker}"
|
||||
timestamp_str = f"[{self._format_milliseconds(utterance.start)}]"
|
||||
|
||||
speaker_text = f"{timestamp_str} {speaker_label}: {utterance.text}\n"
|
||||
|
||||
if len(current_text + speaker_text) > chunk_size and current_text:
|
||||
chunk_metadata = base_metadata.copy()
|
||||
chunk_metadata.update({
|
||||
'speakers': list(set(current_speakers)),
|
||||
'start_timestamp': current_timestamps[0] if current_timestamps else None,
|
||||
'end_timestamp': current_timestamps[-1] if current_timestamps else None,
|
||||
'speaker_count': len(set(current_speakers))
|
||||
})
|
||||
|
||||
chunk = DocumentChunk(
|
||||
content=current_text.strip(),
|
||||
source_file=source_file,
|
||||
source_type='audio',
|
||||
page_number=None,
|
||||
chunk_index=chunk_index,
|
||||
start_char=start_char,
|
||||
end_char=start_char+len(current_text)-1,
|
||||
metadata=chunk_metadata
|
||||
)
|
||||
chunks.append(chunk)
|
||||
|
||||
overlap_text = current_text[-chunk_overlap:] if chunk_overlap > 0 else ""
|
||||
current_text = overlap_text + speaker_text
|
||||
start_char += len(current_text) - len(overlap_text) - len(speaker_text)
|
||||
chunk_index += 1
|
||||
|
||||
current_speakers = [speaker_label]
|
||||
current_timestamps = [utterance.start, utterance.end]
|
||||
else:
|
||||
current_text += speaker_text
|
||||
current_speakers.append(speaker_label)
|
||||
current_timestamps.extend([utterance.start, utterance.end])
|
||||
|
||||
if current_text.strip():
|
||||
chunk_metadata = base_metadata.copy()
|
||||
chunk_metadata.update({
|
||||
'speakers': list(set(current_speakers)),
|
||||
'start_timestamp': current_timestamps[0] if current_timestamps else None,
|
||||
'end_timestamp': current_timestamps[-1] if current_timestamps else None,
|
||||
'speaker_count': len(set(current_speakers))
|
||||
})
|
||||
|
||||
chunk = DocumentChunk(
|
||||
content=current_text.strip(),
|
||||
source_file=source_file,
|
||||
source_type='audio',
|
||||
page_number=None,
|
||||
chunk_index=chunk_index,
|
||||
start_char=start_char,
|
||||
end_char=start_char+len(current_text)-1,
|
||||
metadata=chunk_metadata
|
||||
)
|
||||
chunks.append(chunk)
|
||||
|
||||
return chunks
|
||||
|
||||
# def _create_chunks_without_speakers(
|
||||
# self,
|
||||
# transcript_text: str,
|
||||
# source_file: str,
|
||||
# chunk_size: int,
|
||||
# chunk_overlap: int,
|
||||
# base_metadata: Dict[str, Any]
|
||||
# ) -> List[DocumentChunk]:
|
||||
# if not transcript_text.strip():
|
||||
# return []
|
||||
|
||||
# chunks = []
|
||||
# start = 0
|
||||
# chunk_index = 0
|
||||
|
||||
# while start < len(transcript_text):
|
||||
# end = min(start + chunk_size, len(transcript_text))
|
||||
|
||||
# # Try to break at sentence boundary
|
||||
# if end < len(transcript_text):
|
||||
# last_period = transcript_text.rfind('.', start, end)
|
||||
# last_newline = transcript_text.rfind('\n', start, end)
|
||||
|
||||
# boundary = max(last_period, last_newline)
|
||||
# if boundary > start + chunk_size * 0.5:
|
||||
# end = boundary + 1
|
||||
|
||||
# chunk_text = transcript_text[start:end].strip()
|
||||
|
||||
# if chunk_text:
|
||||
# chunk_metadata = base_metadata.copy()
|
||||
# chunk_metadata.update({
|
||||
# 'speakers': ['Unknown Speaker'],
|
||||
# 'speaker_count': 1
|
||||
# })
|
||||
|
||||
# chunk = DocumentChunk(
|
||||
# content=chunk_text,
|
||||
# source_file=source_file,
|
||||
# source_type='audio',
|
||||
# page_number=None,
|
||||
# chunk_index=chunk_index,
|
||||
# start_char=start,
|
||||
# end_char=end - 1,
|
||||
# metadata=chunk_metadata
|
||||
# )
|
||||
# chunks.append(chunk)
|
||||
# chunk_index += 1
|
||||
|
||||
# start = max(start + chunk_size - chunk_overlap, end)
|
||||
|
||||
# return chunks
|
||||
|
||||
def _format_milliseconds(self, ms: int) -> str:
|
||||
seconds = ms // 1000
|
||||
minutes = seconds // 60
|
||||
seconds = seconds % 60
|
||||
return f"{minutes:02d}:{seconds:02d}"
|
||||
|
||||
def get_transcript_summary(self, audio_path: str) -> Dict[str, Any]:
|
||||
try:
|
||||
config = aai.TranscriptionConfig(
|
||||
speaker_labels=True,
|
||||
summarization=True
|
||||
)
|
||||
|
||||
transcriber = aai.Transcriber(config=config)
|
||||
transcript = transcriber.transcribe(str(audio_path))
|
||||
|
||||
if transcript.status == aai.TranscriptStatus.error:
|
||||
return {"error": transcript.error}
|
||||
|
||||
summary_info = {
|
||||
'id': transcript.id,
|
||||
'file_name': Path(audio_path).name,
|
||||
'duration_seconds': transcript.audio_duration,
|
||||
'confidence': transcript.confidence,
|
||||
'word_count': len(transcript.text.split()) if transcript.text else 0,
|
||||
'character_count': len(transcript.text) if transcript.text else 0,
|
||||
'summary': getattr(transcript, 'summary', 'Not available'),
|
||||
'speaker_count': len(set(u.speaker for u in transcript.utterances)) if hasattr(transcript, 'utterances') and transcript.utterances else 1
|
||||
}
|
||||
|
||||
return summary_info
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting transcript summary: {str(e)}")
|
||||
return {"error": str(e)}
|
||||
|
||||
def batch_transcribe(self, audio_paths: List[str]) -> List[List[DocumentChunk]]:
|
||||
all_chunks = []
|
||||
for audio_path in audio_paths:
|
||||
try:
|
||||
chunks = self.transcribe_audio(audio_path)
|
||||
all_chunks.append(chunks)
|
||||
logger.info(f"Successfully transcribed {audio_path}: {len(chunks)} chunks")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to transcribe {audio_path}: {str(e)}")
|
||||
all_chunks.append([])
|
||||
|
||||
return all_chunks
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
api_key = os.getenv("ASSEMBLYAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Please set ASSEMBLYAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
transcriber = AudioTranscriber(api_key)
|
||||
|
||||
try:
|
||||
audio_file = "data/harvard.wav"
|
||||
|
||||
summary = transcriber.get_transcript_summary(audio_file)
|
||||
print(f"Transcript Summary: {summary}")
|
||||
|
||||
# Full transcription
|
||||
chunks = transcriber.transcribe_audio(audio_file)
|
||||
|
||||
print(f"\nTranscription Results:")
|
||||
print(f"Generated {len(chunks)} chunks")
|
||||
|
||||
for i, chunk in enumerate(chunks[:3]):
|
||||
print(f"\nChunk {i+1}:")
|
||||
print(f"Content: {chunk.content[:200]}...")
|
||||
print(f"Speakers: {chunk.metadata.get('speakers', [])}")
|
||||
print(f"Citation: Source: {chunk.source_file}, Type: Audio Transcript")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in transcription example: {e}")
|
||||
@@ -0,0 +1,153 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
import yt_dlp
|
||||
import assemblyai as aai
|
||||
|
||||
from src.document_processing.doc_processor import DocumentChunk
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class YouTubeTranscriber:
|
||||
def __init__(self, assemblyai_api_key: str):
|
||||
self.assemblyai_api_key = assemblyai_api_key
|
||||
self.temp_dir = Path(tempfile.gettempdir()) / "youtube_transcriber"
|
||||
self.temp_dir.mkdir(exist_ok=True)
|
||||
|
||||
aai.settings.api_key = assemblyai_api_key
|
||||
|
||||
logger.info("YouTubeTranscriber initialized")
|
||||
|
||||
def extract_video_id(self, url: str) -> Optional[str]:
|
||||
if "v=" in url:
|
||||
video_id = url.split("v=")[1].split("&")[0]
|
||||
elif "youtu.be/" in url:
|
||||
video_id = url.split("youtu.be/")[1].split("?")[0]
|
||||
else:
|
||||
video_id = None
|
||||
return video_id
|
||||
|
||||
def download_audio(self, url: str) -> str:
|
||||
video_id = self.extract_video_id(url)
|
||||
if not video_id:
|
||||
raise ValueError("Could not extract video ID from URL")
|
||||
|
||||
expected_path = self.temp_dir / f"{video_id}.m4a"
|
||||
if expected_path.exists():
|
||||
logger.info(f"Audio already exists: {expected_path}")
|
||||
return str(expected_path)
|
||||
|
||||
logger.info(f"Downloading audio from: {url}")
|
||||
|
||||
ydl_opts = {
|
||||
'format': 'm4a/bestaudio/best',
|
||||
'outtmpl': str(self.temp_dir / '%(id)s.%(ext)s'),
|
||||
'postprocessors': [{
|
||||
'key': 'FFmpegExtractAudio',
|
||||
'preferredcodec': 'm4a',
|
||||
}],
|
||||
'quiet': True,
|
||||
'no_warnings': True,
|
||||
}
|
||||
|
||||
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
|
||||
error_code = ydl.download([url])
|
||||
|
||||
if error_code != 0:
|
||||
raise Exception(f"yt-dlp download failed with error code: {error_code}")
|
||||
|
||||
if not expected_path.exists():
|
||||
raise FileNotFoundError(f"Expected audio file not found: {expected_path}")
|
||||
|
||||
logger.info(f"Audio downloaded successfully: {expected_path}")
|
||||
return str(expected_path)
|
||||
|
||||
def transcribe_youtube_video(
|
||||
self,
|
||||
url: str,
|
||||
cleanup_audio: bool = True
|
||||
) -> List[DocumentChunk]:
|
||||
try:
|
||||
audio_path = self.download_audio(url)
|
||||
|
||||
# Configure transcription with speaker diarization
|
||||
config = aai.TranscriptionConfig(
|
||||
speaker_labels=True,
|
||||
punctuate=True
|
||||
)
|
||||
|
||||
logger.info("Starting transcription with speaker diarization...")
|
||||
transcriber = aai.Transcriber(config=config)
|
||||
transcript = transcriber.transcribe(audio_path)
|
||||
|
||||
if transcript.status == aai.TranscriptStatus.error:
|
||||
raise Exception(f"Transcription failed: {transcript.error}")
|
||||
|
||||
chunks = []
|
||||
video_id = self.extract_video_id(url)
|
||||
for i, utterance in enumerate(transcript.utterances):
|
||||
chunk = DocumentChunk(
|
||||
content=f"Speaker {utterance.speaker}: {utterance.text}",
|
||||
source_file=f"YouTube Video {video_id}",
|
||||
source_type="youtube",
|
||||
page_number=None,
|
||||
chunk_index=i,
|
||||
start_char=utterance.start,
|
||||
end_char=utterance.end,
|
||||
metadata={
|
||||
'speaker': utterance.speaker,
|
||||
'start_time': utterance.start,
|
||||
'end_time': utterance.end,
|
||||
'confidence': getattr(utterance, 'confidence', None),
|
||||
'video_url': url,
|
||||
'video_id': video_id
|
||||
}
|
||||
)
|
||||
chunks.append(chunk)
|
||||
|
||||
logger.info(f"Transcription completed: {len(chunks)} utterances")
|
||||
|
||||
if cleanup_audio and os.path.exists(audio_path):
|
||||
os.unlink(audio_path)
|
||||
logger.info("Audio file cleaned up")
|
||||
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error transcribing YouTube video: {str(e)}")
|
||||
raise
|
||||
|
||||
def cleanup_temp_files(self):
|
||||
try:
|
||||
if self.temp_dir.exists():
|
||||
for file in self.temp_dir.glob("*.m4a"):
|
||||
file.unlink()
|
||||
logger.info("Temporary files cleaned up")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not clean up temp files: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
|
||||
api_key = os.getenv("ASSEMBLYAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Please set ASSEMBLYAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
transcriber = YouTubeTranscriber(api_key)
|
||||
|
||||
try:
|
||||
test_url = "https://www.youtube.com/watch?v=D26sUZ6DHNQ"
|
||||
chunks = transcriber.transcribe_youtube_video(test_url)
|
||||
|
||||
print(f"Transcribed {len(chunks)} utterances:")
|
||||
for chunk in chunks[:5]:
|
||||
print(f" {chunk.content}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
@@ -0,0 +1,220 @@
|
||||
import os
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
import hashlib
|
||||
from datetime import datetime
|
||||
|
||||
import pymupdf
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocumentChunk:
|
||||
"""Represents a processed document chunk with metadata for citations"""
|
||||
content: str
|
||||
source_file: str
|
||||
source_type: str # 'pdf', 'txt', 'web', 'audio'
|
||||
page_number: Optional[int] = None
|
||||
chunk_index: int = 0
|
||||
start_char: Optional[int] = None
|
||||
end_char: Optional[int] = None
|
||||
metadata: Dict[str, Any] = None
|
||||
chunk_id: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.chunk_id:
|
||||
self.chunk_id = self._generate_chunk_id()
|
||||
if self.metadata is None:
|
||||
self.metadata = {}
|
||||
|
||||
def _generate_chunk_id(self) -> str:
|
||||
content_hash = hashlib.md5(self.content.encode()).hexdigest()[:8]
|
||||
return f"{self.source_type}_{self.chunk_index}_{content_hash}"
|
||||
|
||||
def get_citation_info(self) -> Dict[str, Any]:
|
||||
citation = {
|
||||
'source': self.source_file,
|
||||
'type': self.source_type,
|
||||
'chunk_id': self.chunk_id,
|
||||
'chunk_index': self.chunk_index
|
||||
}
|
||||
|
||||
if self.page_number:
|
||||
citation['page'] = self.page_number
|
||||
if self.start_char or self.end_char:
|
||||
citation['char_range'] = f"{self.start_char}-{self.end_char}"
|
||||
|
||||
citation.update(self.metadata)
|
||||
return citation
|
||||
|
||||
|
||||
class DocumentProcessor:
|
||||
def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200):
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_overlap = chunk_overlap
|
||||
self.supported_formats = {'.pdf', '.txt', '.md'} # add other formats if need be
|
||||
|
||||
def process_document(self, file_path: str) -> List[DocumentChunk]:
|
||||
file_path = Path(file_path)
|
||||
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
if file_path.suffix.lower() not in self.supported_formats:
|
||||
raise ValueError(f"Unsupported file format: {file_path.suffix}")
|
||||
|
||||
logger.info(f"Processing document: {file_path.name}")
|
||||
|
||||
try:
|
||||
if file_path.suffix.lower() == '.pdf':
|
||||
return self._process_pdf(file_path)
|
||||
elif file_path.suffix.lower() in {'.txt', '.md'}:
|
||||
return self._process_text_file(file_path)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing {file_path.name}: {str(e)}")
|
||||
raise
|
||||
|
||||
def _process_pdf(self, file_path: Path) -> List[DocumentChunk]:
|
||||
chunks = []
|
||||
try:
|
||||
doc = pymupdf.open(file_path)
|
||||
total_pages = len(doc)
|
||||
|
||||
for page_num in range(total_pages):
|
||||
page = doc.load_page(page_num)
|
||||
text = page.get_text()
|
||||
|
||||
if not text.strip():
|
||||
continue
|
||||
|
||||
# Get page metadata
|
||||
page_metadata = {
|
||||
'total_pages': total_pages,
|
||||
'page_width': page.rect.width,
|
||||
'page_height': page.rect.height,
|
||||
'processed_at': datetime.now().isoformat()
|
||||
}
|
||||
|
||||
page_chunks = self._create_chunks_from_text(
|
||||
text,
|
||||
file_path.name,
|
||||
source_type='pdf',
|
||||
page_number=page_num+1,
|
||||
additional_metadata=page_metadata
|
||||
)
|
||||
chunks.extend(page_chunks)
|
||||
|
||||
doc.close()
|
||||
logger.info(f"Processed PDF: {len(chunks)} chunks from {total_pages} pages")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing PDF {file_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
return chunks
|
||||
|
||||
def _process_text_file(self, file_path: Path) -> List[DocumentChunk]:
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
|
||||
metadata = {
|
||||
'file_size': file_path.stat().st_size,
|
||||
'encoding': 'utf-8',
|
||||
'processed_at': datetime.now().isoformat()
|
||||
}
|
||||
|
||||
chunks = self._create_chunks_from_text(
|
||||
content,
|
||||
file_path.name,
|
||||
source_type='txt',
|
||||
page_number=None,
|
||||
additional_metadata=metadata
|
||||
)
|
||||
|
||||
logger.info(f"Processed text file: {len(chunks)} chunks")
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing text file {file_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
def _create_chunks_from_text(
|
||||
self,
|
||||
text: str,
|
||||
source_file: str,
|
||||
source_type: str,
|
||||
page_number: Optional[int] = None,
|
||||
additional_metadata: Dict[str, Any] = None
|
||||
) -> List[DocumentChunk]:
|
||||
|
||||
if not text.strip():
|
||||
return []
|
||||
|
||||
chunks = []
|
||||
start = 0
|
||||
chunk_index = 0
|
||||
while start < len(text):
|
||||
end = min(start + self.chunk_size, len(text))
|
||||
if end < len(text):
|
||||
last_period = text.rfind('.', start, end)
|
||||
last_newline = text.rfind('\n', start, end)
|
||||
boundary = max(last_period, last_newline)
|
||||
if boundary > start + self.chunk_size * 0.5:
|
||||
end = boundary + 1
|
||||
|
||||
chunk_text = text[start:end].strip()
|
||||
|
||||
if chunk_text:
|
||||
chunk_metadata = additional_metadata.copy() if additional_metadata else {}
|
||||
|
||||
chunk = DocumentChunk(
|
||||
content=chunk_text,
|
||||
source_file=source_file,
|
||||
source_type=source_type,
|
||||
page_number=page_number,
|
||||
chunk_index=chunk_index,
|
||||
start_char=start,
|
||||
end_char=end-1,
|
||||
metadata=chunk_metadata
|
||||
)
|
||||
|
||||
chunks.append(chunk)
|
||||
chunk_index += 1
|
||||
|
||||
start = max(start + self.chunk_size - self.chunk_overlap, end)
|
||||
if start >= len(text):
|
||||
break
|
||||
|
||||
return chunks
|
||||
|
||||
def batch_process(self, file_paths: List[str]) -> List[DocumentChunk]:
|
||||
all_chunks = []
|
||||
for file_path in file_paths:
|
||||
try:
|
||||
chunks = self.process_document(file_path)
|
||||
all_chunks.extend(chunks)
|
||||
logger.info(f"Successfully processed {file_path}: {len(chunks)} chunks")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process {file_path}: {str(e)}")
|
||||
continue
|
||||
|
||||
logger.info(f"Batch processing complete: {len(all_chunks)} total chunks from {len(file_paths)} files")
|
||||
return all_chunks
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
processor = DocumentProcessor(chunk_size=800, chunk_overlap=100)
|
||||
|
||||
try:
|
||||
chunks = processor.process_document("data/raft.pdf")
|
||||
sample_chunk = chunks[0]
|
||||
print(f"Sample chunk content: {sample_chunk.content[:200]}...")
|
||||
print(f"Citation info: {sample_chunk.get_citation_info()}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
@@ -0,0 +1,137 @@
|
||||
import logging
|
||||
from typing import List, Dict, Any, Tuple
|
||||
import numpy as np
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastembed import TextEmbedding
|
||||
from src.document_processing.doc_processor import DocumentChunk
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddedChunk:
|
||||
"""Document chunk with its embedding vector"""
|
||||
chunk: DocumentChunk
|
||||
embedding: np.ndarray
|
||||
embedding_model: str
|
||||
|
||||
def to_vector_db_format(self) -> Dict[str, Any]:
|
||||
return {
|
||||
'id': self.chunk.chunk_id,
|
||||
'vector': self.embedding.tolist(),
|
||||
'content': self.chunk.content,
|
||||
'source_file': self.chunk.source_file,
|
||||
'source_type': self.chunk.source_type,
|
||||
'page_number': self.chunk.page_number,
|
||||
'chunk_index': self.chunk.chunk_index,
|
||||
'start_char': self.chunk.start_char,
|
||||
'end_char': self.chunk.end_char,
|
||||
'metadata': self.chunk.metadata,
|
||||
'embedding_model': self.embedding_model
|
||||
}
|
||||
|
||||
|
||||
class EmbeddingGenerator:
|
||||
def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"):
|
||||
self.model_name = model_name
|
||||
self.model = None
|
||||
self.embedding_dim = None
|
||||
self._initialize_model()
|
||||
|
||||
def _initialize_model(self):
|
||||
try:
|
||||
logger.info(f"Initializing embedding model: {self.model_name}")
|
||||
self.model = TextEmbedding(model_name=self.model_name)
|
||||
|
||||
sample_embedding = list(self.model.embed(["test"]))[0]
|
||||
self.embedding_dim = len(sample_embedding)
|
||||
|
||||
logger.info(f"Model initialized successfully. Embedding dimension: {self.embedding_dim}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize embedding model: {str(e)}")
|
||||
raise
|
||||
|
||||
def generate_embeddings(self, chunks: List[DocumentChunk]) -> List[EmbeddedChunk]:
|
||||
if not chunks:
|
||||
return []
|
||||
|
||||
logger.info(f"Generating embeddings for {len(chunks)} chunks")
|
||||
|
||||
try:
|
||||
texts = [chunk.content for chunk in chunks]
|
||||
|
||||
embeddings = list(self.model.embed(texts))
|
||||
embedded_chunks = []
|
||||
for chunk, embedding in zip(chunks, embeddings):
|
||||
embedded_chunk = EmbeddedChunk(
|
||||
chunk=chunk,
|
||||
embedding=np.array(embedding, dtype=np.float32),
|
||||
embedding_model=self.model_name
|
||||
)
|
||||
embedded_chunks.append(embedded_chunk)
|
||||
|
||||
logger.info(f"Successfully generated {len(embedded_chunks)} embeddings")
|
||||
return embedded_chunks
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating embeddings: {str(e)}")
|
||||
raise
|
||||
|
||||
def generate_query_embedding(self, query_text: str) -> np.ndarray:
|
||||
try:
|
||||
embedding = list(self.model.embed([query_text]))[0]
|
||||
return np.array(embedding, dtype=np.float32)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating query embedding: {str(e)}")
|
||||
raise
|
||||
|
||||
def get_embedding_dimension(self) -> int:
|
||||
return self.embedding_dim
|
||||
|
||||
def batch_generate_embeddings(
|
||||
self,
|
||||
chunks_batches: List[List[DocumentChunk]],
|
||||
batch_size: int = 32
|
||||
) -> List[List[EmbeddedChunk]]:
|
||||
|
||||
all_embedded_batches = []
|
||||
for i, chunk_batch in enumerate(chunks_batches):
|
||||
logger.info(f"Processing batch {i+1}/{len(chunks_batches)}")
|
||||
|
||||
embedded_batch = []
|
||||
for j in range(0, len(chunk_batch), batch_size):
|
||||
sub_batch = chunk_batch[j:j + batch_size]
|
||||
embedded_sub_batch = self.generate_embeddings(sub_batch)
|
||||
embedded_batch.extend(embedded_sub_batch)
|
||||
|
||||
all_embedded_batches.append(embedded_batch)
|
||||
|
||||
return all_embedded_batches
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from src.document_processing.doc_processor import DocumentProcessor
|
||||
|
||||
doc_processor = DocumentProcessor()
|
||||
embedding_generator = EmbeddingGenerator()
|
||||
|
||||
try:
|
||||
chunks = doc_processor.process_document("data/raft.pdf")
|
||||
embedded_chunks = embedding_generator.generate_embeddings(chunks)
|
||||
|
||||
if embedded_chunks:
|
||||
sample = embedded_chunks[0]
|
||||
print(f"Sample embedding shape: {sample.embedding.shape}")
|
||||
print(f"Sample content: {sample.chunk.content[:100]}...")
|
||||
print(f"Citation info: {sample.chunk.get_citation_info()}")
|
||||
|
||||
query = "What is the main topic?"
|
||||
query_embedding = embedding_generator.generate_query_embedding(query)
|
||||
print(f"Query embedding shape: {query_embedding.shape}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in example usage: {e}")
|
||||
@@ -0,0 +1,281 @@
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
|
||||
from crewai import LLM
|
||||
from src.vector_database.milvus_vector_db import MilvusVectorDB
|
||||
from src.embeddings.embedding_generator import EmbeddingGenerator
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RAGResult:
|
||||
"""Represents the result of RAG generation with citations"""
|
||||
query: str
|
||||
response: str
|
||||
sources_used: List[Dict[str, Any]]
|
||||
retrieval_count: int
|
||||
generation_tokens: Optional[int] = None
|
||||
|
||||
def get_citation_summary(self) -> str:
|
||||
if not self.sources_used:
|
||||
return "No sources cited"
|
||||
|
||||
source_summary = []
|
||||
for source in self.sources_used:
|
||||
source_info = f"• {source.get('source_file', 'Unknown')} ({source.get('source_type', 'unknown')})"
|
||||
if source.get('page_number'):
|
||||
source_info += f" - Page {source['page_number']}"
|
||||
source_summary.append(source_info)
|
||||
|
||||
return "\n".join(source_summary)
|
||||
|
||||
|
||||
class RAGGenerator:
|
||||
def __init__(
|
||||
self,
|
||||
embedding_generator: EmbeddingGenerator,
|
||||
vector_db: MilvusVectorDB,
|
||||
openai_api_key: str,
|
||||
model_name: str = "gpt-4o-mini",
|
||||
temperature: float = 0.1,
|
||||
max_tokens: int = 2000
|
||||
):
|
||||
self.embedding_generator = embedding_generator
|
||||
self.vector_db = vector_db
|
||||
|
||||
self.llm = LLM(
|
||||
model=f"openai/{model_name}",
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
api_key=openai_api_key
|
||||
)
|
||||
|
||||
self.model_name = model_name
|
||||
logger.info(f"RAG Generator initialized with {model_name}")
|
||||
|
||||
def generate_response(
|
||||
self,
|
||||
query: str,
|
||||
max_chunks: int = 8,
|
||||
max_context_chars: int = 4000,
|
||||
top_k: int = 10,
|
||||
) -> RAGResult:
|
||||
|
||||
if not query.strip():
|
||||
return RAGResult(
|
||||
query=query,
|
||||
response="Please provide a valid question.",
|
||||
sources_used=[],
|
||||
retrieval_count=0
|
||||
)
|
||||
|
||||
try:
|
||||
logger.info(f"Generating response for: '{query[:50]}...'")
|
||||
|
||||
# Step 1: Retrieve relevant chunks
|
||||
query_vector = self.embedding_generator.generate_query_embedding(query)
|
||||
search_results = self.vector_db.search(
|
||||
query_vector=query_vector.tolist(),
|
||||
limit=top_k
|
||||
)
|
||||
|
||||
if not search_results:
|
||||
return RAGResult(
|
||||
query=query,
|
||||
response="I couldn't find any relevant information in the available documents to answer your question.",
|
||||
sources_used=[],
|
||||
retrieval_count=0
|
||||
)
|
||||
|
||||
# Step 2: Format context with citations
|
||||
context, sources_info = self._format_context_with_citations(
|
||||
search_results, max_chunks, max_context_chars
|
||||
)
|
||||
|
||||
# Step 3: Create citation-aware prompt
|
||||
prompt = self._create_rag_prompt(query, context)
|
||||
|
||||
# Step 4: Generate response
|
||||
response = self.llm.call(prompt)
|
||||
|
||||
# Step 5: Create result object
|
||||
rag_result = RAGResult(
|
||||
query=query,
|
||||
response=response,
|
||||
sources_used=sources_info,
|
||||
retrieval_count=len(search_results)
|
||||
)
|
||||
|
||||
logger.info(f"Response generated successfully using {len(sources_info)} sources")
|
||||
return rag_result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating response: {str(e)}")
|
||||
return RAGResult(
|
||||
query=query,
|
||||
response=f"I encountered an error while processing your question: {str(e)}",
|
||||
sources_used=[],
|
||||
retrieval_count=0
|
||||
)
|
||||
|
||||
def _format_context_with_citations(
|
||||
self,
|
||||
search_results: List[Dict[str, Any]],
|
||||
max_chunks: int,
|
||||
max_context_chars: int
|
||||
) -> Tuple[str, List[Dict[str, Any]]]:
|
||||
|
||||
context_parts = []
|
||||
sources_info = []
|
||||
total_chars = 0
|
||||
for i, result in enumerate(search_results[:max_chunks]):
|
||||
citation_info = result['citation']
|
||||
source_file = citation_info.get('source_file', 'Unknown Source')
|
||||
source_type = citation_info.get('source_type', 'unknown')
|
||||
page_number = citation_info.get('page_number')
|
||||
|
||||
citation_ref = f"[{i+1}]"
|
||||
chunk_content = result['content']
|
||||
chunk_text = f"{citation_ref} {chunk_content}"
|
||||
|
||||
if total_chars + len(chunk_text) > max_context_chars and context_parts:
|
||||
break
|
||||
|
||||
context_parts.append(chunk_text)
|
||||
total_chars += len(chunk_text)
|
||||
|
||||
source_info = {
|
||||
'reference': citation_ref,
|
||||
'source_file': source_file,
|
||||
'source_type': source_type,
|
||||
'page_number': page_number,
|
||||
'chunk_id': result['id'],
|
||||
'relevance_score': result['score']
|
||||
}
|
||||
sources_info.append(source_info)
|
||||
|
||||
formatted_context = '\n\n'.join(context_parts)
|
||||
|
||||
return formatted_context, sources_info
|
||||
|
||||
def _create_rag_prompt(self, query: str, context: str) -> str:
|
||||
prompt = f"""You are an AI assistant that answers questions based on provided source material. You must follow these citation rules:
|
||||
|
||||
CITATION REQUIREMENTS:
|
||||
1. For each factual claim in your answer, include the citation reference number in square brackets [1], [2], etc.
|
||||
2. Only use information from the provided context - do not add external knowledge
|
||||
3. If you cannot find relevant information in the context, say so clearly
|
||||
4. Be precise and accurate in your citations
|
||||
5. When multiple sources support the same point, list all relevant citations like this [1], [2], [3].
|
||||
|
||||
CONTEXT (with citation references):
|
||||
{context}
|
||||
|
||||
QUESTION: {query}
|
||||
|
||||
Please provide a comprehensive answer with proper citations. Make sure every factual statement is supported by a citation reference."""
|
||||
|
||||
return prompt
|
||||
|
||||
def generate_summary(
|
||||
self,
|
||||
max_chunks: int = 15,
|
||||
summary_length: str = "medium"
|
||||
) -> RAGResult:
|
||||
try:
|
||||
summary_query = "main topics key findings important information overview"
|
||||
query_vector = self.embedding_generator.generate_query_embedding(summary_query)
|
||||
search_results = self.vector_db.search(
|
||||
query_vector=query_vector.tolist(),
|
||||
limit=max_chunks
|
||||
)
|
||||
|
||||
if not search_results:
|
||||
return RAGResult(
|
||||
query="Document Summary",
|
||||
response="No documents available for summarization.",
|
||||
sources_used=[],
|
||||
retrieval_count=0
|
||||
)
|
||||
|
||||
context, sources_info = self._format_context_with_citations(
|
||||
search_results, max_chunks, 6000
|
||||
)
|
||||
|
||||
length_instructions = {
|
||||
'short': "Provide a concise 2-3 paragraph summary highlighting the most important points.",
|
||||
'medium': "Provide a comprehensive 4-5 paragraph summary covering key topics and findings.",
|
||||
'long': "Provide a detailed summary with multiple sections covering all major topics and supporting details."
|
||||
}
|
||||
|
||||
summary_prompt = f"""You are tasked with creating a summary of the provided document content. Follow these guidelines:
|
||||
|
||||
1. {length_instructions.get(summary_length, length_instructions['medium'])}
|
||||
2. Include citations [1], [2], etc. for all factual claims
|
||||
3. Organize information logically with clear topics
|
||||
4. Focus on the most important and relevant information
|
||||
5. Maintain accuracy and cite sources properly
|
||||
|
||||
DOCUMENT CONTENT (with citation references):
|
||||
{context}
|
||||
|
||||
Please provide a well-structured summary with proper citations:"""
|
||||
|
||||
response = self.llm.call(summary_prompt)
|
||||
|
||||
return RAGResult(
|
||||
query="Document Summary",
|
||||
response=response,
|
||||
sources_used=sources_info,
|
||||
retrieval_count=len(search_results)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating summary: {str(e)}")
|
||||
return RAGResult(
|
||||
query="Document Summary",
|
||||
response=f"Error generating summary: {str(e)}",
|
||||
sources_used=[],
|
||||
retrieval_count=0
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
from src.document_processing.doc_processor import DocumentProcessor
|
||||
from src.embeddings.embedding_generator import EmbeddingGenerator
|
||||
from src.vector_database.milvus_vector_db import MilvusVectorDB
|
||||
|
||||
openai_key = os.getenv("OPENAI_API_KEY")
|
||||
if not openai_key:
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
try:
|
||||
embedding_gen = EmbeddingGenerator()
|
||||
vector_db = MilvusVectorDB()
|
||||
rag_generator = RAGGenerator(
|
||||
embedding_generator=embedding_gen,
|
||||
vector_db=vector_db,
|
||||
openai_api_key=openai_key,
|
||||
model_name="gpt-4o-mini",
|
||||
temperature=0.1
|
||||
)
|
||||
|
||||
test_query = "What are the main findings discussed in the documents?"
|
||||
result = rag_generator.generate_response(test_query)
|
||||
|
||||
print(f"Query: {result.query}")
|
||||
print(f"Response: {result.response}")
|
||||
print(f"\nSources Used ({len(result.sources_used)}):")
|
||||
print(result.get_citation_summary())
|
||||
|
||||
summary_result = rag_generator.generate_summary(summary_length="medium")
|
||||
print(f"\nDocument Summary:")
|
||||
print(summary_result.response)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in RAG pipeline example: {e}")
|
||||
@@ -0,0 +1,326 @@
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Optional, Any, Dict, List
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
||||
from zep_cloud.client import Zep
|
||||
from zep_crewai import ZepUserStorage
|
||||
from crewai.memory.external.external_memory import ExternalMemory
|
||||
from src.generation.rag import RAGResult
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversationTurn:
|
||||
"""Represents a single conversation turn with context"""
|
||||
user_query: str
|
||||
assistant_response: str
|
||||
sources_used: List[Dict[str, Any]]
|
||||
timestamp: str
|
||||
session_id: str
|
||||
|
||||
|
||||
class NotebookMemoryLayer:
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
zep_api_key: Optional[str] = None,
|
||||
mode: str = "summary",
|
||||
indexing_wait_time: int = 10,
|
||||
create_new_session: bool = False
|
||||
):
|
||||
self.user_id = user_id
|
||||
self.session_id = session_id
|
||||
self.indexing_wait_time = indexing_wait_time
|
||||
self.zep_client = Zep(api_key=zep_api_key or os.getenv("ZEP_API_KEY"))
|
||||
|
||||
self._setup_user_and_session(create_new_session)
|
||||
|
||||
self.user_storage = ZepUserStorage(
|
||||
client=self.zep_client,
|
||||
user_id=self.user_id,
|
||||
thread_id=self.session_id,
|
||||
mode=mode,
|
||||
)
|
||||
self.external_memory = ExternalMemory(storage=self.user_storage)
|
||||
|
||||
logger.info(f"NotebookMemoryLayer initialized for user {user_id}, session {session_id}")
|
||||
|
||||
def _setup_user_and_session(self, create_new_session: bool):
|
||||
try:
|
||||
# Ensure user exists
|
||||
try:
|
||||
self.zep_client.user.get(self.user_id)
|
||||
logger.info(f"Using existing user: {self.user_id}")
|
||||
except:
|
||||
self.zep_client.user.add(user_id=self.user_id)
|
||||
logger.info(f"Created new user: {self.user_id}")
|
||||
|
||||
if create_new_session:
|
||||
try:
|
||||
self.zep_client.thread.delete(self.session_id)
|
||||
logger.info(f"Deleted previous session: {self.session_id}")
|
||||
except:
|
||||
pass
|
||||
|
||||
self.zep_client.thread.create(thread_id=self.session_id, user_id=self.user_id)
|
||||
logger.info(f"Created new session: {self.session_id}")
|
||||
else:
|
||||
# Try to use existing session, create if doesn't exist
|
||||
try:
|
||||
self.zep_client.thread.get(self.session_id)
|
||||
logger.info(f"Using existing session: {self.session_id}")
|
||||
except:
|
||||
self.zep_client.thread.create(thread_id=self.session_id, user_id=self.user_id)
|
||||
logger.info(f"Created session: {self.session_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error setting up user/session: {str(e)}")
|
||||
raise
|
||||
|
||||
def save_conversation_turn(
|
||||
self,
|
||||
rag_result: RAGResult,
|
||||
user_metadata: Optional[Dict[str, Any]] = None,
|
||||
assistant_metadata: Optional[Dict[str, Any]] = None
|
||||
):
|
||||
try:
|
||||
user_meta = {
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"session_id": self.session_id,
|
||||
**(user_metadata or {})
|
||||
}
|
||||
|
||||
# Save user message
|
||||
self.external_memory.save(
|
||||
rag_result.query,
|
||||
metadata=user_meta
|
||||
)
|
||||
|
||||
assistant_meta = {
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"session_id": self.session_id,
|
||||
"sources_count": len(rag_result.sources_used),
|
||||
"retrieval_count": rag_result.retrieval_count,
|
||||
"model_used": getattr(rag_result, 'model_name', 'unknown'),
|
||||
"sources_summary": self._create_sources_summary(rag_result.sources_used),
|
||||
**(assistant_metadata or {})
|
||||
}
|
||||
|
||||
# Save assistant response
|
||||
self.external_memory.save(
|
||||
rag_result.response,
|
||||
metadata=assistant_meta
|
||||
)
|
||||
self._save_source_context(rag_result.sources_used)
|
||||
|
||||
logger.info(f"Saved conversation turn with {len(rag_result.sources_used)} sources")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving conversation turn: {str(e)}")
|
||||
raise
|
||||
|
||||
def _create_sources_summary(self, sources_used: List[Dict[str, Any]]) -> str:
|
||||
if not sources_used:
|
||||
return "No sources used"
|
||||
|
||||
source_files = list(set(source.get('source_file', 'Unknown') for source in sources_used))
|
||||
source_types = list(set(source.get('source_type', 'unknown') for source in sources_used))
|
||||
|
||||
summary = f"{len(source_files)} files ({', '.join(source_types)}): {', '.join(source_files[:3])}"
|
||||
if len(source_files) > 3:
|
||||
summary += f" and {len(source_files) - 3} more"
|
||||
|
||||
return summary
|
||||
|
||||
def _save_source_context(self, sources_used: List[Dict[str, Any]]):
|
||||
if not sources_used:
|
||||
return
|
||||
|
||||
source_context = {
|
||||
"referenced_documents": [],
|
||||
"document_types": set(),
|
||||
"key_topics_discussed": []
|
||||
}
|
||||
|
||||
for source in sources_used:
|
||||
doc_info = {
|
||||
"file": source.get('source_file', 'Unknown'),
|
||||
"type": source.get('source_type', 'unknown'),
|
||||
"page": source.get('page_number'),
|
||||
"relevance": source.get('relevance_score', 0)
|
||||
}
|
||||
source_context["referenced_documents"].append(doc_info)
|
||||
source_context["document_types"].add(doc_info["type"])
|
||||
|
||||
source_context["document_types"] = list(source_context["document_types"])
|
||||
|
||||
self.external_memory.save(
|
||||
f"Document sources referenced: {source_context}",
|
||||
metadata={
|
||||
"type": "source_context",
|
||||
"category": "document_usage",
|
||||
"session_id": self.session_id
|
||||
}
|
||||
)
|
||||
|
||||
def save_user_preferences(self, preferences: Dict[str, Any]):
|
||||
try:
|
||||
self.external_memory.save(
|
||||
f"User preferences: {preferences}",
|
||||
metadata={
|
||||
"type": "preferences",
|
||||
"category": "user_settings",
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"session_id": self.session_id
|
||||
}
|
||||
)
|
||||
logger.info("User preferences saved to memory")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving preferences: {str(e)}")
|
||||
|
||||
def save_document_metadata(self, document_info: Dict[str, Any]):
|
||||
try:
|
||||
self.external_memory.save(
|
||||
f"Document processed: {document_info}",
|
||||
metadata={
|
||||
"type": "document_metadata",
|
||||
"category": "system_events",
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"session_id": self.session_id
|
||||
}
|
||||
)
|
||||
logger.info(f"Document metadata saved: {document_info.get('name', 'Unknown')}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving document metadata: {str(e)}")
|
||||
|
||||
def get_conversation_context(self) -> str:
|
||||
try:
|
||||
memory = self.zep_client.thread.get_user_context(thread_id=self.session_id)
|
||||
return memory.context if memory.context else ""
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting conversation context: {str(e)}")
|
||||
return "No conversation context available"
|
||||
|
||||
def get_relevant_memory(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
|
||||
try:
|
||||
# Use Zep's semantic graph search on memory
|
||||
results = self.zep_client.graph.search(
|
||||
user_id=self.user_id,
|
||||
query=query,
|
||||
scope="episodes",
|
||||
)
|
||||
|
||||
relevant_memories = []
|
||||
for ep in results.episodes:
|
||||
memory_info = {
|
||||
"content": ep.content if ep.content else "",
|
||||
"role": ep.role_type if ep.role_type else "unknown",
|
||||
"relevance_score": ep.score if hasattr(ep, 'score') else 0,
|
||||
"thread_id": ep.thread_id if ep.thread_id else None,
|
||||
"session_id": ep.session_id if ep.session_id else None,
|
||||
"timestamp": ep.created_at if ep.created_at else None,
|
||||
}
|
||||
relevant_memories.append(memory_info)
|
||||
|
||||
logger.info(f"Retrieved {len(relevant_memories)} relevant memories for query")
|
||||
return relevant_memories
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting relevant memory: {str(e)}")
|
||||
return []
|
||||
|
||||
def wait_for_indexing(self):
|
||||
logger.info(f"Waiting {self.indexing_wait_time}s for Zep indexing...")
|
||||
time.sleep(self.indexing_wait_time)
|
||||
|
||||
def get_session_summary(self) -> Dict[str, Any]:
|
||||
try:
|
||||
messages = self.zep_client.thread.get(thread_id=self.session_id)
|
||||
|
||||
if not messages or not messages.messages:
|
||||
return {"message_count": 0, "summary": "No messages in session"}
|
||||
|
||||
user_messages = [m for m in messages.messages if m.role == "user"]
|
||||
assistant_messages = [m for m in messages.messages if m.role == "assistant"]
|
||||
|
||||
summary = {
|
||||
"session_id": self.session_id,
|
||||
"user_id": self.user_id,
|
||||
"total_messages": len(messages.messages),
|
||||
"user_messages": len(user_messages),
|
||||
"assistant_messages": len(assistant_messages),
|
||||
"context_available": bool(self.get_conversation_context()),
|
||||
"last_interaction": messages.messages[0].created_at if messages.messages else None
|
||||
}
|
||||
|
||||
return summary
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting session summary: {str(e)}")
|
||||
return {"error": str(e)}
|
||||
|
||||
def clear_session(self):
|
||||
try:
|
||||
self.zep_client.thread.delete(self.session_id)
|
||||
self.zep_client.thread.create(thread_id=self.session_id, user_id=self.user_id)
|
||||
logger.info(f"Session {self.session_id} cleared and recreated")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing session: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from src.generation.rag import RAGGenerator, RAGResult
|
||||
|
||||
memory = NotebookMemoryLayer(
|
||||
user_id="test_user",
|
||||
session_id="test_session_123",
|
||||
create_new_session=True
|
||||
)
|
||||
|
||||
try:
|
||||
mock_rag_result = RAGResult(
|
||||
query="What are the main findings in the research?",
|
||||
response="The research shows three key findings [1, 2]. First, the methodology was effective [1]. Second, the results were significant [2].",
|
||||
sources_used=[
|
||||
{"source_file": "research_paper.pdf", "source_type": "pdf", "page_number": 5},
|
||||
{"source_file": "data_analysis.pdf", "source_type": "pdf", "page_number": 12}
|
||||
],
|
||||
retrieval_count=8
|
||||
)
|
||||
|
||||
memory.save_conversation_turn(mock_rag_result)
|
||||
memory.wait_for_indexing()
|
||||
context = memory.get_conversation_context()
|
||||
print(f"Conversation Context:\n{context}")
|
||||
|
||||
relevant = memory.get_relevant_memory("research findings")
|
||||
print(f"\nRelevant Memories: {len(relevant)} found")
|
||||
|
||||
summary = memory.get_session_summary()
|
||||
print(f"\nSession Summary: {summary}")
|
||||
|
||||
memory.save_user_preferences({
|
||||
"response_length": "detailed",
|
||||
"citation_style": "academic",
|
||||
"preferred_sources": ["pdf", "web"]
|
||||
})
|
||||
|
||||
print("Memory integration test completed successfully")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in memory integration test: {e}")
|
||||
@@ -0,0 +1,299 @@
|
||||
import logging
|
||||
import json
|
||||
from typing import List, Dict, Any
|
||||
from dataclasses import dataclass
|
||||
|
||||
from crewai import LLM
|
||||
from src.document_processing.doc_processor import DocumentProcessor
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PodcastScript:
|
||||
"""Represents a podcast script with metadata"""
|
||||
script: List[Dict[str, str]]
|
||||
source_document: str
|
||||
total_lines: int
|
||||
estimated_duration: str
|
||||
|
||||
def get_speaker_lines(self, speaker: str) -> List[str]:
|
||||
return [item[speaker] for item in self.script if speaker in item]
|
||||
|
||||
def to_json(self) -> str:
|
||||
return json.dumps({
|
||||
'script': self.script,
|
||||
'metadata': {
|
||||
'source_document': self.source_document,
|
||||
'total_lines': self.total_lines,
|
||||
'estimated_duration': self.estimated_duration
|
||||
}
|
||||
}, indent=2)
|
||||
|
||||
|
||||
class PodcastScriptGenerator:
|
||||
def __init__(self, openai_api_key: str, model_name: str = "gpt-4o-mini"):
|
||||
self.llm = LLM(
|
||||
model=f"openai/{model_name}",
|
||||
temperature=0.7,
|
||||
max_tokens=4000
|
||||
)
|
||||
self.doc_processor = DocumentProcessor()
|
||||
logger.info(f"Podcast script generator initialized with {model_name}")
|
||||
|
||||
def generate_script_from_document(
|
||||
self,
|
||||
document_path: str,
|
||||
podcast_style: str = "conversational",
|
||||
target_duration: str = "10 minutes"
|
||||
) -> PodcastScript:
|
||||
|
||||
logger.info(f"Generating podcast script from: {document_path}")
|
||||
|
||||
chunks = self.doc_processor.process_document(document_path)
|
||||
if not chunks:
|
||||
raise ValueError("No content extracted from document")
|
||||
|
||||
document_content = "\n\n".join([chunk.content for chunk in chunks])
|
||||
source_name = chunks[0].source_file
|
||||
script_data = self._generate_conversation_script(
|
||||
document_content,
|
||||
podcast_style,
|
||||
target_duration
|
||||
)
|
||||
|
||||
podcast_script = PodcastScript(
|
||||
script=script_data['script'],
|
||||
source_document=source_name,
|
||||
total_lines=len(script_data['script']),
|
||||
estimated_duration=target_duration
|
||||
)
|
||||
|
||||
logger.info(f"Generated script with {podcast_script.total_lines} lines")
|
||||
return podcast_script
|
||||
|
||||
def generate_script_from_text(
|
||||
self,
|
||||
text_content: str,
|
||||
source_name: str = "Text Input",
|
||||
podcast_style: str = "conversational",
|
||||
target_duration: str = "10 minutes"
|
||||
) -> PodcastScript:
|
||||
|
||||
logger.info("Generating podcast script from text input")
|
||||
|
||||
script_data = self._generate_conversation_script(
|
||||
text_content,
|
||||
podcast_style,
|
||||
target_duration
|
||||
)
|
||||
|
||||
podcast_script = PodcastScript(
|
||||
script=script_data['script'],
|
||||
source_document=source_name,
|
||||
total_lines=len(script_data['script']),
|
||||
estimated_duration=target_duration
|
||||
)
|
||||
|
||||
logger.info(f"Generated script with {podcast_script.total_lines} lines")
|
||||
return podcast_script
|
||||
|
||||
def generate_script_from_website(
|
||||
self,
|
||||
website_chunks: List[Any],
|
||||
source_url: str,
|
||||
podcast_style: str = "conversational",
|
||||
target_duration: str = "10 minutes"
|
||||
) -> PodcastScript:
|
||||
|
||||
logger.info(f"Generating podcast script from website: {source_url}")
|
||||
|
||||
if not website_chunks:
|
||||
raise ValueError("No website content provided")
|
||||
|
||||
website_content = "\n\n".join([chunk.content for chunk in website_chunks])
|
||||
script_data = self._generate_conversation_script(
|
||||
website_content,
|
||||
podcast_style,
|
||||
target_duration
|
||||
)
|
||||
|
||||
podcast_script = PodcastScript(
|
||||
script=script_data['script'],
|
||||
source_document=source_url,
|
||||
total_lines=len(script_data['script']),
|
||||
estimated_duration=target_duration
|
||||
)
|
||||
|
||||
logger.info(f"Generated website script with {podcast_script.total_lines} lines")
|
||||
return podcast_script
|
||||
|
||||
def _generate_conversation_script(
|
||||
self,
|
||||
document_content: str,
|
||||
podcast_style: str,
|
||||
target_duration: str
|
||||
) -> Dict[str, Any]:
|
||||
|
||||
style_prompts = {
|
||||
"conversational": "Create a natural, friendly conversation between two hosts discussing the document. They should build on each other's points and occasionally ask clarifying questions.",
|
||||
"educational": "Create an educational discussion where one speaker explains concepts and the other asks thoughtful questions to help clarify complex topics for listeners.",
|
||||
"interview": "Create an interview format where Speaker 1 acts as the interviewer asking questions and Speaker 2 provides detailed explanations from the document.",
|
||||
"debate": "Create a thoughtful discussion where speakers present different perspectives on the topics, maintaining respect while exploring various viewpoints."
|
||||
}
|
||||
|
||||
style_instruction = style_prompts.get(podcast_style, style_prompts["conversational"])
|
||||
|
||||
duration_guidelines = {
|
||||
"5 minutes": "Keep the conversation concise, focusing on 3-4 main points with brief explanations.",
|
||||
"10 minutes": "Cover the key topics thoroughly with good explanations and examples.",
|
||||
"15 minutes": "Provide comprehensive coverage with detailed discussions and multiple examples.",
|
||||
"20 minutes": "Create an in-depth exploration with extensive analysis and supporting details."
|
||||
}
|
||||
|
||||
duration_guide = duration_guidelines.get(target_duration, duration_guidelines["10 minutes"])
|
||||
|
||||
prompt = f"""Using the following document, create a podcast script for two speakers: 'Speaker 1' and 'Speaker 2'.
|
||||
|
||||
STYLE GUIDELINES:
|
||||
{style_instruction}
|
||||
|
||||
DURATION GUIDELINES:
|
||||
{duration_guide}
|
||||
|
||||
CONVERSATION RULES:
|
||||
1. Each speaker should speak for 2-4 sentences maximum before alternating
|
||||
2. The conversation should flow naturally with smooth transitions
|
||||
3. Use engaging, conversational language that's easy to understand
|
||||
4. Include brief introductions at the start and wrap-up at the end
|
||||
5. Break down complex concepts into digestible explanations
|
||||
6. Maintain professional grammar and punctuation throughout
|
||||
7. Make it engaging for listeners who haven't read the document
|
||||
|
||||
RESPONSE FORMAT:
|
||||
Respond with a valid JSON object containing a 'script' array. Each array element should be an object with either 'Speaker 1' or 'Speaker 2' as the key and their dialogue as the value.
|
||||
|
||||
Example format:
|
||||
{{
|
||||
"script": [
|
||||
{{"Speaker 1": "Welcome everyone to our podcast! Today we're diving into some fascinating insights from this document..."}},
|
||||
{{"Speaker 2": "Thanks for having me! I'm really excited to discuss this topic. The first thing that caught my attention was..."}}
|
||||
]
|
||||
}}
|
||||
|
||||
DOCUMENT CONTENT:
|
||||
{document_content[:8000]}
|
||||
|
||||
Generate an engaging {target_duration} podcast script now:"""
|
||||
|
||||
try:
|
||||
response = self.llm.call(prompt)
|
||||
script_data = json.loads(response)
|
||||
|
||||
if 'script' not in script_data or not isinstance(script_data['script'], list):
|
||||
raise ValueError("Invalid script format returned by LLM")
|
||||
|
||||
validated_script = self._validate_and_clean_script(script_data['script'])
|
||||
|
||||
return {'script': validated_script}
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to parse LLM response as JSON: {e}")
|
||||
response_clean = response.strip()
|
||||
if response_clean.startswith('```json'):
|
||||
response_clean = response_clean[7:-3]
|
||||
elif response_clean.startswith('```'):
|
||||
response_clean = response_clean[3:-3]
|
||||
|
||||
try:
|
||||
script_data = json.loads(response_clean)
|
||||
validated_script = self._validate_and_clean_script(script_data['script'])
|
||||
return {'script': validated_script}
|
||||
except:
|
||||
raise ValueError(f"Could not parse LLM response as valid JSON: {response}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating script: {str(e)}")
|
||||
raise
|
||||
|
||||
def _validate_and_clean_script(self, script: List[Dict[str, str]]) -> List[Dict[str, str]]:
|
||||
cleaned_script = []
|
||||
expected_speaker = "Speaker 1"
|
||||
for item in script:
|
||||
if not isinstance(item, dict) or len(item) != 1:
|
||||
continue
|
||||
|
||||
speaker, dialogue = next(iter(item.items()))
|
||||
speaker = speaker.strip()
|
||||
|
||||
if speaker not in ["Speaker 1", "Speaker 2"]:
|
||||
if "1" in speaker or "one" in speaker.lower():
|
||||
speaker = "Speaker 1"
|
||||
elif "2" in speaker or "two" in speaker.lower():
|
||||
speaker = "Speaker 2"
|
||||
else:
|
||||
speaker = expected_speaker
|
||||
|
||||
dialogue = dialogue.strip()
|
||||
if not dialogue:
|
||||
continue
|
||||
if not dialogue.endswith(('.', '!', '?')):
|
||||
dialogue += '.'
|
||||
|
||||
cleaned_script.append({speaker: dialogue})
|
||||
|
||||
expected_speaker = "Speaker 2" if expected_speaker == "Speaker 1" else "Speaker 1"
|
||||
|
||||
if len(cleaned_script) < 2:
|
||||
raise ValueError("Generated script is too short or invalid")
|
||||
|
||||
return cleaned_script
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
|
||||
openai_key = os.getenv("OPENAI_API_KEY")
|
||||
if not openai_key:
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
generator = PodcastScriptGenerator(openai_key)
|
||||
|
||||
try:
|
||||
sample_text = """
|
||||
Artificial Intelligence (AI) represents one of the most significant technological advances of our time.
|
||||
Machine learning, a subset of AI, enables computers to learn and improve from experience without being
|
||||
explicitly programmed for every task. Deep learning, which uses neural networks with multiple layers,
|
||||
has revolutionized fields like computer vision, natural language processing, and speech recognition.
|
||||
The applications are vast, from autonomous vehicles to medical diagnosis, and the potential impact on
|
||||
society is profound. However, ethical considerations around AI development, including bias, privacy,
|
||||
and job displacement, remain important challenges that need to be addressed as the technology continues to evolve.
|
||||
"""
|
||||
|
||||
script = generator.generate_script_from_text(
|
||||
sample_text,
|
||||
source_name="AI Overview",
|
||||
podcast_style="conversational",
|
||||
target_duration="5 minutes"
|
||||
)
|
||||
|
||||
print("Generated Podcast Script:")
|
||||
print("=" * 50)
|
||||
print(f"Source: {script.source_document}")
|
||||
print(f"Lines: {script.total_lines}")
|
||||
print(f"Duration: {script.estimated_duration}")
|
||||
print("\nScript:")
|
||||
|
||||
for i, line_dict in enumerate(script.script, 1):
|
||||
speaker, dialogue = next(iter(line_dict.items()))
|
||||
print(f"{i}. {speaker}: {dialogue}\n")
|
||||
|
||||
# Save to file
|
||||
with open("sample_podcast_script.json", "w") as f:
|
||||
f.write(script.to_json())
|
||||
print("Script saved to sample_podcast_script.json")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
@@ -0,0 +1,203 @@
|
||||
import logging
|
||||
import os
|
||||
import soundfile as sf
|
||||
from typing import List, Dict, Any
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from kokoro import KPipeline
|
||||
except ImportError:
|
||||
print("Kokoro not installed. Install with: pip install kokoro>=0.9.4")
|
||||
KPipeline = None
|
||||
|
||||
from src.podcast.script_generator import PodcastScript
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AudioSegment:
|
||||
"""Represents a single audio segment with metadata"""
|
||||
speaker: str
|
||||
text: str
|
||||
audio_data: Any
|
||||
duration: float
|
||||
file_path: str
|
||||
|
||||
|
||||
class PodcastTTSGenerator:
|
||||
def __init__(self, lang_code: str = 'a', sample_rate: int = 24000):
|
||||
if KPipeline is None:
|
||||
raise ImportError("Kokoro TTS not available. Install with: pip install kokoro>=0.9.4 soundfile")
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
self.pipeline = KPipeline(lang_code=lang_code)
|
||||
|
||||
self.speaker_voices = {
|
||||
"Speaker 1": "af_heart", # Female voice
|
||||
"Speaker 2": "am_liam" # Male voice
|
||||
}
|
||||
|
||||
logger.info(f"Kokoro TTS initialized with lang_code='{lang_code}', sample_rate={sample_rate}")
|
||||
|
||||
def generate_podcast_audio(
|
||||
self,
|
||||
podcast_script: PodcastScript,
|
||||
output_dir: str = "outputs/podcast_audio",
|
||||
combine_audio: bool = True
|
||||
) -> List[str]:
|
||||
|
||||
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info(f"Generating podcast audio for {podcast_script.total_lines} segments")
|
||||
logger.info(f"Output directory: {output_dir}")
|
||||
|
||||
audio_segments = []
|
||||
output_files = []
|
||||
|
||||
for i, line_dict in enumerate(podcast_script.script):
|
||||
speaker, dialogue = next(iter(line_dict.items()))
|
||||
|
||||
logger.info(f"Processing segment {i+1}/{podcast_script.total_lines}: {speaker}")
|
||||
|
||||
try:
|
||||
segment_audio = self._generate_single_segment(speaker, dialogue)
|
||||
segment_filename = f"segment_{i+1:03d}_{speaker.replace(' ', '_').lower()}.wav"
|
||||
segment_path = os.path.join(output_dir, segment_filename)
|
||||
|
||||
sf.write(segment_path, segment_audio, self.sample_rate)
|
||||
output_files.append(segment_path)
|
||||
|
||||
if combine_audio:
|
||||
audio_segment = AudioSegment(
|
||||
speaker=speaker,
|
||||
text=dialogue,
|
||||
audio_data=segment_audio,
|
||||
duration=len(segment_audio) / self.sample_rate,
|
||||
file_path=segment_path
|
||||
)
|
||||
audio_segments.append(audio_segment)
|
||||
|
||||
logger.info(f"✓ Generated segment {i+1}: {segment_filename}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"✗ Failed to generate segment {i+1}: {str(e)}")
|
||||
continue
|
||||
|
||||
if combine_audio and audio_segments:
|
||||
combined_path = self._combine_audio_segments(audio_segments, output_dir)
|
||||
output_files.append(combined_path)
|
||||
|
||||
logger.info(f"Podcast generation complete! Generated {len(output_files)} files")
|
||||
return output_files
|
||||
|
||||
def _generate_single_segment(self, speaker: str, text: str) -> Any:
|
||||
voice = self.speaker_voices.get(speaker, "af_heart")
|
||||
clean_text = self._clean_text_for_tts(text)
|
||||
|
||||
generator = self.pipeline(clean_text, voice=voice)
|
||||
|
||||
combined_audio = []
|
||||
for i, (gs, ps, audio) in enumerate(generator):
|
||||
combined_audio.append(audio)
|
||||
|
||||
if len(combined_audio) == 1:
|
||||
return combined_audio[0]
|
||||
else:
|
||||
import numpy as np
|
||||
return np.concatenate(combined_audio)
|
||||
|
||||
def _clean_text_for_tts(self, text: str) -> str:
|
||||
clean_text = text.strip()
|
||||
|
||||
clean_text = clean_text.replace("...", ".")
|
||||
clean_text = clean_text.replace("!!", "!")
|
||||
clean_text = clean_text.replace("??", "?")
|
||||
|
||||
if not clean_text.endswith(('.', '!', '?')):
|
||||
clean_text += '.'
|
||||
|
||||
return clean_text
|
||||
|
||||
def _combine_audio_segments(
|
||||
self,
|
||||
segments: List[AudioSegment],
|
||||
output_dir: str
|
||||
) -> str:
|
||||
logger.info(f"Combining {len(segments)} audio segments")
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
|
||||
pause_duration = 0.2 # seconds
|
||||
pause_samples = int(pause_duration * self.sample_rate)
|
||||
pause_audio = np.zeros(pause_samples, dtype=np.float32)
|
||||
|
||||
combined_audio = []
|
||||
for i, segment in enumerate(segments):
|
||||
combined_audio.append(segment.audio_data)
|
||||
|
||||
if i < len(segments) - 1:
|
||||
combined_audio.append(pause_audio)
|
||||
|
||||
final_audio = np.concatenate(combined_audio)
|
||||
|
||||
combined_filename = "complete_podcast.wav"
|
||||
combined_path = os.path.join(output_dir, combined_filename)
|
||||
sf.write(combined_path, final_audio, self.sample_rate)
|
||||
|
||||
duration = len(final_audio) / self.sample_rate
|
||||
logger.info(f"✓ Combined podcast saved: {combined_path} (Duration: {duration:.1f}s)")
|
||||
|
||||
return combined_path
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"✗ Failed to combine audio segments: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import json
|
||||
|
||||
try:
|
||||
tts_generator = PodcastTTSGenerator()
|
||||
|
||||
sample_script_data = {
|
||||
"script": [
|
||||
{"Speaker 1": "Welcome everyone to our podcast! Today we're exploring the fascinating world of artificial intelligence."},
|
||||
{"Speaker 2": "Thanks for having me! AI is indeed one of the most exciting technological developments of our time."},
|
||||
{"Speaker 1": "Let's start with machine learning. Can you explain what makes it so revolutionary?"},
|
||||
{"Speaker 2": "Absolutely! Machine learning allows computers to learn from data without being explicitly programmed for every single task."},
|
||||
{"Speaker 1": "That's incredible! And deep learning takes this even further, doesn't it?"},
|
||||
{"Speaker 2": "Exactly! Deep learning uses neural networks with multiple layers, revolutionizing computer vision and natural language processing."}
|
||||
]
|
||||
}
|
||||
|
||||
from src.podcast.script_generator import PodcastScript
|
||||
test_script = PodcastScript(
|
||||
script=sample_script_data["script"],
|
||||
source_document="AI Overview Test",
|
||||
total_lines=len(sample_script_data["script"]),
|
||||
estimated_duration="2 minutes"
|
||||
)
|
||||
|
||||
print("Generating podcast audio...")
|
||||
output_files = tts_generator.generate_podcast_audio(
|
||||
test_script,
|
||||
output_dir="./podcast_output",
|
||||
combine_audio=True
|
||||
)
|
||||
|
||||
print(f"\nGenerated files:")
|
||||
for file_path in output_files:
|
||||
print(f" - {file_path}")
|
||||
|
||||
print("\nPodcast TTS test completed successfully!")
|
||||
|
||||
except ImportError as e:
|
||||
print(f"Import error: {e}")
|
||||
print("Please install Kokoro TTS: pip install kokoro>=0.9.4")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
@@ -0,0 +1,372 @@
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from pymilvus import MilvusClient, DataType, connections, utility
|
||||
from src.embeddings.embedding_generator import EmbeddedChunk
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MilvusVectorDB:
|
||||
def __init__(
|
||||
self,
|
||||
db_path: str = "./milvus_lite.db",
|
||||
collection_name: str = "notebook_lm",
|
||||
embedding_dim: int = 384
|
||||
):
|
||||
self.db_path = db_path
|
||||
self.collection_name = collection_name
|
||||
self.embedding_dim = embedding_dim
|
||||
self.client = None
|
||||
self.collection_exists = False
|
||||
|
||||
self._initialize_client()
|
||||
self._setup_collection()
|
||||
|
||||
def _initialize_client(self):
|
||||
try:
|
||||
self.client = MilvusClient(uri=self.db_path)
|
||||
logger.info(f"Milvus client initialized with database: {self.db_path}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize Milvus client: {str(e)}")
|
||||
raise
|
||||
|
||||
def _setup_collection(self):
|
||||
try:
|
||||
if self.client.has_collection(collection_name=self.collection_name):
|
||||
logger.info(f"Collection '{self.collection_name}' already exists")
|
||||
self.collection_exists = True
|
||||
return
|
||||
|
||||
schema = self.client.create_schema(
|
||||
auto_id=False,
|
||||
enable_dynamic_field=True # Allow additional metadata fields
|
||||
)
|
||||
|
||||
# Primary key field (chunk_id)
|
||||
schema.add_field(
|
||||
field_name="id",
|
||||
datatype=DataType.VARCHAR,
|
||||
max_length=128,
|
||||
is_primary=True
|
||||
)
|
||||
|
||||
# Vector field for embeddings
|
||||
schema.add_field(
|
||||
field_name="vector",
|
||||
datatype=DataType.FLOAT_VECTOR,
|
||||
dim=self.embedding_dim
|
||||
)
|
||||
|
||||
# Essential fields for citations and content
|
||||
schema.add_field(
|
||||
field_name="content",
|
||||
datatype=DataType.VARCHAR,
|
||||
max_length=8192
|
||||
)
|
||||
|
||||
schema.add_field(
|
||||
field_name="source_file",
|
||||
datatype=DataType.VARCHAR,
|
||||
max_length=512
|
||||
)
|
||||
|
||||
schema.add_field(
|
||||
field_name="source_type",
|
||||
datatype=DataType.VARCHAR,
|
||||
max_length=32
|
||||
)
|
||||
|
||||
schema.add_field(
|
||||
field_name="page_number",
|
||||
datatype=DataType.INT32
|
||||
)
|
||||
|
||||
schema.add_field(
|
||||
field_name="chunk_index",
|
||||
datatype=DataType.INT32
|
||||
)
|
||||
|
||||
schema.add_field(
|
||||
field_name="start_char",
|
||||
datatype=DataType.INT32
|
||||
)
|
||||
|
||||
schema.add_field(
|
||||
field_name="end_char",
|
||||
datatype=DataType.INT32
|
||||
)
|
||||
|
||||
# JSON field for additional metadata
|
||||
schema.add_field(
|
||||
field_name="metadata",
|
||||
datatype=DataType.JSON
|
||||
)
|
||||
|
||||
schema.add_field(
|
||||
field_name="embedding_model",
|
||||
datatype=DataType.VARCHAR,
|
||||
max_length=128
|
||||
)
|
||||
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
schema=schema
|
||||
)
|
||||
|
||||
logger.info(f"Collection '{self.collection_name}' created successfully")
|
||||
self.collection_exists = True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error setting up collection: {str(e)}")
|
||||
raise
|
||||
|
||||
def create_index(
|
||||
self,
|
||||
use_binary_quantization: bool = False,
|
||||
nlist: int = 1024,
|
||||
enable_refine: bool = False,
|
||||
refine_type: str = "SQ8"
|
||||
):
|
||||
try:
|
||||
if not self.collection_exists:
|
||||
raise Exception("Collection does not exist. Setup collection first.")
|
||||
|
||||
index_params = self.client.prepare_index_params()
|
||||
|
||||
if use_binary_quantization:
|
||||
# IVF_RABITQ with binary quantization
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_type="IVF_RABITQ",
|
||||
index_name="vector_index",
|
||||
metric_type="L2",
|
||||
params={
|
||||
"nlist": nlist,
|
||||
"refine": enable_refine,
|
||||
"refine_type": refine_type if enable_refine else None
|
||||
}
|
||||
)
|
||||
logger.info(f"Creating IVF_RABITQ index with nlist={nlist}, refine={enable_refine}")
|
||||
else:
|
||||
# Fallback to IVF_FLAT if BQ not supported
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_type="IVF_FLAT",
|
||||
index_name="vector_index",
|
||||
metric_type="L2",
|
||||
# params={"nlist": nlist}
|
||||
)
|
||||
logger.info(f"Creating IVF_FLAT index with nlist={nlist}")
|
||||
|
||||
self.client.create_index(
|
||||
collection_name=self.collection_name,
|
||||
index_params=index_params
|
||||
)
|
||||
|
||||
logger.info("Index created successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating index: {str(e)}")
|
||||
raise
|
||||
|
||||
def insert_embeddings(self, embedded_chunks: List[EmbeddedChunk]) -> List[str]:
|
||||
if not embedded_chunks:
|
||||
return []
|
||||
try:
|
||||
data = []
|
||||
for embedded_chunk in embedded_chunks:
|
||||
chunk_data = embedded_chunk.to_vector_db_format()
|
||||
chunk_data['page_number'] = chunk_data['page_number'] or -1
|
||||
chunk_data['start_char'] = chunk_data['start_char'] or -1
|
||||
chunk_data['end_char'] = chunk_data['end_char'] or -1
|
||||
|
||||
if isinstance(chunk_data['metadata'], dict):
|
||||
chunk_data['metadata'] = chunk_data['metadata']
|
||||
|
||||
data.append(chunk_data)
|
||||
|
||||
result = self.client.insert(
|
||||
collection_name=self.collection_name,
|
||||
data=data
|
||||
)
|
||||
|
||||
inserted_ids = [item['id'] for item in data]
|
||||
logger.info(f"Inserted {len(inserted_ids)} embeddings into database")
|
||||
|
||||
return inserted_ids
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error inserting embeddings: {str(e)}")
|
||||
raise
|
||||
|
||||
def search(
|
||||
self,
|
||||
query_vector: List[float],
|
||||
limit: int = 10,
|
||||
nprobe: int = 128,
|
||||
rbq_query_bits: int = 0,
|
||||
refine_k: float = 1.0,
|
||||
filter_expr: Optional[str] = None,
|
||||
use_binary_quantization: bool = False
|
||||
) -> List[Dict[str, Any]]:
|
||||
try:
|
||||
if use_binary_quantization:
|
||||
search_params = {
|
||||
"params": {
|
||||
"nprobe": nprobe,
|
||||
"rbq_query_bits": rbq_query_bits,
|
||||
"refine_k": refine_k
|
||||
}
|
||||
}
|
||||
else:
|
||||
search_params = {
|
||||
"params": {
|
||||
"nprobe": nprobe
|
||||
}
|
||||
}
|
||||
|
||||
# Perform vector similarity search
|
||||
results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
data=[query_vector],
|
||||
anns_field="vector",
|
||||
limit=limit,
|
||||
search_params=search_params,
|
||||
filter=filter_expr,
|
||||
output_fields=[
|
||||
"content", "source_file", "source_type", "page_number",
|
||||
"chunk_index", "start_char", "end_char", "metadata", "embedding_model"
|
||||
]
|
||||
)
|
||||
|
||||
formatted_results = []
|
||||
if results and len(results) > 0:
|
||||
for result in results[0]:
|
||||
formatted_result = {
|
||||
'id': result['id'],
|
||||
'score': result['distance'],
|
||||
'content': result['entity']['content'],
|
||||
'citation': {
|
||||
'source_file': result['entity']['source_file'],
|
||||
'source_type': result['entity']['source_type'],
|
||||
'page_number': result['entity']['page_number'] if result['entity']['page_number'] != -1 else None,
|
||||
'chunk_index': result['entity']['chunk_index'],
|
||||
'start_char': result['entity']['start_char'] if result['entity']['start_char'] != -1 else None,
|
||||
'end_char': result['entity']['end_char'] if result['entity']['end_char'] != -1 else None,
|
||||
},
|
||||
'metadata': result['entity']['metadata'],
|
||||
'embedding_model': result['entity']['embedding_model']
|
||||
}
|
||||
formatted_results.append(formatted_result)
|
||||
|
||||
logger.info(f"Search completed: {len(formatted_results)} results found")
|
||||
return formatted_results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during search: {str(e)}")
|
||||
raise
|
||||
|
||||
def delete_collection(self):
|
||||
try:
|
||||
if self.client.has_collection(collection_name=self.collection_name):
|
||||
self.client.drop_collection(collection_name=self.collection_name)
|
||||
logger.info(f"Collection '{self.collection_name}' deleted")
|
||||
self.collection_exists = False
|
||||
else:
|
||||
logger.info(f"Collection '{self.collection_name}' does not exist")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting collection: {str(e)}")
|
||||
raise
|
||||
|
||||
def get_chunk_by_id(self, chunk_id: str) -> Optional[Dict[str, Any]]:
|
||||
try:
|
||||
if not self.collection_exists:
|
||||
logger.warning("Collection does not exist")
|
||||
return None
|
||||
|
||||
logger.info(f"Attempting to retrieve chunk with ID: {chunk_id}")
|
||||
|
||||
results = self.client.query(
|
||||
collection_name=self.collection_name,
|
||||
filter=f'id == "{chunk_id}"',
|
||||
output_fields=["id", "content", "metadata", "source_file", "source_type", "page_number", "chunk_index"]
|
||||
)
|
||||
|
||||
logger.info(f"Query returned {len(results) if results else 0} results")
|
||||
|
||||
if results and len(results) > 0:
|
||||
chunk_data = results[0]
|
||||
logger.info(f"Successfully retrieved chunk: {chunk_data.get('id')}")
|
||||
|
||||
metadata = chunk_data.get("metadata", {})
|
||||
if isinstance(metadata, str):
|
||||
try:
|
||||
metadata = json.loads(metadata)
|
||||
except:
|
||||
metadata = {}
|
||||
|
||||
return {
|
||||
"id": chunk_data.get("id"),
|
||||
"content": chunk_data.get("content"),
|
||||
"metadata": metadata,
|
||||
"source_file": chunk_data.get("source_file"),
|
||||
"source_type": chunk_data.get("source_type"),
|
||||
"page_number": chunk_data.get("page_number"),
|
||||
"chunk_index": chunk_data.get("chunk_index")
|
||||
}
|
||||
|
||||
logger.warning(f"No chunk found with ID: {chunk_id}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving chunk by ID {chunk_id}: {str(e)}")
|
||||
logger.error(f"Exception details: {type(e).__name__}: {str(e)}")
|
||||
return None
|
||||
|
||||
def close(self):
|
||||
try:
|
||||
if self.client:
|
||||
self.client.close()
|
||||
logger.info("Milvus client connection closed")
|
||||
except Exception as e:
|
||||
logger.error(f"Error closing connection: {str(e)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from src.document_processing.doc_processor import DocumentProcessor
|
||||
from src.embeddings.embedding_generator import EmbeddingGenerator
|
||||
|
||||
doc_processor = DocumentProcessor()
|
||||
embedding_generator = EmbeddingGenerator()
|
||||
vector_db = MilvusVectorDB()
|
||||
|
||||
try:
|
||||
chunks = doc_processor.process_document("data/raft.pdf")
|
||||
embedded_chunks = embedding_generator.generate_embeddings(chunks)
|
||||
vector_db.create_index()
|
||||
|
||||
inserted_ids = vector_db.insert_embeddings(embedded_chunks)
|
||||
print(f"Inserted {len(inserted_ids)} embeddings")
|
||||
|
||||
query_text = "What is the main topic?"
|
||||
query_vector = embedding_generator.generate_query_embedding(query_text)
|
||||
|
||||
search_results = vector_db.search(query_vector.tolist(), limit=5)
|
||||
|
||||
for i, result in enumerate(search_results):
|
||||
print(f"\nResult {i+1}:")
|
||||
print(f"Score: {result['score']:.4f}")
|
||||
print(f"Content: {result['content'][:200]}...")
|
||||
print(f"Citation: {result['citation']}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in example: {e}")
|
||||
|
||||
finally:
|
||||
vector_db.close()
|
||||
@@ -0,0 +1,252 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import List, Dict, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse, urljoin
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
from firecrawl import Firecrawl
|
||||
from src.document_processing.doc_processor import DocumentChunk
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WebPageData:
|
||||
"""Represents scraped web page data with additional metadata"""
|
||||
url: str
|
||||
title: str
|
||||
content: str
|
||||
metadata: Dict[str, Any]
|
||||
success: bool
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
class WebScraper:
|
||||
def __init__(self, api_key: str):
|
||||
self.api_key = api_key
|
||||
self.app = Firecrawl(api_key=api_key)
|
||||
|
||||
logger.info("WebScraper initialized with Firecrawl")
|
||||
|
||||
def scrape_url(
|
||||
self,
|
||||
url: str,
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 100,
|
||||
wait_for_results: int = 30
|
||||
) -> List[DocumentChunk]:
|
||||
|
||||
if not self._is_valid_url(url):
|
||||
raise ValueError(f"Invalid URL format: {url}")
|
||||
|
||||
logger.info(f"Scraping URL: {url}")
|
||||
|
||||
try:
|
||||
scrape_params = {
|
||||
'formats': ['markdown', 'html'],
|
||||
'timeout': wait_for_results * 1000
|
||||
}
|
||||
|
||||
result = self.app.scrape(url, **scrape_params)
|
||||
page_data = self._process_firecrawl_result(result, url)
|
||||
|
||||
chunks = self._create_chunks_from_web_content(
|
||||
page_data,
|
||||
chunk_size,
|
||||
chunk_overlap
|
||||
)
|
||||
|
||||
logger.info(f"Successfully scraped {url}: {len(chunks)} chunks created")
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error scraping URL {url}: {str(e)}")
|
||||
raise
|
||||
|
||||
def _process_firecrawl_result(self, result: Dict[str, Any], url: str) -> WebPageData:
|
||||
try:
|
||||
content = result.markdown
|
||||
metadata_dict = result.metadata_dict
|
||||
metadata = {
|
||||
'scraped_at': datetime.now().isoformat(),
|
||||
'original_url': url,
|
||||
'title': metadata_dict.get('title', ''),
|
||||
'description': metadata_dict.get('description', ''),
|
||||
'keywords': metadata_dict.get('keywords', []),
|
||||
'language': metadata_dict.get('language', 'en'),
|
||||
'word_count': len(content.split()) if content else 0,
|
||||
'character_count': len(content) if content else 0,
|
||||
'domain': urlparse(url).netloc
|
||||
}
|
||||
|
||||
return WebPageData(
|
||||
url=url,
|
||||
title=metadata['title'] or f"Web Page - {metadata['domain']}",
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
success=True
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing Firecrawl result: {str(e)}")
|
||||
return WebPageData(
|
||||
url=url,
|
||||
title=f"Error - {urlparse(url).netloc}",
|
||||
content="",
|
||||
metadata={'error': str(e), 'scraped_at': datetime.now().isoformat()},
|
||||
success=False,
|
||||
error=str(e)
|
||||
)
|
||||
|
||||
def _create_chunks_from_web_content(
|
||||
self,
|
||||
page_data: WebPageData,
|
||||
chunk_size: int,
|
||||
chunk_overlap: int
|
||||
) -> List[DocumentChunk]:
|
||||
|
||||
if not page_data.success or not page_data.content.strip():
|
||||
logger.warning(f"No content to process for {page_data.url}")
|
||||
return []
|
||||
|
||||
chunks = []
|
||||
content = page_data.content
|
||||
start = 0
|
||||
chunk_index = 0
|
||||
|
||||
while start < len(content):
|
||||
end = min(start + chunk_size, len(content))
|
||||
if end < len(content):
|
||||
last_double_newline = content.rfind('\n\n', start, end)
|
||||
if last_double_newline > start + chunk_size * 0.3:
|
||||
end = last_double_newline + 2
|
||||
else:
|
||||
last_period = content.rfind('.', start, end)
|
||||
if last_period > start + chunk_size * 0.5:
|
||||
end = last_period + 1
|
||||
|
||||
chunk_text = content[start:end].strip()
|
||||
|
||||
if chunk_text:
|
||||
chunk_metadata = page_data.metadata.copy()
|
||||
chunk_metadata.update({
|
||||
'chunk_character_start': start,
|
||||
'chunk_character_end': end - 1,
|
||||
'url_fragment': f"{page_data.url}#chunk-{chunk_index}"
|
||||
})
|
||||
|
||||
chunk = DocumentChunk(
|
||||
content=chunk_text,
|
||||
source_file=page_data.title,
|
||||
source_type='web',
|
||||
page_number=None,
|
||||
chunk_index=chunk_index,
|
||||
start_char=start,
|
||||
end_char=end-1,
|
||||
metadata=chunk_metadata
|
||||
)
|
||||
|
||||
chunks.append(chunk)
|
||||
chunk_index += 1
|
||||
|
||||
start = max(start + chunk_size - chunk_overlap, end)
|
||||
|
||||
return chunks
|
||||
|
||||
def batch_scrape_urls(
|
||||
self,
|
||||
urls: List[str],
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 100,
|
||||
delay_between_requests: float = 1.0
|
||||
) -> List[List[DocumentChunk]]:
|
||||
|
||||
all_chunks = []
|
||||
for i, url in enumerate(urls):
|
||||
try:
|
||||
chunks = self.scrape_url(url, chunk_size, chunk_overlap)
|
||||
all_chunks.append(chunks)
|
||||
logger.info(f"Successfully scraped {url}: {len(chunks)} chunks")
|
||||
|
||||
if i < len(urls) - 1:
|
||||
time.sleep(delay_between_requests)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to scrape {url}: {str(e)}")
|
||||
all_chunks.append([])
|
||||
|
||||
total_chunks = sum(len(chunks) for chunks in all_chunks)
|
||||
logger.info(f"Batch scraping complete: {total_chunks} total chunks from {len(urls)} URLs")
|
||||
|
||||
return all_chunks
|
||||
|
||||
def get_url_preview(self, url: str) -> Dict[str, Any]:
|
||||
try:
|
||||
result = self.app.scrape(url, **{
|
||||
'formats': ['markdown'],
|
||||
'timeout': 10000
|
||||
})
|
||||
|
||||
content = result.markdown
|
||||
metadata_dict = result.metadata_dict
|
||||
|
||||
preview_info = {
|
||||
'url': url,
|
||||
'title': metadata_dict.get('title', ''),
|
||||
'description': metadata_dict.get('description', ''),
|
||||
'word_count': len(content.split()) if content else 0,
|
||||
'character_count': len(content) if content else 0,
|
||||
'domain': urlparse(url).netloc,
|
||||
'content_preview': content[:500] + '...' if len(content) > 500 else content,
|
||||
'language': metadata_dict.get('language', 'unknown')
|
||||
}
|
||||
return preview_info
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting URL preview: {str(e)}")
|
||||
return {'error': str(e)}
|
||||
|
||||
def _is_valid_url(self, url: str) -> bool:
|
||||
try:
|
||||
result = urlparse(url)
|
||||
return all([result.scheme, result.netloc])
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
api_key = os.getenv("FIRECRAWL_API_KEY")
|
||||
if not api_key:
|
||||
print("Please set FIRECRAWL_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
scraper = WebScraper(api_key)
|
||||
|
||||
try:
|
||||
test_url = "https://blog.dailydoseofds.com/p/5-chunking-strategies-for-rag"
|
||||
preview = scraper.get_url_preview(test_url)
|
||||
print(f"URL Preview: {preview}")
|
||||
|
||||
chunks = scraper.scrape_url(test_url)
|
||||
print(f"\nScraping Results:")
|
||||
print(f"Generated {len(chunks)} chunks")
|
||||
|
||||
for i, chunk in enumerate(chunks[:3]):
|
||||
print(f"\nChunk {i+1}:")
|
||||
print(f"Content: {chunk.content[:200]}...")
|
||||
print(f"Source: {chunk.source_file}")
|
||||
print(f"URL: {chunk.metadata.get('original_url', 'N/A')}")
|
||||
print(f"Citation: [Source: {chunk.source_file}, Type: Web]")
|
||||
|
||||
urls = ["https://example.com/page1", "https://example.com/page2"]
|
||||
batch_results = scraper.batch_scrape_urls(urls)
|
||||
|
||||
total_chunks = sum(len(chunks) for chunks in batch_results)
|
||||
print(f"\nBatch Results: {total_chunks} total chunks from {len(urls)} URLs")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in scraping example: {e}")
|
||||
Reference in New Issue
Block a user