diff --git a/autocog/agents.py b/autocog/agents.py
new file mode 100644
index 0000000..67617a0
--- /dev/null
+++ b/autocog/agents.py
@@ -0,0 +1,388 @@
+from pathlib import Path
+import sys
+import os
+import ast
+import subprocess
+import requests
+import importlib
+import re
+from packaging import version
+from pypi_simple import errors as pypi_errors
+from pypi_simple import PyPISimple
+
+from .ai import AI
+from .retry import retry
+from .testdata import create_empty_file
+from . import prompts
+
+
+class BaseAgent:
+    def __init__(self, ai_provider: str, api_key: str, system_prompt: str, chat_history_path: Path):
+        """
+        Base class for all agents that interact with an AI system.
+
+        :param ai_provider: The provider (e.g., 'openai', 'anthropic')
+        :param api_key: The API key for the AI provider
+        :param system_prompt: The system prompt for this agent
+        :param chat_history_path: The path where this agent's chat history is stored
+        """
+        self.ai = AI(system_prompt=system_prompt, provider=ai_provider, api_key=api_key, chat_history_path=chat_history_path)
+        self.system_prompt = system_prompt
+        self.chat_history_path = chat_history_path
+
+
+class PathOrderingAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for ordering paths of Python files based on importance.
+        """
+        super().__init__(ai_provider, api_key, prompts.order_paths_system, chat_history_path)
+
+    def order_paths(self, repo_path: Path, readme_contents: str = None) -> list[Path]:
+        paths = self._find_python_files(repo_path)
+        if len(paths) == 0:
+            raise ValueError(f"{repo_path} has no Python files")
+
+        if readme_contents is None:
+            _, readme_contents = self._load_readme_contents(repo_path)
+
+        content = self.ai.call(prompts.order_paths(paths=paths, readme_contents=readme_contents))
+        ordered_paths = [repo_path / Path(p) for p in content.strip().splitlines()]
+
+        if set(ordered_paths) - set(paths):
+            raise ValueError("Failed to order paths")
+
+        return ordered_paths
+
+    def _find_python_files(self, repo_path: Path) -> list[Path]:
+        return [path for path in repo_path.rglob("*.py")]
+
+    def _load_readme_contents(self, repo_path: Path) -> tuple[str, str] | tuple[None, None]:
+        readme_filenames = ["README.md", "readme.md", "README.txt", "readme.txt", "README"]
+        for filename in readme_filenames:
+            readme_path = repo_path / filename
+            if readme_path.exists():
+                return filename, readme_path.read_text()
+        return None, None
+
+
+class DocsPullingAgent:
+    def __init__(self, prompts_dir: Path):
+        self.prompts_dir = prompts_dir
+
+    def pull_docs(self):
+        print("Pulling documentation...")
+        base_dir = os.path.dirname(__file__)
+        prompts_dir = os.path.join(base_dir, "prompts")
+
+        self._fetch_and_save(prompts.COG_DOCS, os.path.join(prompts_dir, "cog_yaml_docs.tpl"))
+        self._fetch_and_save(prompts.PREDICT_DOCS, os.path.join(prompts_dir, "cog_python_docs.tpl"))
+
+    def _fetch_and_save(self, url, save_path):
+        response = requests.get(url)
+        if response.status_code == 200:
+            with open(save_path, 'wb') as f:
+                f.write(response.content)
+            print(f"Successfully pulled down documentation from {url}")
+        else:
+            print(f"Failed to download documentation from {url}")
+
+
+class PackageInfoAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for gathering information about required packages from PyPI.
+        """
+        super().__init__(ai_provider, api_key, prompts.package_info_system, chat_history_path)
+
+    def get_packages_info(self, packages: list[str], repo_path: Path):
+        print("Getting package info...")
+        cog_yaml_path = repo_path / "cog.yaml"
+        cog_yaml = cog_yaml_path.read_text() if cog_yaml_path.exists() else None
+        content = self.ai.call(prompts.get_packages(packages=packages, cog_contents=cog_yaml))
+
+        package_info = {}
+        client = PyPISimple()
+        for package in content.strip().split('\n'):
+            versions = self._get_package_versions(client, package)
+            if versions:
+                package_info[package] = sorted(versions, key=version.parse)
+
+        return package_info
+
+    def _get_package_versions(self, client, package):
+        versions = set()
+        if '==' not in package:
+            try:
+                packages_info = client.get_project_page(package).packages
+                for p_info in packages_info:
+                    versions.add(p_info.version)
+            except pypi_errors.NoSuchProjectError:
+                return None
+        else:
+            versions.add(package.split('==')[1])
+        return versions
+
+    def get_imported_packages(self, ordered_paths: list[Path]) -> list[str]:
+        """
+        Reads the import statements from the most important Python files and returns a list of non-standard library packages.
+
+        :param ordered_paths: List of ordered paths to the most important Python files.
+        :return: List of non-standard Python packages used in the project.
+        """
+        imported_packages = set()
+
+        # Iterate over each important file and extract the imports
+        for file_path in ordered_paths:
+            if file_path.suffix == '.py':  # Ensure it's a Python file
+                imports = self._extract_imports_from_file(file_path)
+                imported_packages.update(imports)
+
+        # Filter out standard library packages
+        non_standard_packages = self._filter_standard_libraries(imported_packages)
+
+        return sorted(non_standard_packages)
+
+    def _extract_imports_from_file(self, file_path: Path) -> set[str]:
+        """
+        Extracts import statements from a Python file.
+
+        :param file_path: Path to the Python file.
+        :return: Set of imported modules/packages.
+        """
+        with file_path.open('r', encoding='utf-8') as file:
+            tree = ast.parse(file.read(), filename=str(file_path))
+
+        imports = set()
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    imports.add(alias.name.split('.')[0])  # Get the root module/package
+            elif isinstance(node, ast.ImportFrom):
+                if node.module:  # Handle "from <module> import <name>"
+                    imports.add(node.module.split('.')[0])
+
+        return imports
+
+    def _filter_standard_libraries(self, packages: set[str]) -> list[str]:
+        """
+        Filters out the standard library packages from the set of imported packages.
+
+        :param packages: Set of package names.
+        :return: List of non-standard library package names.
+        """
+        non_standard_packages = []
+
+        for package in packages:
+            if not self._is_standard_library(package):
+                non_standard_packages.append(package)
+
+        return non_standard_packages
+
+    def _is_standard_library(self, package: str) -> bool:
+        """
+        Checks if a given package is part of the Python standard library.
+
+        :param package: Package name to check.
+        :return: True if the package is a standard library package, False otherwise.
+        """
+        # Try to locate the package using importlib to see if it's part of the standard library
+        spec = importlib.util.find_spec(package)
+        if spec is None:
+            return False
+        # If the package exists, check if it's a built-in module
+        return package in sys.builtin_module_names
+
+
+class CogGenerationAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for generating the cog.yaml file.
+        """
+        super().__init__(ai_provider, api_key, prompts.cog_generation_system, chat_history_path)
+
+    def generate_cog_yaml(self, repo_path: Path, predict_py: str | None, cog_yaml: str | None, package_versions: dict | None, tell: str | None = None) -> str:
+        """
+        Generates cog.yaml file based on the current project setup.
+        
+        :param repo_path: Path to the project repository.
+        :param tell: Optional additional information for generating the file.
+        :return: The cog.yaml content.
+        """
+        print("Generating cog.yaml...")
+        files = self._gather_project_files(repo_path)
+        # Call the AI to generate the cog.yaml file
+        content = self.ai.call(prompts.generate_cog_yaml(files=files, predict_py=predict_py, cog_yaml=cog_yaml, package_versions=package_versions, tell=tell))
+
+        # Extract the cog.yaml from the response
+        cog_yaml = self._extract_file_from_gpt_response(content, "cog.yaml")
+        return cog_yaml
+
+    def _gather_project_files(self, repo_path: Path) -> dict[str, str]:
+        """
+        Gathers all relevant project files (e.g., README, requirements) that might help generate cog.yaml.
+        
+        :param repo_path: Path to the repository.
+        :return: Dictionary of file names and their contents.
+        """
+        files = {}
+        # Gather README, requirements.txt, pyproject.toml if they exist
+        for filename in ["README.md", "requirements.txt", "pyproject.toml"]:
+            file_path = repo_path / filename
+            if file_path.exists():
+                files[filename] = file_path.read_text()
+        return files
+
+    def _extract_file_from_gpt_response(self, content: str, filename: str) -> str:
+        """
+        Extracts the content of a specific file from the AI's response.
+        
+        :param content: The full response from the AI.
+        :param filename: The name of the file to extract (e.g., cog.yaml).
+        :return: The extracted content of the file.
+        """
+        pattern = re.compile(rf"(?<=-- FILE_START: {filename})(.*?)(?=-- FILE_END: {filename})", re.DOTALL)
+        match = pattern.search(content)
+        if not match:
+            raise ValueError(f"Failed to generate {filename}")
+        return match.group(1).strip()
+
+
+class PredictGenerationAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for generating the predict.py file.
+        """
+        super().__init__(ai_provider, api_key, prompts.predict_generation_system, chat_history_path)
+
+    def generate_predict_py(self, repo_path: Path, predict_py: str | None, tell: str | None = None) -> str:
+        """
+        Generates predict.py file based on the current project setup.
+        
+        :param repo_path: Path to the project repository.
+        :param tell: Optional additional information for generating the file.
+        :return: The predict.py content.
+        """
+        print("Generating predict.py...")
+        files = self._gather_project_files(repo_path)
+        # Call the AI to generate the predict.py file
+        content = self.ai.call(prompts.generate_predict_py(files=files, predict_py=predict_py, tell=tell))
+
+        # Extract the predict.py from the response
+        predict_py = self._extract_file_from_gpt_response(content, "predict.py")
+        return predict_py
+
+    def _gather_project_files(self, repo_path: Path) -> dict[str, str]:
+        """
+        Gathers all relevant project files (e.g., README, requirements) that might help generate predict.py.
+        
+        :param repo_path: Path to the repository.
+        :return: Dictionary of file names and their contents.
+        """
+        files = {}
+        # Gather README, requirements.txt, pyproject.toml if they exist
+        for filename in ["README.md", "requirements.txt", "pyproject.toml"]:
+            file_path = repo_path / filename
+            if file_path.exists():
+                files[filename] = file_path.read_text()
+        return files
+
+    def _extract_file_from_gpt_response(self, content: str, filename: str) -> str:
+        """
+        Extracts the content of a specific file from the AI's response.
+        
+        :param content: The full response from the AI.
+        :param filename: The name of the file to extract (e.g., predict.py).
+        :return: The extracted content of the file.
+        """
+        pattern = re.compile(rf"(?<=-- FILE_START: {filename})(.*?)(?=-- FILE_END: {filename})", re.DOTALL)
+        match = pattern.search(content)
+        if not match:
+            raise ValueError(f"Failed to generate {filename}")
+        return match.group(1).strip()
+
+
+class ErrorDiagnosisAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for diagnosing errors in the project.
+        """
+        super().__init__(ai_provider, api_key, prompts.error_diagnosis_system, chat_history_path)
+
+    @retry(5)
+    def diagnose_error(self, predict_command: str, error: str) -> tuple[str, bool]:
+        diagnosis = self.ai.call(prompts.diagnose_error(predict_command=predict_command, error=self._truncate_error(error)))
+        package_error = self.ai.call(prompts.package_error(predict_command=predict_command, error=self._truncate_error(error)))
+        return diagnosis, package_error == "True"
+
+    def _truncate_error(self, error, max_length=10000):
+        return error[:max_length]
+
+
+
+class CogFixingAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for fixing issues in configuration files.
+        """
+        super().__init__(ai_provider, api_key, prompts.cog_fixing_system, chat_history_path)
+
+    @retry(5)
+    def fix_cog_yaml(self) -> str:
+        response = self.ai.call(prompts.fix_cog_yaml)
+        return self._file_from_gpt_response(response, "cog.yaml")
+
+    def _file_from_gpt_response(self, content: str, filename: str) -> str:
+        pattern = re.compile(rf"(?<={file_start(filename)})(.*?)(?={file_end(filename)})", re.MULTILINE | re.DOTALL)
+        match = pattern.search(content)
+        if not match:
+            raise ValueError(f"Failed to generate {filename}")
+        return match[1].strip()
+
+
+class PredictFixingAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for fixing issues in configuration files.
+        """
+        super().__init__(ai_provider, api_key, prompts.predict_fixing_system, chat_history_path)
+
+    @retry(5)
+    def fix_predict_py(self) -> str:
+        response = self.ai.call(prompts.fix_predict_py)
+        return self._file_from_gpt_response(response, "predict.py")
+
+    def _file_from_gpt_response(self, content: str, filename: str) -> str:
+        pattern = re.compile(rf"(?<={file_start(filename)})(.*?)(?={file_end(filename)})", re.MULTILINE | re.DOTALL)
+        match = pattern.search(content)
+        if not match:
+            raise ValueError(f"Failed to generate {filename}")
+        return match[1].strip()
+
+
+class CogPredictAgent(BaseAgent):
+    def __init__(self, ai_provider: str, api_key: str, chat_history_path: Path):
+        """
+        Agent responsible for generating and executing the cog predict command.
+        """
+        super().__init__(ai_provider, api_key, prompts.cog_predict_system, chat_history_path)
+
+    def generate_predict_command(self, predict_py: str) -> str:
+        return self.ai.call(prompts.cog_predict(predict_py))
+
+    def run_cog_predict(self, repo_path: Path, predict_command: str) -> tuple[bool, str]:
+        proc = subprocess.Popen(predict_command, cwd=repo_path, stderr=subprocess.PIPE, shell=True)
+        stderr = self._collect_stderr(proc)
+        if proc.returncode == 0 and "Traceback (most recent call last)" not in stderr:
+            return True, stderr
+        return False, stderr
+
+    def _collect_stderr(self, proc):
+        stderr = ""
+        for line in proc.stderr:
+            line = line.decode()
+            sys.stderr.write(line)
+            stderr += line
+        return stderr
+
diff --git a/autocog/autocog.py b/autocog/autocog.py
index 7f5fce4..07b35fe 100644
--- a/autocog/autocog.py
+++ b/autocog/autocog.py
@@ -1,282 +1,57 @@
 from pathlib import Path
 import sys
-import re
-import click
 import os
-import subprocess
-import requests
-from packaging import version
-from pypi_simple import errors as pypi_errors
-from pypi_simple import PyPISimple
-
-from .ai import AI
-from . import prompts
-from .prompts import (
-    COG_DOCS,
-    PREDICT_DOCS,
-    file_start,
-    file_end,
-    COMMAND_START,
-    COMMAND_END,
-    ERROR_COG_PREDICT,
-    ERROR_PREDICT_PY,
-    ERROR_COG_YAML
+import click
+from .agents import (
+    PathOrderingAgent,
+    DocsPullingAgent,
+    PackageInfoAgent,
+    CogGenerationAgent,
+    PredictGenerationAgent,
+    ErrorDiagnosisAgent,
+    CogFixingAgent,
+    PredictFixingAgent,
+    CogPredictAgent
 )
-from .retry import retry
-from .testdata import create_empty_file
-
-
-def truncate_error(error, max_length=10000):
-    return error[:max_length]
-
-
-def order_paths(
-    ai: AI, repo_path: Path, readme_contents: str | None = None
-) -> list[Path]:
-    paths = find_python_files(repo_path)
-    if len(paths) == 0:
-        raise ValueError(f"{repo_path} has no Python files")
-
-    print("Ordering files based on importance...", file=sys.stderr)
-
-    if readme_contents is None:
-        _, readme_contents = load_readme_contents(repo_path)
-
-    content = ai.call(prompts.order_paths(paths=paths, readme_contents=readme_contents))
-
-    ordered_paths = [Path(p) for p in content.strip().splitlines()]
-    
-    for i, path in enumerate(ordered_paths):
-        ordered_paths[i] = repo_path / path
-
-    if set(ordered_paths) - set(paths):
-        raise ValueError("Failed to order paths")
-
-    return ordered_paths
-
-
-def pull_docs():
-    base_dir = os.path.dirname(__file__)
-    prompts_dir = os.path.join(base_dir, "prompts")
-
-    cog_docs = requests.get(COG_DOCS)
-    if cog_docs.status_code == 200:
-        print("Successfully pulled down documentation for cog.yaml")
-        with open(os.path.join(prompts_dir, "cog_yaml_docs.tpl"), 'wb') as f:
-            f.write(cog_docs.content)
-    else:
-        print("Failed to download cog.yaml documentation")
-
-    predict_docs = requests.get(PREDICT_DOCS)
-    if predict_docs.status_code == 200:
-        print("Successfully pulled down documentation for predict.py")
-        with open(os.path.join(prompts_dir, "cog_python_docs.tpl"), 'wb') as f:
-            f.write(predict_docs.content)
-    else:
-        print("Failed to download predict.py documentation")
-
-
-def load_readme_contents(repo_path: Path) -> tuple[str, str] | tuple[None, None]:
-    readme_filenames = ["README.md", "readme.md", "README.txt", "readme.txt", "README"]
-    for filename in readme_filenames:
-        readme_path = repo_path / filename
-        if readme_path.exists():
-            return filename, readme_path.read_text()
-    return None, None
-
-
-def get_packages_info(ai: AI, repo_path: Path):
-    cog_yaml_path = repo_path / "cog.yaml"
-    if cog_yaml_path.exists():
-        cog_yaml = cog_yaml_path.read_text()
-    else:
-        cog_yaml = None
-    content = ai.call(prompts.get_packages(cog_contents=cog_yaml))
-
-    # Initialize PyPI client
-    client = PyPISimple()
-    # Get package information
-    package_info = {}
-    for package in content.strip().split('\n'):
-        valid = True
-        versions = set()
-        if '==' not in package:
-            # If no version is explicitly given, query PyPi
-            try:
-                packages_info = client.get_project_page(package).packages
-                for p_info in packages_info:
-                    versions.add(p_info.version)
-            except pypi_errors.NoSuchProjectError:
-                valid = False
-        else:
-            # If version is explicitly given
-            package_version = package.split('==')[1]
-            versions.add(package_version)
-        if valid:
-            package_info[package] = sorted(versions, key=version.parse)
-    return package_info
-
-
-@retry(3)
-def generate_initial(
-    ai: AI, repo_path: Path, paths: list[Path], tell: str | None
-) -> tuple[str, str]:
-    files = {}
-    readme_filename, readme_contents = load_readme_contents(repo_path)
-    if readme_filename:
-        files[readme_filename] = readme_contents
-
-    requirements_file = repo_path / "requirements.txt"
-    if requirements_file.exists():
-        files["requirements.txt"] = requirements_file.read_text()
-        package_versions = None
-    else:
-        print("Getting package information...")
-        package_versions = get_packages_info(ai, repo_path)
-
-    poetry_file = repo_path / "pyproject.toml"
-    if poetry_file.exists():
-        files["pyproject.toml"] = poetry_file.read_text()
-
-    for path in paths:
-        files[path.name] = path.read_text()
-
-    predict_py_path = repo_path / "predict.py"
-    if predict_py_path.exists():
-        predict_py = predict_py_path.read_text()
-    else:
-        predict_py = None
-    cog_yaml_path = repo_path / "cog.yaml"
-    if cog_yaml_path.exists():
-        cog_yaml = cog_yaml_path.read_text()
-    else:
-        cog_yaml = None
-
-    content = ai.call(
-        prompts.generate_initial(
-            files=files, tell=tell, predict_py=predict_py, cog_yaml=cog_yaml, package_versions=package_versions
-        )
-    )
-    cog_yaml = file_from_gpt_response(content, "cog.yaml")
-    predict_py = file_from_gpt_response(content, "predict.py")
-
-    return cog_yaml, predict_py
-
-
-def find_python_files(repo_path: Path) -> list[Path]:
-    python_files = [path for path in repo_path.rglob("*.py")]
-    return python_files
-
-
-def file_from_gpt_response(content: str, filename: str) -> str:
-    pattern = re.compile(
-        rf"(?<={file_start(filename)})(?:\n```[a-z]*\n)?(.*?)(?:\n```\n)?(?={file_end(filename)})",
-        re.MULTILINE | re.DOTALL,
-    )
-    matches = pattern.search(content)
-    if not matches:
-        raise ValueError(f"Failed to generate {filename}")
-    return matches[1].strip()
-
-
-def write_files(repo_path: Path, files: dict):
-    for filename, content in files.items():
-        file_path = repo_path / filename
-        file_path.write_text(content)
-
-
-def run_cog_predict(repo_path: Path, predict_command: str) -> tuple[bool, str]:
-    print(predict_command, file=sys.stderr)
-
-    proc = subprocess.Popen(
-        predict_command, cwd=repo_path, stderr=subprocess.PIPE, shell=True
-    )
-    stderr = ""
-    assert proc.stderr
-    for line in proc.stderr:
-        line = line.decode()
-        sys.stderr.write(line)
-        stderr += line
-
-        if "Model setup failed" in line:
-            proc.kill()
-            break
-
-    proc.wait()
-    # cog predict will return 0 if the model fails internally
-    if proc.returncode == 0 and "Traceback (most recent call last)" not in stderr:
-        return True, stderr
-
-    return False, stderr
-
-
-def cog_predict_from_gpt_response(content: str) -> str:
-    pattern = re.compile(
-        rf"(?<={re.escape(COMMAND_START)}\n)([\s\S]*?)(?=\n{re.escape(COMMAND_END)})",
-        re.MULTILINE | re.DOTALL,
-    )
-    matches = pattern.search(content)
-    if not matches:
-        raise ValueError(f"Failed to generate cog predict")
-    return matches[1].strip()
-
-
-def create_files_for_predict_command(repo_path: Path, predict_command: str) -> str:
-    print("Parsing predict command...")
-    predict_command = cog_predict_from_gpt_response(predict_command)
-    file_inputs = re.findall(r"@([\w.]+)", predict_command)
-
-    for filename in file_inputs:
-        if not os.path.exists(filename):
-            tmp_path = os.path.join("/tmp", os.path.basename(filename))
-            predict_command = predict_command.replace("@" + filename, "@" + tmp_path)
-            create_empty_file(repo_path, tmp_path)
-
-    return predict_command
-
-
-def parse_cog_predict_error(stderr: str, *, max_length=20000) -> str:
-    if "Running prediction...\n" in stderr:
-        error = stderr.split("Running prediction...\n")[1].split("panic: ")[0]
-    else:
-        error = stderr.split("panic: ")[0]
-
-    return error[-max_length:]
-
-
-@retry(5)
-def diagnose_error(ai: AI, predict_command: str, error: str) -> str:
-    print("Diagnosing source of error: ", file=sys.stderr)
-
-    diagnose_text = ai.call(prompts.diagnose_error(predict_command=predict_command, error=truncate_error(error)))
-    package_error = ai.call(prompts.package_error(predict_command=predict_command, error=truncate_error(error)))
-    package_error = package_error == "True"
-
-    if diagnose_text not in [ERROR_PREDICT_PY, ERROR_COG_PREDICT, ERROR_COG_YAML]:
-        raise ValueError("Failed to diagnose error")
-    return diagnose_text, package_error
-
-
-@retry(5)
-def fix_predict_py(ai: AI) -> str:
-    text = ai.call(prompts.fix_predict_py)
-    return file_from_gpt_response(text, "predict.py")
-
+from . import prompts
 
-@retry(5)
-def fix_cog_yaml(ai: AI) -> str:
-    text = ai.call(prompts.fix_cog_yaml)
-    return file_from_gpt_response(text, "cog.yaml")
 
+def initialize_project(repo_path: Path):
+    """
+    Initializes the project by removing cog.yaml, predict.py, and clearing the chat history for all agents.
 
