Merge branch 'main' into namwoam/cohere

.pre-commit-config.yaml

@@ -6,6 +6,8 @@ repos:
       - id: check-json
       - id: check-merge-conflict
       - id: end-of-file-fixer
+        exclude: tools/compogen/cmd/testdata
+        exclude_types: [svg,mdx]
       - id: trailing-whitespace
       - id: pretty-format-json
         args: [--autofix, --no-sort-keys]

ai/anthropic/v0/README.mdx

@@ -48,7 +48,6 @@ Provide text outputs in response to text inputs.
 | Model Name (required) | `model-name` | string | The Anthropic model to be used. |
 | Prompt (required) | `prompt` | string | The prompt text |
 | System message | `system-message` | string | The system message helps set the behavior of the assistant. For example, you can modify the personality of the assistant or provide specific instructions about how it should behave throughout the conversation. By default, the model’s behavior is using a generic message as "You are a helpful assistant." |
-| Extra Parameters | `extra-params` | object | Extra Parameters |
 | Prompt Images | `prompt-images` | array[string] | The prompt images (Note: The prompt images will be injected in the order they are provided to the 'prompt' message. Anthropic doesn't support sending images via image-url, use this field instead) |
 | Chat history | `chat-history` | array[object] | Incorporate external chat history, specifically previous messages within the conversation. Please note that System Message will be ignored and will not have any effect when this field is populated. Each message should adhere to the format: : \{"role": "The message role, i.e. 'system', 'user' or 'assistant'", "content": "message content"\}. |
 | Seed | `seed` | integer | The seed (Note: Not supported by Anthropic Models) |
@@ -61,3 +60,4 @@ Provide text outputs in response to text inputs.
 | Output | ID | Type | Description |
 | :--- | :--- | :--- | :--- |
 | Text | `text` | string | Model Output |
+| Usage (optional) | `usage` | object | Usage tokens in Anthropic |

ai/anthropic/v0/component_test.go

@@ -42,7 +42,7 @@ func TestComponent_Execute(t *testing.T) {
 	}{
 		{
 			name:        "text generation",
-			task:        textGenerationTask,
+			task:        TextGenerationTask,
 			path:        messagesPath,
 			contentType: httpclient.MIMETypeJSON,
 		},
@@ -191,9 +191,14 @@ func TestComponent_Generation(t *testing.T) {
 		input    map[string]any
 		wantResp messagesOutput
 	}{
-
-		input:    map[string]any{"prompt": "Hi! What's your name?", "chat-history": mockHistory},
-		wantResp: messagesOutput{Text: "Hi! My name is Claude. (messageCount: 3)"},
+		input: map[string]any{"prompt": "Hi! What's your name?", "chat-history": mockHistory},
+		wantResp: messagesOutput{
+			Text: "Hi! My name is Claude. (messageCount: 3)",
+			Usage: messagesUsage{
+				InputTokens:  10,
+				OutputTokens: 25,
+			},
+		},
 	}
 
 	c.Run("ok - generation", func(c *qt.C) {
@@ -203,7 +208,7 @@ func TestComponent_Generation(t *testing.T) {
 		c.Assert(err, qt.IsNil)
 
 		e := &execution{
-			ComponentExecution: base.ComponentExecution{Component: connector, SystemVariables: nil, Setup: setup, Task: textGenerationTask},
+			ComponentExecution: base.ComponentExecution{Component: connector, SystemVariables: nil, Setup: setup, Task: TextGenerationTask},
 			client:             &MockAnthropicClient{},
 		}
 		e.execute = e.generateText

ai/anthropic/v0/config/definition.json

@@ -17,4 +17,3 @@
   "sourceUrl": "https://github.com/instill-ai/component/blob/main/ai/anthropic/v0",
   "releaseStage": "RELEASE_STAGE_ALPHA"
 }
-

ai/anthropic/v0/config/setup.json

@@ -1,29 +1,28 @@
 {
-    "$schema": "http://json-schema.org/draft-07/schema#",
-    "additionalProperties": true,
-    "properties": {
-      "api-key": {
-        "description": "Fill your Anthropic API key. To find your keys, visit the Anthropic console page.",
-        "instillUpstreamTypes": [
-          "reference"
-        ],
-        "instillAcceptFormats": [
-          "string"
-        ],
-        "instillSecret": true,
-        "instillCredential": false,
-        "instillUIOrder": 0,
-        "title": "API Key",
-        "type": "string"
-      }
-    },
-    "required": [
-      "api-key"
-    ],
-    "instillEditOnNodeFields": [
-      "api-key"
-    ],
-    "title": "Anthropic Connection",
-    "type": "object"
-  }
-
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "additionalProperties": true,
+  "properties": {
+    "api-key": {
+      "description": "Fill your Anthropic API key. To find your keys, visit the Anthropic console page.",
+      "instillUpstreamTypes": [
+        "reference"
+      ],
+      "instillAcceptFormats": [
+        "string"
+      ],
+      "instillSecret": true,
+      "instillCredential": true,
+      "instillUIOrder": 0,
+      "title": "API Key",
+      "type": "string"
+    }
+  },
+  "required": [
+    "api-key"
+  ],
+  "instillEditOnNodeFields": [
+    "api-key"
+  ],
+  "title": "Anthropic Connection",
+  "type": "object"
+}

ai/anthropic/v0/config/tasks.json

@@ -25,7 +25,7 @@
             "description": "The type of the content part.",
             "enum": [
               "text",
-              "image_url" 
+              "image_url"
             ],
             "instillFormat": "string",
             "type": "string"
@@ -61,59 +61,30 @@
       "title": "Chat Message",
       "type": "object"
     },
-    "common": {
-      "description": "Input",
-      "instillEditOnNodeFields": [
-        "image-base64",
-        "model-name"
-      ],
-      "instillUIOrder": 0,
+    "usage": {
+      "description": "Usage tokens in Anthropic",
+      "instillUIOrder": 1,
       "properties": {
-        "image-base64": {
-          "description": "Image base64",
-          "instillAcceptFormats": [
-            "image/*"
-          ],
+        "input-tokens": {
+          "description": "The input tokens used by Anthropic",
+          "instillFormat": "number",
           "instillUIOrder": 2,
-          "instillUpstreamTypes": [
-            "reference"
-          ],
-          "title": "Image",
-          "type": "string"
+          "title": "Input Tokens",
+          "type": "number"
         },
-        "model-name": {
-          "description": "The Instill Model model to be used.",
-          "instillAcceptFormats": [
-            "string"
-          ],
-          "instillUIOrder": 0,
-          "instillUpstreamTypes": [
-            "value",
-            "reference",
-            "template"
-          ],
-          "title": "Model Name",
-          "type": "string"
+        "output-tokens": {
+          "description": "The output tokens used by Anthropic",
+          "instillFormat": "number",
+          "instillUIOrder": 3,
+          "title": "Output Tokens",
+          "type": "number"
         }
       },
       "required": [
-        "image-base64",
-        "model-name"
+        "input-tokens",
+        "output-tokens"
       ],
-      "title": "Input",
-      "type": "object"
-    },
-    "extra-params": {
-      "description": "Extra Parameters",
-      "instillAcceptFormats": [
-        "semi-structured/object"
-      ],
-      "instillUIOrder": 3,
-      "instillUpstreamTypes": [
-        "reference"
-      ],
-      "required": [],
-      "title": "Extra Parameters",
+      "title": "Usage",
       "type": "object"
     }
   },
@@ -144,9 +115,6 @@
           "title": "Chat history",
           "type": "array"
         },
-        "extra-params": {
-          "$ref": "#/$defs/extra-params"
-        },
         "max-new-tokens": {
           "default": 50,
           "description": "The maximum number of tokens for model to generate",
@@ -179,6 +147,17 @@
             "reference",
             "template"
           ],
+          "instillCredentialMap": {
+            "values": [
+              "claude-3-5-sonnet-20240620",
+              "claude-3-opus-20240229",
+              "claude-3-sonnet-20240229",
+              "claude-3-haiku-20240307"
+            ],
+            "targets": [
+              "setup.api-key"
+            ]
+          },
           "title": "Model Name",
           "type": "string"
         },
@@ -289,6 +268,9 @@
           "instillUIMultiline": true,
           "title": "Text",
           "type": "string"
+        },
+        "usage": {
+          "$ref": "#/$defs/usage"
         }
       },
       "required": [
@@ -298,4 +280,4 @@
       "type": "object"
     }
   }
