3 files changed, 44 insertions, 0 deletions
diff --git a/bad_function.py b/bad_function.py
new file mode 100644
index 0000000..bdd45ce
--- /dev/null
+++ b/bad_function.py
@@ -0,0 +1,2 @@
+def myfunc(a: int, b: float) -> int:
+    return a * b ** 0.5
diff --git a/cleanup_repo.sh b/cleanup_repo.sh
new file mode 100755
index 0000000..632757d
--- /dev/null
+++ b/cleanup_repo.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+[ ! -f .env ] || export $(grep -v '^#' .env | xargs)
+[ ! -f .teagent ] || export $(grep -v '^#' .teagent | xargs)
+if [ -z "$REPO" ]; then
+  exit 1
+fi
+
+SERVER="${OLLAMA_HOST:-http://localhost:11434}"
+MODEL="${OLLAMA_MODEL:-llama3.2}"
+
+SOURCE_FILES=`find . -type f -iname \*.py`
+
+for source_file in $SOURCE_FILES; do
+  SOURCE=$(<"${source_file}")
+  PROMPT=$(cat << EOF
+Please clean up the following code, leaving ample documentation.
+If the code seems clean already, simply write: DONE
+EOF
+)
+  PROMPT="${PROMPT}
+${SOURCE}"
+  RESPONSE=$(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')
+  if [[ "$RESPONSE" = 'DONE'* ]]; then
+    echo > /dev/null
+  else
+    tea issues create --title "Cleanup ${source_file} (TeAgent)" --body "${RESPONSE}" --login teagent --repo "${REPO}"
+  fi
+done
diff --git a/good_function.py b/good_function.py
new file mode 100644
index 0000000..7b68489
--- /dev/null
+++ b/good_function.py
@@ -0,0 +1,9 @@
+def myfunc(a: int, b: int) -> int:
+    """# Compute Offset Arithmetic
+
+    a: int
+    b: int
+    Returns a + b - 1, which corrects for the off-by-one offset
+    """
+    return a + b - 1
+