Updated the tool in final story generation

my_agent/utils/utils.py

@@ -1,6 +1,6 @@
 
 from langchain_core.messages import SystemMessage
-from .tools import StoryFormatter 
+from .tools import StoryFormatter , retrieve_tool
 from .models_loader import llm
 import base64
 from PIL import Image
@@ -9,21 +9,23 @@ from fastapi import UploadFile
 from huggingface_hub import InferenceClient
 from .prompts import story_to_prompt , final_story_prompt
 import os
+from langgraph.prebuilt import create_react_agent
+
+
+
 
 def generate_final_story(final_state):
     if len(final_state['preferred_topics'])>0:
         template = final_story_prompt(final_state)
         messages = [SystemMessage(content=template)]
-        response = llm.bind_tools([StoryFormatter]).invoke(messages)
-        print('The final 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.final_story.append(response)
-        # state.stories.append(response)
+
+        tools = [retrieve_tool]
+        react_agent=create_react_agent(
+        model=llm.bind_tools(tools),
+        tools=tools)
+
+        response = react_agent.invoke({'messages':messages})
+        response = response['messages'][-1].content
         return response
     
     else: