Handled the captioning node

my_agent/utils/nodes.py

@@ -11,32 +11,34 @@ from .prompts import image_captioning_prompt , initial_story_prompt , refined_st
 
 
 def caption_image(state: State) -> State:
-    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},
-                      {
-                          "type": "image_url",
-                          "image_url": {
-                              "url": f"data:image/jpeg;base64,{state.images[-1]}",
+    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},
+                          {
+                              "type": "image_url",
+                              "image_url": {
+                                  "url": f"data:image/jpeg;base64,{state.images[-1]}",
+                              },
                           },
-                      },
-                  ],
-              }
-          ],
-          model="meta-llama/llama-4-scout-17b-16e-instruct",
-      )
-      response=chat_completion.choices[0].message.content
-      state.image_captions.append(response)
-      return state
+                      ],
+                  }
+              ],
+              model="meta-llama/llama-4-scout-17b-16e-instruct",
+          )
+          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
     
@@ -173,24 +175,7 @@ def route_after_selection(state:State):
   elif len(state.latest_preferred_topics)>0:
     return True
 
-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)
-        return response
-    
-    else:
-        return final_state['stories'][-1]
+