Deal with summarization Error for images (#910)

Author: tawsifkamal

src/codegen/extensions/langchain/graph.py

@@ -155,15 +155,27 @@ def format_header(header_type: str) -> str:
 
         # Format messages with appropriate headers
         formatted_messages = []
-        for msg in to_summarize:  # No need for slice when iterating full list
+        image_urls = []  # Track image URLs for the summary prompt
+
+        for msg in to_summarize:
             if isinstance(msg, HumanMessage):
-                formatted_messages.append(format_header("human") + msg.content)
+                # Now we know content is always a list
+                for item in msg.content:
+                    if item.get("type") == "text":
+                        text_content = item.get("text", "")
+                        if text_content:
+                            formatted_messages.append(format_header("human") + text_content)
+                    elif item.get("type") == "image_url":
+                        image_url = item.get("image_url", {}).get("url")
+                        if image_url:
+                            # We are not including any string data in the summary for image. The image will be present itself!
+                            image_urls.append({"type": "image_url", "image_url": {"url": image_url}})
             elif isinstance(msg, AIMessage):
                 # Check for summary message using additional_kwargs
                 if msg.additional_kwargs.get("is_summary"):
                     formatted_messages.append(format_header("summary") + msg.content)
                 elif isinstance(msg.content, list) and len(msg.content) > 0 and isinstance(msg.content[0], dict):
-                    for item in msg.content:  # No need for slice when iterating full list
+                    for item in msg.content:
                         if item.get("type") == "text":
                             formatted_messages.append(format_header("ai") + item["text"])
                         elif item.get("type") == "tool_use":
@@ -173,16 +185,25 @@ def format_header(header_type: str) -> str:
             elif isinstance(msg, ToolMessage):
                 formatted_messages.append(format_header("tool_response") + msg.content)
 
-        conversation = "\n".join(formatted_messages)  # No need for slice when joining full list
+        conversation = "\n".join(formatted_messages)
 
         summary_llm = LLM(
             model_provider="anthropic",
             model_name="claude-3-5-sonnet-latest",
             temperature=0.3,
         )
 
-        chain = ChatPromptTemplate.from_template(SUMMARIZE_CONVERSATION_PROMPT) | summary_llm
-        new_summary = chain.invoke({"conversation": conversation}).content
+        # Choose template based on whether we have images
+        summarizer_content = [{"type": "text", "text": SUMMARIZE_CONVERSATION_PROMPT}]
+        for image_url in image_urls:
+            summarizer_content.append(image_url)
+
+        chain = ChatPromptTemplate([("human", summarizer_content)]) | summary_llm
+        new_summary = chain.invoke(
+            {
+                "conversation": conversation,
+            }
+        ).content
 
         return {"messages": {"type": "summarize", "summary": new_summary, "tail": tail, "head": head}}
 
@@ -199,7 +220,7 @@ def should_continue(self, state: GraphState) -> Literal["tools", "summarize_conv
             return "summarize_conversation"
 
         # Summarize if the last message exceeds the max input tokens of the model - 10000 tokens
-        elif isinstance(last_message, AIMessage) and not just_summarized and curr_input_tokens > (max_input_tokens - 10000):
+        elif isinstance(last_message, AIMessage) and not just_summarized and curr_input_tokens > (max_input_tokens - 30000):
             return "summarize_conversation"
 
         elif hasattr(last_message, "tool_calls") and last_message.tool_calls: