Add SyncAgentBase

[pingsutw]

Tracking issue

flyteorg/flyte#3936

Why are the changes needed?

Not all the external services have async API (create/get/delete). For example, Openai API returns the response after the job is completed. Thus, we need to add another API (do) in the agent to run the task synchronously.

What changes were proposed in this pull request?

• Add the SyncAgentBase class for the sync task (ChatGPT task).
• Rename AgentBase to AsyncAgentBase
• Add SyncAgentExecutorMixin for executing the openAI task locally.

class SyncAgentBase(AgentBase):
    @abstractmethod
    def do(self, task_template: TaskTemplate, inputs: Optional[LiteralMap], **kwargs) -> Resource:
        pass

Mock OpenAI agent

class MockAsyncOpenAIAgent(SyncAgentBase):
    name = "mock async openAI Agent"

    def __init__(self):
        super().__init__(task_type_name="async_openai")

    async def do(self, task_template: TaskTemplate, inputs: LiteralMap = None, **kwargs) -> Resource:
        assert inputs.literals["a"].scalar.primitive.integer == 1
        return Resource(phase=TaskExecution.SUCCEEDED, outputs={"o0": 1})

I also update the async agent API.

Before:

class AsyncAgentBase(ABC):
    def __init__(self, task_type: str, **kwargs):
        self._task_type = task_type

    def create(
        self,
        output_prefix: str,
        task_template: TaskTemplate,
        inputs: typing.Optional[LiteralMap] = None,
        **kwargs,
    ) -> CreateTaskResponse:  
        ...
        raise CreateTaskResponse(resource_meta=json.dumps(asdict(metadata)).encode("utf-8"))

    def get(self, resource_meta: bytes, **kwargs) -> GetTaskResponse:
        metadata = Metadata(**json.loads(resource_meta.decode("utf-8")))
        raise GetTaskResponse(....)

    def delete(self, resource_meta: bytes, **kwargs) -> DeleteTaskResponse:
        metadata = Metadata(**json.loads(resource_meta.decode("utf-8")))
        raise NotImplementedError

After:

class AsyncAgentBase(AgentBase, typing.Generic[T]):
    @abstractmethod
    def create(self, task_template: TaskTemplate, inputs: Optional[LiteralMap], **kwargs) -> ResourceMeta:
        raise NotImplementedError

    @abstractmethod
    def get(self, resource_meta: ResourceMeta, **kwargs) -> Resource:
        raise NotImplementedError

    @abstractmethod
    def delete(self, resource_meta: ResourceMeta, **kwargs):
        raise NotImplementedError

How was this patch tested?

local/remote

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

https://github.com/flyteorg/flyte/pull/4874/files

Docs link

NA

[codecov]

Codecov Report

Attention: Patch coverage is 86.18926% with 54 lines in your changes are missing coverage. Please review.

Project coverage is 85.96%. Comparing base (90324e0) to head (f098681).

Files	Patch %	Lines
flytekit/extend/backend/agent_service.py	58.90%	24 Missing and 6 partials ⚠️
flytekit/extend/backend/base_agent.py	90.24%	5 Missing and 7 partials ⚠️
flytekit/core/type_engine.py	58.33%	2 Missing and 3 partials ⚠️
.../flytekit-airflow/flytekitplugins/airflow/agent.py	84.21%	3 Missing ⚠️
...gins/flytekit-spark/flytekitplugins/spark/agent.py	87.50%	2 Missing ⚠️
...lytekit-bigquery/flytekitplugins/bigquery/agent.py	95.00%	1 Missing ⚠️
...tekit-snowflake/flytekitplugins/snowflake/agent.py	95.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2146      +/-   ##
==========================================
+ Coverage   85.86%   85.96%   +0.10%     
==========================================
  Files         314      312       -2     
  Lines       24042    24025      -17     
  Branches     3643     3665      +22     
==========================================
+ Hits        20643    20653      +10     
+ Misses       2759     2722      -37     
- Partials      640      650      +10     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.