Add successful request count to OpenAI callback (#2128)

I've found it useful to track the number of successful requests to
OpenAI. This gives me a better sense of the efficiency of my prompts and
helps compare map_reduce/refine on a cheaper model vs. stuffing on a
more expensive model with higher capacity.
This commit is contained in:
Jonathan Page 2023-03-28 22:56:17 -07:00 committed by GitHub
parent 4ab66c4f52
commit 8441cbfc03
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 2 deletions

View File

@ -45,7 +45,8 @@
"text": [ "text": [
"Total Tokens: 42\n", "Total Tokens: 42\n",
"Prompt Tokens: 4\n", "Prompt Tokens: 4\n",
"Completion Tokens: 38\n" "Completion Tokens: 38\n",
"Successful Requests: 1\n"
] ]
} }
], ],
@ -54,7 +55,8 @@
" result = llm(\"Tell me a joke\")\n", " result = llm(\"Tell me a joke\")\n",
" print(f\"Total Tokens: {cb.total_tokens}\")\n", " print(f\"Total Tokens: {cb.total_tokens}\")\n",
" print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n", " print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
" print(f\"Completion Tokens: {cb.completion_tokens}\")" " print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
" print(f\"Successful Requests: {cb.successful_requests}\")"
] ]
}, },
{ {

View File

@ -11,6 +11,7 @@ class OpenAICallbackHandler(BaseCallbackHandler):
total_tokens: int = 0 total_tokens: int = 0
prompt_tokens: int = 0 prompt_tokens: int = 0
completion_tokens: int = 0 completion_tokens: int = 0
successful_requests: int = 0
@property @property
def always_verbose(self) -> bool: def always_verbose(self) -> bool:
@ -30,6 +31,7 @@ class OpenAICallbackHandler(BaseCallbackHandler):
def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
"""Collect token usage.""" """Collect token usage."""
if response.llm_output is not None: if response.llm_output is not None:
self.successful_requests += 1
if "token_usage" in response.llm_output: if "token_usage" in response.llm_output:
token_usage = response.llm_output["token_usage"] token_usage = response.llm_output["token_usage"]
if "total_tokens" in token_usage: if "total_tokens" in token_usage: