完全基于 AWS 的视频内容搜索引擎 VideoSearch

# Author:Li Kaiyuan
# Version: 0.2

import boto3
import time

#lambda entry point
def lambda_handler(event, context):
    bucket_name = event['Records'][0]['s3']['bucket']['name']
    source_s3_key = event['Records'][0]['s3']['object']['key']
    delimiter = source_s3_key.find('/')
    file_name = source_s3_key[delimiter+1:]
    job_uri = 's3://'+ bucket_name + '/' + source_s3_key
    
    timestamp = time.time()
    local_time = time.localtime(timestamp)
    formatted_time = time.strftime("%Y-%m-%d-%H-%M-%S", local_time)
    job_name = "conv-"+file_name[0:-4]+'-'+formatted_time
    
    transcribe = boto3.client('transcribe')
    transcribe.start_transcription_job(
        TranscriptionJobName=job_name,
        MediaFormat='mp4',
        Media={
            'MediaFileUri': job_uri
        },
        OutputBucketName=bucket_name,
        OutputKey='output/'+file_name[0:-4]+formatted_time,
        Settings={
            'ShowSpeakerLabels': True,
            'MaxSpeakerLabels': 10,
            'ChannelIdentification': False
        },
        IdentifyMultipleLanguages=True,
        Subtitles={
            'Formats': [
                'srt',
            ],
            'OutputStartIndex': 1
        },
        Tags=[
            {
                'Key': 'transcription job',
                'Value': job_name
            },
        ]
    )
    return {
        'statusCode': 200,
        'body': 'Transcription job started'
    }

<persona> 
1. You are an expert subtitle editor, skilled in subtitle proofreading. 
2. The input <srt> is an array containing SRT fragments, which may include spelling errors and are divided into overly short sections.
3. Your task is to identify the sentence most likely to be the beginning of a statement in each group.
4. The output must only be an array matching the specified <output format> and should not include any other content.
</persona> 

<srt> 
{{srt}} 
</srt> 

<task> 
1. The <srt> input is an array, with each item being a group of SRT fragments. 
   - Each group has the structure: 
     {
       "srt_num_1": {"Srt": "Value", "Timestamp": "Value" },
       "srt_num_2": {"Srt": "Value", "Timestamp": "Value" },
       ...
     }
2. The SRT fragments within each group represent chopped segments from a continuous context.
3. For each group, select the fragment most likely to be the beginning of a sentence among all fragments.
4. Output the selected fragments in an array format, ensuring only the selected items appear in the output.
5. Confirm the output strictly adheres to the array format.

</task> 

<output format> 
[
  {
    "srt_num_A": { "Srt": "Value", "Timestamp": "Value" },
    "srt_num_B": { "Srt": "Value", "Timestamp": "Value" }
  }
]
</output format> 

<example> 
If the input is: 
[
  { 
    "67": { "Srt": "That's a little bit faster than G BT 40, and way faster than cloud", "Timestamp": "00:03:15,020 --> 00:03:20,199" },
    "68": { "Srt": "opus. All right.", "Timestamp": "00:03:20,490 --> 00:03:21,429" },
    "69": { "Srt": "So now it's time for the showdown side by side tests between cloud 3.5 Sonnet", "Timestamp": "00:03:21,440 --> 00:03:26,699" },
    "70": { "Srt": "and GP T 40,", "Timestamp": "00:03:26,880 --> 00:03:28,270" },
    "71": { "Srt": "I'll present each model with the same prompt", "Timestamp": "00:03:28,460 --> 00:03:31,100" },
    "72": { "Srt": "and evaluate their responses for each test.", "Timestamp": "00:03:31,440 --> 00:03:34,210" },
    "73": { "Srt": "I'll choose a winner based on my somewhat subjective criteria", "Timestamp": "00:03:34,220 --> 00:03:38,500" }
  }
]

Thinking steps:
1. There is only one group in the input array.
2. Within this group, the fragment "That's a little bit faster than G BT 40, and way faster than cloud" is most likely the beginning of a sentence.
3. The SRT number of this fragment is 67.

The output should be:
[
  { 
    "67": { "Srt": "That's a little bit faster than G BT 40, and way faster than cloud", "Timestamp": "00:03:15,020 --> 00:03:20,199" }
  }
]
</example>

<persona> 
You are an expert subtitle editor, highly skilled at proofreading and correcting subtitles. Using your knowledge and context, ensure the final output is an array, containing only the corrected subtitles and no additional content.
</persona> 

<srt> 
{srt} 
</srt> 

