Update awsLib.py

awsLib.py

@@ -1,11 +1,15 @@
 import json, asyncio
 import requests
+from requests.models import PreparedRequest
 from requests_auth_aws_sigv4 import AWSSigV4
+from datetime import datetime
+from datetime import UTC
+from dateutil.relativedelta import relativedelta
 
 # lib by drago
 async def bedrock_model_available(session, key_id="", secret_key="", region="us-east-1", model="anthropic.claude-3-sonnet-20240229-v1%3A"):
     try:
-        model = model.replace(":","%3A") # issue with urlencode just this so should be fine just doint it with replace 
+        model = model.replace(":","%3A")
     except Exception as e:
         pass
     endpoint = f'https://bedrock.{region}.amazonaws.com/foundation-model-availability/'+model
@@ -81,4 +85,98 @@ async def bedrock_send_fake_form(session, key_id="", secret_key="", region="us-e
             response_data = await response.text()
             return json.loads(response_data)
             
-        return await response.json()
+        return await response.json()
+
+async def send_signed_request_bedrock(session, payload, model="", key_id="", secret_key="", region="us-east-1"):
+    try:
+        model = model.replace(":","%3A")
+    except Exception as e:
+        pass
+        
+    endpoint = f'https://bedrock-runtime.{region}.amazonaws.com/model/{model}/invoke'
+    
+    headers = {
+        "Content-Type": "application/json"
+    }
+    
+    aws_auth = AWSSigV4("bedrock",
+        aws_access_key_id=key_id,
+        aws_secret_access_key=secret_key,
+        region=region
+    )
+    
+    loop = asyncio.get_event_loop()
+    req = requests.Request('POST', endpoint, data=payload, headers=headers)
+    prepared_req = req.prepare()
+    signed_req = await loop.run_in_executor(None, aws_auth, prepared_req)
+
+    async with session.post(endpoint, headers=signed_req.headers, data=signed_req.body) as response:
+        if response.status != 200:
+            response_data = await response.text()
+            return json.loads(response_data)
+            
+        return await response.json()
+
+async def send_signed_request_ce(session, service, payload, key_id="", secret_key=""):
+    ce_endpoint = 'https://ce.us-east-1.amazonaws.com'
+    headers = {
+        "Content-Type": "application/x-www-form-urlencoded",
+        "Accept": "application/json",
+    }
+
+    aws_auth = AWSSigV4(service,
+        aws_access_key_id=key_id,
+        aws_secret_access_key=secret_key,
+        region="us-east-1"
+    )
+
+    # Creating the request with required headers for AWS Cost Explorer
+    headers = {
+        "Content-Type": "application/x-amz-json-1.1",
+        "Accept": "application/json",
+        "X-Amz-Target": f"AWSInsightsIndexService.GetCostAndUsage"  # This header is specific for the GetCostAndUsage operation
+    }
+
+    loop = asyncio.get_event_loop()
+    req = PreparedRequest()
+    req.prepare_url(ce_endpoint, {})
+    req.prepare_method('POST')
+    req.prepare_headers(headers)
+    req.prepare_body(data=payload, files=None)
+
+    signed_req = await loop.run_in_executor(None, aws_auth, req)
+    
+    async with session.post(ce_endpoint, headers=signed_req.headers, data=signed_req.body) as response:
+        response_text = await response.text()
+        
+        if response.status != 200:
+            try:
+                return json.loads(response_text)
+            except:
+                return {}
+        return json.loads(response_text)
+
+async def get_service_cost_and_usage(session, key, secret, service="Claude 3 Sonnet (Amazon Bedrock Edition)"):
+    now = datetime.now(UTC)  # Use UTC for AWS requests
+    six_months_ago = now - relativedelta(months=6)
+    start_date = six_months_ago.replace(day=1).strftime('%Y-%m-%d')
+    end_date = now.replace(day=30).strftime('%Y-%m-%d')
+    
+    payload = {
+        "TimePeriod": {
+            "Start": start_date,
+            "End": end_date
+        },
+        "Granularity": "MONTHLY",
+        "Metrics": ["BlendedCost"],
+        "Filter": {
+            "Dimensions": {
+                "Key": "SERVICE",
+                "Values": [service]          
+            }
+        }
+    }
+
+    json_payload = json.dumps(payload)
+
+    return await send_signed_request_ce(session, "ce", json_payload, key_id=key, secret_key=secret), service