Agential workflow

autocog/prompts.py

@@ -53,18 +53,18 @@ def render(template_name, **kwargs):
 
 
 def order_paths(paths: list[Path], readme_contents: str | None) -> str:
-    return render("order_paths", paths=paths, readme_contents=readme_contents)
+    return render("order_paths/order_paths", paths=paths, readme_contents=readme_contents)
 
 
-def generate_initial(
+def generate_cog_yaml(
     files: dict[str, str],
     tell: str | None,
     predict_py: str | None,
     cog_yaml: str | None,
     package_versions: dict[set] | None,
 ) -> str:
     return render(
-        "generate_initial",
+        "cog_generation/generate",
         files=files,
         tell=tell,
         predict_py=predict_py,
@@ -73,26 +73,43 @@ def generate_initial(
     )
 
 
+def generate_predict_py(
+    files: dict[str, str],
+    tell: str | None,
+    predict_py: str | None,
+) -> str:
+    return render(
+        "predict_generation/generate",
+        files=files,
+        tell=tell,
+        predict_py=predict_py
+    )
+
+
 def diagnose_error(predict_command: str, error: str) -> str:
-    return render("diagnose_error", predict_command=predict_command, error=error)
+    return render("error_diagnosis/diagnose_error", predict_command=predict_command, error=error)
 
 
 def package_error(predict_command: str, error: str) -> str:
-    return render("package_error", predict_command=predict_command, error=error)
+    return render("error_diagnosis/package_error", predict_command=predict_command, error=error)
 
 
-def get_packages(cog_contents: str | None) -> str:
-    return render("get_packages", cog_contents=cog_contents)
+def get_packages(packages: list[str] | None, cog_contents: str | None) -> str:
+    return render("package_info/get_packages", packages=packages, cog_contents=cog_contents)
 
 
-def get_packages_versions(packages: str | None) -> str:
-    return render("get_packages_versions", packages=packages)
+def cog_predict(predict_py: str) -> str:
+    return render("cog_predict/generate", predict_py=predict_py)
 
 
-system = render("system")
-cog_predict = render("cog_predict")
-fix_predict_py = render("fix_predict_py")
-fix_cog_yaml = render("fix_cog_yaml")
+order_paths_system = render("order_paths/system")
+package_info_system = render("package_info/system")
+cog_generation_system = render("cog_generation/system")
+predict_generation_system = render("predict_generation/system")
+error_diagnosis_system = render("error_diagnosis/system")
+cog_fixing_system = render("cog_fixing/system")
+predict_fixing_system = render("predict_fixing/system")
+cog_predict_system = render("cog_predict/system")
 
 
 def file_start(filename):

autocog/prompts/generate_initial.tpl → autocog/prompts/cog_generation/generate.tpl

@@ -25,9 +25,9 @@ Here is a list of packages and valid versions:
 {% endfor %}
 {% endif %}
 
-Given the files above, {% if predict_py %}update predict.py{% else %}generate a predict.py file{% endif %} and {% if cog_yaml %}update cog.yaml{% else %}generate a cog.yaml file{% endif %}.
+Given the files above, {% if cog_yaml %}update cog.yaml{% else %}generate a cog.yaml file{% endif %}. Ensure that all Python packages must have pinned versions. Add short comments to describe what parts of the code made you decide on the different parts of cog.yaml. 
 
-In cog.yaml, ensure that all Python packages must have pinned versions. Also in cog.yaml, add short comments to describe what parts of the code made you decide on the different parts of cog.yaml. Wrap the contents of both files in the strings '{% file_start "filename" %}' and '{% file_end "filename" %}'. Don't output any other text before or after the files since I intend to execute the output that you generate in a Python programming environment.
+Wrap the contents of both files in the strings '{% file_start "filename" %}' and '{% file_end "filename" %}'. Don't output any other text before or after the files since I intend to execute the output that you generate in a Python programming environment.
 
 {% if tell %}
 Also make sure to follow these additional instructions: {{ tell }}

autocog/prompts/cog_generation/system.tpl

@@ -0,0 +1,25 @@
+You are an expert Python machine learning developer. Your task is to "cogify" a repository of Python files into a Cog model that can be deployed on Replicate.com.
+
+These are the Cog YAML docs for cog.yaml:
+
+{% include "cog_yaml_docs.tpl" %}
+
+---
+
+Below is an example cog.yaml:
+
+{% file_start "cog.yaml" %}
+
+build:
+  gpu: true
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+  python_version: "3.8"
+  python_packages:
+    - "torch==1.8.1"
+predict: "predict.py:Predictor"
+
+{% file_end "cog.yaml" %}
+
+Your job is to generate a cog.yaml file.

autocog/prompts/cog_predict/generate.tpl

@@ -0,0 +1,2 @@
+Here is the predict.py you generated:
+{{ predict_py }}

autocog/prompts/cog_predict/system.tpl

@@ -0,0 +1,5 @@
+Below is an example of a cog predict command:
+
+cog predict -i input1=@input.jpg -i input2=foo
+
+Return a cog predict command for the predict.py file that you generated. Wrap the command in the strings {{ command_start }} and {{ command_end }}. Don't output any other text before or after the strings since I intend to execute the output that you generate in a Python programming environment.

autocog/prompts/error_diagnosis/system.tpl

@@ -0,0 +1 @@
+You are an expert debugger.

autocog/prompts/order_paths/order_paths.tpl

@@ -0,0 +1,12 @@
+Here are the paths:
+
+{% for path in paths %}
+{{ path }}
+{% endfor %}
+
+End of paths.
+{% if readme_contents %}
+Below is the readme:
+
+{{ readme_contents }}
+{% endif %}

autocog/prompts/order_paths.tpl → autocog/prompts/order_paths/system.tpl

@@ -7,16 +7,3 @@ second_most_relevant.py
 third_most_relevant.py
 [...]
 least_relevant.py
-
-Here are the paths:
-
-{% for path in paths %}
-{{ path }}
-{% endfor %}
-
-End of paths.
-{% if readme_contents %}
-Below is the readme:
-
-{{ readme_contents }}
-{% endif %}

autocog/prompts/package_info/get_packages.tpl

@@ -0,0 +1,13 @@
+{% if packages %}
+Here are all of the imported packages:
+{% for package in packages %}
+- {{ package }}
+{% endfor %}
+{% else %}
+No packages were found.
+{% endif %}
+
+{% if cog_content %}
+Here are the contents of cog.yaml: 
+{{ cog_contents }}
+{% endif %}

autocog/prompts/get_packages.tpl → autocog/prompts/package_info/system.tpl

@@ -6,9 +6,4 @@ package3==v3.v3.v3
 [...]
 packagen
 
-{% if cog_content %}
-Here are the contents of cog.yaml: 
-{{ cog_contents }}
-{% endif %}
-
 Don't output anything else since I intend to parse the output and use it in a programmatic pipeline.

autocog/prompts/predict_generation/generate.tpl

@@ -0,0 +1,21 @@
+Below are the contents of the relevant files in the repository:
+
+{% for filename, contents in files.items() %}
+{% file_start filename %}
+{{ contents }}
+{% file_end filename %}
+{% endfor %}
+
+{% if predict_py %}
+{% file_start "predict.py" %}
+{{ predict_py }}
+{% file_end "predict.py" %}
+{% endif %}
+
+Given the files above, {% if predict_py %}update predict.py{% else %}generate a predict.py file{% endif %}.
+
+Wrap the contents of both files in the strings '{% file_start "filename" %}' and '{% file_end "filename" %}'. Don't output any other text before or after the files since I intend to execute the output that you generate in a Python programming environment.
+
+{% if tell %}
+Also make sure to follow these additional instructions: {{ tell }}
+{% endif %}

autocog/prompts/system.tpl → autocog/prompts/predict_generation/system.tpl

@@ -6,31 +6,7 @@ These are the Cog Python docs for predict.py:
 
 ---
 
-These are the Cog YAML docs for cog.yaml:
-
-{% include "cog_yaml_docs.tpl" %}
-
----
-
-Below is an example cog.yaml:
-
-{% file_start "cog.yaml" %}
-
-build:
-  gpu: true
-  system_packages:
-    - "libgl1-mesa-glx"
-    - "libglib2.0-0"
-  python_version: "3.8"
-  python_packages:
-    - "torch==1.8.1"
-predict: "predict.py:Predictor"
-
-{% file_end "cog.yaml" %}
-
----
-
-Below is an example predit.py:
+Below is an example predict.py:
 
 {% file_start "predict.py" %}
 
@@ -93,3 +69,5 @@ class Predictor(BasePredictor):
         return output_path
 
 {% file_end "predict.py" %}
+
+Your job is to generate a predict.py file.

setup.py

@@ -9,7 +9,7 @@
     name="autocog",
     long_description=long_description,
     long_description_content_type="text/markdown",
-    version="0.0.11",
+    version="0.0.12",
     url="https://github.com/andreasjansson/AutoCog",
     packages=find_packages(),
     install_requires=[