From 3b13e8345ae96438fd841044b558cc98cf1bcf31 Mon Sep 17 00:00:00 2001
From: Sunny Modi <sunnymodi21@gmail.com>
Date: Sat, 4 Jan 2025 11:52:01 -0500
Subject: [PATCH] Feature/Add Groq for followup prompts (#3806)

Feature/Add-Groq-for-followup-prompts
---
 packages/components/src/Interface.ts          |  3 +-
 packages/components/src/followUpPrompts.ts    | 11 ++++++
 .../ui-component/extended/FollowUpPrompts.jsx | 38 +++++++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/packages/components/src/Interface.ts b/packages/components/src/Interface.ts
index 1b7f4a2e..17ad1b5e 100644
--- a/packages/components/src/Interface.ts
+++ b/packages/components/src/Interface.ts
@@ -427,7 +427,8 @@ export enum FollowUpPromptProvider {
     AZURE_OPENAI = 'azureChatOpenAI',
     GOOGLE_GENAI = 'chatGoogleGenerativeAI',
     MISTRALAI = 'chatMistralAI',
-    OPENAI = 'chatOpenAI'
+    OPENAI = 'chatOpenAI',
+    GROQ = 'groqChat'
 }
 
 export type FollowUpPromptProviderConfig = {
diff --git a/packages/components/src/followUpPrompts.ts b/packages/components/src/followUpPrompts.ts
index 91be9b69..23d5df74 100644
--- a/packages/components/src/followUpPrompts.ts
+++ b/packages/components/src/followUpPrompts.ts
@@ -7,6 +7,7 @@ import { ChatOpenAI } from '@langchain/openai'
 import { z } from 'zod'
 import { PromptTemplate } from '@langchain/core/prompts'
 import { StructuredOutputParser } from '@langchain/core/output_parsers'
+import { ChatGroq } from '@langchain/groq'
 
 const FollowUpPromptType = z
     .object({
@@ -108,6 +109,16 @@ export const generateFollowUpPrompts = async (
                 const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
                 return structuredResponse
             }
+            case FollowUpPromptProvider.GROQ: {
+                const llm = new ChatGroq({
+                    apiKey: credentialData.groqApiKey,
+                    model: providerConfig.modelName,
+                    temperature: parseFloat(`${providerConfig.temperature}`)
+                })
+                const structuredLLM = llm.withStructuredOutput(FollowUpPromptType)
+                const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
+                return structuredResponse
+            }
         }
     } else {
         return undefined
diff --git a/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx b/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx
index 5d0c703a..6e2a822f 100644
--- a/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx
+++ b/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx
@@ -13,6 +13,7 @@ import anthropicIcon from '@/assets/images/anthropic.svg'
 import azureOpenAiIcon from '@/assets/images/azure_openai.svg'
 import mistralAiIcon from '@/assets/images/mistralai.svg'
 import openAiIcon from '@/assets/images/openai.svg'
+import groqIcon from '@/assets/images/groq.png'
 import { TooltipWithParser } from '@/ui-component/tooltip/TooltipWithParser'
 import CredentialInputHandler from '@/views/canvas/CredentialInputHandler'
 import { Input } from '@/ui-component/input/Input'
@@ -32,6 +33,7 @@ const FollowUpPromptProviders = {
     ANTHROPIC: 'chatAnthropic',
     AZURE_OPENAI: 'azureChatOpenAI',
     GOOGLE_GENAI: 'chatGoogleGenerativeAI',
+    GROQ: 'groqChat',
     MISTRALAI: 'chatMistralAI',
     OPENAI: 'chatOpenAI'
 }
@@ -149,6 +151,42 @@ const followUpPromptsOptions = {
             }
         ]
     },
+    [FollowUpPromptProviders.GROQ]: {
+        label: 'Groq',
+        name: FollowUpPromptProviders.GROQ,
+        icon: groqIcon,
+        inputs: [
+            {
+                label: 'Connect Credential',
+                name: 'credential',
+                type: 'credential',
+                credentialNames: ['groqApi']
+            },
+            {
+                label: 'Model Name',
+                name: 'modelName',
+                type: 'asyncOptions',
+                loadMethod: 'listModels'
+            },
+            {
+                label: 'Prompt',
+                name: 'prompt',
+                type: 'string',
+                rows: 4,
+                description: promptDescription,
+                optional: true,
+                default: defaultPrompt
+            },
+            {
+                label: 'Temperature',
+                name: 'temperature',
+                type: 'number',
+                step: 0.1,
+                optional: true,
+                default: 0.9
+            }
+        ]
+    },
     [FollowUpPromptProviders.MISTRALAI]: {
         label: 'Mistral AI',
         name: FollowUpPromptProviders.MISTRALAI,