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