feat: add portkey provider (promptfoo#819)

examples/portkey-test/README.md

@@ -1,6 +1,7 @@
-To get started, set your OPENAI_API_KEY and PORTKEY_API_KEY environment variables.
+There are two examples:
 
-Next, edit promptfooconfig.yaml and replace the portkey prompt with your own portkey prompt id.
+- prompt_example.yaml shows how to pull from portkey's prompt management platform. It requires you to set PORTKEY_API_KEY and OPENAI_API_KEY environment variables. Replace the portkey prompt with your own portkey prompt id.
+- provider_example.yaml shows how to use portkey's gateway.  It requires the PORTKEY_API_KEY environment variable.
 
 Then run:
 ```

examples/portkey-test/promptfooconfig.yaml → examples/portkey-test/prompt_example.yaml

@@ -1,4 +1,4 @@
-description: 'Portkey.ai integration'
+description: 'Portkey.ai prompt integration'
 
 prompts:
   - 'portkey://pp-test-promp-669f48'

examples/portkey-test/provider_example.yaml

@@ -0,0 +1,13 @@
+description: 'Portkey.ai provider integration'
+
+prompts:
+  - 'Write a short tweet about {{topic}}'
+
+providers:
+  - id: portkey:gpt-3.5-turbo-0613
+    config:
+      portkeyProvider: openai
+
+tests:
+  - vars:
+      topic: bananas

site/docs/integrations/portkey.md

@@ -27,3 +27,30 @@ To reference prompts in Portkey:
 Variables from your promptfoo test cases will be automatically plugged into the Portkey prompt as variables. The resulting prompt will be rendered and returned to promptfoo, and used as the prompt for the test case.
 
 Note that promptfoo does not follow the temperature, model, and other parameters set in Portkey. You must set them in the `providers` configuration yourself.
+
+## Using Portkey gateway
+
+The Portkey AI gateway is directly supported by promptfoo.  See also [portkey's documentation on integrating promptfoo](https://portkey.ai/docs/welcome/integration-guides/promptfoo).
+
+Example:
+
+```yaml
+providers:
+  id: portkey:gpt-4o
+  config:
+    portkeyProvider: openai
+```
+
+More complex portkey configurations are also supported.
+
+```yaml
+providers:
+  id: portkey:gpt-3.5-turbo-0613
+  # Can alternative set environment variables, e.g. PORTKEY_API_KEY
+  portkeyApiKey: xxx
+  portkeyVirtualKey: xxx
+  portkeyMetadata:
+    team: xxx
+  portkeyConfig: xxx
+  portkeyProvider: xxx
+```

src/providers.ts

@@ -48,6 +48,7 @@ import { AwsBedrockCompletionProvider, AwsBedrockEmbeddingProvider } from './pro
 import { PythonProvider } from './providers/pythonCompletion';
 import { CohereChatCompletionProvider } from './providers/cohere';
 import { BAMChatProvider, BAMEmbeddingProvider } from './providers/bam';
+import { PortkeyChatCompletionProvider } from './providers/portkey';
 import { HttpProvider } from './providers/http';
 import { importModule } from './esm';
 
@@ -126,7 +127,10 @@ export async function loadApiProvider(
     env,
   };
   let ret: ApiProvider;
-  if (providerPath.startsWith('file://') && (providerPath.endsWith('.yaml')|| providerPath.endsWith('.yml'))){
+  if (
+    providerPath.startsWith('file://') &&
+    (providerPath.endsWith('.yaml') || providerPath.endsWith('.yml'))
+  ) {
     const filePath = providerPath.slice('file://'.length);
     const yamlContent = yaml.load(fs.readFileSync(filePath, 'utf8')) as ProviderOptions;
     invariant(yamlContent, `Provider config ${filePath} is undefined`);
@@ -205,6 +209,10 @@ export async function loadApiProvider(
         apiKeyEnvar: 'OPENROUTER_API_KEY',
       },
     });
+  } else if (providerPath.startsWith('portkey:')) {
+    const splits = providerPath.split(':');
+    const modelName = splits.slice(1).join(':');
+    ret = new PortkeyChatCompletionProvider(modelName, providerOptions);
   } else if (providerPath.startsWith('anthropic:')) {
     const splits = providerPath.split(':');
     const modelType = splits[1];
@@ -336,7 +344,8 @@ export async function loadApiProvider(
   } else if (providerPath.startsWith('http:') || providerPath.startsWith('https:')) {
     ret = new HttpProvider(providerPath, providerOptions);
   } else if (providerPath === 'promptfoo:redteam:iterative') {
-    const RedteamIterativeProvider = (await import(path.join(__dirname, './redteam/iterative'))).default;
+    const RedteamIterativeProvider = (await import(path.join(__dirname, './redteam/iterative')))
+      .default;
     ret = new RedteamIterativeProvider(providerOptions);
   } else {
     if (providerPath.startsWith('file://')) {

src/providers/portkey.ts

@@ -0,0 +1,46 @@
+import { OpenAiChatCompletionProvider } from './openai';
+
+import type { ProviderOptions } from '../types';
+
+export class PortkeyChatCompletionProvider extends OpenAiChatCompletionProvider {
+  constructor(modelName: string, providerOptions: ProviderOptions) {
+    const headers = [
+      {
+        key: 'x-portkey-api-key',
+        value: process.env.PORTKEY_API_KEY || providerOptions.config?.portkeyApiKey,
+      },
+      {
+        key: 'x-portkey-virtual-key',
+        value: process.env.PORTKEY_VIRTUAL_KEY || providerOptions.config?.portkeyVirtualKey,
+      },
+      {
+        key: 'x-portkey-metadata',
+        value:
+          process.env.PORTKEY_METADATA ||
+          (providerOptions.config?.portkeyMetadata
+            ? JSON.stringify(providerOptions.config?.portkeyMetadata)
+            : undefined),
+      },
+      {
+        key: 'x-portkey-config',
+        value: process.env.PORTKEY_CONFIG || providerOptions.config?.portkeyConfig,
+      },
+      {
+        key: 'x-portkey-provider',
+        value: process.env.PORTKEY_PROVIDER || providerOptions.config?.portkeyProvider,
+      },
+    ].reduce((acc: Record<string, string>, { key, value }) => {
+      if (value) acc[key] = value;
+      return acc;
+    }, {});
+
+    super(modelName, {
+      ...providerOptions,
+      config: {
+        ...providerOptions.config,
+        apiBaseUrl: 'https://api.portkey.ai/v1',
+        headers,
+      },
+    });
+  }
+}