lattifai 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

lattifai/bin/agent.py

@@ -210,7 +210,7 @@ async def _run_youtube_workflow(
     click.echo()
 
     # Import workflow components
-    from lattifai import AsyncLattifAI
+    from lattifai.client import AsyncLattifAI
     from lattifai.workflows import YouTubeSubtitleAgent
     from lattifai.workflows.gemini import GeminiTranscriber
     from lattifai.workflows.youtube import YouTubeDownloader
@@ -309,12 +309,13 @@ if __name__ == '__main__':
 
     asyncio.run(
         _run_youtube_workflow(
+            # url='https://www.youtube.com/watch?v=7nv1snJRCEI',
             url='https://www.youtube.com/watch?v=DQacCB9tDaw',
             lattifai_api_key=os.getenv('LATTIFAI_API_KEY'),
             gemini_api_key=os.getenv('GEMINI_API_KEY', ''),
-            device='cpu',
+            device='mps',
             model_name_or_path='Lattifai/Lattice-1-Alpha',
-            media_format='mp4',
+            media_format='mp3',
             output_format='TextGrid',
             output_dir='~/Downloads/lattifai_youtube',
             max_retries=0,

lattifai/io/reader.py

@@ -5,6 +5,7 @@ from typing import List, Literal, Optional, Union
 from lhotse.utils import Pathlike
 
 from .supervision import Supervision
+from .text_parser import parse_speaker_text
 
 SubtitleFormat = Literal['txt', 'srt', 'vtt', 'ass', 'auto']
 
@@ -72,9 +73,11 @@ class SubtitleReader(ABCMeta):
 
         supervisions = []
         for event in subs.events:
+            speaker, text = parse_speaker_text(event.text)
             supervisions.append(
                 Supervision(
-                    text=event.text,
+                    text=text,
+                    speaker=speaker,
                     start=event.start / 1000.0 if event.start is not None else None,
                     duration=(event.end - event.start) / 1000.0 if event.end is not None else None,
                 )

lattifai/io/text_parser.py

@@ -0,0 +1,75 @@
+import logging
+import re
+from typing import Optional, Tuple
+
+# 来自于字幕中常见的说话人标记格式
+SPEAKER_PATTERN = re.compile(r'((?:>>|>>|>|>).*?[:：])\s*(.*)')
+
+# Transcriber Output Example:
+# 26:19.919 --> 26:34.921
+# [SPEAKER_01]: 越来越多的科技巨头入...
+SPEAKER_LATTIFAI = re.compile(r'(^\[SPEAKER_.*?\]:)\s*(.*)')
+
+# NISHTHA BHATIA: Hey, everyone.
+# DIETER: Oh, hey, Nishtha.
+# GEMINI: That might
+SPEAKER_PATTERN2 = re.compile(r'^([A-Z]{1,15}(?:\s+[A-Z]{1,15})?)[:：]\s*(.*)$')
+
+
+def parse_speaker_text(line) -> Tuple[Optional[str], str]:
+    line = line.replace('\\N', ' ')
+
+    if ':' not in line and '：' not in line:
+        return None, line
+
+    # 匹配以 >> 开头的行，并去除开头的名字和冒号
+    match = SPEAKER_PATTERN.match(line)
+    if match:
+        return match.group(1).strip(), match.group(2).strip()
+
+    match = SPEAKER_LATTIFAI.match(line)
+    if match:
+        assert len(match.groups()) == 2, match.groups()
+        if not match.group(1):
+            logging.error(f'ParseSub LINE [{line}]')
+        else:
+            return match.group(1).strip(), match.group(2).strip()
+
+    match = SPEAKER_PATTERN2.match(line)
+    if match:
+        assert len(match.groups()) == 2, match.groups()
+        return match.group(1).strip(), match.group(2).strip()
+
+    return None, line
+
+
+if __name__ == '__main__':
+    pattern = re.compile(r'>>\s*(.*?)\s*[:：]\s*(.*)')
+    pattern = re.compile(r'(>>.*?[:：])\s*(.*)')
+
+    test_strings = [
+        '>>Key: Value',
+        '>>  Key with space : Value with space ',
+        '>>  全角键 ： 全角值',
+        '>>Key：Value xxx. >>Key：Value',
+    ]
+
+    for text in test_strings:
+        match = pattern.match(text)
+        if match:
+            print(f"Input: '{text}'")
+            print(f"  Key:   '{match.group(1)}'")
+            print(f"  Value: '{match.group(2)}'")
+            print('-------------')
+
+    # pattern2
+    test_strings2 = ['NISHTHA BHATIA: Hey, everyone.', 'DIETER: Oh, hey, Nishtha.', 'GEMINI: That might']
+    for text in test_strings2:
+        match = SPEAKER_PATTERN2.match(text)
+        if match:
+            print(f"  Input: '{text}'")
+            print(f"Speaker: '{match.group(1)}'")
+            print(f"Content: '{match.group(2)}'")
+            print('-------------')
+        else:
+            raise ValueError(f"No match for: '{text}'")

lattifai/tokenizer/tokenizer.py

@@ -192,12 +192,15 @@ class LatticeTokenizer:
         for s, supervision in enumerate(supervisions):
             text_len += len(supervision.text)
             if supervision.speaker:
-                speakers.append(supervision.speaker)
                 if sidx < s:
+                    if len(speakers) < len(texts) + 1:
+                        speakers.append(None)
                     text = ' '.join([sup.text for sup in supervisions[sidx:s]])
                     texts.append(text)
                     sidx = s
                     text_len = len(supervision.text)
+                speakers.append(supervision.speaker)
+
             else:
                 if text_len >= 2000 or s == len(supervisions) - 1:
                     if len(speakers) < len(texts) + 1:
@@ -238,6 +241,12 @@ class LatticeTokenizer:
                     processed_sentences.extend(resplit_parts)
             _sentences = processed_sentences
 
+            if not _sentences:
+                if remainder:
+                    _sentences, remainder = [remainder.strip()], ''
+                else:
+                    continue
+
             if any(_sentences[-1].endswith(ep) for ep in END_PUNCTUATION):
                 supervisions.extend(
                     Supervision(text=text, speaker=(_speaker if s == 0 else None)) for s, text in enumerate(_sentences)

lattifai/workflows/__init__.py

@@ -0,0 +1,34 @@
+"""
+LattifAI Agentic Workflows
+
+This module provides agentic workflow capabilities for automated processing
+of multimedia content through intelligent agent-based pipelines.
+"""
+
+# Import transcript processing functionality
+from lattifai.io import (
+    ALL_SUBTITLE_FORMATS,
+    INPUT_SUBTITLE_FORMATS,
+    OUTPUT_SUBTITLE_FORMATS,
+    SUBTITLE_FORMATS,
+    GeminiReader,
+    GeminiWriter,
+)
+
+from .agents import YouTubeSubtitleAgent
+from .base import WorkflowAgent, WorkflowResult, WorkflowStep
+from .file_manager import FileExistenceManager
+
+__all__ = [
+    'WorkflowAgent',
+    'WorkflowStep',
+    'WorkflowResult',
+    'YouTubeSubtitleAgent',
+    'FileExistenceManager',
+    'GeminiReader',
+    'GeminiWriter',
+    'SUBTITLE_FORMATS',
+    'INPUT_SUBTITLE_FORMATS',
+    'OUTPUT_SUBTITLE_FORMATS',
+    'ALL_SUBTITLE_FORMATS',
+]

lattifai/workflows/agents.py

@@ -0,0 +1,10 @@
+"""
+Subtitle Agents
+
+An agentic workflow for processing YouTube(or more) videos through:
+1. URL processing and audio download
+2. Gemini 2.5 Pro transcription
+3. LattifAI alignment
+"""
+
+from .youtube import YouTubeSubtitleAgent

lattifai/workflows/base.py

@@ -0,0 +1,192 @@
+"""
+Base classes for agentic workflows
+"""
+
+import abc
+import logging
+import time
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+import colorful
+
+
+def setup_workflow_logger(name: str) -> logging.Logger:
+    """Setup a logger with consistent formatting for workflow modules"""
+    logger = logging.getLogger(f'workflows.{name}')
+
+    # Only add handler if it doesn't exist
+    if not logger.handlers:
+        handler = logging.StreamHandler()
+        formatter = logging.Formatter(
+            '%(asctime)s - %(name)+17s.py:%(lineno)-4d - %(levelname)-8s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S'
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+        logger.setLevel(logging.INFO)
+        logger.propagate = False
+
+    return logger
+
+
+logger = setup_workflow_logger('base')
+
+
+class WorkflowStatus(Enum):
+    """Workflow execution status"""
+
+    PENDING = 'pending'
+    RUNNING = 'running'
+    COMPLETED = 'completed'
+    FAILED = 'failed'
+    RETRYING = 'retrying'
+
+
+@dataclass
+class WorkflowResult:
+    """Result of a workflow execution"""
+
+    status: WorkflowStatus
+    data: Optional[Any] = None
+    error: Optional[str] = None
+    exception: Optional[Exception] = None  # Store the original exception object
+    execution_time: Optional[float] = None
+    step_results: Optional[List[Dict[str, Any]]] = None
+
+    @property
+    def is_success(self) -> bool:
+        return self.status == WorkflowStatus.COMPLETED
+
+    @property
+    def is_error(self) -> bool:
+        return self.status == WorkflowStatus.FAILED
+
+
+@dataclass
+class WorkflowStep:
+    """Individual step in a workflow"""
+
+    name: str
+    description: str
+    required: bool = True
+    retry_count: int = 0
+    max_retries: int = 1
+
+    def should_retry(self) -> bool:
+        return self.retry_count < self.max_retries
+
+
+class WorkflowAgent(abc.ABC):
+    """Base class for agentic workflows"""
+
+    def __init__(self, name: str, max_retries: int = 0):
+        self.name = name
+        self.max_retries = max_retries
+        self.steps: List[WorkflowStep] = []
+        self.logger = setup_workflow_logger('agent')
+
+    @abc.abstractmethod
+    def define_steps(self) -> List[WorkflowStep]:
+        """Define the workflow steps"""
+        pass
+
+    @abc.abstractmethod
+    async def execute_step(self, step: WorkflowStep, context: Dict[str, Any]) -> Any:
+        """Execute a single workflow step"""
+        pass
+
+    def setup(self):
+        """Setup the workflow"""
+        self.steps = self.define_steps()
+        for step in self.steps:
+            step.max_retries = self.max_retries
+
+    async def execute(self, **kwargs) -> WorkflowResult:
+        """Execute the complete workflow"""
+        if not self.steps:
+            self.setup()
+
+        start_time = time.time()
+        context = kwargs.copy()
+        step_results = []
+
+        self.logger.info(colorful.bold_white_on_green(f'🚀 Starting workflow: {self.name}'))
+
+        try:
+            for i, step in enumerate(self.steps):
+                step_info = f'📋 Step {i + 1}/{len(self.steps)}: {step.name}'
+                self.logger.info(colorful.bold_white_on_green(step_info))
+
+                step_start = time.time()
+                step_result = await self._execute_step_with_retry(step, context)
+                step_duration = time.time() - step_start
+
+                step_results.append(
+                    {'step_name': step.name, 'status': 'completed', 'duration': step_duration, 'result': step_result}
+                )
+
+                # Update context with step result
+                context[f'step_{i}_result'] = step_result
+                context[f'{step.name.lower().replace(" ", "_")}_result'] = step_result
+
+                self.logger.info(f'✅ Step {i + 1} completed in {step_duration:.2f}s')
+
+            execution_time = time.time() - start_time
+            self.logger.info(f'🎉 Workflow completed in {execution_time:.2f}s')
+
+            return WorkflowResult(
+                status=WorkflowStatus.COMPLETED, data=context, execution_time=execution_time, step_results=step_results
+            )
+
+        except Exception as e:
+            execution_time = time.time() - start_time
+            # For LattifAI errors, just log the error code and basic message
+            from lattifai.errors import LattifAIError
+
+            if isinstance(e, LattifAIError):
+                self.logger.error(f'❌ Workflow failed after {execution_time:.2f}s: [{e.error_code}] {e.message}')
+            else:
+                self.logger.error(f'❌ Workflow failed after {execution_time:.2f}s: {str(e)}')
+
+            return WorkflowResult(
+                status=WorkflowStatus.FAILED,
+                error=str(e),
+                exception=e,  # Store the original exception
+                execution_time=execution_time,
+                step_results=step_results,
+            )
+
+    async def _execute_step_with_retry(self, step: WorkflowStep, context: Dict[str, Any]) -> Any:
+        """Execute a step with retry logic"""
+        last_error = None
+
+        for attempt in range(step.max_retries + 1):
+            try:
+                if attempt > 0:
+                    self.logger.info(f'🔄 Retrying step {step.name} (attempt {attempt + 1}/{step.max_retries + 1})')
+
+                result = await self.execute_step(step, context)
+                return result
+
+            except Exception as e:
+                last_error = e
+                step.retry_count += 1
+
+                # For LattifAI errors, show simplified message in logs
+                from lattifai.errors import LattifAIError
+
+                error_summary = f'[{e.error_code}]' if isinstance(e, LattifAIError) else str(e)[:100]
+
+                if step.should_retry():
+                    self.logger.warning(f'⚠️ Step {step.name} failed: {error_summary}. Retrying...')
+                    continue
+                else:
+                    self.logger.error(
+                        f'❌ Step {step.name} failed after {step.max_retries + 1} attempts: {error_summary}'
+                    )
+                    raise e
+
+        # This should never be reached, but just in case
+        if last_error:
+            raise last_error