app/workers/summarization_worker.py

from datetime import datetime, timezone
import time
from typing import Optional

from sqlalchemy import select
from sqlalchemy.orm import Session

import app.models.smart_models  # noqa: F401 — registers UserPreference and related ORM classes
from app.models.database import (
    SessionLocal,
    SummarizationJob,
    Conversation,
    Message,
    Tenant,
)
from app.utils.config import config_manager, get_logger
from app.models.conversation_manager import ConversationManager, CHARS_PER_TOKEN
from app.utils.token_counter import count_tokens

config = config_manager.get_config()
logger = get_logger(__name__)


def fetch_pending_job(db: Session) -> Optional[SummarizationJob]:

    job = db.execute(
        select(SummarizationJob)
        .where(SummarizationJob.status == "pending")
        .order_by(SummarizationJob.created_at.asc())
        .limit(1)
    ).scalar_one_or_none()

    return job


def process_job(
    job: SummarizationJob, db: Session, conversation_manager: ConversationManager
) -> None:

    ##########################################################################
    # Split verbatim and displaced messages with _split_verbatim_and_displaced
    ##########################################################################

    messages = db.scalars(
        select(Message)
        .where(Message.conversation_id == job.conversation_id)
        .order_by(Message.timestamp.asc())
    ).all()

    formatted_messages = [
        {"user_input": msg.user_input, "bot_response": msg.bot_response}
        for msg in messages
    ]

    conversation = db.scalar(
        select(Conversation).where(Conversation.id == job.conversation_id)
    )

    if not conversation:
        raise ValueError(f"Conversation not found for job {job.id}")

    tenant = db.get(Tenant, conversation.tenant_id)

    if not tenant:
        raise ValueError(f"Tenant not found for conversation {conversation.id}")

    # Calculate the budget tokens
    system_prompt = (tenant.customization or {}).get("system_prompt")
    system_tokens = count_tokens(
        [{"role": "system", "content": system_prompt or config.system_prompt}]
    )

    summary_tokens = (
        len(conversation.context_summary) // CHARS_PER_TOKEN
        if conversation.context_summary
        else 0
    )

    max_input_tokens = config.nlp["max_input_tokens"]
    context_reserve_tokens = config.nlp["context_reserve_tokens"]

    budget = max_input_tokens - system_tokens - summary_tokens - context_reserve_tokens

    # Split messages into `verbatim` and `displaced`
    _, displaced = (
        conversation_manager._split_verbatim_and_displaced(  # pylint: disable=protected-access
            messages=formatted_messages, budget=budget
        )
    )

    ##########################################################################
    # Summarize the displaced messages
    ##########################################################################
    context_summary = (
        conversation_manager._summarize_messages(  # pylint: disable=protected-access
            messages=displaced
        )
    )

    ##########################################################################
    # Update the context_summary inside `conversations` table
    ##########################################################################
    conversation.context_summary = context_summary

    # Delete the summarization_job
    db.delete(job)
    db.commit()


def handle_failure(job: SummarizationJob, db: Session) -> None:

    max_attempt_count = config.worker["max_attempts"]

    job.attempt_count += 1
    job.last_attempted_at = datetime.now(timezone.utc)

    if job.attempt_count >= max_attempt_count:
        job.status = "dead"

    db.commit()


def main() -> None:
    conversation_manager = ConversationManager(config=config)

    while True:

        with SessionLocal() as db:
            job = fetch_pending_job(db=db)

            if job:
                try:
                    process_job(
                        job=job,
                        db=db,
                        conversation_manager=conversation_manager,
                    )

                except Exception as e:  # pylint: disable=broad-except
                    handle_failure(job=job, db=db)
                    logger.info(
                        "Some error has occured: %s: %s", type(e).__name__, str(e)
                    )

            else:
                time.sleep(config.worker["poll_interval"])


if __name__ == "__main__":
    main()