Refactor file operations

api/db/logs.py

@@ -1,17 +1,21 @@
 import os
 import time
-import asyncio
 
 from dotenv import load_dotenv
 from motor.motor_asyncio import AsyncIOMotorClient
 
-try:
-    from helpers import network
-except ImportError:
-    pass
+from helpers import network
 
 load_dotenv()
 
+UA_SIMPLIFY = {
+    'Windows NT': 'W',
+    'Mozilla/5.0': 'M',
+    'Win64; x64': '64',
+    'Safari/537.36': 'S',
+    'AppleWebKit/537.36 (KHTML, like Gecko)': 'K',
+}
+
 ## MONGODB Setup
 
 conn = AsyncIOMotorClient(os.environ['MONGO_URI'])
@@ -26,7 +30,18 @@ async def replacer(text: str, dict_: dict) -> str:
     return text
 
 async def log_api_request(user: dict, incoming_request, target_url: str):
-    """Logs the API Request into the database."""
+    """Logs the API Request into the database.
+    No input prompt is logged, however data such as IP & useragent is noted.
+    This would be useful for security reasons. Other minor data is also collected.
+
+    Args:
+        user (dict): User dict object
+        incoming_request (_type_): Request
+        target_url (str): The URL the api request was targetted to.
+
+    Returns:
+        _type_: _description_
+    """
     db = await _get_collection('logs')
     payload = {}
 
@@ -38,6 +53,7 @@ async def log_api_request(user: dict, incoming_request, target_url: str):
 
     model = payload.get('model')
     ip_address = await network.get_ip(incoming_request)
+    useragent = await replacer(incoming_request.headers.get('User-Agent', ''), UA_SIMPLIFY)
 
     new_log_item = {
         'timestamp': time.time(),
@@ -46,6 +62,7 @@ async def log_api_request(user: dict, incoming_request, target_url: str):
         'user_id': str(user['_id']),
         'security': {
             'ip': ip_address,
+            'useragent': useragent,
         },
         'details': {
             'model': model,
@@ -73,21 +90,5 @@ async def delete_by_user_id(user_id: str):
     db = await _get_collection('logs')
     return await db.delete_many({'user_id': user_id})
 
-async def get_logs_time_range(start: int, end: int):
-    db = await _get_collection('logs')
-
-    entries = []
-    async for entry in db.find({'timestamp': {'$gte': start, '$lte': end}}):
-        entries.append(entry)
-
-    return entries
-
-async def main():
-    # how many requests in last 24 hours?
-    last_24_hours = time.time() - 86400
-    logs = await get_logs_time_range(last_24_hours, time.time())
-
-    print(f'Number of logs in last 24 hours: {len(logs)}')
-
 if __name__ == '__main__':
-    asyncio.run(main())
+    pass

api/db/providerkeys.py

@@ -100,7 +100,7 @@ manager = KeyManager()
 
 async def main():
     keys = await manager.get_possible_keys('closed')
-    print(keys)
+    print(len(keys))
 
 if __name__ == '__main__':
     asyncio.run(main())

api/db/stats.py

@@ -2,6 +2,8 @@ import os
 import pytz
 import asyncio
 import datetime
+import json
+import time
 
 from dotenv import load_dotenv
 from motor.motor_asyncio import AsyncIOMotorClient
@@ -13,6 +15,13 @@ load_dotenv()
 class StatsManager:
     """
     ### The manager for all statistics tracking
+    Stats tracked:
+    - Dates
+    - IPs
+    - Target URLs
+    - Tokens
+    - Models
+    - URL Paths
     """
 
     def __init__(self):

api/db/tester.py

@@ -0,0 +1,6 @@
+from stats import * 
+import asyncio
+
+manager = StatsManager()
+
+asyncio.run(manager.get_model_usage())

api/providers/__init__.py

@@ -1,12 +1,12 @@
 from . import \
-    azure \
-    # closed, \
-    # closed4
+    azure, \
+    closed, \
+    closed4
     # closed432
 
 MODULES = [
     azure,
-    # closed,
-    # closed4,
+    closed,
+    closed4,
     # closed432,
 ]

api/providers/azure.py

@@ -12,7 +12,7 @@ MODELS = [
     'gpt-4',
     'gpt-4-32k'
 ]
-# MODELS = [f'{model}-azure' for model in MODELS]
+MODELS = [f'{model}-azure' for model in MODELS]
 
 AZURE_API = '2023-07-01-preview'
 

api/providers/helpers/utils.py

@@ -34,4 +34,7 @@ async def conversation_to_prompt(conversation: list) -> str:
     return text
 
 async def random_secret_for(name: str) -> str:
-    return await providerkeys.manager.get_key(name)
+    try:
+        return await providerkeys.manager.get_key(name)
+    except ValueError:
+        raise ValueError(f'Keys missing for "{name}" <no_keys>')

api/responder.py

@@ -2,7 +2,7 @@
 
 import os
 import json
-import ujson
+import logging
 import aiohttp
 import asyncio
 import starlette
@@ -49,7 +49,9 @@ async def respond(
         'Content-Type': 'application/json'
     }
 
-    for i in range(1):
+    for i in range(20):
+        print(i)
+        # Load balancing: randomly selecting a suitable provider
         try:
             if is_chat:
                 target_request = await load_balancing.balance_chat_request(payload)
@@ -149,21 +151,13 @@ async def respond(
 
                         async for chunk in response.content.iter_any():
                             chunk = chunk.decode('utf8').strip()
-
-                            if 'azure' in provider_name:
-                                chunk = chunk.strip().replace('data: ', '')
-
-                                if not chunk or 'prompt_filter_results' in chunk:
-                                    continue
-
                             yield chunk + '\n\n'
 
                     break
 
             except Exception as exc:
                 print('[!] exception', exc)
-                # continue
-                raise exc
+                continue
 
     else:
         yield await errors.yield_error(500, 'Sorry, our API seems to have issues connecting to our provider(s).', 'This most likely isn\'t your fault. Please try again later.')

checks/client.py

@@ -215,10 +215,7 @@ async def demo():
         else:
             raise ConnectionError('API Server is not running.')
 
-        for func in [
-            # test_chat_non_stream_gpt4,
-            test_chat_stream_gpt3
-        ]:
+        for func in [test_chat_non_stream_gpt4, test_chat_stream_gpt3]:
             print(f'[*] {func.__name__}')
             result = await func()
             print(f'[+] {func.__name__} - {result}')

playground/notes.txt

@@ -1,10 +0,0 @@
---- EXPECTED ---
-
-data: {"id":"custom-chatcmpl-nUSiapqELukaPT7vEnGcXkbvrS1fR","object":"chat.completion.chunk","created":1696716717,"model":"gpt-3.5-turbo-0613","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
-
-data: {"id":"custom-chatcmpl-nUSiapqELukaPT7vEnGcXkbvrS1fR","object":"chat.completion.chunk","created":1696716717,"model":"gpt-3.5-turbo-0613","choices":[{"index":0,"delta":{"content":"123"},"finish_reason":null}]}
-
-data: {"id":"custom-chatcmpl-nUSiapqELukaPT7vEnGcXkbvrS1fR","object":"chat.completion.chunk","created":1696716717,"model":"gpt-3.5-turbo-0613","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
-
-data: [DONE]
-