langchain.py

# -*- coding: utf-8 -*-
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Iterable,
    Mapping,
    Optional,
    Union,
)

if TYPE_CHECKING:
    try:
        from langchain_core import runnables
        from langchain_core import tools as lc_tools
        from langchain_core.language_models import base as lc_language_models

        BaseTool = lc_tools.BaseTool
        BaseLanguageModel = lc_language_models.BaseLanguageModel
        GetSessionHistoryCallable = runnables.history.GetSessionHistoryCallable
        RunnableConfig = runnables.RunnableConfig
        RunnableSerializable = runnables.RunnableSerializable
    except ImportError:
        BaseTool = Any
        BaseLanguageModel = Any
        GetSessionHistoryCallable = Any
        RunnableConfig = Any
        RunnableSerializable = Any

    try:
        from langchain_google_genai.functions_utils import _ToolsType
    except ImportError:
        try:
            from langchain_google_vertexai.functions_utils import _ToolsType
        except ImportError:
            _ToolsType = Any

    try:
        from opentelemetry.sdk import trace

        TracerProvider = trace.TracerProvider
        SpanProcessor = trace.SpanProcessor
        SynchronousMultiSpanProcessor = trace.SynchronousMultiSpanProcessor
    except ImportError:
        TracerProvider = Any
        SpanProcessor = Any
        SynchronousMultiSpanProcessor = Any


def _default_runnable_kwargs(has_history: bool) -> Mapping[str, Any]:
    # https://github.com/langchain-ai/langchain/blob/5784dfed001730530637793bea1795d9d5a7c244/libs/core/langchain_core/runnables/history.py#L237-L241
    runnable_kwargs = {
        # input_messages_key (str): Must be specified if the underlying
        # agent accepts a dict as input.
        "input_messages_key": "input",
        # output_messages_key (str): Must be specified if the underlying
        # agent returns a dict as output.
        "output_messages_key": "output",
    }
    if has_history:
        # history_messages_key (str): Must be specified if the underlying
        # agent accepts a dict as input and a separate key for historical
        # messages.
        runnable_kwargs["history_messages_key"] = "history"
    return runnable_kwargs


def _default_output_parser():
    try:
        from langchain_classic.agents.output_parsers.tools import ToolsAgentOutputParser
    except (ModuleNotFoundError, ImportError):
        try:
            from langchain.agents.output_parsers.tools import ToolsAgentOutputParser
        except (ModuleNotFoundError, ImportError):
            # Fallback to an older version if needed.
            from langchain.agents.output_parsers.openai_tools import (
                OpenAIToolsAgentOutputParser as ToolsAgentOutputParser,
            )
    return ToolsAgentOutputParser()


def _default_model_builder(
    model_name: str,
    *,
    project: str,
    location: str,
    model_kwargs: Optional[Mapping[str, Any]] = None,
) -> "BaseLanguageModel":
    model_kwargs = model_kwargs or {}
    try:
        from langchain_google_genai import ChatGoogleGenerativeAI

        model = ChatGoogleGenerativeAI(
            model=model_name,
            project=project,
            location=location,
            vertexai=True,
            **model_kwargs,
        )
        return model
    except ImportError:
        import agentplatform
        from google.cloud.aiplatform import initializer
        from langchain_google_vertexai import ChatVertexAI

        current_project = initializer.global_config.project
        current_location = initializer.global_config.location
        agentplatform.init(project=project, location=location)
        model = ChatVertexAI(model_name=model_name, **model_kwargs)
        agentplatform.init(project=current_project, location=current_location)
        return model


