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()