diff options
| -rwxr-xr-x | python_agent/refactor.sh | 65 | ||||
| -rw-r--r-- | python_agent/to_ast.py | 63 | ||||
| -rwxr-xr-x | python_agent/write_test.sh | 67 |
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 |