<srt_instructions> 
1. The <srt> provided is a JSON format subtitle file with subtitle numbers, timestamps, and content. 
2. Due to limitations in subtitle generation tools, there are many word errors in the subtitle content.
3. The time segments for splitting subtitles are too short, often dividing a single sentence into multiple fragments.
4. Example of the <srt> format:
   {
     "129": { 
       "Timestamp": "00:06:22,850 --> 00:06:24,720", 
       "Srt": "So let's see if these two models can do it." 
     },
     "130": { 
       "Timestamp": "00:06:24,929 --> 00:06:25,640", 
       "Srt": "Um" 
     },
     "131": { 
       "Timestamp": "00:06:25,799 --> 00:06:27,790", 
       "Srt": "I was thinking the prompt of" 
     },
     "132": { 
       "Timestamp": "00:06:28,640 --> 00:06:31,500", 
       "Srt": "create a dialogue between a dragon and a knight." 
     }
   }
</srt_instructions> 

<task> 
Your tasks:
1. Merge and re-segment subtitles so that each segment contains **one complete sentence**. Merge segments with adjacent subtitles if they form a single sentence.
2. Correct any errors within these sentences using your knowledge of the context.
3. Translate each sentence into Chinese.
4. If segments are merged, combine the corresponding timings and renumber the segments sequentially, starting from the first subtitle number in <srt>.
5. Ensure the output is an array and contains only the formatted subtitles. Do not include any additional content.
</task> 

<example> 
Example output format:
[
  {
    "xxxxx": { 
      "Timestamp": "00:06:25,799 --> 00:06:31,500", 
      "Srt": "I was thinking about creating a dialogue between a dragon and a knight.", 
      "TranslatedSrt": "我在考虑创作一个龙与骑士之间的对话。" 
    }
  }
]
</example>

# Author:Li Kaiyuan
# Version: 0.8
from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth
import os
import boto3 
import json
import logging
from botocore.exceptions import ClientError

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

srt_index_name = os.environ["QUERY_INDEX_NAME"]
host = os.environ["OPENSEARCH_HOST"]
region = os.environ["AWS_REGION"]

def generate_text_embeddings(event):
    try:
        content = json.loads(event["Body"]["content"][0]["text"])
        query = content["key_words"]
        
        body = json.dumps({
            "texts": [query],
            "input_type": "search_query",  
            "embedding_types": ["float"]
        })
        
        bedrock = boto3.client(service_name='bedrock-runtime')
        response = bedrock.invoke_model(
            body=body,
            modelId="cohere.embed-english-v3",
            accept='*/*',
            contentType='application/json'
        )
        
        response_body = json.loads(response.get('body').read())
        query_embedding = response_body["embeddings"]["float"][0]

        return { 
            "question": query,
            "query_embedding": query_embedding
        }
        
    except ClientError as e:
        logger.error(f"Error invoking the model: {e}")
        return {"error": str(e)}


def query_opensearch(query_embedding, host=host, index_name=srt_index_name):
    credentials = boto3.Session().get_credentials()
    auth = AWSV4SignerAuth(credentials, region, "aoss")
    
    # Initialize OpenSearch client
    search = OpenSearch(
        hosts=[{'host': host, 'port': 443}],
        http_auth=auth,
        use_ssl=True,
        verify_certs=True,
        connection_class=RequestsHttpConnection
    )

    # Main query to retrieve relevant docs
    main_query = {
        "_source": ["Number", "Srt", "Timestamp", "TranslatedSrt", "VideoName"],
        "size": 10,
        "query": {
            "knn": {
                "embeddings": {
                    "vector": query_embedding,
                    "k": 100
                }
            }
        }
    }

    response = search.search(body=main_query, index=index_name)
    retrieves = response["hits"]["hits"]

    related_srts = []

    for retrieve in retrieves:
        retrieve_srt = retrieve["_source"]
        Number = retrieve_srt["Number"]
        VideoName = retrieve_srt["VideoName"]

        # Adjust "Numbers" to a range of IDs
        Numbers = [Number + i for i in range(6)]
        
        # Query to retrieve related SRTs
        related_query = {
            "_source": ["Srt"],
            "query": {
                "bool": {
                    "filter": [
                        {"terms": {"Number": Numbers}},
                        {"term": {"VideoName": VideoName}}
                    ]
                }
            }
        }

        related_response = search.search(body=related_query, index=index_name)
        doc = [srt["_source"].get("Srt", "No Srt found") for srt in related_response.get("hits", {}).get("hits", [])]

        retrieve_srt["doc"] = doc  # Attach related SRTs
        related_srts.append(retrieve_srt)

    return related_srts


def lambda_handler(event, context):
    query_info = generate_text_embeddings(event)
    related_srts = query_opensearch(query_info["query_embedding"])
    
    rerank_info_str = {
        "rerank_info": json.dumps({
            "Question": query_info["question"],
            "Related_srts": related_srts
        })
    }

    return rerank_info_str

# Author:Li Kaiyuan
# Version: 0.14

import streamlit as st
import time
import json
import boto3
import logging
import requests
from typing import List, Dict, Optional
import base64

# 配置类
class Config:
    LOGGING_LEVEL = logging.INFO
    PAGE_TITLE = "VideoSearch Engine"
    PAGE_ICON = ":robot:"
    LOGO_PATH = "images/logo.png"
    LAMBDA_FUNCTION_NAME = "apiGatewayEntry"  # 替换为你的 Lambda 函数名称

# UI组件类
class UIComponents:
    @staticmethod
    def get_user_input() -> str:
        user_input = st.text_input("您的问题:", key="input")
        if user_input and len(user_input.strip()) < 3:
            st.warning("请输入至少3个字符")
            return ""
        return user_input.strip()
        
    @staticmethod
    def setup_page():
        st.set_page_config(
            page_title=Config.PAGE_TITLE,
            page_icon=Config.PAGE_ICON,
            layout="wide"
        )

    @staticmethod
    def show_header():
        col_logo, col_title = st.columns([1, 7])
        with col_logo:
            st.image(Config.LOGO_PATH, width=80)
        with col_title:
            st.title(Config.PAGE_TITLE)

# 添加缓存装饰器
@st.cache_data(ttl=3600)
def cached_lambda_invoke(user_input: str):
    return APIHandler.invoke_lambda(user_input)

# 问卷显示类
class QuestionnaireDisplay:
    @staticmethod
    def reset_quiz(video_id: str, questions: List[Dict]):
        if st.button(f"重置测验", key=f"reset_{video_id}"):
            st.session_state[f'score_{video_id}'] = 0
            st.session_state[f'answered_{video_id}'] = [False] * len(questions)
            st.experimental_rerun()
            
    @staticmethod
    def display(questions: List[Dict], video_id: str):
        if not questions:
            st.write("问卷内容无法加载")
            return
        
        QuestionnaireDisplay.reset_quiz(video_id, questions)

        # 初始化会话状态 - 使用video_id作为key的一部分
        score_key = f'score_{video_id}'
        answered_key = f'answered_{video_id}'
        
        if score_key not in st.session_state:
            st.session_state[score_key] = 0
        if answered_key not in st.session_state or len(st.session_state[answered_key]) != len(questions):
            st.session_state[answered_key] = [False] * len(questions)

        for i, question in enumerate(questions):
            QuestionnaireDisplay._display_question(i, question, video_id)

        QuestionnaireDisplay._show_score(len(questions), video_id)

    @staticmethod
    def _display_question(index: int, question: Dict, video_id: str):
        st.markdown(f"##### Question {index + 1}: {question['question']}")
        for option in question['options']:
            # 使用video_id创建唯一的key
            key = f"q{index}_o{option['index']}_{video_id}"
            checked = st.checkbox(
                option['content'], 
                key=key, 
                disabled=st.session_state[f'answered_{video_id}'][index]
            )

            if checked and not st.session_state[f'answered_{video_id}'][index]:
                QuestionnaireDisplay._handle_answer(index, option, question, video_id)
        st.write("---")

    @staticmethod
    def _handle_answer(index: int, option: Dict, question: Dict, video_id: str):
        st.session_state[f'answered_{video_id}'][index] = True
        if option['index'] == question['answer']:
            st.success("Correct!")
            st.session_state[f'score_{video_id}'] += 1
        else:
            correct_answer = question['options'][question['answer']]['content']
            st.error(f"Wrong. Correct answer: {correct_answer}")

    @staticmethod
    def _show_score(total_questions: int, video_id: str):
        answered_questions = sum(st.session_state[f'answered_{video_id}'])
        score = st.session_state[f'score_{video_id}']
        if answered_questions == total_questions:
            st.success(f"Quiz completed! Your score: {score}/{total_questions}")
        else:
            st.info(f"Questions answered: {answered_questions}/{total_questions}")
            st.info(f"Current score: {score}/{total_questions}")

# 视频显示类
class VideoDisplay:
    @staticmethod
    def display(item: Dict):
        """Display video and questionnaire with VTT subtitles in memory to avoid CORS issues."""
        try:
            video_id = f"{item['title']}_{item['start_time']}".replace(':', '_')
            
            # Display video title
            if item.get('title'):
                st.markdown(f"### Video Name: {item['title']}")

            # Download VTT content and encode it as base64
            vtt_content = VideoDisplay._download_vtt_content(item.get('vtt_url'))

            if vtt_content:
                vtt_data_uri = f"data:text/vtt;base64,{vtt_content}"
                
                # Calculate start time in seconds
                start_time = VideoDisplay._timestamp_to_seconds(item.get('start_time', '00:00:00'))

                # Construct HTML with embedded subtitles
                video_html = f"""
                    <video id="myVideo_{video_id}" width="800" height="450" controls>
                        <source src="{item['video_url']}" type="video/mp4">
                        <track src="{vtt_data_uri}" kind="subtitles" srclang="en" label="English" default>
                        Your browser does not support the video tag.
                    </video>
                    <script>
                        var video = document.getElementById("myVideo_{video_id}");
                        video.addEventListener('loadedmetadata', function() {{
                            video.currentTime = {start_time};
                        }}, false);
                    </script>
                """
                st.components.v1.html(video_html, height=500)
            else:
                st.error("Failed to load subtitles due to CORS or other issues.")

            if item.get('description'):
                st.markdown("#### Related SRT")
                st.caption(item['description'])
            
            if item.get('relevance'):
                st.caption("Relevance: " + str(item['relevance']))

                
            if item.get('summary'):
                st.markdown("#### Summary")
                st.caption(item['summary'])
                
                
            st.write("---")
            with st.expander("查看相关的问卷"):
                st.markdown("### 问卷内容")
                QuestionnaireDisplay.display(item['checklist'], video_id)

        except Exception as e:
            st.error(f"Error displaying content: {e}")

    @staticmethod
    def _download_vtt_content(vtt_url: str) -> Optional[str]:
        """Download VTT content and return it as a base64-encoded string."""
        try:
            response = requests.get(vtt_url)
            response.raise_for_status()
            vtt_content = base64.b64encode(response.content).decode('utf-8')
            return vtt_content
        except requests.RequestException as e:
            st.error(f"Error downloading VTT content: {e}")
            return None

    @staticmethod
    def _timestamp_to_seconds(timestamp: str) -> int:
        """将时间戳转换为秒数"""
        try:
            hours, minutes, seconds = map(int, timestamp.split(':'))
            return hours * 3600 + minutes * 60 + seconds
        except:
            return 0
                
                

# API处理类
class APIHandler:
    @staticmethod
    def invoke_lambda(user_input: str) -> Optional[Dict]:
        session = boto3.Session(profile_name="skydev")
        lambda_client = session.client(
            service_name = "lambda",
            region_name = "us-west-2")
        payload = {"input": user_input}  # 传递用户输入作为 payload

        try:
            # 调用 Lambda 函数
            response = lambda_client.invoke(
                FunctionName=Config.LAMBDA_FUNCTION_NAME,
                InvocationType="RequestResponse",  # 同步调用
                Payload=json.dumps(payload)
            )
            
            # 读取并解析 Lambda 返回的内容
            response_payload = response['Payload'].read()
            result = json.loads(response_payload)

            # 检查响应状态码和内容
            if "statusCode" in result and result["statusCode"] == 200:
                body_content = result["body"]
                if isinstance(body_content, str):  # 处理嵌套的 JSON
                    body_content = json.loads(body_content)
                return body_content  # 返回解析后的 JSON 数据
            else:
                st.error("Lambda 返回了错误信息")
                return None
        except Exception as e:
            st.error(f"调用 Lambda 时出现异常: {str(e)}")
            return None

# 主应用逻辑
def main():
    # 设置日志
    logging.basicConfig(level=Config.LOGGING_LEVEL)
    logger = logging.getLogger(__name__)
    
    try:
        UIComponents.setup_page()
        UIComponents.show_header()
        
        user_input = UIComponents.get_user_input()

        if user_input:
            # 开始计时
            start_time = time.time()
            with st.spinner('正在搜索相关内容...'):
                response = cached_lambda_invoke(user_input)
            
            # 结束计时
            end_time = time.time()
            elapsed_time = end_time - start_time
            
            # 显示响应时间
            st.write(f"搜索耗时: {elapsed_time:.2f} 秒")
                
            if response and 'res' in response:
                if len(response['res']) == 0:
                    st.info("没有找到相关内容")
                else:
                    for item in response['res']:
                        VideoDisplay.display(item)
                        st.write("---")
            else:
                st.error("搜索出现错误，请稍后重试")
                
    except Exception as e:
        logger.error(f"Application error: {str(e)}")
        st.error("应用程序出现错误，请刷新页面重试")

if __name__ == "__main__":
    main()