src/genai/context_analysis_agent/utils/utils.py

import pandas as pd
from src.genai.utils.data_loader import caption_df
from src.genai.utils.models_loader import llm_gpt
from .prompts import details_extract_prompt
from langchain_core.messages import SystemMessage, HumanMessage
from .state import DetailsFormatter
from langsmith import traceable

class DetailsExtractorNode:
    def __init__(self, interactions):
        self.llm = llm_gpt
        self.interactions = interactions
    
    @traceable(name="details extraction")
    def run(self):
        template = details_extract_prompt()
        messages = [SystemMessage(content=template), HumanMessage(content=str(self.interactions))]
        response=llm_gpt.with_structured_output(DetailsFormatter).invoke(messages)
        return response.model_dump()


class SaveToDB:
    def __init__(self, caption_df):
        self.df = caption_df.drop(columns=['embeddings'], errors='ignore')

    def _prepare_values(self, business_details):
        """Extract lowercase string values from business_details dict."""
        all_values = set()
        for v in business_details.values():
            if isinstance(v, str):
                all_values.add(v.lower())
            elif isinstance(v, list):
                all_values.update(map(str.lower, map(str, v)))
        return all_values

    def _row_matches(self, row, all_values):
        """Check if any value in all_values exists in the row."""
        return any(
            str(cell).lower().find(val) != -1
            for cell in row
            for val in all_values
        )

    def save_to_csv(self, business_details, output_file='extracted_data.csv'):
        """Filter dataframe rows based on business_details and save to CSV."""
        all_values = self._prepare_values(business_details)
        matched_df = self.df[self.df.apply(self._row_matches, axis=1, args=(all_values,))]
        matched_df.to_csv(output_file, index=False)