Comment
03-31-2025
06:17 PM
03-31-2025
06:17 PM
@s-udhaya and @jiayi-wu really excellent work, thank you for the write up.
@Ben_H good question and thank you for providing your code. I used your above in place of the existing chain in the Chat History Extractor Chain command of the llm-rag-chatbot/test_03-advanced-app/02-Advanced-Chatbot-Chain notebook and saw that it worked in the notebook as you described, but not in the Review App.
Please add this cell after the one where you invoke the chain locally:
# Extract and handle both filter information and chat history
def create_configurable_with_filters(input_data):
"""Create configurable object with filters from the input data structure."""
# Handle the review app format where custom_inputs come directly in the input
if "custom_inputs" in input_data and "filters" in input_data["custom_inputs"]:
filters = input_data["custom_inputs"]["filters"]
# Handle the notebook testing format
elif "messages" in input_data and isinstance(input_data["messages"], list):
if "custom_inputs" in input_data and "filters" in input_data["custom_inputs"]:
filters = input_data["custom_inputs"]["filters"]
else:
filters = {} # Default empty filter
else:
filters = {} # Default to empty filters if no valid format is found
# Create the configurable object with the extracted filters
configurable = {
"configurable": {
"search_kwargs": {
"k": model_config.get("retriever_config").get("parameters")["k"],
"query_type": model_config.get("retriever_config").get("parameters")["query_type"],
"filters": filters,
}
}
}
return configurable
# Extract user query for context retrieval
def get_retrieval_query(input_data):
"""Get the query for retrieval, using query rewriting when chat history exists."""
messages = input_data.get("messages", [])
chat_history = extract_chat_history(messages)
if len(chat_history) > 0:
# Use query rewriting when chat history exists
user_query = extract_user_query_string(messages)
rewrite_input = {"question": user_query, "chat_history": chat_history}
return query_rewrite_prompt.format_messages(**rewrite_input) | model | StrOutputParser()
else:
# Use direct query when no chat history
return extract_user_query_string(messages)
# Function to retrieve context with filters
def retrieve_with_filters(input_data):
"""Retrieve context using vector search with appropriate filters."""
# Get the query for retrieval
query = get_retrieval_query(input_data)
# Create the configurable object with filters
config = create_configurable_with_filters(input_data)
# Use the configurable retriever
return configurable_vs_retriever.invoke(query, config=config)
# Define the language model
model = ChatDatabricks(
endpoint=model_config.get("databricks_resources").get("llm_endpoint_name"),
extra_params=model_config.get("llm_config").get("llm_parameters"),
)
# Prompt Template for generation
prompt = ChatPromptTemplate.from_messages(
[
("system", model_config.get("llm_config").get("llm_prompt_template")),
# Note: This chain does not compress the history, so very long conversations can overflow the context window.
MessagesPlaceholder(variable_name="formatted_chat_history"),
# User's most current question
("user", "{question}"),
]
)
# Define the integrated chain
chain = (
RunnablePassthrough()
| {
"question": itemgetter("messages") | RunnableLambda(extract_user_query_string),
"formatted_chat_history": itemgetter("messages") | RunnableLambda(format_chat_history_for_prompt),
"context": RunnableLambda(retrieve_with_filters) | RunnableLambda(format_context)
}
| prompt
| model
| StrOutputParser()
)
# Create a wrapper function for MLflow to simplify the input handling
def run_chain(input_data):
"""Run the chain with standardized input handling."""
return chain.invoke(input_data)
# Register this as the main entry point for the model
run_chain
I was able to run the remainder of the notebook as is and see history and the custom_inputs flow through to the request field in Unity Catalog. Please let us know if you are able to reproduce.
Warm regards,
-Austin Zaccor