aicmo 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

aicmo/__init__.py

@@ -134,7 +134,6 @@ class AICMOClient:
                     tools=tools,
                     tool_choice={"type": "function", "function": {"name": tool_name}}
                 )
-                tokens = self.get_gpt_tokens(completion, model, tokens)
                 ret_val['completion'] = completion
                 ret_val['status'] = "success"
                 ret_val['tokens'] = self.get_gpt_tokens(completion, model, tokens)
@@ -198,25 +197,34 @@ class AICMOClient:
             self,
             data: types.chat.chat_completion.ChatCompletion,
             model: str,
-            tokens: dict
+            tokens: dict,
+            use_openrouter: bool=True
         ) -> dict:
         """
         Calculate the token usage and cost.
         """
-        prompt_tokens = data.usage.prompt_tokens
-        completion_tokens = data.usage.completion_tokens
-        cost = self.COST.get('openai', {}).get("texts", {}).get(model, {})
-        openai_input_cost = cost.get('input', None)
-        openai_output_cost = cost.get('output', None)
-        if openai_input_cost and openai_output_cost:
-            tokens['input_cost'] += prompt_tokens * openai_input_cost
-            tokens['output_cost'] += completion_tokens * openai_output_cost
-            tokens['total_cost'] = round(tokens['input_cost'] + tokens['output_cost'], 4)
-        tokens['prompt_tokens'] += prompt_tokens
-        tokens['completion_tokens'] += completion_tokens
-        tokens['total_tokens'] += prompt_tokens + completion_tokens
-        tokens['openai_input_cost'] = openai_input_cost
-        tokens['openai_output_cost'] = openai_output_cost
+        if use_openrouter:
+            tokens['total_cost'] += data.usage.cost
+            tokens['input_cost'] += data.usage.cost_details['upstream_inference_prompt_cost']
+            tokens['output_cost'] += data.usage.cost_details['upstream_inference_completions_cost']
+            tokens["prompt_tokens"] += data.usage.prompt_tokens
+            tokens["completion_tokens"] += data.usage.completion_tokens
+            tokens["total_tokens"] += data.usage.total_tokens
+        else:
+            prompt_tokens = data.usage.prompt_tokens
+            completion_tokens = data.usage.completion_tokens
+            cost = self.COST.get('openai', {}).get("texts", {}).get(model, {})
+            openai_input_cost = cost.get('input', None)
+            openai_output_cost = cost.get('output', None)
+            if openai_input_cost and openai_output_cost:
+                tokens['input_cost'] += prompt_tokens * openai_input_cost
+                tokens['output_cost'] += completion_tokens * openai_output_cost
+                tokens['total_cost'] = round(tokens['input_cost'] + tokens['output_cost'], 4)
+            tokens['prompt_tokens'] += prompt_tokens
+            tokens['completion_tokens'] += completion_tokens
+            tokens['total_tokens'] += prompt_tokens + completion_tokens
+            tokens['openai_input_cost'] = openai_input_cost
+            tokens['openai_output_cost'] = openai_output_cost
         return tokens
     
     def s3_upload_pickle(

{aicmo-0.0.2.dist-info → aicmo-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aicmo
-Version: 0.0.2
+Version: 0.0.4
 Summary: A package for using aicmo functions and tools
 Author: Jayr Castro
 Author-email: jayrcastro.py@gmail.com

aicmo-0.0.4.dist-info/RECORD

@@ -0,0 +1,5 @@
+aicmo/__init__.py,sha256=kqXJeT_TiN5ywY4fHF1DM6TcG6WlQx5-L9t8ih36XII,34238
+aicmo-0.0.4.dist-info/METADATA,sha256=D7_AemHeR9AWAowHIP5kttR8IAtvKHwdzb-LUVoW8vU,1187
+aicmo-0.0.4.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+aicmo-0.0.4.dist-info/top_level.txt,sha256=uqMJv53xrmaqWkkPb9xI38tTYnSuPoplqyCzyKygG60,6
+aicmo-0.0.4.dist-info/RECORD,,

{aicmo-0.0.2.dist-info → aicmo-0.0.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.1)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

aicmo-0.0.2.dist-info/RECORD

@@ -1,5 +0,0 @@
-aicmo/__init__.py,sha256=4hgGmwqW9btLI6Ik3prZSJ0JrC0jHv0LOEZA1aOns_c,33730
-aicmo-0.0.2.dist-info/METADATA,sha256=PIStOBiJ79wqhgAXEffY0lqcGc4xWyfTi1KGUrkJIQI,1187
-aicmo-0.0.2.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-aicmo-0.0.2.dist-info/top_level.txt,sha256=uqMJv53xrmaqWkkPb9xI38tTYnSuPoplqyCzyKygG60,6
-aicmo-0.0.2.dist-info/RECORD,,