From 8879ff7f9393dde489457da19dbae398d765e1d5 Mon Sep 17 00:00:00 2001 From: Kishan Patel Date: Mon, 21 Apr 2025 01:11:56 +0530 Subject: [PATCH] feat: Security System for LLM --- src/definition/helper/userPreference.ts | 2 + src/enum/helper/Language.ts | 12 + src/enum/modals/UserPreferenceModal.ts | 2 + src/handlers/AIHandler.ts | 284 +++++++------- src/handlers/ExecuteViewSubmitHandler.ts | 7 + src/helper/AISecurity.ts | 467 +++++++++++++++++++++++ src/lib/Translation/locales/en.ts | 10 +- src/modal/UserPreferenceModal.ts | 44 +++ src/storage/userPreferenceStorage.ts | 7 +- 9 files changed, 699 insertions(+), 136 deletions(-) create mode 100644 src/enum/helper/Language.ts create mode 100644 src/helper/AISecurity.ts diff --git a/src/definition/helper/userPreference.ts b/src/definition/helper/userPreference.ts index 1f88b2b..15ae96d 100644 --- a/src/definition/helper/userPreference.ts +++ b/src/definition/helper/userPreference.ts @@ -1,4 +1,5 @@ import { Language } from '../../lib/Translation/translation'; +import { SecurityLevel } from '../../helper/AISecurity'; export type AIusagePreference = | AIusagePreferenceEnum.Personal @@ -27,6 +28,7 @@ export interface IPreference { AIconfiguration: { AIPrompt: string; AIProvider: AIProviderType; + securityLevel: SecurityLevel; selfHosted: { url: string; }; diff --git a/src/enum/helper/Language.ts b/src/enum/helper/Language.ts new file mode 100644 index 0000000..e8f4bac --- /dev/null +++ b/src/enum/helper/Language.ts @@ -0,0 +1,12 @@ +export enum Language { + EN = 'en', + ES = 'es', + FR = 'fr', + DE = 'de', + IT = 'it', + PT = 'pt', + RU = 'ru', + JA = 'ja', + KO = 'ko', + ZH = 'zh' +} \ No newline at end of file diff --git a/src/enum/modals/UserPreferenceModal.ts b/src/enum/modals/UserPreferenceModal.ts index 957b5db..29d27b6 100644 --- a/src/enum/modals/UserPreferenceModal.ts +++ b/src/enum/modals/UserPreferenceModal.ts @@ -10,6 +10,8 @@ export enum UserPreferenceModalEnum { AI_PREFERENCE_DROPDOWN_BLOCK_ID = 'AI-preference-drop-down-block-id', AI_OPTION_DROPDOWN_BLOCK_ID = 'AI-options-drop-down-block-id', AI_OPTION_DROPDOWN_ACTION_ID = 'AI-options-drop-down-action-id', + SECURITY_LEVEL_DROPDOWN_BLOCK_ID = 'security-level-drop-down-block-id', + SECURITY_LEVEL_DROPDOWN_ACTION_ID = 'security-level-drop-down-action-id', CLOSE_ACTION_ID = 'set-user-preference-language-modal-action-id', CLOSE_BLOCK_ID = 'set-user-preference-language-modal-block-id', OPEN_AI_API_KEY_ACTION_ID = 'open-ai-api-key-action-id', diff --git a/src/handlers/AIHandler.ts b/src/handlers/AIHandler.ts index e4ddf52..cfb54b0 100644 --- a/src/handlers/AIHandler.ts +++ b/src/handlers/AIHandler.ts @@ -11,6 +11,7 @@ import { IPreference, } from '../definition/helper/userPreference'; import { t } from '../lib/Translation/translation'; +import { AISecurity, SecurityLevel } from '../helper/AISecurity'; class AIHandler { constructor( @@ -19,138 +20,163 @@ class AIHandler { private userPreference: IPreference, ) {} private language = this.userPreference.language; + private aiSecurity = new AISecurity( + this.userPreference.AIconfiguration.securityLevel, + this.language + ); public async handleResponse( user: IUser, message: string, prompt: string, ): Promise { - let aiProvider: string; - if ( - this.userPreference.AIusagePreference === - AIusagePreferenceEnum.Personal - ) { - aiProvider = this.userPreference.AIconfiguration.AIProvider; - } else { - aiProvider = await this.app - .getAccessors() - .environmentReader.getSettings() - .getValueById(SettingEnum.AI_PROVIDER_OPTOIN_ID); - } + try { + let aiProvider: string; + if ( + this.userPreference.AIusagePreference === + AIusagePreferenceEnum.Personal + ) { + aiProvider = this.userPreference.AIconfiguration.AIProvider; + } else { + aiProvider = await this.app + .getAccessors() + .environmentReader.getSettings() + .getValueById(SettingEnum.AI_PROVIDER_OPTOIN_ID); + } - switch (aiProvider) { - case AIProviderEnum.SelfHosted: - case SettingEnum.SELF_HOSTED_MODEL: - return this.handleSelfHostedModel(user, message, prompt); + this.app.getLogger().log(`[AIHandler] Using AI Provider: ${aiProvider}`); - case AIProviderEnum.OpenAI: - case SettingEnum.OPEN_AI: - return this.handleOpenAI(user, message, prompt); + const combinedInput = this.getPrompt(message, prompt); + this.app.getLogger().log(`[AIHandler] Combined Input: ${combinedInput}`); - case AIProviderEnum.Gemini: - case SettingEnum.GEMINI: - return this.handleGemini(user, message, prompt); + const validation = this.aiSecurity.validateInput(combinedInput); + if (!validation.isValid) { + const errorMessage = validation.error || t('AI_Security_Invalid_Input', this.language); + this.app.getLogger().log(`[AIHandler] Security validation failed: ${errorMessage}`); + return errorMessage; + } - default: - const errorMsg = - this.userPreference.AIusagePreference === - AIusagePreferenceEnum.Personal - ? t('AI_Not_Configured_Personal', this.language) - : t('AI_Not_Configured_Admin', this.language); + let response: string; + switch (aiProvider) { + case AIProviderEnum.SelfHosted: + case SettingEnum.SELF_HOSTED_MODEL: + this.app.getLogger().log('[AIHandler] Using SelfHosted Model'); + response = await this.handleSelfHostedModel(user, combinedInput); + break; - this.app.getLogger().log(errorMsg); - return errorMsg; + case AIProviderEnum.OpenAI: + case SettingEnum.OPEN_AI: + this.app.getLogger().log('[AIHandler] Using OpenAI'); + response = await this.handleOpenAI(user, combinedInput); + break; + + case AIProviderEnum.Gemini: + case SettingEnum.GEMINI: + this.app.getLogger().log('[AIHandler] Using Gemini'); + response = await this.handleGemini(user, combinedInput); + break; + + default: + const errorMsg = + this.userPreference.AIusagePreference === + AIusagePreferenceEnum.Personal + ? t('AI_Not_Configured_Personal', this.language) + : t('AI_Not_Configured_Admin', this.language); + + this.app.getLogger().log(`[AIHandler] ${errorMsg}`); + return errorMsg; + } + + const responseValidation = this.aiSecurity.validateResponse(response); + if (!responseValidation.isValid) { + const errorMessage = responseValidation.error || t('AI_Security_Invalid_Input', this.language); + this.app.getLogger().log(`[AIHandler] Response validation failed: ${errorMessage}`); + return errorMessage; + } + + this.app.getLogger().log(`[AIHandler] Response received and validated`); + return response; + } catch (error) { + this.app.getLogger().log(`[AIHandler] Error in handleResponse: ${error.message}`); + this.app.getLogger().log(`[AIHandler] Error stack: ${error.stack}`); + return t('AI_Something_Went_Wrong', this.language); } } private getPrompt(message: string, prompt: string): string { - return `Write a reply to this message: "${message}". ${this.userPreference.AIconfiguration.AIPrompt} and Use the following as a prompt or response reply: "${prompt}" and make sure you respond with just the reply without quotes.`; + const basePrompt = `You are a helpful assistant in a chat application. + Please provide a professional and appropriate response to the following message. + + Original message: "${message}" + + Additional context: "${prompt}" + + Guidelines for your response: + 1. Keep the response concise and relevant + 2. Maintain a professional tone + 3. Avoid any inappropriate or offensive content + 4. If the request is inappropriate, respond with a standardized rejection message + 5. Focus on being helpful while maintaining appropriate boundaries + + Please provide your response without any additional formatting or quotes.`; + + return basePrompt; } private async handleSelfHostedModel( user: IUser, message: string, - prompt: string, ): Promise { try { - const url = await this.getSelfHostedModelUrl(); + const selfHostedModelUrl = await this.app + .getAccessors() + .environmentReader.getSettings() + .getValueById(SettingEnum.SELF_HOSTED_MODEL); - if (!url) { - this.app.getLogger().log('Self Hosted Model address not set.'); - if ( - this.userPreference.AIusagePreference === - AIusagePreferenceEnum.Personal - ) { - return t( - 'AI_Self_Hosted_Model_Not_Configured', - this.language, - ); - } else { - return t( - 'AI_Workspace_Model_Not_Configured', - this.language, - ); - } + if (!selfHostedModelUrl) { + this.app.getLogger().log('[AIHandler] Self-hosted model URL not configured'); + return t('AI_Not_Configured_Admin', this.language); } - const requestBody = { - messages: [ - { - role: 'system', - content: this.getPrompt(message, prompt), - }, - ], - temperature: 0, - }; + const combinedInput = this.getPrompt(message, ''); + this.app.getLogger().log(`[AIHandler] Sending request to self-hosted model: ${combinedInput}`); - const response: IHttpResponse = await this.http.post( - `${url}/chat/completions`, + const result = await this.aiSecurity.secureAIRequest( + combinedInput, + this.http, + `${selfHostedModelUrl}/chat/completions`, { - headers: { - 'Content-Type': 'application/json', - }, - content: JSON.stringify(requestBody), - }, + model: 'gpt-3.5-turbo' // Default model for self-hosted + } ); - if (!response || !response.data) { - this.app.getLogger().log('No response data received from AI.'); + if (!result.success) { + this.app.getLogger().log(`[AIHandler] Self-hosted model error: ${result.error}`); + return result.error || t('AI_Something_Went_Wrong', this.language); + } + + if (!result.response) { + this.app.getLogger().log('[AIHandler] Self-hosted model returned empty response'); return t('AI_Something_Went_Wrong', this.language); } - return response.data.choices[0].message.content; + this.app.getLogger().log(`[AIHandler] Self-hosted model response: ${result.response}`); + return result.response; } catch (error) { - this.app - .getLogger() - .log(`Error in handleSelfHostedModel: ${error.message}`); + this.app.getLogger().log(`[AIHandler] Self-hosted model exception: ${error.message}`); return t('AI_Something_Went_Wrong', this.language); } } - private async getSelfHostedModelUrl(): Promise { - if ( - this.userPreference.AIusagePreference === - AIusagePreferenceEnum.Personal - ) { - return this.userPreference.AIconfiguration.selfHosted.url; - } else { - return await this.app - .getAccessors() - .environmentReader.getSettings() - .getValueById(SettingEnum.SELF_HOSTED_MODEL_ADDRESS_ID); - } - } - private async handleOpenAI( user: IUser, - message: string, - prompt: string, + input: string, ): Promise { try { const { openaikey, openaimodel } = await this.getOpenAIConfig(); if (!openaikey || !openaimodel) { - this.app.getLogger().log('OpenAI settings not set properly.'); + this.app.getLogger().log('[AIHandler] OpenAI settings not configured'); const errorMsg = this.userPreference.AIusagePreference === AIusagePreferenceEnum.Personal @@ -160,34 +186,31 @@ class AIHandler { return errorMsg; } - const response: IHttpResponse = await this.http.post( + this.app.getLogger().log(`[AIHandler] Sending request to OpenAI: ${input}`); + + const result = await this.aiSecurity.secureAIRequest( + input, + this.http, 'https://api.openai.com/v1/chat/completions', { - headers: { - 'Content-Type': 'application/json', - Authorization: `Bearer ${openaikey}`, - }, - content: JSON.stringify({ - model: openaimodel, - messages: [ - { - role: 'system', - content: this.getPrompt(message, prompt), - }, - ], - }), - }, + model: openaimodel + } ); - if (!response || !response.data) { - this.app.getLogger().log('No response data received from AI.'); + if (!result.success) { + this.app.getLogger().log(`[AIHandler] OpenAI error: ${result.error}`); + return result.error || t('AI_Something_Went_Wrong', this.language); + } + + if (!result.response) { + this.app.getLogger().log('[AIHandler] OpenAI returned empty response'); return t('AI_Something_Went_Wrong', this.language); } - const { choices } = response.data; - return choices[0].message.content; + this.app.getLogger().log(`[AIHandler] OpenAI response: ${result.response}`); + return result.response; } catch (error) { - this.app.getLogger().log(`Error in handleOpenAI: ${error.message}`); + this.app.getLogger().log(`[AIHandler] OpenAI exception: ${error.message}`); return t('AI_Something_Went_Wrong', this.language); } } @@ -218,17 +241,16 @@ class AIHandler { return { openaikey: apikey, openaimodel: model }; } } + private async handleGemini( user: IUser, - message: string, - prompt: string, + input: string, ): Promise { try { const geminiAPIkey = await this.getGeminiAPIKey(); if (!geminiAPIkey) { - this.app.getLogger().log('Gemini API key not set Properly'); - + this.app.getLogger().log('[AIHandler] Gemini API key not configured'); const errorMsg = this.userPreference.AIusagePreference === AIusagePreferenceEnum.Personal @@ -238,35 +260,29 @@ class AIHandler { return errorMsg; } - const response: IHttpResponse = await this.http.post( + this.app.getLogger().log(`[AIHandler] Sending request to Gemini: ${input}`); + + const result = await this.aiSecurity.secureAIRequest( + input, + this.http, `https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key=${geminiAPIkey}`, - { - headers: { - 'Content-Type': 'application/json', - }, - content: JSON.stringify({ - contents: [ - { - parts: { - text: this.getPrompt(message, prompt), - }, - }, - ], - }), - }, + {} // Empty config as we'll format it in secureAIRequest ); - if (!response || !response.content) { - this.app - .getLogger() - .log('No response content received from AI.'); + if (!result.success) { + this.app.getLogger().log(`[AIHandler] Gemini error: ${result.error}`); + return result.error || t('AI_Something_Went_Wrong', this.language); + } + + if (!result.response) { + this.app.getLogger().log('[AIHandler] Gemini returned empty response'); return t('AI_Something_Went_Wrong', this.language); } - const data = response.data; - return data.candidates[0].content.parts[0].text; + this.app.getLogger().log(`[AIHandler] Gemini response: ${result.response}`); + return result.response; } catch (error) { - this.app.getLogger().log(`Error in handleGemini: ${error.message}`); + this.app.getLogger().log(`[AIHandler] Gemini exception: ${error.message}`); return t('AI_Something_Went_Wrong', this.language); } } diff --git a/src/handlers/ExecuteViewSubmitHandler.ts b/src/handlers/ExecuteViewSubmitHandler.ts index 53334e5..cfc4d90 100644 --- a/src/handlers/ExecuteViewSubmitHandler.ts +++ b/src/handlers/ExecuteViewSubmitHandler.ts @@ -34,6 +34,7 @@ import { AIusagePreference } from '../definition/helper/userPreference'; import { listReplyContextualBar } from '../modal/listContextualBar'; import { Receiverstorage } from '../storage/ReceiverStorage'; import { Replacements } from '../definition/helper/message'; +import { SecurityLevel } from '../helper/AISecurity'; export class ExecuteViewSubmitHandler { private context: UIKitViewSubmitInteractionContext; @@ -193,6 +194,11 @@ export class ExecuteViewSubmitHandler { UserPreferenceModalEnum.AI_OPTION_DROPDOWN_ACTION_ID ]; + const securityLevelInput = + view.state?.[UserPreferenceModalEnum.SECURITY_LEVEL_DROPDOWN_BLOCK_ID]?.[ + UserPreferenceModalEnum.SECURITY_LEVEL_DROPDOWN_ACTION_ID + ] as SecurityLevel; + const OpenAIAPIKeyInput = view.state?.[UserPreferenceModalEnum.OPEN_AI_API_KEY_BLOCK_ID]?.[ UserPreferenceModalEnum.OPEN_AI_API_KEY_ACTION_ID @@ -228,6 +234,7 @@ export class ExecuteViewSubmitHandler { AIconfiguration: { AIPrompt: PromptConfigurationInput, AIProvider: AIoptionInput, + securityLevel: securityLevelInput, openAI: { apiKey: OpenAIAPIKeyInput, model: OpenAImodelInput, diff --git a/src/helper/AISecurity.ts b/src/helper/AISecurity.ts new file mode 100644 index 0000000..12a2492 --- /dev/null +++ b/src/helper/AISecurity.ts @@ -0,0 +1,467 @@ +import { IHttp } from '@rocket.chat/apps-engine/definition/accessors'; +import { Language, t } from '../lib/Translation/translation'; + +export enum SecurityLevel { + STRICT = 'strict', + MODERATE = 'moderate', + RELAXED = 'relaxed' +} + +interface ValidationResult { + isValid: boolean; + error?: string; +} + +interface ValidationContext { + securityLevel: SecurityLevel; + language?: string; + userId?: string; +} + +export class AISecurity { + private securityLevel: SecurityLevel; + private language: Language; + + constructor(securityLevel: SecurityLevel = SecurityLevel.STRICT, language: Language = 'en' as Language) { + this.securityLevel = securityLevel; + this.language = language; + } + + private getSystemPrompt(): string { + const basePrompt = `You are a helpful assistant integrated into a chat application. + Your responses must be professional, appropriate, and helpful. + + Rules for responses: + 1. Never provide instructions for illegal, harmful, or unethical activities + 2. Never generate explicit, offensive, or inappropriate content + 3. For sensitive topics, provide factual information and resources + 4. Keep responses concise and relevant to the user's query + 5. If a request is inappropriate, respond with a standardized rejection message + + Standardized rejection messages: + - For illegal/harmful requests: "I'm sorry, I cannot assist with that request as it goes against our guidelines." + - For inappropriate content: "Let's keep our conversation professional and appropriate." + - For sensitive topics: "I recommend consulting with a qualified professional about this matter." + + Current security level: ${this.securityLevel}`; + + switch (this.securityLevel) { + case SecurityLevel.STRICT: + return `${basePrompt} + Strict guidelines: + - Reject any request that could be considered harmful or inappropriate + - Provide minimal information on sensitive topics + - Focus on professional and constructive responses`; + case SecurityLevel.MODERATE: + return `${basePrompt} + Moderate guidelines: + - Allow discussion of sensitive topics with appropriate context + - Provide balanced information while maintaining professionalism + - Use discretion when handling potentially controversial topics`; + case SecurityLevel.RELAXED: + return `${basePrompt} + Relaxed guidelines: + - Allow more open discussion while avoiding extreme content + - Provide general guidance on most topics`; + default: + return basePrompt; + } + } + + private getRejectionMessage(type: 'illegal' | 'inappropriate' | 'sensitive'): string { + const messages = { + illegal: t('AI_Security_Rejection_Message', this.language), + inappropriate: t('AI_Security_Respectful_Message', this.language), + sensitive: t('AI_Security_Sensitive_Topic_Message', this.language) + }; + return messages[type]; + } + + private validateContent(input: string): ValidationResult { + // Enhanced pattern matching with categories + const contentPatterns = { + illegal: [ + /how\s+to\s+hack/i, + /how\s+to\s+exploit/i, + /unauthorized\s+access/i, + /create\s+malware/i, + /create\s+virus/i, + /phishing\s+attack/i, + /identity\s+theft/i, + /cyber\s+attack/i, + /ddos\s+attack/i, + /sql\s+injection/i, + /xss\s+attack/i, + /csrf\s+attack/i, + ], + inappropriate: [ + /you're\s+stupid/i, + /you're\s+useless/i, + /you're\s+dumb/i, + /you're\s+idiot/i, + /hate\s+you/i, + /fuck\s+you/i, + /asshole/i, + /bitch/i, + /cunt/i, + /nigger/i, + /faggot/i, + ], + sensitive: [ + /credit\s+card\s+number/i, + /social\s+security\s+number/i, + /ssn/i, + /password/i, + /secret\s+key/i, + /api\s+key/i, + ] + }; + + // Check patterns in order of severity + for (const [category, patterns] of Object.entries(contentPatterns)) { + for (const pattern of patterns) { + if (pattern.test(input)) { + console.log(`[AISecurity] Content validation failed: ${category} pattern matched`); + return { + isValid: false, + error: this.getRejectionMessage(category as 'illegal' | 'inappropriate' | 'sensitive') + }; + } + } + } + + console.log('[AISecurity] Content validation passed'); + return { isValid: true }; + } + + static async validateInput(request: string, context: ValidationContext): Promise { + if (!request || request.trim().length === 0) { + return { + isValid: false, + error: 'Empty request' + }; + } + + const { securityLevel } = context; + if (this.containsSensitiveData(request)) { + return { + isValid: false, + error: 'Request contains sensitive data' + }; + } + + switch (securityLevel) { + case SecurityLevel.STRICT: + if (this.containsProfanity(request) || this.containsInappropriateContent(request)) { + return { + isValid: false, + error: 'Request contains inappropriate content' + }; + } + break; + case SecurityLevel.MODERATE: + if (this.containsProfanity(request)) { + return { + isValid: false, + error: 'Request contains profanity' + }; + } + break; + case SecurityLevel.RELAXED: + break; + } + + return { isValid: true }; + } + + private static containsSensitiveData(text: string): boolean { + const sensitivePatterns = [ + /\b\d{16}\b/, + /\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b/, + /\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/, + /\b\d{3}[-]?\d{2}[-]?\d{4}\b/ + ]; + + return sensitivePatterns.some(pattern => pattern.test(text)); + } + + private static containsProfanity(text: string): boolean { + const profanityList = [ + 'profanity1', + 'profanity2', + ]; + + return profanityList.some(word => + text.toLowerCase().includes(word.toLowerCase()) + ); + } + + private static containsInappropriateContent(text: string): boolean { + const inappropriatePatterns = [ + /\b(hack|exploit|crack)\b/i, + ]; + + return inappropriatePatterns.some(pattern => pattern.test(text)); + } + + public validateInput(input: string): ValidationResult { + if (!input || input.trim().length === 0) { + console.log('[AISecurity] Empty input validation failed'); + return { + isValid: false, + error: t('AI_Security_Invalid_Input', this.language) + }; + } + + // Use the content validation first + const contentValidation = this.validateContent(input); + if (!contentValidation.isValid) { + return contentValidation; + } + + // Then apply security level specific validation + switch (this.securityLevel) { + case SecurityLevel.STRICT: + if (AISecurity.containsProfanity(input) || AISecurity.containsInappropriateContent(input)) { + console.log('[AISecurity] Strict validation failed'); + return { + isValid: false, + error: this.getRejectionMessage('inappropriate') + }; + } + break; + case SecurityLevel.MODERATE: + if (AISecurity.containsProfanity(input)) { + console.log('[AISecurity] Moderate validation failed'); + return { + isValid: false, + error: this.getRejectionMessage('inappropriate') + }; + } + break; + case SecurityLevel.RELAXED: + // Only check for illegal content + if (AISecurity.containsInappropriateContent(input)) { + console.log('[AISecurity] Relaxed validation failed'); + return { + isValid: false, + error: this.getRejectionMessage('illegal') + }; + } + break; + } + + // Check for sensitive data last + if (AISecurity.containsSensitiveData(input)) { + console.log('[AISecurity] Sensitive data validation failed'); + return { + isValid: false, + error: this.getRejectionMessage('sensitive') + }; + } + + console.log('[AISecurity] All validations passed'); + return { isValid: true }; + } + + public validateResponse(response: string): ValidationResult { + if (!response || response.trim().length === 0) { + return { + isValid: false, + error: t('AI_Security_Invalid_Input', this.language) + }; + } + + const contentValidation = this.validateContent(response); + if (!contentValidation.isValid) { + return contentValidation; + } + + if (AISecurity.containsProfanity(response) || AISecurity.containsInappropriateContent(response)) { + return { + isValid: false, + error: this.getRejectionMessage('inappropriate') + }; + } + + if (AISecurity.containsSensitiveData(response)) { + return { + isValid: false, + error: this.getRejectionMessage('sensitive') + }; + } + + response = this.filterResponse(response); + return { isValid: true }; + } + + private filterResponse(response: string): string { + const filteredResponse = response.replace( + /(?:^|\s)(?:fuck|shit|asshole|bitch|cunt|nigger|faggot)(?:\s|$)/gi, + '****' + ); + return filteredResponse; + } + + public async secureAIRequest( + input: string, + http: IHttp, + apiUrl: string, + requestBody: any + ): Promise<{ success: boolean; response?: string; error?: string }> { + try { + // 1. Validate input + const validation = this.validateInput(input); + if (!validation.isValid) { + return { + success: false, + error: validation.error + }; + } + + // 2. Add enhanced security system prompt + const systemPrompt = this.getSystemPrompt(); + + // 3. Format the request body based on the API type + let formattedRequestBody; + if (apiUrl.includes('openai.com')) { + formattedRequestBody = { + model: requestBody.model, + messages: [ + { + role: 'system', + content: systemPrompt + }, + { + role: 'user', + content: input + } + ], + temperature: 0.7, + max_tokens: 500 + }; + } else if (apiUrl.includes('generativelanguage.googleapis.com')) { + formattedRequestBody = { + contents: [ + { + role: 'user', + parts: [ + { + text: systemPrompt + }, + { + text: input + } + ] + } + ], + generationConfig: { + temperature: 0.7, + maxOutputTokens: 500 + } + }; + } else { + // For self-hosted models + formattedRequestBody = { + messages: [ + { + role: 'system', + content: systemPrompt + }, + { + role: 'user', + content: input + } + ], + temperature: 0.7, + max_tokens: 500 + }; + } + + // 4. Make the API request with enhanced error handling + const response = await http.post(apiUrl, { + headers: { + 'Content-Type': 'application/json', + 'Accept': 'application/json', + 'Cache-Control': 'no-cache' + }, + content: JSON.stringify(formattedRequestBody), + timeout: 30000 + }); + + if (!response) { + return { + success: false, + error: t('AI_Something_Went_Wrong', this.language) + }; + } + + // 5. Enhanced response handling + if (response.statusCode !== 200) { + const errorMessage = response.data?.error?.message || 'Unknown error'; + console.log(`[AISecurity] API Error: ${errorMessage}`); + return { + success: false, + error: `API Error (${response.statusCode}): ${errorMessage}` + }; + } + + // 6. Parse and validate response + let aiResponse: string; + try { + if (response.data.choices?.[0]?.message?.content) { + aiResponse = response.data.choices[0].message.content; + } else if (response.data.candidates?.[0]?.content?.parts?.[0]?.text) { + aiResponse = response.data.candidates[0].content.parts[0].text; + } else { + return { + success: false, + error: t('AI_Something_Went_Wrong', this.language) + }; + } + + // Validate the AI response + const responseValidation = this.validateResponse(aiResponse); + if (!responseValidation.isValid) { + return { + success: false, + error: responseValidation.error + }; + } + } catch (parseError) { + return { + success: false, + error: `Failed to parse AI response: ${parseError.message}` + }; + } + + // 7. Filter and return response + const filteredResponse = this.filterResponse(aiResponse); + return { + success: true, + response: filteredResponse + }; + } catch (error) { + console.log(`[AISecurity] Request failed: ${error.message}`); + return { + success: false, + error: `Request failed: ${error.message}` + }; + } + } + + public async processRequest(request: string, context: ValidationContext): Promise<{ success: boolean; response?: string; error?: string }> { + const validation = await AISecurity.validateInput(request, context); + if (!validation.isValid) { + return { + success: false, + error: validation.error + }; + } + + return { + success: true, + response: 'Processed successfully' + }; + } +} \ No newline at end of file diff --git a/src/lib/Translation/locales/en.ts b/src/lib/Translation/locales/en.ts index 5e140db..6b436e2 100644 --- a/src/lib/Translation/locales/en.ts +++ b/src/lib/Translation/locales/en.ts @@ -86,4 +86,12 @@ export const en = { AI_Gemini_Model_Not_Configured: "Your Gemini Model is not set up properly. Please check your configuration", AI_Workspace_Model_Not_Configured: "Your Workspace AI is not set up properly. Please contact your administrator", AI_Something_Went_Wrong: "Something went wrong. Please try again later.", - Refresh_Button_Text: "Refresh"}; + Refresh_Button_Text: "Refresh", + AI_Security_Invalid_Input: "Sorry, I can't help with that request. Let's keep this conversation respectful and appropriate.", + AI_Security_Rejection_Message: "Sorry, I can't help with that request.", + AI_Security_Respectful_Message: "Let's keep this conversation respectful and appropriate.", + AI_Security_Sensitive_Topic_Message: "That topic goes against our community guidelines.", + AI_Security_Community_Guidelines: "I'm here to help with constructive and safe queries.", + Choose_Security_Level_Placeholder: "Choose security level", + Choose_Security_Level_Label: "Security Level" +}; diff --git a/src/modal/UserPreferenceModal.ts b/src/modal/UserPreferenceModal.ts index 89db9e6..0b5b855 100644 --- a/src/modal/UserPreferenceModal.ts +++ b/src/modal/UserPreferenceModal.ts @@ -23,6 +23,7 @@ import { IPreference, } from '../definition/helper/userPreference'; import { inputElementComponent } from './common/inputElementComponent'; +import { SecurityLevel } from '../helper/AISecurity'; export async function UserPreferenceModal({ app, @@ -276,6 +277,49 @@ export async function UserPreferenceModal({ break; } } + + const securityLevelOptions = [ + { + text: SecurityLevel.STRICT, + value: SecurityLevel.STRICT, + }, + { + text: SecurityLevel.MODERATE, + value: SecurityLevel.MODERATE, + }, + { + text: SecurityLevel.RELAXED, + value: SecurityLevel.RELAXED, + }, + ]; + + const securityLevelDropDownOption = + elementBuilder.createDropDownOptions(securityLevelOptions); + + const securityLevelDropDown = elementBuilder.addDropDown( + { + placeholder: t('Choose_Security_Level_Placeholder', language), + options: securityLevelDropDownOption, + dispatchActionConfig: [Modals.dispatchActionConfigOnSelect], + initialOption: securityLevelDropDownOption.find( + (option) => + option.value === + existingPreference.AIconfiguration?.securityLevel, + ), + }, + { + blockId: UserPreferenceModalEnum.SECURITY_LEVEL_DROPDOWN_BLOCK_ID, + actionId: UserPreferenceModalEnum.SECURITY_LEVEL_DROPDOWN_ACTION_ID, + }, + ); + + blocks.push( + blockBuilder.createInputBlock({ + text: t('Choose_Security_Level_Label', language), + element: securityLevelDropDown, + optional: false, + }), + ); } const submitButton = elementBuilder.addButton( diff --git a/src/storage/userPreferenceStorage.ts b/src/storage/userPreferenceStorage.ts index 2491cc4..fe4e395 100644 --- a/src/storage/userPreferenceStorage.ts +++ b/src/storage/userPreferenceStorage.ts @@ -13,6 +13,7 @@ import { IPreference, } from '../definition/helper/userPreference'; import { Language } from '../lib/Translation/translation'; +import { SecurityLevel } from '../helper/AISecurity'; export class UserPreferenceStorage implements IuserPreferenceStorage { private userId: string; @@ -41,6 +42,9 @@ export class UserPreferenceStorage implements IuserPreferenceStorage { AIProvider: preference.AIconfiguration.AIProvider || currentPreference.AIconfiguration.AIProvider, + securityLevel: + preference.AIconfiguration.securityLevel || + currentPreference.AIconfiguration.securityLevel, openAI: { apiKey: preference.AIconfiguration.openAI.apiKey || @@ -89,8 +93,9 @@ export class UserPreferenceStorage implements IuserPreferenceStorage { language: Language.en, AIusagePreference: AIusagePreferenceEnum.Workspace, AIconfiguration: { - AIPrompt: `Keep the comprehensive clear and concise reply, and ensure it's well-articulated and helpfull`, + AIPrompt: `Keep the comprehensive clear and concise reply, and ensure it's well-articulated and helpful`, AIProvider: AIProviderEnum.SelfHosted, + securityLevel: SecurityLevel.STRICT, gemini: { apiKey: '', },