app/categorization/categorizer.py

import re
import ast
import json
import logging
import os
from typing import Any, List, Tuple, Optional, Dict, Union

import numpy as np
import pandas as pd
import asyncio
from rapidfuzz import process
from tenacity import retry, wait_random_exponential, stop_after_attempt
from pydantic import ValidationError

from fastapi import HTTPException

from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from langchain.prompts import PromptTemplate
from app.categorization.template import CATEGORY_TEMPLATE
from app.categorization.config import CATEGORY_REFERENCE_OUTPUT_FILE, TX_PER_LLM_RUN


def fuzzy_match_list_categorizer(
    description: str,
    descriptions: np.ndarray,
    description_category_pairs: pd.DataFrame,
    threshold: int = 90,
) -> Optional[str]:
    """Find the most similar transaction description and return its category.

    This function uses fuzzy string matching to compare the input description
    against a list of known descriptions. If a sufficient match is found,
    the function returns the category associated with the matched description.

    Args:
        description (str): The transaction description to categorize.
        descriptions (np.ndarray): Known descriptions to compare against.
        description_category_pairs (pd.DataFrame): DataFrame mapping descriptions to categories.
        threshold (int): Minimum similarity score to consider a match.

    Returns:
        str or None: Category of the matched description, or None if no match found.
    """
    result = {"description": description, "descriptions": descriptions, "description_category_pairs": description_category_pairs, "match_results": "", "error": ""}
    try:
        if description_category_pairs.empty:
            return None

         # Fuzzy-match this description against the reference descriptions
        description_stripped = description.strip()
        descriptions_stripped = [s.strip() for s in descriptions]

        match_results = process.extractOne(
            description_stripped, descriptions_stripped, score_cutoff=threshold)
        
        result["match_results"] = match_results
        # print(f"\pairs_index: {description_category_pairs.columns}\n nmatch_results0: {match_results[0]}")

        # If a match is found, return the category of the matched description
        if match_results and (description_category_pairs.index == match_results[0]).any:
            return description_category_pairs.at[match_results[0], 'category']

        return None
        
    except Exception as e:
        result["error"] = str(e)
        raise HTTPException(status_code = 500, detail=f"fuzzy_match_list_categorizer result: {result}")


async def llm_list_categorizer(df: pd.DataFrame) -> pd.DataFrame:
    """Categorize a list of transactions using a language model.

    This function uses a Language Model (LLM) to categorize a list of transaction descriptions.
    It splits the input DataFrame into chunks and processes each chunk asynchronously to improve performance.

    Args:
        df (pd.DataFrame): DataFrame containing the transaction descriptions to categorize.

    Returns:
        pd.DataFrame: DataFrame mapping transaction descriptions to their inferred categories.
    """
    
    result = {"df": "", "results": "", "valid results": "", "valid outputs": "", "error": ""}

    try:
        # Initialize language model and prompt
        openai_api_key = os.environ['OPENAI_API_KEY']
        llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0125",
                        api_key=openai_api_key)
        prompt = PromptTemplate.from_template(template=CATEGORY_TEMPLATE)
        chain = prompt | llm
        # chain = LLMChain(llm=llm, prompt=prompt)

        # Iterate over the DataFrame in batches of TX_PER_LLM_RUN transactions
        tasks = [llm_sublist_categorizer(df.attrs['file_name'], chain=chain, tx_descriptions="\n".join(chunk['name_description']).strip())
                for chunk in np.array_split(df, df.shape[0] // TX_PER_LLM_RUN + 1)]

        # Gather results and extract (valid) outputs
        # The results variable is a list of 'results', each 'result' being the output of a single LLM run
        results = await asyncio.gather(*tasks)
        result["results"] = results

        # Extract valid results (each valid result is a list of description-category pairs)
        valid_results = [result['output'] for result in results if result['valid']]
        result["valid results"] = valid_results

        # Flatten the list of valid results to obtain a single list of description-category pairs
        valid_outputs = [
            output for valid_result in valid_results for output in valid_result]
        result["valid outputs"] = valid_outputs

        # Return a DataFrame with the valid outputs
        return pd.DataFrame(valid_outputs, columns=['name_description', 'category'])
    
    except Exception as e:
        result["error"] = str(e)
        raise HTTPException(status_code = 500, detail=f"llm_list_categorizer result: {result}")



@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
async def llm_sublist_categorizer(
    file_name: str,
    chain: LLMChain,
    tx_descriptions: str,
) -> Dict[str, Union[bool, List[Tuple[str, str]]]]:
    """Categorize a batch of transactions using a language model.

    This function takes a batch of transaction descriptions and passes them to a language model
    for categorization. The function retries on failure, with an exponential backoff.

    Args:
        file_name (str): Name of the file the transaction descriptions were extracted from.
        chain (LLMChain): Language model chain to use for categorization.
        tx_descriptions (str): Concatenated transaction descriptions to categorize.

    Returns:
        dict: Dictionary containing a 'valid' flag and a list of categorized descriptions.
    """

    raw_result = await chain.ainvoke(tx_descriptions)
    # raw_result = await chain.arun(input_data=tx_descriptions)

    logger = logging.getLogger(__name__)
    result = {'valid': True, 'output': []}
    try:
        # Create a pattern to match a list Description-Category pairs (List[Tuple[str, str]])
        pattern = r"\['([^']+)', '([^']+)'\]"

        # Use it to extract all the correctly formatted pairs from the raw result
        matches = re.findall(pattern, raw_result.content.replace("\\'", "'"))

        # print(f"llm_sublist_categorizer type_matches: {type(matches)} match: {matches}\n")

        # Loop over the matches, and try to parse them to ensure the content is valid
        valid_outputs = []
        for match in matches:
            try:
                parsed_pair = ast.literal_eval(str(list(match)))
                # print(f"llm_sublist_categorizer \n parsed_pair: {parsed_pair} \n match: {match}\n list(match): {list(match)}\n str(list(match): {str(list(match))} \n\n")
                valid_outputs.append(parsed_pair)
            except Exception as e:
                logger.log(logging.ERROR,
                           f"Parsing Error: {e}\nMatch: {match}\n")
                result['valid'] = False

        # print(f"llm_sublist_categorizer valid_outputs \n{valid_outputs}\n")

        result['output'] = valid_outputs

    except Exception as e:
        # logging.log(
        #     logging.ERROR, f"| File: {file_name} | Unexpected Error: {e}\nRaw Result: {raw_result}")
        result['valid'] = False
        raise HTTPException(status_code = 500, detail=f"llm_sublist_categorizer result: {result}")

    return result