Skip to content

Refactor/ai model registry centralized #324

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
urjitc merged 3 commits into from
Apr 9, 2026
Merged

Refactor/ai model registry centralized #324

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e132aba
refactor(ai): centralize model definitions into single registry
urjitc Apr 9, 2026
4296980
Remove unused text selection worker
urjitc Apr 9, 2026
ff9a243
Centralize AI Gateway options; migrate autogen to gateway
urjitc Apr 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 18 additions & 56 deletions src/app/api/chat/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import { gateway } from "ai";
import {
streamText,
smoothStream,
Expand All @@ -22,6 +21,12 @@ import {
import { withServerObservability } from "@/lib/with-server-observability";
import { normalizeLegacyToolMessages } from "@/lib/ai/legacy-tool-message-compat";
import type { ReplySelection } from "@/lib/stores/ui-store";
import { getDefaultChatModelId, resolveGatewayModelId } from "@/lib/ai/models";
import {
buildGatewayProviderOptions,
createGatewayLanguageModel,
getGatewayAttributionHeaders,
} from "@/lib/ai/gateway-provider-options";

/**
* Extract workspaceId from system context or request body
Expand Down Expand Up @@ -158,7 +163,9 @@ async function handlePOST(req: Request) {
// Convert messages (pass tools so toModelOutput strips event from historical tool results)
let convertedMessages;
try {
convertedMessages = await convertToModelMessages(validatedMessages, { tools });
convertedMessages = await convertToModelMessages(validatedMessages, {
tools,
});
} catch (convertError) {
logger.error("❌ [CHAT-API] convertToModelMessages FAILED:", {
error:
Expand All @@ -181,19 +188,9 @@ async function handlePOST(req: Request) {
// Get pre-formatted selected cards context from client (no DB fetch needed)
const selectedCardsContext = getSelectedCardsContext(body);

// Get model ID and ensure it has the correct prefix for Gateway
let modelId = body.modelId || "gemini-3-flash-preview";

// Auto-prefix with google/ if it looks like a gemini model and lacks prefix
// This allows existing client code to work without changes
if (modelId.startsWith("gemini-") && !modelId.startsWith("google/")) {
modelId = `google/${modelId}`;
}

// Auto-prefix with anthropic/ if it looks like a Claude model and lacks prefix
if (modelId.includes("claude") && !modelId.startsWith("anthropic/")) {
modelId = `anthropic/${modelId}`;
}
const modelId = resolveGatewayModelId(
body.modelId || getDefaultChatModelId(),
);

// Inject selected cards + reply selections into the last user message
injectSelectionContext(
Expand All @@ -203,16 +200,17 @@ async function handlePOST(req: Request) {
);

const posthogClient = getPostHogServerClient();
const baseGatewayModel = createGatewayLanguageModel(modelId);
const tracedModel = posthogClient
? withTracing(gateway(modelId) as any, posthogClient, {
? withTracing(baseGatewayModel as any, posthogClient, {
posthogDistinctId: userId || "anonymous",
posthogProperties: {
workspaceId,
activeFolderId,
modelId,
},
})
: (gateway(modelId) as any);
: (baseGatewayModel as any);

// Use AI Gateway
const model = wrapLanguageModel({
Expand All @@ -227,53 +225,17 @@ async function handlePOST(req: Request) {
modelId,
});

// Configure Google Thinking capabilities
const googleConfig: any = {
grounding: {
// googleSearchRetrieval removed to force usage of explicit web_search tool
},
thinkingConfig: {
includeThoughts: true,
},
};

// Gemini 3 Flash: set thinkingLevel (Gemini 2.5 uses default dynamic budget)
if (modelId.includes("gemini-3-flash")) {
googleConfig.thinkingConfig.thinkingLevel = "minimal";
}
const providerOptions = buildGatewayProviderOptions(modelId, { userId });

// Prepare provider options.
// Prefer Bedrock first for Claude models, then fall back to Anthropic.
// Non-Claude models stay on their native providers.
const gatewayOptions: any = {
caching: "auto",
models: [modelId],
...(userId ? { user: userId } : {}),
};

if (modelId.startsWith("anthropic/")) {
gatewayOptions.order = ["bedrock", "anthropic"];
gatewayOptions.only = ["bedrock", "anthropic"];
}

const providerOptions: any = {
gateway: gatewayOptions,
google: googleConfig,
};

const appUrl = process.env.NEXT_PUBLIC_APP_URL || "https://thinkex.app";
const result = streamText({
model: model,
temperature: 1.0,
system,
messages: convertedMessages,
stopWhen: stepCountIs(25),
tools,
providerOptions,
headers: {
"http-referer": appUrl,
"x-title": "ThinkEx",
},
providerOptions: providerOptions as any,
headers: getGatewayAttributionHeaders(),
experimental_telemetry: {
isEnabled: true,
metadata: {
Expand Down
146 changes: 74 additions & 72 deletions src/app/api/threads/[id]/title/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,7 @@ import {
} from "@/lib/api/workspace-helpers";
import { eq } from "drizzle-orm";
import { withServerObservability } from "@/lib/with-server-observability";

/** Model ID used for lightweight background text tasks */
const GEMINI_FLASH_LITE_MODEL = "gemini-2.5-flash-lite";
import { getModelForPurpose } from "@/lib/ai/models";

function extractTextFromMessage(msg: { content?: unknown[] }): string {
if (!msg.content || !Array.isArray(msg.content)) return "";
Expand All @@ -28,84 +26,88 @@ function extractTextFromMessage(msg: { content?: unknown[] }): string {
* Generate a title from messages using Gemini Flash Lite.
* Body: { messages: ThreadMessage[] }
*/
export const POST = withServerObservability(async function POST(
req: NextRequest,
{ params }: { params: Promise<{ id: string }> }
) {
try {
const userId = await requireAuth();
const { id } = await params;
const body = await req.json().catch(() => ({}));
const { messages } = body;
export const POST = withServerObservability(
async function POST(
req: NextRequest,
{ params }: { params: Promise<{ id: string }> },
) {
try {
const userId = await requireAuth();
const { id } = await params;
const body = await req.json().catch(() => ({}));
const { messages } = body;

const [thread] = await db
.select()
.from(chatThreads)
.where(eq(chatThreads.id, id))
.limit(1);
const [thread] = await db
.select()
.from(chatThreads)
.where(eq(chatThreads.id, id))
.limit(1);

if (!thread) {
return NextResponse.json({ error: "Thread not found" }, { status: 404 });
}
if (!thread) {
return NextResponse.json(
{ error: "Thread not found" },
{ status: 404 },
);
}

await verifyWorkspaceAccess(thread.workspaceId, userId);
verifyThreadOwnership(thread, userId);
await verifyWorkspaceAccess(thread.workspaceId, userId);
verifyThreadOwnership(thread, userId);

let title = "New Chat";
let title = "New Chat";

if (messages && Array.isArray(messages) && messages.length > 0) {
const conversationText = messages
.slice(0, 6)
.map((m: { role?: string; content?: unknown[] }) => {
const text = extractTextFromMessage(m);
if (!text) return "";
const role = m.role === "user" ? "User" : "Assistant";
return `${role}: ${text}`;
})
.filter(Boolean)
.join("\n\n");
if (messages && Array.isArray(messages) && messages.length > 0) {
const conversationText = messages
.slice(0, 6)
.map((m: { role?: string; content?: unknown[] }) => {
const text = extractTextFromMessage(m);
if (!text) return "";
const role = m.role === "user" ? "User" : "Assistant";
return `${role}: ${text}`;
})
.filter(Boolean)
.join("\n\n");

if (conversationText.trim()) {
try {
const { text } = await generateText({
model: google(GEMINI_FLASH_LITE_MODEL),
system: `Generate a very short chat title (2-6 words) that captures the topic. Output ONLY the title, no quotes or punctuation.`,
prompt: `Conversation:\n\n${conversationText}\n\nTitle:`,
experimental_telemetry: {
isEnabled: true,
metadata: {
"tcc.sessionId": id,
...(userId ? { userId } : {}),
if (conversationText.trim()) {
try {
const { text } = await generateText({
model: google(getModelForPurpose("title-generation")),

@cubic-dev-ai cubic-dev-ai Bot Apr 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: This still hardcodes the Google provider, so title-generation cannot safely move to a non-Google model in the registry.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At src/app/api/threads/[id]/title/route.ts, line 73:

<comment>This still hardcodes the Google provider, so `title-generation` cannot safely move to a non-Google model in the registry.</comment>

<file context>
@@ -28,84 +26,88 @@ function extractTextFromMessage(msg: { content?: unknown[] }): string {
+        if (conversationText.trim()) {
+          try {
+            const { text } = await generateText({
+              model: google(getModelForPurpose("title-generation")),
+              system: `Generate a very short chat title (2-6 words) that captures the topic. Output ONLY the title, no quotes or punctuation.`,
+              prompt: `Conversation:\n\n${conversationText}\n\nTitle:`,
</file context>
Fix with Cubic

system: `Generate a very short chat title (2-6 words) that captures the topic. Output ONLY the title, no quotes or punctuation.`,
prompt: `Conversation:\n\n${conversationText}\n\nTitle:`,
experimental_telemetry: {
isEnabled: true,
metadata: {
"tcc.sessionId": id,
...(userId ? { userId } : {}),
},
},
},
});
const generated = text.trim().slice(0, 60);
if (generated) title = generated;
} catch (err) {
console.warn("[threads] title Gemini fallback:", err);
const firstUser = messages.find(
(m: { role?: string }) => m.role === "user"
);
const fallback = extractTextFromMessage(firstUser ?? {});
if (fallback) {
title = fallback.slice(0, 50) + (fallback.length > 50 ? "..." : "");
});
const generated = text.trim().slice(0, 60);
if (generated) title = generated;
} catch (err) {
console.warn("[threads] title Gemini fallback:", err);
const firstUser = messages.find(
(m: { role?: string }) => m.role === "user",
);
const fallback = extractTextFromMessage(firstUser ?? {});
if (fallback) {
title =
fallback.slice(0, 50) + (fallback.length > 50 ? "..." : "");
}
}
}
}
}

await db
.update(chatThreads)
.set({ title })
.where(eq(chatThreads.id, id));
await db.update(chatThreads).set({ title }).where(eq(chatThreads.id, id));

return NextResponse.json({ title });
} catch (error) {
if (error instanceof Response) return error;
console.error("[threads] title error:", error);
return NextResponse.json(
{ error: "Internal server error" },
{ status: 500 }
);
}
}, { routeName: "POST /api/threads/[id]/title" });
return NextResponse.json({ title });
} catch (error) {
if (error instanceof Response) return error;
console.error("[threads] title error:", error);
return NextResponse.json(
{ error: "Internal server error" },
{ status: 500 },
);
}
},
{ routeName: "POST /api/threads/[id]/title" },
);
Loading
Loading