forked from Yifan-Song793/RestGPT
-
Notifications
You must be signed in to change notification settings - Fork 0
/
run_tmdb.py
79 lines (59 loc) · 2.52 KB
/
run_tmdb.py
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
import os
import json
import logging
import time
from dotenv import load_dotenv
from llm_commons.langchain.proxy import ChatOpenAI
import yaml
from langchain.requests import Requests
from gen_ai_hub import GenAIHubProxyClient
from gen_ai_hub.proxy import set_proxy_version
from utils import reduce_openapi_spec, ColorPrint, MyRotatingFileHandler
from model import RestGPT
logger = logging.getLogger()
load_dotenv(dotenv_path='.env', verbose=True, override=False)
set_proxy_version('gen-ai-hub')
proxy_client=GenAIHubProxyClient()
def run(query, api_spec, requests_wrapper, simple_parser=False):
llm = ChatOpenAI(proxy_model_name="gpt-4", proxy_client=proxy_client)
#llm = OpenAI(model_name="text-davinci-003", temperature=0.0, max_tokens=256)
# llm = OpenAI(model_name="gpt-3.5-turbo-0301", temperature=0.0, max_tokens=256)
rest_gpt = RestGPT(llm, api_spec=api_spec, scenario='tmdb', requests_wrapper=requests_wrapper, simple_parser=simple_parser)
logger.info(f"Query: {query}")
start_time = time.time()
rest_gpt.run(query)
logger.info(f"Execution Time: {int(time.time() - start_time)} seconds")
def main():
config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)
os.environ["OPENAI_API_KEY"] = config['openai_api_key']
os.environ["TMDB_ACCESS_TOKEN"] = config['tmdb_access_token']
print(config['tmdb_access_token'])
log_dir = os.path.join("logs", "restgpt_tmdb")
if not os.path.exists(log_dir):
os.mkdir(log_dir)
file_handler = MyRotatingFileHandler(os.path.join(log_dir, f"tmdb.log"), encoding='utf-8')
logging.basicConfig(
format="%(message)s",
handlers=[logging.StreamHandler(ColorPrint()), file_handler],
)
logger.setLevel(logging.INFO)
with open("specs/tmdb_oas.json") as f:
raw_tmdb_api_spec = json.load(f)
api_spec = reduce_openapi_spec(raw_tmdb_api_spec, only_required=False)
access_token = os.environ["TMDB_ACCESS_TOKEN"]
headers = {
'Authorization': f'Bearer {access_token}'
}
requests_wrapper = Requests(headers=headers)
queries = json.load(open('datasets/tmdb.json', 'r'))[:10]
queries = [item['query'] for item in queries]
for idx, query in enumerate(queries, 1):
try:
print('#' * 20 + f" Query-{idx} " + '#' * 20)
run(query, api_spec, requests_wrapper, simple_parser=False)
except Exception as e:
print(f"TMDB Query: {query}\nError: {e}")
finally:
file_handler.doRollover()
if __name__ == '__main__':
main()