diff --git a/examples/chatbot-simulation-evaluation/simulation_utils.py b/examples/chatbot-simulation-evaluation/simulation_utils.py new file mode 100644 index 000000000..f09e39276 --- /dev/null +++ b/examples/chatbot-simulation-evaluation/simulation_utils.py @@ -0,0 +1,203 @@ +import functools +from typing import Annotated, Any, Callable, Dict, List, Optional, Union + +from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, HumanMessage +from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder +from langchain_core.runnables import Runnable, RunnableLambda +from langchain_core.runnables import chain as as_runnable +from langchain_openai import ChatOpenAI +from typing_extensions import TypedDict + +from langgraph.graph import END, StateGraph + + +def langchain_to_openai_messages(messages: List[BaseMessage]): + """ + Convert a list of langchain base messages to a list of openai messages. + + Parameters: + messages (List[BaseMessage]): A list of langchain base messages. + + Returns: + List[dict]: A list of openai messages. + """ + from langchain.adapters.openai import convert_message_to_dict # noqa: I001 + + return [ + convert_message_to_dict(m) if isinstance(m, BaseMessage) else m + for m in messages + ] + + +def create_simulated_user( + system_prompt: str, llm: Runnable | None = None +) -> Runnable[Dict, AIMessage]: + """ + Creates a simulated user for chatbot simulation. + + Args: + system_prompt (str): The system prompt to be used by the simulated user. + llm (Runnable | None, optional): The language model to be used for the simulation. + Defaults to gpt-3.5-turbo. + + Returns: + Runnable[Dict, AIMessage]: The simulated user for chatbot simulation. + """ + return ChatPromptTemplate.from_messages( + [ + ("system", system_prompt), + MessagesPlaceholder(variable_name="messages"), + ] + ) | (llm or ChatOpenAI(model="gpt-3.5-turbo")).with_config( + run_name="simulated_user" + ) + + +Messages = Union[list[AnyMessage], AnyMessage] + + +def add_messages(left: Messages, right: Messages) -> Messages: + if not isinstance(left, list): + left = [left] + if not isinstance(right, list): + right = [right] + return left + right + + +class SimulationState(TypedDict): + """ + Represents the state of a simulation. + + Attributes: + messages (List[AnyMessage]): A list of messages in the simulation. + inputs (Optional[dict[str, Any]]): Optional inputs for the simulation. + """ + + messages: Annotated[List[AnyMessage], add_messages] + inputs: Optional[dict[str, Any]] + + +def create_chat_simulator( + assistant: ( + Callable[[List[AnyMessage]], str | AIMessage] + | Runnable[List[AnyMessage], str | AIMessage] + ), + simulated_user: Runnable[Dict, AIMessage], + *, + input_key: str, + max_turns: int = 6, + should_continue: Optional[Callable[[SimulationState], str]] = None, +): + """Creates a chat simulator for evaluating a chatbot. + + Args: + assistant: The chatbot assistant function or runnable object. + simulated_user: The simulated user object. + input_key: The key for the input to the chat simulation. + max_turns: The maximum number of turns in the chat simulation. Default is 6. + should_continue: Optional function to determine if the simulation should continue. + If not provided, a default function will be used. + + Returns: + The compiled chat simulation graph. + + """ + graph_builder = StateGraph(SimulationState) + graph_builder.add_node( + "user", + _create_simulated_user_node(simulated_user), + ) + graph_builder.add_node( + "assistant", _fetch_messages | assistant | _coerce_to_message + ) + graph_builder.add_edge("assistant", "user") + graph_builder.add_conditional_edges( + "user", + should_continue or functools.partial(_should_continue, max_turns=max_turns), + ) + # If your dataset has a 'leading question/input', then we route first to the assistant, otherwise, we let the user take the lead. + graph_builder.set_entry_point("assistant" if input_key is not None else "user") + + return ( + RunnableLambda(_prepare_example).bind(input_key=input_key) + | graph_builder.compile() + ) + + +## Private methods + + +def _prepare_example(inputs: dict[str, Any], input_key: Optional[str] = None): + if input_key is not None: + if input_key not in inputs: + raise ValueError( + f"Dataset's example input must contain the provided input key: '{input_key}'.\nFound: {list(inputs.keys())}" + ) + messages = [HumanMessage(content=inputs[input_key])] + return { + "inputs": {k: v for k, v in inputs.items() if k != input_key}, + "messages": messages, + } + return {"inputs": inputs, "messages": []} + + +def _invoke_simulated_user(state: SimulationState, simulated_user: Runnable): + """Invoke the simulated user node.""" + runnable = ( + simulated_user + if isinstance(simulated_user, Runnable) + else RunnableLambda(simulated_user) + ) + inputs = state.get("inputs", {}) + inputs["messages"] = state["messages"] + return runnable.invoke(inputs) + + +def _swap_roles(state: SimulationState): + new_messages = [] + for m in state["messages"]: + if isinstance(m, AIMessage): + new_messages.append(HumanMessage(content=m.content)) + else: + new_messages.append(AIMessage(content=m.content)) + return { + "inputs": state.get("inputs", {}), + "messages": new_messages, + } + + +@as_runnable +def _fetch_messages(state: SimulationState): + """Invoke the simulated user node.""" + return state["messages"] + + +def _convert_to_human_message(message: BaseMessage): + return {"messages": [HumanMessage(content=message.content)]} + + +def _create_simulated_user_node(simulated_user: Runnable): + """Simulated user accepts a {"messages": [...]} argument and returns a single message.""" + return ( + _swap_roles + | RunnableLambda(_invoke_simulated_user).bind(simulated_user=simulated_user) + | _convert_to_human_message + ) + + +def _coerce_to_message(assistant_output: str | BaseMessage): + if isinstance(assistant_output, str): + return {"messages": [AIMessage(content=assistant_output)]} + else: + return {"messages": [assistant_output]} + + +def _should_continue(state: SimulationState, max_turns: int = 6): + messages = state["messages"] + # TODO support other stop criteria + if len(messages) > max_turns: + return END + elif messages[-1].content.strip() == "FINISHED": + return END + else: + return "assistant"