pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/processor/core/topic_splitter.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import re
+
+from langchain_core.prompts import ChatPromptTemplate
+
+from pulse_engine.processor.core.prompts import (
+    TRANSCRIPT_SPLITTER_EXAMPLES,
+    TRANSCRIPT_SPLITTER_PROMPT,
+)
+from pulse_engine.processor.llm.provider import get_llm
+from pulse_engine.processor.schemas import TopicSplitResult
+
+
+class TopicSplitter:
+    """LLM-backed topic splitter that detects topic shift points.
+
+    Accepts either a plain string (split into paragraphs) or a list of
+    ``(index, text)`` tuples.  Returns a :class:`TopicSplitResult` with
+    the detected shift indexes and reasons.
+    """
+
+    def __init__(
+        self,
+        provider: str | None = None,
+        model: str | None = None,
+        api_key: str | None = None,
+        temperature: float | None = None,
+        interaction_title: str = "",
+        account: str = "",
+        client: str = "",
+        account_attendees: str = "",
+        client_attendees: str = "",
+    ) -> None:
+        self._provider = provider
+        self._model = model
+        self._api_key = api_key
+        self._temperature = temperature
+        self.interaction_title = interaction_title
+        self.account = account
+        self.client = client
+        self.account_attendees = account_attendees
+        self.client_attendees = client_attendees
+
+    def split(self, content: str | list[tuple[int, str]]) -> TopicSplitResult:
+        """Detect topic shifts in *content*.
+
+        Parameters
+        ----------
+        content:
+            Either a plain string (paragraphs separated by blank lines)
+            or a list of ``(index, text)`` tuples.
+
+        Returns
+        -------
+        TopicSplitResult
+            Detected topic shifts with indexes and reasons.
+        """
+        tuples = self._normalize_input(content)
+        if len(tuples) <= 1:
+            return TopicSplitResult()
+        return self._call_llm(tuples)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _normalize_input(
+        content: str | list[tuple[int, str]],
+    ) -> list[tuple[int, str]]:
+        if isinstance(content, list):
+            return content
+        paragraphs = re.split(r"\n\n+", content.strip())
+        paragraphs = [p.strip() for p in paragraphs if p.strip()]
+        return [(i + 1, p) for i, p in enumerate(paragraphs)]
+
+    def _call_llm(self, tuples: list[tuple[int, str]]) -> TopicSplitResult:
+        llm = get_llm(
+            provider=self._provider,
+            model=self._model,
+            api_key=self._api_key,
+            temperature=self._temperature,
+        )
+
+        prompt = ChatPromptTemplate.from_template(TRANSCRIPT_SPLITTER_PROMPT)
+        structured_llm = llm.with_structured_output(TopicSplitResult)
+        chain = prompt | structured_llm
+
+        transcript_str = "\n".join(f'({idx}, "{text}")' for idx, text in tuples)
+        max_idx = max(idx for idx, _ in tuples)
+
+        result: TopicSplitResult = chain.invoke(
+            {
+                "examples": TRANSCRIPT_SPLITTER_EXAMPLES,
+                "transcript": transcript_str,
+                "max_idx": max_idx,
+                "interaction_title": self.interaction_title,
+                "account": self.account,
+                "client": self.client,
+                "account_attendees": self.account_attendees,
+                "client_attendees": self.client_attendees,
+            }
+        )
+        return result

pulse_engine/processor/defaults/__init__.py

@@ -0,0 +1,11 @@
+"""Default implementations of the pipeline stage ABCs."""
+
+from pulse_engine.processor.defaults.core_processor import DefaultCoreProcessor
+from pulse_engine.processor.defaults.postprocessor import DefaultPostprocessor
+from pulse_engine.processor.defaults.preprocessor import DefaultPreprocessor
+
+__all__ = [
+    "DefaultPreprocessor",
+    "DefaultCoreProcessor",
+    "DefaultPostprocessor",
+]

pulse_engine/processor/defaults/core_processor.py

@@ -0,0 +1,12 @@
+"""Default core processor — wraps the existing run_core_processing function."""
+
+from __future__ import annotations
+
+from pulse_engine.processor.base import BaseCoreProcessor
+from pulse_engine.processor.core.analysis import run_core_processing
+from pulse_engine.processor.schemas import ProcessingContext
+
+
+class DefaultCoreProcessor(BaseCoreProcessor):
+    def process(self, ctx: ProcessingContext) -> ProcessingContext:
+        return run_core_processing(ctx)

pulse_engine/processor/defaults/postprocessor.py

@@ -0,0 +1,12 @@
+"""Default postprocessor — wraps the existing run_postprocessing function."""
+
+from __future__ import annotations
+
+from pulse_engine.processor.base import BasePostprocessor
+from pulse_engine.processor.postprocessor.tasks import run_postprocessing
+from pulse_engine.processor.schemas import ProcessingContext
+
+
+class DefaultPostprocessor(BasePostprocessor):
+    def process(self, ctx: ProcessingContext) -> ProcessingContext:
+        return run_postprocessing(ctx)

pulse_engine/processor/defaults/preprocessor.py

@@ -0,0 +1,12 @@
+"""Default preprocessor — wraps the existing run_preprocessing function."""
+
+from __future__ import annotations
+
+from pulse_engine.processor.base import BasePreprocessor
+from pulse_engine.processor.preprocessor.tasks import run_preprocessing
+from pulse_engine.processor.schemas import ProcessingContext
+
+
+class DefaultPreprocessor(BasePreprocessor):
+    def process(self, ctx: ProcessingContext) -> ProcessingContext:
+        return run_preprocessing(ctx)

pulse_engine/processor/llm/provider.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pulse_engine.config import get_settings
+
+_llm_instance: Any = None
+
+
+def get_llm(
+    provider: str | None = None,
+    model: str | None = None,
+    api_key: str | None = None,
+    temperature: float | None = None,
+) -> Any:
+    """Return a LangChain chat model, using Settings as defaults.
+
+    Constructor params override env-var defaults.  The instance is cached
+    as a module singleton on first call.
+    """
+    global _llm_instance
+    if _llm_instance is not None:
+        return _llm_instance
+
+    settings = get_settings()
+    provider = provider or settings.pulse_llm_provider
+    model = model or settings.pulse_llm_model
+    api_key = api_key or settings.pulse_llm_api_key
+    if temperature is None:
+        temperature = settings.pulse_llm_temperature
+
+    if provider == "openai":
+        from langchain_openai import ChatOpenAI
+
+        _llm_instance = ChatOpenAI(
+            model=model,
+            api_key=api_key,  # type: ignore[arg-type]
+            temperature=temperature,
+        )
+    elif provider == "anthropic":
+        from langchain_anthropic import ChatAnthropic
+
+        _llm_instance = ChatAnthropic(  # type: ignore[call-arg]
+            model=model,
+            api_key=api_key,  # type: ignore[arg-type]
+            temperature=temperature,
+        )
+    else:
+        msg = f"Unsupported LLM provider: {provider!r}. Use 'openai' or 'anthropic'."
+        raise ValueError(msg)
+
+    return _llm_instance
+
+
+def reset_llm() -> None:
+    """Reset the cached LLM instance (useful for testing)."""
+    global _llm_instance
+    _llm_instance = None

pulse_engine/processor/ocr/gemini.py

@@ -0,0 +1,52 @@
+"""Gemini OCR provider using Google Generative AI."""
+
+from pathlib import Path
+from typing import Any
+
+from google import genai
+from google.genai import types
+
+from pulse_engine.processor.base import BaseOCRProvider
+from pulse_engine.processor.schemas import OCRInput
+
+
+class GeminiOCRProvider(BaseOCRProvider):
+    """OCR provider powered by Google Gemini's vision capabilities."""
+
+    def _get_client(self, api_key: str) -> Any:
+        return genai.Client(api_key=api_key)
+
+    def _read_file_bytes(self, ocr_input: OCRInput) -> bytes:
+        if ocr_input.file_bytes:
+            return ocr_input.file_bytes
+        if ocr_input.file_path:
+            return Path(ocr_input.file_path).read_bytes()
+        msg = "OCRInput must provide either file_bytes or file_path."
+        raise ValueError(msg)
+
+    def extract(self, ocr_input: OCRInput) -> Any:
+        file_bytes = self._read_file_bytes(ocr_input)
+        client = self._get_client(ocr_input.api_key)
+
+        response = client.models.generate_content(
+            model=ocr_input.model,
+            contents=[
+                types.Content(
+                    role="user",
+                    parts=[
+                        types.Part.from_bytes(
+                            data=file_bytes,
+                            mime_type=ocr_input.mime_type,
+                        ),
+                        types.Part.from_text(text=ocr_input.prompt),
+                    ],
+                ),
+            ],
+            config=types.GenerateContentConfig(
+                temperature=ocr_input.temperature,
+                max_output_tokens=ocr_input.max_output_tokens,
+                response_mime_type="application/json",
+            ),
+        )
+
+        return response

