nova-api/api/proxies.py

126 lines
3.6 KiB
Python

"""This module contains the Proxy class, which represents a proxy."""
import os
import httpx
import socket
import asyncio
from sockslib import socks
from rich import print
from dotenv import load_dotenv
load_dotenv()
class Proxy:
"""Represents a proxy. The type can be either http, https, socks4 or socks5."""
def __init__(self,
proxy_type: str='http',
host: str='127.0.0.1',
port: int=8080,
username: str=None,
password: str=None
):
self.proxy_type = proxy_type
self.ip_address = socket.gethostbyname(host)
self.host = host
self.port = port
self.username = username
self.password = password
@property
def auth(self):
"""Returns the authentication part of the proxy URL, if the proxy has a username and password."""
return f'{self.username}:{self.password}@' if all([self.username, self.password]) else ''
@property
def protocol(self):
"""Makes sure the hostnames are resolved correctly using the proxy.
See https://stackoverflow.com/a/43266186
"""
return self.proxy_type# + 'h' if self.proxy_type.startswith('socks') else self.proxy_type
@property
def proxies(self):
"""Returns a dictionary of proxies, ready to be used with the requests library or httpx.
"""
url = f'{self.protocol}://{self.auth}{self.host}:{self.port}'
proxies_dict = {
'http://': url.replace('https', 'http') if self.proxy_type == 'https' else url,
'https://': url.replace('http', 'https') if self.proxy_type == 'http' else url
}
return proxies_dict
@property
def url(self):
"""Returns the proxy URL."""
return f'{self.protocol}://{self.auth}{self.host}:{self.port}'
def __str__(self):
return f'{self.proxy_type}://{len(self.auth) * "*"}{self.host}:{self.port}'
def __repr__(self):
return f'<Proxy type={self.proxy_type} host={self.host} port={self.port} username={self.username} password={len(self.password) * "*"}>'
active_proxy = Proxy(
proxy_type=os.getenv('PROXY_TYPE', 'http'),
host=os.getenv('PROXY_HOST', '127.0.0.1'),
port=int(os.getenv('PROXY_PORT', 8080)),
username=os.getenv('PROXY_USER'),
password=os.getenv('PROXY_PASS')
)
def activate_proxy() -> None:
socks.set_default_proxy(
proxy_type=socks.PROXY_TYPES[active_proxy.proxy_type.upper()],
addr=active_proxy.host,
port=active_proxy.port,
username=active_proxy.username,
password=active_proxy.password
)
socket.socket = socks.socksocket
def check_proxy():
"""Checks if the proxy is working."""
resp = httpx.get(
url='https://echo.hoppscotch.io/',
timeout=20,
proxies=active_proxy.proxies
)
resp.raise_for_status()
return resp.ok
async def check_api():
model = 'gpt-3.5-turbo'
messages = [
{
'role': 'user',
'content': 'Explain what a wormhole is.'
},
]
headers = {
'Content-Type': 'application/json',
'Authorization': 'Bearer ' + os.getenv('CLOSEDAI_KEY')
}
json_data = {
'model': model,
'messages': messages,
'stream': True
}
async with httpx.AsyncClient(timeout=20) as client:
async with client.stream("POST", 'https://api.openai.com/v1/chat/completions', headers=headers, json=json_data) as response:
response.raise_for_status()
async for chunk in response.aiter_text():
print(chunk.strip())
if __name__ == '__main__':
asyncio.run(check_api())