Customize an application template

In Develop the application, we used a prebuilt template (i.e. reasoning_engines.LangchainAgent) for developing an application. In this section, we go through the steps to customize your own application template. This might be useful if you have needs that go beyond what the prebuilt template provides.

An application template in Reasoning Engine is defined as a Python class. To give an example, the following Python code is an example of a LangChain application that is deployable on Vertex AI (you can give the CLASS_NAME variable a value such as MyAgent):

from typing import Callable, Sequence

class CLASS_NAME:
    def __init__(
            self,
            model: str,
            tools: Sequence[Callable],
            project: str,
            location: str,
        ):
        self.model_name = model
        self.tools = tools
        self.project = project
        self.location = location

    def set_up(self):
        """All unpickle-able logic should go here.

        The .set_up() method should not be called for an object that is being
        prepared for deployment.
        """
        import vertexai
        from langchain_google_vertexai import ChatVertexAI
        from langchain.agents import AgentExecutor
        from langchain.agents.format_scratchpad.tools import format_to_tool_messages
        from langchain.agents.output_parsers.tools import ToolsAgentOutputParser
        from langchain.tools.base import StructuredTool
        from langchain_core import prompts

        vertexai.init(project=self.project, location=self.location)

        prompt = {
            "input": lambda x: x["input"],
            "agent_scratchpad": (
                lambda x: format_to_tool_messages(x["intermediate_steps"])
            ),
        } | prompts.ChatPromptTemplate.from_messages([
            ("user", "{input}"),
            prompts.MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

        llm = ChatVertexAI(model_name=self.model_name)
        if self.tools:
            llm = llm.bind_tools(tools=self.tools)

        self.agent_executor = AgentExecutor(
            agent=prompt | llm | ToolsAgentOutputParser(),
            tools=[StructuredTool.from_function(tool) for tool in self.tools],
        )

    def query(self, input: str):
        """Query the application.

        Args:
            input: The user prompt.

        Returns:
            The output of querying the application with the given input.
        """
        return self.agent_executor.invoke(input={"input": input})

When writing your Python class, the following three methods are important for reasoning engine:

1. __init__():
   • Use this method for only application configuration parameters. For example, you can use this method to collect the model parameters and safety attributes as input arguments from your users. You can also use this method to collect parameters such as the project ID, the region, application credentials, and API keys.
   • The constructor returns an object that must be "pickle-able" for it to be deployable to reasoning engine. Therefore, you should initialize service clients and establish connections to databases in the .set_up method instead of in the __init__ method.
   • This method is optional. If it's not specified, Vertex AI uses the default Python constructor for the class.
2. set_up():
   • You must use this method to define application initialization logic. For example, you use this method to establish connections to databases or dependent services, import dependent packages, or precompute data that's used for serving queries.
   • This method is optional. If it's not specified, Vertex AI assumes that the application doesn't need to call a .set_up method before serving user queries.
3. query():
   • You must use this method to define runtime logic that serves user queries. For example, you use this method to generate content with Generative AI models or retrieve real-time data from external APIs.
   • This method is required. If it's not specified, the ReasoningEngine service returns an error when you attempt to create a remote instance of the application.
   • You should give this method a clear docstring that defines what it does, documents its attributes, and provides type annotations for its inputs. Avoid variable arguments in the query method.

Test the application locally

Instantiate the application in local memory using the following code:

agent = CLASS_NAME(
    model=model,  # Required.
    tools=[get_exchange_rate],  # Optional.
    project=PROJECT_ID,
    location=LOCATION,
)
agent.set_up()

You can test the application by sending test queries to the local instance:

response = agent.query(
    input="What is the exchange rate from US dollars to Swedish currency?"
)

The response is a dictionary that's similar to the following:

{"input": "What is the exchange rate from US dollars to Swedish currency?",
 # ...
 "output": "For 1 US dollar you will get 10.7345 Swedish Krona."}

What's next

• Deploy the application.