feat: v0.2.0 — sync with upstream Python SDK v2.0.4

Add ToolReferenceChunk, ToolFileChunk, BuiltInConnector enum, ReferenceID union type (int|string), GuardrailConfig with v1/v2 moderation, ConnectorTool for custom connectors, and guardrails field on chat/agents/conversation requests. Add AudioTranscriptionRealtime and AudioSpeech to ModelCapabilities. Move GuardrailConfig from agents/ to chat/ as shared base type. Remove bundled OpenAPI spec; SDK now tracks upstream Python SDK. BREAKING: ReferenceChunk.ReferenceIDs changed from []int to []ReferenceID. Use IntRef(n) / StringRef(s) constructors.
2026-03-17 11:23:58 +01:00
parent d2a40d17af
commit ea47d14b38
17 changed files with 472 additions and 9553 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -29,3 +29,4 @@ vendor/
 # OS
 .DS_Store
 Thumbs.db
+client-python/
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,40 @@
+## v0.2.0 — 2026-03-17
+
+Sync with upstream Python SDK v2.0.4. Upstream reference changed from OpenAPI
+spec to official Python SDK (https://github.com/mistralai/client-python).
+
+### Breaking Changes
+
+- **`ReferenceChunk.ReferenceIDs`** changed from `[]int` to `[]ReferenceID`.
+  The API now returns mixed integer and string identifiers. Use `IntRef(n)` and
+  `StringRef(s)` constructors; read back with `.Int()` and `.IsString()`.
+- **`agents.GuardrailConfig` and `agents.ModerationLLMV1Config`** moved to
+  `chat.GuardrailConfig` and `chat.ModerationLLMV1Config`. The types are now
+  shared across chat, agents, and conversation packages.
+
+### Added
+
+- **`ToolReferenceChunk`** — new content chunk type for tool references
+  returned by built-in connectors (web search, code interpreter, etc.).
+- **`ToolFileChunk`** — new content chunk type for tool-generated files.
+- **`BuiltInConnector`** constants — `ConnectorWebSearch`,
+  `ConnectorWebSearchPremium`, `ConnectorCodeInterpreter`,
+  `ConnectorImageGeneration`, `ConnectorDocumentLibrary`.
+- **`ModerationLLMV2Config`** — v2 moderation guardrail with split
+  `dangerous`/`criminal` categories and new `jailbreaking` category.
+- **`GuardrailConfig`** on `chat.CompletionRequest`,
+  `agents.CompletionRequest`, `conversation.StartRequest`, and
+  `conversation.RestartRequest`.
+- **`ConnectorTool`** — new agent tool type for custom connectors with
+  `ConnectorAuth` (api-key / oauth2-token authorization).
+- **`ModelCapabilities`** — added `AudioTranscriptionRealtime` and
+  `AudioSpeech` fields.
+
+### Removed
+
+- Bundled `docs/openapi.yaml`. The SDK now tracks the upstream Python SDK
+  directly as its reference implementation.
+
 ## v0.1.0 — 2026-03-05

 Initial release.
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ The most complete Go client for the [Mistral AI API](https://docs.mistral.ai/).

 **Hand-written, not generated.** Idiomatic Go with sealed interfaces, discriminated unions, and functional options — not a Speakeasy/OpenAPI auto-gen dump with `any` everywhere.

-**Test-driven.** 126 tests with race detection clean. Every endpoint tested against mock servers; integration tests against the real API.
+**Test-driven.** 193 tests with race detection clean. Every endpoint tested against mock servers; integration tests against the real API.

 ## Install

@@ -213,6 +213,11 @@ if err != nil {
 }
 ```

+## Upstream Reference
+
+This SDK tracks the [official Mistral Python SDK](https://github.com/mistralai/client-python)
+as its upstream reference for API surface and type definitions.
+
 ## License

 [MIT](LICENSE)
--- a/agents/agent.go
+++ b/agents/agent.go
@@ -3,6 +3,8 @@ package agents
 import (
 	"encoding/json"
 	"fmt"
+
+	"somegit.dev/vikingowl/mistral-go-sdk/chat"
 )

 // AgentTool is a sealed interface for agent tool types.
@@ -59,6 +61,22 @@ type DocumentLibraryTool struct {

 func (*DocumentLibraryTool) agentToolType() string { return "document_library" }

+// ConnectorAuth holds authorization for a custom connector.
+type ConnectorAuth struct {
+	Type  string `json:"type"`
+	Value string `json:"value"`
+}
+
+// ConnectorTool represents a custom connector tool.
+type ConnectorTool struct {
+	Type              string             `json:"type"`
+	ConnectorID       string             `json:"connector_id"`
+	Authorization     *ConnectorAuth     `json:"authorization,omitempty"`
+	ToolConfiguration *ToolConfiguration `json:"tool_configuration,omitempty"`
+}
+
+func (*ConnectorTool) agentToolType() string { return "connector" }
+
 // UnknownAgentTool holds an unrecognized tool type.
 type UnknownAgentTool struct {
 	Type string
@@ -99,6 +117,9 @@ func UnmarshalAgentTool(data []byte) (AgentTool, error) {
 	case "document_library":
 		var t DocumentLibraryTool
 		return &t, json.Unmarshal(data, &t)
+	case "connector":
+		var t ConnectorTool
+		return &t, json.Unmarshal(data, &t)
 	default:
 		return &UnknownAgentTool{Type: probe.Type, Raw: json.RawMessage(data)}, nil
 	}
@@ -151,7 +172,7 @@ type Agent struct {
 	Description    *string           `json:"description,omitempty"`
 	Tools          AgentTools        `json:"tools,omitempty"`
 	CompletionArgs *CompletionArgs   `json:"completion_args,omitempty"`
-	Guardrails     []GuardrailConfig `json:"guardrails,omitempty"`
+	Guardrails     []chat.GuardrailConfig `json:"guardrails,omitempty"`
 	Handoffs       []string          `json:"handoffs,omitempty"`
 	Metadata       map[string]any    `json:"metadata,omitempty"`
 	VersionMessage *string           `json:"version_message,omitempty"`
@@ -165,7 +186,7 @@ type CreateRequest struct {
 	Description    *string           `json:"description,omitempty"`
 	Tools          AgentTools        `json:"tools,omitempty"`
 	CompletionArgs *CompletionArgs   `json:"completion_args,omitempty"`
-	Guardrails     []GuardrailConfig `json:"guardrails,omitempty"`
+	Guardrails     []chat.GuardrailConfig `json:"guardrails,omitempty"`
 	Handoffs       []string          `json:"handoffs,omitempty"`
 	Metadata       map[string]any    `json:"metadata,omitempty"`
 	VersionMessage *string           `json:"version_message,omitempty"`
@@ -179,7 +200,7 @@ type UpdateRequest struct {
 	Description    *string           `json:"description,omitempty"`
 	Tools          AgentTools        `json:"tools,omitempty"`
 	CompletionArgs *CompletionArgs   `json:"completion_args,omitempty"`
-	Guardrails     []GuardrailConfig `json:"guardrails,omitempty"`
+	Guardrails     []chat.GuardrailConfig `json:"guardrails,omitempty"`
 	Handoffs       []string          `json:"handoffs,omitempty"`
 	DeploymentChat *bool             `json:"deployment_chat,omitempty"`
 	Metadata       map[string]any    `json:"metadata,omitempty"`
@@ -223,20 +244,6 @@ type CompletionArgs struct {
 	ToolChoice       *string         `json:"tool_choice,omitempty"`
 }

-// GuardrailConfig configures moderation guardrails for an agent.
-type GuardrailConfig struct {
-	BlockOnError    bool                   `json:"block_on_error"`
-	ModerationLLMV1 *ModerationLLMV1Config `json:"moderation_llm_v1"`
-}
-
-// ModerationLLMV1Config configures the moderation LLM guardrail.
-type ModerationLLMV1Config struct {
-	ModelName                string          `json:"model_name,omitempty"`
-	CustomCategoryThresholds json.RawMessage `json:"custom_category_thresholds,omitempty"`
-	IgnoreOtherCategories    bool            `json:"ignore_other_categories,omitempty"`
-	Action                   string          `json:"action,omitempty"`
-}
-
 // ToolConfiguration holds include/exclude/confirmation lists for tools.
 type ToolConfiguration struct {
 	Exclude              []string `json:"exclude,omitempty"`
--- a/agents/agent_test.go
+++ b/agents/agent_test.go
@@ -66,6 +66,30 @@ func TestAgentTools_RoundTrip(t *testing.T) {
 	}
 }

+func TestUnmarshalAgentTool_Connector(t *testing.T) {
+	data := []byte(`{"type":"connector","connector_id":"my-connector","authorization":{"type":"api-key","value":"sk-test"}}`)
+	tool, err := UnmarshalAgentTool(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ct, ok := tool.(*ConnectorTool)
+	if !ok {
+		t.Fatalf("expected *ConnectorTool, got %T", tool)
+	}
+	if ct.ConnectorID != "my-connector" {
+		t.Errorf("got connector_id %q", ct.ConnectorID)
+	}
+	if ct.Authorization == nil {
+		t.Fatal("expected authorization")
+	}
+	if ct.Authorization.Type != "api-key" {
+		t.Errorf("got auth type %q", ct.Authorization.Type)
+	}
+	if ct.Authorization.Value != "sk-test" {
+		t.Errorf("got auth value %q", ct.Authorization.Value)
+	}
+}
+
 func TestAgent_UnmarshalWithTools(t *testing.T) {
 	data := []byte(`{
 		"id":"ag-1","object":"agent","name":"A","model":"m",
--- a/agents/doc.go
+++ b/agents/doc.go
@@ -1,3 +1,10 @@
 // Package agents provides types for the Mistral agents API,
 // including agent CRUD operations and agent chat completions.
+//
+// # Tool Types
+//
+// Agents support multiple tool types via the [AgentTool] sealed interface:
+// [FunctionTool], [WebSearchTool], [WebSearchPremiumTool],
+// [CodeInterpreterTool], [ImageGenerationTool], [DocumentLibraryTool],
+// and [ConnectorTool] for custom connectors.
 package agents
--- a/agents/request.go
+++ b/agents/request.go
@@ -23,8 +23,9 @@ type CompletionRequest struct {
 	N                 *int                 `json:"n,omitempty"`
 	ParallelToolCalls *bool                `json:"parallel_tool_calls,omitempty"`
 	Metadata          map[string]any       `json:"metadata,omitempty"`
-	Prediction        *chat.Prediction     `json:"prediction,omitempty"`
-	PromptMode        *chat.PromptMode     `json:"prompt_mode,omitempty"`
+	Prediction        *chat.Prediction       `json:"prediction,omitempty"`
+	PromptMode        *chat.PromptMode       `json:"prompt_mode,omitempty"`
+	Guardrails        []chat.GuardrailConfig `json:"guardrails,omitempty"`
 	stream            bool
 }

--- a/chat/content.go
+++ b/chat/content.go
@@ -3,6 +3,18 @@ package chat
 import (
 	"encoding/json"
 	"fmt"
+	"strconv"
+)
+
+// BuiltInConnector identifies a built-in connector type.
+type BuiltInConnector string
+
+const (
+	ConnectorWebSearch        BuiltInConnector = "web_search"
+	ConnectorWebSearchPremium BuiltInConnector = "web_search_premium"
+	ConnectorCodeInterpreter  BuiltInConnector = "code_interpreter"
+	ConnectorImageGeneration  BuiltInConnector = "image_generation"
+	ConnectorDocumentLibrary  BuiltInConnector = "document_library"
 )

 // ContentChunk is a sealed interface for message content parts.
@@ -112,9 +124,65 @@ func (c *FileChunk) MarshalJSON() ([]byte, error) {
 	})
 }

+// ReferenceID is a reference identifier that can be an integer or string.
+// Use [IntRef] or [StringRef] constructors.
+type ReferenceID struct {
+	raw      string
+	isString bool
+}
+
+// IntRef creates an integer reference ID.
+func IntRef(n int) ReferenceID {
+	return ReferenceID{raw: strconv.Itoa(n)}
+}
+
+// StringRef creates a string reference ID.
+func StringRef(s string) ReferenceID {
+	return ReferenceID{raw: s, isString: true}
+}
+
+// String returns the string representation.
+func (id ReferenceID) String() string { return id.raw }
+
+// Int returns the integer value and true if this is a numeric reference.
+func (id ReferenceID) Int() (int, bool) {
+	if id.isString {
+		return 0, false
+	}
+	n, err := strconv.Atoi(id.raw)
+	return n, err == nil
+}
+
+// IsString reports whether this is a string reference.
+func (id ReferenceID) IsString() bool { return id.isString }
+
+func (id ReferenceID) MarshalJSON() ([]byte, error) {
+	if id.isString {
+		return json.Marshal(id.raw)
+	}
+	return []byte(id.raw), nil
+}
+
+func (id *ReferenceID) UnmarshalJSON(data []byte) error {
+	if len(data) == 0 {
+		return nil
+	}
+	if data[0] == '"' {
+		var s string
+		if err := json.Unmarshal(data, &s); err != nil {
+			return err
+		}
+		id.raw = s
+		id.isString = true
+		return nil
+	}
+	id.raw = string(data)
+	return nil
+}
+
 // ReferenceChunk represents a reference content part.
 type ReferenceChunk struct {
-	ReferenceIDs []int `json:"reference_ids"`
+	ReferenceIDs []ReferenceID `json:"reference_ids"`
 }

 func (*ReferenceChunk) contentChunk() {}
@@ -192,6 +260,49 @@ func (c *AudioChunk) MarshalJSON() ([]byte, error) {
 	})
 }

