#!/bin/bash
NORMAL_REQ='{model: $model, messages:[{role: "user", content: $prompt}]}'
PREFILL_REQ='{model: $model, chat_template_kwargs: {"enable_thinking": false}, messages:[{role: "user", content: $prompt}, {role: "assistant", content: $prefill}]}'

ENDPOINT='http://localhost:8080/v1/chat/completions'
WORDLIST='/etc/dictionaries-common/words'

get_response() {
	model=$1
	prompt_file=$2
	prefill_file=$3
	if [ -z $prefill_file ]; then
		req=$(jq --null-input --arg model $model --rawfile prompt $prompt_file "$NORMAL_REQ")
	else
		cat $prefill_file
		req=$(jq --null-input --arg model $model --rawfile prompt $prompt_file --rawfile prefill $prefill_file "$PREFILL_REQ")
	fi

	resp=$(curl -X POST --url $ENDPOINT --silent --data @- <<< $req)
	jq --raw-output --exit-status .choices[0].message.content <<< $resp 2>&1
	
	if [ $? -ne 0 ]; then
		echo Error: $(jq --raw-output .error.message <<< $resp)
	fi
}

sort --random-sort models.txt | while read model; do
	random_name=$(grep -Po '^[A-Z][a-z]+$' $WORDLIST | sort --random-sort | head -n1)
	random_name=${random_name,}
	echo $random_name : $model >> key.txt
	for prompt_file in prompts/*.txt; do
		promptname=$(basename ${prompt_file} .txt)
		mkdir -p responses/$promptname
		for prefill_file in prefills/*.txt; do
			if [ ! -f $prefill_file ]; then
				# no prefill
				prefill_file=""
				prefill="."
			else
				# prefill included
				prefill=$(basename $prefill_file .txt)
			fi
			echo $(date +%T) $random_name: $promptname $prefill
			mkdir -p responses/$promptname/$prefill
			get_response $model $prompt_file $prefill_file > responses/$promptname/$prefill/$random_name.txt
		done
	done
done