workflow.py

from collections import deque
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Deque, List, Dict, Any, Optional

from temporalio import workflow
from temporalio.common import RetryPolicy

with workflow.unsafe.imports_passed_through():
    from config import settings
    from activities import (
        ContextualizeInput,
        RetrieveInput,
        RetrieveResult,
        GenerateInput,
        GenerateResult,
        contextualize_activity,
        retrieve_activity,
        generate_activity,
    )


NO_INFO_RESPONSE = "I am sorry, but I do not have that information."


def _activity_retry_policy(max_attempts: Optional[int] = None) -> RetryPolicy:
    return RetryPolicy(
        initial_interval=timedelta(seconds=settings.retries.initial_interval),
        backoff_coefficient=settings.retries.backoff_coefficient,
        maximum_interval=timedelta(seconds=settings.retries.maximum_interval),
        maximum_attempts=max_attempts or settings.retries.max_attempts,
        non_retryable_error_types=[
            "ContextualizationError",
            "RetrievalError",
            "GenerationError",
            "EmbeddingError",
            "IngestionError",
        ],
    )


@dataclass
class SendMessageInput:
    question: str
    thread_id: str


@dataclass
class SendMessageResult:
    answer: str
    contextualized_question: Optional[str] = None
    used_fallback_contextualize: bool = False
    skipped_generation: bool = False
    skip_reason: Optional[str] = None
    error: Optional[str] = None
    continued_as_new: bool = False


@dataclass
class ConversationWorkflowState:
    chat_history: List[Dict[str, Any]] = field(default_factory=list)
    thread_id: str = ""
    turns_completed: int = 0


@dataclass
class InFlightRequest:
    message: SendMessageInput
    result_slot: Optional[SendMessageResult] = None
    ready: bool = False


@workflow.defn
class ConversationWorkflow:
    def __init__(self) -> None:
        self._state = ConversationWorkflowState()
        self._queue: Deque[InFlightRequest] = deque()

    @workflow.run
    async def run(
        self,
        thread_id: str,
        carryover_history: Optional[List[Dict[str, Any]]] = None,
        carryover_turns: int = 0,
    ) -> None:
        self._state.thread_id = thread_id
        if carryover_history:
            self._state.chat_history = list(carryover_history)
        self._state.turns_completed = carryover_turns

        workflow.logger.info(
            f"ConversationWorkflow started for thread_id={thread_id} "
            f"(run={workflow.info().workflow_id}, turns_carried={carryover_turns}, "
            f"history_len={len(self._state.chat_history)})"
        )

        while True:
            await workflow.wait_condition(lambda: len(self._queue) > 0)

            inflight = self._queue.popleft()
            workflow.logger.info(
                f"Processing queued message for thread_id={thread_id}, "
                f"q_len={len(message.question) if (message := inflight.message) else 0}, "
                f"pending_queue={len(self._queue)}"
            )

            try:
                result = await self._process_message(inflight.message)
            except Exception as e:
                workflow.logger.error(
                    f"Unhandled exception in _process_message thread_id={thread_id}: {e}",
                    exc_info=True,
                )
                result = SendMessageResult(
                    answer=NO_INFO_RESPONSE,
                    error=f"Internal workflow error: {e}",
                )

            self._state.turns_completed += 1

            should_continue_as_new = (
                self._state.turns_completed > 0
                and self._state.turns_completed % settings.temporal.continue_as_new_turns == 0
            )

            if should_continue_as_new:
                workflow.logger.info(
                    f"Continue-as-new triggered at {self._state.turns_completed} turns "
                    f"(threshold={settings.temporal.continue_as_new_turns}). "
                    f"Carrying forward {len(self._state.chat_history)} history entries."
                )
                result.continued_as_new = True

                inflight.result_slot = result
                inflight.ready = True

                workflow.continue_as_new(
                    thread_id,
                    carryover_history=list(self._state.chat_history),
                    carryover_turns=self._state.turns_completed,
                )
                return

            inflight.result_slot = result
            inflight.ready = True

    async def _process_message(self, message: Optional[SendMessageInput]) -> SendMessageResult:
        if message is None:
            return SendMessageResult(answer=NO_INFO_RESPONSE, error="No message")

        thread_id = message.thread_id
        question = message.question
        current_history = list(self._state.chat_history)

        contextualized_question = question
        used_fallback = False

        try:
            ctx_input = ContextualizeInput(
                thread_id=thread_id,
                question=question,
                chat_history=current_history,
            )
            contextualized_question = await workflow.execute_activity(
                contextualize_activity,
                ctx_input,
                start_to_close_timeout=settings.timeouts.contextualize_td,
                retry_policy=_activity_retry_policy(),
            )
        except Exception as e:
            workflow.logger.warning(
                f"contextualize_activity failed for thread_id={thread_id}, falling back to raw question: {e}"
            )
            used_fallback = True
            contextualized_question = question

        try:
            retrieve_input = RetrieveInput(
                thread_id=thread_id,
                question=contextualized_question,
            )
            retrieve_result: RetrieveResult = await workflow.execute_activity(
                retrieve_activity,
                retrieve_input,
                start_to_close_timeout=settings.timeouts.retrieve_td,
                retry_policy=_activity_retry_policy(),
            )
        except Exception as e:
            workflow.logger.error(f"retrieve_activity failed for thread_id={thread_id}: {e}")
            return SendMessageResult(
                answer=NO_INFO_RESPONSE,
                contextualized_question=contextualized_question,
                used_fallback_contextualize=used_fallback,
                error=f"Retrieval failed after retries: {e}",
            )

        if retrieve_result.skipped:
            return SendMessageResult(
                answer=NO_INFO_RESPONSE,
                contextualized_question=contextualized_question,
                used_fallback_contextualize=used_fallback,
                skipped_generation=True,
                skip_reason=retrieve_result.skip_reason,
            )

        if not retrieve_result.docs:
            return SendMessageResult(
                answer=NO_INFO_RESPONSE,
                contextualized_question=contextualized_question,
                used_fallback_contextualize=used_fallback,
            )

        try:
            gen_input = GenerateInput(
                thread_id=thread_id,
                question=contextualized_question,
                chat_history=current_history,
                docs=retrieve_result.docs,
                original_question=question,
            )
            gen_result: GenerateResult = await workflow.execute_activity(
                generate_activity,
                gen_input,
                start_to_close_timeout=settings.timeouts.generate_td,
                retry_policy=_activity_retry_policy(),
            )
        except Exception as e:
            workflow.logger.error(f"generate_activity failed for thread_id={thread_id}: {e}")
            return SendMessageResult(
                answer=NO_INFO_RESPONSE,
                contextualized_question=contextualized_question,
                used_fallback_contextualize=used_fallback,
                error=f"Generation failed after retries: {e}",
            )

        self._state.chat_history.append(gen_result.human_message)
        self._state.chat_history.append(gen_result.ai_message)

        return SendMessageResult(
            answer=gen_result.answer,
            contextualized_question=contextualized_question,
            used_fallback_contextualize=used_fallback,
        )

    @workflow.update
    async def send_message_update(self, message: SendMessageInput) -> SendMessageResult:
        inflight = InFlightRequest(message=message)
        self._queue.append(inflight)

        await workflow.wait_condition(lambda: inflight.ready)

        return inflight.result_slot or SendMessageResult(
            answer=NO_INFO_RESPONSE, error="No result produced"
        )

    @workflow.query
    def get_chat_history(self) -> List[Dict[str, Any]]:
        return list(self._state.chat_history)

    @workflow.query
    def get_queue_depth(self) -> int:
        return len(self._queue)

    @workflow.query
    def get_turns_completed(self) -> int:
        return self._state.turns_completed