feat: better error handling for ollama reverse proxy

This commit is contained in:
Timothy J. Baek 2023-12-13 17:37:29 -08:00
parent 346b0df811
commit 25987fe3c9
1 changed files with 22 additions and 17 deletions

View File

@ -59,27 +59,32 @@ def proxy(path):
else: else:
pass pass
# Make a request to the target server try:
target_response = requests.request( # Make a request to the target server
method=request.method, target_response = requests.request(
url=target_url, method=request.method,
data=data, url=target_url,
headers=headers, data=data,
stream=True, # Enable streaming for server-sent events headers=headers,
) stream=True, # Enable streaming for server-sent events
)
# Proxy the target server's response to the client target_response.raise_for_status()
def generate():
for chunk in target_response.iter_content(chunk_size=8192):
yield chunk
response = Response(generate(), status=target_response.status_code) # Proxy the target server's response to the client
def generate():
for chunk in target_response.iter_content(chunk_size=8192):
yield chunk
# Copy headers from the target server's response to the client's response response = Response(generate(), status=target_response.status_code)
for key, value in target_response.headers.items():
response.headers[key] = value
return response # Copy headers from the target server's response to the client's response
for key, value in target_response.headers.items():
response.headers[key] = value
return response
except Exception as e:
return jsonify({"detail": "Server Connection Error", "message": str(e)}), 400
if __name__ == "__main__": if __name__ == "__main__":