From 23ebb2b09fe63ebe2961bd9cfb28e14dfd7703d7 Mon Sep 17 00:00:00 2001
From: Rene <rene@done.ph>
Date: Sat, 23 May 2026 23:06:07 +0800
Subject: [PATCH] support deepseek 1M

---
 cmd/ocgo/main.go      | 31 ++++++++++++++++++++++++++-----
 cmd/ocgo/main_test.go | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/cmd/ocgo/main.go b/cmd/ocgo/main.go
index 53aa358..8aabb40 100644
--- a/cmd/ocgo/main.go
+++ b/cmd/ocgo/main.go
@@ -184,7 +184,22 @@ func listCmd() *cobra.Command {
 }
 
 func knownModelIDs() []string {
-	return []string{"glm-5.1", "glm-5", "kimi-k2.6", "kimi-k2.5", "mimo-v2.5-pro", "mimo-v2.5", "mimo-v2-pro", "mimo-v2-omni", "minimax-m2.7", "minimax-m2.5", "deepseek-v4-pro", "deepseek-v4-flash", "qwen3.6-plus", "qwen3.5-plus"}
+	return []string{"glm-5.1", "glm-5", "kimi-k2.6", "kimi-k2.5", "mimo-v2.5-pro", "mimo-v2.5", "mimo-v2-pro", "mimo-v2-omni", "minimax-m2.7", "minimax-m2.5", "deepseek-v4-pro[1M]", "deepseek-v4-flash[1M]", "qwen3.6-plus", "qwen3.5-plus"}
+}
+
+func cleanModelName(model string) string {
+	model = strings.TrimSuffix(model, "[1M]")
+	model = strings.TrimSuffix(model, "[128K]")
+	return model
+}
+
+func modelContextWindow(model string) int {
+	switch cleanModelName(model) {
+	case "deepseek-v4-pro", "deepseek-v4-flash":
+		return 1000000
+	default:
+		return 128000
+	}
 }
 
 func modelSupportsImages(model string) bool {
@@ -489,6 +504,11 @@ func prepareChatBody(body []byte) ([]byte, error) {
 	}
 	changed := requestStreamingUsage(req)
 	model, _ := req["model"].(string)
+	if clean := cleanModelName(model); clean != model {
+		req["model"] = clean
+		changed = true
+		model = clean
+	}
 	if rawChatBodyHasImages(req) {
 		if !modelSupportsImages(model) {
 			return nil, unsupportedImageModelError(model)
@@ -582,7 +602,7 @@ func stripRawChatImageDetails(req map[string]any) bool {
 }
 
 func convertRequest(ar AnthropicRequest) OAIRequest {
-	model := ar.Model
+	model := cleanModelName(ar.Model)
 	if model == "" || strings.HasPrefix(model, "claude-") {
 		model = "kimi-k2.6"
 	}
@@ -600,7 +620,7 @@ func convertRequest(ar AnthropicRequest) OAIRequest {
 }
 
 func responsesToChat(rr ResponsesRequest) OAIRequest {
-	model := rr.Model
+	model := cleanModelName(rr.Model)
 	if model == "" {
 		model = "kimi-k2.6"
 	}
@@ -1522,6 +1542,7 @@ func writeCodexProfile(path, baseURL string) error {
 func writeCodexModelCatalog(path string) error {
 	models := make([]map[string]any, 0, len(knownModelIDs()))
 	for i, id := range knownModelIDs() {
+		cw := modelContextWindow(id)
 		models = append(models, map[string]any{
 			"slug":                             id,
 			"display_name":                     id,
@@ -1544,8 +1565,8 @@ func writeCodexModelCatalog(path string) error {
 			"truncation_policy":                map[string]any{"mode": "tokens", "limit": 10000},
 			"supports_parallel_tool_calls":     false,
 			"supports_image_detail_original":   false,
-			"context_window":                   128000,
-			"max_context_window":               128000,
+			"context_window":                   cw,
+			"max_context_window":               cw,
 			"auto_compact_token_limit":         nil,
 			"effective_context_window_percent": 95,
 			"experimental_supported_tools":     []any{},
diff --git a/cmd/ocgo/main_test.go b/cmd/ocgo/main_test.go
index e541806..f802032 100644
--- a/cmd/ocgo/main_test.go
+++ b/cmd/ocgo/main_test.go
@@ -60,6 +60,36 @@ func TestWriteCodexProfileReplacesExistingSections(t *testing.T) {
 	}
 }
 
+func TestModelContextWindow(t *testing.T) {
+	if got := modelContextWindow("deepseek-v4-flash[1M]"); got != 1000000 {
+		t.Fatalf("deepseek-v4-flash[1M]: got %d, want 1000000", got)
+	}
+	if got := modelContextWindow("deepseek-v4-pro[1M]"); got != 1000000 {
+		t.Fatalf("deepseek-v4-pro[1M]: got %d, want 1000000", got)
+	}
+	if got := modelContextWindow("kimi-k2.6"); got != 128000 {
+		t.Fatalf("kimi-k2.6: got %d, want 128000", got)
+	}
+	if got := modelContextWindow("unknown-model"); got != 128000 {
+		t.Fatalf("unknown-model: got %d, want 128000", got)
+	}
+}
+
+func TestCleanModelName(t *testing.T) {
+	tests := []struct{ in, want string }{
+		{"deepseek-v4-flash[1M]", "deepseek-v4-flash"},
+		{"deepseek-v4-pro[1M]", "deepseek-v4-pro"},
+		{"deepseek-v4-flash", "deepseek-v4-flash"},
+		{"kimi-k2.6", "kimi-k2.6"},
+		{"kimi-k2.6[128K]", "kimi-k2.6"},
+	}
+	for _, tt := range tests {
+		if got := cleanModelName(tt.in); got != tt.want {
+			t.Errorf("cleanModelName(%q) = %q, want %q", tt.in, got, tt.want)
+		}
+	}
+}
+
 func TestWriteCodexModelCatalog(t *testing.T) {
 	path := filepath.Join(t.TempDir(), "ocgo-models.json")
 	if err := writeCodexModelCatalog(path); err != nil {
@@ -70,7 +100,7 @@ func TestWriteCodexModelCatalog(t *testing.T) {
 		t.Fatal(err)
 	}
 	content := string(b)
-	for _, want := range []string{`"models"`, `"slug": "deepseek-v4-pro"`, `"context_window": 128000`, `"truncation_policy"`, `"supports_image_detail_original": false`, `"image"`} {
+	for _, want := range []string{`"models"`, `"slug": "deepseek-v4-pro[1M]"`, `"context_window": 1000000`, `"truncation_policy"`, `"supports_image_detail_original": false`, `"image"`} {
 		if !strings.Contains(content, want) {
 			t.Fatalf("missing %q in:\n%s", want, content)
 		}