+// ToolReferenceChunk represents a tool reference content part.
+type ToolReferenceChunk struct {
+	Tool        string  `json:"tool"`
+	Title       string  `json:"title"`
+	URL         *string `json:"url,omitempty"`
+	Favicon     *string `json:"favicon,omitempty"`
+	Description *string `json:"description,omitempty"`
+}
+
+func (*ToolReferenceChunk) contentChunk() {}
+
+func (c *ToolReferenceChunk) MarshalJSON() ([]byte, error) {
+	type alias ToolReferenceChunk
+	return json.Marshal(&struct {
+		Type string `json:"type"`
+		*alias
+	}{
+		Type:  "tool_reference",
+		alias: (*alias)(c),
+	})
+}
+
+// ToolFileChunk represents a tool-generated file content part.
+type ToolFileChunk struct {
+	Tool     string  `json:"tool"`
+	FileID   string  `json:"file_id"`
+	FileName *string `json:"file_name,omitempty"`
+	FileType *string `json:"file_type,omitempty"`
+}
+
+func (*ToolFileChunk) contentChunk() {}
+
+func (c *ToolFileChunk) MarshalJSON() ([]byte, error) {
+	type alias ToolFileChunk
+	return json.Marshal(&struct {
+		Type string `json:"type"`
+		*alias
+	}{
+		Type:  "tool_file",
+		alias: (*alias)(c),
+	})
+}
+
 // UnmarshalContentChunk dispatches to the concrete ContentChunk type
 // based on the "type" discriminator field.
 func UnmarshalContentChunk(data []byte) (ContentChunk, error) {
@@ -223,6 +334,12 @@ func UnmarshalContentChunk(data []byte) (ContentChunk, error) {
 	case "input_audio":
 		var c AudioChunk
 		return &c, json.Unmarshal(data, &c)
+	case "tool_reference":
+		var c ToolReferenceChunk
+		return &c, json.Unmarshal(data, &c)
+	case "tool_file":
+		var c ToolFileChunk
+		return &c, json.Unmarshal(data, &c)
 	default:
 		return &UnknownChunk{Type: probe.Type, Raw: json.RawMessage(data)}, nil
 	}
--- a/chat/content_test.go
+++ b/chat/content_test.go
@@ -150,12 +150,16 @@ func TestFileChunk_RoundTrip(t *testing.T) {
 	}
 }

