ZK-AI
/
ollama_proxy_server
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update main.py

This commit is contained in:
Saifeddine ALOUI 2024-01-16 11:18:23 +01:00 committed by GitHub
parent 18261b625b
commit da68e7f1cc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
ollama_proxy_server/main.py
View File

@ -44,19 +44,19 @@ def main():
authorized_users = get_authorized_users(args.users_list) authorized_users = get_authorized_users(args.users_list)
class RequestHandler(BaseHTTPRequestHandler): class RequestHandler(BaseHTTPRequestHandler):
def add_access_log_entry(self, user, ip_address, access): def add_access_log_entry(self, user, ip_address, access, server):
log_file_path = Path(sys.argv[1]) log_file_path = Path(sys.argv[1])
if not log_file_path.exists(): if not log_file_path.exists():
with open(log_file_path, mode='w', newline='') as csvfile: with open(log_file_path, mode='w', newline='') as csvfile:
fieldnames = ['time_stamp', 'user_name', 'ip_address','access'] fieldnames = ['time_stamp', 'user_name', 'ip_address','access','server']
writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader() writer.writeheader()
with open(log_file_path, mode='a', newline='') as csvfile: with open(log_file_path, mode='a', newline='') as csvfile:
fieldnames = ['time_stamp', 'user_name', 'ip_address'] fieldnames = ['time_stamp', 'user_name', 'ip_address']
writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
row = {'time_stamp': str(datetime.datetime.now()), 'user_name': user, 'ip_address': ip_address, 'access': access} row = {'time_stamp': str(datetime.datetime.now()), 'user_name': user, 'ip_address': ip_address, 'access': access, 'server': server}
writer.writerow(row) writer.writerow(row)
def _send_response(self, response): def _send_response(self, response):
@ -82,12 +82,18 @@ def main():
user, key = token.split(':') user, key = token.split(':')
# Check if the user and key are in the list of authorized users # Check if the user and key are in the list of authorized users
return authorized_users.get(user) == key if authorized_users.get(user) == key:
self.user = user
return True
else:
self.user = "unknown"
return False
def proxy(self): def proxy(self):
if not self._validate_user_and_key(): if not self._validate_user_and_key():
ASCIIColors.red(f'User is not authorized') ASCIIColors.red(f'User is not authorized')
# self.add_access_log_entry("unknown") client_ip, client_port = self.client_address
self.add_access_log_entry(user="unknown", ip_address=client_ip, "Denied", "None")
self.send_response(403) self.send_response(403)
self.end_headers() self.end_headers()
return return
@ -114,6 +120,7 @@ def main():
# Apply the queuing mechanism only for a specific endpoint. # Apply the queuing mechanism only for a specific endpoint.
if path == '/api/generate': if path == '/api/generate':
que = min_queued_server[1]['queue'] que = min_queued_server[1]['queue']
self.add_access_log_entry(user=self.user, ip_address=client_ip, "Authorized", min_queued_server[0])
que.put_nowait(1) que.put_nowait(1)
try: try:
response = requests.request(self.command, min_queued_server[1]['url'] + path, params=get_params, data=post_params) response = requests.request(self.command, min_queued_server[1]['url'] + path, params=get_params, data=post_params)