aicmo 0.0.1__tar.gz → 0.0.3__tar.gz

{aicmo-0.0.1 → aicmo-0.0.3}/PKG-INFO

@@ -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.1 → aicmo-0.0.3}/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 → aicmo-0.0.3}/aicmo.egg-info/PKG-INFO

@@ -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.1 → aicmo-0.0.3}/setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-VERSION = '0.0.1'
+VERSION = '0.0.3'
 
 # Setting up
 setup(