lattifai 0.4.5__py3-none-any.whl → 1.0.0__py3-none-any.whl

lattifai/transcription/gemini.py

@@ -0,0 +1,219 @@
+"""Gemini 2.5 Pro transcription module with config-driven architecture."""
+
+import asyncio
+from pathlib import Path
+from typing import Optional, Union
+
+from google import genai
+from google.genai.types import GenerateContentConfig, Part, ThinkingConfig
+
+from lattifai.audio2 import AudioData
+from lattifai.config import TranscriptionConfig
+from lattifai.transcription.base import BaseTranscriber
+from lattifai.transcription.prompts import get_prompt_loader
+
+
+class GeminiTranscriber(BaseTranscriber):
+    """
+    Gemini 2.5/3 Pro audio transcription with config-driven architecture.
+
+    Uses TranscriptionConfig for all behavioral settings.
+    """
+
+    # Transcriber metadata
+    file_suffix = ".md"
+
+    # The specific Gem URL
+    GEM_URL = "https://gemini.google.com/gem/1870ly7xvW2hU_umtv-LedGsjywT0sQiN"
+
+    def __init__(
+        self,
+        transcription_config: Optional[TranscriptionConfig] = None,
+    ):
+        """
+        Initialize Gemini transcriber.
+
+        Args:
+            transcription_config: Transcription configuration. If None, uses default.
+        """
+        super().__init__(config=transcription_config)
+
+        self._client: Optional[genai.Client] = None
+        self._generation_config: Optional[GenerateContentConfig] = None
+        self._system_prompt: Optional[str] = None
+
+        # Warn if API key not available
+        if not self.config.gemini_api_key:
+            self.logger.warning(
+                "⚠️ Gemini API key not provided. API key will be required when calling transcription methods."
+            )
+
+    @property
+    def name(self) -> str:
+        """Human-readable name of the transcriber."""
+        return f"{self.config.model_name}"
+
+    async def transcribe_url(self, url: str, language: Optional[str] = None) -> str:
+        """
+        Transcribe audio from URL using Gemini 2.5 Pro.
+
+        Args:
+            url: URL to transcribe (e.g., YouTube)
+            language: Optional language code for transcription (overrides config)
+
+        Returns:
+            Transcribed text
+
+        Raises:
+            ValueError: If API key not provided
+            RuntimeError: If transcription fails
+        """
+        if self.config.verbose:
+            self.logger.info(f"🎤 Starting Gemini transcription for: {url}")
+
+        try:
+            contents = Part.from_uri(file_uri=url, mime_type="video/*")
+            return await self._run_generation(contents, source=url)
+
+        except ImportError:
+            raise RuntimeError("Google GenAI SDK not installed. Please install with: pip install google-genai")
+        except Exception as e:
+            self.logger.error(f"Gemini transcription failed: {str(e)}")
+            raise RuntimeError(f"Gemini transcription failed: {str(e)}")
+
+    async def transcribe_file(self, media_file: Union[str, Path, AudioData], language: Optional[str] = None) -> str:
+        """
+        Transcribe audio/video from local file using Gemini 2.5 Pro.
+
+        Args:
+            media_file: Path to local audio/video file
+            language: Optional language code for transcription (overrides config)
+
+        Returns:
+            Transcribed text
+
+        Raises:
+            ValueError: If API key not provided
+            RuntimeError: If transcription fails
+        """
+        media_file = str(media_file)
+
+        if self.config.verbose:
+            self.logger.info(f"🎤 Starting Gemini transcription for file: {media_file}")
+
+        try:
+            client = self._get_client()
+
+            # Upload audio file
+            if self.config.verbose:
+                self.logger.info("📤 Uploading audio file to Gemini...")
+            media_file = client.files.upload(path=media_file)
+
+            contents = Part.from_uri(file_uri=media_file.uri, mime_type=media_file.mime_type)
+            return await self._run_generation(contents, source=media_file, client=client)
+
+        except ImportError:
+            raise RuntimeError("Google GenAI SDK not installed. Please install with: pip install google-genai")
+        except Exception as e:
+            self.logger.error(f"Gemini transcription failed: {str(e)}")
+            raise RuntimeError(f"Gemini transcription failed: {str(e)}")
+
+    def _get_transcription_prompt(self) -> str:
+        """Get (and cache) transcription system prompt from prompts module."""
+        if self._system_prompt is not None:
+            return self._system_prompt
+
+        # Load prompt from prompts/gemini/transcription_gem.txt
+        prompt_loader = get_prompt_loader()
+        base_prompt = prompt_loader.get_gemini_transcription_prompt()
+
+        # Add language-specific instruction if configured
+        if self.config.language:
+            base_prompt += f"\n\n* Use {self.config.language} language for transcription."
+
+        self._system_prompt = base_prompt
+        return self._system_prompt
+
+    def get_gem_info(self) -> dict:
+        """Get information about the Gem being used."""
+        return {
+            "gem_name": "Media Transcription Gem",
+            "gem_url": self.GEM_URL,
+            "model": self.config.model_name,
+            "description": "Specialized Gem for media content transcription",
+        }
+
+    def _build_result(self, transcript: str, output_file: Path) -> dict:
+        """Augment the base result with Gemini-specific metadata."""
+        base_result = super()._build_result(transcript, output_file)
+        base_result.update({"model": self.config.model_name, "language": self.config.language})
+        return base_result
+
+    def _get_client(self) -> genai.Client:
+        """Lazily create the Gemini client when first needed."""
+        if not self.config.gemini_api_key:
+            raise ValueError("Gemini API key is required for transcription")
+
+        if self._client is None:
+            self._client = genai.Client(api_key=self.config.gemini_api_key)
+        return self._client
+
+    def _get_generation_config(self) -> GenerateContentConfig:
+        """Lazily build the generation config since it rarely changes."""
+        if self._generation_config is None:
+            self._generation_config = GenerateContentConfig(
+                system_instruction=self._get_transcription_prompt(),
+                response_modalities=["TEXT"],
+                thinking_config=ThinkingConfig(
+                    include_thoughts=False,
+                    thinking_budget=-1,
+                    # thinking_level="high",  # "low", "medium"
+                ),
+            )
+        return self._generation_config
+
+    async def _run_generation(
+        self,
+        contents: Part,
+        *,
+        source: str,
+        client: Optional[genai.Client] = None,
+    ) -> str:
+        """
+        Shared helper for sending generation requests and handling the response.
+        """
+        client = client or self._get_client()
+        config = self._get_generation_config()
+
+        if self.config.verbose:
+            self.logger.info(f"🔄 Sending transcription request to {self.config.model_name} ({source})...")
+
+        response = await asyncio.get_event_loop().run_in_executor(
+            None,
+            lambda: client.models.generate_content(
+                model=self.config.model_name,
+                contents=contents,
+                config=config,
+            ),
+        )
+
+        if not response.text:
+            raise RuntimeError("Empty response from Gemini API")
+
+        transcript = response.text.strip()
+
+        if self.config.verbose:
+            self.logger.info(f"✅ Transcription completed ({source}): {len(transcript)} characters")
+
+        return transcript
+
+    def write(
+        self, transcript: str, output_file: Path, encoding: str = "utf-8", cache_audio_events: bool = True
+    ) -> Path:
+        """
+        Persist transcript text to disk and return the file path.
+        """
+        if isinstance(output_file, str):
+            output_file = Path(output_file)
+        output_file.write_text(transcript, encoding=encoding)
+        return output_file

