format new app

Agenta-AI · Apr 2, 2024 · f90c3eb · f90c3eb
1 parent 23af353
commit f90c3eb
Showing 1 changed file with 23 additions and 11 deletions.
diff --git a/examples/what_to_do_today/app.py b/examples/what_to_do_today/app.py
@@ -1,7 +1,7 @@
 import agenta as ag
 import litellm
 
-litellm.drop_params=True
+litellm.drop_params = True
 
 ag.init()
 
@@ -11,7 +11,6 @@
 I am {marital_status} and I have {duration} free time?
 I am not interested in activities like: {not_desired_activities}
     """,
-
 }
 
 provider_models = {
@@ -33,27 +32,31 @@
     ],
     "Cohere": [
         "command-nightly",
-    ]
+    ],
 }
 
 GPT_FORMAT_RESPONSE = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview"]
 
 ag.config.default(
     temperature=ag.FloatParam(default=1, minval=0.0, maxval=2.0),
-    model=ag.GroupedMultipleChoiceParam(default="gpt-3.5-turbo", choices=provider_models),
+    model=ag.GroupedMultipleChoiceParam(
+        default="gpt-3.5-turbo", choices=provider_models
+    ),
     max_tokens=ag.IntParam(-1, -1, 4000),
     prompt_system=ag.TextParam(prompts["system_prompt"]),
     prompt_user=ag.TextParam(prompts["user_prompt"]),
     top_p=ag.FloatParam(1),
     frequence_penalty=ag.FloatParam(default=0.0, minval=-2.0, maxval=2.0),
     presence_penalty=ag.FloatParam(default=0.0, minval=-2.0, maxval=2.0),
-    force_json=ag.BinaryParam(False)
+    force_json=ag.BinaryParam(False),
 )
 
 
 @ag.entrypoint
 async def generate(
-    inputs: ag.DictInput = ag.DictInput(default_keys=["country", "marital_status", "duration", "not_desired_activities"]),
+    inputs: ag.DictInput = ag.DictInput(
+        default_keys=["country", "marital_status", "duration", "not_desired_activities"]
+    ),
 ):
     try:
         prompt_user = ag.config.prompt_user.format(**inputs)
@@ -67,13 +70,22 @@ async def generate(
     max_tokens = ag.config.max_tokens if ag.config.max_tokens != -1 else None
 
     if ag.config.force_json and ag.config.model not in GPT_FORMAT_RESPONSE:
-        raise ValueError("Model {} does not support JSON response format".format(ag.config.model))
+        raise ValueError(
+            "Model {} does not support JSON response format".format(ag.config.model)
+        )
 
-    response_format = {"type": "json_object"} if ag.config.force_json and ag.config.model in GPT_FORMAT_RESPONSE else {"type": "text"}
+    response_format = (
+        {"type": "json_object"}
+        if ag.config.force_json and ag.config.model in GPT_FORMAT_RESPONSE
+        else {"type": "text"}
+    )
 
     completion_params = {
         "model": ag.config.model,
-        "messages": [{"content": prompt_system, "role": "system"}, {"content": prompt_user, "role": "user"}],
+        "messages": [
+            {"content": prompt_system, "role": "system"},
+            {"content": prompt_user, "role": "user"},
+        ],
         "temperature": ag.config.temperature,
         "max_tokens": max_tokens,
         "top_p": ag.config.top_p,
@@ -89,7 +101,7 @@ async def generate(
 
     token_usage = response.usage.dict()
     return {
-        "message": response.choices[0]. message.content,
+        "message": response.choices[0].message.content,
         **{"usage": token_usage},
         "cost": ag.calculate_token_usage(ag.config.model, token_usage),
-    }
+    }