mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-04 22:02:16 +02:00
feat(youtube): integrate YouTube video processing connector
- Added support for processing YouTube videos, including transcript extraction and document creation. - Implemented a new background task for adding YouTube video documents. - Enhanced the connector service to search for YouTube videos and return relevant results. - Updated frontend components to include YouTube video options in the dashboard and connector sources. - Added necessary dependencies for YouTube transcript API.
This commit is contained in:
DESKTOP-RTLN3BA\$punk
parent
753f40dfea
commit
b43272a115
13 changed files with 608 additions and 18 deletions
|
|
@ -6,7 +6,7 @@ from app.db import get_async_session, User, SearchSpace, Document, DocumentType
|
|||
from app.schemas import DocumentsCreate, DocumentUpdate, DocumentRead
|
||||
from app.users import current_active_user
|
||||
from app.utils.check_ownership import check_ownership
|
||||
from app.tasks.background_tasks import add_extension_received_document, add_received_file_document, add_crawled_url_document
|
||||
from app.tasks.background_tasks import add_extension_received_document, add_received_file_document, add_crawled_url_document, add_youtube_video_document
|
||||
# Force asyncio to use standard event loop before unstructured imports
|
||||
import asyncio
|
||||
try:
|
||||
|
|
@ -368,8 +368,7 @@ async def process_youtube_video_with_new_session(
|
|||
|
||||
async with async_session_maker() as session:
|
||||
try:
|
||||
# TODO: Implement YouTube video processing
|
||||
print("Processing YouTube video with new session")
|
||||
await add_youtube_video_document(session, url, search_space_id)
|
||||
except Exception as e:
|
||||
import logging
|
||||
logging.error(f"Error processing YouTube video: {str(e)}")
|
||||
|
|
|
|||
|
|
@ -244,3 +244,142 @@ async def add_received_file_document(
|
|||
except Exception as e:
|
||||
await session.rollback()
|
||||
raise RuntimeError(f"Failed to process file document: {str(e)}")
|
||||
|
||||
|
||||
async def add_youtube_video_document(
|
||||
session: AsyncSession,
|
||||
url: str,
|
||||
search_space_id: int
|
||||
):
|
||||
"""
|
||||
Process a YouTube video URL, extract transcripts, and add as document.
|
||||
"""
|
||||
try:
|
||||
from youtube_transcript_api import YouTubeTranscriptApi
|
||||
|
||||
# Extract video ID from URL
|
||||
def get_youtube_video_id(url: str):
|
||||
from urllib.parse import urlparse, parse_qs
|
||||
|
||||
parsed_url = urlparse(url)
|
||||
hostname = parsed_url.hostname
|
||||
|
||||
if hostname == "youtu.be":
|
||||
return parsed_url.path[1:]
|
||||
if hostname in ("www.youtube.com", "youtube.com"):
|
||||
if parsed_url.path == "/watch":
|
||||
query_params = parse_qs(parsed_url.query)
|
||||
return query_params.get("v", [None])[0]
|
||||
if parsed_url.path.startswith("/embed/"):
|
||||
return parsed_url.path.split("/")[2]
|
||||
if parsed_url.path.startswith("/v/"):
|
||||
return parsed_url.path.split("/")[2]
|
||||
return None
|
||||
|
||||
# Get video ID
|
||||
video_id = get_youtube_video_id(url)
|
||||
if not video_id:
|
||||
raise ValueError(f"Could not extract video ID from URL: {url}")
|
||||
|
||||
# Get video metadata
|
||||
import json
|
||||
from urllib.parse import urlencode
|
||||
from urllib.request import urlopen
|
||||
|
||||
params = {"format": "json", "url": f"https://www.youtube.com/watch?v={video_id}"}
|
||||
oembed_url = "https://www.youtube.com/oembed"
|
||||
query_string = urlencode(params)
|
||||
full_url = oembed_url + "?" + query_string
|
||||
|
||||
with urlopen(full_url) as response:
|
||||
response_text = response.read()
|
||||
video_data = json.loads(response_text.decode())
|
||||
|
||||
# Get video transcript
|
||||
try:
|
||||
captions = YouTubeTranscriptApi.get_transcript(video_id)
|
||||
# Include complete caption information with timestamps
|
||||
transcript_segments = []
|
||||
for line in captions:
|
||||
start_time = line.get("start", 0)
|
||||
duration = line.get("duration", 0)
|
||||
text = line.get("text", "")
|
||||
timestamp = f"[{start_time:.2f}s-{start_time + duration:.2f}s]"
|
||||
transcript_segments.append(f"{timestamp} {text}")
|
||||
transcript_text = "\n".join(transcript_segments)
|
||||
except Exception as e:
|
||||
transcript_text = f"No captions available for this video. Error: {str(e)}"
|
||||
|
||||
# Format document metadata in a more maintainable way
|
||||
metadata_sections = [
|
||||
("METADATA", [
|
||||
f"TITLE: {video_data.get('title', 'YouTube Video')}",
|
||||
f"URL: {url}",
|
||||
f"VIDEO_ID: {video_id}",
|
||||
f"AUTHOR: {video_data.get('author_name', 'Unknown')}",
|
||||
f"THUMBNAIL: {video_data.get('thumbnail_url', '')}"
|
||||
]),
|
||||
("CONTENT", [
|
||||
"FORMAT: transcript",
|
||||
"TEXT_START",
|
||||
transcript_text,
|
||||
"TEXT_END"
|
||||
])
|
||||
]
|
||||
|
||||
# Build the document string more efficiently
|
||||
document_parts = []
|
||||
document_parts.append("<DOCUMENT>")
|
||||
|
||||
for section_title, section_content in metadata_sections:
|
||||
document_parts.append(f"<{section_title}>")
|
||||
document_parts.extend(section_content)
|
||||
document_parts.append(f"</{section_title}>")
|
||||
|
||||
document_parts.append("</DOCUMENT>")
|
||||
combined_document_string = '\n'.join(document_parts)
|
||||
|
||||
# Generate summary
|
||||
summary_chain = SUMMARY_PROMPT_TEMPLATE | config.long_context_llm_instance
|
||||
summary_result = await summary_chain.ainvoke({"document": combined_document_string})
|
||||
summary_content = summary_result.content
|
||||
summary_embedding = config.embedding_model_instance.embed(summary_content)
|
||||
|
||||
# Process chunks
|
||||
chunks = [
|
||||
Chunk(content=chunk.text, embedding=chunk.embedding)
|
||||
for chunk in config.chunker_instance.chunk(transcript_text)
|
||||
]
|
||||
|
||||
# Create document
|
||||
from app.db import Document, DocumentType
|
||||
|
||||
document = Document(
|
||||
title=video_data.get("title", "YouTube Video"),
|
||||
document_type=DocumentType.YOUTUBE_VIDEO,
|
||||
document_metadata={
|
||||
"url": url,
|
||||
"video_id": video_id,
|
||||
"video_title": video_data.get("title", "YouTube Video"),
|
||||
"author": video_data.get("author_name", "Unknown"),
|
||||
"thumbnail": video_data.get("thumbnail_url", "")
|
||||
},
|
||||
content=summary_content,
|
||||
embedding=summary_embedding,
|
||||
chunks=chunks,
|
||||
search_space_id=search_space_id
|
||||
)
|
||||
|
||||
session.add(document)
|
||||
await session.commit()
|
||||
await session.refresh(document)
|
||||
|
||||
return document
|
||||
except SQLAlchemyError as db_error:
|
||||
await session.rollback()
|
||||
raise db_error
|
||||
except Exception as e:
|
||||
await session.rollback()
|
||||
import logging
|
||||
logging.error(f"Failed to process YouTube video: {str(e)}")
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -59,6 +59,32 @@ async def stream_connector_search_results(
|
|||
|
||||
# Process each selected connector
|
||||
for connector in selected_connectors:
|
||||
if connector == "YOUTUBE_VIDEO":
|
||||
# Send terminal message about starting search
|
||||
yield streaming_service.add_terminal_message("Starting to search for youtube videos...")
|
||||
|
||||
# Search for YouTube videos using reformulated query
|
||||
result_object, youtube_chunks = await connector_service.search_youtube(
|
||||
user_query=reformulated_query,
|
||||
user_id=user_id,
|
||||
search_space_id=search_space_id,
|
||||
top_k=TOP_K
|
||||
)
|
||||
|
||||
# Send terminal message about search results
|
||||
yield streaming_service.add_terminal_message(
|
||||
f"Found {len(result_object['sources'])} relevant YouTube videos",
|
||||
"success"
|
||||
)
|
||||
|
||||
# Update sources
|
||||
all_sources.append(result_object)
|
||||
yield streaming_service.update_sources(all_sources)
|
||||
|
||||
# Add documents to collection
|
||||
all_raw_documents.extend(youtube_chunks)
|
||||
|
||||
|
||||
# Extension Docs
|
||||
if connector == "EXTENSION":
|
||||
# Send terminal message about starting search
|
||||
|
|
|
|||
|
|
@ -479,4 +479,83 @@ class ConnectorService:
|
|||
"sources": sources_list,
|
||||
}
|
||||
|
||||
return result_object, extension_chunks
|
||||
return result_object, extension_chunks
|
||||
|
||||
async def search_youtube(self, user_query: str, user_id: int, search_space_id: int, top_k: int = 20) -> tuple:
|
||||
"""
|
||||
Search for YouTube videos and return both the source information and langchain documents
|
||||
|
||||
Args:
|
||||
user_query: The user's query
|
||||
user_id: The user's ID
|
||||
search_space_id: The search space ID to search in
|
||||
top_k: Maximum number of results to return
|
||||
|
||||
Returns:
|
||||
tuple: (sources_info, langchain_documents)
|
||||
"""
|
||||
youtube_chunks = await self.retriever.hybrid_search(
|
||||
query_text=user_query,
|
||||
top_k=top_k,
|
||||
user_id=user_id,
|
||||
search_space_id=search_space_id,
|
||||
document_type="YOUTUBE_VIDEO"
|
||||
)
|
||||
|
||||
# Map youtube_chunks to the required format
|
||||
mapped_sources = {}
|
||||
for i, chunk in enumerate(youtube_chunks):
|
||||
# Fix for UI
|
||||
youtube_chunks[i]['document']['id'] = self.source_id_counter
|
||||
|
||||
# Extract document metadata
|
||||
document = chunk.get('document', {})
|
||||
metadata = document.get('metadata', {})
|
||||
|
||||
# Extract YouTube-specific metadata
|
||||
video_title = metadata.get('video_title', 'Untitled Video')
|
||||
video_id = metadata.get('video_id', '')
|
||||
channel_name = metadata.get('channel_name', '')
|
||||
published_date = metadata.get('published_date', '')
|
||||
|
||||
# Create a more descriptive title for YouTube videos
|
||||
title = video_title
|
||||
if channel_name:
|
||||
title += f" - {channel_name}"
|
||||
|
||||
# Create a more descriptive description for YouTube videos
|
||||
description = metadata.get('description', chunk.get('content', '')[:100])
|
||||
if len(description) == 100:
|
||||
description += "..."
|
||||
|
||||
# For URL, construct a URL to the YouTube video
|
||||
url = f"https://www.youtube.com/watch?v={video_id}" if video_id else ""
|
||||
|
||||
source = {
|
||||
"id": self.source_id_counter,
|
||||
"title": title,
|
||||
"description": description,
|
||||
"url": url,
|
||||
"video_id": video_id, # Additional field for YouTube videos
|
||||
"channel_name": channel_name # Additional field for YouTube videos
|
||||
}
|
||||
|
||||
self.source_id_counter += 1
|
||||
|
||||
# Use video_id as a unique identifier for tracking unique sources
|
||||
source_key = video_id or f"youtube_{i}"
|
||||
if source_key and source_key not in mapped_sources:
|
||||
mapped_sources[source_key] = source
|
||||
|
||||
# Convert to list of sources
|
||||
sources_list = list(mapped_sources.values())
|
||||
|
||||
# Create result object
|
||||
result_object = {
|
||||
"id": 6, # Assign a unique ID for the YouTube connector
|
||||
"name": "YouTube Videos",
|
||||
"type": "YOUTUBE_VIDEO",
|
||||
"sources": sources_list,
|
||||
}
|
||||
|
||||
return result_object, youtube_chunks
|
||||
|
|
@ -26,4 +26,5 @@ dependencies = [
|
|||
"unstructured[all-docs]>=0.16.25",
|
||||
"uvicorn[standard]>=0.34.0",
|
||||
"validators>=0.34.0",
|
||||
"youtube-transcript-api>=1.0.3",
|
||||
]
|
||||
|
|
|
|||
24
surfsense_backend/uv.lock
generated
24
surfsense_backend/uv.lock
generated
|
|
@ -580,6 +580,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/c3/be/d0d44e092656fe7a06b55e6103cbce807cdbdee17884a5367c68c9860853/dataclasses_json-0.6.7-py3-none-any.whl", hash = "sha256:0dbf33f26c8d5305befd61b39d2b3414e8a407bedc2834dea9b8d642666fb40a", size = 28686 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "defusedxml"
|
||||
version = "0.7.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/0f/d5/c66da9b79e5bdb124974bfe172b4daf3c984ebd9c2a06e2b8a4dc7331c72/defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69", size = 75520 }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl", hash = "sha256:a352e7e428770286cc899e2542b6cdaedb2b4953ff269a210103ec58f6198a61", size = 25604 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "deprecated"
|
||||
version = "1.2.18"
|
||||
|
|
@ -3240,6 +3249,7 @@ dependencies = [
|
|||
{ name = "unstructured-client" },
|
||||
{ name = "uvicorn", extra = ["standard"] },
|
||||
{ name = "validators" },
|
||||
{ name = "youtube-transcript-api" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
|
|
@ -3265,6 +3275,7 @@ requires-dist = [
|
|||
{ name = "unstructured-client", specifier = ">=0.30.0" },
|
||||
{ name = "uvicorn", extras = ["standard"], specifier = ">=0.34.0" },
|
||||
{ name = "validators", specifier = ">=0.34.0" },
|
||||
{ name = "youtube-transcript-api", specifier = ">=1.0.3" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3919,6 +3930,19 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/f5/4b/a06e0ec3d155924f77835ed2d167ebd3b211a7b0853da1cf8d8414d784ef/yarl-1.18.3-py3-none-any.whl", hash = "sha256:b57f4f58099328dfb26c6a771d09fb20dbbae81d20cfb66141251ea063bd101b", size = 45109 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "youtube-transcript-api"
|
||||
version = "1.0.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "defusedxml" },
|
||||
{ name = "requests" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b0/32/f60d87a99c05a53604c58f20f670c7ea6262b55e0bbeb836ffe4550b248b/youtube_transcript_api-1.0.3.tar.gz", hash = "sha256:902baf90e7840a42e1e148335e09fe5575dbff64c81414957aea7038e8a4db46", size = 2153252 }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f0/44/40c03bb0f8bddfb9d2beff2ed31641f52d96c287ba881d20e0c074784ac2/youtube_transcript_api-1.0.3-py3-none-any.whl", hash = "sha256:d1874e57de65cf14c9d7d09b2b37c814d6287fa0e770d4922c4cd32a5b3f6c47", size = 2169911 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zipp"
|
||||
version = "3.21.0"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue