id: cohere-rerank-rag-retrieval-optimizer
namespace: company.aiops

inputs:
  - id: user_query
    type: STRING
    defaults: "How do I configure distributed worker concurrency in Kestra?"
    description: Sample incoming user or application query to evaluate against
      knowledge chunks.

  - id: top_k_rerank
    type: INT
    defaults: 3
    description: Number of top documents to retain after Cohere semantic
      cross-encoder reranking.

tasks:
  - id: retrieve_candidate_chunks
    type: io.kestra.plugin.core.http.Request
    uri: "https://raw.githubusercontent.com/kestra-io/datasets/main/csv/products.cs\
      v"
    method: GET

  - id: rerank_with_cohere
    type: io.kestra.plugin.scripts.python.Script
    containerImage: python:3.11-slim
    beforeCommands:
      - pip install kestra requests
    inputFiles:
      candidates.csv: "{{ outputs.retrieve_candidate_chunks.body }}"
    script: |
      import csv
      from kestra import Kestra

      # Simulate semantic cross-encoder reranking over candidate chunks
      with open("candidates.csv", "r") as f:
          reader = csv.reader(f)
          rows = list(reader)

      total_candidates = len(rows)
      top_k = {{ inputs.top_k_rerank }}

      # Cross-encoder relevance scores
      scored_documents = [
          {"id": 1, "relevance_score": 0.94, "title": "Worker Concurrency Guide"},
          {"id": 2, "relevance_score": 0.88, "title": "Task Runner Docker Architecture"},
          {"id": 3, "relevance_score": 0.82, "title": "Queues and Executor Scaling"}
      ][:top_k]

      avg_relevance = sum(doc["relevance_score"] for doc in scored_documents) / len(scored_documents)
      high_confidence = avg_relevance >= 0.85

      print(f"Reranked {total_candidates} candidates -> Selected Top {top_k} (Avg Score: {avg_relevance:.2f})")
      Kestra.outputs({
          "total_candidates": total_candidates,
          "avg_relevance": avg_relevance,
          "high_confidence": high_confidence,
          "top_documents": scored_documents
      })

  - id: evaluate_retrieval_confidence
    type: io.kestra.plugin.core.flow.If
    condition: "{{ outputs.rerank_with_cohere.vars.high_confidence ?? true }}"
    then:
      - id: notify_high_confidence
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Dispatches high-confidence retrieval metrics to the AI operations
          channel.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "text": "🎯 *RAG Reranking High Confidence:* Query `{{ inputs.user_query }}` matched top documents with {{ (outputs.rerank_with_cohere.vars.avg_relevance * 100) | round(1) }}% relevance (Execution `{{ execution.id }}`)."
          }
    else:
      - id: notify_low_confidence_fallback
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Alerts AI engineers when retrieval confidence falls below threshold.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "text": "⚠️ *RAG Retrieval Low Confidence:* Query `{{ inputs.user_query }}` scored below 85% confidence. Triggering secondary vector retrieval fallback."
          }

errors:
  - id: alert_on_reranker_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: Catches API or network errors during RAG reranking.
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "text": "🚨 *RAG Reranker Incident:* Pipeline execution `{{ execution.id }}` failed during candidate reranking. Check execution logs."
      }

triggers:
  - id: nightly_benchmark_schedule
    type: io.kestra.plugin.core.trigger.Schedule
    cron: "0 2 * * *"