Function classify_problem

optillm/plugins/spl/generation.py:25–90

Use the LLM to classify the problem type, ensuring the result is one of the valid types. Args: content: The query/problem to classify client: LLM client for making API calls model: Model identifier Returns: str: The problem type classification (

(content: str, client, model: str)

logger = logging.getLogger(__name__)

def classify_problem(content: str, client, model: str) -> str:
    """
    Use the LLM to classify the problem type, ensuring the result is one of the valid types.
    
    Args:
        content: The query/problem to classify
        client: LLM client for making API calls
        model: Model identifier
    
    Returns:
        str: The problem type classification (always a valid type)
    """
    # Format problem types as a comma-separated list
    problem_types_str = ", ".join(VALID_PROBLEM_TYPES[:-1])  # Exclude the general_problem fallback
    
    try:
        messages = [
            {
                "role": "system", 
                "content": PROBLEM_CLASSIFICATION_PROMPT.format(problem_types=problem_types_str)
            },
            {
                "role": "user", 
                "content": (
                    f"Classify the following problem into ONE of these types: {problem_types_str}\n\n"
                    f"Problem: {content}"
                )
            }
        ]
        
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0.1,  # Low temperature for more deterministic output
            max_tokens=DEFAULT_MAX_TOKENS  # Increased token limit for reasoning LLMs
        )
        
        # Extract final response and thinking content
        raw_response = response.choices[0].message.content
        final_response, thinking = extract_thinking(raw_response)
        
        # Clean and normalize the response
        final_response = final_response.strip().lower()
        
        logger.debug(f"Problem classification - raw response: '{raw_response}'")
        logger.debug(f"Problem classification - final response after removing thinking: '{final_response}'")
        
        # Find the exact match from our list of valid types
        for valid_type in VALID_PROBLEM_TYPES:
            if valid_type.lower() == final_response:
                logger.info(f"Classified problem as '{valid_type}' (exact match)")
                return valid_type
        
        # If no exact match, look for partial matches
        for valid_type in VALID_PROBLEM_TYPES:
            if valid_type.lower() in final_response:
                logger.info(f"Classified problem as '{valid_type}' (partial match from '{final_response}')")
                return valid_type