topos-node 0.1.0__py3-none-any.whl

topos/enrichment/jobs/__init__.py

@@ -0,0 +1,30 @@
+"""Enrichment job registry."""
+
+from .base import BaseEnrichmentJob
+from .canonical.embeddings_job import EmbeddingsJob
+from .canonical.entities_job import EntitiesJob
+from .canonical.emo_27_job import Emo27Job
+from .canonical.sentiment_job import SentimentJob
+from .canonical.topics_job import TopicsJob
+from .raw.attachments_job import AttachmentsJob
+from .raw.language_job import LanguageJob
+from .raw.time_normalization_job import TimeNormalizationJob
+from .raw.tool_calls_job import ToolCallsJob
+
+CANONICAL_JOBS = [EntitiesJob(), TopicsJob(), SentimentJob(), EmbeddingsJob(), Emo27Job()]
+RAW_JOBS = [AttachmentsJob(), ToolCallsJob(), LanguageJob(), TimeNormalizationJob()]
+
+__all__ = [
+    "BaseEnrichmentJob",
+    "EntitiesJob",
+    "TopicsJob",
+    "SentimentJob",
+    "EmbeddingsJob",
+    "Emo27Job",
+    "AttachmentsJob",
+    "ToolCallsJob",
+    "LanguageJob",
+    "TimeNormalizationJob",
+    "CANONICAL_JOBS",
+    "RAW_JOBS",
+]

topos/enrichment/jobs/base.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+from ...utils.base_object import BaseObject
+
+
+@dataclass(frozen=True)
+class EnrichmentResult:
+    result_id: str
+    payload: Dict[str, str]
+
+
+class EnrichmentJob:
+    def run(self, input_ref: str) -> EnrichmentResult:
+        raise NotImplementedError
+
+
+class BaseEnrichmentJob(BaseObject):
+    """Base enrichment job for canonical messages."""
+
+    def __init__(self, *, name: Optional[str] = None) -> None:
+        """Initialize enrichment job with optional name.
+        
+        Args:
+            name: Optional custom name. Defaults to `ClassName#N`
+        """
+        super().__init__(name=name)
+
+    def get_job_name(self) -> str:
+        raise NotImplementedError
+
+    def get_derived_table(self) -> str:
+        raise NotImplementedError
+
+    async def enrich(
+        self, 
+        canonical_messages: List[Dict[str, Any]],
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> List[Dict[str, Any]]:
+        """Enrich canonical messages.
+        
+        Args:
+            canonical_messages: List of canonical message dictionaries
+            progress_callback: Optional callback(current_count, total_count) called during processing
+            
+        Returns:
+            List of enrichment result dictionaries
+        """
+        raise NotImplementedError
+
+    def should_run(self, canonical_messages: List[Dict[str, Any]]) -> bool:
+        return bool(canonical_messages)

topos/enrichment/jobs/canonical/__init__.py

@@ -0,0 +1 @@
+"""Canonical enrichment jobs."""

topos/enrichment/jobs/canonical/embeddings_job.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from ..base import BaseEnrichmentJob
+
+logger = logging.getLogger("topos.enrichment.jobs.embeddings")
+
+
+class EmbeddingsJob(BaseEnrichmentJob):
+    def get_derived_table(self) -> str:
+        return "message_embeddings"
+
+    async def enrich(
+        self, 
+        canonical_messages: List[Dict[str, Any]],
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> List[Dict[str, Any]]:
+        logger.debug("%s: Embeddings enrichment stub: %d messages", self, len(canonical_messages))
+        # Call progress callback to indicate completion (stub jobs complete instantly)
+        if progress_callback:
+            progress_callback(len(canonical_messages), len(canonical_messages))
+        return []
+
+    def get_job_name(self) -> str:
+        return "embeddings"

topos/enrichment/jobs/canonical/emo_27_job.py

@@ -0,0 +1,97 @@
+"""Emotion classification enrichment via the Engine (HF or Ollama adapter)."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from ..base import BaseEnrichmentJob
+from ...progress_bar import ProgressBar
+from ....engine import Engine
+from ....engine.tasks import ModelRequest, ProcessingTask
+
+logger = logging.getLogger("topos.enrichment.jobs.emo_27")
+
+
+class Emo27Job(BaseEnrichmentJob):
+    """Emotion classification enrichment using the Engine (HF or Ollama)."""
+
+    def __init__(self, *, name: Optional[str] = None):
+        super().__init__(name=name)
+        self._engine = Engine()
+
+    def get_derived_table(self) -> str:
+        return "message_emotions"
+
+    def get_job_name(self) -> str:
+        return "emo_27"
+
+    async def enrich(
+        self,
+        canonical_messages: List[Dict[str, Any]],
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> List[Dict[str, Any]]:
+        """Enrich messages with emotion classifications via Engine.run(task)."""
+        logger.debug("[PIPELINE:ENRICHMENT] %s: processing %d messages", self, len(canonical_messages))
+        results = []
+        total_messages = len(canonical_messages)
+
+        with ProgressBar(total=total_messages, desc=str(self)) as pbar:
+            for msg_idx, msg in enumerate(canonical_messages):
+                if msg_idx % 10 == 0:
+                    await asyncio.sleep(0)
+                message_id = msg.get("message_id") or msg.get("id")
+                content = msg.get("content", "")
+                source_id = msg.get("source_id")
+
+                if not message_id or not content:
+                    pbar.update(1)
+                    if progress_callback:
+                        progress_callback(msg_idx + 1, total_messages)
+                    continue
+
+                try:
+                    task = ProcessingTask(
+                        id=f"emo27_{message_id}",
+                        type="enrichment",
+                        subtype="emotion_classification",
+                        source_id=source_id,
+                        record_ids=[message_id],
+                        input={"text": content},
+                        model_request=ModelRequest(provider="huggingface"),
+                    )
+                    result = await asyncio.to_thread(self._engine.run, task)
+                    if result.status != "completed":
+                        logger.warning(
+                            "[PIPELINE:ENRICHMENT] %s: Engine failed for message %s: %s",
+                            self, message_id, result.error or result.status,
+                        )
+                        pbar.update(1)
+                        if progress_callback:
+                            progress_callback(msg_idx + 1, total_messages)
+                        continue
+                    out = result.output
+                    results.append({
+                        "message_id": message_id,
+                        "source_id": source_id,
+                        "emotion_label": out.get("emotion_label"),
+                        "confidence": out.get("confidence"),
+                        "all_emotions": out.get("all_emotions", []),
+                        "model": out.get("model", ""),
+                    })
+                except Exception as e:
+                    logger.error(
+                        "[PIPELINE:ENRICHMENT] %s: Failed to enrich message %s: %s",
+                        self, message_id, e,
+                    )
+                    pbar.update(1)
+                    if progress_callback:
+                        progress_callback(msg_idx + 1, total_messages)
+                    continue
+                pbar.update(1)
+                if progress_callback:
+                    progress_callback(msg_idx + 1, total_messages)
+
+        logger.debug("[PIPELINE:ENRICHMENT] %s: created %d results", self, len(results))
+        return results

topos/enrichment/jobs/canonical/entities_job.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from ..base import BaseEnrichmentJob
+
+logger = logging.getLogger("topos.enrichment.jobs.entities")
+
+
+class EntitiesJob(BaseEnrichmentJob):
+    def get_derived_table(self) -> str:
+        return "message_entities"
+
+    async def enrich(
+        self, 
+        canonical_messages: List[Dict[str, Any]],
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> List[Dict[str, Any]]:
+        logger.debug("%s: Entities enrichment stub: %d messages", self, len(canonical_messages))
+        # Call progress callback to indicate completion (stub jobs complete instantly)
+        if progress_callback:
+            progress_callback(len(canonical_messages), len(canonical_messages))
+        return []
+
+    def get_job_name(self) -> str:
+        return "entities"

topos/enrichment/jobs/canonical/sentiment_job.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from ..base import BaseEnrichmentJob
+
+logger = logging.getLogger("topos.enrichment.jobs.sentiment")
+
+
+class SentimentJob(BaseEnrichmentJob):
+    def get_derived_table(self) -> str:
+        return "message_sentiment"
+
+    async def enrich(
+        self, 
+        canonical_messages: List[Dict[str, Any]],
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> List[Dict[str, Any]]:
+        logger.debug("%s: Sentiment enrichment stub: %d messages", self, len(canonical_messages))
+        # Call progress callback to indicate completion (stub jobs complete instantly)
+        if progress_callback:
+            progress_callback(len(canonical_messages), len(canonical_messages))
+        return []
+
+    def get_job_name(self) -> str:
+        return "sentiment"

topos/enrichment/jobs/canonical/topics_job.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from ..base import BaseEnrichmentJob
+
+logger = logging.getLogger("topos.enrichment.jobs.topics")
+
+
+class TopicsJob(BaseEnrichmentJob):
+    def get_derived_table(self) -> str:
+        return "message_topics"
+
+    async def enrich(
+        self, 
+        canonical_messages: List[Dict[str, Any]],
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> List[Dict[str, Any]]:
+        logger.debug("%s: Topics enrichment stub: %d messages", self, len(canonical_messages))
+        # Call progress callback to indicate completion (stub jobs complete instantly)
+        if progress_callback:
+            progress_callback(len(canonical_messages), len(canonical_messages))
+        return []
+
+    def get_job_name(self) -> str:
+        return "topics"

topos/enrichment/jobs/raw/__init__.py

@@ -0,0 +1 @@
+"""Raw enrichment jobs."""

topos/enrichment/jobs/raw/attachments_job.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+
+class AttachmentsJob:
+    def get_job_name(self) -> str:
+        return "attachments"
+
+    async def run(self, raw_records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        _ = raw_records
+        return []

topos/enrichment/jobs/raw/language_job.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+
+class LanguageJob:
+    def get_job_name(self) -> str:
+        return "language"
+
+    async def run(self, raw_records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        _ = raw_records
+        return []

topos/enrichment/jobs/raw/time_normalization_job.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+
+class TimeNormalizationJob:
+    def get_job_name(self) -> str:
+        return "time_normalization"
+
+    async def run(self, raw_records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        _ = raw_records
+        return []

topos/enrichment/jobs/raw/tool_calls_job.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+
+class ToolCallsJob:
+    def get_job_name(self) -> str:
+        return "tool_calls"
+
+    async def run(self, raw_records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        _ = raw_records
+        return []

topos/enrichment/models/__init__.py

@@ -0,0 +1 @@
+"""Enrichment model registry scaffolding."""

topos/enrichment/models/manager.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+
+class ModelManager:
+    """Placeholder for model loading/management."""
+
+    def __init__(self):
+        self._loaded = {}

topos/enrichment/models/registry.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional
+
+
+class ModelRegistry:
+    """In-memory model registry. Supports HuggingFace and Ollama providers."""
+
+    def __init__(self):
+        self._models: Dict[str, Dict[str, Any]] = {}
+
+    def register_model(
+        self,
+        model_id: str,
+        model_name: str,
+        model_version: str,
+        model_type: str,
+        task_name: str,
+        huggingface_path: str = "",
+        is_preferred: bool = False,
+        metadata: Optional[Dict[str, Any]] = None,
+        provider: Literal["ollama", "huggingface"] = "huggingface",
+        ollama_model: Optional[str] = None,
+    ) -> None:
+        self._models[model_id] = {
+            "model_id": model_id,
+            "model_name": model_name,
+            "model_version": model_version,
+            "model_type": model_type,
+            "task_name": task_name,
+            "huggingface_path": huggingface_path or "",
+            "is_preferred": is_preferred,
+            "metadata": metadata or {},
+            "provider": provider,
+            "ollama_model": ollama_model,
+        }
+
+    def get_model(self, model_id: str) -> Optional[Dict[str, Any]]:
+        return self._models.get(model_id)
+
+    def list_models(self, task_name: Optional[str] = None) -> List[Dict[str, Any]]:
+        models = list(self._models.values())
+        if task_name:
+            models = [m for m in models if m.get("task_name") == task_name]
+        return models
+
+    def get_preferred_model(self, task_name: str, model_type: str) -> Optional[Dict[str, Any]]:
+        for model in self._models.values():
+            if model.get("task_name") == task_name and model.get("model_type") == model_type and model.get("is_preferred"):
+                return model
+        return None
+
+    def get_model_for_task(
+        self,
+        task_type: str,
+        subtype: Optional[str] = None,
+        source_id: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Return model spec for a task (e.g. enrichment + url_classification).
+        Prefers is_preferred; otherwise first match by task_name or subtype.
+        Returns dict with huggingface_path (and later provider, ollama_model).
+        """
+        candidates = [
+            m for m in self._models.values()
+            if m.get("task_name") == task_type or (subtype and m.get("task_name") == subtype)
+        ]
+        if not candidates:
+            return None
+        preferred = [m for m in candidates if m.get("is_preferred")]
+        return (preferred[0] if preferred else candidates[0])

topos/enrichment/models/versioning.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+
+class ModelVersioning:
+    """Placeholder for model versioning helpers."""
+
+    def __init__(self):
+        self._versions = {}

topos/enrichment/orchestrator.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from ..utils.base_object import BaseObject
+from .derived_tables import DerivedTablesManager
+from .jobs import CANONICAL_JOBS, RAW_JOBS
+from .jobs.base import BaseEnrichmentJob
+
+logger = logging.getLogger("topos.enrichment.orchestrator")
+
+
+class EnrichmentOrchestrator(BaseObject):
+    def __init__(self, tables_manager: Optional[DerivedTablesManager] = None, *, name: Optional[str] = None):
+        super().__init__(name=name)
+        self.raw_jobs = list(RAW_JOBS)
+        self.canonical_jobs = list(CANONICAL_JOBS)
+        self.tables_manager = tables_manager or DerivedTablesManager()
+
+    def register_raw_job(self, job) -> None:
+        self.raw_jobs.append(job)
+
+    def register_canonical_job(self, job: BaseEnrichmentJob) -> None:
+        self.canonical_jobs.append(job)
+
+    async def run_raw(self, raw_records: List[Dict[str, Any]]) -> Dict[str, Any]:
+        results = {"jobs_run": 0, "records_created": {}, "errors": []}
+        for job in self.raw_jobs:
+            try:
+                records = await job.run(raw_records)
+                results["records_created"][job.get_job_name()] = len(records)
+                results["jobs_run"] += 1
+            except Exception as exc:
+                results["errors"].append({"job": job.get_job_name(), "error": str(exc)})
+        return results
+
+    async def run_canonical(
+        self,
+        canonical_messages: List[Dict[str, Any]],
+        job_names: Optional[List[str]] = None,
+        progress_callback: Optional[Callable[[int, int, str, float, float], None]] = None,
+    ) -> Dict[str, Any]:
+        """Run canonical enrichment jobs.
+        
+        Args:
+            canonical_messages: List of canonical message dictionaries
+            job_names: Optional list of specific job names to run
+            progress_callback: Optional callback function(processed_count, total_count, job_name) called during execution
+            
+        Returns:
+            Results dictionary with jobs_run, records_created, errors
+        """
+        results = {"jobs_run": 0, "records_created": {}, "errors": []}
+        jobs_to_run = self.canonical_jobs
+        if job_names:
+            jobs_to_run = [job for job in self.canonical_jobs if job.get_job_name() in job_names]
+        
+        total_messages = len(canonical_messages)
+        total_jobs = len(jobs_to_run)
+        logger.info(
+            "[PIPELINE:ENRICHMENT] %s: Starting enrichment: %d messages, %d jobs to run",
+            self,
+            total_messages,
+            total_jobs,
+        )
+        
+        # Track messages processed across all jobs
+        # For progress calculation: each job processes all messages, so we track cumulative progress
+        messages_processed_so_far = 0
+        
+        for job_idx, job in enumerate(jobs_to_run, 1):
+            if not job.should_run(canonical_messages):
+                logger.debug("[PIPELINE:ENRICHMENT] %s: Skipping job %s (should_run=False)", self, job.get_job_name())
+                continue
+            try:
+                job_name = job.get_job_name()
+                logger.info(
+                    "[PIPELINE:ENRICHMENT] %s: Running job %d/%d: %s (%d messages, %.1f%% of jobs complete)",
+                    self,
+                    job_idx,
+                    total_jobs,
+                    job_name,
+                    total_messages,
+                    ((job_idx - 1) / total_jobs * 100) if total_jobs > 0 else 0,
+                )
+                
+                # Create job-level progress callback
+                def job_progress_callback(current_count: int, total_count: int):
+                    """Callback for job-level progress updates."""
+                    if progress_callback:
+                        # Calculate job progress percent
+                        job_progress = (current_count / total_count * 100) if total_count > 0 else 0.0
+                        # Calculate jobs completion percent: (job_idx - 1) means previous jobs are done
+                        jobs_percent = ((job_idx - 1) / total_jobs * 100) if total_jobs > 0 else 0
+                        # Call orchestrator progress callback with job-level info
+                        progress_callback(
+                            processed_count=0,  # Not used for job-level tracking
+                            total_count=total_count,
+                            job_name=job_name,
+                            job_percent=jobs_percent,
+                            current_job_progress=job_progress,
+                        )
+                
+                # Call progress callback at start of job
+                if progress_callback:
+                    jobs_percent = ((job_idx - 1) / total_jobs * 100) if total_jobs > 0 else 0
+                    progress_callback(0, total_messages, job_name, jobs_percent, 0.0)
+                
+                records = await job.enrich(canonical_messages, progress_callback=job_progress_callback)
+                
+                # After job completes, calculate how many messages were effectively processed
+                # For jobs that create records, assume all messages were processed
+                # For jobs that return 0 records, they still "processed" the messages (just didn't create output)
+                messages_processed_this_job = total_messages if records else 0
+                messages_processed_so_far += messages_processed_this_job
+                
+                if records:
+                    # Write to derived table
+                    table_name = job.get_derived_table()
+                    records_written = self.tables_manager.write_enrichment_batch(
+                        records, table_name
+                    )
+                    results["records_created"][table_name] = records_written
+                    logger.info(
+                        "[PIPELINE:ENRICHMENT] %s → %s: %d records written to %s (job %d/%d, %.1f%% complete)",
+                        self,
+                        job,
+                        records_written,
+                        table_name,
+                        job_idx,
+                        total_jobs,
+                        (job_idx / total_jobs * 100) if total_jobs > 0 else 100,
+                    )
+                else:
+                    results["records_created"][job.get_derived_table()] = 0
+                    logger.info(
+                        "[PIPELINE:ENRICHMENT] %s → %s: completed with 0 records (job %d/%d, %.1f%% complete)",
+                        self,
+                        job,
+                        job_idx,
+                        total_jobs,
+                        (job_idx / total_jobs * 100) if total_jobs > 0 else 100,
+                    )
+                
+                # Call progress callback after job completes
+                if progress_callback:
+                    job_progress_percent = (job_idx / total_jobs * 100) if total_jobs > 0 else 100
+                    # Update with messages processed so far (cumulative across jobs)
+                    # Job is 100% complete
+                    progress_callback(messages_processed_so_far, total_messages, job_name, job_progress_percent, 100.0)
+                
+                results["jobs_run"] += 1
+            except Exception as exc:
+                logger.error(
+                    "[PIPELINE:ENRICHMENT] %s → %s: failed: %s (job %d/%d)",
+                    self,
+                    job,
+                    exc,
+                    job_idx,
+                    total_jobs,
+                )
+                results["errors"].append({"job": job.get_job_name(), "error": str(exc)})
+                # Still count this job's messages as "processed" (even if failed) for progress tracking
+                messages_processed_so_far += total_messages
+                if progress_callback:
+                    job_progress_percent = (job_idx / total_jobs * 100) if total_jobs > 0 else 100
+                    # Mark job as 100% complete (even if failed, we've moved past it)
+                    progress_callback(messages_processed_so_far, total_messages, job.get_job_name(), job_progress_percent, 100.0)
+        
+        logger.info(
+            "[PIPELINE:ENRICHMENT] %s: Enrichment complete: %d jobs run, %d total records created",
+            self,
+            results["jobs_run"],
+            sum(results["records_created"].values()),
+        )
+        return results

topos/enrichment/processor.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from .orchestrator import EnrichmentOrchestrator
+
+
+class EnrichmentProcessor:
+    def __init__(self, orchestrator: Optional[EnrichmentOrchestrator] = None):
+        self.orchestrator = orchestrator or EnrichmentOrchestrator()
+
+    async def process(
+        self,
+        canonical_messages: List[Dict[str, Any]],
+        job_names: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        return await self.orchestrator.run_canonical(canonical_messages, job_names=job_names)