feat: log openai prompt usage tokens

src/planai/openai.py

@@ -11,9 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 from typing import Any, Dict, List, Literal, Mapping
 
-from openai import ContentFilterFinishReasonError, OpenAI, LengthFinishReasonError
+from openai import ContentFilterFinishReasonError, LengthFinishReasonError, OpenAI
 
 
 class OpenAIWrapper:
@@ -130,6 +131,25 @@ def chat(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
                 response = self.client.chat.completions.create(**api_params)
                 content = response.choices[0].message.content
 
+            # Log the usage details
+            usage = response.usage
+            logging.info(
+                "Usage details - Prompt tokens: %d, Completion tokens: %d, Total tokens: %d, Cached tokens: %d, Reasoning tokens: %d",
+                usage.prompt_tokens,
+                usage.completion_tokens,
+                usage.total_tokens,
+                (
+                    usage.prompt_tokens_details.cached_tokens
+                    if usage.prompt_tokens_details
+                    else 0
+                ),
+                (
+                    usage.completion_tokens_details.reasoning_tokens
+                    if usage.completion_tokens_details
+                    else 0
+                ),
+            )
+
             return {"message": {"content": content}}
         except LengthFinishReasonError:
             # Handle the length error

tests/planai/test_openai.py

@@ -2,6 +2,7 @@
 from unittest.mock import MagicMock, Mock, patch
 
 from openai import ContentFilterFinishReasonError, LengthFinishReasonError
+from openai.types import CompletionUsage
 
 from planai.openai import OpenAIWrapper
 
@@ -74,7 +75,10 @@ def test_generate_with_json_format(self):
 
     def test_chat_basic(self):
         self.mock_client.chat.completions.create.return_value = Mock(
-            choices=[Mock(message=Mock(content="Chat response content"))]
+            choices=[Mock(message=Mock(content="Chat response content"))],
+            usage=CompletionUsage(
+                completion_tokens=9, prompt_tokens=10, total_tokens=19
+            ),
         )
 
         messages = [{"role": "user", "content": "Hello, assistant!"}]
@@ -127,7 +131,12 @@ def test_chat_with_response_schema(self):
         # Configure __contains__ to allow 'in' checks
         mock_message.__contains__.side_effect = lambda key: key in mock_message.__dict__
 
-        mock_response = Mock(choices=[Mock(message=mock_message)])
+        mock_response = Mock(
+            choices=[Mock(message=mock_message)],
+            usage=CompletionUsage(
+                completion_tokens=9, prompt_tokens=10, total_tokens=19
+            ),
+        )
         self.mock_client.beta.chat.completions.parse.return_value = mock_response
 
         messages = [{"role": "user", "content": "Test message"}]
@@ -139,7 +148,10 @@ def test_chat_with_response_schema(self):
 
     def test_chat_with_options(self):
         self.mock_client.chat.completions.create.return_value = Mock(
-            choices=[Mock(message=Mock(content="Chat response with options"))]
+            choices=[Mock(message=Mock(content="Chat response with options"))],
+            usage=CompletionUsage(
+                completion_tokens=9, prompt_tokens=10, total_tokens=19
+            ),
         )
 
         messages = [{"role": "user", "content": "Test message"}]