diff --git a/examples/dict_type.py b/examples/dict_type.py
index 4a45c5e..31552f5 100644
--- a/examples/dict_type.py
+++ b/examples/dict_type.py
@@ -5,7 +5,12 @@
 settings.llm = "llamacpp/openchat-3.5-0106:Q3_K_M"
 
 
-def get_kwargs(user_query: str) -> dict:
+class Output(BaseModel):
+    plot: str
+    kwargs: dict
+
+
+def get_kwargs(user_query: str) -> Output:
     """
     Translate the user query to a dictionary of kwargs
     """