aicmo 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

aicmo/__init__.py

@@ -23,6 +23,7 @@ class AICMOClient:
             ts_host: str=None,
             ts_port: int=None,
             ts_api_key: str=None,
+            use_openrouter: bool=True
         ) -> None:
         """
         Initialize the AICMOClient with AWS credentials and OpenAI model.
@@ -44,9 +45,15 @@ class AICMOClient:
         else:
             self.aws_s3_bucket = self.secret_dict.get('AWS_S3_BUCKET', None)
         
-        # Initialize OpenAI client
-        openai_dict = {x:self.secret_dict[y] for x,y in (("api_key", "OPENAI_API_KEY"), ("organization", "OPENAI_ORG_KEY")) if self.secret_dict.get(y, None)}
-        self.openai_client = OpenAI(**openai_dict)
+        if use_openrouter:
+            # OPENROUTER_API_KEY
+            # OPENROUTER_BASE_URL
+            openai_dict = {x:self.secret_dict[y] for x,y in (("api_key", "OPENROUTER_API_KEY"), ("base_url", "OPENROUTER_BASE_URL")) if self.secret_dict.get(y, None)}
+            self.openai_client = OpenAI(**openai_dict)
+        else:    
+            # Initialize OpenAI client
+            openai_dict = {x:self.secret_dict[y] for x,y in (("api_key", "OPENAI_API_KEY"), ("organization", "OPENAI_ORG_KEY")) if self.secret_dict.get(y, None)}
+            self.openai_client = OpenAI(**openai_dict)
 
         # Initialize ScrapingBee client
         self.SCRAPINGBEE_API_KEY = self.secret_dict.get('SCRAPINGBEE_API_KEY', None)
@@ -75,7 +82,10 @@ class AICMOClient:
             self.ts_client = None
 
         # Initialize OpenAI model
-        self.OPENAI_MODEL = self.secret_dict.get('OPENAI_MODEL', None)
+        if use_openrouter:
+            self.OPENAI_MODEL = self.secret_dict.get('OPENROUTER_MODEL', None)
+        else:
+            self.OPENAI_MODEL = self.secret_dict.get('OPENAI_MODEL', None)
 
         # Initialize Costing for per APIs
         self.COST = json.loads(self.secret_dict['COST'])
@@ -188,25 +198,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(
@@ -300,9 +319,9 @@ class AICMOClient:
                     "text": {
                         "type": "mrkdwn",
                         "text": (
-                            f"Event ID: {event['event_id']}"
-                            f"User ID: {event['user_id']}"
-                            f"{tb}"
+                            f"Event ID: {event['event_id']}\n"
+                            f"User ID: {event['user_id']}\n"
+                            f"Error:\n{tb}\n"
                         )
                     }
                 }

{aicmo-0.0.1.dist-info → aicmo-0.0.3.dist-info}/METADATA

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

aicmo-0.0.3.dist-info/RECORD

@@ -0,0 +1,5 @@
+aicmo/__init__.py,sha256=mpBZRf4wGqigidr8cLuMW2Z20sWmAhmf3Sduv4zpK5Y,34310
+aicmo-0.0.3.dist-info/METADATA,sha256=R28YKWI7HOuJuJiJPwPOkObVQ4PFK4HpkRocsDT0p7M,1187
+aicmo-0.0.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+aicmo-0.0.3.dist-info/top_level.txt,sha256=uqMJv53xrmaqWkkPb9xI38tTYnSuPoplqyCzyKygG60,6
+aicmo-0.0.3.dist-info/RECORD,,

{aicmo-0.0.1.dist-info → aicmo-0.0.3.dist-info}/WHEEL

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

aicmo-0.0.1.dist-info/RECORD

@@ -1,5 +0,0 @@
-aicmo/__init__.py,sha256=mgPNc3BeqwTi0JvQknoBKwn6HOpk1Sb-ktiCdWlYsYc,33208
-aicmo-0.0.1.dist-info/METADATA,sha256=xwzhMlJ16IQepkdEmt22XIbY91MYO8vUeqA8LteHu_Y,1187
-aicmo-0.0.1.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-aicmo-0.0.1.dist-info/top_level.txt,sha256=uqMJv53xrmaqWkkPb9xI38tTYnSuPoplqyCzyKygG60,6
-aicmo-0.0.1.dist-info/RECORD,,