-
Notifications
You must be signed in to change notification settings - Fork 281
/
agent.py
57 lines (43 loc) · 1.86 KB
/
agent.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
import asyncio
import logging
from typing import Optional
from flyteidl.core.execution_pb2 import TaskExecution
from flytekit import FlyteContextManager, lazy_module
from flytekit.core.type_engine import TypeEngine
from flytekit.extend.backend.base_agent import AgentRegistry, Resource, SyncAgentBase
from flytekit.extend.backend.utils import get_agent_secret
from flytekit.models.literals import LiteralMap
from flytekit.models.task import TaskTemplate
openai = lazy_module("openai")
TIMEOUT_SECONDS = 10
OPENAI_API_KEY = "FLYTE_OPENAI_API_KEY"
class ChatGPTAgent(SyncAgentBase):
name = "ChatGPT Agent"
def __init__(self):
super().__init__(task_type_name="chatgpt")
async def do(
self,
task_template: TaskTemplate,
inputs: Optional[LiteralMap] = None,
) -> Resource:
ctx = FlyteContextManager.current_context()
input_python_value = TypeEngine.literal_map_to_kwargs(ctx, inputs, {"message": str})
message = input_python_value["message"]
custom = task_template.custom
custom["chatgpt_config"]["messages"] = [{"role": "user", "content": message}]
client = openai.AsyncOpenAI(
organization=custom["openai_organization"],
api_key=get_agent_secret(secret_key=OPENAI_API_KEY),
)
logger = logging.getLogger("httpx")
logger.setLevel(logging.WARNING)
try:
completion = await asyncio.wait_for(
client.chat.completions.create(**custom["chatgpt_config"]), TIMEOUT_SECONDS
)
message = completion.choices[0].message.content
outputs = {"o0": message}
return Resource(phase=TaskExecution.SUCCEEDED, outputs=outputs)
except Exception as error_message:
return Resource(phase=TaskExecution.FAILED, message=str(error_message))
AgentRegistry.register(ChatGPTAgent())