main.py

import asyncio
import logging
import os
import re
import signal
import subprocess
import sys
import time
import uuid
from datetime import timedelta
from pathlib import Path

from temporalio.client import Client, WorkflowHandle
from temporalio.common import RetryPolicy

from config import settings
from workflow import ConversationWorkflow, SendMessageInput, SendMessageResult
from test_drivers import AcceptanceIngestDriverWorkflow
from activities import IngestInput

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("lcqdrag.client")

_worker_subprocess: "subprocess.Popen | None" = None


def _start_background_worker() -> subprocess.Popen:
    """Launch the native Temporal worker as a detached daemon subprocess using
    the same venv/interpreter that is running main.py.

    Why this exists:
      `docker compose up` brings up the *infra* containers (postgres,
      temporal, qdrant) but the worker itself (which runs Python activities
      and has the Temporal SDK + Groq + Ollama + Qdrant client installed in
      the venv) runs natively on the host for two reasons:
        1. It avoids rebuilding the Dockerfile just to test code changes.
        2. Host-native Ollama is already running on :11434 with bge-large
           pulled, so the worker should talk to it directly (no extra
           networking gymnastics).
      If we forget this step, main.py looks like it "hangs" on the first
      execute_update() because no one is polling the task queue and Temporal
      patiently waits for a worker forever.
    """
    python_exe = sys.executable
    worker_py = Path(__file__).resolve().parent / "worker.py"
    log_dir = Path(__file__).resolve().parent / ".logs"
    log_dir.mkdir(exist_ok=True)
    log_path = log_dir / "worker.log"
    logf = open(log_path, "ab")

    env = os.environ.copy()
    # Prevent the child worker from inheriting asyncio / threading state from
    # the parent CLI process; start clean.
    env.pop("PYTHONASYNCIODEBUG", None)

    logger.info(f"Starting background worker (pid-file log -> {log_path})")
    proc = subprocess.Popen(
        [python_exe, str(worker_py)],
        stdout=logf,
        stderr=subprocess.STDOUT,
        cwd=str(Path(__file__).resolve().parent),
        env=env,
        start_new_session=True,
    )
    # Give the worker a moment to boot and connect to Temporal before we try
    # to schedule workflow tasks against its queue.
    time.sleep(3)
    rc = proc.poll()
    if rc is not None:
        raise RuntimeError(
            f"Background worker exited immediately with code {rc}. "
            f"Check tail of {log_path} for details."
        )
    logger.info(f"Background worker running: PID={proc.pid}")
    return proc


def _stop_background_worker():
    global _worker_subprocess
    if _worker_subprocess is None:
        return
    try:
        rc = _worker_subprocess.poll()
        if rc is None:
            logger.info(f"Stopping background worker PID={_worker_subprocess.pid}")
            _worker_subprocess.terminate()
            try:
                _worker_subprocess.wait(timeout=10)
            except subprocess.TimeoutExpired:
                logger.warning(
                    "Worker didn't exit in 10s after SIGTERM; sending SIGKILL"
                )
                _worker_subprocess.kill()
                _worker_subprocess.wait(timeout=5)
    except Exception as e:
        logger.warning(f"Ignoring error stopping background worker: {e}")
    finally:
        _worker_subprocess = None


class InputValidationError(Exception):
    pass


def validate_input(question: str) -> str:
    if question is None:
        raise InputValidationError("Question cannot be None")

    stripped = question.strip()
    if len(stripped) < settings.input.min_length:
        raise InputValidationError(
            f"Question too short (min {settings.input.min_length} characters)"
        )
    if len(stripped) > settings.input.max_length:
        raise InputValidationError(
            f"Question too long (max {settings.input.max_length} characters, got {len(stripped)})"
        )

    sanitized = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', stripped)
    return sanitized


async def ensure_ingestion(client: Client):
    logger.info(f"Ensuring documents ingested from `{settings.ingestion_url}`")
    try:
        ingest_wf_id = f"ingest-driver-{uuid.uuid4().hex[:8]}"
        result = await client.execute_workflow(
            AcceptanceIngestDriverWorkflow.run,
            IngestInput(url=settings.ingestion_url, thread_id="startup-ingestion"),
            id=ingest_wf_id,
            task_queue=settings.temporal.task_queue,
            execution_timeout=timedelta(seconds=600),
            retry_policy=RetryPolicy(
                initial_interval=timedelta(seconds=2.0),
                backoff_coefficient=2.0,
                maximum_interval=timedelta(seconds=60.0),
                maximum_attempts=3,
                non_retryable_error_types=["IngestionError"],
            ),
        )
        logger.info(f"Ingestion complete: {result}")
    except Exception as e:
        logger.warning(f"Ingestion failed (continuing anyway): {e}")


async def get_or_create_workflow(client: Client, thread_id: str) -> WorkflowHandle:
    workflow_id = f"{settings.temporal.workflow_id_prefix}{thread_id}"

    try:
        handle = client.get_workflow_handle(workflow_id)
        desc = await handle.describe()
        if desc.status.value == 1:
            logger.info(f"Attaching to existing running workflow: {workflow_id}")
            return handle
    except Exception:
        pass

    logger.info(f"Starting new workflow: {workflow_id}")
    handle = await client.start_workflow(
        ConversationWorkflow.run,
        thread_id,
        id=workflow_id,
        task_queue=settings.temporal.task_queue,
    )
    return handle


async def main():
    global _worker_subprocess

    # Start the background worker FIRST, so that (a) the startup ingestion
    # workflow has someone to run on and (b) the conversation workflow has a
    # task-queue consumer listening before main() dispatches updates.
    _worker_subprocess = _start_background_worker()

    def _install_shutdown_handlers():
        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                signal.signal(sig, lambda *_: sys.exit(0))
            except ValueError:
                pass  # not main thread; ignore

    _install_shutdown_handlers()

    try:
        logger.info(f"Connecting to Temporal at {settings.temporal.host}")
        client = await Client.connect(
            settings.temporal.host,
            namespace=settings.temporal.namespace,
        )
        logger.info("Connected to Temporal")

        await ensure_ingestion(client)

        thread_id = f"session_{uuid.uuid4().hex[:8]}"
        logger.info(f"Session thread_id: {thread_id}")

        workflow_handle = await get_or_create_workflow(client, thread_id)

        print("\n" + "=" * 60)
        print("  LCQDRAG — Production RAG Chatbot (Temporal-orchestrated)")
        print(f"  Thread ID: {thread_id}")
        print("  Type 'exit' to quit, 'history' to see chat history")
        print("=" * 60 + "\n")

        while True:
            try:
                question = input("Ask: ")
            except (EOFError, KeyboardInterrupt):
                print("\nGoodbye!")
                break

            if question.lower() == "exit":
                print("Goodbye!")
                break

            if question.lower() == "history":
                try:
                    history = await workflow_handle.query(ConversationWorkflow.get_chat_history)
                    print("\n--- Chat History ---")
                    for i, msg in enumerate(history):
                        role = "You" if msg["type"] == "human" else "Bot"
                        print(f"  [{role}] {msg['content'][:200]}")
                    print("--------------------\n")
                except Exception as e:
                    logger.error(f"Failed to fetch history: {e}")
                continue

            try:
                clean_question = validate_input(question)
            except InputValidationError as e:
                print(f"\n[Invalid input] {e}\n")
                continue

            if not clean_question:
                continue

            try:
                result: SendMessageResult = await workflow_handle.execute_update(
                    ConversationWorkflow.send_message_update,
                    SendMessageInput(question=clean_question, thread_id=thread_id),
                    rpc_timeout=timedelta(seconds=180),
                )
            except Exception as e:
                logger.error(f"Workflow update failed: {e}", exc_info=True)
                print(f"\n[Error] Failed to process message: {e}\n")
                continue

            print("\nAnswer:\n", result.answer)

            notes = []
            if result.used_fallback_contextualize:
                notes.append("(contextualize fallback used)")
            if result.skipped_generation and result.skip_reason:
                notes.append(f"(retrieval skipped: {result.skip_reason})")
            if result.error:
                notes.append(f"(error: {result.error})")
            if os.environ.get("VERBOSE_CONTEXTUALIZE", "0") in ("1", "true", "yes", "on"):
                if result.contextualized_question and result.contextualized_question != clean_question:
                    notes.append(f"(contextualized as: {result.contextualized_question[:120]})")

            if notes:
                print("  " + " ".join(notes))
            print()
    finally:
        _stop_background_worker()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nInterrupted, exiting.")
    except Exception as e:
        logger.error(f"Client exited with error: {e}", exc_info=True)
        sys.exit(1)
    finally:
        _stop_background_worker()
        sys.exit(0)