Updated the react agent with tool in node

my_agent/utils/nodes.py

@@ -1,13 +1,17 @@
 import pandas as pd
 import ast
 from .state import State
-from .tools import StoryFormatter, BrainstromTopicFormatter
+from .tools import StoryFormatter, BrainstromTopicFormatter , retrieve_tool
 from langchain_core.messages import SystemMessage
 from .models_loader import llm , ST
 from .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 , final_story_prompt
+from langgraph.prebuilt import create_react_agent
+from pydantic import BaseModel , Field
+from langchain_core.tools import tool
+
 
 
 def caption_image(state: State) -> State:
@@ -79,36 +83,27 @@ def retrieve(state: State) -> State:
   return state
 
 def generate_story(state:State)-> State:
-  print('The state retrieval is:',state.retrievals)
-  retrieval_list= state.retrievals[-1]
-  agentic_stories = []
-
-  for item in retrieval_list:
-      print('item:', item[-1].values())
-
-      agentic_stories.extend(item[-1].values())  # Add all stories to the list
-
-  retrieval = " ".join(agentic_stories)
-
-  if len(state.preferred_topics)==0:
-    template = initial_story_prompt(retrieval , state)
-  else:
-    template = refined_story_prompt(retrieval , state)
+    tools=[retrieve_tool]
+
+    react_agent=create_react_agent(
+        model=llm.bind_tools(tools),
+        tools=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)]
-  response = llm.bind_tools([StoryFormatter]).invoke(messages)
-  print('The response is:',response)
-  if hasattr(response, 'tool_calls') and response.tool_calls:
-      response = response.tool_calls[0]['args']
-  elif hasattr(response, 'content'):
-      response = response.content
-  else:
-    response = "No response"
-  state.stories.append(response)
-  # return State(messages="Story generated", topic=state.topic,stories=state.stories)
-  return state
+    messages = [SystemMessage(content=template)]
+
+    response = react_agent.invoke({'messages':messages})
+    response = response['messages'][-1]
+    state.stories.append(response)
+    # return State(messages="Story generated", topic=state.topic,stories=state.stories)
+    return state