Integrating Human-in-the-loop components to your tools

# /// script
# dependencies = [
#   "requests",
# ]
# ///

from typing import Any, Optional, List, Dict
from shinkai_local_tools import (
    shinkai_llm_prompt_processor,
    clarification_dialog_human_in_the_loop,
    confirmation_dialog_human_in_the_loop,
)

class CONFIG:
    show_individual_answers: bool

class INPUTS:
    initial_question: str

class OUTPUT:
    final_answer: str
    researched_subquestions_summary: List[str]
    individual_answers: Optional[Dict[str, str]]

async def ask_number_of_subquestions(initial_question: str) -> int:
    clarification = await clarification_dialog_human_in_the_loop({
        "title": "Number of Sub-questions",
        "question": (
            f"Your research question is:\n\n\"{initial_question}\"\n\n"
            "How many sub-questions would you like to generate to help research this question? "
            "Please enter a number (for example, 3)."
        ),
        "label_submit_button": "Set Number",
        "label_ignore_button": "Cancel",
    })
    # Try to extract a positive integer from the user's answer
    response = clarification.get("clarification_response", "").strip()
    try:
        n = int(response)
        if n < 1:
            n = 1
    except Exception:
        n = 3
    return n

async def generate_subquestions(initial_question: str, n: int) -> List[str]:
    prompt = (
        f"Given the research question:\n\"{initial_question}\"\n"
        f"Generate {n} concise and relevant sub-questions that, if answered, would help address the main question. "
        "Output each sub-question as a separate line, without numbering."
    )
    result = await shinkai_llm_prompt_processor({
        "prompt": prompt,
        "format": "text"
    })
    message = result.get("message", "")
    subquestions = [line.strip("-•. ").strip() for line in message.splitlines() if line.strip()]
    # Only keep up to n subquestions, in case of LLM over-generation
    return subquestions[:n]

async def edit_subquestion(subq: str) -> str:
    clarification = await clarification_dialog_human_in_the_loop({
        "title": "Edit Sub-question",
        "question": (
            f"The following sub-question was generated:\n\n\"{subq}\"\n\n"
            "If you would like to edit or improve this sub-question, please enter your feedback. "
            "Otherwise, you can leave this blank and submit to keep it as is."
        ),
        "label_submit_button": "Edit/Keep",
        "label_ignore_button": "Skip Editing"
    })
    feedback = clarification.get("clarification_response", "").strip()
    if not feedback:
        return subq
    # Use LLM to process feedback and generate improved sub-question
    prompt = (
        f"Original sub-question:\n\"{subq}\"\n\n"
        f"User feedback for improvement or rewording:\n\"{feedback}\"\n\n"
        "Rewrite the sub-question based on the user's feedback. Output only the improved sub-question."
    )
    result = await shinkai_llm_prompt_processor({
        "prompt": prompt,
        "format": "text"
    })
    edited = result.get("message", "").strip()
    return edited if edited else subq

async def confirm_subquestion(subq: str) -> bool:
    confirmation = await confirmation_dialog_human_in_the_loop({
        "title": "Confirm Sub-question for Research",
        "text": (
            f"Do you want to include the following sub-question for research?\n\n\"{subq}\""
        ),
        "label_yes_button": "Include",
        "label_no_button": "Exclude"
    })
    choice = confirmation.get("user_choice", "").strip().lower()
    return choice == "include"

async def answer_question(question: str) -> str:
    prompt = (
        f"Research and answer the following question in a detailed and informative manner:\n\n\"{question}\""
    )
    result = await shinkai_llm_prompt_processor({
        "prompt": prompt,
        "format": "text"
    })
    answer = result.get("message", "").strip()
    return answer

async def run(config: CONFIG, inputs: INPUTS) -> OUTPUT:
    output = OUTPUT()
    initial_question = inputs.initial_question.strip()

    # Step 1: Ask how many sub-questions to generate
    num_subq = await ask_number_of_subquestions(initial_question)

    # Step 2: Generate sub-questions
    subquestions = await generate_subquestions(initial_question, num_subq)

    # Step 3: For each sub-question, allow user to edit and confirm inclusion
    final_subquestions = []
    for sq in subquestions:
        # Allow user to edit the sub-question
        edited_sq = await edit_subquestion(sq)
        # Ask for confirmation to include for research
        include = await confirm_subquestion(edited_sq)
        if include:
            final_subquestions.append(edited_sq)

    # Always include the initial question as the "main" question to research
    questions_to_research = [initial_question] + final_subquestions

    # Step 4: Research all questions (main + confirmed sub-questions)
    answers = {}
    for q in questions_to_research:
        answers[q] = await answer_question(q)

    # Step 5: Build final answer and output structure
    # Compose a detailed answer using the initial question + sub-questions' answers
    main_answer = answers[initial_question]
    subq_answers = [
        f"Sub-question: \"{sq}\"\nAnswer: {answers[sq]}"
        for sq in final_subquestions
    ]
    detailed_answer = (
        f"Main Question:\n\"{initial_question}\"\n\n"
        f"Comprehensive Answer:\n{main_answer}\n\n"
    )
    if final_subquestions:
        detailed_answer += "Researched Sub-questions and their Answers:\n\n"
        for sq in final_subquestions:
            detailed_answer += (
                f"Sub-question:\n\"{sq}\"\n"
                f"Answer:\n{answers[sq]}\n\n"
            )

    output.final_answer = detailed_answer
    output.researched_subquestions_summary = final_subquestions
    if config.show_individual_answers:
        output.individual_answers = answers
    else:
        output.individual_answers = None
    return output