Merge pull request #79 from ServiceNow/fix_test_step_kind

Fix, update tape improver test

examples/tape_improver/tape_improver.py

@@ -145,7 +145,7 @@ def improver_tape_view(tape: Tape) -> str:
         data[-1]["index"] = index
         if isinstance(step, AgentStep):
             data[-1]["metadata"] = {"agent": step.metadata.agent}
-    return json.dumps(data, indent=2)
+    return json.dumps(data, indent=2, sort_keys=True)
 
 
 ### Agents ###

requirements.txt

@@ -10,7 +10,7 @@ fastapi==0.112.4
 langchain-community==0.0.38
 langchain-core==0.1.52
 Levenshtein==0.25.1
-litellm==1.37.9
+litellm==1.51.2
 matplotlib==3.9.0
 pathvalidate==3.2.0
 podman==5.0

tapeagents/llms.py

@@ -7,6 +7,7 @@
 import os
 import time
 from abc import ABC, abstractmethod
+from itertools import zip_longest
 from typing import Any, Callable, Generator
 
 import litellm
@@ -496,7 +497,9 @@ def _implementation():
                 known_prompts = list(self.outputs.keys())
                 closest, score = closest_prompt(prompt_key, known_prompts)
                 if score >= 0.7:
-                    logger.warning(f"Closest prompt score {score:.3f}:\n{diff_strings(closest, prompt_key)}")
+                    logger.warning(f"Closest prompt score {score:.3f}")
+                    for i, (a, b) in enumerate(zip_longest(prompt.messages, json.loads(closest), fillvalue={})):
+                        logger.warning(f"STEP{i}: {diff_strings(a.get('content', str(a)), b.get('content', str(b)))}\n")
                 raise FatalError("prompt not found")
             yield LLMEvent(output=LLMOutput(content=output))
 

tests/make_test_data.py

@@ -9,8 +9,8 @@
 
 import tapeagents.observe
 from examples import delegate_stack
-from examples.tape_improver import tape_improver
 from examples.data_science import data_science
+from examples.tape_improver import tape_improver
 
 
 @contextlib.contextmanager
@@ -80,6 +80,4 @@ def run_in_tmp_dir_to_make_test_data(test_name: str, keep_llm_cache=False):
             with run_in_tmp_dir_to_make_test_data("data_science"):
                 data_science.main(studio=False)
         case _:
-            raise Exception(
-                "Usage: python -m examples.make_test_data [delegate_stack | intro_notebook | tape_improver]"
-            )
+            raise Exception("Usage: python -m tests.make_test_data [delegate_stack | intro_notebook | tape_improver]")