Merge pull request #588 from linear-b/add-role-option

added role to askAI plugin + examples
linear-b · Sep 9, 2024 · d485818 · d485818
2 parents a8555b2 + 1f8d7c7
commit d485818
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 27 deletions.
diff --git a/docs/downloads/automation-library/integrations/askAI/askAI_CR.cm b/docs/downloads/automation-library/integrations/askAI/askAI_CR.cm
@@ -15,8 +15,8 @@ automations:
         args:
           comment: |
             # ✨ gitStream Review ✨
-            {{ source | askAI("
-              Review the PR code diff.
+            {{ source | askAI("Code reviewer",
+              "Review the PR code diff.
               - Identify bugs, security risks, and performance issues
               - Check for deprecated methods and style guide violations
               - Provide specific improvement suggestions based on the changes", env.OPEN_AI_TOKEN) | encode }}
diff --git a/docs/downloads/automation-library/integrations/askAI/askAI_summarize.cm b/docs/downloads/automation-library/integrations/askAI/askAI_summarize.cm
@@ -15,5 +15,5 @@ automations:
         args:
           comment: |
             # 📜 PR Summary 📜
-            {{ source | askAI("
-              Summarize the changes in this PR in bullet points.", env.OPEN_AI_TOKEN) | encode }}
+            {{ source | askAI("Experienced developer",
+              "Summarize the changes in this PR in bullet points.", env.OPEN_AI_TOKEN) | encode }}
diff --git a/docs/downloads/automation-library/integrations/askAI/askAI_tests.cm b/docs/downloads/automation-library/integrations/askAI/askAI_tests.cm
@@ -15,7 +15,7 @@ automations:
         args:
           comment: |
             # 🧪 Test Suggestions 🧪
-            {{ source | askAI("
-              Identify any new or modified functions without test coverage in this PR.
+            {{ source | askAI("QA tester",
+              "Identify any new or modified functions without test coverage in this PR.
               Suggest specific test cases to add, including edge cases.
               If all functions are covered, return 'No additional tests needed.'", env.OPEN_AI_TOKEN) | encode }}
diff --git a/plugins/filters/askAI/askAI.cm b/plugins/filters/askAI/askAI.cm
@@ -13,7 +13,7 @@ automations:
       - action: add-comment@v1
         args:
           comment: |
-            {{ source | askAI("Based on the given context, search for new functions without tests and suggest the tests to add. If all functions are covered completely, return 'no tests to suggest.'", env.OPEN_AI_TOKEN) | encode }}
+            {{ source | askAI("Experienced developer", "Summarize in simple terms the changes in this PR using bullet points.", env.OPEN_AI_TOKEN) | encode }}
 
   generate_pr_desc_on_ask_ai_label:
     on:
@@ -24,4 +24,4 @@ automations:
       - action: add-comment@v1
         args:
           comment: |
-            {{ source | askAI("Based on the given context, search for new functions without tests and suggest the tests to add. If all functions are covered completely, return 'no tests to suggest.'", env.OPEN_AI_TOKEN) | encode }}
+            {{ source | askAI("qa tester", "Based on the given context, search for new functions without tests and suggest the tests to add. If all functions are covered completely, return 'no tests to suggest.'", env.OPEN_AI_TOKEN) | encode }}
diff --git a/plugins/filters/askAI/index.js b/plugins/filters/askAI/index.js
@@ -2,6 +2,7 @@
  * @module askAI
  * @description A gitStream plugin to interact with AI models. Currently works with `ChatGPR-4o-mini`.
  * @param {Object} context - The context that will be attached to the prompt .
+ * @param {string} role - Role instructions for the conversation.
  * @param {string} prompt - The prompt string.
  * @param {Object} token - The token to the AI model.
  * @returns {Object} Returns the response from the AI model.
@@ -82,8 +83,8 @@ const buildContextForGPT = context => {
   return context;
 };
 
-const askAI = async (context, prompt, token, callback) => {
-  const cacheKey = `${__filename}${prompt}`;
+const askAI = async (context, role = '', prompt, token, callback) => {
+  const cacheKey = `${__filename}${role}${prompt}`;
 
   if (process.env[cacheKey]) {
     return callback(null, process.env[cacheKey]);
@@ -103,19 +104,17 @@ const askAI = async (context, prompt, token, callback) => {
     body: JSON.stringify({
       model: 'gpt-4o-2024-08-06',
       messages: [
+        ...(role ? 
+        [
+          {
+            role: 'system',
+            content: `You are a ${role}. Answer only to the request, without any introductory or conclusion text.`
+          }] 
+          : []),
         {
-          role: 'system',
-          content: 'You are a code reviewer.'
-        },
-        {
-          role: 'system',
+          role: 'user',
           content: JSON.stringify(formattedContext)
         },
-        {
-          role: 'assistant',
-          content:
-            'You are code reviewer for a project. please answer without introductory phrases.'
-        },
         { role: 'user', content: prompt }
       ],
       max_tokens: maxTokens

diff --git a/plugins/filters/askAI/reference.md b/plugins/filters/askAI/reference.md
@@ -5,14 +5,15 @@ A gitStream plugin to interact with AI models. Currently works with `ChatGPR-4o-
 
 ![Example PR description](screenshots/askAI-describe-PR.png)
 
-**Returns**: <code>Object</code> - Returns the response from the AI model
+**Returns**: <code>Object</code> -  Returns the AI-generated response based on the provided context and prompt.
 **License**: MIT
 
-| Param   | Type     | Description                                     |
-| ------- | -------- | ----------------------------------------------- |
-| context | `Object` | The context that will be attached to the prompt |
-| prompt  | `string` | The prompt string                               |
-| token   | `Object` | The token to the AI model                       |
+| Param   | Type     | Description                                                                                                        |
+| ------- | -------- | ------------------------------------------------------------------------------------------------------------------ |
+| context | `Object` | The context that needs to be sent to the AI model for analysis.                                                    |
+| role    | `string` |  Free text. If not empty, Defines the role or persona for the AI to adopt while generating the response.           |
+| prompt  | `string` | The specific request or question you want the AI to respond to, after the context has been provided.               |
+| token   | `Object` | The token to the AI model                                                                                          |
 
 
 **Example**
@@ -21,5 +22,5 @@ A gitStream plugin to interact with AI models. Currently works with `ChatGPR-4o-
     The [`add-comment`](./automation-actions.md#add-comment) action automatically decodes encoded strings.
 
 ```yaml
-{{ source | askAI("Based on the given context, search for new functions without tests and suggest the tests to add. If all functions are covered completely, return 'no tests to suggest.'", env.OPEN_AI_TOKEN) | encode }}
+{{ source | askAI("QA tester", "Based on the given context, search for new functions without tests and suggest the tests to add. If all functions are covered completely, return 'no tests to suggest.'", env.OPEN_AI_TOKEN) | encode }}
 ```