def _default_runnable_builder(
    model: "BaseLanguageModel",
    *,
    system_instruction: Optional[str] = None,
    tools: Optional["_ToolsType"] = None,
    prompt: Optional["RunnableSerializable"] = None,
    output_parser: Optional["RunnableSerializable"] = None,
    chat_history: Optional["GetSessionHistoryCallable"] = None,
    model_tool_kwargs: Optional[Mapping[str, Any]] = None,
    agent_executor_kwargs: Optional[Mapping[str, Any]] = None,
    runnable_kwargs: Optional[Mapping[str, Any]] = None,
) -> "RunnableSerializable":
    from langchain_core import tools as lc_tools

    try:
        from langchain_classic.agents import AgentExecutor
    except ImportError:
        from langchain.agents import AgentExecutor

    try:
        from langchain_core.tools import StructuredTool
    except ImportError:
        from langchain.tools.base import StructuredTool

    # The prompt template and runnable_kwargs needs to be customized depending
    # on whether the user intends for the agent to have history. The way the
    # user would reflect that is by setting chat_history (which defaults to
    # None).
    has_history: bool = chat_history is not None
    prompt = prompt or _default_prompt(
        has_history=has_history,
        system_instruction=system_instruction,
    )
    output_parser = output_parser or _default_output_parser()
    model_tool_kwargs = model_tool_kwargs or {}
    agent_executor_kwargs = agent_executor_kwargs or {}
    runnable_kwargs = runnable_kwargs or _default_runnable_kwargs(has_history)
    if tools:
        model = model.bind_tools(tools=tools, **model_tool_kwargs)
    else:
        tools = []
    agent_executor = AgentExecutor(
        agent=prompt | model | output_parser,
        tools=[
            (
                tool
                if isinstance(tool, lc_tools.BaseTool)
                else StructuredTool.from_function(tool)
            )
            for tool in tools
            if isinstance(tool, (Callable, lc_tools.BaseTool))
        ],
        **agent_executor_kwargs,
    )
    if has_history:
        from langchain_core.runnables.history import RunnableWithMessageHistory

        return RunnableWithMessageHistory(
            runnable=agent_executor,
            get_session_history=chat_history,
            **runnable_kwargs,
        )
    return agent_executor


def _default_instrumentor_builder(project_id: str):
    from agentplatform._genai import _agent_engines_utils

    openinference_langchain = (
        _agent_engines_utils._import_openinference_langchain_or_warn()
    )
    opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn()
    opentelemetry_sdk_trace = (
        _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn()
    )
    if all(
        (
            openinference_langchain,
            opentelemetry,
            opentelemetry_sdk_trace,
        )
    ):
        try:
            import opentelemetry.exporter.otlp.proto.http.trace_exporter
            import opentelemetry.sdk.resources
            import google.auth
            import google.auth.transport.requests
        except (ImportError, AttributeError):
            from google.cloud.aiplatform import base

            _LOGGER = base.Logger(__name__)
            _LOGGER.warning(
                "enable_tracing=True but proceeding with tracing disabled "
                "because opentelemetry-exporter-otlp-proto-http is not installed."
            )
            return None

        credentials, _ = google.auth.default()
        resource = opentelemetry.sdk.resources.Resource.create(
            attributes={"gcp.project_id": project_id},
        ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect())
        span_exporter = (
            opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter(
                session=google.auth.transport.requests.AuthorizedSession(
                    credentials=credentials
                ),
                endpoint="https://telemetry.googleapis.com/v1/traces",
            )
        )
        span_processor: SpanProcessor = (
            opentelemetry_sdk_trace.export.SimpleSpanProcessor(
                span_exporter=span_exporter,
            )
        )
        tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
        # Get the appropriate tracer provider:
        # 1. If _TRACER_PROVIDER is already set, use that.
        # 2. Otherwise, if the OTEL_PYTHON_TRACER_PROVIDER environment
        # variable is set, use that.
        # 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
        # If none of the above is set, we log a warning, and
        # create a tracer provider.
        if not tracer_provider:
            from google.cloud.aiplatform import base

            _LOGGER = base.Logger(__name__)
            _LOGGER.warning(
                "No tracer provider. By default, "
                "we should get one of the following providers: "
                "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
                "or _PROXY_TRACER_PROVIDER."
            )
            tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource)
            opentelemetry.trace.set_tracer_provider(tracer_provider)
        # Avoids AttributeError:
        # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
        # attribute 'add_span_processor'.
        if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider):
            tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource)
            opentelemetry.trace.set_tracer_provider(tracer_provider)
        # Avoids OpenTelemetry client already exists error.
        _override_active_span_processor(
            tracer_provider,
            opentelemetry_sdk_trace.SynchronousMultiSpanProcessor(),
        )
        tracer_provider.add_span_processor(span_processor)
        # Keep the instrumentation up-to-date.
        # When creating multiple LangchainAgents,
        # we need to keep the instrumentation up-to-date.
        # We deliberately override the instrument each time,
        # so that if different agents end up using different
        # instrumentations, we guarantee that the user is always
        # working with the most recent agent's instrumentation.
        instrumentor = openinference_langchain.LangChainInstrumentor()
        if instrumentor.is_instrumented_by_opentelemetry:
            instrumentor.uninstrument()
        instrumentor.instrument()
        return instrumentor
    else:
        from google.cloud.aiplatform import base

        _LOGGER = base.Logger(__name__)
        _LOGGER.warning(
            "enable_tracing=True but proceeding with tracing disabled "
            "because not all packages for tracing have been installed"
        )
        return None


