Improve issue tracker example

Author: fjsj

example/issue_tracker/ai_assistants.py

@@ -1,4 +1,4 @@
-from typing import Sequence
+import json
 
 from django.contrib.auth.models import User
 
@@ -11,9 +11,9 @@ class IssueTrackerAIAssistant(AIAssistant):
     name = "Issue Tracker Assistant"
     instructions = (
         "You are a issue tracker assistant. "
-        "Help the user manage issues using the provided tools. "
-        "Issue IDs are unique and auto-incremented, they are represented as #<id>. "
-        "Make sure to include it in your responses, "
+        "Interact with the issue DB by using the provided tools. "
+        "Always refer to issue IDs as #<id>. "
+        "Make sure to include issue IDs in your responses, "
         "to know which issue you or the user are referring to. "
     )
     model = "gpt-4o"
@@ -24,26 +24,34 @@ def get_current_assignee_email(self) -> str:
         """Get the current user's email"""
         return self._user.email
 
-    def _format_issues(self, issues: Sequence[Issue]) -> str:
-        if not issues:
-            return "No issues found"
-        return "\n\n".join(
-            [f"- {issue.title} #{issue.id}\n{issue.description}" for issue in issues]
-        )
-
     @method_tool
     def list_issues(self) -> str:
         """List all issues"""
-        return self._format_issues(list(Issue.objects.all()))
+        return json.dumps(
+            {
+                "issues": list(
+                    Issue.objects.values("id", "title", "description", "assignee__email")
+                ),
+            }
+        )
 
     @method_tool
-    def list_user_assigned_issues(self, assignee_email: str) -> str:
-        """List the issues assigned to the provided user"""
-        return self._format_issues(list(Issue.objects.filter(assignee__email=assignee_email)))
+    def list_assigned_issues(self, assignee_email: str) -> str:
+        """List the issues assigned to the user with the email from `assignee_email` param."""
+        return json.dumps(
+            {
+                "issues": list(
+                    Issue.objects.filter(
+                        assignee__email=assignee_email,
+                    ).values("id", "title", "description", "assignee__email")
+                ),
+            }
+        )
 
     @method_tool
-    def assign_user_to_issue(self, issue_id: int, assignee_email: str = "") -> str:
-        """Assign a user to an issue. When assignee_email is empty, the issue assignment is removed."""
+    def assign_user_to_issue(self, issue_id: int, assignee_email: str) -> str:
+        """Assign a user to an issue. Set the `assignee_email` param to an user's email.
+        Or pass an empty string to remove the issue assignment."""
         try:
             issue = Issue.objects.get(id=issue_id)
             if assignee_email:
@@ -56,11 +64,12 @@ def assign_user_to_issue(self, issue_id: int, assignee_email: str = "") -> str:
             return f"ERROR: User {assignee_email} does not exist"
         issue.assignee = assignee
         issue.save()
-        return f"Assigned {assignee_email} to issue {issue.title} #{issue.id}"
+        return f"Assigned {assignee_email} to issue #{issue.id} - {issue.title}"
 
     @method_tool
     def create_issue(self, title: str, description: str = "", assignee_email: str = "") -> str:
-        """Create a new issue. Title is required. Description is optional. Assignee is optional."""
+        """Create a new issue.
+        Assign it to a user by passing the `assignee_email` param."""
         if assignee_email:
             try:
                 assignee = User.objects.get(email=assignee_email)
@@ -69,19 +78,19 @@ def create_issue(self, title: str, description: str = "", assignee_email: str =
         else:
             assignee = None
         issue = Issue.objects.create(title=title, description=description, assignee=assignee)
-        return f"Created issue {issue.title} #{issue.id}"
+        return f"Created issue #{issue.id} - {issue.title}"
 
     @method_tool
     def update_issue(self, issue_id: int, title: str, description: str = "") -> str:
-        """Update an issue"""
+        """Update an issue."""
         try:
             issue = Issue.objects.get(id=issue_id)
         except Issue.DoesNotExist:
             return f"ERROR: Issue {issue_id} does not exist"
         issue.title = title
         issue.description = description
         issue.save()
-        return f"Updated issue {issue.title} #{issue.id}"
+        return f"Updated issue #{issue_id} - {issue.title}"
 
     @method_tool
     def delete_issue(self, issue_id: int) -> str:
@@ -91,4 +100,4 @@ def delete_issue(self, issue_id: int) -> str:
         except Issue.DoesNotExist:
             return f"ERROR: Issue {issue_id} does not exist"
         issue.delete()
-        return f"Deleted issue {issue.title} #{issue.id}"
+        return f"Deleted issue #{issue_id} - {issue.title}"