From 27a9d2bc689852e4266998ca049558ce945a01bb Mon Sep 17 00:00:00 2001
From: Chan Jun Shern <chanjunshern@gmail.com>
Date: Wed, 3 Jan 2024 11:59:24 +0800
Subject: [PATCH] Log model and usage stats in sample logger

---
 evals/completion_fns/openai.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/evals/completion_fns/openai.py b/evals/completion_fns/openai.py
index ed50818630..f3075f64bb 100644
--- a/evals/completion_fns/openai.py
+++ b/evals/completion_fns/openai.py
@@ -88,7 +88,12 @@ def __call__(
             **{**kwargs, **self.extra_options},
         )
         result = OpenAICompletionResult(raw_data=result, prompt=openai_create_prompt)
-        record_sampling(prompt=result.prompt, sampled=result.get_completions())
+        record_sampling(
+            prompt=result.prompt,
+            sampled=result.get_completions(),
+            model=result.raw_data.model,
+            usage=result.raw_data.usage,
+        )
         return result
 
 
@@ -133,5 +138,10 @@ def __call__(
             **{**kwargs, **self.extra_options},
         )
         result = OpenAIChatCompletionResult(raw_data=result, prompt=openai_create_prompt)
-        record_sampling(prompt=result.prompt, sampled=result.get_completions())
+        record_sampling(
+            prompt=result.prompt,
+            sampled=result.get_completions(),
+            model=result.raw_data.model,
+            usage=result.raw_data.usage,
+        )
         return result