diff --git a/docs/docs/developers/build/connectors/services/openai.md b/docs/docs/developers/build/connectors/services/openai.md index 64ddde07f630..45e07ef45a3e 100644 --- a/docs/docs/developers/build/connectors/services/openai.md +++ b/docs/docs/developers/build/connectors/services/openai.md @@ -48,6 +48,29 @@ For details on managing credentials across environments, see [Configure Local Cr For additional configuration options (model, base URL, API type, etc.), see the [OpenAI connector reference](/reference/project-files/connectors#openai). +### OpenAI-compatible APIs + +The connector can also target APIs that implement the OpenAI chat completions protocol. Configure the provider's URL and model on the connector. If the provider supports JSON mode but not OpenAI's JSON Schema response format, set `structured_output_mode` to `json_object`: + +```yaml +type: connector +driver: openai +api_key: "{{ .env.PROVIDER_API_KEY }}" +base_url: https://llm.example.com/v1 +model: example-model +structured_output_mode: json_object +``` + +For provider-specific request extensions, use the advanced `extra_body` map. Its values must be JSON-serializable, and it cannot override core request fields such as `model`, `messages`, `tools`, `response_format`, or streaming controls. For example, an OpenAI-compatible endpoint can receive a chat-template option as follows: + +```yaml +extra_body: + chat_template_kwargs: + enable_thinking: false +``` + +These options belong to the connector, so different connectors in the same Rill process can use different provider behavior. + ## Deploy to Rill Cloud Rill requires you to explicitly provide an OpenAI API key to use the OpenAI connector. See the [connector reference](/reference/project-files/connectors#openai) for details. diff --git a/runtime/connection_cache.go b/runtime/connection_cache.go index 0c5126d93eee..be580fe38564 100644 --- a/runtime/connection_cache.go +++ b/runtime/connection_cache.go @@ -2,9 +2,10 @@ package runtime import ( "context" + "crypto/sha256" + "encoding/json" "errors" "fmt" - "slices" "strings" "time" @@ -200,24 +201,23 @@ func generateKey(cfg cachedConnectionConfig) string { sb.WriteString(":") sb.WriteString(cfg.driver) sb.WriteString(":") - keys := maps.Keys(cfg.config) - slices.Sort(keys) - for _, key := range keys { - sb.WriteString(key) - sb.WriteString(":") - sb.WriteString(fmt.Sprint(cfg.config[key])) - sb.WriteString(" ") - } + writeConfigHash(&sb, cfg.config) if cfg.provision { sb.WriteString(":provision=true:") - keys := maps.Keys(cfg.provisionArgs) - slices.Sort(keys) - for _, key := range keys { - sb.WriteString(key) - sb.WriteString(":") - sb.WriteString(fmt.Sprint(cfg.provisionArgs[key])) - sb.WriteString(" ") - } + writeConfigHash(&sb, cfg.provisionArgs) } return sb.String() } + +// writeConfigHash adds a deterministic, type-preserving identity for a connector configuration without embedding +// credentials in the cache key. JSON is canonical for the JSON-shaped connector maps produced by the parser (map +// keys are sorted by encoding/json, and strings/maps/slices remain distinct). The typed Go representation is a +// compatibility fallback for legacy driver configs containing values JSON cannot encode. +func writeConfigHash(sb *strings.Builder, config map[string]any) { + canonical, err := json.Marshal(config) + if err != nil { + canonical = []byte(fmt.Sprintf("%#v", config)) + } + sum := sha256.Sum256(canonical) + fmt.Fprintf(sb, "%x", sum) +} diff --git a/runtime/connection_cache_test.go b/runtime/connection_cache_test.go new file mode 100644 index 000000000000..242178b31456 --- /dev/null +++ b/runtime/connection_cache_test.go @@ -0,0 +1,46 @@ +package runtime + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGenerateConnectionKeyPreservesNestedJSONTypes(t *testing.T) { + base := cachedConnectionConfig{instanceID: "instance", name: "connector", driver: "openai"} + + stringConfig := base + stringConfig.config = map[string]any{"extra_body": map[string]any{"extension": "[END]"}} + sliceConfig := base + sliceConfig.config = map[string]any{"extra_body": map[string]any{"extension": []any{"END"}}} + require.NotEqual(t, generateKey(stringConfig), generateKey(sliceConfig), + "a string and a JSON array must never reuse one connector handle") + + mapLikeString := base + mapLikeString.config = map[string]any{"extra_body": "map[a:b]"} + nestedMap := base + nestedMap.config = map[string]any{"extra_body": map[string]any{"a": "b"}} + require.NotEqual(t, generateKey(mapLikeString), generateKey(nestedMap), + "a string and a JSON object must never reuse one connector handle") +} + +func TestGenerateConnectionKeyIsCanonicalAndDoesNotExposeSecrets(t *testing.T) { + left := cachedConnectionConfig{ + instanceID: "instance", name: "connector", driver: "openai", + config: map[string]any{ + "api_key": "super-secret", + "extra_body": map[string]any{"thinking": map[string]any{"type": "disabled"}, "seed": float64(1)}, + }, + } + right := cachedConnectionConfig{ + instanceID: "instance", name: "connector", driver: "openai", + config: map[string]any{ + "extra_body": map[string]any{"seed": float64(1), "thinking": map[string]any{"type": "disabled"}}, + "api_key": "super-secret", + }, + } + + leftKey := generateKey(left) + require.Equal(t, leftKey, generateKey(right), "map insertion order must not change connector identity") + require.NotContains(t, leftKey, "super-secret", "cache keys must not embed credentials") +} diff --git a/runtime/drivers/openai/openai.go b/runtime/drivers/openai/openai.go index 2a5d2691e665..2a58631f7b79 100644 --- a/runtime/drivers/openai/openai.go +++ b/runtime/drivers/openai/openai.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "sort" "strings" "github.com/mitchellh/mapstructure" @@ -20,7 +21,40 @@ import ( "google.golang.org/protobuf/types/known/structpb" ) -const defaultTemperature = 0.1 +const ( + defaultTemperature = 0.1 + structuredOutputModeJSONSchema = "json_schema" + structuredOutputModeJSONObject = "json_object" +) + +// reservedExtraBodyFields are owned by Rill or would change assumptions made +// by Complete (for example that a non-streaming request returns one choice). +// Provider-specific extensions such as thinking and chat_template_kwargs are +// intentionally not reserved. +var reservedExtraBodyFields = map[string]struct{}{ + "audio": {}, + "api_key": {}, + "api_type": {}, + "api_version": {}, + "base_url": {}, + "function_call": {}, + "functions": {}, + "max_completion_tokens": {}, + "max_output_tokens": {}, + "max_tokens": {}, + "messages": {}, + "modalities": {}, + "model": {}, + "n": {}, + "parallel_tool_calls": {}, + "reasoning_effort": {}, + "response_format": {}, + "stream": {}, + "stream_options": {}, + "temperature": {}, + "tool_choice": {}, + "tools": {}, +} func init() { drivers.Register("openai", driver{}) @@ -86,6 +120,22 @@ var spec = drivers.Spec{ Description: "The version of the OpenAI API to use (e.g., '2023-05-15'). Required when APIType is APITypeAzure or APITypeAzureAD", Placeholder: "", }, + { + Key: "structured_output_mode", + Type: drivers.StringPropertyType, + Required: false, + DisplayName: "Structured Output Mode", + Description: "How output schemas are requested: json_schema (default) or json_object for compatible providers that do not support JSON Schema.", + Default: structuredOutputModeJSONSchema, + }, + { + Key: "extra_body", + Type: drivers.UnspecifiedPropertyType, + Required: false, + DisplayName: "Extra Request Body", + Description: "Advanced map of provider-specific JSON fields added to chat completion requests. Core request and response-shape fields cannot be overridden.", + NoPrompt: true, + }, }, ImplementsAI: true, } @@ -110,6 +160,9 @@ func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Cl if conf.APIKey == "" { return nil, errors.New("API key is required") } + if err := conf.validate(); err != nil { + return nil, err + } var opts []option.RequestOption switch strings.ToLower(conf.APIType) { @@ -151,14 +204,47 @@ func (d driver) TertiarySourceConnectors(ctx context.Context, srcProps map[strin } type configProperties struct { - APIKey string `mapstructure:"api_key"` - Model string `mapstructure:"model"` - MaxOutputTokens int64 `mapstructure:"max_output_tokens"` - ReasoningEffort string `mapstructure:"reasoning_effort"` - Temperature *float64 `mapstructure:"temperature"` - BaseURL string `mapstructure:"base_url"` - APIType string `mapstructure:"api_type"` - APIVersion string `mapstructure:"api_version"` + APIKey string `mapstructure:"api_key"` + Model string `mapstructure:"model"` + MaxOutputTokens int64 `mapstructure:"max_output_tokens"` + ReasoningEffort string `mapstructure:"reasoning_effort"` + Temperature *float64 `mapstructure:"temperature"` + BaseURL string `mapstructure:"base_url"` + APIType string `mapstructure:"api_type"` + APIVersion string `mapstructure:"api_version"` + StructuredOutputMode string `mapstructure:"structured_output_mode"` + ExtraBody map[string]any `mapstructure:"extra_body"` +} + +func (c *configProperties) validate() error { + switch c.getStructuredOutputMode() { + case structuredOutputModeJSONSchema, structuredOutputModeJSONObject: + default: + return fmt.Errorf("invalid structured_output_mode %q: must be %q or %q", c.StructuredOutputMode, structuredOutputModeJSONSchema, structuredOutputModeJSONObject) + } + + var reserved []string + for key := range c.ExtraBody { + if _, ok := reservedExtraBodyFields[strings.ToLower(key)]; ok { + reserved = append(reserved, key) + } + } + if len(reserved) > 0 { + sort.Strings(reserved) + return fmt.Errorf("extra_body cannot override core request fields: %s", strings.Join(reserved, ", ")) + } + + if _, err := json.Marshal(c.ExtraBody); err != nil { + return fmt.Errorf("extra_body must contain JSON-serializable values: %w", err) + } + return nil +} + +func (c *configProperties) getStructuredOutputMode() string { + if c.StructuredOutputMode != "" { + return strings.ToLower(c.StructuredOutputMode) + } + return structuredOutputModeJSONSchema } func (c *configProperties) getModel() string { @@ -326,16 +412,34 @@ func (o *openaiHandle) Complete(ctx context.Context, opts *drivers.CompleteOptio if o.config.ReasoningEffort != "" { params.ReasoningEffort = shared.ReasoningEffort(o.config.ReasoningEffort) } + if len(o.config.ExtraBody) > 0 { + params.SetExtraFields(o.config.ExtraBody) + } // Set response format based on output schema if opts.OutputSchema != nil { - params.ResponseFormat = openai.ChatCompletionNewParamsResponseFormatUnion{ - OfJSONSchema: &openai.ResponseFormatJSONSchemaParam{ - JSONSchema: openai.ResponseFormatJSONSchemaJSONSchemaParam{ - Name: "llm_completion_result", - Schema: opts.OutputSchema, + // Fallback for OpenAI-compatible providers without json_schema support: + // degrade to json_object and inject the schema as an explicit instruction. + if o.config.getStructuredOutputMode() == structuredOutputModeJSONObject { + schemaJSON, err := json.Marshal(opts.OutputSchema) + if err != nil { + return nil, fmt.Errorf("failed to marshal output schema: %w", err) + } + schemaInstruction := openai.SystemMessage( + "Return ONLY a single valid JSON object that conforms exactly to this JSON Schema (no prose, no markdown fences): " + string(schemaJSON)) + params.Messages = append([]openai.ChatCompletionMessageParamUnion{schemaInstruction}, params.Messages...) + params.ResponseFormat = openai.ChatCompletionNewParamsResponseFormatUnion{ + OfJSONObject: &shared.ResponseFormatJSONObjectParam{}, + } + } else { + params.ResponseFormat = openai.ChatCompletionNewParamsResponseFormatUnion{ + OfJSONSchema: &openai.ResponseFormatJSONSchemaParam{ + JSONSchema: openai.ResponseFormatJSONSchemaJSONSchemaParam{ + Name: "llm_completion_result", + Schema: opts.OutputSchema, + }, }, - }, + } } } diff --git a/runtime/drivers/openai/openai_request_test.go b/runtime/drivers/openai/openai_request_test.go new file mode 100644 index 000000000000..48b255158a9e --- /dev/null +++ b/runtime/drivers/openai/openai_request_test.go @@ -0,0 +1,219 @@ +package openai + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/google/jsonschema-go/jsonschema" + aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1" + "github.com/rilldata/rill/runtime/drivers" + "github.com/stretchr/testify/require" +) + +func TestCompleteAppliesConnectorRequestBehavior(t *testing.T) { + fake := newFakeChatCompletionsServer([]string{completionResponse(`{"answer":"ok"}`)}) + defer fake.Close() + + ai := openTestAI(t, fake.URL, map[string]any{ + "structured_output_mode": structuredOutputModeJSONObject, + "extra_body": map[string]any{ + "thinking": map[string]any{"type": "disabled"}, + }, + }) + + _, err := ai.Complete(t.Context(), &drivers.CompleteOptions{ + Messages: []*aiv1.CompletionMessage{textMessage("user", "answer as JSON")}, + OutputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "answer": {Type: "string"}, + }, + Required: []string{"answer"}, + }, + }) + require.NoError(t, err) + + body := fake.request(t, 0) + require.Equal(t, map[string]any{"type": "disabled"}, body["thinking"]) + require.Equal(t, map[string]any{"type": "json_object"}, body["response_format"]) + + messages := requireJSONArray(t, body["messages"]) + require.Len(t, messages, 2) + schemaInstruction := requireJSONObject(t, messages[0]) + require.Equal(t, "system", schemaInstruction["role"]) + require.Contains(t, schemaInstruction["content"], "Return ONLY a single valid JSON object") + require.Contains(t, schemaInstruction["content"], `"required":["answer"]`) + require.Equal(t, "user", requireJSONObject(t, messages[1])["role"]) +} + +func TestCompletePassesNestedExtraBody(t *testing.T) { + fake := newFakeChatCompletionsServer([]string{completionResponse("ok")}) + defer fake.Close() + + ai := openTestAI(t, fake.URL, map[string]any{ + "extra_body": map[string]any{ + "chat_template_kwargs": map[string]any{"enable_thinking": false}, + }, + }) + + _, err := ai.Complete(t.Context(), &drivers.CompleteOptions{ + Messages: []*aiv1.CompletionMessage{textMessage("user", "hello")}, + }) + require.NoError(t, err) + + body := fake.request(t, 0) + require.Equal(t, map[string]any{"enable_thinking": false}, body["chat_template_kwargs"]) + require.NotContains(t, body, "response_format") +} + +func TestCompleteDefaultsToJSONSchema(t *testing.T) { + fake := newFakeChatCompletionsServer([]string{completionResponse(`{"answer":"ok"}`)}) + defer fake.Close() + + ai := openTestAI(t, fake.URL, nil) + _, err := ai.Complete(t.Context(), &drivers.CompleteOptions{ + Messages: []*aiv1.CompletionMessage{textMessage("user", "answer as JSON")}, + OutputSchema: &jsonschema.Schema{Type: "object"}, + }) + require.NoError(t, err) + + body := fake.request(t, 0) + require.Equal(t, "json_schema", requireJSONObject(t, body["response_format"])["type"]) + require.Len(t, requireJSONArray(t, body["messages"]), 1) +} + +func TestOpenValidatesProviderRequestBehavior(t *testing.T) { + t.Run("structured output mode", func(t *testing.T) { + _, err := (driver{}).Open("", "", map[string]any{ + "api_key": "test-key", + "structured_output_mode": "xml", + }, nil, nil, nil) + require.ErrorContains(t, err, `invalid structured_output_mode "xml"`) + }) + + t.Run("reserved extra body fields", func(t *testing.T) { + _, err := (driver{}).Open("", "", map[string]any{ + "api_key": "test-key", + "extra_body": map[string]any{ + "audio": map[string]any{"format": "wav"}, + "modalities": []any{"text", "audio"}, + "Tools": []any{}, + "model": "other-model", + "stream": true, + }, + }, nil, nil, nil) + require.EqualError(t, err, "extra_body cannot override core request fields: Tools, audio, modalities, model, stream") + }) + + t.Run("non JSON extra body", func(t *testing.T) { + _, err := (driver{}).Open("", "", map[string]any{ + "api_key": "test-key", + "extra_body": map[string]any{ + "extension": make(chan int), + }, + }, nil, nil, nil) + require.ErrorContains(t, err, "extra_body must contain JSON-serializable values") + }) +} + +func textMessage(role, text string) *aiv1.CompletionMessage { + return &aiv1.CompletionMessage{ + Role: role, + Content: []*aiv1.ContentBlock{{ + BlockType: &aiv1.ContentBlock_Text{Text: text}, + }}, + } +} + +func openTestAI(t *testing.T, serverURL string, config map[string]any) drivers.AIService { + t.Helper() + if config == nil { + config = make(map[string]any) + } + config["api_key"] = "test-key" + config["base_url"] = serverURL + "/v1" + config["model"] = "test-model" + + handle, err := (driver{}).Open("", "", config, nil, nil, nil) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, handle.Close()) }) + ai, ok := handle.AsAI("") + require.True(t, ok) + return ai +} + +type fakeChatCompletionsServer struct { + *httptest.Server + mu sync.Mutex + requests []map[string]any + responses []string +} + +func newFakeChatCompletionsServer(responses []string) *fakeChatCompletionsServer { + fake := &fakeChatCompletionsServer{responses: responses} + fake.Server = httptest.NewServer(http.HandlerFunc(fake.serveHTTP)) + return fake +} + +func (f *fakeChatCompletionsServer) serveHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "unexpected method "+r.Method, http.StatusMethodNotAllowed) + return + } + if r.URL.Path != "/v1/chat/completions" { + http.Error(w, "unexpected path "+r.URL.Path, http.StatusNotFound) + return + } + + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, "invalid JSON request: "+err.Error(), http.StatusBadRequest) + return + } + + f.mu.Lock() + idx := len(f.requests) + f.requests = append(f.requests, body) + if idx >= len(f.responses) { + f.mu.Unlock() + http.Error(w, "unexpected request", http.StatusInternalServerError) + return + } + response := f.responses[idx] + f.mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(response)) +} + +func (f *fakeChatCompletionsServer) request(t *testing.T, idx int) map[string]any { + t.Helper() + f.mu.Lock() + defer f.mu.Unlock() + require.Greater(t, len(f.requests), idx) + return f.requests[idx] +} + +func completionResponse(content string) string { + encoded, _ := json.Marshal(content) + return `{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"test-model","choices":[{"index":0,"message":{"role":"assistant","content":` + + string(encoded) + + `},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}` +} + +func requireJSONObject(t *testing.T, value any) map[string]any { + t.Helper() + result, ok := value.(map[string]any) + require.True(t, ok, "expected JSON object, got %T", value) + return result +} + +func requireJSONArray(t *testing.T, value any) []any { + t.Helper() + result, ok := value.([]any) + require.True(t, ok, "expected JSON array, got %T", value) + return result +}