lattifai/transcription/lattifai.py

@@ -0,0 +1,103 @@
+"""Transcription module with config-driven architecture."""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from lattifai.audio2 import AudioData
+from lattifai.caption import Caption
+from lattifai.config import TranscriptionConfig
+from lattifai.transcription.base import BaseTranscriber
+from lattifai.transcription.prompts import get_prompt_loader  # noqa: F401
+
+
+class LattifAITranscriber(BaseTranscriber):
+    """
+    LattifAI local transcription with config-driven architecture.
+
+    Uses TranscriptionConfig for all behavioral settings.
+    Note: This transcriber only supports local file transcription, not URLs.
+    """
+
+    # Transcriber metadata
+    file_suffix = ".ass"
+    supports_url = False
+
+    def __init__(
+        self,
+        transcription_config: TranscriptionConfig,
+    ):
+        """
+        Initialize Gemini transcriber.
+
+        Args:
+            transcription_config: Transcription configuration. If None, uses default.
+        """
+        super().__init__(
+            config=transcription_config,
+        )
+
+        self._system_prompt: Optional[str] = None
+        self._transcriber = None
+
+    @property
+    def name(self) -> str:
+        return f"{self.config.model_name}"
+
+    async def transcribe_url(self, url: str, language: Optional[str] = None) -> str:
+        """
+        URL transcription not supported for LattifAI local models.
+
+        This method exists to satisfy the BaseTranscriber interface but
+        will never be called because supports_url = False and the base
+        class checks this flag before calling this method.
+
+        Args:
+            url: URL to transcribe (not supported)
+            language: Optional language code (not used)
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support URL transcription. "
+            f"Please download the file first and use transcribe_file()."
+        )
+
+    async def transcribe_file(self, media_file: Union[str, Path, AudioData], language: Optional[str] = None) -> Caption:
+        if self._transcriber is None:
+            from lattifai_core.transcription import LattifAITranscriber as CoreLattifAITranscriber
+
+            self._transcriber = CoreLattifAITranscriber.from_pretrained(model_config=self.config)
+
+        transcription, audio_events = self._transcriber.transcribe(media_file, language=language, num_workers=2)
+        caption = Caption.from_transcription_results(
+            transcription=transcription,
+            audio_events=audio_events,
+        )
+
+        return caption
+
+    def write(
+        self, transcript: Caption, output_file: Path, encoding: str = "utf-8", cache_audio_events: bool = True
+    ) -> Path:
+        """
+        Persist transcript text to disk and return the file path.
+        """
+        transcript.write(
+            output_file,
+            include_speaker_in_text=False,
+        )
+        if cache_audio_events and transcript.audio_events:
+            from tgt import write_to_file
+
+            events_file = output_file.with_suffix(".AED")
+            write_to_file(transcript.audio_events, events_file, format="long")
+
+        return output_file
+
+    def _get_transcription_prompt(self) -> str:
+        """Get (and cache) transcription system prompt from prompts module."""
+        if self._system_prompt is not None:
+            return self._system_prompt
+
+        base_prompt = ""  # TODO
+
+        self._system_prompt = base_prompt
+        return self._system_prompt

lattifai/{workflows → transcription}/prompts/__init__.py

@@ -26,16 +26,16 @@ class PromptLoader:
         Raises:
             FileNotFoundError: If prompt file doesn't exist
         """