-def initialize_project(ai: AI, repo_path: Path):
+    :param repo_path: Path to the project repository.
+    """
+    # Define the paths for cog.yaml and predict.py
     cog_yaml_path = repo_path / "cog.yaml"
     predict_py_path = repo_path / "predict.py"
+    
+    # Remove cog.yaml and predict.py if they exist
     if cog_yaml_path.exists():
         cog_yaml_path.unlink()
+        print(f"Removed {cog_yaml_path}")
+    
     if predict_py_path.exists():
         predict_py_path.unlink()
-    ai.clear_history()
+        print(f"Removed {predict_py_path}")
+
+    # Define the chat history files for each agent
+    chat_history_files = [
+        repo_path / "path_ordering.chat",
+        repo_path / "package_info.chat",
+        repo_path / "cog_generation.chat",
+        repo_path / "predict_generation.chat",
+        repo_path / "error_diagnosis.chat",
+        repo_path / "cog_fixing.chat",
+        repo_path / "predict_fixing.chat",
+        repo_path / "cog_predict.chat"
+    ]
+
+    # Remove the chat history files if they exist
+    for history_file in chat_history_files:
+        if history_file.exists():
+            history_file.unlink()
+            print(f"Cleared chat history: {history_file}")
 
 
 @click.command()
@@ -333,42 +108,105 @@ def autocog(
     initialize: bool,
 ):
     repo_path = repo or Path(os.getcwd())
-    pull_docs()
-    ai = AI(
-        system_prompt=prompts.system,
-        provider=ai_provider,
+    
+    # 1. Initialize project if specified
+    if initialize:
+        initialize_project(repo_path)
+
+    # 2. Instantiate the agents with separate system prompts and chat history paths
+    path_ordering_agent = PathOrderingAgent(
+        ai_provider=ai_provider,
         api_key=api_key,
-        chat_history_path=repo_path / "autocog.chat",
+        chat_history_path=repo_path / "path_ordering.chat"
     )
 
-    if initialize:
-        initialize_project(ai, repo_path)
+    package_agent = PackageInfoAgent(
+        ai_provider=ai_provider,
+        api_key=api_key,
+        chat_history_path=repo_path / "package_info.chat"
+    )
+
+    cog_generation_agent = CogGenerationAgent(
+        ai_provider=ai_provider,
+        api_key=api_key,
+        chat_history_path=repo_path / "cog_generation.chat"
+    )
+    
+    predict_generation_agent = PredictGenerationAgent(
+        ai_provider=ai_provider,
+        api_key=api_key,
+        chat_history_path=repo_path / "predict_generation.chat"
+    )
+
+    error_diagnosis_agent = ErrorDiagnosisAgent(
+        ai_provider=ai_provider,
+        api_key=api_key,
+        chat_history_path=repo_path / "error_diagnosis.chat"
+    )
+
+    cog_fixing_agent = CogFixingAgent(
+        ai_provider=ai_provider,
+        api_key=api_key,
+        chat_history_path=repo_path / "cog_fixing.chat"
+    )
+    
+    predict_fixing_agent = PredictFixingAgent(
+        ai_provider=ai_provider,
+        api_key=api_key,
+        chat_history_path=repo_path / "predict_fixing.chat"
+    )
 
+    cog_predict_agent = CogPredictAgent(
+        ai_provider=ai_provider,
+        api_key=api_key,
+        chat_history_path=repo_path / "cog_predict.chat"
+    )
+
+    # 3. Use DocsPullingAgent (no AI required) to pull documentation
+    docs_agent = DocsPullingAgent(prompts_dir=repo_path / "prompts")
+    docs_agent.pull_docs()
+
+    if initialize: 
+        initialize_project(repo_path)
+
+    # 4. Check for the existence of cog.yaml, predict.py, and chat history
     cog_yaml_exists = (repo_path / "cog.yaml").exists()
     predict_py_exists = (repo_path / "predict.py").exists()
-    chat_history_exists = ai.chat_history_path.exists()
 
-    if chat_history_exists and (not cog_yaml_exists or not predict_py_exists):
-        raise ValueError(
-            f"AutoCog is in a semi-initialized state in {repo_path}, because one of cog.yaml or predict.py have been deleted. Run `autocog --initialize` to re-initialize the project"
-        )
+    # Check if project is in a semi-initialized state
+    if not cog_yaml_exists or not predict_py_exists:
+        if any(ai.chat_history_path.exists() for ai in [
+            path_ordering_agent.ai, package_agent.ai, cog_generation_agent.ai, 
+            predict_generation_agent.ai, error_diagnosis_agent.ai, 
+            cog_fixing_agent.ai, predict_fixing_agent.ai, cog_predict_agent.ai
+        ]):
+            raise ValueError(
+                f"AutoCog is in a semi-initialized state in {repo_path}, because one of cog.yaml or predict.py have been deleted. Run `autocog --initialize` to re-initialize the project"
+            )
 
-    if chat_history_exists:
-        ai.load_chat_history()
-    else:
-        paths = order_paths(ai, repo_path)
-        cog_yaml, predict_py = generate_initial(ai, repo_path, paths=paths, tell=tell)
-        (repo_path / "cog.yaml").write_text(cog_yaml)
+    # 5. Get package information from PackageInfoAgent
+    paths = path_ordering_agent.order_paths(repo_path)
+    packages = package_agent.get_imported_packages(paths)
+    package_versions = package_agent.get_packages_info(packages, repo_path)
+    
+    # 6. Generate initial `cog.yaml` and `predict.py` if necessary
+    if not predict_py_exists:
+        predict_py = predict_generation_agent.generate_predict_py(repo_path, predict_py=None, tell=tell)
         (repo_path / "predict.py").write_text(predict_py)
+    
+    if not cog_yaml_exists:
+        cog_yaml = cog_generation_agent.generate_cog_yaml(repo_path, predict_py=predict_py, cog_yaml=None, package_versions=package_versions, tell=tell)
+        (repo_path / "cog.yaml").write_text(cog_yaml)
 
+    # 7. Generate or use the provided predict command
     if not predict_command:
-        predict_command = ai.call(prompts.cog_predict)
+        predict_command = cog_predict_agent.generate_predict_command(predict_py)
 
-    predict_command = create_files_for_predict_command(repo_path, predict_command)
+    # 8. Attempt to run the prediction command and fix errors if necessary
     for attempt in range(attempts):
-        print("Predict command")
-        print(predict_command)
-        success, stderr = run_cog_predict(repo_path, predict_command)
+        print("Predict command:", predict_command)
+        success, stderr = cog_predict_agent.run_cog_predict(repo_path, predict_command)
+
         if success:
             return
 
@@ -376,32 +214,27 @@ def autocog(
             print(f"Failed after {attempts} attempts, giving up :'(")
             sys.exit(1)
 
-        print(
-            f"Attempt {attempt + 1}/{attempts} failed, trying to fix...",
-            file=sys.stderr,
-        )
+        print(f"Attempt {attempt + 1}/{attempts} failed, trying to fix...")
+
+        error = stderr.split("Traceback (most recent call last)")[-1] if "Traceback" in stderr else stderr
+        error_source, package_error = error_diagnosis_agent.diagnose_error(predict_command, error)
+
+        print("Error source:", error_source)
+        print("Package error:", package_error)
 
-        error = parse_cog_predict_error(stderr)
-        error_source, package_error = diagnose_error(ai, predict_command, error)
-        print("Error source")
-        print(error_source)
-        print("Package error")
-        print(package_error)
         if package_error:
-            get_packages_info(ai, repo_path)
+            package_agent.get_packages_info(repo_path)
 
-        if error_source == ERROR_PREDICT_PY:
-            predict_py = fix_predict_py(ai)
+        if error_source == "predict.py":
+            predict_py = predict_fixing_agent.fix_predict_py()
             (repo_path / "predict.py").write_text(predict_py)
-        elif error_source == ERROR_COG_YAML:
-            cog_yaml = fix_cog_yaml(ai)
+        elif error_source == "cog.yaml":
+            cog_yaml = cog_fixing_agent.fix_cog_yaml()
             (repo_path / "cog.yaml").write_text(cog_yaml)
-        elif error_source == ERROR_COG_PREDICT:
-            predict_command = ai.call(prompts.cog_predict)
-            predict_command = create_files_for_predict_command(
-                repo_path, predict_command
-            )
+        elif error_source == "cog_predict":
+            predict_command = cog_predict_agent.generate_predict_command()
 
 
 if __name__ == "__main__":
     autocog()
+
diff --git a/autocog/prompts.py b/autocog/prompts.py
index b33e911..29c9334 100644
--- a/autocog/prompts.py
+++ b/autocog/prompts.py
@@ -53,10 +53,10 @@ 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,
@@ -64,7 +64,7 @@ def generate_initial(
     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):
diff --git a/autocog/prompts/fix_cog_yaml.tpl b/autocog/prompts/cog_fixing/system.tpl
similarity index 100%
rename from autocog/prompts/fix_cog_yaml.tpl
rename to autocog/prompts/cog_fixing/system.tpl
diff --git a/autocog/prompts/generate_initial.tpl b/autocog/prompts/cog_generation/generate.tpl
similarity index 51%
rename from autocog/prompts/generate_initial.tpl
rename to autocog/prompts/cog_generation/generate.tpl
index 61d2432..09fbfe5 100644
--- a/autocog/prompts/generate_initial.tpl
+++ b/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 }}
diff --git a/autocog/prompts/cog_generation/system.tpl b/autocog/prompts/cog_generation/system.tpl
new file mode 100644
index 0000000..9264475
--- /dev/null
+++ b/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.
diff --git a/autocog/prompts/cog_predict.tpl b/autocog/prompts/cog_predict.tpl
deleted file mode 100644
index b6347ea..0000000
--- a/autocog/prompts/cog_predict.tpl
+++ /dev/null
@@ -1,5 +0,0 @@
-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 latest version of the predict.py file that you generated above. 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.
diff --git a/autocog/prompts/cog_predict/generate.tpl b/autocog/prompts/cog_predict/generate.tpl
new file mode 100644
index 0000000..c1b954e
--- /dev/null
+++ b/autocog/prompts/cog_predict/generate.tpl
@@ -0,0 +1,2 @@
+Here is the predict.py you generated:
+{{ predict_py }}
diff --git a/autocog/prompts/cog_predict/system.tpl b/autocog/prompts/cog_predict/system.tpl
new file mode 100644
index 0000000..fc6cefa
--- /dev/null
+++ b/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.
diff --git a/autocog/prompts/diagnose_error.tpl b/autocog/prompts/error_diagnosis/diagnose_error.tpl
similarity index 100%
rename from autocog/prompts/diagnose_error.tpl
rename to autocog/prompts/error_diagnosis/diagnose_error.tpl
diff --git a/autocog/prompts/package_error.tpl b/autocog/prompts/error_diagnosis/package_error.tpl
similarity index 100%
rename from autocog/prompts/package_error.tpl
rename to autocog/prompts/error_diagnosis/package_error.tpl
diff --git a/autocog/prompts/error_diagnosis/system.tpl b/autocog/prompts/error_diagnosis/system.tpl
new file mode 100644
index 0000000..4eb4ab4
--- /dev/null
+++ b/autocog/prompts/error_diagnosis/system.tpl
@@ -0,0 +1 @@
+You are an expert debugger.
diff --git a/autocog/prompts/order_paths/order_paths.tpl b/autocog/prompts/order_paths/order_paths.tpl
new file mode 100644
index 0000000..9476ac0
--- /dev/null
+++ b/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 %}
diff --git a/autocog/prompts/order_paths.tpl b/autocog/prompts/order_paths/system.tpl
similarity index 80%
rename from autocog/prompts/order_paths.tpl
rename to autocog/prompts/order_paths/system.tpl
index 3b386ca..0ce3fab 100644
--- a/autocog/prompts/order_paths.tpl
+++ b/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 %}
diff --git a/autocog/prompts/package_info/get_packages.tpl b/autocog/prompts/package_info/get_packages.tpl
new file mode 100644
index 0000000..fb181a8
--- /dev/null
+++ b/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 %}
diff --git a/autocog/prompts/get_packages.tpl b/autocog/prompts/package_info/system.tpl
similarity index 82%
rename from autocog/prompts/get_packages.tpl
rename to autocog/prompts/package_info/system.tpl
index 7402f24..4f54e2a 100644
--- a/autocog/prompts/get_packages.tpl
+++ b/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.
diff --git a/autocog/prompts/fix_predict_py.tpl b/autocog/prompts/predict_fixing/system.tpl
similarity index 100%
rename from autocog/prompts/fix_predict_py.tpl
rename to autocog/prompts/predict_fixing/system.tpl
diff --git a/autocog/prompts/predict_generation/generate.tpl b/autocog/prompts/predict_generation/generate.tpl
new file mode 100644
index 0000000..d6544c8
--- /dev/null
+++ b/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 %}
diff --git a/autocog/prompts/system.tpl b/autocog/prompts/predict_generation/system.tpl
similarity index 87%
rename from autocog/prompts/system.tpl
rename to autocog/prompts/predict_generation/system.tpl
index fad254a..d3f7a1a 100644
--- a/autocog/prompts/system.tpl
+++ b/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.
diff --git a/setup.py b/setup.py
index 775c2fa..4572529 100644
--- a/setup.py
+++ b/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=[