Merge pull request #9 from petritavd/fix-stream

fix stream issue

README.md

@@ -54,7 +54,7 @@ ollama_proxy_add_user --users_list [path to the authorized `authorized_users.txt
 ### Starting the server
 Start the Ollama Proxy Server by running the following command in your terminal:
 ```bash
-ollama_proxy_server --config [configuration file path] --users_list [users list file path] --port [port number to access the proxy]
+python3 ollama_proxy_server/main.py --config [configuration file path] --users_list [users list file path] --port [port number to access the proxy]
 ```
 The server will listen on port 808x, with x being the number of available ports starting from 0 (e.g., 8080, 8081, etc.). The first available port will be automatically selected if no other instance is running.
 

ollama_proxy_server/main.py

@@ -7,13 +7,12 @@ description: This is a proxy server that adds a security layer to one or multipl
 
 import configparser
 from http.server import BaseHTTPRequestHandler, HTTPServer
+import json
 from socketserver import ThreadingMixIn
 from urllib.parse import urlparse, parse_qs
 from queue import Queue
 import requests
-import threading
 import argparse
-import base64
 from ascii_colors import ASCIIColors
 from pathlib import Path
 import csv
@@ -70,23 +69,23 @@ def main():
                 writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                 row = {'time_stamp': str(datetime.datetime.now()), 'event':event, 'user_name': user, 'ip_address': ip_address, 'access': access, 'server': server, 'nb_queued_requests_on_server': nb_queued_requests_on_server, 'error': error}
                 writer.writerow(row)
-                
+
         def _send_response(self, response):
             self.send_response(response.status_code)
-            self.send_header('Content-type', response.headers['content-type'])
+            for key, value in response.headers.items():
+                if key.lower() not in ['content-length', 'transfer-encoding', 'content-encoding']:
+                    self.send_header(key, value)
+            self.send_header('Transfer-Encoding', 'chunked')
             self.end_headers()
-            self.wfile.write(response.content)
 
-        def _send_response_stream(self, response):
+            try:
-            self.send_response(response.status_code)
+                for chunk in response.iter_content(chunk_size=1024):
-            self.send_header('Content-type', response.headers['content-type'])
+                    if chunk:
-            self.send_header('Stream', True)
+                        self.wfile.write(b"%X\r\n%s\r\n" % (len(chunk), chunk))
-            self.end_headers()
+                        self.wfile.flush()
-            for line in response.iter_lines():
+                self.wfile.write(b"0\r\n\r\n")
-                if line:
+            except BrokenPipeError:
-                    chunk = line.decode('utf-8') + '\r\n'
+                pass
-                    self.wfile.write(chunk.encode('utf-8'))
-                    self.wfile.flush()
 
         def do_GET(self):
             self.log_request()
@@ -157,8 +156,14 @@ def main():
                 self.add_access_log_entry(event="gen_request", user=self.user, ip_address=client_ip, access="Authorized", server=min_queued_server[0], nb_queued_requests_on_server=que.qsize())
                 que.put_nowait(1)
                 try:
-                    response = requests.request(self.command, min_queued_server[1]['url'] + path, params=get_params, data=post_params, stream=True)
+                    post_data_dict = {}
-                    self._send_response_stream(response)
+
+                    if isinstance(post_data, bytes):
+                        post_data_str = post_data.decode('utf-8')
+                        post_data_dict = json.loads(post_data_str)
+
+                    response = requests.request(self.command, min_queued_server[1]['url'] + path, params=get_params, data=post_params, stream=post_data_dict.get("stream", False))
+                    self._send_response(response)
                 except Exception as ex:
                     self.add_access_log_entry(event="gen_error",user=self.user, ip_address=client_ip, access="Authorized", server=min_queued_server[0], nb_queued_requests_on_server=que.qsize(),error=ex)                    
                 finally:

requirements.txt

@@ -1,6 +1,8 @@
-configparser
+ascii-colors==0.2.2
-queues
+certifi==2024.2.2
-requests
+charset-normalizer==3.3.2
-urllib3
+configparser==6.0.1
-requests
+idna==3.6
-ascii_colors
+queues==0.6.3
+requests==2.31.0
+urllib3==2.2.1