-        prompt_path = self.prompts_dir / category / f'{name}.txt'
+        prompt_path = self.prompts_dir / category / f"{name}.txt"
 
         if not prompt_path.exists():
-            raise FileNotFoundError(f'Prompt not found: {prompt_path}')
+            raise FileNotFoundError(f"Prompt not found: {prompt_path}")
 
-        return prompt_path.read_text(encoding='utf-8').strip()
+        return prompt_path.read_text(encoding="utf-8").strip()
 
     def get_gemini_transcription_prompt(self) -> str:
         """Get the Gemini transcription Gem prompt"""
-        return self.load_prompt('gemini', 'transcription_gem')
+        return self.load_prompt("gemini", "transcription_gem")
 
 
 # Global instance

lattifai/types.py

@@ -0,0 +1,30 @@
+"""Common type definitions for LattifAI."""
+
+from pathlib import Path
+from typing import List, TypeAlias, Union
+
+from lhotse.utils import Pathlike
+
+from .caption import Supervision
+
+# Path-like types
+PathLike: TypeAlias = Pathlike  # Re-export for convenience (str | Path)
+
+# Caption types
+SupervisionList: TypeAlias = List[Supervision]
+"""List of caption segments with timing and text information."""
+
+# Media format types
+MediaFormat: TypeAlias = str
+"""Media format string (e.g., 'mp3', 'wav', 'mp4')."""
+
+# URL types
+URL: TypeAlias = str
+"""String representing a URL."""
+
+__all__ = [
+    "PathLike",
+    "SupervisionList",
+    "MediaFormat",
+    "URL",
+]

lattifai/utils.py

@@ -6,14 +6,12 @@ from pathlib import Path
 from typing import Any, Optional, Type
 
 from lattifai.errors import ModelLoadError
