File size: 3,335 Bytes
3943768 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import os
import pickle
import requests
import json
def execute_function_on_server(host: str, port: int, function_name: str, args: tuple, kwargs: dict, use_disk: bool,
use_pickle: bool, function_api_key: str):
url = f"http://{host}:{port}/execute_function/"
payload = {
"function_name": function_name,
"args": args,
"kwargs": kwargs,
"use_disk": use_disk,
"use_pickle": use_pickle,
}
headers = {
"Authorization": f"Bearer {function_api_key}"
}
response = requests.post(url, json=payload, headers=headers)
if response.status_code == 200:
return response.json()
else:
return {"error": response.json()["detail"]}
def read_result_from_disk(file_path: str, use_pickle: bool, verbose=False):
if verbose:
print(f"Size of {file_path} is {os.path.getsize(file_path)}")
try:
if use_pickle:
with open(file_path, "rb") as f:
result = pickle.load(f)
else:
with open(file_path, "r") as f:
result = json.load(f)
except Exception as e:
raise IOError(f"Error reading file {file_path}: {e}")
finally:
try:
os.remove(file_path)
except OSError as e:
print(f"Error deleting file {file_path}: {e}")
return result
def call_function_server(host, port, function_name, args, kwargs, use_disk=False, use_pickle=False,
function_api_key='EMPTY', verbose=False):
execute_result = execute_function_on_server(host, port, function_name, args, kwargs, use_disk, use_pickle,
function_api_key)
if "error" in execute_result:
raise RuntimeError(execute_result['error'])
else:
if use_disk or use_pickle:
file_path = execute_result["file_path"]
result_from_disk = read_result_from_disk(file_path, use_pickle, verbose=verbose)
return result_from_disk
else:
return execute_result["result"]
def get_data_h2ogpt(file_path, verbose=False, is_url=False, **kwargs):
"""
Simple function for Open Web UI
"""
function_server_host = os.getenv('H2OGPT_FUNCTION_SERVER_HOST', '0.0.0.0')
function_server_port = int(os.getenv('H2OGPT_FUNCTION_SERVER_PORT', '5002'))
function_api_key = os.getenv('H2OGPT_FUNCTION_SERVER_API_KEY', 'EMPTY')
# could set other things:
# https://github.com/h2oai/h2ogpt/blob/d2fa3d7ce507e8fb141c78ff92a83a8e27cf8b31/src/gpt_langchain.py#L9498
simple_kwargs = kwargs
if is_url:
simple_kwargs.update(dict(filei=None, url=file_path, text=None))
file_path = None
function_name = 'path_to_docs'
use_disk = False
use_pickle = True
sources = call_function_server(function_server_host,
function_server_port,
function_name,
(file_path,),
simple_kwargs,
use_disk=use_disk, use_pickle=use_pickle,
function_api_key=function_api_key,
verbose=verbose)
known_type = len(sources) > 0
return sources, known_type
|