Feat : Integrate hooks in LiteLLM to modify/reject requests and responses

apps/account-job-executor/src/main/java/com/akto/account_job_executor/executor/executors/AIAgentConnectorExecutor.java

@@ -69,6 +69,10 @@ protected void runJob(AccountJob job) throws Exception {
                 executeCopilotStudioConnector(job, config);
                 break;
 
+            case "LITELLM":
+                executeLiteLLMConnector(job, config);
+                break;
+
             default:
                 logger.warn("Unknown AI Agent Connector subType: {}. Skipping job execution.", subType);
                 throw new IllegalArgumentException("Unsupported AI Agent Connector subType: " + subType);
@@ -116,6 +120,20 @@ private void executeCopilotStudioConnector(AccountJob job, Map<String, Object> c
         logger.info("Copilot Studio connector execution completed: jobId={}", job.getId());
     }
 
+    /**
+     * Execute LiteLLM connector logic.
+     * Downloads the LiteLLM shield binary from Azure Storage and executes it with
+     * config as env vars.
+     */
+    private void executeLiteLLMConnector(AccountJob job, Map<String, Object> config) throws Exception {
+        logger.info("Executing LiteLLM connector: jobId={}", job.getId());
+
+        // Execute connector binary with config
+        executeBinaryConnector(job, config, BINARY_NAME_LITELLM);
+
+        logger.info("LiteLLM connector execution completed: jobId={}", job.getId());
+    }
+
     /**
      * Common method to execute any AI Agent Connector binary.
      * Downloads binary from Azure Storage and executes it with config as environment variables.

apps/dashboard/src/main/java/com/akto/action/AIAgentConnectorImportAction.java

@@ -47,6 +47,10 @@ public class AIAgentConnectorImportAction extends UserAction {
     private String dataverseClientId;
     private String dataverseClientSecret;
 
+    // LiteLLM-specific parameters
+    private String litellmUrl;
+    private String litellmApiKey;
+
     /**
      * Unified method to initiate import for any AI Agent Connector.
      * The connector type is determined by the connectorType parameter.
@@ -147,6 +151,16 @@ private Map<String, String> buildConfig() {
                 config.put(CONFIG_DATAVERSE_CLIENT_SECRET, dataverseClientSecret);
                 break;
 
+            case CONNECTOR_TYPE_LITELLM:
+                if (litellmUrl == null || litellmUrl.isEmpty() || litellmApiKey == null
+                        || litellmApiKey.isEmpty()) {
+                    loggerMaker.error("Missing required LiteLLM configuration", LogDb.DASHBOARD);
+                    return null;
+                }
+                config.put(CONFIG_LITELLM_BASE_URL, litellmUrl);
+                config.put(CONFIG_LITELLM_API_KEY, litellmApiKey);
+                break;
+
             default:
                 loggerMaker.error("Unsupported connector type: " + connectorType, LogDb.DASHBOARD);
                 return null;

apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/quick_start/constants/aiAgentConnectorConstants.js

@@ -7,21 +7,25 @@
 export const CONNECTOR_TYPE_N8N = 'N8N';
 export const CONNECTOR_TYPE_LANGCHAIN = 'LANGCHAIN';
 export const CONNECTOR_TYPE_COPILOT_STUDIO = 'COPILOT_STUDIO';
+export const CONNECTOR_TYPE_LITELLM = 'LITELLM';
 
 // Connector Names (Display)
 export const CONNECTOR_NAME_N8N = 'N8N';
 export const CONNECTOR_NAME_LANGCHAIN = 'Langchain';
 export const CONNECTOR_NAME_COPILOT_STUDIO = 'Copilot Studio';
+export const CONNECTOR_NAME_LITELLM = 'LiteLLM';
 
 // Documentation URLs
 export const DOCS_URL_N8N = 'https://docs.akto.io/traffic-connector/workflow-automation/n8n';
 export const DOCS_URL_LANGCHAIN = 'https://docs.akto.io/traffic-connector/workflow-automation/langchain';
 export const DOCS_URL_COPILOT_STUDIO = 'https://docs.akto.io/traffic-connector/workflow-automation/copilot-studio';
+export const DOCS_URL_LITELLM = 'https://docs.akto.io/traffic-connector/workflow-automation/litellm';
 
 // Recurring Interval Seconds (in seconds)
 export const INTERVAL_N8N = 300; // 5 minutes
 export const INTERVAL_LANGCHAIN = 300; // 5 minutes
 export const INTERVAL_COPILOT_STUDIO = 300; // 5 minutes
+export const INTERVAL_LITELLM = 300; // 5 minutes
 
 // Field Names
 export const FIELD_N8N_URL = 'n8nUrl';
@@ -32,6 +36,8 @@ export const FIELD_DATAVERSE_ENVIRONMENT_URL = 'dataverseEnvironmentUrl';
 export const FIELD_DATAVERSE_TENANT_ID = 'dataverseTenantId';
 export const FIELD_DATAVERSE_CLIENT_ID = 'dataverseClientId';
 export const FIELD_DATAVERSE_CLIENT_SECRET = 'dataverseClientSecret';
+export const FIELD_LITELLM_URL = 'litellmUrl';
+export const FIELD_LITELLM_API_KEY = 'litellmApiKey';
 export const FIELD_DATA_INGESTION_URL = 'dataIngestionUrl';
 
 // Field Types
@@ -43,6 +49,7 @@ export const FIELD_TYPE_PASSWORD = 'password';
 export const DESCRIPTION_N8N = 'Use our N8N feature to capture traffic and instantly send it to your dashboard for real-time insights.';
 export const DESCRIPTION_LANGCHAIN = 'Use our Langchain feature to capture traffic from LangSmith and instantly send it to your dashboard for real-time insights.';
 export const DESCRIPTION_COPILOT_STUDIO = 'Use our Copilot Studio feature to capture conversation data from Azure Dataverse API and instantly send it to your dashboard for real-time insights.';
+export const DESCRIPTION_LITELLM = 'Use our LiteLLM feature to capture traffic from your LiteLLM proxy and instantly send it to your dashboard for real-time insights.';
 
 // N8N Field Configuration
 export const N8N_FIELDS = [
@@ -129,3 +136,28 @@ export const COPILOT_STUDIO_FIELDS = [
         placeholder: 'https://ingestion.example.com'
     }
 ];
+
+// LiteLLM Field Configuration
+export const LITELLM_FIELDS = [
+    {
+        name: FIELD_LITELLM_URL,
+        label: 'LiteLLM Proxy URL',
+        type: FIELD_TYPE_URL,
+        placeholder: 'http://localhost:4000',
+        configKey: FIELD_LITELLM_URL
+    },
+    {
+        name: FIELD_LITELLM_API_KEY,
+        label: 'LiteLLM API Key',
+        type: FIELD_TYPE_PASSWORD,
+        placeholder: '*******',
+        configKey: FIELD_LITELLM_API_KEY
+    },
+    {
+        name: FIELD_DATA_INGESTION_URL,
+        label: 'URL for Data Ingestion Service',
+        type: FIELD_TYPE_URL,
+        placeholder: 'https://ingestion.example.com'
+    }
+];
+

apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/quick_start/transform.js

@@ -44,7 +44,13 @@ import {
     DESCRIPTION_COPILOT_STUDIO,
     DOCS_URL_COPILOT_STUDIO,
     INTERVAL_COPILOT_STUDIO,
-    COPILOT_STUDIO_FIELDS
+    COPILOT_STUDIO_FIELDS,
+    CONNECTOR_TYPE_LITELLM,
+    CONNECTOR_NAME_LITELLM,
+    DESCRIPTION_LITELLM,
+    DOCS_URL_LITELLM,
+    INTERVAL_LITELLM,
+    LITELLM_FIELDS
 } from "./constants/aiAgentConnectorConstants"
 
 const mirroringObj = {
@@ -929,6 +935,21 @@ const copilotStudioImportObj = {
     />
 }
 
+const litellmImportObj = {
+    icon: '/public/litellm.svg',
+    label: "LiteLLM",
+    text: "Use our LiteLLM feature to capture traffic from your LiteLLM proxy and instantly send it to your dashboard for real-time insights.",
+    docsUrl: 'https://docs.akto.io/litellm-import',
+    key: "LITELLM_IMPORT",
+    component: <AIAgentConnectorImport
+        connectorType={CONNECTOR_TYPE_LITELLM}
+        connectorName={CONNECTOR_NAME_LITELLM}
+        description={DESCRIPTION_LITELLM}
+        fields={LITELLM_FIELDS}
+        docsUrl={DOCS_URL_LITELLM}
+        recurringIntervalSeconds={INTERVAL_LITELLM}
+    />
+}
 
 const quick_start_policy_lines= [
     `{`,
@@ -1530,7 +1551,7 @@ const quickStartFunc = {
 
         const aiAgentConnectors = [
             awsBedrockObj, azureAIFoundryObj, databricksObj, googleVertexAIObj, ibmWatsonxObj, customAgentObj,
-            n8nImportObj, langchainImportObj, copilotStudioImportObj
+            n8nImportObj, langchainImportObj, copilotStudioImportObj, litellmImportObj
         ]
 
         // MCP Scan
@@ -1589,7 +1610,7 @@ const quickStartFunc = {
                 connectors["AI Agent Security"] = [
                     awsBedrockObj, azureAIFoundryObj, databricksObj, googleVertexAIObj,
                     ibmWatsonxObj, customAgentObj, agenticShieldObj,  // Include agenticShieldObj for non-demo
-                    n8nImportObj, langchainImportObj, copilotStudioImportObj
+                    n8nImportObj, langchainImportObj, copilotStudioImportObj, litellmImportObj
                 ]
                 connectors["AI Model Security"] = aiScanConnectors
                 connectors["Browser Extension"] = browserExtensions
@@ -1636,13 +1657,13 @@ const quickStartFunc = {
     getConnectorsList: function () {
 
         if(func.checkLocal() || func.isLimitedAccount()){
-            return [burpObj, postmanObj, openApiObj, harFileUploadObj, impervaImportObj, n8nImportObj, langchainImportObj, copilotStudioImportObj]
+            return [burpObj, postmanObj, openApiObj, harFileUploadObj, impervaImportObj, n8nImportObj, langchainImportObj, copilotStudioImportObj, litellmImportObj]
         }
 
         // Combine all categories into connectorsList
         let connectorsList = [
             gcpObj, kubernetesObj, fargateObj, nginxObj, burpObj, postmanObj,
-            openApiObj, beanStalkObj, eksObj, dockerObj, envoyObj, mcpScanObj, mcpProxyObj, mcpGateway, mcpWrapperObj, impervaImportObj, n8nImportObj, langchainImportObj, copilotStudioImportObj,
+            openApiObj, beanStalkObj, eksObj, dockerObj, envoyObj, mcpScanObj, mcpProxyObj, mcpGateway, mcpWrapperObj, impervaImportObj, n8nImportObj, langchainImportObj, copilotStudioImportObj, litellmImportObj,
             harFileUploadObj, kongObj, tcpObj, mirroringObj, hybridSaasObj, apiInventoryFromSourceCodeObj,
             ebpfObj, ebpfMTLSObj, istioObj, pythonObj, awsApiGatewayObj, awsLambdaObj,
             apigeeObj, iisObj, azureObj, cloudflareObj, f5Obj, goObj, haproxyObj, javaObj, kongmeshObj, layer7Obj, nodejsObj, openshiftObj, threescaleObj, githubObj, gitlabObj, bitbucketObj, aktoJaxObj,
@@ -1653,7 +1674,7 @@ const quickStartFunc = {
             connectorsList = connectorsList.concat([
                 geminiObj, openAIObj, claudeObj, deepseekObj, llamaObj, grokObj, customAIObj,
                 awsBedrockObj, azureAIFoundryObj, databricksObj, googleVertexAIObj, ibmWatsonxObj, customAgentObj, agenticShieldObj,
-                n8nImportObj, langchainImportObj, copilotStudioImportObj
+                n8nImportObj, langchainImportObj, copilotStudioImportObj, litellmImportObj
             ])
         }
 

apps/guardrails-service/litellm/.env.example

@@ -0,0 +1,20 @@
+# LiteLLM Configuration
+LITELLM_MASTER_KEY=sk-1234
+ANTHROPIC_API_KEY=your_anthropic_api_key_here
+
+# Database Abstractor Service
+DATABASE_ABSTRACTOR_SERVICE_URL=https://cyborg.akto.io
+DATABASE_ABSTRACTOR_SERVICE_TOKEN=your_jwt_token_here
+
+# Guardrails Configuration
+SYNC_MODE=true
+GUARDRAILS_TIMEOUT=10
+GUARDRAILS_BASE_URL=http://guardrails-service:8080
+
+# Common
+LOG_LEVEL=info
+
+# Database Credentials
+POSTGRES_DB=litellm
+POSTGRES_USER=llmproxy
+POSTGRES_PASSWORD=your_db_password_here

apps/guardrails-service/litellm/config.yaml

@@ -0,0 +1,21 @@
+model_list:
+  - model_name: claude-sonnet-4-20250514
+    litellm_params:
+      model: claude-sonnet-4-20250514
+      api_key: os.environ/ANTHROPIC_API_KEY
+
+litellm_settings:
+  callbacks: [custom_hooks.proxy_handler_instance]
+  drop_params: true
+  set_verbose: false
+  request_timeout: 600
+  num_retries: 2
+
+general_settings:
+  master_key: os.environ/LITELLM_MASTER_KEY
+
+router_settings:
+  routing_strategy: simple-shuffle
+  num_retries: 2
+  timeout: 120
+  retry_after: 5