address cmt

This commit is contained in:
cotran 2024-11-07 11:15:03 -08:00
parent dd07ba2cd0
commit 1f383eafc4
4 changed files with 27 additions and 23 deletions

View file

@ -10,6 +10,16 @@ LLM_GATEWAY_ENDPOINT = os.getenv(
)
ARCH_STATE_HEADER = "x-arch-state"
PREFILL_LIST = [
"May",
"Could",
"Sure",
"Definitely",
"Certainly",
"Of course",
"Can",
]
def get_data_chunks(stream, n=1):
chunks = []

View file

@ -3,20 +3,12 @@ import pytest
import requests
from deepdiff import DeepDiff
from common import PROMPT_GATEWAY_ENDPOINT, get_arch_messages, get_data_chunks
def get_prefill_list():
prefill_list = [
"May",
"Could",
"Sure",
"Definitely",
"Certainly",
"Of course",
"Can",
]
return prefill_list
from common import (
PROMPT_GATEWAY_ENDPOINT,
PREFILL_LIST,
get_arch_messages,
get_data_chunks,
)
@pytest.mark.parametrize("stream", [True, False])
@ -126,10 +118,9 @@ def test_prompt_gateway_arch_direct_response(stream):
message = choices[0]["message"]["content"]
assert "Could you provide the following details days" not in message
prefill_list = get_prefill_list()
assert any(
message.startswith(word) for word in prefill_list
), f"Expected assistant message to start with one of {prefill_list}, but got '{assistant_message}'"
message.startswith(word) for word in PREFILL_LIST
), f"Expected assistant message to start with one of {PREFILL_LIST}, but got '{assistant_message}'"
@pytest.mark.parametrize("stream", [True, False])