def _default_prompt(
    has_history: bool,
    system_instruction: Optional[str] = None,
) -> "RunnableSerializable":
    from langchain_core import prompts

    try:
        from langchain_classic.agents.format_scratchpad.tools import (
            format_to_tool_messages,
        )
    except (ModuleNotFoundError, ImportError):
        try:
            from langchain.agents.format_scratchpad.tools import format_to_tool_messages
        except (ModuleNotFoundError, ImportError):
            from langchain.agents.format_scratchpad.openai_tools import (
                format_to_openai_tool_messages as format_to_tool_messages,
            )

    system_instructions = []
    if system_instruction:
        system_instructions = [("system", system_instruction)]

    if has_history:
        return {
            "history": lambda x: x["history"],
            "input": lambda x: x["input"],
            "agent_scratchpad": (
                lambda x: format_to_tool_messages(x["intermediate_steps"])
            ),
        } | prompts.ChatPromptTemplate.from_messages(
            system_instructions
            + [
                prompts.MessagesPlaceholder(variable_name="history"),
                ("user", "{input}"),
                prompts.MessagesPlaceholder(variable_name="agent_scratchpad"),
            ]
        )
    else:
        return {
            "input": lambda x: x["input"],
            "agent_scratchpad": (
                lambda x: format_to_tool_messages(x["intermediate_steps"])
            ),
        } | prompts.ChatPromptTemplate.from_messages(
            system_instructions
            + [
                ("user", "{input}"),
                prompts.MessagesPlaceholder(variable_name="agent_scratchpad"),
            ]
        )


def _validate_callable_parameters_are_annotated(callable: Callable):
    """Validates that the parameters of the callable have type annotations.

    This ensures that they can be used for constructing LangChain tools that are
    usable with Gemini function calling.
    """
    import inspect

    parameters = dict(inspect.signature(callable).parameters)
    for name, parameter in parameters.items():
        if parameter.annotation == inspect.Parameter.empty:
            raise TypeError(
                f"Callable={callable.__name__} has untyped input_arg={name}. "
                f"Please specify a type when defining it, e.g. `{name}: str`."
            )


def _validate_tools(tools: "_ToolsType"):
    """Validates that the tools are usable for tool calling."""
    for tool in tools:
        if isinstance(tool, Callable):
            _validate_callable_parameters_are_annotated(tool)


def _override_active_span_processor(
    tracer_provider: "TracerProvider",
    active_span_processor: "SynchronousMultiSpanProcessor",
):
    """Overrides the active span processor.

    When working with multiple LangchainAgents in the same environment,
    it's crucial to manage trace exports carefully.
    Each agent needs its own span processor tied to a unique project ID.
    While we add a new span processor for each agent, this can lead to
    unexpected behavior.
    For instance, with two agents linked to different projects, traces from the
    second agent might be sent to both projects.
    To prevent this and guarantee traces go to the correct project, we overwrite
    the active span processor whenever a new LangchainAgent is created.

    Args:
        tracer_provider (TracerProvider):
            The tracer provider to use for the project.
        active_span_processor (SynchronousMultiSpanProcessor):
            The active span processor overrides the tracer provider's
            active span processor.
    """
    if tracer_provider._active_span_processor:
        tracer_provider._active_span_processor.shutdown()
    tracer_provider._active_span_processor = active_span_processor


