Resource & Resource Provider Design

[wenjin272]

1. Introduce

This doc will introduce Resource & Resource Provider, which is an unified abstraction in Flink-Agents for all kinds of resources such as models, tools, vector_stores and so on. This abstraction can make the Runtime of Flink-Agents not need be aware of the difference among various resources, and supports lazy initialization and reuse of resources.

Sections 2, 3, and 4 will cover:

• The concept and interface of Resource & Resource Provider
• Example about how to defined specific Resource & Resource Provider, and use them in an agent.
• Explanation of how does Resource & Resource Provider work internally, includes compile stage and runtime stage.

2. Concept & Interface

Resource

Resource is the abstraction of object defined in user codes, like model, tools and so on. It

• provides unified type declaration for various built-in or user-defined resource
• represented as resource with different types in runtime

class ResourceType(Enum):
    chat_model = "chat_model"
    embedding = "embedding_model"
    prompt = "prompt"
    tool = "tool"
    vector_store = "vector_store"
    mcp_server = "mcp_server"

class Resource:
    """Base interface for user defined resources, like model, prompt, tool."""
    name: str
    type: ResourceType

Resource Provider

Resource Provider is the factory to create Resource object, it

• supports lazy initialization and reuse of Resource object.
• should be serializable to support use in other process.
• represented as resource provider provides specific type resource in runtime

class ResourceProvider(BaseModel, ABC):
    """Resource provider carries resource meta to construct
     Resource object in runtime.
     
     Attributes:
    ----------
    name : str
        The name of the resource
    module : str
        The module name of the resource.
    clazz : str
        The class name of the resource.
    kwargs : Dict[str, Any]
        The initialization arguments of the resource.
    """
    name: str
    module: str
    clazz: str
    kwargs: Dict[str, Any]

    @abstractmethod
    def provide(self) -> Resource:
        module = importlib.import_module(self.module)
        cls = getattr(module, self.clazz)
        return cls(**self.kwargs)

3. API

This section will introduce how can user define their own Resource by implements Resource interface, and decalare a ResourceProvider in agent to use it.
Taking chat_model as an example.

• Flink-Agents will provide the ChatModel interface, which implements Resource to decalare it is Resource, and provides chat model specific interface

class ChatModel(Resource, ABC):
    """Abstract base class defining the interface for chat models."""

    tools: List[str] = []

    @abstractmethod
    def chat(self, messages: List[ChatMessage]) -> ChatMessage | List[ToolCall]:

    @abstractmethod
    def bind_tools(self, tools: List[ToolDescriptor]) -> None:

• User can implement ChatModel interface to defined own chat model

class UserDefinedChatModel(ChatModel):
    model: str
    host: str

    def chat(self, messages: List[ChatMessage]) -> ChatMessage | List[ToolCall]:
      """user defined logic"""

    def bind_tools(self, tools: List[ToolDescriptor]) -> None:
      """user defined logic"""

• User can add ResourceProvider in agent to use Resource in actions.

class MockChatAgent(Agent):
  
    @model
    @staticmethod
    def llm() -> Tuple[Type[ChatModel], Dict[str, Any]]:
        return UserDefinedChatModel, 
               {'model': 'qwen2.5:7b', 
                'host': 'http://localhost:9999'}

• • The @model decorator indicates that the llm method defines a model resource provider.
  • The method returns a factory class and a dictionary of arguments, which the framework will use to create the model resource.
  • The method name llm will be the name of ResourceProvider and Resource.

• User can get Resource instance by name and type in action, and use it.

def call_chat_model(
    event: ChatModelEvent | ToolEvent, state: State, ctx: RunnerContext
) -> None:
  
  model = ctx.get_resource(event.model, ResourceType.chat_model)
  response = model.chat(event.messages)
  ctx.send_event(ChatModelResponseEvent(request=event, response=response))

• we can provide some syntactic sugar like get_model, get_tool to improve usability

4. How does Resource & Resource Provider work internally

This section will introduce how does Flink-Agents framework use Resource & Resource Provider, includes:

• • In compile stage, convert user defined factory function in agent to Resource Provider object in agent plan.
  • In runtime stage, get resource provider from agent plan by name and resource type
  • In runtime stage, create resource by resource provider and supports reuse in Context wise.

• Convert factory function defined in agent to a two-level dict 'type->name->ResourceProvider' in agent plan.

class AgentPlan(BaseModel):
    """Agent plan managing actions, models, and tools."""

    actions: Dict[str, action]
    resource_providers: Dict[ResourceType, Dict[str, ResourceProvider]]
    actions_by_event: Dict[Type[Event], List[Action]]
    resources: Dict[ResourceType, Dict[str, Resource]]

• There will be an agent Runner Context to manage resources, events, state in agent execution. The agent plan will be an property of Context.
• In action, user will get resource from Context, Context will return resource from agent plan.

class RunnerContextImpl(RunnerContext):
    __agent_plan: AgentPlan
    __id: UUID
    state: State
    events: deque[Event]
    
    """get resource according to name and type."""
    def get_resource(self, name: str, resource_type: ResourceType) -> Resource:
        return self.__agent_plan.get_resource(name, resource_type)

• agent plan will return cached resource directly or crate resource object by resource provider and return it.

class AgentPlan(BaseModel):
    """Agent plan managing actions, models, and tools."""

    actions: Dict[str, action]
    resource_providers: Dict[ResourceType, Dict[str, ResourceProvider]]
    actions_by_event: Dict[Type[Event], List[Action]]
    resources: Dict[ResourceType, Dict[str, Resource]]

    def get_resource(name: str, resource_type: ResourceType) -> Resource:
        if name not in self.resources[resource_type][name]:
            resource_provider = self.__agent.get_resource_provider(name, resource_type)
            self.resources[resource_type][name] = resource_provider.provide()
        return self.resources[resource_type][name]

