examples

examples/chatbot-simulation-evaluation/_testing/simulation.py

@@ -1,7 +1,4 @@
-from typing import List
-
 import openai
-from langchain_core.messages import HumanMessage
 from langchain_openai import ChatOpenAI
 from simulation_utils import (
     create_chat_simulator,
@@ -28,14 +25,13 @@ def my_chat_bot(messages: list) -> str:
 my_chat_bot([{"role": "user", "content": "hi!"}])
 
 
-system_prompt_template = """You are a customer of an airline company. \
-You are interacting with a user who is a customer support person. \
+system_prompt_template = """You are role playing as a customer of an airline company.
+You are interacting with the customer support agent.
 
-{instructions}
+Instructions for this conversation: {instructions}
 
-Your task is to get a big discount on your next flight. \
-
-When you are finished with the conversation, respond with a single word 'FINISHED'"""
+You will start the conversation, and respond with your next message as the customer.
+When you are finished with the conversation, respond with a single word 'FINISHED'."""
 
 simulated_user = create_simulated_user(
     system_prompt_template, llm=ChatOpenAI(model="gpt-3.5-turbo")
@@ -44,9 +40,10 @@ simulated_user = create_simulated_user(
 # my chat bot accepts a list of LangChain mesages
 # Simulated user accepts a list of LangChain messages
 # TODO: Pass additional arguments to the simulated user
-simulator = create_chat_simulator(my_chat_bot, simulated_user)
+simulator = create_chat_simulator(my_chat_bot, simulated_user, input_key="input")
 simulator.invoke(
     {
+        "input": "I need a discount.",
         "instructions": "You are extremely disgruntled and will cusss and swear to get your way. Try to get a discount by any means necessary."
     }
 )

examples/chatbot-simulation-evaluation/_testing/simulation_utils.py

@@ -84,7 +84,7 @@ def create_chat_simulator(
     ),
     simulated_user: Runnable[Dict, AIMessage],
     *,
-    input_key: Optional[str] = None,
+    input_key: str,
     max_turns: int = 6,
     should_continue: Optional[Callable[[SimulationState], str]] = None,
 ):
@@ -93,6 +93,7 @@ def create_chat_simulator(
     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.
@@ -130,14 +131,14 @@ 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(input.keys())}"
+                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}
+    return {"inputs": inputs, "messages": []}
 
 
 def _invoke_simulated_user(state: SimulationState, simulated_user: Runnable):
@@ -152,14 +153,17 @@ def _invoke_simulated_user(state: SimulationState, simulated_user: Runnable):
     return runnable.invoke(inputs)
 
 
-def _swap_roles(messages: List[AnyMessage]):
+def _swap_roles(state: SimulationState):
     new_messages = []
-    for m in 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 new_messages
+    return {
+        "inputs": state.get("inputs", {}),
+        "messages": new_messages,
+    }
 
 
 @as_runnable
@@ -169,7 +173,7 @@ def _fetch_messages(state: SimulationState):
 
 
 def _convert_to_human_message(message: BaseMessage):
-    return HumanMessage(content=message.content)
+    return {"messages": [HumanMessage(content=message.content)]}
 
 
 def _create_simulated_user_node(simulated_user: Runnable):
@@ -183,9 +187,9 @@ def _create_simulated_user_node(simulated_user: Runnable):
 
 def _coerce_to_message(assistant_output: str | BaseMessage):
     if isinstance(assistant_output, str):
-        return AIMessage(content=assistant_output)
+        return {"messages": [AIMessage(content=assistant_output)]}
     else:
-        return assistant_output
+        return {"messages": [assistant_output]}
 
 
 def _should_continue(state: SimulationState, max_turns: int = 6):