phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) dynamoDB import issues

Browse source
This commit is contained in:
ishaan-jaff 2024-01-13 11:38:35 -08:00
parent 7fc8ccf854
commit b931841e60
1 changed files with 83 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

93
litellm/proxy/db/dynamo_db.py
View file

@ -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.db.base_client import CustomDB
from litellm.proxy._types import ( from litellm.proxy._types import (
DynamoDBArgs, DynamoDBArgs,
@ -13,17 +7,32 @@ from litellm.proxy._types import (
) )
from litellm import get_secret from litellm import get_secret
from typing import Any, List, Literal, Optional, Union from typing import Any, List, Literal, Optional, Union
from aiodynamo.expressions import UpdateExpression, F, Value import json
from aiodynamo.models import ReturnValues
from aiodynamo.http.aiohttp import AIOHTTP
from aiohttp import ClientSession
from datetime import datetime from datetime import datetime
class DynamoDBWrapper(CustomDB): class DynamoDBWrapper(CustomDB):
from aiodynamo.credentials import Credentials, StaticCredentials
credentials: Credentials credentials: Credentials
def __init__(self, database_arguments: DynamoDBArgs): 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 self.throughput_type = None
if database_arguments.billing_mode == "PAY_PER_REQUEST": if database_arguments.billing_mode == "PAY_PER_REQUEST":
self.throughput_type = PayPerRequest() self.throughput_type = PayPerRequest()
@ -46,6 +55,22 @@ class DynamoDBWrapper(CustomDB):
""" """
Connect to DB, and creating / updating any tables 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: async with ClientSession() as session:
client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
## User ## User
@ -97,6 +122,22 @@ class DynamoDBWrapper(CustomDB):
async def insert_data( async def insert_data(
self, value: Any, table_name: Literal["user", "key", "config"] 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: async with ClientSession() as session:
client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
table = None table = None
@ -114,6 +155,22 @@ class DynamoDBWrapper(CustomDB):
await table.put_item(item=value) await table.put_item(item=value)
async def get_data(self, key: str, table_name: Literal["user", "key", "config"]): 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: async with ClientSession() as session:
client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
table = None table = None
@ -152,6 +209,22 @@ class DynamoDBWrapper(CustomDB):
async def update_data( async def update_data(
self, key: str, value: dict, table_name: Literal["user", "key", "config"] 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: async with ClientSession() as session:
client = Client(AIOHTTP(session), Credentials.auto(), self.region_name) client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
table = None table = None