src/genai/brainstroming_agent/utils/nodes.py

import pandas as pd
import ast
from .state import State
from .tools import retrieve_tool
from langchain_core.messages import SystemMessage ,HumanMessage, ToolMessage
from src.genai.utils.models_loader import llm , ST
from src.genai.utils.data_loader import load_influencer_data
from groq import Groq
import os
from .prompts import image_captioning_prompt , initial_story_prompt , refined_story_prompt , brainstroming_prompt 
from langgraph.prebuilt import create_react_agent
from .state import BrainstromTopicFormatter



def caption_image(state: State) -> State:
    if len(state.images)>0:
      if state.images[-1]!=None:
          print('Captioning image')
          client = Groq(api_key=os.environ.get('GROQ_API_KEY'))

          chat_completion = client.chat.completions.create(
              messages=[
                  {
                      "role": "user",
                      "content": [
                          {"type": "text", "text": image_captioning_prompt(state)},
                          {
                              "type": "image_url",
                              "image_url": {
                                  "url": f"data:image/jpeg;base64,{state.images[-1]}",
                              },
                          },
                      ],
                  }
              ],
              model="meta-llama/llama-4-scout-17b-16e-instruct",
              max_completion_tokens=50,
              temperature = 1
          )
          response=chat_completion.choices[0].message.content
          state.image_captions.append(response)
          return state
    
    else:
       state.images.append(None)
       state.image_captions.append(None)
       return state
    


def retrieve(state: State) -> State:
  print('Moving to retrieval process')
  retrievals=[]
  query_prompt = 'Represent this sentence for searching relevant passages: '
  if len(state.latest_preferred_topics)==0:
      for idea in state.idea:  
          print('The idea for retrieval:', idea)
          result = retrieve_tool(idea+query_prompt)
          retrievals.append(result)
          print('Retrieval process completed......')
      state.retrievals.append(retrievals)

  if len (state.latest_preferred_topics)>0:
      print('The preferred_topics are:',state.latest_preferred_topics)
      state.preferred_topics.append(state.latest_preferred_topics)
      for idea in state.preferred_topics[-1]:  
          result = retrieve_tool(idea+query_prompt)
          retrievals.append(result)
          print('Retrieval process completed for preferred_topics......')
      state.latest_preferred_topics=[]
      state.retrievals.append(retrievals)
  return state

def generate_story(state:State)-> State:
    react_agent=create_react_agent(
        model=llm,
        tools=[]
        
    )
    if len(state.preferred_topics)==0:
        template = initial_story_prompt(state)
    else:
        template = refined_story_prompt(state)

# and {state.image_captions[-1]}

    messages = [SystemMessage(content=template),
                HumanMessage(content=f'''The idea of the video is:\n{state.idea}\n'''),
                ToolMessage(content=f'''The business details is:\n{state.business_details}\n
                            The retrieved data of influencers is:\n{state.retrievals[-1]}\n
                            The information from the image is:\n{state.image_captions[-1]} ''', tool_call_id='generate_story_tool')]
    
    print('Messages:',messages)

    response = react_agent.invoke({'messages':messages})
    response = response['messages'][-1].content
    print('The genrated story: ', response)
    state.stories.append(response)
    return state



def generate_brainstroming(state:State)-> State:

  template= brainstroming_prompt(state)

  messages = [SystemMessage(content=template),
              HumanMessage(content=f'''Here is the story to you for brainstorming:\n{state.stories[-1]}'''),
              ToolMessage(content=f'''The details of business is:\n{state.business_details}\n''', tool_call_id="brainstorm_tool")]
  print('Message for brainstorming:',messages)
  response = llm.with_structured_output(BrainstromTopicFormatter).invoke(messages)
  response = response.model_dump()
  state.brainstroming_topics.append(response)
  print('The brainstroming topics are:',state.brainstroming_topics)
  # return State(messages="Story generated",topic=state.topic,brainstroming_topics=state.brainstroming_topics)
  return state