-from lattifai.tokenizer import LatticeTokenizer
-from lattifai.workers import Lattice1AlphaWorker
 
 
 def _get_cache_marker_path(cache_dir: Path) -> Path:
     """Get the path for the cache marker file with current date."""
-    today = datetime.now().strftime('%Y%m%d')
-    return cache_dir / f'.done{today}'
+    today = datetime.now().strftime("%Y%m%d")
+    return cache_dir / f".done{today}"
 
 
 def _is_cache_valid(cache_dir: Path) -> bool:
@@ -22,7 +20,7 @@ def _is_cache_valid(cache_dir: Path) -> bool:
         return False
 
     # Find any .done* marker files
-    marker_files = list(cache_dir.glob('.done*'))
+    marker_files = list(cache_dir.glob(".done*"))
     if not marker_files:
         return False
 
@@ -31,8 +29,8 @@ def _is_cache_valid(cache_dir: Path) -> bool:
 
     # Extract date from marker filename (format: .doneYYYYMMDD)
     try:
-        date_str = latest_marker.name.replace('.done', '')
-        marker_date = datetime.strptime(date_str, '%Y%m%d')
+        date_str = latest_marker.name.replace(".done", "")
+        marker_date = datetime.strptime(date_str, "%Y%m%d")
         # Check if marker is older than 1 days
         if datetime.now() - marker_date > timedelta(days=1):
             return False
@@ -45,7 +43,7 @@ def _is_cache_valid(cache_dir: Path) -> bool:
 def _create_cache_marker(cache_dir: Path) -> None:
     """Create a cache marker file with current date and clean old markers."""
     # Remove old marker files
-    for old_marker in cache_dir.glob('.done*'):
+    for old_marker in cache_dir.glob(".done*"):
         old_marker.unlink(missing_ok=True)
 
     # Create new marker file
@@ -55,8 +53,8 @@ def _create_cache_marker(cache_dir: Path) -> None:
 
 def _resolve_model_path(model_name_or_path: str) -> str:
     """Resolve model path, downloading from Hugging Face when necessary."""
-    if Path(model_name_or_path).exists():
-        return model_name_or_path
+    if Path(model_name_or_path).expanduser().exists():
+        return str(Path(model_name_or_path).expanduser())
 
     from huggingface_hub import snapshot_download
     from huggingface_hub.constants import HF_HUB_CACHE
@@ -68,7 +66,7 @@ def _resolve_model_path(model_name_or_path: str) -> str:
     # Check if we have a valid cached version
     if _is_cache_valid(cache_dir):
         # Return the snapshot path (latest version)
-        snapshots_dir = cache_dir / 'snapshots'
+        snapshots_dir = cache_dir / "snapshots"
         if snapshots_dir.exists():
             snapshot_dirs = [d for d in snapshots_dir.iterdir() if d.is_dir()]
             if snapshot_dirs:
@@ -77,13 +75,13 @@ def _resolve_model_path(model_name_or_path: str) -> str:
                 return str(latest_snapshot)
 
     try:
-        downloaded_path = snapshot_download(repo_id=model_name_or_path, repo_type='model')
+        downloaded_path = snapshot_download(repo_id=model_name_or_path, repo_type="model")
         _create_cache_marker(cache_dir)
         return downloaded_path
     except LocalEntryNotFoundError:
         try:
-            os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
-            downloaded_path = snapshot_download(repo_id=model_name_or_path, repo_type='model')
+            os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+            downloaded_path = snapshot_download(repo_id=model_name_or_path, repo_type="model")
             _create_cache_marker(cache_dir)
             return downloaded_path
         except Exception as e:  # pragma: no cover - bubble up for caller context
@@ -94,40 +92,14 @@ def _resolve_model_path(model_name_or_path: str) -> str:
 
 def _select_device(device: Optional[str]) -> str:
     """Select best available torch device when not explicitly provided."""
-    if device:
+    if device and device != "auto":
         return device
 
     import torch
 
-    detected = 'cpu'
+    detected = "cpu"
     if torch.backends.mps.is_available():
-        detected = 'mps'
+        detected = "mps"
     elif torch.cuda.is_available():
-        detected = 'cuda'
+        detected = "cuda"
     return detected
