diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 0196b16cf..c79fcd106 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -1,9 +1,3 @@ -import json -from aiodynamo.client import Client -from aiodynamo.credentials import Credentials, StaticCredentials -from aiodynamo.http.httpx import HTTPX -from aiodynamo.models import Throughput, KeySchema, KeySpec, KeyType, PayPerRequest -from yarl import URL from litellm.proxy.db.base_client import CustomDB from litellm.proxy._types import ( DynamoDBArgs, @@ -13,17 +7,32 @@ from litellm.proxy._types import ( ) from litellm import get_secret from typing import Any, List, Literal, Optional, Union -from aiodynamo.expressions import UpdateExpression, F, Value -from aiodynamo.models import ReturnValues -from aiodynamo.http.aiohttp import AIOHTTP -from aiohttp import ClientSession +import json from datetime import datetime class DynamoDBWrapper(CustomDB): + from aiodynamo.credentials import Credentials, StaticCredentials + credentials: Credentials def __init__(self, database_arguments: DynamoDBArgs): + from aiodynamo.client import Client + from aiodynamo.credentials import Credentials, StaticCredentials + from aiodynamo.http.httpx import HTTPX + from aiodynamo.models import ( + Throughput, + KeySchema, + KeySpec, + KeyType, + PayPerRequest, + ) + from yarl import URL + from aiodynamo.expressions import UpdateExpression, F, Value + from aiodynamo.models import ReturnValues + from aiodynamo.http.aiohttp import AIOHTTP + from aiohttp import ClientSession + self.throughput_type = None if database_arguments.billing_mode == "PAY_PER_REQUEST": self.throughput_type = PayPerRequest() @@ -46,6 +55,22 @@ class DynamoDBWrapper(CustomDB): """ Connect to DB, and creating / updating any tables """ + from aiodynamo.client import Client + from aiodynamo.credentials import Credentials, StaticCredentials + from aiodynamo.http.httpx import HTTPX + from aiodynamo.models import ( + Throughput, + KeySchema, + KeySpec, + KeyType, + PayPerRequest, + ) + from yarl import URL + from aiodynamo.expressions import UpdateExpression, F, Value + from aiodynamo.models import ReturnValues + from aiodynamo.http.aiohttp import AIOHTTP + from aiohttp import ClientSession + async with ClientSession() as session: client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) ## User @@ -97,6 +122,22 @@ class DynamoDBWrapper(CustomDB): async def insert_data( self, value: Any, table_name: Literal["user", "key", "config"] ): + from aiodynamo.client import Client + from aiodynamo.credentials import Credentials, StaticCredentials + from aiodynamo.http.httpx import HTTPX + from aiodynamo.models import ( + Throughput, + KeySchema, + KeySpec, + KeyType, + PayPerRequest, + ) + from yarl import URL + from aiodynamo.expressions import UpdateExpression, F, Value + from aiodynamo.models import ReturnValues + from aiodynamo.http.aiohttp import AIOHTTP + from aiohttp import ClientSession + async with ClientSession() as session: client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) table = None @@ -114,6 +155,22 @@ class DynamoDBWrapper(CustomDB): await table.put_item(item=value) async def get_data(self, key: str, table_name: Literal["user", "key", "config"]): + from aiodynamo.client import Client + from aiodynamo.credentials import Credentials, StaticCredentials + from aiodynamo.http.httpx import HTTPX + from aiodynamo.models import ( + Throughput, + KeySchema, + KeySpec, + KeyType, + PayPerRequest, + ) + from yarl import URL + from aiodynamo.expressions import UpdateExpression, F, Value + from aiodynamo.models import ReturnValues + from aiodynamo.http.aiohttp import AIOHTTP + from aiohttp import ClientSession + async with ClientSession() as session: client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) table = None @@ -152,6 +209,22 @@ class DynamoDBWrapper(CustomDB): async def update_data( self, key: str, value: dict, table_name: Literal["user", "key", "config"] ): + from aiodynamo.client import Client + from aiodynamo.credentials import Credentials, StaticCredentials + from aiodynamo.http.httpx import HTTPX + from aiodynamo.models import ( + Throughput, + KeySchema, + KeySpec, + KeyType, + PayPerRequest, + ) + from yarl import URL + from aiodynamo.expressions import UpdateExpression, F, Value + from aiodynamo.models import ReturnValues + from aiodynamo.http.aiohttp import AIOHTTP + from aiohttp import ClientSession + async with ClientSession() as session: client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) table = None