From 1ee0eae51231baf245904e40331565a10aac07d5 Mon Sep 17 00:00:00 2001
From: brandon <brandon@groundlight.ai>
Date: Tue, 5 Sep 2023 16:47:34 -0700
Subject: [PATCH] Improved type annotation assignment logic, fixed test

---
 src/groundlight/cli.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/groundlight/cli.py b/src/groundlight/cli.py
index 880260f4..5d80a2a3 100644
--- a/src/groundlight/cli.py
+++ b/src/groundlight/cli.py
@@ -22,8 +22,9 @@ def wrapper(*args, **kwargs):
     for name, annotation in method.__annotations__.items():
         if get_origin(annotation) is Union:
             if str in annotation.__args__:
-                new_annotation = str
-        wrapper.__annotations__[name] = new_annotation
+                wrapper.__annotations__[name] = str
+            else:
+                wrapper.__annotations__[name] = annotation
 
     return wrapper