1 files changed, 116 insertions, 0 deletions
diff --git a/agent.sh b/agent.sh
new file mode 100755
index 0000000..597571c
--- /dev/null
+++ b/agent.sh
@@ -0,0 +1,116 @@
+#!/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}"
+N_RETRIES="${RETRY_MAX:-5}"
+
+teaissue() {
+  echo tea issues create --title "$1 (TeAgent)" --body "$2" --login teagent --repo "${REPO}"
+}
+
+SERVER="${OLLAMA_HOST:-http://localhost:11434}"
+MODEL="${OLLAMA_MODEL:-llama3.2}"
+
+read -p "Source file listing cmd: " -e -i "find . -type f -iname \\\\*.py" source_file_fn
+SOURCE_FILES=`eval "${source_file_fn}"`
+
+# Select what files
+FILE_SELECTIONS=()
+INDEX=1
+for source_file in ${SOURCE_FILES}; do
+  if [ ! -s "${source_file}" ]; then
+    continue
+  fi
+  FILE_SELECTIONS+=("${source_file}" "" off)
+  INDEX=$((INDEX+1))
+done
+
+FILE_CHOICES=$(dialog --clear \
+  --title "Files to process" \
+  --checklist "Choose files:" \
+  10 40 3 \
+  "${FILE_SELECTIONS[@]}" \
+  2>&1 >/dev/tty)
+if [ "$?" -ne "0" ]; then
+  exit 1
+fi
+
+TASK_SELECTIONS=(
+  1 "Clean" off
+  2 "UnitTest" off
+  3 "Cancel" on
+)
+
+TASK_CHOICE=$(dialog --clear \
+  --title "Agent Task" \
+  --radiolist "Select Task:" \
+  15 40 3 \
+  "${TASK_SELECTIONS[@]}" \
+  2>&1 >/dev/tty)
+if [ "$?" -ne "0" ]; then
+  exit 1
+fi
+
+case $TASK_CHOICE in
+  1)
+    TASK_NAME="Tests"
+    PROMPT_BASE=$(cat << EOF
+Please clean up the following code, leaving ample documentation.
+If the code seems clean already, simply write: DONE
+EOF
+)
+    ;;
+  2)
+    TASK_NAME="Cleanup"
+    PROMPT_BASE=$(cat << EOF
+Please write tests for the following source, leaving ample documentation.
+If there don't seem to be any natural tests, simply write: DONE
+EOF
+)
+      ;;
+  3)
+    echo Bye!
+    exit 0
+    ;;
+  *)
+    echo "Invalid choice!"
+    exit 1
+    ;;
+esac
+echo "$PROMPT_BASE"
+
+for source_file in $FILE_CHOICES; do
+  if [ ! -s "${source_file}" ]; then
+    continue
+  fi
+  echo $source_file;
+  SOURCE=$(<"${source_file}")
+  PROMPT="${PROMPT_BASE}
+${SOURCE}"
+  TRIES=0
+  while :; do
+    echo Sending over for "${source_file}"
+    RESPONSE=$(timeout "${TIMEOUT_MAX}" curl -s "$SERVER/api/generate" \
+      -H "Content-Type: application/json" \
+      -d "$(jq -n --arg model "$MODEL" --arg prompt "$PROMPT" \
+            '{model:$model, prompt:$prompt, stream:false}')" \
+            | jq -r '.response')
+    TRIES="$((TRIES+1))"
+    if [ "$?" -eq "0" ]; then
+      echo "Response success!"
+      break
+    fi
+    if [ "$TRIES" -ge "$N_RETRIES" ]; then
+      echo "Response timeout!"
+      break
+    fi
+  done
+  if [ ! -z "${RESPONSE}" ]; then
+    teaissue "${TASK_NAME} ${source_file}" "${RESPONSE}"
+  fi
+done