diff --git a/README.md b/README.md
index 5081369..2d19a50 100644
--- a/README.md
+++ b/README.md
@@ -38,6 +38,7 @@ View the [Change Logs](https://github.com/csgoh/roadmapper/wiki/Change-Logs) to
* drawsvg >= 2.3.0
* rich >= 13.7.1
* ruff >= 0.2.1
+* openai >= 1.0.0
@@ -45,6 +46,30 @@ Any feedback or suggestions are welcome. Please feel free to create an issue or
+## AI Roadmap Generation
+You can now generate roadmap code from natural language descriptions using the `generate_roadmap_code` function.
+
+```python
+from roadmapper import generate_roadmap_code
+
+description = """
+Create a roadmap for a new product launch.
+The timeline should be monthly starting from 2024-01-01.
+Add a group for 'Development'.
+Add a task 'MVP' from 2024-01-01 to 2024-04-01 with a milestone 'Alpha Release' on 2024-03-01.
+Add a task 'Beta' from 2024-04-01 to 2024-06-01 with a milestone 'Public Beta' on 2024-05-15.
+"""
+
+api_key = "your-openai-api-key"
+code = generate_roadmap_code(description, api_key=api_key)
+print(code)
+
+# You can then execute the code:
+exec(code)
+```
+
+
+
## Installation
### Install from PyPI
diff --git a/pyproject.toml b/pyproject.toml
index abbb901..4a8725b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ classifiers = [
"Programming Language :: Python :: 3.10",
"Operating System :: OS Independent",
]
-dependencies = ['Pillow>=10.0.0', 'python-dateutil>=2.8.2', 'drawsvg>=2.2.0']
+dependencies = ['Pillow>=10.0.0', 'python-dateutil>=2.8.2', 'drawsvg>=2.2.0', 'openai>=1.0.0']
[project.urls]
diff --git a/requirements.txt b/requirements.txt
index a41f1e5..3396c3f 100644
Binary files a/requirements.txt and b/requirements.txt differ
diff --git a/src/roadmapper/__init__.py b/src/roadmapper/__init__.py
index e69de29..bdf66cd 100644
--- a/src/roadmapper/__init__.py
+++ b/src/roadmapper/__init__.py
@@ -0,0 +1,5 @@
+from roadmapper.roadmap import Roadmap
+from roadmapper.timelinemode import TimelineMode
+from roadmapper.ai import generate_roadmap_code
+
+__all__ = ["Roadmap", "TimelineMode", "generate_roadmap_code"]
diff --git a/src/roadmapper/ai.py b/src/roadmapper/ai.py
new file mode 100644
index 0000000..2e54907
--- /dev/null
+++ b/src/roadmapper/ai.py
@@ -0,0 +1,116 @@
+# MIT License
+
+# Copyright (c) 2024 CS Goh
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import re
+from openai import OpenAI
+
+
+def generate_roadmap_code(
+ prompt: str, api_key: str, model: str = "gpt-4o"
+) -> str:
+ """
+ Generates Python code for a roadmap using the roadmapper library based on a natural language description.
+
+ Args:
+ prompt (str): A description of the roadmap you want to generate.
+ api_key (str): Your OpenAI API key.
+ model (str, optional): The OpenAI model to use. Defaults to "gpt-4o".
+
+ Returns:
+ str: Valid Python code that uses roadmapper to generate the roadmap.
+ """
+
+ system_prompt = """
+You are an expert Python developer specialized in using the `roadmapper` library.
+Your task is to convert the user's natural language description of a roadmap into valid Python code using the `roadmapper` library.
+
+### Instructions:
+1. Return ONLY valid Python code.
+2. Do not include markdown formatting (like ```python ... ```) if possible, or I will strip it.
+3. The code must be complete and runnable.
+4. Use the provided examples to understand the API.
+
+### Roadmapper API Examples:
+
+```python
+from roadmapper.roadmap import Roadmap
+from roadmapper.timelinemode import TimelineMode
+
+# Create roadmap
+roadmap = Roadmap(1200, 400, colour_theme="BLUEMOUNTAIN")
+roadmap.set_title("My Demo Roadmap")
+roadmap.set_subtitle("Matariki Technologies Ltd")
+roadmap.set_timeline(TimelineMode.MONTHLY, start="2025-01-01", number_of_items=18)
+roadmap.set_footer("Generated by Roadmapper")
+
+# Add Logo (optional)
+# roadmap.add_logo("logo.png", position="top-right", width=50, height=50)
+
+# Add Group
+group = roadmap.add_group("Core Product Work Stream")
+
+# Add Task
+task = group.add_task("Base Functionality", "2025-01-01", "2025-10-31")
+task.add_milestone("v.1.0", "2025-02-15")
+
+# Add Parallel Task
+parallel_task = task.add_parallel_task("Enhancements", "2025-11-15", "2026-03-31")
+parallel_task.add_milestone("v.2.0", "2026-03-30")
+
+# Draw and Save
+roadmap.draw()
+roadmap.save("roadmap.png")
+```
+
+### Timeline Modes:
+- TimelineMode.WEEKLY
+- TimelineMode.MONTHLY
+- TimelineMode.QUARTERLY
+- TimelineMode.HALF_YEARLY
+- TimelineMode.YEARLY
+
+### Colour Themes:
+- "DEFAULT", "BLUEMOUNTAIN", "ORANGE", "GREYWOOLF", "TEAL"
+
+### Output Format:
+Provide only the python code. No explanations.
+"""
+
+ client = OpenAI(api_key=api_key)
+
+ response = client.chat.completions.create(
+ model=model,
+ messages=[
+ {"role": "system", "content": system_prompt},
+ {"role": "user", "content": prompt},
+ ],
+ temperature=0.2,
+ )
+
+ content = response.choices[0].message.content
+
+ # Strip markdown code blocks if present
+ content = re.sub(r"^```python\n", "", content)
+ content = re.sub(r"^```\n", "", content)
+ content = re.sub(r"\n```$", "", content)
+
+ return content.strip()
diff --git a/tests/test_ai.py b/tests/test_ai.py
new file mode 100644
index 0000000..5383301
--- /dev/null
+++ b/tests/test_ai.py
@@ -0,0 +1,35 @@
+import unittest
+from unittest.mock import MagicMock, patch
+from roadmapper.ai import generate_roadmap_code
+
+
+class TestAI(unittest.TestCase):
+ @patch("roadmapper.ai.OpenAI")
+ def test_generate_roadmap_code(self, mock_openai):
+ # Setup mock
+ mock_client = MagicMock()
+ mock_openai.return_value = mock_client
+
+ mock_completion = MagicMock()
+ mock_completion.choices = [
+ MagicMock(message=MagicMock(content="```python\nprint('Hello World')\n```"))
+ ]
+ mock_client.chat.completions.create.return_value = mock_completion
+
+ # Execute
+ code = generate_roadmap_code("Create a roadmap", "fake-api-key")
+
+ # Verify
+ self.assertEqual(code, "print('Hello World')")
+ mock_client.chat.completions.create.assert_called_once()
+
+ # Verify prompt structure
+ args, kwargs = mock_client.chat.completions.create.call_args
+ self.assertEqual(kwargs["model"], "gpt-4o")
+ self.assertEqual(len(kwargs["messages"]), 2)
+ self.assertEqual(kwargs["messages"][0]["role"], "system")
+ self.assertEqual(kwargs["messages"][1]["role"], "user")
+ self.assertIn("Roadmapper API Examples", kwargs["messages"][0]["content"])
+
+if __name__ == "__main__":
+ unittest.main()