1 files changed, 67 insertions, 0 deletions
diff --git a/python_agent/write_test.sh b/python_agent/write_test.sh
new file mode 100755
index 0000000..053ce3e
--- /dev/null
+++ b/python_agent/write_test.sh
@@ -0,0 +1,67 @@
+#!/bin/bash
+
+[ ! -f .env ] || export $(grep -v '^#' .env | xargs)
+[ ! -f .teagent ] || export $(grep -v '^#' .teagent | xargs)
+if [ -z "$REPO" ]; then
+  exit 1
+fi
+# Default 5 minute maximum
+TIMEOUT_MAX="${TIMEOUT_MAX:-300}"
+SERVER="${OLLAMA_HOST:-http://localhost:11434}"
+MODEL="${OLLAMA_MODEL:-llama3.2}"
+
+echo "${SERVER}"
+for FILE in agent/ast/*.py; do
+  SOURCE_FILE="`head -n1 \"${FILE}\"`"
+  SOURCE_FILE="${SOURCE_FILE### Function derived from }"
+  SOURCE_IMPORT="${SOURCE_FILE##src/}"
+  PROMPT=$(cat << EOF
+You are an engineer in charge of freshening up old code. For this task, you will be given a source file. Please write tests for the function to make sure any future changes continue to cover expectations. The code is in Python and the test suite is pytest.
+To do this task, you will be drafting a PR. Therefore, you will output three pieces: a 'title' for the PR, the 'discussion' of the PR, and the 'code' of the PR.
+(You may treat any called functions from the histclass package as correctly working functions. Also, imports were removed so you can assume all symbols are properly imported already)
+The function/ class can be imported from ${SOURCE_IMPORT}
+
+File contents:
+`cat "${FILE}"`
+EOF
+)
+  echo "${PROMPT}"
+  SCHEMA="json"
+  OPTS='{"num_ctx": 8192}'
+  SCHEMA='{"type": "object", "properties": {"code": {"type": "string"}, "title": {"type": "string"}, "discussion": {"type": "string"}}, "required": ["code", "title", "dicussion"]}'
+  CDATA="$(jq -n --arg model "$MODEL" --arg prompt "$PROMPT" --argjson options "$OPTS" --argjson schema "$SCHEMA" \
+    '{model:$model, prompt:$prompt, options: $options, stream:false, format: $schema}')"
+  echo "${CDATA}"
+  RESPONSE=$(timeout "${TIMEOUT_MAX}" curl -s "$SERVER/api/generate" \
+    -H "Content-Type: application/json" \
+    -d "${CDATA}")
+  #        '{model:$model, prompt:$prompt, stream:false}')" \
+  #        | jq -r '.response')
+  echo "${RESPONSE}"
+  #echo "${RESPONSE}" | jq -r '.response' | jq
+  echo "${RESPONSE}" | jq -r '.response' | jq -r '.title'
+  echo "${RESPONSE}" | jq -r '.response' | jq -r '.discussion'
+  echo "${RESPONSE}" | jq -r '.response' | jq -r '.code'
+  
+  
+  TITLE="`echo "${RESPONSE}" | jq -r '.response' | jq -r '.title'`"
+  DISCUSSION="`echo "${RESPONSE}" | jq -r '.response' | jq -r '.discussion'`"
+  CODE="`echo "${RESPONSE}" | jq -r '.response' | jq -r '.code'`"
+
+  BODY=$(cat << EOF
+File Source: ${SOURCE_FILE}
+
+# Details
+${DISCUSSION}
+
+# Code
+\`\`\`python3
+${CODE}
+\`\`\`
+EOF
+)
+  echo "${BODY}"
+  echo "${TITLE}"
+
+  tea issues create --title "${TITLE}" --body "${BODY}" --login teagent --repo "cc/HistClass"
+done