Add python agent codeHEADmaster

Preliminary

3 files changed, 195 insertions, 0 deletions

diff --git a/python_agent/refactor.sh b/python_agent/refactor.sh
new file mode 100755
index 0000000..81cca70
--- /dev/null
+++ b/python_agent/refactor.sh
@@ -0,0 +1,65 @@
+#!/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
+  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 update documentation as necessary and propose any changes to better suit the code. The code is in Python.
+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)
+
+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'`"
+  SOURCE_FILE="`head -n1 \"${FILE}\"`"
+  SOURCE_FILE="${SOURCE_FILE### Function derived from }"
+
+  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
diff --git a/python_agent/to_ast.py b/python_agent/to_ast.py
new file mode 100644
index 0000000..3f5aadb
--- /dev/null
+++ b/python_agent/to_ast.py
@@ -0,0 +1,63 @@
+import hashlib
+from os.path import exists
+from pathlib import Path
+from ast import parse, unparse, FunctionDef, ClassDef
+
+def rglob(directory, extension):
+    """
+    Recursively glob all files with the given extension in the specified directory.
+
+    Args:
+    directory (str or Path): The root directory to start the search.
+    extension (str): The file extension (e.g., '.txt', '.py').
+                     The leading dot is optional but good practice
+
+    Returns:
+        list: A list of Path objects for the matching files
+    """
+    # Ensure the extension starts with a dot if not already provided
+    if not extension.startswith("."):
+        extension = "." + extension
+
+    # Use rglob to recursively find files matching the pattern
+    # The pattern should be the extension itself as rglob operates recursively
+    files = list(Path(directory).rglob(f"*{extension}"))
+    return files
+
+for file in rglob("src/", "py"):
+    source = ""
+    with open(file, "r") as fhandle:
+        source = fhandle.read()
+    tree = parse(source)
+    nodes = [node for node in tree.body if isinstance(node, (FunctionDef, ClassDef))]
+    for node in nodes:
+        src = unparse(node)
+        src = f"# Function derived from {file}\n" + src
+        if not src.endswith("\n"):
+            src = src + "\n"
+        fout = ""
+        while True:
+            srce = src.encode('utf-8')
+            src_hobj = hashlib.sha256(srce)
+            src_dig = src_hobj.hexdigest()
+            fout = f"agent/ast/{src_dig}.py"
+            if exists(fout):
+                # Check hash to see if it is the same
+                old_src = ""
+                with open(fout, "r") as fhandle:
+                    old_src = fhandle.read()
+                old_srce = old_src.encode('utf-8')
+                osrc_hobj = hashlib.sha256(old_srce)
+                osrc_dig = osrc_hobj.hexdigest()
+                if osrc_dig == src_dig:
+                    break
+                # Add something to clear up collision
+                src = src + "# Added due to hash collision\n"
+                print("Collision detected!")
+                continue
+            break
+        if exists(fout):
+            continue
+        with open(fout, "w+") as fhandle:
+            fhandle.write(src)
+
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