Refactor (examples/): remove response structure of llm example apps

Agenta-AI · May 20, 2024 · 264dbbc · 264dbbc
1 parent abaa09c
commit 264dbbc
Show file tree

Hide file tree

Showing 8 changed files with 9 additions and 50 deletions.
diff --git a/examples/async_chat_sdk_output_format/app.py b/examples/async_chat_sdk_output_format/app.py
@@ -33,9 +33,4 @@ async def chat(inputs: MessagesInput = MessagesInput()):
         temperature=ag.config.temperature,
         max_tokens=max_tokens,
     )
-    token_usage = chat_completion.usage.dict()
-    return {
-        "message": chat_completion.choices[0].message.content,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    return chat_completion.choices[0].message.content
diff --git a/examples/async_startup_technical_ideas/app.py b/examples/async_startup_technical_ideas/app.py
@@ -33,9 +33,4 @@ async def chat(inputs: MessagesInput = MessagesInput()) -> str:
         temperature=ag.config.temperature,
         max_tokens=max_tokens,
     )
-    token_usage = chat_completion.usage.dict()
-    return {
-        "message": chat_completion.choices[0].message.content,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    return chat_completion.choices[0].message.content
diff --git a/examples/baby_name_generator/app.py b/examples/baby_name_generator/app.py
@@ -8,7 +8,7 @@
     "Give me 10 names for a baby from this country {country} with gender {gender}!!!!"
 )
 
-ag.init(app_name="test", base_name="app")
+ag.init()
 ag.config.default(
     temperature=FloatParam(0.2), prompt_template=TextParam(default_prompt)
 )
@@ -31,9 +31,4 @@ def generate(country: str, gender: str) -> str:
     chat_completion = client.chat.completions.create(
         model="gpt-3.5-turbo", messages=[{"role": "user", "content": prompt}]
     )
-    token_usage = chat_completion.usage.dict()
-    return {
-        "message": chat_completion.choices[0].message.content,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage("gpt-3.5-turbo", token_usage),
-    }
+    return chat_completion.choices[0].message.content
diff --git a/examples/chat_json_format/app.py b/examples/chat_json_format/app.py
@@ -39,9 +39,4 @@ def chat(inputs: ag.MessagesInput = ag.MessagesInput()):
         max_tokens=max_tokens,
         response_format=response_format,
     )
-    token_usage = chat_completion.usage.dict()
-    return {
-        "message": chat_completion.choices[0].message.content,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    return chat_completion.choices[0].message.content
diff --git a/examples/extract_data_to_json/app.py b/examples/extract_data_to_json/app.py
@@ -48,10 +48,5 @@ def generate(
         temperature=ag.config.temperature,
         functions=[function],
     )
-    token_usage = response.usage.dict()
     output = str(response["choices"][0]["message"]["function_call"])
-    return {
-        "message": output,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    return output
diff --git a/examples/startup_technical_ideas/app.py b/examples/startup_technical_ideas/app.py
@@ -32,9 +32,4 @@ def chat(inputs: MessagesInput = MessagesInput()) -> str:
         temperature=ag.config.temperature,
         max_tokens=max_tokens,
     )
-    token_usage = chat_completion.usage.dict()
-    return {
-        "message": chat_completion.choices[0].message.content,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    return chat_completion.choices[0].message.content
diff --git a/examples/vision_gpt_explain_image/app.py b/examples/vision_gpt_explain_image/app.py
@@ -47,9 +47,4 @@ def explain(
         messages=messages,
         max_tokens=ag.config.max_tokens,
     )
-    token_usage = chat_completion.usage.dict()
-    return {
-        "message": chat_completion.choices[0].message.content,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    return chat_completion.choices[0].message.content
diff --git a/examples/what_to_do_today/app.py b/examples/what_to_do_today/app.py
@@ -99,10 +99,4 @@ async def generate(
         completion_params["presence_penalty"] = ag.config.presence_penalty
 
     response = await litellm.acompletion(**completion_params)
-
-    token_usage = response.usage.dict()
-    return {
-        "message": response.choices[0].message.content,
-        **{"usage": token_usage},
-        "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    return response.choices[0].message.content