Thicket/thicket/pipeline_worker.py

"""Qt workers — threads that keep the UI responsive.

  * ``ProbeWorker``      — runs env_probe.probe_environment() off-thread.
  * ``PipelineWorker``   — wraps IngestPipeline, forwarding progress as
                           Qt signals (same signal contract as
                           OpenTranscode's EncoderWorker).
  * ``SearchWorker``     — embeds a query and searches Qdrant off-thread.
"""

from __future__ import annotations

from pathlib import Path

from PySide6.QtCore import QThread, Signal

from .embedder import DEFAULT_EMBED_MODEL, EmbeddingEngine, model_dim
from .env_probe import probe_environment
from .pipeline_core import IngestConfig, IngestPipeline, PipelineCallbacks


class ProbeWorker(QThread):
    """Background environment probe — service pings + module checks."""

    log_msg = Signal(str)
    probe_done = Signal(object)  # EnvProbe

    def __init__(self, qdrant_host: str = "localhost", qdrant_port: int = 6333,
                 parent=None):
        super().__init__(parent)
        self._host = qdrant_host
        self._port = qdrant_port

    def run(self):
        self.log_msg.emit("Probing environment (modules, Qdrant, Ollama)...")
        env = probe_environment(qdrant_host=self._host, qdrant_port=self._port)
        self.probe_done.emit(env)


class PipelineWorker(QThread):
    """Drives IngestPipeline on a worker thread.

    Signal contract mirrors OpenTranscode's EncoderWorker:
      log_msg(str)                — human-readable progress lines
      file_status(str, str)       — (absolute path, stage status)
      file_detail(str, str)       — (filename, detail text for table row)
      progress_msg(str, int, int) — (filename, current, total)
      finished_queue(int, int)    — (ok count, fail count)
    """

    log_msg = Signal(str)
    file_status = Signal(str, str)
    file_detail = Signal(str, str)
    progress_msg = Signal(str, int, int)
    finished_queue = Signal(int, int)

    def __init__(self, config: IngestConfig, parent=None):
        super().__init__(parent)
        self.config = config
        self._pipeline: IngestPipeline | None = None

    def stop(self) -> None:
        """Cooperative stop — the pipeline exits after the current file."""
        if self._pipeline is not None:
            self._pipeline.request_stop()

    def run(self):
        callbacks = PipelineCallbacks(
            log=self.log_msg.emit,
            file_status=self.file_status.emit,
            file_detail=self.file_detail.emit,
            progress=self.progress_msg.emit,
        )
        self._pipeline = IngestPipeline(self.config, callbacks)
        try:
            ok, fail = self._pipeline.run()
        except Exception as e:  # noqa: BLE001 — fatal pipeline errors
            self.log_msg.emit(f"FATAL: pipeline aborted — {e}")
            ok, fail = 0, 0
        self.finished_queue.emit(ok, fail)


class SearchWorker(QThread):
    """Semantic retrieval against the ingested collection of the
    selected vector target."""

    log_msg = Signal(str)
    search_done = Signal(int)  # number of hits returned

    def __init__(self, query: str, vault_dir: Path, target: str,
                 host: str = "localhost", port: int = 6333,
                 collection: str = "second_brain",
                 embed_model: str = DEFAULT_EMBED_MODEL, top_k: int = 5,
                 parent=None):
        super().__init__(parent)
        self._query = query
        self._vault_dir = vault_dir
        self._target = target
        self._host = host
        self._port = port
        self._collection = collection
        self._embed_model = embed_model
        self._top_k = top_k

    def run(self):
        try:
            hits = self._search()
        except Exception as e:  # noqa: BLE001 — report any failure to the log
            self.log_msg.emit(f"SEARCH ERROR: {e}")
            self.search_done.emit(-1)
            return
        self._report(hits)
        self.search_done.emit(len(hits))

    def _search(self) -> list[dict]:
        """Embed the query and search the collection (worker thread)."""
        from .vector_stores import create_store

        engine = EmbeddingEngine(self._embed_model)
        self.log_msg.emit(f"Embedding query with '{self._embed_model}'...")
        store = create_store(
            self._target,
            collection=self._collection,
            dim=model_dim(self._embed_model),
            host=self._host, port=self._port,
            data_dir=self._vault_dir / ".thicket" / self._target,
        )
        try:
            store.ensure_collection()
            return store.search(engine.embed_query(self._query),
                                limit=self._top_k)
        finally:
            store.close()

    def _report(self, hits: list[dict]) -> None:
        """Emit ranked results to the log, best first."""
        if not hits:
            self.log_msg.emit(
                f"No matches in '{self._collection}' — is anything ingested?"
            )
            return
        for rank, hit in enumerate(hits, start=1):
            payload = hit["payload"]
            snippet = (payload.get("content") or "")[:160].replace("\n", " ")
            tags = payload.get("chunk_kind", "prose")
            if payload.get("lang"):
                tags += f":{payload['lang']}"
            if payload.get("source_path"):
                tags += f"  src={payload['source_path']}"
            self.log_msg.emit(
                f"[{rank}] {hit['score']:.3f}  "
                f"{payload.get('document_title', '?')} "
                f"§ {payload.get('section_header', '?')}  ({tags})"
            )
            self.log_msg.emit(f"      {snippet}...")


class AskWorker(QThread):
    """Natural-language SQL over the corpus (Vanna 2 + Ollama)."""

    log_msg = Signal(str)
    ask_done = Signal(bool)  # success

    def __init__(self, question: str, target: str, collection: str,
                 llm_model: str, parent=None):
        super().__init__(parent)
        self._question = question
        self._target = target
        self._collection = collection
        self._llm_model = llm_model

    def run(self):
        from .ask_vanna import ask

        try:
            answer = ask(self._question, target=self._target,
                         collection=self._collection, llm_model=self._llm_model,
                         log=self.log_msg.emit)
        except Exception as e:  # noqa: BLE001 — report, never crash the UI
            self.log_msg.emit(f"ASK ERROR: {e}")
            self.ask_done.emit(False)
            return
        self.log_msg.emit(f"ANSWER: {answer}")
        self.ask_done.emit(True)