Fix metrics for newer models and with new trace types

src/generative_ai_toolkit/metrics/modules/conciseness.py

@@ -68,7 +68,7 @@ def evaluate_conversation(self, conversation_traces, **kwargs):
                                 Example output:
                                 {{ "score": 9, "reasoning": "The agent's responses are concise, and it does not provide superfluous examples or useless encouragements."}}
 
-                                Only return the valid JSON object.
+                                Only return the valid JSON object. Do not wrap it in markdown code blocks or any other formatting.
                                 """
                             )
                             .format(conversation=json.dumps(user_conversation))

src/generative_ai_toolkit/metrics/modules/conversation.py

@@ -83,7 +83,7 @@ def evaluate_conversation(self, conversation_traces, **kwargs):
                                 Example output:
                                 {{ "score": 9, "reasoning": "The agent succeeded in helping the user as expected"}}
 
-                                Only return the JSON object.
+                                Only return the JSON object. Do not wrap it in markdown code blocks or any other formatting.
                                 """
                             )
                             .format(

src/generative_ai_toolkit/metrics/modules/latency.py

@@ -42,6 +42,8 @@ def evaluate_trace(self, trace, **kwargs):
             dimensions.append({"ConversationHistory": "add-message"})
         elif trace_type in {"converse", "converse-stream"}:
             dimensions.append({"Converse": trace_type})
+        elif trace_type == "cycle":
+            dimensions.append({"Cycle": "agent-cycle"})
         else:
             logger.warn("Unknown trace type", trace_type=trace_type)
 

src/generative_ai_toolkit/utils/llm_response.py

@@ -30,6 +30,22 @@ def get_text(response: "ConverseResponseTypeDef"):
 
 def json_parse(response: "ConverseResponseTypeDef"):
     text = get_text(response).strip()
+
+    # Handle markdown code blocks
+    if text.startswith("```json"):
+        # Find the closing ``` and extract content between
+        end_marker = text.rfind("```")
+        if end_marker > 7:  # Make sure we found a closing marker after ```json
+            text = text[7:end_marker].strip()  # Remove ```json and closing ```
+    elif text.startswith("```") and text.count("```") >= 2:
+        # Handle generic code blocks that might contain JSON
+        first_newline = text.find('\n')
+        if first_newline != -1:
+            # Skip the opening ``` line
+            end_marker = text.rfind("```")
+            if end_marker > first_newline:
+                text = text[first_newline+1:end_marker].strip()
+
     try:
         return json.loads(text.replace("\n", " "))
     except json.decoder.JSONDecodeError as e: