From fb29548333515a94533621fb9b8d71e6e21addb1 Mon Sep 17 00:00:00 2001
From: "Rossdan Craig rossdan@lastmileai.dev" <>
Date: Mon, 18 Mar 2024 13:09:53 -0400
Subject: [PATCH] Small minor fixes to functions in Prompt functions

Just a couple of minor fixes.

No tests exist for this so just gotta trust me I guess
---
 python/src/aiconfig/schema.py  | 10 +++++++---
 python/tests/test_dataclass.py |  3 +--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/python/src/aiconfig/schema.py b/python/src/aiconfig/schema.py
index 0a7a465b7..5f1bc339c 100644
--- a/python/src/aiconfig/schema.py
+++ b/python/src/aiconfig/schema.py
@@ -196,14 +196,18 @@ class Config:
         extra = "allow"
 
     def add_output(self, output: Output):
+        """
+        Add the output to the prompt's output list
+        """
+        if self.outputs is None:
+            self.outputs = [output]
         self.outputs.append(output)
 
-    def get_raw_prompt_from_config(self):
+    def get_raw_prompt_from_config(self) -> str:
         """Gets raw prompt from config"""
         if isinstance(self.input, str):
             return self.input
-        else:
-            return self.input.prompt
+        return str(self.input.data)
 
 
 class SchemaVersion(BaseModel):
diff --git a/python/tests/test_dataclass.py b/python/tests/test_dataclass.py
index cbbb0169f..8d6f0e3c8 100644
--- a/python/tests/test_dataclass.py
+++ b/python/tests/test_dataclass.py
@@ -18,8 +18,7 @@ def test_get_raw_prompt_from_prompt_object_prompt_input():
     prompt = Prompt(
         name="test",
         input=PromptInput(
-            prompt="This is a basic prompt",
-            data={},
+            data="This is a basic prompt",
         ),
         metadata=PromptMetadata(
             model=ModelMetadata(name="FakeMetadata", settings={}),