diff --git a/llms/count_tokens.go b/llms/count_tokens.go
index 45bb52d5..844a3a1a 100644
--- a/llms/count_tokens.go
+++ b/llms/count_tokens.go
@@ -11,25 +11,47 @@ const (
 )
 
 const (
-	_gpt35TurboContextSize   = 4096
-	_gpt432KContextSize      = 32768
-	_gpt4ContextSize         = 8192
-	_textDavinci3ContextSize = 4097
-	_textBabbage1ContextSize = 2048
-	_textAda1ContextSize     = 2048
-	_textCurie1ContextSize   = 2048
-	_codeDavinci2ContextSize = 8000
-	_codeCushman1ContextSize = 2048
-	_textBisonContextSize    = 2048
-	_chatBisonContextSize    = 2048
-	_defaultContextSize      = 2048
+	_gpt35TurboContextSize    = 16385  // gpt-3.5-turbo default context
+	_gpt35Turbo16KContextSize = 16385  // gpt-3.5-turbo-16k
+	_gpt4ContextSize          = 8192   // gpt-4
+	_gpt432KContextSize       = 32768  // gpt-4-32k
+	_gpt4TurboContextSize     = 128000 // gpt-4-turbo models
+	_gpt4oContextSize         = 128000 // gpt-4o models
+	_gpt4oMiniContextSize     = 128000 // gpt-4o-mini
+	_textDavinci3ContextSize  = 4097
+	_textBabbage1ContextSize  = 2048
+	_textAda1ContextSize      = 2048
+	_textCurie1ContextSize    = 2048
+	_codeDavinci2ContextSize  = 8000
+	_codeCushman1ContextSize  = 2048
+	_defaultContextSize       = 2048
 )
 
 // nolint:gochecknoglobals
 var modelToContextSize = map[string]int{
-	"gpt-3.5-turbo":    _gpt35TurboContextSize,
-	"gpt-4-32k":        _gpt432KContextSize,
-	"gpt-4":            _gpt4ContextSize,
+	// GPT-3.5 models
+	"gpt-3.5-turbo":      _gpt35TurboContextSize,
+	"gpt-3.5-turbo-16k":  _gpt35Turbo16KContextSize,
+	"gpt-3.5-turbo-0125": _gpt35TurboContextSize,
+	"gpt-3.5-turbo-1106": _gpt35TurboContextSize,
+	// GPT-4 models
+	"gpt-4":          _gpt4ContextSize,
+	"gpt-4-32k":      _gpt432KContextSize,
+	"gpt-4-0613":     _gpt4ContextSize,
+	"gpt-4-32k-0613": _gpt432KContextSize,
+	// GPT-4 Turbo models
+	"gpt-4-turbo":            _gpt4TurboContextSize,
+	"gpt-4-turbo-preview":    _gpt4TurboContextSize,
+	"gpt-4-turbo-2024-04-09": _gpt4TurboContextSize,
+	"gpt-4-1106-preview":     _gpt4TurboContextSize,
+	"gpt-4-0125-preview":     _gpt4TurboContextSize,
+	// GPT-4o models
+	"gpt-4o":                 _gpt4oContextSize,
+	"gpt-4o-2024-05-13":      _gpt4oContextSize,
+	"gpt-4o-2024-08-06":      _gpt4oContextSize,
+	"gpt-4o-mini":            _gpt4oMiniContextSize,
+	"gpt-4o-mini-2024-07-18": _gpt4oMiniContextSize,
+	// Legacy models
 	"text-davinci-003": _textDavinci3ContextSize,
 	"text-curie-001":   _textCurie1ContextSize,
 	"text-babbage-001": _textBabbage1ContextSize,
diff --git a/llms/count_tokens_test.go b/llms/count_tokens_test.go
index b7089340..7faa5abe 100644
--- a/llms/count_tokens_test.go
+++ b/llms/count_tokens_test.go
@@ -12,3 +12,50 @@ func TestCountTokens(t *testing.T) {
 	expectedNumTokens := 4
 	assert.Equal(t, expectedNumTokens, numTokens)
 }
+
+func TestGetModelContextSize(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		model        string
+		expectedSize int
+	}{
+		// GPT-3.5 models
+		{"gpt-3.5-turbo", 16385},
+		{"gpt-3.5-turbo-16k", 16385},
+		{"gpt-3.5-turbo-0125", 16385},
+		{"gpt-3.5-turbo-1106", 16385},
+		// GPT-4 models
+		{"gpt-4", 8192},
+		{"gpt-4-32k", 32768},
+		{"gpt-4-0613", 8192},
+		{"gpt-4-32k-0613", 32768},
+		// GPT-4 Turbo models
+		{"gpt-4-turbo", 128000},
+		{"gpt-4-turbo-preview", 128000},
+		{"gpt-4-turbo-2024-04-09", 128000},
+		{"gpt-4-1106-preview", 128000},
+		{"gpt-4-0125-preview", 128000},
+		// GPT-4o models
+		{"gpt-4o", 128000},
+		{"gpt-4o-2024-05-13", 128000},
+		{"gpt-4o-2024-08-06", 128000},
+		{"gpt-4o-mini", 128000},
+		{"gpt-4o-mini-2024-07-18", 128000},
+		// Legacy models
+		{"text-davinci-003", 4097},
+		{"text-curie-001", 2048},
+		{"text-babbage-001", 2048},
+		{"text-ada-001", 2048},
+		{"code-davinci-002", 8000},
+		{"code-cushman-001", 2048},
+		// Unknown model should return default
+		{"unknown-model", 2048},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.model, func(t *testing.T) {
+			size := GetModelContextSize(tt.model)
+			assert.Equal(t, tt.expectedSize, size, "Context size for model %s", tt.model)
+		})
+	}
+}