# How To Develop Agent Plugin Service Efficiently? (Revision)

## Outline
- How flyteplugins work?
- How flyteplugins agent work?
- How to setup Agent Service in Development Mode?
- How to Develop your agent efficiently?

## How flyteplugins work?
This is a section about how flyteplugins work, and why we use agent in flyteplugins.
Flyteplugins plays a really important role in Flyte.
It makes the task with specific APIs or modules doesn't need to create a pod to execute it, which saves lots of time and resources.
Without flyteplugins, every task's lifecycle will be like this.

However, there are some APIs are used frequently, we can reduce the lifecycle to this.

This is why flyteplugins exist, it make the task use less resources and faster.
(Since we don't need to create the pod.)

That's why flyteplugins exist, it make the lifecycle like this.
Before the Agent Framework exists, programmers have to do the followings to create flyteplugins:
1. Create a plugin in `flytekit/plugins` directory.
2. Specify the task config in (1), which will be passed to task metadata for flytepropeller
3. Write the plugin handler in `flyteplugins` repo.
Note: we write python code in (1) and write golang code in (3).

Let me briefly explain how flyteplugins work.
The flyteplugins is a part of flytepropeller, so the overhead of execute flyteplugins will happen at the flytepropeller.
Note: flytepropeller is the k8s operator of flyte.

Here is a picture about how it works.

This is already an amazing mechanism, which can create lots of posiibilities.
However, it's not friendly to most of data scientist and ML engineers, since they have to learn how to write golang codes, and even study the flytepropeller architecture, which is time-consuming and hard.

## How flyteplugins agent work?
Hence, we introduce an agent framework to make it more friendly.
The agent server is a stateless server, written by python, it brings 2 main benefits.

1. The overhead to execute plugin task transfer to the agent server, now flytepropeller only has to monitor the status of the task, which is cost less resouces.
2. Plugins become far more easier to be created by users.


Note: We can have lots of agent server to handle lots of requests.

## How to setup Agent Service in Development Mode?
Let's develop the powerful agent service!
1. Use the dev mode through flytectl.

flytectl demo start --dev

2. Start the agent grpc server.

pyflyte serve

3. Set the config in the yaml file (Bigquery for example)

cd flyte
vim ./flyte-single-binary-local-dev.yaml

tasks:
  task-plugins:
    enabled-plugins:
      - agent-service
      - container
      - sidecar
      - K8S-ARRAY
    default-for-task-types:
      - bigquery_query_job_task: agent-service
      - container: container
      - container_array: K8S-ARRAY

plugins:
  # Registered Task Types
  agent-service:
    supportedTaskTypes:
      - bigquery_query_job_task
    defaultAgent:
      endpoint: "dns:///localhost:8000" # your grpc agent server port
      insecure: true
      timeouts:
        GetTask: 100s
      defaultTimeout: 10s

4. Start the Flyte server with config yaml file

flyte start --config ./flyte-single-binary-local-dev.yaml

## How to Develop your agent efficiently?
Starting with the development of an agent can be challenging, so I recommend split the process step by step.
Here's the recommended process.
1. Study how `task.py` in plugin works to pass parameters to agent
2. Study how local agent works
3. How to test it efficiently in remote environment

### How task.py in plugin pass parameters to agent
Take BigQuery plugin as an example again.
The [get_custom](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py#L70-L79) function will create a dictionary object in [task.py](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py) will be passed to [task_template.custom](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py#L70-L72) in [agent.py](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py).

Here's a more comprehensive diagram about how parameters pass to `agent` from `task.py`.

### How local agent works
Using BigQuery as an example.
You can see the code [here](https://github.com/flyteorg/flytekit/blob/master/flytekit/extend/backend/base_agent.py#L168-L181) to know how local agent works.
Let's take BigQuery as example again.

DogeCoinDataset = Annotated[StructuredDataset, kwtypes(hash=str, size=int, block_number=int)]

bigquery_task_templatized_query = BigQueryTask(
    name="sql.bigquery.w_io",
    # Define inputs as well as their types that can be used to customize the query.
    inputs=kwtypes(version=int),
    output_structured_dataset_type=DogeCoinDataset,
    task_config=BigQueryConfig(ProjectID="flyte"),
    query_template="SELECT * FROM `bigquery-public-data.crypto_dogecoin.transactions` WHERE version = @version LIMIT 10;",
)


@task
def convert_bq_table_to_pandas_dataframe(sd: DogeCoinDataset) -> pd.DataFrame:
    return sd.open(pd.DataFrame).all()


@workflow
def full_bigquery_wf(version: int) -> pd.DataFrame:
    sd = bigquery_task_templatized_query(version=version)
    return convert_bq_table_to_pandas_dataframe(sd=sd)

if __name__ == "__main__":
    full_bigquery_wf(version=1)

To test your BigQuery code locally, execute:

python bigquery example.py

This will process the code through the local agent using the [AsyncAgentExecutorMixin](https://github.com/flyteorg/flytekit/blob/master/flytekit/extend/backend/base_agent.py#L158-L167) class, and use the code [agent.py](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py#L48-L114).

Before moving to a remote environment, I strongly advise testing your code locally.

### How to test it efficiently in remote environment
To execute your agent remotely, use:
`pyflyte run --remote` or `pyflyte register`

I recommend starting by registering and invoking your agent tasks in FlyteConsole.
This initial step doesn’t require building an image, as you’ll only be invoking the web API, without initiating a pod for task execution.
Take ChatGPT plugin for example.
You can write the code like this to invoke the task.
For example, this is my sample code.

from flytekit import task, workflow, ImageSpec
from flytekitplugins.chatgpt import ChatGPTTask

chatgpt_job = ChatGPTTask(
    name="chatgpt",
    config={
        "openai_organization": "org-NayNG68kGnVXMJ8Ak4PMgQv7",
        "chatgpt_conf": {
            "model": "gpt-3.5-turbo",
            "temperature": 0.7,
        },
    },
)

First, let's register the workflow and task to the flyte console.

pyflyte register chatgpt_example.py

The console will have these logs.

Loading packages ['chatgpt_example'] under source root /mnt/c/code/dev/example/plugins
Successfully serialized 1 flyte objects
[✔] Registration chatgpt type TASK successful with version iopF5N9M7ABb6gEe3Chpsg==
Successfully registered 1 entities

You can invoke the task in the flyte console!

Ultimately, you must test it with a simple workflow and define the image to mock the real world use case.
There are two methods for this: using ImageSpec or a Dockerfile.
This [PR](https://github.com/flyteorg/flytekit/pull/1822) provides guidance on both approaches.
Here's a briefly example.

@task()
def t1(s: str) -> str:
    s = "Repsonse: " + s
    return s

@workflow
def wf() -> str:
    message = chatgpt_job(message="hi")
    return t1(s=message)

if __name__ == "__main__":
    print(wf())

## Additional Resources
If you want to study more about flyteplugins and agent, you can refer these.

1. Flyte School: Enrich your AI pipelines - A deep dive into Flyte plugins
https://www.youtube.com/watch?v=ah8Q5mSeikE&t=1145s&ab_channel=Union-ai
2. Writing Agents in Python
https://docs.flyte.org/projects/cookbook/en/latest/auto_examples/development_lifecycle/agent_service.html
3. Demystifying Flyte Agents
https://youtu.be/nD98GQ-pyAE?si=mI_s7DG8LBMt92Zp

## Special Thanks
- Kevin Su, an amazing flyte maintainer and a good mentor, he taught me lots of things about flyte, help me go through the process from zero to hero in Flyte.
- Yi Cheng Lu, spend times discuss flyteplugins, flytepropeller and agent with me, which helps me have a deeper understanding on those.
- Da Yi Wu, ask me how flyteplugins and how agent works, give me a chance to practice to explain how it works.
- David Espejo, the reviewer of this article, gave me tons of awesome advices, I added lots of diagram because of him.