From 436010560ef6630f4908d20df441b23a6c94aa9f Mon Sep 17 00:00:00 2001
From: megatron0000 <vitor_pimenta97@hotmail.com>
Date: Sat, 3 May 2025 02:15:03 +0000
Subject: [PATCH] feat(transcribe): add --speech_model flag

---
 cmd/transcribe.go | 4 ++++
 main_test.go      | 4 ++++
 schemas/types.go  | 2 ++
 3 files changed, 10 insertions(+)

diff --git a/cmd/transcribe.go b/cmd/transcribe.go
index 6a16c85..b5ca8ae 100644
--- a/cmd/transcribe.go
+++ b/cmd/transcribe.go
@@ -129,6 +129,9 @@ var transcribeCmd = &cobra.Command{
 				}
 				params.SummaryModel = summaryModel
 			}
+
+			speechModel, _ := cmd.Flags().GetString("speech_model")
+			params.SpeechModel = &speechModel
 		}
 
 		if params.RedactPii {
@@ -278,6 +281,7 @@ func init() {
 	transcribeCmd.PersistentFlags().StringP("webhook_url", "w", "", "Receive a webhook once your transcript is complete.")
 	transcribeCmd.PersistentFlags().StringP("word_boost", "k", "", "The value of this flag MUST be used surrounded by quotes. Any term included will have its likelihood of being transcribed boosted.")
 	transcribeCmd.PersistentFlags().StringP("summary_model", "q", "informative", "The model used to generate the summary.")
+	transcribeCmd.PersistentFlags().StringP("speech_model", "", "", "The speech model to use for the transcription. Allowed values: see https://www.assemblyai.com/docs/api-reference/transcripts/submit#request.body.speech_model")
 
 	transcribeCmd.Flags().Bool("test", false, "Flag for test executing purpose")
 	transcribeCmd.Flags().MarkHidden("test")
diff --git a/main_test.go b/main_test.go
index 52acb1d..4d576bb 100644
--- a/main_test.go
+++ b/main_test.go
@@ -103,6 +103,7 @@ func TestTranscribeWithFlags(t *testing.T) {
 		"--speaker_labels",
 		"--summarization",
 		"--topic_detection",
+		"--speech_model=nano",
 		"-p=false",
 		"-j",
 		"--test",
@@ -147,6 +148,9 @@ func TestTranscribeWithFlags(t *testing.T) {
 	if *result.IabCategories != true {
 		t.Errorf("Expected IAB Categories(Topic detection) true, got false.")
 	}
+	if result.SpeechModel == nil || *result.SpeechModel != "nano" {
+		t.Errorf("Expected speech_model to be 'nano', got '%v'.", result.SpeechModel)
+	}
 }
 
 func TestTranscribeRestrictions(t *testing.T) {
diff --git a/schemas/types.go b/schemas/types.go
index f285655..95dbdcf 100644
--- a/schemas/types.go
+++ b/schemas/types.go
@@ -62,6 +62,7 @@ type TranscriptResponse struct {
 	SentimentAnalysis        *bool                      `json:"sentiment_analysis,omitempty"`
 	SentimentAnalysisResults *[]SentimentAnalysisResult `json:"sentiment_analysis_results,omitempty"`
 	SpeakerLabels            bool                       `json:"speaker_labels,omitempty"`
+	SpeechModel              *string                    `json:"speech_model,omitempty"`
 	SpeedBoost               *bool                      `json:"speed_boost,omitempty"`
 	Status                   *string                    `json:"status,omitempty"`
 	Summarization            *bool                      `json:"summarization,omitempty"`
@@ -208,6 +209,7 @@ type TranscribeParams struct {
 	WebhookAuthHeaderValue string           `json:"webhook_auth_header_value,omitempty"`
 	WebhookURL             string           `json:"webhook_url,omitempty"`
 	WordBoost              []string         `json:"word_boost,omitempty"`
+	SpeechModel            *string          `json:"speech_model,omitempty"`
 }
 
 type CustomSpelling struct {