tools/reasoning_tools.py

from pydantic import BaseModel, Field
from langchain_core.tools import tool
from src.helper import load_catalog

catalog = load_catalog()
catalog_by_id = { cat['entity_id'] : cat for cat in catalog }
keys_mapping = {'Ability & Aptitude' : 'A', 
                'Assessment Exercises' : 'E', 
                'Biodata & Situational Judgment' : 'B',
                'Competencies' : 'C', 
                'Development & 360' : 'D', 
                'Knowledge & Skills' : 'K',
                'Personality & Behavior' : 'P',
                'Simulations' : 'S'}


class ClarificationInput(BaseModel):

    clarification_message: str = Field(
        ...,
        description=(
            "A concise clarification question or decision "
            "branch based on retrieved SHL catalog evidence."
        )
    )

@tool(
    args_schema=ClarificationInput,
    description="""
Use this tool when retrieved assessments reveal
important ambiguity or decision branches that
must be resolved before final recommendations.

Examples:
- language/accent variants
- seniority calibration differences
- development vs hiring usage
- industry-specific calibration
- personality-only vs full battery tradeoffs
- simulation variant selection

The clarification should reference retrieved
catalog evidence whenever useful.
"""
)
def clarification(clarification_message: str):

    return {
        "type": "clarification",
        "message": clarification_message
    }


class ComparisonInput(BaseModel):

    comparison_response: str = Field(
        ...,
        description=(
            "A focused comparison between assessments "
            "explicitly requested by the user."
            """
            Comparison responses should:
            - explain practical differences
            - explain intended usage differences
            - explain calibration differences
            - explain standalone vs bundled behavior
            - explain stage-of-hiring differences
            """
        )
    )

    compared_assessments: list[str] = Field(
        ...,
        description=(
            "Exact assessment names involved in the comparison."
        )
    )


@tool(
    args_schema=ComparisonInput,
    description="""
        Use comparison ONLY when the user explicitly asks:
        - differences
        - comparison
        - tradeoffs
        - which is better
        - whether two assessments overlap

        Compare ONLY:
        - assessments explicitly referenced by the user
        OR
        - assessments central to the discussion

        Ignore unrelated retrieved documents.

        Comparison responses should:
        - explain practical differences
        - explain intended usage differences
        - explain calibration differences
        - explain standalone vs bundled behavior
        - explain stage-of-hiring differences

        Examples:
        - DSI vs Safety & Dependability 8.0
        - OPQ32r vs OPQ MQ Sales Report
        - Contact Center Simulation vs Customer Service Phone Simulation
"""
)
def comparison(
    comparison_response: str,
    compared_assessments: list[str]
):

    return {
        "type": "comparison",
        "message": comparison_response,
        "assessments": compared_assessments
    }



class RecommenderInput(BaseModel):

    recommendation_summary: str = Field(
        ...,
        description=(
            "A concise explanation of why the selected "
            "assessments fit the user's hiring or "
            "development needs."
            """
            The `recommendation_summary` MUST be written as a direct conversational response to the user.

                Speak TO the user.

                Do NOT describe the user in third person.
                Do NOT mention any entity id in recommendation summary.

                BAD:
                - "The user seeks..."
                - "The candidate requires..."
                - "This request involves..."

                GOOD:
                - "For your senior leadership benchmarking use case..."
                - "Since you're hiring CXOs and directors..."
                - "I'd recommend..."
                - "You could combine..."
            """
        )
    )

    recommended_assessments_ids: list[str] = Field(
        ...,
        description=(
            "Exact assessment entity ids selected from the "
            "retrieved catalog entries."
            """
            You do NOT need to recommend all retrieved documents.
            But try to recommend atleast 3-5 assessment everytime.

            Good recommendation counts:
            - sometimes 3
            - sometimes 4
            - sometimes 7

            Selection quality matters more than quantity.
            If there are two versions of same assessment name like 1.0 and 2.0 then only select 2.0 version.

            Recommendations should:
            - align to role
            - align to seniority
            - align to hiring stage
            - align to assessment goals
            - align to language constraints
            - align to technical specialization

            IMPORTANT:
            Use the EXACT assessment entity id from retrieved docs.
            Never invent Entity id.
            """
        )
    )

    
    end_of_conversation: bool = Field(
        ...,
        description=(
            "Whether the conversation is truly complete "
            "based on explicit user confirmation or acceptance "
            "of the recommendation shortlist.\n\n"
            
            "IMPORTANT:\n"
            "Recommendations alone do NOT mean the "
            "conversation is complete.\n\n"
            
            "Set to True ONLY if the user clearly confirms "
            "or accepts the recommendations.\n\n"
            
            "Examples where True is appropriate:\n"
            "- Perfect" 
            "- Looks good" 
            "- That works" 
            "- Confirmed" 
            "- Thanks" 
            "- Please confirm these final recommended assessments" 
            "- These recommendations work for us" 
            "- Let's proceed with these" 
            "- This is exactly what we need"
            
            "Examples where False is appropriate:\n"
            "- first-time recommendations\n"
            "- recommendations followed by possible refinements\n"
            "- unresolved recommendation branches\n"
            "- situations where follow-up clarification may still help\n\n"
            
            "Initial recommendation responses should usually "
            "set this to False."
        )
    )

@tool(
    args_schema=RecommenderInput,
    description="""
Use this tool when:
- enough information exists
- no more clarification is required
- the best assessments have been identified

Select ONLY the best matching assessments.

Do NOT recommend every retrieved document.
If there are two versions of same assessment name like 1.0 and 2.0 then only select 2.0 version.
Do NOT mention any entity id in recommendation summary.

Use ONLY exact assessment entity ids from retrieved docs.
"""
)
def recommender(
    recommendation_summary: str,
    recommended_assessments_ids: list[str],
    end_of_conversation: bool
):
    recommendations = []
    for entity_id in recommended_assessments_ids:
        cat = catalog_by_id[entity_id]
        name = cat['name']
        url = cat['link']
        test_type = ",".join([keys_mapping[key] for key in cat['keys']])
        recommendations.append({
            "name": name,
            "url": url,
            "test_type": test_type
        })  
        
    return {
        "type": "recommendation",
        "message": recommendation_summary,
        "recommendations": recommendations,
        "end_of_conversation": end_of_conversation
    }