Repository: lastmile-ai/aiconfig (language: python)

Issue:
Do not overwrite outputs in `add_output` if value passed is None
In the Python SDK, the add_output method currently contains a critical flaw. 

As observed at [this link](https://github.com/lastmile-ai/aiconfig/blob/62c8e7d664fb9fa5bc52168937406d738ece3b37/python/src/aiconfig/schema.py#L663), the method erroneously overwrites the prompt's output, even in scenarios where the Output parameter is None. This behavior contradicts the method's design, as Output is not intended to be an optional attribute. 

To align with the intended functionality, the method should be modified to raise an exception when Output is None, ensuring robust and predictable behavior.


Discussion / hints:
I don't fully understand the issue. Can you describe what's expected vs. actual behavior with an example?
Looks like added code changed the expected path of the link. Updated

> I don't fully understand the issue. Can you describe what's expected vs. actual behavior with an example?

The issue is that add_output incorrectly adds None to the output list when output = None is passed. The expected behavior is for the method to reject None as an invalid Output, not add it (or alternatively raise an exception). The current implementation doesn't align with this expectation.
Here's a link to my original discussion with @Ankush-lastmile about this. 
https://github.com/lastmile-ai/aiconfig/issues/153#issuecomment-1859596880

Fix patch:
diff --git a/python/src/aiconfig/schema.py b/python/src/aiconfig/schema.py
index df23faa43..aeca7c2ba 100644
--- a/python/src/aiconfig/schema.py
+++ b/python/src/aiconfig/schema.py
@@ -681,7 +681,7 @@ def delete_metadata(self, key: str, prompt_name: Optional[str] = None):
 
     def add_output(self, prompt_name: str, output: Output, overwrite: bool = False):
         """
-        Add an output to the [rompt with the given name in the AIConfig
+        Add an output to the prompt with the given name in the AIConfig
 
         Args:
             prompt_name (str): The name of the prompt to add the output to.
@@ -690,8 +690,10 @@ def add_output(self, prompt_name: str, output: Output, overwrite: bool = False):
         """
         prompt = self.get_prompt(prompt_name)
         if not prompt:
-            raise IndexError(f"Cannot out output. Prompt '{prompt_name}' not found in config.")
-        if overwrite or not output:
+            raise IndexError(f"Cannot add output. Prompt '{prompt_name}' not found in config.")
+        if not output:
+            raise ValueError(f"Cannot add output to prompt '{prompt_name}'. Output is not defined.")
+        if overwrite:
             prompt.outputs = [output]
         else:
             prompt.outputs.append(output)


Test patch:
diff --git a/python/tests/test_programmatically_create_an_AIConfig.py b/python/tests/test_programmatically_create_an_AIConfig.py
index 43bb8624b..4aa04a4b9 100644
--- a/python/tests/test_programmatically_create_an_AIConfig.py
+++ b/python/tests/test_programmatically_create_an_AIConfig.py
@@ -536,12 +536,12 @@ def test_set_and_delete_metadata_ai_config(ai_config_runtime: AIConfigRuntime):
 
 def test_set_and_delete_metadata_ai_config_prompt(ai_config_runtime: AIConfigRuntime):
     """Test deleting a non-existent metadata key at the AIConfig level."""
-    prompt1 = Prompt(
+    prompt = Prompt(
         name="GreetingPrompt",
         input="Hello, how are you?",
         metadata=PromptMetadata(model="fakemodel"),
     )
-    ai_config_runtime.add_prompt(prompt1.name, prompt1)
+    ai_config_runtime.add_prompt(prompt.name, prompt)
     ai_config_runtime.set_metadata("testkey", "testvalue", "GreetingPrompt")
 
     assert (
@@ -558,17 +558,19 @@ def test_set_and_delete_metadata_ai_config_prompt(ai_config_runtime: AIConfigRun
 
 def test_add_output_existing_prompt_no_overwrite(ai_config_runtime: AIConfigRuntime):
     """Test adding an output to an existing prompt without overwriting."""
-    prompt1 = Prompt(
+    prompt = Prompt(
         name="GreetingPrompt",
         input="Hello, how are you?",
         metadata=PromptMetadata(model="fakemodel"),
     )
-    ai_config_runtime.add_prompt(prompt1.name, prompt1)
+    ai_config_runtime.add_prompt(prompt.name, prompt)
     test_result = ExecuteResult(
         output_type="execute_result",
-        execution_count=0.0,
-        data={"role": "assistant", "content": "test output"},
-        metadata={"finish_reason": "stop"},
+        execution_count=0,
+        data="test output",
+        metadata={
+            "raw_response": {"role": "assistant", "content": "test output"}
+        },
     )
     ai_config_runtime.add_output("GreetingPrompt", test_result)
 
@@ -576,18 +578,72 @@ def test_add_output_existing_prompt_no_overwrite(ai_config_runtime: AIConfigRunt
 
     test_result2 = ExecuteResult(
         output_type="execute_result",
-        execution_count=0.0,
-        data={"role": "assistant", "content": "test output"},
-        metadata={"finish_reason": "stop"},
-    )
+        execution_count=0,
+        data="test output",
+        metadata={
+            "raw_response": {"role": "assistant", "content": "test output for second time"}
+        },
+    ) 
+    
     ai_config_runtime.add_output("GreetingPrompt", test_result2)
     assert ai_config_runtime.get_latest_output("GreetingPrompt") == test_result2
 
     ai_config_runtime.delete_output("GreetingPrompt")
-
     assert ai_config_runtime.get_latest_output("GreetingPrompt") == None
 
 
+def test_add_output_existing_prompt_overwrite(ai_config_runtime: AIConfigRuntime):
+    """Test adding an output to an existing prompt with overwriting."""
+    original_output = ExecuteResult(
+        output_type="execute_result",
+        execution_count=0,
+        data="original output",
+        metadata={
+            "raw_response": {"role": "assistant", "content": "original output"}
+        },
+    )
+    prompt = Prompt(
+        name="GreetingPrompt",
+        input="Hello, how are you?",
+        metadata=PromptMetadata(model="fakemodel"),
+        outputs=[original_output],
+    )
+    ai_config_runtime.add_prompt(prompt.name, prompt)
+    # check that the original_output is there
+    assert ai_config_runtime.get_latest_output("GreetingPrompt") == original_output
+    expected_output = ExecuteResult(
+        output_type="execute_result",
+        execution_count=0,
+        data="original output",
+        metadata={
+            "raw_response": {"role": "assistant", "content": "original output"}
+        },
+    )
+    # overwrite the original_output
+    ai_config_runtime.add_output("GreetingPrompt", expected_output, True)
+    assert ai_config_runtime.get_latest_output("GreetingPrompt") == expected_output
+
+def test_add_undefined_output_to_prompt(ai_config_runtime: AIConfigRuntime):
+    """Test for adding an undefined output to a prompt with/without overwriting. Should result in an error."""
+    prompt = Prompt(
+        name="GreetingPrompt",
+        input="Hello, how are you?",
+        metadata=PromptMetadata(model="fakemodel"),
+    )
+    ai_config_runtime.add_prompt(prompt.name, prompt)
+    # Case 1: No output, overwrite param not defined
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot add output to prompt 'GreetingPrompt'. Output is not defined.",
+    ):
+        ai_config_runtime.add_output("GreetingPrompt", None)
+    # Case 2: No output, overwrite param set to True
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot add output to prompt 'GreetingPrompt'. Output is not defined.",
+    ):
+        ai_config_runtime.add_output("GreetingPrompt", None, True)
+    
 def test_extract_override_settings(ai_config_runtime: AIConfigRuntime):
     initial_settings = {"topP": 0.9}