TwelveLabs와 Perplexity Sonar로 구축하는 비디오 딥 리서치(Deep Research)

class TwelveLabsService:
    def __init__(self, api_key=None):
        # API 키가 주어지지 않은 경우 환경 변수 조회를 시도합니다.
        if api_key is None:
            api_key = os.environ.get('TWELVELABS_API_KEY', '')
        self.api_key = api_key # 인스턴스에 API 키 저장
        self.client = TwelveLabs(api_key=api_key) # TwelveLabs 클라이언트 초기화

def get_indexes(self):
        try:
            print("인덱스를 가져오는 중...")
            # API 키 사용 가능 여부 확인
            if not self.api_key:
                print("사용 가능한 API 키가 없습니다.")
                return []
            
            # TwelveLabs 클라이언트를 사용해 인덱스 목록 획득
            indexes = self.client.indexes.list()
            
            result = []
            for index in indexes:
                # 인덱스 상세 구조 구성
                result.append({
                    "id": index.id,
                    "name": index.index_name
                })
                print(f"ID: {index.id}")
                print(f"  Name: {index.index_name}")
            
            return result
        except Exception as e:
            print(f"인덱스를 가져오는 중 오류 발생: {e}")
            return []

def get_videos(self, index_id, page=1):
        try:
            # API 키 사용 가능 여부 확인
            if not self.api_key:
                print("사용 가능한 API 키가 없습니다.")
                return []
            
            # TwelveLabs 클라이언트를 사용해 비디오 조회
            videos_response = self.client.indexes.videos.list(index_id=index_id, page=page)
            
            result = []
            for video in videos_response.items:
                system_metadata = video.system_metadata
                hls_data = video.hls
                thumbnail_urls = hls_data.get('thumbnail_urls', []) if hls_data else []
                thumbnail_url = thumbnail_urls[0] if thumbnail_urls else None
                video_url = hls_data.get('video_url') if hls_data else None
                # 비디오 상세 인스턴스 사전 매핑
                result.append({
                    "id": video.id,
                    "name": system_metadata.filename if system_metadata and system_metadata.filename else f'Video {video.id}',
                    "duration": system_metadata.duration if system_metadata else 0,
                    "thumbnail_url": thumbnail_url,
                    "video_url": video_url,
                    "width": system_metadata.width if system_metadata else 0,
                    "height": system_metadata.height if system_metadata else 0,
                    "fps": system_metadata.fps if system_metadata else 0,
                    "size": system_metadata.size if system_metadata else 0
                })
            
            return result
        except Exception as e:
            print(f"인덱스 {index_id}의 비디오를 가져오는 중 오류 발생: {e}")
            return []

def analyze_video(self, video_id, prompt):
        try:
            # TwelveLabs 클라이언트를 통한 비디오 정보 생성
            analysis_response = self.client.analyze(
                video_id=video_id,
                prompt=prompt
            )
            return analysis_response.data
        except Exception as e:
            print(f"비디오 {video_id} 분석 중 오류 발생: {e}")
            raise e

def deep_research(self, query, timeout=180):
        try:
            # API 키 사용 가능 여부 확인
            if not self.api_key:
                raise ValueError("API 키가 필요합니다.")
            
            payload = {
                "model": "sonar",
                "messages": [
                    {"role": "user", "content": query}
                ]
            }
            
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            
            # 타임아웃 설정을 포함해 API 엔드포인트로 POST 요청 전송
            response = requests.post(
                self.base_url, 
                json=payload, 
                headers=headers, 
                timeout=timeout
            )
            
            # 요청 결과 승인 확인
            if response.status_code == 200:
                return response.json()
            else:
                print(f"오류: {response.status_code} - {response.text}")
                return {"error": f"API 요청 실패 (상태 코드: {response.status_code})"}
                
        except requests.exceptions.Timeout:
            print("요청 시간 초과")
            return {"error": "요청 시간 초과 - Sonar 리서치가 너무 오래 걸립니다."}
        except requests.exceptions.RequestException as e:
            print(f"요청 중 오류 발생: {e}")
            return {"error": f"네트워크 오류: {str(e)}"}
        except Exception as e:
            print(f"deep_research 실행 중 오류 발생: {e}")
            return {"error": str(e)}

def generate_workflow(twelvelabs_api_key, index_id, video_id, analysis_prompt, research_query, research_prompt_template):
    # 입력 값 유효성 체크
    if not twelvelabs_api_key:
        yield json.dumps({'type': 'error', 'message': 'TwelveLabs API 키가 필요합니다.'}) + '\n'
        return
    if not index_id or not video_id:
        yield json.dumps({'type': 'error', 'message': 'Index ID와 Video ID가 필요합니다.'}) + '\n'
        return
    if not research_query:
        yield json.dumps({'type': 'error', 'message': '리서치 질의가 필요합니다.'}) + '\n'
        return

    try:
        twelvelabs_service = TwelveLabsService(api_key=twelvelabs_api_key)
        
        # 1단계: 비디오 상세 정보 가져오기
        yield safe_json_dumps({
            'type': 'progress',
            'step': 'video_details',
            'message': '비디오 상세 정보를 가져오는 중...',
            'progress': 0
        }) + '\n'

        video_details = twelvelabs_service.get_video_details(index_id, video_id)
        if not video_details:
            yield safe_json_dumps({'type': 'error', 'message': '비디오 상세 정보를 검색할 수 없습니다.'}) + '\n'
            return

        yield safe_json_dumps({
            'type': 'data',
            'step': 'video_details',
            'data': {
                'id': video_details.get('_id', ''),
                'filename': video_details.get('system_metadata', {}).get('filename', ''),
                'duration': video_details.get('system_metadata', {}).get('duration', 0)
            },
            'progress': 33
        }) + '\n'

        # 2단계: 비디오 분석
        yield safe_json_dumps({
            'type': 'progress',
            'step': 'analysis',
            'message': '비디오 콘텐츠 분석 중...',
            'progress': 33
        }) + '\n'

        analysis_result = twelvelabs_service.analyze_video(video_id, analysis_prompt)
        yield safe_json_dumps({
            'type': 'data',
            'step': 'analysis',
            'data': analysis_result,
            'progress': 66
        }) + '\n'

        # 3단계: 컨텍스트 기반 리서치
        yield safe_json_dumps({
            'type': 'progress',
            'step': 'research',
            'message': '심층 리서치 진행 중...',
            'progress': 66
        }) + '\n'

        enhanced_query = research_prompt_template.format(
            analysis_result=analysis_result,
            research_query=research_query
        )
        
        sonar_service = SonarService()
        research_result = sonar_service.deep_research(enhanced_query, timeout=180)

        if 'error' in research_result:
            yield safe_json_dumps({'type': 'error', 'message': f'리서치 실패: {research_result["error"]}'}) + '\n'
            return

        # 리서치 결과 콘텐츠 추출
        research_content = ""
        if research_result and research_result.get('choices'):
            research_content = research_result['choices'][0].get('message', {}).get('content', '')
        
        # 결과 데이터가 클 경우 청크 단위로 나누어 전송
        max_chunk_size = 10000
        if len(research_content) > max_chunk_size:
            for i in range(0, len(research_content), max_chunk_size):
                chunk = research_content[i:i + max_chunk_size]
                is_final = (i + max_chunk_size) >= len(research_content)
                
                yield safe_json_dumps({
                    'type': 'research_chunk',
                    'content': chunk,
                    'is_final': is_final,
                    'progress': 80 + (i / len(research_content)) * 20
                }) + '\n'
            
            # 청크 분할 전송 완료 플래그 적용
            yield safe_json_dumps({
                'type': 'complete',
                'data': {
                    'research': {
                        'choices': [{
                            'message': {
                                'content': '[CHUNKED_CONTENT]'
                            }
                        }],
                        'citations': research_result.get('citations', [])[:10],
                        'usage': research_result.get('usage', {})
                    },
                    'sources': research_result.get('search_results', [])[:10]
                },
                'progress': 100
            }) + '\n'
        else:
            # 단일 전송 전개 수렴 완료 처리
            yield safe_json_dumps({
                'type': 'complete',
                'data': {
                    'research': {
                        'choices': [{
                            'message': {
                                'content': research_content
                            }
                        }],
                        'citations': research_result.get('citations', [])[:10],
                        'usage': research_result.get('usage', {})
                    },
                    'sources': research_result.get('search_results', [])[:10]
                },
                'progress': 100
            }) + '\n'

    except Exception as e:
        yield safe_json_dumps({'type': 'error', 'message': str(e)}) + '\n'

def upload_video_file(self, index_id: str, file_path: str, timeout_seconds: int = 900):
        import sys
        try:
            if not self.api_key:
                return {"error": "TwelveLabs API 키가 없습니다."}
            if not index_id:
                return {"error": "index_id가 없습니다."}
            if not os.path.exists(file_path):
                return {"error": f"파일을 찾을 수 없습니다: {file_path}"}
            
            print(f"파일 업로드 시작: {file_path}", file=sys.stderr)

            tasks_url = "https://api.twelvelabs.io/v1.3/tasks"
            headers = {
                "x-api-key": self.api_key
            }

            # 업로드 태스크 생성
            with open(file_path, "rb") as f:
                files = {
                    "video_file": (os.path.basename(file_path), f)
                }
                data = {
                    "index_id": index_id
                }
                resp = requests.post(tasks_url, headers=headers, files=files, data=data)

            if resp.status_code not in (200, 201):
                return {"error": f"업로드 태스크 생성 실패: {resp.status_code} {resp.text}"}

            resp_json = resp.json() if resp.text else {}
            task_id = resp_json.get("id") or resp_json.get("task_id") or resp_json.get("_id")
            if not task_id:
                return {"error": f"반환된 태스크 id가 없습니다: {resp_json}"}

            # 작업이 완결될 때까지 폴링(Polling) 수행
            import time
            start_time = time.time()
            print(f"태스크 {task_id} 완료 여부 확인 중...", file=sys.stderr)
            
            while time.time() - start_time < timeout_seconds:
                r = requests.get(f"{tasks_url}/{task_id}", headers=headers)
                if r.status_code != 200:
                    time.sleep(2)
                    continue
                task = r.json() if r.text else {}
                status = task.get("status")
                print(f"태스크 {task_id} 상태: {status}", file=sys.stderr)
                
                if status in ("ready", "completed"):
                    video_id = task.get("video_id") or (task.get("data") or {}).get("video_id")
                    print(f"인덱싱이 성공적으로 완료되었습니다! Video ID: {video_id}", file=sys.stderr)
                    return {"status": status, "video_id": video_id, "task": task}
                if status in ("failed", "error"):
                    print(f"인덱싱 실패, 상태: {status}", file=sys.stderr)
                    return {"error": f"인덱싱 실패와 관련된 상태값 {status}", "task": task}
                time.sleep(2)

            print(f"{timeout_seconds}초 초과로 업로드 타임아웃 발생", file=sys.stderr)
            return {"error": "업로드 시간 초과"}