-func TestReferenceChunk_RoundTrip(t *testing.T) {
-	original := &ReferenceChunk{ReferenceIDs: []int{1, 2, 3}}
+func TestReferenceChunk_RoundTrip_IntIDs(t *testing.T) {
+	original := &ReferenceChunk{ReferenceIDs: []ReferenceID{IntRef(1), IntRef(2), IntRef(3)}}
 	data, err := json.Marshal(original)
 	if err != nil {
 		t.Fatal(err)
 	}
+	want := `{"type":"reference","reference_ids":[1,2,3]}`
+	if string(data) != want {
+		t.Errorf("marshal: got %s, want %s", data, want)
+	}
 	chunk, err := UnmarshalContentChunk(data)
 	if err != nil {
 		t.Fatal(err)
@@ -164,8 +168,64 @@ func TestReferenceChunk_RoundTrip(t *testing.T) {
 	if !ok {
 		t.Fatalf("expected *ReferenceChunk, got %T", chunk)
 	}
-	if len(rc.ReferenceIDs) != 3 || rc.ReferenceIDs[0] != 1 {
-		t.Errorf("got %v, want [1 2 3]", rc.ReferenceIDs)
+	if len(rc.ReferenceIDs) != 3 {
+		t.Fatalf("got %d IDs, want 3", len(rc.ReferenceIDs))
+	}
+	n, ok := rc.ReferenceIDs[0].Int()
+	if !ok || n != 1 {
+		t.Errorf("got %v (ok=%v), want 1", n, ok)
+	}
+}
+
+func TestReferenceChunk_RoundTrip_MixedIDs(t *testing.T) {
+	data := []byte(`{"type":"reference","reference_ids":[1,"abc",42]}`)
+	chunk, err := UnmarshalContentChunk(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+	rc, ok := chunk.(*ReferenceChunk)
+	if !ok {
+		t.Fatalf("expected *ReferenceChunk, got %T", chunk)
+	}
+	if len(rc.ReferenceIDs) != 3 {
+		t.Fatalf("got %d IDs, want 3", len(rc.ReferenceIDs))
+	}
+	// First: int 1
+	if n, ok := rc.ReferenceIDs[0].Int(); !ok || n != 1 {
+		t.Errorf("IDs[0]: got %v (ok=%v), want int 1", n, ok)
+	}
+	// Second: string "abc"
+	if !rc.ReferenceIDs[1].IsString() || rc.ReferenceIDs[1].String() != "abc" {
+		t.Errorf("IDs[1]: got %q (isString=%v), want string abc", rc.ReferenceIDs[1].String(), rc.ReferenceIDs[1].IsString())
+	}
+	// Third: int 42
+	if n, ok := rc.ReferenceIDs[2].Int(); !ok || n != 42 {
+		t.Errorf("IDs[2]: got %v (ok=%v), want int 42", n, ok)
+	}
+	// Round-trip preserves types
+	out, err := json.Marshal(rc)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(out) != `{"type":"reference","reference_ids":[1,"abc",42]}` {
+		t.Errorf("round-trip: got %s", out)
+	}
+}
+
+func TestReferenceID_StringRef(t *testing.T) {
+	id := StringRef("doc-123")
+	if !id.IsString() {
+		t.Error("expected IsString=true")
+	}
+	if id.String() != "doc-123" {
+		t.Errorf("got %q", id.String())
+	}
+	if _, ok := id.Int(); ok {
+		t.Error("Int() should return false for string ref")
+	}
+	data, _ := json.Marshal(id)
+	if string(data) != `"doc-123"` {
+		t.Errorf("marshal: got %s", data)
 	}
 }

@@ -174,7 +234,7 @@ func TestThinkChunk_RoundTrip(t *testing.T) {
 	original := &ThinkChunk{
 		Thinking: []ContentChunk{
 			&TextChunk{Text: "reasoning step"},
-			&ReferenceChunk{ReferenceIDs: []int{42}},
+			&ReferenceChunk{ReferenceIDs: []ReferenceID{IntRef(42)}},
 		},
 		Closed: &closed,
 	}
@@ -365,3 +425,76 @@ func TestContent_IsNull(t *testing.T) {
 		t.Error("chunks content should not be null")
 	}
 }
+
+func TestToolReferenceChunk_RoundTrip(t *testing.T) {
+	url := "https://example.com/result"
+	desc := "A search result"
+	original := &ToolReferenceChunk{
+		Tool:        string(ConnectorWebSearch),
+		Title:       "Example Result",
+		URL:         &url,
+		Description: &desc,
+	}
+	data, err := json.Marshal(original)
+	if err != nil {
+		t.Fatal(err)
+	}
+	chunk, err := UnmarshalContentChunk(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+	tr, ok := chunk.(*ToolReferenceChunk)
+	if !ok {
+		t.Fatalf("expected *ToolReferenceChunk, got %T", chunk)
+	}
+	if tr.Tool != "web_search" {
+		t.Errorf("got tool %q, want web_search", tr.Tool)
+	}
+	if tr.Title != "Example Result" {
+		t.Errorf("got title %q", tr.Title)
+	}
+	if tr.URL == nil || *tr.URL != url {
+		t.Errorf("got url %v, want %q", tr.URL, url)
+	}
+	if tr.Description == nil || *tr.Description != desc {
+		t.Errorf("got description %v, want %q", tr.Description, desc)
+	}
+	if tr.Favicon != nil {
+		t.Errorf("expected nil favicon, got %v", tr.Favicon)
+	}
+}
+
+func TestToolFileChunk_RoundTrip(t *testing.T) {
+	fname := "output.csv"
+	ftype := "text/csv"
+	original := &ToolFileChunk{
+		Tool:     string(ConnectorCodeInterpreter),
+		FileID:   "file-abc123",
+		FileName: &fname,
+		FileType: &ftype,
+	}
+	data, err := json.Marshal(original)
+	if err != nil {
+		t.Fatal(err)
+	}
+	chunk, err := UnmarshalContentChunk(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+	tf, ok := chunk.(*ToolFileChunk)
+	if !ok {
+		t.Fatalf("expected *ToolFileChunk, got %T", chunk)
+	}
+	if tf.Tool != "code_interpreter" {
+		t.Errorf("got tool %q", tf.Tool)
+	}
+	if tf.FileID != "file-abc123" {
+		t.Errorf("got file_id %q", tf.FileID)
+	}
+	if tf.FileName == nil || *tf.FileName != fname {
+		t.Errorf("got file_name %v", tf.FileName)
+	}
+	if tf.FileType == nil || *tf.FileType != ftype {
+		t.Errorf("got file_type %v", tf.FileType)
+	}
+}
--- a/chat/doc.go
+++ b/chat/doc.go
@@ -5,5 +5,22 @@
 // struct literals.
 //
 // Content is polymorphic: it can be a plain string (via [TextContent]),
-// nil, or a slice of [ContentChunk] values (text, image URL, document URL, audio).
+// nil, or a slice of [ContentChunk] values (text, image URL, document URL,
+// audio, tool reference, tool file).
+//
+// # Guardrails
+//
+// [GuardrailConfig] configures moderation guardrails on completion requests.
+// Both v1 ([ModerationLLMV1Config]) and v2 ([ModerationLLMV2Config]) moderation
+// configs are supported.
+//
+// # Content Chunks
+//
+// The following chunk types are supported: [TextChunk], [ImageURLChunk],
+// [DocumentURLChunk], [FileChunk], [ReferenceChunk], [ThinkChunk],
+// [AudioChunk], [ToolReferenceChunk], [ToolFileChunk].
+// Unrecognized types are preserved as [UnknownChunk].
+//
+// [ReferenceChunk] uses [ReferenceID] values that can hold either integer
+// or string identifiers. Use [IntRef] and [StringRef] constructors.
 package chat
--- a/chat/guardrail.go
+++ b/chat/guardrail.go
@@ -0,0 +1,62 @@
+package chat
+
+// ModerationLLMAction specifies the action to take when content exceeds thresholds.
+type ModerationLLMAction string
+
+const (
+	ModerationActionNone  ModerationLLMAction = "none"
+	ModerationActionBlock ModerationLLMAction = "block"
+)
+
+// ModerationLLMV1CategoryThresholds defines per-category score thresholds for v1 moderation.
+type ModerationLLMV1CategoryThresholds struct {
+	Sexual                      *float64 `json:"sexual,omitempty"`
+	HateAndDiscrimination       *float64 `json:"hate_and_discrimination,omitempty"`
+	ViolenceAndThreats          *float64 `json:"violence_and_threats,omitempty"`
+	DangerousAndCriminalContent *float64 `json:"dangerous_and_criminal_content,omitempty"`
+	Selfharm                    *float64 `json:"selfharm,omitempty"`
+	Health                      *float64 `json:"health,omitempty"`
+	Financial                   *float64 `json:"financial,omitempty"`
+	Law                         *float64 `json:"law,omitempty"`
+	PII                         *float64 `json:"pii,omitempty"`
+}
+
+// ModerationLLMV1Config configures the v1 moderation LLM guardrail.
+type ModerationLLMV1Config struct {
+	ModelName                string                              `json:"model_name,omitempty"`
+	CustomCategoryThresholds *ModerationLLMV1CategoryThresholds  `json:"custom_category_thresholds,omitempty"`
+	IgnoreOtherCategories    bool                                `json:"ignore_other_categories,omitempty"`
+	Action                   ModerationLLMAction                 `json:"action,omitempty"`
+}
+
+// ModerationLLMV2CategoryThresholds defines per-category score thresholds for v2 moderation.
+// V2 splits "dangerous_and_criminal_content" into separate "dangerous" and "criminal"
+// categories and adds "jailbreaking".
+type ModerationLLMV2CategoryThresholds struct {
+	Sexual                *float64 `json:"sexual,omitempty"`
+	HateAndDiscrimination *float64 `json:"hate_and_discrimination,omitempty"`
+	ViolenceAndThreats    *float64 `json:"violence_and_threats,omitempty"`
+	Dangerous             *float64 `json:"dangerous,omitempty"`
+	Criminal              *float64 `json:"criminal,omitempty"`
+	Selfharm              *float64 `json:"selfharm,omitempty"`
+	Health                *float64 `json:"health,omitempty"`
+	Financial             *float64 `json:"financial,omitempty"`
+	Law                   *float64 `json:"law,omitempty"`
+	PII                   *float64 `json:"pii,omitempty"`
+	Jailbreaking          *float64 `json:"jailbreaking,omitempty"`
+}
+
+// ModerationLLMV2Config configures the v2 moderation LLM guardrail.
+type ModerationLLMV2Config struct {
+	ModelName                string                             `json:"model_name,omitempty"`
+	CustomCategoryThresholds *ModerationLLMV2CategoryThresholds `json:"custom_category_thresholds,omitempty"`
+	IgnoreOtherCategories    bool                               `json:"ignore_other_categories,omitempty"`
+	Action                   ModerationLLMAction                `json:"action,omitempty"`
+}
+
+// GuardrailConfig configures moderation guardrails for requests.
+type GuardrailConfig struct {
+	BlockOnError    bool                   `json:"block_on_error"`
+	ModerationLLMV1 *ModerationLLMV1Config `json:"moderation_llm_v1,omitempty"`
+	ModerationLLMV2 *ModerationLLMV2Config `json:"moderation_llm_v2,omitempty"`
+}
--- a/chat/request.go
+++ b/chat/request.go
@@ -30,11 +30,12 @@ type CompletionRequest struct {
 	PresencePenalty   *float64        `json:"presence_penalty,omitempty"`
 	FrequencyPenalty  *float64        `json:"frequency_penalty,omitempty"`
 	N                 *int            `json:"n,omitempty"`
-	SafePrompt        bool            `json:"safe_prompt,omitempty"`
-	ParallelToolCalls *bool           `json:"parallel_tool_calls,omitempty"`
-	Metadata          map[string]any  `json:"metadata,omitempty"`
-	Prediction        *Prediction     `json:"prediction,omitempty"`
-	PromptMode        *PromptMode     `json:"prompt_mode,omitempty"`
+	SafePrompt        bool              `json:"safe_prompt,omitempty"`
+	ParallelToolCalls *bool             `json:"parallel_tool_calls,omitempty"`
+	Metadata          map[string]any    `json:"metadata,omitempty"`
+	Prediction        *Prediction       `json:"prediction,omitempty"`
+	PromptMode        *PromptMode       `json:"prompt_mode,omitempty"`
+	Guardrails        []GuardrailConfig `json:"guardrails,omitempty"`
 	stream            bool
 }

--- a/conversation/request.go
+++ b/conversation/request.go
@@ -1,6 +1,10 @@
 package conversation

-import "encoding/json"
+import (
+	"encoding/json"
+
+	"somegit.dev/vikingowl/mistral-go-sdk/chat"
+)

 // StartRequest starts a new conversation.
 type StartRequest struct {
@@ -11,9 +15,10 @@ type StartRequest struct {
 	Instructions     *string          `json:"instructions,omitempty"`
 	Tools            []Tool           `json:"tools,omitempty"`
 	CompletionArgs   *CompletionArgs  `json:"completion_args,omitempty"`
-	Store            *bool            `json:"store,omitempty"`
-	HandoffExecution *HandoffExecution `json:"handoff_execution,omitempty"`
-	Name             *string          `json:"name,omitempty"`
+	Store            *bool                  `json:"store,omitempty"`
+	HandoffExecution *HandoffExecution      `json:"handoff_execution,omitempty"`
+	Guardrails       []chat.GuardrailConfig `json:"guardrails,omitempty"`
+	Name             *string                `json:"name,omitempty"`
 	Description      *string          `json:"description,omitempty"`
 	Metadata         map[string]any   `json:"metadata,omitempty"`
 	stream           bool
@@ -61,13 +66,14 @@ func (r *AppendRequest) MarshalJSON() ([]byte, error) {

 // RestartRequest restarts a conversation from a specific entry.
 type RestartRequest struct {
-	Inputs           Inputs           `json:"inputs"`
-	FromEntryID      string           `json:"from_entry_id"`
-	CompletionArgs   *CompletionArgs  `json:"completion_args,omitempty"`
-	Store            *bool            `json:"store,omitempty"`
-	HandoffExecution *HandoffExecution `json:"handoff_execution,omitempty"`
-	AgentVersion     json.RawMessage  `json:"agent_version,omitempty"`
-	Metadata         map[string]any   `json:"metadata,omitempty"`
+	Inputs           Inputs                 `json:"inputs"`
+	FromEntryID      string                 `json:"from_entry_id"`
+	CompletionArgs   *CompletionArgs        `json:"completion_args,omitempty"`
+	Store            *bool                  `json:"store,omitempty"`
+	HandoffExecution *HandoffExecution      `json:"handoff_execution,omitempty"`
+	Guardrails       []chat.GuardrailConfig `json:"guardrails,omitempty"`
+	AgentVersion     json.RawMessage        `json:"agent_version,omitempty"`
+	Metadata         map[string]any         `json:"metadata,omitempty"`
 	stream           bool
 }

--- a/doc.go
+++ b/doc.go
@@ -42,4 +42,10 @@
 // [conversation], [embedding], [model], [file], [finetune], [batch],
 // [ocr], [audio], [library], [moderation], [classification], and [fim].
 // All service methods live directly on [Client].
+//
+// # Reference
+//
+// This SDK tracks the official Mistral Python SDK
+// (https://github.com/mistralai/client-python) as its upstream reference
+// for API surface and type definitions.
 package mistral
--- a/docs/openapi.yaml
+++ b/docs/openapi.yaml
--- a/mistral.go
+++ b/mistral.go
@@ -6,7 +6,7 @@ import (
 )

 // Version is the SDK version string.
-const Version = "0.1.0"
+const Version = "0.2.0"

 const (
 	defaultBaseURL = "https://api.mistral.ai"
--- a/model/model.go
+++ b/model/model.go
@@ -34,8 +34,10 @@ type ModelCapabilities struct {
 	OCR                bool `json:"ocr"`
 	Classification     bool `json:"classification"`
 	Moderation         bool `json:"moderation"`
-	Audio              bool `json:"audio"`
-	AudioTranscription bool `json:"audio_transcription"`
+	Audio                      bool `json:"audio"`
+	AudioTranscription         bool `json:"audio_transcription"`
+	AudioTranscriptionRealtime bool `json:"audio_transcription_realtime"`
+	AudioSpeech                bool `json:"audio_speech"`
 }

 // ModelList is the response from listing models.