From 4df14268abdc235651d365a8c965383f0dbdb78b Mon Sep 17 00:00:00 2001 From: Filip Michalsky Date: Thu, 6 Jul 2023 09:53:38 -0400 Subject: [PATCH] add streaming test --- tests/test_salesgpt.py | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/tests/test_salesgpt.py b/tests/test_salesgpt.py index f142283c..3a4749a1 100644 --- a/tests/test_salesgpt.py +++ b/tests/test_salesgpt.py @@ -39,4 +39,37 @@ def test_valid_inference(self, load_env): def test_valid_inference_stream(self, load_env): - pass + """Test that the agent will start and generate the first utterance when streaming.""" + + llm = ChatOpenAI(temperature=0.9) + model_name = 'gpt-3.5-turbo' + + sales_agent = SalesGPT.from_llm( + llm, + verbose=False, + salesperson_name="Ted Lasso", + salesperson_role="Sales Representative", + company_name="Sleep Haven", + company_business="""Sleep Haven + is a premium mattress company that provides + customers with the most comfortable and + supportive sleeping experience possible. + We offer a range of high-quality mattresses, + pillows, and bedding accessories + that are designed to meet the unique + needs of our customers.""", + ) + + sales_agent.seed_agent() + sales_agent.determine_conversation_stage() # optional for demonstration, built into the prompt + + # agent output sample + stream_generator = sales_agent.step(return_streaming_generator=True, model_name=model_name) + agent_output='' + for chunk in stream_generator: + token = chunk["choices"][0]["delta"].get("content", "") + agent_output += token + + assert agent_output is not None, "Agent output cannot be None." + assert isinstance(agent_output, str), "Agent output needs to be of type str" + assert len(agent_output) > 0, "Length of output needs to be greater than 0."