-
-
-def _load_tokenizer(
-    client_wrapper: Any,
-    model_path: str,
-    device: str,
-    *,
-    tokenizer_cls: Type[LatticeTokenizer] = LatticeTokenizer,
-) -> LatticeTokenizer:
-    """Instantiate tokenizer with consistent error handling."""
-    try:
-        return tokenizer_cls.from_pretrained(
-            client_wrapper=client_wrapper,
-            model_path=model_path,
-            device=device,
-        )
-    except Exception as e:
-        raise ModelLoadError(f'tokenizer from {model_path}', original_error=e)
-
-
-def _load_worker(model_path: str, device: str) -> Lattice1AlphaWorker:
-    """Instantiate lattice worker with consistent error handling."""
-    try:
-        return Lattice1AlphaWorker(model_path, device=device, num_threads=8)
-    except Exception as e:
-        raise ModelLoadError(f'worker from {model_path}', original_error=e)

lattifai/workflow/__init__.py

@@ -0,0 +1,22 @@
+"""
+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 .base import WorkflowAgent, WorkflowResult, WorkflowStep
+from .file_manager import TRANSCRIBE_CHOICE, FileExistenceManager
+from .youtube import YouTubeDownloader
+
+__all__ = [
+    "WorkflowAgent",
+    "WorkflowStep",
+    "WorkflowResult",
+    "FileExistenceManager",
+    "YouTubeDownloader",
+    "TRANSCRIBE_CHOICE",
+]

lattifai/workflow/agents.py

@@ -0,0 +1,6 @@
+"""
+Caption Agents
+
+"""
+
+__all__ = []

lattifai/{workflows → workflow}/base.py

@@ -7,20 +7,20 @@ import logging
 import time
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional
 
 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}')
+    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'
+            "%(asctime)s - %(name)+17s.py:%(lineno)-4d - %(levelname)-8s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
         )
         handler.setFormatter(formatter)
         logger.addHandler(handler)
@@ -30,17 +30,17 @@ def setup_workflow_logger(name: str) -> logging.Logger:
     return logger
 
 
-logger = setup_workflow_logger('base')
+logger = setup_workflow_logger("base")
 
 
 class WorkflowStatus(Enum):
     """Workflow execution status"""
 
-    PENDING = 'pending'
-    RUNNING = 'running'
-    COMPLETED = 'completed'
-    FAILED = 'failed'
-    RETRYING = 'retrying'
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    RETRYING = "retrying"
 
 
 @dataclass
@@ -84,7 +84,7 @@ class WorkflowAgent(abc.ABC):
         self.name = name
         self.max_retries = max_retries
         self.steps: List[WorkflowStep] = []
-        self.logger = setup_workflow_logger('agent')
+        self.logger = setup_workflow_logger("agent")
 
     @abc.abstractmethod
     def define_steps(self) -> List[WorkflowStep]:
@@ -111,11 +111,11 @@ class WorkflowAgent(abc.ABC):
         context = kwargs.copy()
         step_results = []
 
-        self.logger.info(colorful.bold_white_on_green(f'🚀 Starting workflow: {self.name}'))
+        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}'
+                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()
@@ -123,17 +123,17 @@ class WorkflowAgent(abc.ABC):
                 step_duration = time.time() - step_start
 
                 step_results.append(
-                    {'step_name': step.name, 'status': 'completed', 'duration': step_duration, 'result': step_result}
+                    {"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_{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')
+                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')
+            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
@@ -145,9 +145,9 @@ class WorkflowAgent(abc.ABC):
             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}')
+                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)}')
+                self.logger.error(f"❌ Workflow failed after {execution_time:.2f}s: {str(e)}")
 
             return WorkflowResult(
                 status=WorkflowStatus.FAILED,
@@ -164,7 +164,7 @@ class WorkflowAgent(abc.ABC):
         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})')
+                    self.logger.info(f"🔄 Retrying step {step.name} (attempt {attempt + 1}/{step.max_retries + 1})")
 
                 result = await self.execute_step(step, context)
                 return result
@@ -176,14 +176,14 @@ class WorkflowAgent(abc.ABC):
                 # 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]
+                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...')
+                    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}'
+                        f"❌ Step {step.name} failed after {step.max_retries + 1} attempts: {error_summary}"
                     )
                     raise e