src/genai/analytics_chatbot/utils/nodes.py

import requests
from langchain_core.messages import SystemMessage , HumanMessage , FunctionMessage
from .state import State
from .tools import RetrieverBackup
from .schemas import ParameterFormatter, EndpointFormatter
from .prompts import query_check_prompt, fetch_last_message_prompt , fetch_parameters_prompt, fetch_endpoint_prompt
from .utils import  process_query, get_endpoint_info
from src.genai.utils.models_loader import llm_gpt
import numpy as np
from src.genai.utils.data_loader import api_knowledge_df, api_index, caption_df , caption_index
from src.genai.utils.models_loader import embedding_model
from ..handlers import (
    compare,
    get_posting_time,
    get_peak_comment_hour,
    get_emoji_count,
    get_comment_quality,
    get_bot_and_diversity,
)


class FetchLastMessage:
    def __init__(self):
        self.llm = llm_gpt
    
    def run (self, state:State):
        print('Message:',state['messages'])
        template = fetch_last_message_prompt
        messages=[SystemMessage(content=template)]+state['messages']
        result = self.llm.invoke(messages)
        print('Latest Message:', process_query(result.content))
        if len(state['messages'])>11:
            state["messages"] = state["messages"][-9:]
        return {
            'latest_message': process_query(result.content)
        }
        

class RetrievePossibleEndpoints:
    def __init__(self):
        self.df = api_knowledge_df
        self.index = api_index
        self.results = []

    def run(self,state:State):
        print('Gone to retrieve possible endpoints')
        query_embedding = np.array(embedding_model.embed_query(state['latest_message'])).reshape(1, -1).astype('float32')
        distances, indices = self.index.search(query_embedding,10)
        for idx in indices[0]:
            row = self.df.iloc[idx]
            print('Endpoint:',row['endpoint'])
            self.results.append(row['endpoint'])
        print('The possible endpoints are:', self.results)
        return {
            "possible_endpoints": self.results,
        }

class RetrieveExactEndpoint:
    def __init__(self):
        self.llm = llm_gpt
    
    def run(self,state:State):
        print('Gone to retrieve exact endpoint')
        messages = [SystemMessage(content=fetch_endpoint_prompt),
                    FunctionMessage(name='possible_endpoints',content=f'''The possible endpoints are: {state['possible_endpoints']}'''),
                    HumanMessage(content=f'''The user query is: {state['latest_message']}''')]
        
        result = self.llm.with_structured_output(EndpointFormatter, method='function_calling').invoke(messages)
        print('The exact endpoint is:', result.endpoint)
        endpoint_info=get_endpoint_info(result.endpoint)
        print('The endpoint info is:', endpoint_info)

        return {
            "messages":[{"role": "assistant", "content": f'''The endpoint is: {result.endpoint}'''}],
            "endpoint": result.endpoint,
            "method": endpoint_info['method'],
            "needed_parameters": endpoint_info["parameters"]
        }
        
class QueryCheckNode:
    def __init__(self):
        self.llm = llm_gpt

    def run(self, state:State):
        try:
            print('Entered to query checking')
            messages = [SystemMessage(content=query_check_prompt),
            HumanMessage(content=f'''The user query is: {state['latest_message']}''')]
            result = self.llm.invoke(messages)
            print(result.content)
            return{'query_type': result.content}
        
        except Exception as e:
            print('Error occoured:', e)
            return {'error_message': str(e)}


class FetchParametersNode:
    def __init__(self):
        self.llm = llm_gpt
        self.complex_endpoints=['/api/v1/compare/','/api/v1/engagement/posting-time-analysis','/api/v1/audience/peak-comment-hour','/api/v1/audience/emoji-count','/api/v1/audience/comment-quality']

    def run(self , state:State):
        try:
            print('Entered to fetch parameters')

            if state['endpoint'] not in self.complex_endpoints:
                template = fetch_parameters_prompt
                messages=[SystemMessage(content=template),HumanMessage(content=f'''The query is: {state['latest_message']}\n. The needed parameters: {str(state['needed_parameters'])}''')]
                # print('messages:', messages)
                result = self.llm.with_structured_output(ParameterFormatter, method='function_calling').invoke(messages)
                parameters_values = {k: (process_query(v) if isinstance(v, str) else v) for k, v in result.parameters_values.items()}
                            
                # if 'single_influencer_query' in state['query_type']:
                #     print('The parameter values:', parameters_values)
                #     return {
                #         'parameters_values':parameters_values
                #     }
                # elif 'aggregate_query' in state['query_type']:
                #     parameters_values['influencer_username'] = ['divyadhakal_','munachiya','mydarlingfood','_its.me.muskan_']
                #     print('The parameter values:', parameters_values)
                #     return{
                #         'parameters_values': parameters_values
                #     }
                print('The parameter values:', parameters_values)
                return {'parameters_values': parameters_values}
        except Exception as e:
            print('Error occoured:', e)
            return {'error_message': str(e)}

    

class FetchDataNode:
    def __init__(self):
        self.llm = llm_gpt
        self.base_url = 'https://reveltrends.vercel.app'
        self.headers = {
            "Authorization": "Bearer YOUR_API_KEY",  # replace with your API key if needed
            "Content-Type": "application/json"
            }
        
        self.endpoint_handlers = {
            '/api/v1/compare/': compare,
            '/api/v1/engagement/posting-time-analysis': get_posting_time,
            '/api/v1/audience/peak-comment-hour': get_peak_comment_hour,
            '/api/v1/audience/emoji-count': get_emoji_count,
            '/api/v1/audience/comment-quality': get_comment_quality,
            '/api/v1/audience/bot-and-diversity': get_bot_and_diversity
        }

    def run(self, state:State):
        try:
            state['query_type']='single_influencer_query'

            print('Entered to fetch data')
            url = f'''{self.base_url}{state['endpoint']}'''

            if state['endpoint'] in self.endpoint_handlers:
                print('Entered to handler.')
                handler = self.endpoint_handlers[state['endpoint']]
                response = handler(state, llm_gpt, url)
                print('Returned by handler.')
                return {'response':response.json()}


            elif 'single_influencer_query' in  state['query_type']:
                response = requests.get(url, params=state['parameters_values'],headers=self.headers)
                print('Data from api:', response)
                return {'response':response.json()}
            
            # elif 'aggregate_query' in state['query_type']:
            #     print('Entered to aggregrated query execution')
            #     print(state['parameters_values'])
            #     params = state["parameters_values"]
            #     if "influencer_username" in params and isinstance(params["influencer_username"], list):
            #         results = {}

            #         # Iterate through each influencer username
            #         for username in params["influencer_username"]:
            #             current_params = params.copy()
            #             current_params["influencer_username"] = username

            #             response = requests.get(url, params=current_params, headers=self.headers)
            #             results[username] = response.json()  # Store influencer-wise response
            #         print('Data from api:', response)
            #         return {"response": results}
            

        except Exception as e:
            print('Error occoured:', e)
            return {'error_message': str(e), 'response': 'No response'}


class BackupRetrievalNode:
    def __init__(self):
        self.llm = llm_gpt
    
    def run(self, state:State):
        retrieval=RetrieverBackup().retrieve(state['latest_message'])
        return {'backup_data': retrieval}



class BackupRoutingNode:
    def __init__(self):
        pass
    
    def run(self,state:State):
        if state.get('error_message') is not None:
            return 'execute_backup'
        else:
            return 'go_on'