Compare commits

..

5 commits

Author SHA1 Message Date
nsde f0fee481db Models endpoint is now free and very quick 2023-08-09 11:43:05 +02:00
nsde a3d51ed92e Merge branch 'main' of https://github.com/Luna-OSS/nova-api 2023-08-09 11:15:50 +02:00
nsde 0b86f7c26c Added /v1/models and fixed key invalidation 2023-08-09 11:15:49 +02:00
Leander 0337eed472
Update users.py
made by_discord_id use strings
2023-08-09 11:06:55 +02:00
henceiusegentoo c741c03bae Made role-querying use a string for id instead of an int 2023-08-09 10:20:31 +02:00
8 changed files with 1401 additions and 37 deletions

View file

@ -15,7 +15,7 @@ with open('config/credits.yml', encoding='utf8') as f:
async def _get_mongo(collection_name: str):
return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name]
async def create(discord_id: int=0) -> dict:
async def create(discord_id: str='') -> dict:
"""Adds a new user to the MongoDB collection."""
chars = string.ascii_letters + string.digits
@ -36,7 +36,7 @@ async def create(discord_id: int=0) -> dict:
'ban_reason': '',
},
'auth': {
'discord': discord_id,
'discord': str(discord_id),
'github': None
}
}
@ -52,7 +52,7 @@ async def by_id(user_id: str):
async def by_discord_id(discord_id: str):
db = await _get_mongo('users')
return await db.find_one({'auth.discord': discord_id})
return await db.find_one({'auth.discord': str(int(discord_id))})
async def by_api_key(key: str):
db = await _get_mongo('users')

1349
api/models.json Normal file

File diff suppressed because it is too large Load diff

View file

@ -47,7 +47,8 @@ async def is_policy_violated(inp) -> bool:
return False
except Exception as exc:
# await provider_auth.invalidate_key(req.get('provider_auth'))
if '401' in str(exc):
await provider_auth.invalidate_key(req.get('provider_auth'))
print('[!] moderation error:', type(exc), exc)
continue

View file

@ -10,6 +10,7 @@ from dotenv import load_dotenv
from python_socks._errors import ProxyError
import proxies
import provider_auth
import load_balancing
from db import logs, users, stats
@ -36,8 +37,6 @@ async def stream(
input_tokens: int=0,
incoming_request: starlette.requests.Request=None,
):
payload = payload or DEMO_PAYLOAD
is_chat = False
is_stream = payload.get('stream', False)
@ -45,7 +44,6 @@ async def stream(
is_chat = True
model = payload['model']
if is_chat and is_stream:
chat_id = await chat.create_chat_id()
@ -88,7 +86,7 @@ async def stream(
webhook = dhooks.Webhook(os.getenv('DISCORD_WEBHOOK__API_ISSUE'))
webhook.send(content=f'API Issue: **`{exc}`**\nhttps://i.imgflip.com/7uv122.jpg')
error = errors.yield_error(
error = await errors.yield_error(
500,
'Sorry, the API has no working keys anymore.',
'The admins have been messaged automatically.'
@ -97,7 +95,10 @@ async def stream(
return
for k, v in target_request.get('headers', {}).items():
headers[k] = v
target_request['headers'][k] = v
if target_request['method'] == 'GET' and not payload:
target_request['payload'] = None
async with aiohttp.ClientSession(connector=proxies.default_proxy.connector) as session:
try:
@ -108,24 +109,30 @@ async def stream(
data=target_request.get('data'),
json=target_request.get('payload'),
headers=headers,
headers=target_request.get('headers', {}),
cookies=target_request.get('cookies'),
ssl=False,
timeout=aiohttp.ClientTimeout(total=float(os.getenv('TRANSFER_TIMEOUT', '120'))),
) as response:
if response.content_type == 'application/json':
data = await response.json()
if not is_stream:
json_response = await response.json()
if data.get('code') == 'invalid_api_key':
await provider_auth.invalidate_key(target_request.get('provider_auth'))
continue
if response.ok:
json_response = data
if is_stream:
try:
response.raise_for_status()
except Exception as exc:
if 'Too Many Requests' in str(exc):
continue
if is_stream:
try:
async for chunk in response.content.iter_any():
send = False
@ -159,7 +166,7 @@ async def stream(
except Exception as exc:
if 'Connection closed' in str(exc):
error = errors.yield_error(
error = await errors.yield_error(
500,
'Sorry, there was an issue with the connection.',
'Please first check if the issue on your end. If this error repeats, please don\'t heistate to contact the staff!.'
@ -180,10 +187,9 @@ async def stream(
content=chat.CompletionStop
)
yield chunk
yield 'data: [DONE]\n\n'
if not is_stream:
if not is_stream and json_response:
yield json.dumps(json_response)
# DONE =========================================================

View file

@ -14,6 +14,8 @@ from helpers import tokens, errors
load_dotenv()
models_list = json.load(open('models.json'))
with open('config/credits.yml', encoding='utf8') as f:
credits_config = yaml.safe_load(f)
@ -65,6 +67,9 @@ async def handle(incoming_request):
error = await errors.error(418, 'Your NovaAI account is not active (paused).', 'Simply re-activate your account using a Discord command or the web panel.')
return error
if '/models' in path:
return fastapi.responses.JSONResponse(content=models_list)
# COST
costs = credits_config['costs']
cost = costs['other']
@ -101,8 +106,6 @@ async def handle(incoming_request):
# READY
# payload['user'] = str(user['_id'])
if 'chat/completions' in path and not payload.get('stream') is True:
payload['stream'] = False

View file

@ -1 +1 @@
1691460004.7354248
1691546405.042006

View file

@ -35,7 +35,7 @@ async def update_roles(mongo):
for role in level_role_names:
if role in role_names:
users.update_one(
{'auth.discord': int(discord)},
{'auth.discord': discord},
{'$set': {'level': role}}
)
print(f'Updated {discord} to {role}')

View file

@ -40,11 +40,6 @@ def test_server():
def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
"""Tests an API api_endpoint."""
headers = {
'Content-Type': 'application/json',
'Authorization': 'Bearer ' + closedai.api_key
}
json_data = {
'model': model,
'messages': messages or MESSAGES,
@ -53,7 +48,7 @@ def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
response = httpx.post(
url=f'{api_endpoint}/chat/completions',
headers=headers,
headers=HEADERS,
json=json_data,
timeout=20
)
@ -74,25 +69,30 @@ def test_library():
def test_library_moderation():
return closedai.Moderation.create("I wanna kill myself, I wanna kill myself; It's all I hear right now, it's all I hear right now")
def test_models():
response = httpx.get(
url=f'{api_endpoint}/models',
headers=HEADERS,
timeout=5
)
response.raise_for_status()
return response.json()
def test_all():
"""Runs all tests."""
# print(test_server())
# print(test_api())
print(test_library())
# print(test_library())
# print(test_library_moderation())
print(test_models())
def test_api_moderation(model: str=MODEL, messages: List[dict]=None) -> dict:
"""Tests an API api_endpoint."""
headers = {
'Authorization': 'Bearer ' + closedai.api_key
}
response = httpx.get(
url=f'{api_endpoint}/moderations',
headers=headers,
headers=HEADERS,
timeout=20
)
response.raise_for_status()
@ -104,4 +104,9 @@ if __name__ == '__main__':
closedai.api_base = api_endpoint
closedai.api_key = os.getenv('TEST_NOVA_KEY')
HEADERS = {
'Content-Type': 'application/json',
'Authorization': 'Bearer ' + closedai.api_key
}
test_all()