From ee77bd61b142dd63f8070763cf087d0aaf97fdd3 Mon Sep 17 00:00:00 2001 From: tolak Date: Mon, 9 Sep 2024 14:45:30 +0800 Subject: [PATCH 1/5] mech tool for tee ai agent --- .../customs/tee_openai_request/__init__.py | 20 +++ .../customs/tee_openai_request/component.yaml | 22 +++ .../tee_openai_request/tee_openai_request.py | 136 ++++++++++++++++++ 3 files changed, 178 insertions(+) create mode 100644 packages/valory/customs/tee_openai_request/__init__.py create mode 100644 packages/valory/customs/tee_openai_request/component.yaml create mode 100644 packages/valory/customs/tee_openai_request/tee_openai_request.py diff --git a/packages/valory/customs/tee_openai_request/__init__.py b/packages/valory/customs/tee_openai_request/__init__.py new file mode 100644 index 00000000..690f8f4f --- /dev/null +++ b/packages/valory/customs/tee_openai_request/__init__.py @@ -0,0 +1,20 @@ +# -*- coding: utf-8 -*- +# ------------------------------------------------------------------------------ +# +# Copyright 2023-2024 Valory AG +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# ------------------------------------------------------------------------------ + +"""This module contains the tool to request openAI agent in TEE.""" diff --git a/packages/valory/customs/tee_openai_request/component.yaml b/packages/valory/customs/tee_openai_request/component.yaml new file mode 100644 index 00000000..c078a249 --- /dev/null +++ b/packages/valory/customs/tee_openai_request/component.yaml @@ -0,0 +1,22 @@ +name: tee_openai_request +author: valory +version: 0.1.0 +type: custom +description: A tool that runs a prompt against the OpenAI API hosted in TEE. +license: Apache-2.0 +aea_version: '>=1.0.0, <2.0.0' +fingerprint: + __init__.py: 3b70694a1ad269bdb357c4b42dd55bc50911e8f4bf43aa6d8d026925f85f2154 + tee_openai_request.py: b9ddadcf9bc4d398c37d37e8df5ac8e4e74488f6232f03ffadd4fd8fcd9d5086 +fingerprint_ignore_patterns: [] +entry_point: tee_openai_request.py +callable: run +dependencies: + openai: + version: ==1.30.2 + tiktoken: + version: ==0.7.0 + anthropic: + version: ==0.21.3 + google-api-python-client: + version: ==2.95.0 diff --git a/packages/valory/customs/tee_openai_request/tee_openai_request.py b/packages/valory/customs/tee_openai_request/tee_openai_request.py new file mode 100644 index 00000000..28639b27 --- /dev/null +++ b/packages/valory/customs/tee_openai_request/tee_openai_request.py @@ -0,0 +1,136 @@ +# -*- coding: utf-8 -*- +# ------------------------------------------------------------------------------ +# +# Copyright 2023-2024 Valory AG +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# ------------------------------------------------------------------------------ +"""Contains the job definitions""" +import functools +from typing import Any, Dict, Optional, Tuple, Callable + +import anthropic +import googleapiclient +import openai +from tiktoken import encoding_for_model +import requests + +MechResponse = Tuple[str, Optional[str], Optional[Dict[str, Any]], Any, Any] + +def with_key_rotation(func: Callable): + @functools.wraps(func) + def wrapper(*args, **kwargs) -> MechResponse: + # this is expected to be a KeyChain object, + # although it is not explicitly typed as such + api_keys = kwargs["api_keys"] + retries_left: Dict[str, int] = api_keys.max_retries() + + def execute() -> MechResponse: + """Retry the function with a new key.""" + try: + result = func(*args, **kwargs) + return result + (api_keys,) + except anthropic.RateLimitError as e: + # try with a new key again + service = "anthropic" + if retries_left[service] <= 0: + raise e + retries_left[service] -= 1 + api_keys.rotate(service) + return execute() + except openai.RateLimitError as e: + # try with a new key again + if retries_left["openai"] <= 0 and retries_left["openrouter"] <= 0: + raise e + retries_left["openai"] -= 1 + retries_left["openrouter"] -= 1 + api_keys.rotate("openai") + api_keys.rotate("openrouter") + return execute() + except googleapiclient.errors.HttpError as e: + # try with a new key again + rate_limit_exceeded_code = 429 + if e.status_code != rate_limit_exceeded_code: + raise e + service = "google_api_key" + if retries_left[service] <= 0: + raise e + retries_left[service] -= 1 + api_keys.rotate(service) + return execute() + except Exception as e: + return str(e), "", None, None, api_keys + + mech_response = execute() + return mech_response + + return wrapper + +def count_tokens(text: str, model: str) -> int: + """Count the number of tokens in a text.""" + enc = encoding_for_model(model) + return len(enc.encode(text)) + +PREFIX = "tee-openai-" +ENGINES = { + "chat": ["gpt-3.5-turbo", "gpt-4o-2024-08-06"], + "completion": ["gpt-3.5-turbo-instruct"], +} +ALLOWED_TOOLS = [PREFIX + value for values in ENGINES.values() for value in values] +AGENT_URL = "https://wapo-testnet.phala.network/ipfs/QmeUiNKgsHiAK3WM57XYd7ssqMwVNbcGwtm8gKLD2pVXiP" + +@with_key_rotation +def run(**kwargs) -> Tuple[Optional[str], Optional[Dict[str, Any]], Any, Any]: + """Run the task""" + api_key = kwargs["api_keys"]["openai"] + prompt = kwargs["prompt"] + tool = kwargs["tool"] + counter_callback = kwargs.get("counter_callback", None) + if tool not in ALLOWED_TOOLS: + return ( + f"Tool {tool} is not in the list of supported tools.", + None, + None, + None, + ) + + engine = tool.replace(PREFIX, "") + + params = { + "openaiApiKey": api_key, + "chatQuery": prompt, + "openAiModel": engine + } + + # Request to agent contract in TEE + response = requests.get(AGENT_URL, params=params) + + if response.status_code == 200: + json_response = response.json() + if 'message' in json_response: + return json_response['message'], prompt, None, counter_callback + else: + return ( + "The 'message' field is not present in the response.", + None, + None, + None, + ) + else: + return ( + f"Failed to retrieve data: {response.status_code}, {response.text}.", + None, + None, + None, + ) From 0056061c3d2d913a6bb55f2aab9a34628b860ec7 Mon Sep 17 00:00:00 2001 From: tolak Date: Tue, 10 Sep 2024 11:10:19 +0800 Subject: [PATCH 2/5] fix Copyright declaration --- packages/valory/customs/tee_openai_request/__init__.py | 2 +- .../valory/customs/tee_openai_request/tee_openai_request.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/valory/customs/tee_openai_request/__init__.py b/packages/valory/customs/tee_openai_request/__init__.py index 690f8f4f..593767e0 100644 --- a/packages/valory/customs/tee_openai_request/__init__.py +++ b/packages/valory/customs/tee_openai_request/__init__.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- # ------------------------------------------------------------------------------ # -# Copyright 2023-2024 Valory AG +# Copyright 2024 Valory AG # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/packages/valory/customs/tee_openai_request/tee_openai_request.py b/packages/valory/customs/tee_openai_request/tee_openai_request.py index 28639b27..031bcf9f 100644 --- a/packages/valory/customs/tee_openai_request/tee_openai_request.py +++ b/packages/valory/customs/tee_openai_request/tee_openai_request.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- # ------------------------------------------------------------------------------ # -# Copyright 2023-2024 Valory AG +# Copyright 2024 Valory AG # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. From abdf54de986923b292e9108cd07332ab6cdba4ec Mon Sep 17 00:00:00 2001 From: tolak Date: Tue, 10 Sep 2024 11:20:18 +0800 Subject: [PATCH 3/5] force type cast --- .../valory/customs/tee_openai_request/tee_openai_request.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/valory/customs/tee_openai_request/tee_openai_request.py b/packages/valory/customs/tee_openai_request/tee_openai_request.py index 031bcf9f..62d51be4 100644 --- a/packages/valory/customs/tee_openai_request/tee_openai_request.py +++ b/packages/valory/customs/tee_openai_request/tee_openai_request.py @@ -119,7 +119,7 @@ def run(**kwargs) -> Tuple[Optional[str], Optional[Dict[str, Any]], Any, Any]: if response.status_code == 200: json_response = response.json() if 'message' in json_response: - return json_response['message'], prompt, None, counter_callback + return str(json_response['message']), prompt, None, counter_callback else: return ( "The 'message' field is not present in the response.", From 3ee1ff90b52b0b5c67e02d4213310d80dc22e9e0 Mon Sep 17 00:00:00 2001 From: tolak Date: Tue, 10 Sep 2024 11:25:01 +0800 Subject: [PATCH 4/5] regenerate fingerprint --- .../valory/customs/tee_openai_request/component.yaml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/packages/valory/customs/tee_openai_request/component.yaml b/packages/valory/customs/tee_openai_request/component.yaml index c078a249..e634fa63 100644 --- a/packages/valory/customs/tee_openai_request/component.yaml +++ b/packages/valory/customs/tee_openai_request/component.yaml @@ -6,17 +6,17 @@ description: A tool that runs a prompt against the OpenAI API hosted in TEE. license: Apache-2.0 aea_version: '>=1.0.0, <2.0.0' fingerprint: - __init__.py: 3b70694a1ad269bdb357c4b42dd55bc50911e8f4bf43aa6d8d026925f85f2154 - tee_openai_request.py: b9ddadcf9bc4d398c37d37e8df5ac8e4e74488f6232f03ffadd4fd8fcd9d5086 + __init__.py: bafybeibaxl4f33swhl5vk63v75ssw5mk3mhlgcwarok242ivgk5hphwpey + tee_openai_request.py: bafybeibhmihak7bexecww7mxxrtbsl54q77zl2t2r722u44yfmxv25j6bq fingerprint_ignore_patterns: [] entry_point: tee_openai_request.py callable: run dependencies: - openai: - version: ==1.30.2 - tiktoken: - version: ==0.7.0 anthropic: version: ==0.21.3 google-api-python-client: version: ==2.95.0 + openai: + version: ==1.30.2 + tiktoken: + version: ==0.7.0 From 45ca6483a60c4ab3e399a8d767cdca7cbb381013 Mon Sep 17 00:00:00 2001 From: tolak Date: Wed, 11 Sep 2024 21:06:32 +0800 Subject: [PATCH 5/5] fix package lock check --- packages/packages.json | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/packages.json b/packages/packages.json index 19c85fde..0025aa1d 100644 --- a/packages/packages.json +++ b/packages/packages.json @@ -31,6 +31,7 @@ "custom/victorpolisetty/dalle_request/0.1.0": "bafybeieqqtd6gtlry7vheix54nj3ok4cag3uy47yoxlufhi6y3u5i6doti", "custom/jhehemann/prediction_sentence_embeddings/0.1.0": "bafybeifyyb2wpa77tl7a7fs3fabns45llivhgccbnrpupubojmq2fwe4si", "custom/gnosis/ofv_market_resolver/0.1.0": "bafybeiemvfq6uxiz3wvdplnxg7wloy6siuggejerlkfkchks6ytgk27uqa", + "custom/valory/tee_openai_request/0.1.0": "bafybeictmezaorzxelsy4dztbxh5n2343zio3rk6vo7wc5lptxlobhdnku", "protocol/valory/acn_data_share/0.1.0": "bafybeih5ydonnvrwvy2ygfqgfabkr47s4yw3uqxztmwyfprulwfsoe7ipq", "protocol/valory/websocket_client/0.1.0": "bafybeifjk254sy65rna2k32kynzenutujwqndap2r222afvr3zezi27mx4", "contract/valory/agent_mech/0.1.0": "bafybeiah6b5epo2hlvzg5rr2cydgpp2waausoyrpnoarf7oa7bw33rex34",