@Ramana - can you please be more specific with the changes that were required? I am also receiving the same error and was originally using the langchain_core.runnables library and re-worked the code to not rely on it and I am still receiving the same issue when deploying. Agents works fine when running it in my notebook. My original code (listed below) stemmed from the Multi-Agent Genie system example in this link below. I originally had additional nodes including Genie but removed them to try and get the deployment to work for now.
https://docs.databricks.com/aws/en/generative-ai/agent-framework/multi-agent-genie?scid=701Vp000004h...
import functools
import os
from typing import Any, Generator, Literal, Optional
import mlflow
from databricks.sdk import WorkspaceClient
from databricks_langchain import ChatDatabricks, VectorSearchRetrieverTool
from databricks_langchain.uc_ai import (
DatabricksFunctionClient,
UCFunctionToolkit,
set_uc_function_client
)
from databricks_langchain.genie import GenieAgent
from langchain_core.runnables import RunnableLambda
from langgraph.graph import END, StateGraph
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt import create_react_agent
from mlflow.langchain.chat_agent_langgraph import ChatAgentState
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
ChatAgentChunk,
ChatAgentMessage,
ChatAgentResponse,
ChatContext,
)
from pydantic import BaseModel
from langchain_openai import OpenAIEmbeddings
mlflow.langchain.autolog()
############################################
############################################
LLM_ENDPOINT_NAME = "XXXX"
llm = ChatDatabricks(endpoint=LLM_ENDPOINT_NAME)
assert LLM_ENDPOINT_NAME is not None
client = DatabricksFunctionClient()
set_uc_function_client(client)
tools = []
# # TODO if desired, add additional tools and update the description of this agent
uc_tool_names = ["system.ai.python_exec"]#,"dev_catalog.tmp.liquidity_calculations"]
uc_toolkit = UCFunctionToolkit(function_names=uc_tool_names)
tools.extend(uc_toolkit.tools)
tools_agent_description = (
"This agent can execute python code and perform data analysis.",
)
embedding_model = OpenAIEmbeddings(model="text-embedding-3-large")
index_name = "dev_catalog.default.earnings_index"
endpoint_name = "earnings_index"
vs_tools = []
vs_agent_description = ("""
The Earnings Vector Search agent has access to a knowledge base of earnings report data related to the company ABC
The knowledge base includes information from 10Ks and Investor Presentations.
Users will want to be able to access numerical data from these reports. This includes servicing related metrics such as delinquency
""")
vs_tool_description=(
"Provide users information from company earnings reports"
"Returns numerical results related to earnings, company performance, delinquency rates, origination volume"
"Avoid any explanation or commentary unless you are unsure "
)
vs_tool = [VectorSearchRetrieverTool(
index_name=index_name, # Index name in the format 'catalog.schema.index'
num_results=4, # Max number of documents to return
query_type="ANN", # Query type ("ANN" or "HYBRID").
tool_name="earnings_reports_vector_search", # Used by the LLM to understand the purpose of the tool
tool_description=vs_tool_description, # Used by the LLM to understand the purpose of the tool
text_column="text", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
embedding=embedding_model # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)]
vs_tools.extend(vs_tool)
#tools.extend(earnings_vs_tool)
tools_agent = create_react_agent(llm, tools=tools)
vs_agent = create_react_agent(llm, tools=vs_tools)
worker_descriptions = {
"Earnings Vector Search": vs_agent_description,
}
formatted_descriptions = "\n".join(
f"- {name}: {desc}" for name, desc in worker_descriptions.items()
)
system_prompt = f"""You are the supervisor in a multi-agent system. Your job is to route the user's question to the appropriate specialist agent(s).
You may choose from the following workers, or select FINISH if the question has already been fully answered based on data that has been retrieved from Genie queries.
- If multiple agents are needed, route them one at a time and collect their answers before selecting FINISH.
- Maintain awareness of which agents have already responded by reviewing the message history.
- IMPORTANT: DO NOT CALL AN AGENT MORE THAN ONCE
Available agents:
{formatted_descriptions}
"""
options = ["FINISH"] + list(worker_descriptions.keys())
def supervisor_agent(state):
class nextNode(BaseModel):
next_node: Literal[tuple(options)]
preprocessor = RunnableLambda(
lambda state: [{"role": "system", "content": system_prompt}] + state["messages"]
)
supervisor_chain = preprocessor | llm.with_structured_output(nextNode)
return supervisor_chain.invoke(state)
def agent_node(state, agent, name):
result = agent.invoke(state)
return {
"messages": [
{
"role": "assistant",
"content": result["messages"][-1].content,
"name": name,
}
]
}
def final_answer(state):
system_prompt = f'''
Using only the content in the messages, respond to the user's question using the answer given by the other agents.
- You should be trying to create a report framework that has an "Introduction"
- It should then have a section of a high level summary, title this as "Overview"
- Then have a final section that says "Figures" and includes just sub-bullets with any numerical value for each requested metric
'''
preprocessor = RunnableLambda(
lambda state: [{"role": "system", "content": system_prompt}] + state["messages"]
)
final_answer_chain = preprocessor | llm
return {"messages": [final_answer_chain.invoke(state)]}
class AgentState(ChatAgentState):
next_node: str
vs_node = functools.partial(agent_node, agent=vs_agent, name="Earnings Vector Search")
workflow = StateGraph(AgentState)
workflow.add_node("Earnings Vector Search", vs_node)
workflow.add_node("supervisor", supervisor_agent)
workflow.add_node("final_answer", final_answer)
workflow.set_entry_point("supervisor")
# We want our workers to ALWAYS "report back" to the supervisor when done
for worker in worker_descriptions.keys():
workflow.add_edge(worker, "supervisor")
# Let the supervisor decide which next node to go
workflow.add_conditional_edges(
"supervisor",
lambda x: x["next_node"],
{**{k: k for k in worker_descriptions.keys()}, "FINISH": "final_answer"},
)
workflow.add_edge("final_answer", END)
multi_agent = workflow.compile()
class LangGraphChatAgent(ChatAgent):
def __init__(self, agent: CompiledStateGraph):
self.agent = agent
def predict(
self,
messages: list[ChatAgentMessage],
context: Optional[ChatContext] = None,
custom_inputs: Optional[dict[str, Any]] = None,
) -> ChatAgentResponse:
request = {
"messages": [m.model_dump_compat(exclude_none=True) for m in messages]
}
messages = []
for event in self.agent.stream(request, stream_mode="updates"):
for node_data in event.values():
messages.extend(
ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
)
return ChatAgentResponse(messages=messages)
def predict_stream(
self,
messages: list[ChatAgentMessage],
context: Optional[ChatContext] = None,
custom_inputs: Optional[dict[str, Any]] = None,
) -> Generator[ChatAgentChunk, None, None]:
request = {
"messages": [m.model_dump_compat(exclude_none=True) for m in messages]
}
for event in self.agent.stream(request, stream_mode="updates"):
for node_data in event.values():
yield from (
ChatAgentChunk(**{"delta": msg})
for msg in node_data.get("messages", [])
)
# Create the agent object, and specify it as the agent object to use when
# loading the agent back for inference via mlflow.models.set_model()
AGENT = LangGraphChatAgent(multi_agent)
mlflow.models.set_model(AGENT)
