update pricing model to work with azure models

NicolasAG · NicolasAG · commit b3179258e7c8 · 2025-11-20T16:10:31.000Z
diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py
@@ -433,7 +433,7 @@ def __init__(
             min_retry_wait_time=min_retry_wait_time,
             client_class=OpenAI,
             client_args=client_args,
-            pricing_func=tracking.get_pricing_openai,
+            pricing_func=partial(tracking.get_pricing_litellm, model_name=model_name),
             log_probs=log_probs,
         )
 

Original file line number	Diff line number	Diff line change
`@@ -433,7 +433,7 @@ def __init__(`
`433`	`433`	`min_retry_wait_time=min_retry_wait_time,`
`434`	`434`	`client_class=OpenAI,`
`435`	`435`	`client_args=client_args,`
`436`		`- pricing_func=tracking.get_pricing_openai,`
	`436`	`+ pricing_func=partial(tracking.get_pricing_litellm, model_name=model_name),`
`437`	`437`	`log_probs=log_probs,`
`438`	`438`	`)`
`439`	`439`