class LangchainAgent:
    """A Langchain Agent.

    See https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/develop
    for details.
    """

    agent_framework = "langchain"

    def __init__(
        self,
        model: str,
        *,
        system_instruction: Optional[str] = None,
        prompt: Optional["RunnableSerializable"] = None,
        tools: Optional["_ToolsType"] = None,
        output_parser: Optional["RunnableSerializable"] = None,
        chat_history: Optional["GetSessionHistoryCallable"] = None,
        model_kwargs: Optional[Mapping[str, Any]] = None,
        model_tool_kwargs: Optional[Mapping[str, Any]] = None,
        agent_executor_kwargs: Optional[Mapping[str, Any]] = None,
        runnable_kwargs: Optional[Mapping[str, Any]] = None,
        model_builder: Optional[Callable] = None,
        runnable_builder: Optional[Callable] = None,
        enable_tracing: bool = False,
        instrumentor_builder: Optional[Callable[..., Any]] = None,
    ):
        """Initializes the LangchainAgent.

        Under-the-hood, assuming .set_up() is called, this will correspond to

        ```
        model = model_builder(model_name=model, model_kwargs=model_kwargs)
        runnable = runnable_builder(
            prompt=prompt,
            model=model,
            tools=tools,
            output_parser=output_parser,
            chat_history=chat_history,
            agent_executor_kwargs=agent_executor_kwargs,
            runnable_kwargs=runnable_kwargs,
        )
        ```

        When everything is based on their default values, this corresponds to
        ```
        # model_builder
        from langchain_google_vertexai import ChatVertexAI
        llm = ChatVertexAI(model_name=model, **model_kwargs)

        # runnable_builder
        from langchain import agents
        from langchain_core.runnables.history import RunnableWithMessageHistory
        llm_with_tools = llm.bind_tools(tools=tools, **model_tool_kwargs)
        agent_executor = agents.AgentExecutor(
            agent=prompt | llm_with_tools | output_parser,
            tools=tools,
            **agent_executor_kwargs,
        )
        runnable = RunnableWithMessageHistory(
            runnable=agent_executor,
            get_session_history=chat_history,
            **runnable_kwargs,
        )
        ```

        Args:
            model (str):
                Optional. The name of the model (e.g. "gemini-1.0-pro").
            system_instruction (str):
                Optional. The system instruction to use for the agent. This
                argument should not be specified if `prompt` is specified.
            prompt (langchain_core.runnables.RunnableSerializable):
                Optional. The prompt template for the model. Defaults to a
                ChatPromptTemplate.
            tools (Sequence[langchain_core.tools.BaseTool, Callable]):
                Optional. The tools for the agent to be able to use. All input
                callables (e.g. function or class method) will be converted
                to a langchain.tools.base.StructuredTool. Defaults to None.
            output_parser (langchain_core.runnables.RunnableSerializable):
                Optional. The output parser for the model. Defaults to an
                output parser that works with Gemini function-calling.
            chat_history (langchain_core.runnables.history.GetSessionHistoryCallable):
                Optional. Callable that returns a new BaseChatMessageHistory.
                Defaults to None, i.e. chat_history is not preserved.
            model_kwargs (Mapping[str, Any]):
                Optional. Additional keyword arguments for the constructor of
                chat_models.ChatVertexAI. An example would be
                ```
                {
                    # temperature (float): Sampling temperature, it controls the
                    # degree of randomness in token selection.
                    "temperature": 0.28,
                    # max_output_tokens (int): Token limit determines the
                    # maximum amount of text output from one prompt.
                    "max_output_tokens": 1000,
                    # top_p (float): Tokens are selected from most probable to
                    # least, until the sum of their probabilities equals the
                    # top_p value.
                    "top_p": 0.95,
                    # top_k (int): How the model selects tokens for output, the
                    # next token is selected from among the top_k most probable
                    # tokens.
                    "top_k": 40,
                }
                ```
            model_tool_kwargs (Mapping[str, Any]):
                Optional. Additional keyword arguments when binding tools to the
                model using `model.bind_tools()`.
            agent_executor_kwargs (Mapping[str, Any]):
                Optional. Additional keyword arguments for the constructor of
                langchain.agents.AgentExecutor. An example would be
                ```
                {
                    # Whether to return the agent's trajectory of intermediate
                    # steps at the end in addition to the final output.
                    "return_intermediate_steps": False,
                    # The maximum number of steps to take before ending the
                    # execution loop.
                    "max_iterations": 15,
                    # The method to use for early stopping if the agent never
                    # returns `AgentFinish`. Either 'force' or 'generate'.
                    "early_stopping_method": "force",
                    # How to handle errors raised by the agent's output parser.
                    # Defaults to `False`, which raises the error.
                    "handle_parsing_errors": False,
                }
                ```
            runnable_kwargs (Mapping[str, Any]):
                Optional. Additional keyword arguments for the constructor of
                langchain.runnables.history.RunnableWithMessageHistory if
                chat_history is specified. If chat_history is None, this will be
                ignored.
            model_builder (Callable):
                Optional. Callable that returns a new language model. Defaults
                to a a callable that returns ChatVertexAI based on `model`,
                `model_kwargs` and the parameters in `agentplatform.init`.
            runnable_builder (Callable):
                Optional. Callable that returns a new runnable. This can be used
                for customizing the orchestration logic of the Agent based on
                the model returned by `model_builder` and the rest of the input
                arguments.
            enable_tracing (bool):
                Optional. Whether to enable tracing in Cloud Trace. Defaults to
                False.
            instrumentor_builder (Callable[..., Any]):
                Optional. Callable that returns a new instrumentor. This can be
                used for customizing the instrumentation logic of the Agent.
                If not provided, a default instrumentor builder will be used.
                This parameter is ignored if `enable_tracing` is False.

        Raises:
            ValueError: If both `prompt` and `system_instruction` are specified.
            TypeError: If there is an invalid tool (e.g. function with an input
            that did not specify its type).
        """
        from google.cloud.aiplatform import initializer

        self._tmpl_attrs: dict[str, Any] = {
            "project": initializer.global_config.project,
            "location": initializer.global_config.location,
            "tools": [],
            "model_name": model,
            "system_instruction": system_instruction,
            "prompt": prompt,
            "output_parser": output_parser,
            "chat_history": chat_history,
            "model_kwargs": model_kwargs,
            "model_tool_kwargs": model_tool_kwargs,
            "agent_executor_kwargs": agent_executor_kwargs,
            "runnable_kwargs": runnable_kwargs,
            "model_builder": model_builder,
            "runnable_builder": runnable_builder,
            "enable_tracing": enable_tracing,
            "model": None,
            "runnable": None,
            "instrumentor": None,
            "instrumentor_builder": instrumentor_builder,
        }
        if tools:
            # We validate tools at initialization for actionable feedback before
            # they are deployed.
            _validate_tools(tools)
            self._tmpl_attrs["tools"] = tools
        if prompt and system_instruction:
            raise ValueError(
                "Only one of `prompt` or `system_instruction` should be specified. "
                "Consider incorporating the system instruction into the prompt "
                "rather than passing it separately as an argument."
            )

    def set_up(self):
        """Sets up the agent for execution of queries at runtime.

        It initializes the model, binds the model with tools, and connects it
        with the prompt template and output parser.

        This method should not be called for an object being passed to the
        service for deployment, as it might initialize clients that can not be
        serialized.
        """
        if self._tmpl_attrs.get("enable_tracing"):
            instrumentor_builder = (
                self._tmpl_attrs.get("instrumentor_builder")
                or _default_instrumentor_builder
            )
            self._tmpl_attrs["instrumentor"] = instrumentor_builder(
                project_id=self._tmpl_attrs.get("project")
            )
        model_builder = self._tmpl_attrs.get("model_builder") or _default_model_builder
        self._tmpl_attrs["model"] = model_builder(
            model_name=self._tmpl_attrs.get("model_name"),
            model_kwargs=self._tmpl_attrs.get("model_kwargs"),
            project=self._tmpl_attrs.get("project"),
            location=self._tmpl_attrs.get("location"),
        )
        runnable_builder = (
            self._tmpl_attrs.get("runnable_builder") or _default_runnable_builder
        )
        self._tmpl_attrs["runnable"] = runnable_builder(
            prompt=self._tmpl_attrs.get("prompt"),
            model=self._tmpl_attrs.get("model"),
            tools=self._tmpl_attrs.get("tools"),
            system_instruction=self._tmpl_attrs.get("system_instruction"),
            output_parser=self._tmpl_attrs.get("output_parser"),
            chat_history=self._tmpl_attrs.get("chat_history"),
            model_tool_kwargs=self._tmpl_attrs.get("model_tool_kwargs"),
            agent_executor_kwargs=self._tmpl_attrs.get("agent_executor_kwargs"),
            runnable_kwargs=self._tmpl_attrs.get("runnable_kwargs"),
        )

    def clone(self) -> "LangchainAgent":
        """Returns a clone of the LangchainAgent."""
        import copy

        return LangchainAgent(
            model=self._tmpl_attrs.get("model_name"),
            system_instruction=self._tmpl_attrs.get("system_instruction"),
            prompt=copy.deepcopy(self._tmpl_attrs.get("prompt")),
            tools=copy.deepcopy(self._tmpl_attrs.get("tools")),
            output_parser=copy.deepcopy(self._tmpl_attrs.get("output_parser")),
            chat_history=copy.deepcopy(self._tmpl_attrs.get("chat_history")),
            model_kwargs=copy.deepcopy(self._tmpl_attrs.get("model_kwargs")),
            model_tool_kwargs=copy.deepcopy(self._tmpl_attrs.get("model_tool_kwargs")),
            agent_executor_kwargs=copy.deepcopy(
                self._tmpl_attrs.get("agent_executor_kwargs")
            ),
            runnable_kwargs=copy.deepcopy(self._tmpl_attrs.get("runnable_kwargs")),
            model_builder=self._tmpl_attrs.get("model_builder"),
            runnable_builder=self._tmpl_attrs.get("runnable_builder"),
            enable_tracing=self._tmpl_attrs.get("enable_tracing"),
            instrumentor_builder=self._tmpl_attrs.get("instrumentor_builder"),
        )

    def query(
        self,
        *,
        input: Union[str, Mapping[str, Any]],
        config: Optional["RunnableConfig"] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """Queries the Agent with the given input and config.

        Args:
            input (Union[str, Mapping[str, Any]]):
                Required. The input to be passed to the Agent.
            config (langchain_core.runnables.RunnableConfig):
                Optional. The config (if any) to be used for invoking the Agent.
            **kwargs:
                Optional. Any additional keyword arguments to be passed to the
                `.invoke()` method of the corresponding AgentExecutor.

        Returns:
            The output of querying the Agent with the given input and config.
        """
        try:
            from langchain_core.load import dumpd
        except ImportError:
            from langchain.load import dump as langchain_load_dump

            dumpd = langchain_load_dump.dumpd

        if isinstance(input, str):
            input = {"input": input}
        if not self._tmpl_attrs.get("runnable"):
            self.set_up()
        return dumpd(
            self._tmpl_attrs.get("runnable").invoke(
                input=input, config=config, **kwargs
            )
        )

    def stream_query(
        self,
        *,
        input: Union[str, Mapping[str, Any]],
        config: Optional["RunnableConfig"] = None,
        **kwargs,
    ) -> Iterable[Any]:
        """Stream queries the Agent with the given input and config.

        Args:
            input (Union[str, Mapping[str, Any]]):
                Required. The input to be passed to the Agent.
            config (langchain_core.runnables.RunnableConfig):
                Optional. The config (if any) to be used for invoking the Agent.
            **kwargs:
                Optional. Any additional keyword arguments to be passed to the
                `.invoke()` method of the corresponding AgentExecutor.

        Yields:
            The output of querying the Agent with the given input and config.
        """
        try:
            from langchain_core.load import dumpd
        except ImportError:
            from langchain.load import dump as langchain_load_dump

            dumpd = langchain_load_dump.dumpd

        if isinstance(input, str):
            input = {"input": input}
        if not self._tmpl_attrs.get("runnable"):
            self.set_up()
        for chunk in self._tmpl_attrs.get("runnable").stream(
            input=input,
            config=config,
            **kwargs,
        ):
            yield dumpd(chunk)