pulse_engine/processor/pipeline.py

@@ -0,0 +1,107 @@
+"""Processing pipeline — orchestrates pre-processing, core, and post-processing.
+
+Stages are pluggable via the base ABCs.  Pass ``None`` to skip a stage,
+or omit (use the sentinel default) to get the engine's built-in implementation.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import structlog
+
+from pulse_engine.processor.base import (
+    BaseCoreProcessor,
+    BasePostprocessor,
+    BasePreprocessor,
+)
+from pulse_engine.processor.schemas import ProcessingContext, ProcessingError
+from pulse_engine.storage.knowledge_base import KnowledgeBaseService
+
+logger = structlog.get_logger()
+
+_SENTINEL: Any = object()
+
+
+class ProcessingPipeline:
+    def __init__(
+        self,
+        kb_service: KnowledgeBaseService | None = None,
+        preprocessor: BasePreprocessor | None = _SENTINEL,
+        core_processor: BaseCoreProcessor | None = _SENTINEL,
+        postprocessor: BasePostprocessor | None = _SENTINEL,
+    ) -> None:
+        from pulse_engine.processor.defaults import (
+            DefaultCoreProcessor,
+            DefaultPostprocessor,
+            DefaultPreprocessor,
+        )
+
+        self._kb_service = kb_service
+        self._preprocessor = (
+            DefaultPreprocessor() if preprocessor is _SENTINEL else preprocessor
+        )
+        self._core_processor = (
+            DefaultCoreProcessor() if core_processor is _SENTINEL else core_processor
+        )
+        self._postprocessor = (
+            DefaultPostprocessor() if postprocessor is _SENTINEL else postprocessor
+        )
+
+    async def run(self, ctx: ProcessingContext) -> ProcessingContext:
+        """Run the full 3-stage pipeline."""
+        ctx = await self.run_preprocess(ctx)
+        ctx = await self.run_core(ctx)
+        ctx = await self.run_postprocess(ctx)
+
+        if ctx.options.store_results and self._kb_service and ctx.documents:
+            try:
+                await self._kb_service.store_documents(ctx.tenant_id, ctx.documents)
+            except Exception as e:
+                logger.warning("storage_failed", error=str(e))
+                ctx.errors.append(
+                    ProcessingError(
+                        stage="storage", task="store_documents", message=str(e)
+                    )
+                )
+
+        return ctx
+
+    async def run_preprocess(self, ctx: ProcessingContext) -> ProcessingContext:
+        if self._preprocessor is None:
+            return ctx
+        try:
+            ctx = self._preprocessor.process(ctx)
+        except Exception as e:
+            ctx.errors.append(
+                ProcessingError(
+                    stage="preprocessing", task="run_preprocess", message=str(e)
+                )
+            )
+        return ctx
+
+    async def run_core(self, ctx: ProcessingContext) -> ProcessingContext:
+        if self._core_processor is None:
+            return ctx
+        try:
+            ctx = self._core_processor.process(ctx)
+        except Exception as e:
+            ctx.errors.append(
+                ProcessingError(
+                    stage="core_processing", task="run_core", message=str(e)
+                )
+            )
+        return ctx
+
+    async def run_postprocess(self, ctx: ProcessingContext) -> ProcessingContext:
+        if self._postprocessor is None:
+            return ctx
+        try:
+            ctx = self._postprocessor.process(ctx)
+        except Exception as e:
+            ctx.errors.append(
+                ProcessingError(
+                    stage="postprocessing", task="run_postprocess", message=str(e)
+                )
+            )
+        return ctx

pulse_engine/processor/postprocessor/embeddings.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pulse_engine.config import get_settings
+
+
+def generate_embedding(
+    text: str,
+    model_name: str = "text-embedding-3-small",
+    provider: str | None = None,
+) -> list[float]:
+    """Generate embedding for *text* using OpenAI.
+
+    Parameters
+    ----------
+    text:
+        The text to embed.
+    model_name:
+        Ignored (kept for backward compatibility). Model is read from settings.
+    provider:
+        Ignored (kept for backward compatibility). Always uses OpenAI.
+    """
+    settings = get_settings()
+    return _openai_embedding(text, settings)
+
+
+def _openai_embedding(text: str, settings: Any) -> list[float]:
+    from langchain_openai import OpenAIEmbeddings
+
+    api_key = settings.pulse_openai_api_key or settings.pulse_llm_api_key
+    model = settings.pulse_openai_embedding_model
+    embedder = OpenAIEmbeddings(model=model, api_key=api_key)
+    return list(embedder.embed_query(text))

pulse_engine/processor/postprocessor/tasks.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from pulse_engine.processor.postprocessor.embeddings import generate_embedding
+from pulse_engine.processor.schemas import (
+    ContentChunk,
+    ProcessingContext,
+    ProcessingError,
+)
+from pulse_engine.storage.schemas import Document, DocumentSource
+
+
+def format_for_storage(ctx: ProcessingContext) -> list[Document]:
+    """Convert processing context chunks into Document objects for storage."""
+    # Map source types to DocumentSource enum
+    source_map: dict[str, DocumentSource] = {
+        "speech": DocumentSource.speech,
+        "tweet": DocumentSource.tweet,
+        "youtube": DocumentSource.youtube,
+        "processed": DocumentSource.processed,
+    }
+    doc_source = source_map.get(ctx.source_type, DocumentSource.processed)
+
+    documents: list[Document] = []
+    for i, chunk in enumerate(ctx.chunks):
+        metadata: dict[str, Any] = {
+            "source_type": ctx.source_type,
+            "product": ctx.product,
+            "tenant_id": ctx.tenant_id,
+            "language": ctx.language,
+        }
+
+        metadata_chunk: dict[str, Any] = {
+            "chunk_index": chunk.chunk_index,
+            "token_count": chunk.token_count,
+            "sentiment": chunk.sentiment,
+            "topics": chunk.topics,
+            "entities": [e.model_dump() for e in chunk.entities],
+        }
+
+        # Get embeddings if available
+        content_embedding = ctx.embeddings[i] if i < len(ctx.embeddings) else None
+
+        doc = Document(
+            doc_id=ctx.source_id,
+            source=doc_source,
+            chunk_id=str(chunk.chunk_index),
+            content=chunk.content,
+            metadata=metadata,
+            metadata_chunk=metadata_chunk,
+            content_summary=ctx.summary,
+            content_title=None,
+            embedding=content_embedding,
+            content_embeddings=content_embedding,
+            content_summary_embeddings=None,
+            content_title_embeddings=None,
+        )
+        documents.append(doc)
+
+    return documents
+
+
+def compute_quality_score(chunk: ContentChunk) -> float:
+    """Compute a heuristic quality score for a chunk (0.0 - 1.0)."""
+    score = 0.0
+
+    # Content length factor (0 - 0.4)
+    length = len(chunk.content)
+    if length >= 100:
+        score += 0.4
+    elif length >= 50:
+        score += 0.3
+    elif length >= 20:
+        score += 0.2
+    else:
+        score += 0.1
+
+    # Entity count factor (0 - 0.2)
+    entity_count = len(chunk.entities)
+    score += min(entity_count * 0.05, 0.2)
+
+    # Topic count factor (0 - 0.2)
+    topic_count = len(chunk.topics)
+    score += min(topic_count * 0.04, 0.2)
+
+    # Token count factor (0 - 0.2) — prefer chunks with reasonable token counts
+    if 50 <= chunk.token_count <= 512:
+        score += 0.2
+    elif chunk.token_count > 0:
+        score += 0.1
+
+    return min(score, 1.0)
+
+
+def _cosine_similarity(a: list[float], b: list[float]) -> float:
+    """Compute cosine similarity between two vectors."""
+    dot = sum(x * y for x, y in zip(a, b))
+    norm_a = math.sqrt(sum(x * x for x in a))
+    norm_b = math.sqrt(sum(x * x for x in b))
+    if norm_a == 0 or norm_b == 0:
+        return 0.0
+    return dot / (norm_a * norm_b)
+
+
+def deduplicate(embeddings: list[list[float]], threshold: float = 0.95) -> list[bool]:
+    """
+    Check for duplicate embeddings using cosine similarity.
+    Returns list of bools: True = keep, False = duplicate.
+    """
+    if not embeddings:
+        return []
+
+    n = len(embeddings)
+    keep = [True] * n
+
+    if n <= 1:
+        return keep
+
+    for i in range(n):
+        if not keep[i]:
+            continue
+        for j in range(i + 1, n):
+            if not keep[j]:
+                continue
+            if _cosine_similarity(embeddings[i], embeddings[j]) >= threshold:
+                keep[j] = False
+
+    return keep
+
+
+def run_postprocessing(ctx: ProcessingContext) -> ProcessingContext:
+    """Run all postprocessing tasks on the context."""
+    try:
+        # Generate embeddings
+        if ctx.options.enable_embedding:
+            for chunk in ctx.chunks:
+                emb = generate_embedding(
+                    chunk.content,
+                    ctx.options.embedding_model,
+                    provider=ctx.options.embedding_provider,
+                )
+                ctx.embeddings.append(emb)
+
+        # Compute quality scores
+        for chunk in ctx.chunks:
+            score = compute_quality_score(chunk)
+            ctx.quality_scores.append(score)
+
+        # Deduplication
+        if ctx.options.enable_dedup and ctx.embeddings:
+            keep_flags = deduplicate(ctx.embeddings)
+            # Filter out duplicates
+            new_chunks = []
+            new_embeddings = []
+            new_scores = []
+            for i, keep in enumerate(keep_flags):
+                if keep:
+                    new_chunks.append(ctx.chunks[i])
+                    if i < len(ctx.embeddings):
+                        new_embeddings.append(ctx.embeddings[i])
+                    if i < len(ctx.quality_scores):
+                        new_scores.append(ctx.quality_scores[i])
+            ctx.chunks = new_chunks
+            ctx.embeddings = new_embeddings
+            ctx.quality_scores = new_scores
+
+        # Format for storage
+        ctx.documents = format_for_storage(ctx)
+
+        ctx.stages_completed.append("postprocessing")
+    except Exception as e:
+        ctx.errors.append(
+            ProcessingError(
+                stage="postprocessing", task="run_postprocessing", message=str(e)
+            )
+        )
+
+    return ctx

pulse_engine/processor/preprocessor/tasks.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import re
+import unicodedata
+
+from bs4 import BeautifulSoup
+from langdetect import detect
+from langdetect.lang_detect_exception import LangDetectException
+
+from pulse_engine.processor.schemas import ProcessingContext, ProcessingError
+
+
+def clean_html(text: str) -> str:
+    """Strip HTML tags, decode entities, and normalize whitespace."""
+    soup = BeautifulSoup(text, "html.parser")
+    cleaned = str(soup.get_text(separator=" "))
+    # Normalize whitespace
+    cleaned = re.sub(r"\s+", " ", cleaned).strip()
+    return cleaned
+
+
+def normalize_text(text: str) -> str:
+    """Apply Unicode NFC normalization and strip control characters."""
+    text = unicodedata.normalize("NFC", text)
+    # Remove control characters except newlines and tabs
+    text = "".join(
+        ch
+        for ch in text
+        if not unicodedata.category(ch).startswith("C") or ch in "\n\t"
+    )
+    return text
+
+
+def detect_language(text: str) -> str:
+    """Detect language using langdetect, return ISO 639-1 code."""
+    try:
+        return str(detect(text))
+    except LangDetectException:
+        return "unknown"
+
+
+def validate_content(text: str, min_length: int = 10) -> bool:
+    """Reject empty or too-short content."""
+    if not text or not text.strip():
+        return False
+    return len(text.strip()) >= min_length
+
+
+def run_preprocessing(ctx: ProcessingContext) -> ProcessingContext:
+    """Run all preprocessing tasks in sequence on the context."""
+    try:
+        cleaned = clean_html(ctx.raw_content)
+        cleaned = normalize_text(cleaned)
+        ctx.cleaned_content = cleaned
+
+        if validate_content(cleaned):
+            ctx.language = detect_language(cleaned)
+        else:
+            ctx.language = "unknown"
+
+        ctx.stages_completed.append("preprocessing")
+    except Exception as e:
+        ctx.errors.append(
+            ProcessingError(
+                stage="preprocessing",
+                task="run_preprocessing",
+                message=str(e),
+            )
+        )
+
+    return ctx