From 77be8d59a5836c612ddc955ae2299c3dbe79ced6 Mon Sep 17 00:00:00 2001 From: Aditya Gupta Date: Sat, 11 Jul 2026 18:29:34 +0530 Subject: [PATCH] feat(llm): Added max context length configuration and statistics --- apps/gui/next-env.d.ts | 2 +- apps/gui/src/app/config/page.tsx | 28 +++- apps/gui/src/app/play/actions.ts | 6 +- apps/gui/src/components/play/PlayView.tsx | 164 ++++++++++++++------- apps/gui/src/lib/simulation-types.ts | 2 + apps/gui/src/lib/simulation.ts | 9 +- packages/llm/src/llm.ts | 3 + packages/llm/src/provider-manager.ts | 78 ++++++---- packages/llm/src/providers/google-genai.ts | 8 +- packages/llm/src/providers/openrouter.ts | 8 +- packages/llm/tests/openrouter.test.ts | 2 + 11 files changed, 218 insertions(+), 92 deletions(-) diff --git a/apps/gui/next-env.d.ts b/apps/gui/next-env.d.ts index c4b7818..9edff1c 100644 --- a/apps/gui/next-env.d.ts +++ b/apps/gui/next-env.d.ts @@ -1,6 +1,6 @@ /// /// -import "./.next/dev/types/routes.d.ts"; +import "./.next/types/routes.d.ts"; // NOTE: This file should not be edited // see https://nextjs.org/docs/app/api-reference/config/typescript for more information. diff --git a/apps/gui/src/app/config/page.tsx b/apps/gui/src/app/config/page.tsx index a8c8afb..2a3ebea 100644 --- a/apps/gui/src/app/config/page.tsx +++ b/apps/gui/src/app/config/page.tsx @@ -37,6 +37,7 @@ export default function ConfigPage() { const [editModel, setEditModel] = useState("gemini-2.5-flash"); const [editIsActive, setEditIsActive] = useState(false); const [editType, setEditType] = useState<"generative" | "embedding">("generative"); + const [editMaxContext, setEditMaxContext] = useState(32768); useEffect(() => { if (selectedInstanceId === null) { @@ -46,6 +47,7 @@ export default function ConfigPage() { setEditModel("gemini-2.5-flash"); setEditIsActive(false); setEditType("generative"); + setEditMaxContext(32768); } else if (selectedInstanceId === "new") { setEditName(""); const defaultProvider = "google-genai"; @@ -55,6 +57,7 @@ export default function ConfigPage() { const pMeta = availableProviders.find((p) => p.id === defaultProvider); setEditModel(pMeta?.defaultModel || "gemini-2.5-flash"); setEditIsActive(false); + setEditMaxContext(32768); } else { const inst = instances.find((i) => i.id === selectedInstanceId); if (inst) { @@ -65,6 +68,7 @@ export default function ConfigPage() { const pMeta = availableProviders.find((p) => p.id === inst.providerName); setEditModel(inst.modelName || (inst.type === "embedding" ? pMeta?.defaultEmbeddingModel : pMeta?.defaultModel) || "gemini-2.5-flash"); setEditIsActive(inst.isActive); + setEditMaxContext(inst.maxContext !== undefined && inst.maxContext !== null ? inst.maxContext : 32768); } } }, [selectedInstanceId, instances, availableProviders]); @@ -132,13 +136,14 @@ export default function ConfigPage() { setLoading(false); return; } - const created = await createProviderInstance(editName, editProvider, editKey, editModel || undefined, editType); + const created = await createProviderInstance(editName, editProvider, editKey, editModel || undefined, editType, editType === "generative" ? editMaxContext : 0); if (editIsActive) { await setActiveProviderInstance(created.id); } targetInstanceId = created.id; setSelectedInstanceId(created.id); } else { + if (!selectedInstanceId) return; const inst = instances.find((i) => i.id === selectedInstanceId); if (inst && inst.type === "embedding") { const isMapped = mappings["embeddings"] === selectedInstanceId; @@ -158,7 +163,7 @@ export default function ConfigPage() { } } - await updateProviderInstance(selectedInstanceId, editName, editProvider, editKey || undefined, editModel || undefined, editType); + await updateProviderInstance(selectedInstanceId, editName, editProvider, editKey || undefined, editModel || undefined, editType, editType === "generative" ? editMaxContext : 0); if (editIsActive) { await setActiveProviderInstance(selectedInstanceId); } @@ -167,7 +172,7 @@ export default function ConfigPage() { await loadInstances(); await loadMappings(); - if (shouldRegenerate && targetInstanceId !== "new") { + if (shouldRegenerate && targetInstanceId && targetInstanceId !== "new") { await regenerateEmbeddings(targetInstanceId); } } catch (err) { @@ -381,6 +386,23 @@ export default function ConfigPage() { /> + {editType === "generative" && ( +
+ + setEditMaxContext(parseInt(e.target.value) || 0)} + min={0} + placeholder="e.g. 32768" + className="w-full rounded-md border border-gray-300 bg-white px-3 py-2 text-sm outline-none transition-[border-color,box-shadow] focus:border-blue-500 focus:ring-3 focus:ring-blue-500/15" + /> +
+ )} +
{ - return ProviderManager.create(name, providerName, apiKey, modelName, type); + return ProviderManager.create(name, providerName, apiKey, modelName, type, maxContext); } export async function deleteProviderInstance(id: string): Promise { @@ -262,8 +263,9 @@ export async function updateProviderInstance( apiKey?: string, modelName?: string, type: "generative" | "embedding" = "generative", + maxContext?: number, ): Promise { - ProviderManager.update(id, name, providerName, apiKey, modelName, type); + ProviderManager.update(id, name, providerName, apiKey, modelName, type, maxContext); } export async function getProviderMappings(): Promise> { diff --git a/apps/gui/src/components/play/PlayView.tsx b/apps/gui/src/components/play/PlayView.tsx index 8bbb5fe..d7a8ee5 100644 --- a/apps/gui/src/components/play/PlayView.tsx +++ b/apps/gui/src/components/play/PlayView.tsx @@ -13,7 +13,7 @@ import { listProviderInstances, } from "@/app/play/actions"; import type { SimSnapshot } from "@/lib/simulation-types"; -import type { LLMProviderInstance } from "@omnia/llm"; +import type { ModelProviderInstance } from "@omnia/llm"; function IntentTag({ intent, @@ -90,9 +90,9 @@ function PromptModal({ const memTokens = Math.max(0, inputTokens - sysTokens - worldTokens); return [ - { label: "System Prompt", pct: sysPct, tokens: sysTokens, type: "system", content: systemPrompt }, - { label: "World Info", pct: worldPct, tokens: worldTokens, type: "world", content: worldStr }, - { label: "Recent Memories", pct: memPct, tokens: memTokens, type: "memories", content: memStr || "(No memories yet.)" }, + { label: "System Prompt", pct: sysPct, relativePct: sysPct, tokens: sysTokens, type: "system", content: systemPrompt }, + { label: "World Info", pct: worldPct, relativePct: worldPct, tokens: worldTokens, type: "world", content: worldStr }, + { label: "Recent Memories", pct: memPct, relativePct: memPct, tokens: memTokens, type: "memories", content: memStr || "(No memories yet.)" }, ]; }; @@ -124,15 +124,35 @@ function PromptModal({ const proseTokens = Math.max(0, inputTokens - sysTokens - worldTokens); return [ - { label: "System Prompt", pct: sysPct, tokens: sysTokens, type: "system", content: systemPrompt }, - { label: "Decoder Context", pct: worldPct, tokens: worldTokens, type: "world", content: worldStr }, - { label: "Narrative Prose", pct: prosePct, tokens: proseTokens, type: "memories", content: proseStr }, + { label: "System Prompt", pct: sysPct, relativePct: sysPct, tokens: sysTokens, type: "system", content: systemPrompt }, + { label: "Decoder Context", pct: worldPct, relativePct: worldPct, tokens: worldTokens, type: "world", content: worldStr }, + { label: "Narrative Prose", pct: prosePct, relativePct: prosePct, tokens: proseTokens, type: "memories", content: proseStr }, ]; }; const actorBreakdown = (entry.rawPrompt && entry.usage) ? parseActorPrompt(entry.rawPrompt.systemPrompt, entry.rawPrompt.userContext, entry.usage.inputTokens) : null; const decoderBreakdown = (entry.decoderPrompt && entry.decoderUsage) ? parseDecoderPrompt(entry.decoderPrompt.systemPrompt, entry.decoderPrompt.userContext, entry.decoderUsage.inputTokens) : null; + const actorMaxContext = entry.usage?.maxContext !== undefined ? entry.usage.maxContext : 32768; + const actorUsedTokens = entry.usage?.inputTokens || 0; + const actorUsagePctOfContext = actorMaxContext > 0 ? (actorUsedTokens / actorMaxContext) * 100 : 0; + const isActorAbsolute = actorMaxContext > 0 && actorUsagePctOfContext >= 20; + + const scaledActorBreakdown = actorBreakdown ? actorBreakdown.map((item) => ({ + ...item, + pct: isActorAbsolute ? item.relativePct * (actorUsedTokens / actorMaxContext) : item.relativePct + })) : null; + + const decoderMaxContext = entry.decoderUsage?.maxContext !== undefined ? entry.decoderUsage.maxContext : 32768; + const decoderUsedTokens = entry.decoderUsage?.inputTokens || 0; + const decoderUsagePctOfContext = decoderMaxContext > 0 ? (decoderUsedTokens / decoderMaxContext) * 100 : 0; + const isDecoderAbsolute = decoderMaxContext > 0 && decoderUsagePctOfContext >= 20; + + const scaledDecoderBreakdown = decoderBreakdown ? decoderBreakdown.map((item) => ({ + ...item, + pct: isDecoderAbsolute ? item.relativePct * (decoderUsedTokens / decoderMaxContext) : item.relativePct + })) : null; + useEffect(() => { if (!entry.rawPrompt && entry.decoderPrompt) { setActiveTab("decoder"); @@ -180,37 +200,57 @@ function PromptModal({
)} - {actorBreakdown && ( + {scaledActorBreakdown && (
Input Prompt Breakdown - Total Input Tokens: {entry.usage?.inputTokens} + + Total Input Tokens: {actorUsedTokens} + {actorMaxContext > 0 ? ( + / {actorMaxContext} ({actorUsagePctOfContext.toFixed(1)}% used) + ) : ( + (infinite context) + )} +
- {actorBreakdown.map((item, idx) => ( + {scaledActorBreakdown.map((item, idx) => { + const displayPct = actorMaxContext > 0 ? (item.tokens / actorMaxContext) * 100 : item.relativePct; + return ( +
+ ); + })} + {isActorAbsolute && (
- ))} + )}
- {actorBreakdown.map((item, idx) => ( -
- - - - {item.label}: {item.tokens} tokens ({item.pct.toFixed(0)}%) - - ▼ - -
- 
{item.content}
-
-
- ))} + {scaledActorBreakdown.map((item, idx) => { + const displayPct = actorMaxContext > 0 ? (item.tokens / actorMaxContext) * 100 : item.relativePct; + return ( +
+ + + + {item.label}: {item.tokens} tokens ({displayPct.toFixed(0)}%) + + ▼ + +
+ 
{item.content}
+
+
+ ); + })}
)} @@ -240,37 +280,57 @@ function PromptModal({
)} - {decoderBreakdown && ( + {scaledDecoderBreakdown && (
Input Prompt Breakdown - Total Input Tokens: {entry.decoderUsage?.inputTokens} + + Total Input Tokens: {decoderUsedTokens} + {decoderMaxContext > 0 ? ( + / {decoderMaxContext} ({decoderUsagePctOfContext.toFixed(1)}% used) + ) : ( + (infinite context) + )} +
- {decoderBreakdown.map((item, idx) => ( + {scaledDecoderBreakdown.map((item, idx) => { + const displayPct = decoderMaxContext > 0 ? (item.tokens / decoderMaxContext) * 100 : item.relativePct; + return ( +
+ ); + })} + {isDecoderAbsolute && (
- ))} + )}
- {decoderBreakdown.map((item, idx) => ( -
- - - - {item.label}: {item.tokens} tokens ({item.pct.toFixed(0)}%) - - ▼ - -
- 
{item.content}
-
-
- ))} + {scaledDecoderBreakdown.map((item, idx) => { + const displayPct = decoderMaxContext > 0 ? (item.tokens / decoderMaxContext) * 100 : item.relativePct; + return ( +
+ + + + {item.label}: {item.tokens} tokens ({displayPct.toFixed(0)}%) + + ▼ + +
+ 
{item.content}
+
+
+ ); + })}
)} @@ -488,7 +548,7 @@ export function PlayView() { const [availableEntities, setAvailableEntities] = useState<{ id: string; name: string }[]>([]); const [selectedEntity, setSelectedEntity] = useState(""); - const [providerInstances, setProviderInstances] = useState([]); + const [providerInstances, setProviderInstances] = useState([]); // Load scenarios and provider instances on mount useEffect(() => { diff --git a/apps/gui/src/lib/simulation-types.ts b/apps/gui/src/lib/simulation-types.ts index 9cd87ce..f67e264 100644 --- a/apps/gui/src/lib/simulation-types.ts +++ b/apps/gui/src/lib/simulation-types.ts @@ -26,6 +26,7 @@ export interface LogEntry { totalTokens: number; modelName?: string; providerInstanceName?: string; + maxContext?: number; }; decoderPrompt?: { systemPrompt: string; @@ -37,6 +38,7 @@ export interface LogEntry { totalTokens: number; modelName?: string; providerInstanceName?: string; + maxContext?: number; }; } diff --git a/apps/gui/src/lib/simulation.ts b/apps/gui/src/lib/simulation.ts index dabd212..9cc3197 100644 --- a/apps/gui/src/lib/simulation.ts +++ b/apps/gui/src/lib/simulation.ts @@ -235,11 +235,12 @@ class SimulationManager { const providerName = inst ? inst.providerName : "google-genai"; const modelName = inst ? inst.modelName : undefined; const instanceName = inst ? inst.name : undefined; + const maxContext = inst ? inst.maxContext : undefined; if (providerName === "google-genai") { - return new GeminiProvider(key, modelName, instanceName); + return new GeminiProvider(key, modelName, instanceName, maxContext); } else if (providerName === "openrouter") { - return new OpenRouterProvider(key, modelName, instanceName); + return new OpenRouterProvider(key, modelName, instanceName, maxContext); } else { return new MockLLMProvider([]); } @@ -702,9 +703,9 @@ class SimulationManager { } if (inst.providerName === "google-genai") { - return new GeminiProvider(inst.apiKey, inst.modelName, inst.name); + return new GeminiProvider(inst.apiKey, inst.modelName, inst.name, inst.maxContext); } else if (inst.providerName === "openrouter") { - return new OpenRouterProvider(inst.apiKey, inst.modelName, inst.name); + return new OpenRouterProvider(inst.apiKey, inst.modelName, inst.name, inst.maxContext); } else { return new MockLLMProvider([]); } diff --git a/packages/llm/src/llm.ts b/packages/llm/src/llm.ts index 7a2f621..98b12e8 100644 --- a/packages/llm/src/llm.ts +++ b/packages/llm/src/llm.ts @@ -17,6 +17,7 @@ export interface LLMResponse { totalTokens: number; modelName?: string; providerInstanceName?: string; + maxContext?: number; }; } @@ -29,6 +30,7 @@ export interface LLMCallRecord { totalTokens: number; modelName?: string; providerInstanceName?: string; + maxContext?: number; }; } @@ -53,6 +55,7 @@ export interface ModelProviderInstance { isActive: boolean; modelName?: string; type: "generative" | "embedding"; + maxContext?: number; } export interface ModelProviderMeta { diff --git a/packages/llm/src/provider-manager.ts b/packages/llm/src/provider-manager.ts index 87ad1bc..26a4d26 100644 --- a/packages/llm/src/provider-manager.ts +++ b/packages/llm/src/provider-manager.ts @@ -68,6 +68,12 @@ function getSettingsDb() { // ignore } + try { + db.prepare(`ALTER TABLE provider_instances ADD COLUMN maxContext INTEGER`).run(); + } catch { + // ignore + } + // Auto-bootstrap environment variables if DB contains 0 instances try { if (!hasBootstrapped) { @@ -80,25 +86,25 @@ function getSettingsDb() { if (googleKey && googleKey.trim()) { const id = "provider-default-google"; db.prepare(` - INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type) - VALUES (?, ?, ?, ?, ?, ?, ?) - `).run(id, "Gemini (Env)", "google-genai", googleKey.trim(), 1, "gemini-2.5-flash", "generative"); + INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type, maxContext) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + `).run(id, "Gemini (Env)", "google-genai", googleKey.trim(), 1, "gemini-2.5-flash", "generative", 32768); hasInsertedGenerative = true; const embedId = "provider-default-google-embed"; db.prepare(` - INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type) - VALUES (?, ?, ?, ?, ?, ?, ?) - `).run(embedId, "Gemini Embed (Env)", "google-genai", googleKey.trim(), 1, "gemini-embedding-001", "embedding"); + INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type, maxContext) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + `).run(embedId, "Gemini Embed (Env)", "google-genai", googleKey.trim(), 1, "gemini-embedding-001", "embedding", 0); } if (openRouterKey && openRouterKey.trim()) { const id = "provider-default-openrouter"; const isActive = hasInsertedGenerative ? 0 : 1; db.prepare(` - INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type) - VALUES (?, ?, ?, ?, ?, ?, ?) - `).run(id, "OpenRouter (Env)", "openrouter", openRouterKey.trim(), isActive, "google/gemini-2.5-flash", "generative"); + INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type, maxContext) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + `).run(id, "OpenRouter (Env)", "openrouter", openRouterKey.trim(), isActive, "google/gemini-2.5-flash", "generative", 32768); } } hasBootstrapped = true; @@ -122,6 +128,7 @@ export class ProviderManager { isActive: number; modelName?: string; type: string; + maxContext?: number; }[]; return rows.map((r) => ({ id: r.id, @@ -131,6 +138,7 @@ export class ProviderManager { isActive: r.isActive === 1, modelName: r.modelName || undefined, type: (r.type as "generative" | "embedding") || "generative", + maxContext: r.maxContext !== undefined && r.maxContext !== null ? r.maxContext : (r.type === "embedding" ? 0 : 32768), })); } finally { db.close(); @@ -142,7 +150,8 @@ export class ProviderManager { providerName: string, apiKey: string, modelName?: string, - type: "generative" | "embedding" = "generative" + type: "generative" | "embedding" = "generative", + maxContext?: number ): ModelProviderInstance { const db = getSettingsDb(); try { @@ -152,12 +161,14 @@ export class ProviderManager { .get(type) as { count: number }; const isActive = activeCount.count === 0 ? 1 : 0; + const actualMaxContext = maxContext !== undefined ? maxContext : (type === "generative" ? 32768 : 0); + db.prepare(` - INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type) - VALUES (?, ?, ?, ?, ?, ?, ?) - `).run(id, name, providerName, apiKey, isActive, modelName || null, type); + INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type, maxContext) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + `).run(id, name, providerName, apiKey, isActive, modelName || null, type, actualMaxContext); - return { id, name, providerName, apiKey, isActive: isActive === 1, modelName, type }; + return { id, name, providerName, apiKey, isActive: isActive === 1, modelName, type, maxContext: actualMaxContext }; } finally { db.close(); } @@ -201,22 +212,24 @@ export class ProviderManager { providerName: string, apiKey?: string, modelName?: string, - type: "generative" | "embedding" = "generative" + type: "generative" | "embedding" = "generative", + maxContext?: number ): void { const db = getSettingsDb(); try { + const actualMaxContext = maxContext !== undefined ? maxContext : (type === "generative" ? 32768 : 0); if (apiKey && apiKey.trim()) { db.prepare(` UPDATE provider_instances - SET name = ?, providerName = ?, apiKey = ?, modelName = ?, type = ? + SET name = ?, providerName = ?, apiKey = ?, modelName = ?, type = ?, maxContext = ? WHERE id = ? - `).run(name, providerName, apiKey, modelName || null, type, id); + `).run(name, providerName, apiKey, modelName || null, type, actualMaxContext, id); } else { db.prepare(` UPDATE provider_instances - SET name = ?, providerName = ?, modelName = ?, type = ? + SET name = ?, providerName = ?, modelName = ?, type = ?, maxContext = ? WHERE id = ? - `).run(name, providerName, modelName || null, type, id); + `).run(name, providerName, modelName || null, type, actualMaxContext, id); } } finally { db.close(); @@ -234,6 +247,7 @@ export class ProviderManager { isActive: number; modelName?: string; type: string; + maxContext?: number; } | undefined; if (!row) { @@ -246,25 +260,25 @@ export class ProviderManager { if (googleKey && googleKey.trim()) { const id = "provider-default-google"; db.prepare(` - INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type) - VALUES (?, ?, ?, ?, ?, ?, ?) - `).run(id, "Gemini (Env)", "google-genai", googleKey.trim(), 1, "gemini-2.5-flash", "generative"); + INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type, maxContext) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + `).run(id, "Gemini (Env)", "google-genai", googleKey.trim(), 1, "gemini-2.5-flash", "generative", 32768); hasInsertedGenerative = true; const embedId = "provider-default-google-embed"; db.prepare(` - INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type) - VALUES (?, ?, ?, ?, ?, ?, ?) - `).run(embedId, "Gemini Embed (Env)", "google-genai", googleKey.trim(), 1, "gemini-embedding-001", "embedding"); + INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type, maxContext) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + `).run(embedId, "Gemini Embed (Env)", "google-genai", googleKey.trim(), 1, "gemini-embedding-001", "embedding", 0); } if (openRouterKey && openRouterKey.trim()) { const id = "provider-default-openrouter"; const isActive = hasInsertedGenerative ? 0 : 1; db.prepare(` - INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type) - VALUES (?, ?, ?, ?, ?, ?, ?) - `).run(id, "OpenRouter (Env)", "openrouter", openRouterKey.trim(), isActive, "google/gemini-2.5-flash", "generative"); + INSERT INTO provider_instances (id, name, providerName, apiKey, isActive, modelName, type, maxContext) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + `).run(id, "OpenRouter (Env)", "openrouter", openRouterKey.trim(), isActive, "google/gemini-2.5-flash", "generative", 32768); } const retryRow = db.prepare(`SELECT * FROM provider_instances WHERE isActive = 1 AND type = ?`).get(type) as { @@ -275,6 +289,7 @@ export class ProviderManager { isActive: number; modelName?: string; type: string; + maxContext?: number; } | undefined; if (retryRow) { @@ -286,6 +301,7 @@ export class ProviderManager { isActive: true, modelName: retryRow.modelName || undefined, type: retryRow.type as "generative" | "embedding", + maxContext: retryRow.maxContext !== undefined && retryRow.maxContext !== null ? retryRow.maxContext : (retryRow.type === "embedding" ? 0 : 32768), }; } } @@ -299,6 +315,7 @@ export class ProviderManager { isActive: number; modelName?: string; type: string; + maxContext?: number; } | undefined; if (firstRow) { db.prepare(`UPDATE provider_instances SET isActive = 1 WHERE id = ?`).run(firstRow.id); @@ -310,6 +327,7 @@ export class ProviderManager { isActive: true, modelName: firstRow.modelName || undefined, type: firstRow.type as "generative" | "embedding", + maxContext: firstRow.maxContext !== undefined && firstRow.maxContext !== null ? firstRow.maxContext : (firstRow.type === "embedding" ? 0 : 32768), }; } return null; @@ -323,6 +341,7 @@ export class ProviderManager { isActive: true, modelName: row.modelName || undefined, type: (row.type as "generative" | "embedding") || "generative", + maxContext: row.maxContext !== undefined && row.maxContext !== null ? row.maxContext : (row.type === "embedding" ? 0 : 32768), }; } catch { const googleKey = process.env.GOOGLE_API_KEY; @@ -336,6 +355,7 @@ export class ProviderManager { isActive: true, modelName: "gemini-embedding-001", type: "embedding", + maxContext: 0, }; } return null; @@ -351,6 +371,7 @@ export class ProviderManager { isActive: true, modelName: "gemini-2.5-flash", type: "generative", + maxContext: 32768, }; } const openRouterKey = process.env.OPENROUTER_API_KEY; @@ -363,6 +384,7 @@ export class ProviderManager { isActive: true, modelName: "google/gemini-2.5-flash", type: "generative", + maxContext: 32768, }; } return null; diff --git a/packages/llm/src/providers/google-genai.ts b/packages/llm/src/providers/google-genai.ts index c64ae5a..b1b013b 100644 --- a/packages/llm/src/providers/google-genai.ts +++ b/packages/llm/src/providers/google-genai.ts @@ -14,12 +14,14 @@ export class GeminiProvider implements ILLMProvider { private model: ChatGoogleGenerativeAI; private modelNameUsed: string; private providerInstanceName?: string; + private maxContextUsed?: number; lastCalls: LLMCallRecord[] = []; - constructor(apiKey?: string, modelName?: string, providerInstanceName?: string) { + constructor(apiKey?: string, modelName?: string, providerInstanceName?: string, maxContext?: number) { let key = apiKey; let model = modelName; this.providerInstanceName = providerInstanceName; + this.maxContextUsed = maxContext; if (!key) { const active = ProviderManager.getActive("generative"); @@ -31,6 +33,9 @@ export class GeminiProvider implements ILLMProvider { if (!this.providerInstanceName) { this.providerInstanceName = active.name; } + if (this.maxContextUsed === undefined) { + this.maxContextUsed = active.maxContext; + } } } @@ -79,6 +84,7 @@ export class GeminiProvider implements ILLMProvider { totalTokens: raw?.usage_metadata?.total_tokens || 0, modelName: this.modelNameUsed, providerInstanceName: this.providerInstanceName || "Default", + maxContext: this.maxContextUsed !== undefined ? this.maxContextUsed : 32768, }; this.lastCalls.push({ diff --git a/packages/llm/src/providers/openrouter.ts b/packages/llm/src/providers/openrouter.ts index 2cc970e..c166bb2 100644 --- a/packages/llm/src/providers/openrouter.ts +++ b/packages/llm/src/providers/openrouter.ts @@ -14,12 +14,14 @@ export class OpenRouterProvider implements ILLMProvider { private model: ChatOpenRouter; private modelNameUsed: string; private providerInstanceName?: string; + private maxContextUsed?: number; lastCalls: LLMCallRecord[] = []; - constructor(apiKey?: string, modelName?: string, providerInstanceName?: string) { + constructor(apiKey?: string, modelName?: string, providerInstanceName?: string, maxContext?: number) { let key = apiKey; let model = modelName; this.providerInstanceName = providerInstanceName; + this.maxContextUsed = maxContext; if (!key) { const active = ProviderManager.getActive("generative"); @@ -31,6 +33,9 @@ export class OpenRouterProvider implements ILLMProvider { if (!this.providerInstanceName) { this.providerInstanceName = active.name; } + if (this.maxContextUsed === undefined) { + this.maxContextUsed = active.maxContext; + } } } @@ -79,6 +84,7 @@ export class OpenRouterProvider implements ILLMProvider { totalTokens: raw?.usage_metadata?.total_tokens || 0, modelName: this.modelNameUsed, providerInstanceName: this.providerInstanceName || "Default", + maxContext: this.maxContextUsed !== undefined ? this.maxContextUsed : 32768, }; this.lastCalls.push({ diff --git a/packages/llm/tests/openrouter.test.ts b/packages/llm/tests/openrouter.test.ts index 354de1a..4e2a675 100644 --- a/packages/llm/tests/openrouter.test.ts +++ b/packages/llm/tests/openrouter.test.ts @@ -93,6 +93,7 @@ describe("OpenRouterProvider Unit Tests (Tier 1)", () => { totalTokens: 15, modelName: "google/gemini-2.5-flash", providerInstanceName: "Default", + maxContext: 32768, }); expect(provider.lastCalls.length).toBe(1); @@ -105,6 +106,7 @@ describe("OpenRouterProvider Unit Tests (Tier 1)", () => { totalTokens: 15, modelName: "google/gemini-2.5-flash", providerInstanceName: "Default", + maxContext: 32768, }, }); });