5. An implementation of Resource & Resource Provider for special resource

For some kinds of resources, user may prefer define resource directly in Agent, rather than define a factory function.
Take tool as an example, user defines a python function, and the function itself is a Resource. Require user to define a factory function which returns the tool function is unreasonable.

class MockChatAgent(Agent):
  
    @tool
    @staticmethod
    def add(a: int, b: int) -> int:
        print(f"calling `add`: ({a}, {b})")
        return a + b

For these kinds of resources, we provide SerializableResource & SerializableResourceProvider interface

SerializableResource

SerializableResource object should be serializable

class SerializableResource(Resource, BaseModel):
    """Resource which is serializable"""

SerializableResourceProvider

SerializableResourceProvider will return the SerializableResource object directly.

class SerializableResourceProvider(ResourceProvider):
    """Resource Provider which persist Resource Object.
    Attributes:
    ----------
    module : str
        module of the resource class.
    classname : str
        class name of the resource class
    serialized : Dict[str, Any]
        serilaized resource object
    resource : Optional[SerializableResource]
        SerializableResource object
    """
    
    module: str
    classname: str
    serialized: Dict[str, Any]
    resource: Optional[SerializableResource] = None

    def provide(self) -> Resource:
        if self.resource is None:
            module = importlib.import_module(self.module)
            clazz = getattr(module, self.classname)
            self.resource = clazz(**self.serialized)
        return self.resource

6. Resource Access Cross Language

This section will introduce the use case which may access resource cross language, and talk about the brief thought of how to solve it. In our early stage, we don't need implement this.

What is Resource Access Cross Language

Flink-Agents will provide two API, python and java, and user may mix it to reuse some action or resource defined previous. And this may lead

• Use Java Resource object in Python Action
• Use Python Resource object in Java Action
  

How to solve this problem (overall)

• Use pemja
  • pemja (https://github.com/alibaba/pemja) is a cross language call framework based on FFI, and is used by PyFlink. It supports Java call Python, and Java call Python call Java. Because our dispatcher is running in JVM, so pemja can satisfy our requirement.
• The python API and java API should provide identical interface for each kind of Resource, includes ChatModel, Tool, VectorStore and so on.
• And, for each kind of Resource, there will be an correspond Wrapper class in Python and Java. Take ChatModel as an example.
  • python wrapper class

from pemja import findClass
JavaChatModelClass = findClass("flink.agents.api.ChatModel")
class JavaChatModel(ChatModel):
    tools: List[str] = []
    chat_model: JavaChatModelClass

    def __init__(self, clazz: Str, *args) -> None:
        """
        Args:
          clazz: fully qualified name of java resource
          args: init parameters of java resource
        """
        JavaChatModelImplementationClass = findClass(clazz);
        self.chat_modle=JavaChatModelImplementationClass(*args);

    def chat(self, messages: List[ChatMessage]) -> ChatMessage | List[ToolCall]:
        return self.chat_model.chat(messages)

    @abstractmethod
    def bind_tools(self, tools: List[ToolDescriptor]) -> None:
        self.chat_model.bindTools(tools)

• • java wrapper class

import pemja.core.object.PyObject;
public class PyChatModel extends ChatModel {
    private final PythonInterpreter interpreter
    private final String module;
    private final String clazz;
    private PyObject chatModel;
    
    public PyChatModel(PythonInterpreter interpreter, 
                       String module, 
                       String clazz,
                       Map<String, Object> kwargs) {
        this.module = module;
        this.clazz = clazz;
        this.interpreter = interpreter;
        interpreter.exec(String.format("import %s", module));
        this.chatModel = interpreter.invoke(String.format("%s.__init__"), kwargs)
    }
    
    public ChatMessage chat(List<ChatMessage> messages) {
        return this.chatModel.invokeMethod("chat", messages)
    }

    public void bindTools(List<ToolDescriptor> tools) {
        this.chatModel.invokeMethod("bind_tools", tools)
    }
}

How to use other language defined resource in Agent

User may already implements some resource in Python or Java, and want reuse it when create and Agent use Java or Python API
Take define Java ChatModel in Python API as an example

• User declare a java ChatModel in Agent just like declare a python ChatModel. The difference is
  • The ChatModel class must be JavaChatModel
  • The arguments dict must contain the Java ChatModel fully qualified name

class MockChatAgent(Agent):
    @model
    @staticmethod
    def llm() -> Tuple[Type[ChatModel], Dict[str, Any]]:
    return JavaChatModel, 
               {'clazz': 'org.apache.flink.agents.example.UserDefinedChatModel'
                'model': 'qwen2.5:7b', 
                'host': 'http://localhost:9999'}

• There will be no difference in compile stage between using Python Resource and Java Resource.
• When get this resource in action, according to the actually executor (Pyhon or Java) of action , there will be two case
  • execute an python action: we need create python JavaChatModel object according to the info in ResourceProvider
  • execute an java action: we create java ChatModel object directly according to the info in ResourceProvider

Overall, we just store the cross language resource meta info and init arguments in compile stage, and decide whether need to create Wrapper class in runtime stage.

[istavrak]

Great write-up and design! The unified declaration between built-in and user-defined resource types is a great proposal for the extendability of the Flink-Agents. At the same time, a list of the most used resource types in agentic pipelines could already be part of the built-in resource types as interfaces.

[istavrak]

The use of pemja seems to be an important enabler for the cross-language compatibility of steps in the agents. Did you consider any alternative designs or frameworks to achieve that?

[xintongsong]

We are considering pemja mostly because 1) we are familiar with it (it's creaged by @HuangXingBo) and 2) pyflink uses it. Do you have any alternative to suggest?