improve e2e tests

This commit is contained in:
Adil Hafeez 2024-10-28 12:12:15 -07:00
parent 9b1e9ba49d
commit 2517120eeb
8 changed files with 169 additions and 23 deletions

View file

@ -1,3 +1,4 @@
import json
import os
@ -7,6 +8,7 @@ PROMPT_GATEWAY_ENDPOINT = os.getenv(
LLM_GATEWAY_ENDPOINT = os.getenv(
"LLM_GATEWAY_ENDPOINT", "http://localhost:12000/v1/chat/completions"
)
ARCH_STATE_HEADER = "x-arch-state"
def get_data_chunks(stream, n=1):
@ -22,3 +24,19 @@ def get_data_chunks(stream, n=1):
if len(chunks) >= n:
break
return chunks
def get_arch_messages(response_json):
arch_messages = []
if response_json and "metadata" in response_json:
# load arch_state from metadata
arch_state_str = response_json.get("metadata", {}).get(ARCH_STATE_HEADER, "{}")
# parse arch_state into json object
arch_state = json.loads(arch_state_str)
# load messages from arch_state
arch_messages_str = arch_state.get("messages", "[]")
# parse messages into json object
arch_messages = json.loads(arch_messages_str)
# append messages from arch gateway to history
return arch_messages
return []