-}
+}

ai/anthropic/v0/main.go

@@ -16,7 +16,7 @@ import (
 )
 
 const (
-	textGenerationTask = "TASK_TEXT_GENERATION_CHAT"
+	TextGenerationTask = "TASK_TEXT_GENERATION_CHAT"
 	cfgAPIKey          = "api-key"
 	host               = "https://api.anthropic.com"
 	messagesPath       = "/v1/messages"
@@ -38,6 +38,9 @@ var (
 
 type component struct {
 	base.Component
+
+	usageHandlerCreator base.UsageHandlerCreator
+	secretAPIKey        string
 }
 
 type AnthropicClient interface {
@@ -70,7 +73,13 @@ type messagesReq struct {
 }
 
 type messagesOutput struct {
-	Text string `json:"text"`
+	Text  string        `json:"text"`
+	Usage messagesUsage `json:"usage"`
+}
+
+type messagesUsage struct {
+	InputTokens  int `json:"input-tokens"`
+	OutputTokens int `json:"output-tokens"`
 }
 
 type message struct {
@@ -109,24 +118,75 @@ func Init(bc base.Component) *component {
 
 type execution struct {
 	base.ComponentExecution
-	execute func(*structpb.Struct) (*structpb.Struct, error)
-	client  AnthropicClient
+
+	execute    func(*structpb.Struct) (*structpb.Struct, error)
+	client     AnthropicClient
+	usesSecret bool
+}
+
+// WithSecrets loads secrets into the connector, which can be used to configure
+// it with globaly defined parameters.
+func (c *component) WithSecrets(s map[string]any) *component {
+	c.secretAPIKey = base.ReadFromSecrets(cfgAPIKey, s)
+	return c
+}
+
+// WithUsageHandlerCreator overrides the UsageHandlerCreator method.
+func (c *component) WithUsageHandlerCreator(newUH base.UsageHandlerCreator) *component {
+	c.usageHandlerCreator = newUH
+	return c
+}
+
+// UsageHandlerCreator returns a function to initialize a UsageHandler.
+func (c *component) UsageHandlerCreator() base.UsageHandlerCreator {
+	if c.usageHandlerCreator == nil {
+		return c.Component.UsageHandlerCreator()
+	}
+	return c.usageHandlerCreator
 }
 
 func (c *component) CreateExecution(sysVars map[string]any, setup *structpb.Struct, task string) (*base.ExecutionWrapper, error) {
+
+	resolvedSetup, resolved, err := c.resolveSecrets(setup)
+	if err != nil {
+		return nil, err
+	}
+
 	e := &execution{
 		ComponentExecution: base.ComponentExecution{Component: c, SystemVariables: sysVars, Task: task, Setup: setup},
-		client:             newClient(getAPIKey(setup), getBasePath(setup), c.GetLogger()),
+		client:             newClient(getAPIKey(resolvedSetup), getBasePath(resolvedSetup), c.GetLogger()),
+		usesSecret:         resolved,
 	}
 	switch task {
-	case textGenerationTask:
+	case TextGenerationTask:
 		e.execute = e.generateText
 	default:
 		return nil, fmt.Errorf("unsupported task")
 	}
 	return &base.ExecutionWrapper{Execution: e}, nil
 }
 
+// resolveSecrets looks for references to a global secret in the setup
+// and replaces them by the global secret injected during initialization.
+func (c *component) resolveSecrets(conn *structpb.Struct) (*structpb.Struct, bool, error) {
+
+	apiKey := conn.GetFields()[cfgAPIKey].GetStringValue()
+	if apiKey != base.SecretKeyword {
+		return conn, false, nil
+	}
+
+	if c.secretAPIKey == "" {
+		return nil, false, base.NewUnresolvedSecret(cfgAPIKey)
+	}
+
+	conn.GetFields()[cfgAPIKey] = structpb.NewStringValue(c.secretAPIKey)
+	return conn, true, nil
+}
+
+func (e *execution) UsesSecret() bool {
+	return e.usesSecret
+}
+
 func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	outputs := make([]*structpb.Struct, len(inputs))
 
@@ -214,6 +274,10 @@ func (e *execution) generateText(in *structpb.Struct) (*structpb.Struct, error)
 
 	outputStruct := messagesOutput{
 		Text: "",
+		Usage: messagesUsage{
+			InputTokens:  resp.Usage.InputTokens,
+			OutputTokens: resp.Usage.OutputTokens,
+		},
 	}
 	for _, c := range resp.Content {
 		outputStruct.Text += c.Text