Skip to content

Context aware chat #2876

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

Closed
wants to merge 5 commits into from
+296 −7
Closed

Context aware chat #2876

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
aa786c2
context-aware chat component
Rish-it Sep 19, 2025
0307408
cleanup
Rish-it Sep 19, 2025
4d34d65
removed color indicators
Rish-it Sep 19, 2025
81f60a6
refactor
Rish-it Sep 19, 2025
8be2c85
Merge branch 'main' into context-aware-chat
Kitenite Sep 19, 2025
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
17 changes: 13 additions & 4 deletions apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-input/index.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import { InputContextPills } from '../context-pills/input-context-pills';
import { Suggestions, type SuggestionsRef } from '../suggestions';
import { ActionButtons } from './action-buttons';
import { ChatModeToggle } from './chat-mode-toggle';
import { ContextIndicator } from '../context-indicator';

interface ChatInputProps {
messages: ChatMessage[];
Expand Down Expand Up @@ -349,9 +350,17 @@ export const ChatInput = observer(({
}}
/>
<div className="flex flex-col w-full p-4">
<div className="flex flex-row flex-wrap items-center gap-1.5 mb-1">
{/* <ContextWheel /> */}
<InputContextPills />
<div className="flex flex-row flex-wrap items-center justify-between gap-1.5 mb-1">
<div className="flex flex-row flex-wrap items-center gap-1.5">
{editorEngine.chat.context.context.length > 0 && (
<ContextIndicator messages={messages} />
)}
{/* <ContextWheel /> */}
<InputContextPills />
</div>
{editorEngine.chat.context.context.length === 0 && (
<ContextIndicator messages={messages} />
)}
</div>
<Textarea
ref={textareaRef}
Expand All @@ -369,7 +378,7 @@ export const ChatInput = observer(({
onKeyDown={handleKeyDown}
onPaste={handlePaste}
onCompositionStart={() => setIsComposing(true)}
onCompositionEnd={(e) => {
onCompositionEnd={() => {
setIsComposing(false);
}}
onDragEnter={(e) => {
Expand Down
92 changes: 92 additions & 0 deletions apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/context-indicator.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
'use client'

import { useContextTracking } from '@/app/project/[id]/_hooks/use-context-tracking';
import { Tooltip, TooltipContent, TooltipTrigger } from '@onlook/ui/tooltip';
import { cn } from '@onlook/ui/utils';
import { observer } from 'mobx-react-lite';
import type { ChatMessage } from '@onlook/models';

interface ContextIndicatorProps {
messages: ChatMessage[];
modelId?: string;
}

function formatTokens(tokens: number): string {
if (tokens >= 1000000) {
return `${(tokens / 1000000).toFixed(1)}M`;
} else if (tokens >= 1000) {
return `${(tokens / 1000).toFixed(1)}K`;
}
return tokens.toString();
}

export const ContextIndicator = observer(({ messages, modelId = 'openai:gpt-4' }: ContextIndicatorProps) => {
const contextTracking = useContextTracking(messages, modelId);

if (contextTracking.usage.totalTokens === 0) {
return null;
}

const percentage = Math.min(contextTracking.percentage, 100);

// Color logic based on percentage thresholds
let colors = { stroke: '#6b7280', text: 'text-gray-400', bg: '' };

if (percentage >= 90) {
colors = { stroke: '#ef4444', text: 'text-red-500', bg: 'bg-red-900' };
} else if (percentage >= 85) {
colors = { stroke: '#f97316', text: 'text-orange-500', bg: 'bg-orange-900' };
} else if (percentage >= 75) {
colors = { stroke: '#fbbf24', text: 'text-amber-200', bg: 'bg-amber-900' };
} else if (percentage >= 50) {
colors = { stroke: '#6b7280', text: 'text-gray-400', bg: '' };
}

// Format percentage display
const displayPercentage = percentage < 50
? Math.round(percentage).toString() + '%'
: percentage.toFixed(1) + '%';

return (
<Tooltip>
<TooltipTrigger asChild>
<div className={cn(
"flex items-center gap-1.5 px-1.5 py-1 rounded-md hover:bg-muted/30 transition-colors cursor-pointer shadow-md hover:shadow-lg transition-shadow",
colors.bg
)}>
<div className="relative">
<svg width="16" height="16" className="transform -rotate-90">
<circle
cx="8"
cy="8"
r="6"
stroke="currentColor"
strokeWidth="1.5"
fill="none"
className="text-muted-foreground/20"
/>
<circle
cx="8"
cy="8"
r="6"
stroke={colors.stroke}
strokeWidth="1.5"
fill="none"
strokeDasharray={2 * Math.PI * 6}
strokeDashoffset={2 * Math.PI * 6 - (percentage / 100) * 2 * Math.PI * 6}
strokeLinecap="round"
className="transition-all duration-300"
/>
</svg>
</div>
<span className={cn("text-xs font-medium tabular-nums", colors.text)}>
{displayPercentage}
</span>
</div>
</TooltipTrigger>
<TooltipContent side="top" className="hideArrow">
{formatTokens(contextTracking.usage.totalTokens)}/{formatTokens(contextTracking.limits.contextWindow)} context used
</TooltipContent>
</Tooltip>
);
});
33 changes: 33 additions & 0 deletions apps/web/client/src/app/project/[id]/_hooks/use-context-tracking.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
'use client'

import { getContextUsage, getModelLimits } from '@onlook/ai/src/tokens';
import type { TokenUsage, ModelLimits } from '@onlook/ai/src/tokens';
import type { ChatMessage } from '@onlook/models';
import { useEffect, useState } from 'react';

interface ContextTrackingState {
usage: TokenUsage;
limits: ModelLimits;
percentage: number;
}

export const useContextTracking = (messages: ChatMessage[], modelId: string = 'openai:gpt-4') => {
const [contextState, setContextState] = useState<ContextTrackingState>({
usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
limits: getModelLimits(modelId),
percentage: 0
});

useEffect(() => {
const updateContextUsage = async () => {
const contextUsage = await getContextUsage(messages, modelId);
setContextState(contextUsage);
};

updateContextUsage();
}, [messages, modelId]);
Comment on lines +22 to +28
Copy link
Contributor

@graphite-app graphite-app bot Sep 19, 2025

Choose a reason for hiding this comment

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

Async setup race condition: The useEffect calls an async function updateContextUsage() but doesn't await it or handle the Promise properly. This can cause the component to render with stale state while the async operation is still pending. The async function should be properly awaited or the Promise should be handled to prevent race conditions between state updates and renders.

Spotted by Diamond

Fix in Graphite

Is this helpful? React 👍 or 👎 to let us know.


return {
...contextState
};
};
25 changes: 25 additions & 0 deletions apps/web/client/src/components/store/editor/chat/context.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import {
import { assertNever, type ParsedError } from '@onlook/utility';
import { makeAutoObservable, reaction } from 'mobx';
import type { EditorEngine } from '../engine';
import { countTokensInString } from '@onlook/ai/src/tokens';

export class ChatContext {
private _context: MessageContext[] = [];
Expand Down Expand Up @@ -312,6 +313,30 @@ export class ChatContext {
this.context = this.context.filter((context) => context.type !== MessageContextType.IMAGE);
}

getContextTokenCount(): number {
return this.context.reduce((total, ctx) => {
return total + countTokensInString(ctx.content || '');
}, 0);
}

getContextSummary(): {
totalContexts: number;
totalTokens: number;
byType: Record<string, number>;
} {
const summary = {
totalContexts: this.context.length,
totalTokens: this.getContextTokenCount(),
byType: {} as Record<string, number>
};

this.context.forEach(ctx => {
summary.byType[ctx.type] = (summary.byType[ctx.type] || 0) + 1;
});

return summary;
}

clear() {
this.selectedReactionDisposer?.();
this.selectedReactionDisposer = undefined;
Expand Down
11 changes: 11 additions & 0 deletions bun.lock
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,7 @@
"name": "@onlook/ui",
"version": "0.0.0",
"dependencies": {
"@ai-sdk/ui-utils": "^1.2.11",
"@emotion/react": "^11.13.3",
"@emotion/styled": "^11.13.0",
"@hookform/resolvers": "^5.0.1",
Expand Down Expand Up @@ -531,6 +532,8 @@

"@ai-sdk/react": ["@ai-sdk/[email protected]", "", { "dependencies": { "@ai-sdk/provider-utils": "3.0.7", "ai": "5.0.25", "swr": "^2.2.5", "throttleit": "2.1.0" }, "peerDependencies": { "react": "^18 || ^19 || ^19.0.0-rc", "zod": "^3.25.76 || ^4" }, "optionalPeers": ["zod"] }, "sha512-y7yhQbiImAvOcNm75uf4RnRqJvKB/gj790m5Dct+W63seOjEDNYnQhfwxvF0Py3snlcFO+q3SzvUSH+1yG+jlQ=="],

"@ai-sdk/ui-utils": ["@ai-sdk/[email protected]", "", { "dependencies": { "@ai-sdk/provider": "1.1.3", "@ai-sdk/provider-utils": "2.2.8", "zod-to-json-schema": "^3.24.1" }, "peerDependencies": { "zod": "^3.23.8" } }, "sha512-3zcwCc8ezzFlwp3ZD15wAPjf2Au4s3vAbKsXQVyhxODHcmu0iyPO2Eua6D/vicq/AUm/BAo60r97O6HU+EI0+w=="],

"@alloc/quick-lru": ["@alloc/[email protected]", "", {}, "sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw=="],

"@aws-crypto/crc32": ["@aws-crypto/[email protected]", "", { "dependencies": { "@aws-crypto/util": "^5.2.0", "@aws-sdk/types": "^3.222.0", "tslib": "^2.6.2" } }, "sha512-nLbCWqQNgUiwwtFsen1AdzAtvuLRsQS8rYgMuxCrdKf9kOssamGLuPwyTY9wyYblNr9+1XM8v6zoDTPPSIeANg=="],
Expand Down Expand Up @@ -4057,6 +4060,10 @@

"@ai-sdk/react/ai": ["[email protected]", "", { "dependencies": { "@ai-sdk/gateway": "1.0.14", "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.7", "@opentelemetry/api": "1.9.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-c7WeAeGnWk+aKUCMgUal2Qa7POcC1/75SrrYHPhlYLWE+DJ2G0gde5TftwCEMwG+pQ6BhaTQHIAQoBR2GGNyAw=="],

"@ai-sdk/ui-utils/@ai-sdk/provider-utils": ["@ai-sdk/[email protected]", "", { "dependencies": { "@ai-sdk/provider": "1.1.3", "nanoid": "^3.3.8", "secure-json-parse": "^2.7.0" }, "peerDependencies": { "zod": "^3.23.8" } }, "sha512-fqhG+4sCVv8x7nFzYnFo19ryhAa3w096Kmc3hWxMQfW/TubPOmt3A6tYZhl4mUfQWWQMsuSkLrtjlWuXBVSGQA=="],

"@ai-sdk/ui-utils/zod": ["[email protected]", "", {}, "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ=="],

"@aws-crypto/util/@smithy/util-utf8": ["@smithy/[email protected]", "", { "dependencies": { "@smithy/util-buffer-from": "^2.2.0", "tslib": "^2.6.2" } }, "sha512-R8Rdn8Hy72KKcebgLiv8jQcQkXoLMOGGv5uI1/k0l+snqkOzQ1R0ChUBCxWMlBsFMekWjq0wRudIweFs7sKT5A=="],

"@babel/core/convert-source-map": ["[email protected]", "", {}, "sha512-Kvp459HrV2FEJ1CAsi1Ku+MY3kasH19TFykTz2xWmMeq6bk2NU3XXvfJ+Q61m0xktWwt+1HSYf3JZsTms3aRJg=="],
Expand Down Expand Up @@ -4737,6 +4744,10 @@

"@ai-sdk/react/ai/@ai-sdk/provider": ["@ai-sdk/[email protected]", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="],

"@ai-sdk/ui-utils/@ai-sdk/provider-utils/nanoid": ["[email protected]", "", { "bin": { "nanoid": "bin/nanoid.cjs" } }, "sha512-N8SpfPUnUp1bK+PMYW8qSWdl9U+wwNWI4QKxOYDy9JAro3WMX7p2OeVRF9v+347pnakNevPmiHhNmZ2HbFA76w=="],

"@ai-sdk/ui-utils/@ai-sdk/provider-utils/secure-json-parse": ["[email protected]", "", {}, "sha512-6aU+Rwsezw7VR8/nyvKTx8QpWH9FrcYiXXlqC4z5d5XQBDRqtbfsRjnwGyqbi3gddNtWHuEk9OANUotL26qKUw=="],

"@aws-crypto/util/@smithy/util-utf8/@smithy/util-buffer-from": ["@smithy/[email protected]", "", { "dependencies": { "@smithy/is-array-buffer": "^2.2.0", "tslib": "^2.6.2" } }, "sha512-IJdWBbTcMQ6DA0gdNhh/BwrLkDR+ADW5Kr1aZmd4k3DIF6ezMV4R2NIAmT08wQJ3yUK82thHWmC/TnK/wpMMIA=="],

"@babel/helper-compilation-targets/lru-cache/yallist": ["[email protected]", "", {}, "sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g=="],
Expand Down
124 changes: 121 additions & 3 deletions packages/ai/src/tokens/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,97 @@ import { type ChatMessage } from '@onlook/models';
import type { TextUIPart, ToolUIPart } from 'ai';
import { encode } from 'gpt-tokenizer';

export interface ModelLimits {
maxTokens: number;
contextWindow: number;
inputLimit: number;
outputLimit: number;
}

export const MODEL_LIMITS: Record<string, ModelLimits> = {
'claude-sonnet-4-20250514': {
maxTokens: 1000000,
contextWindow: 1000000,
inputLimit: 800000,
outputLimit: 200000,
},
'claude-3-5-haiku-20241022': {
maxTokens: 200000,
contextWindow: 200000,
inputLimit: 180000,
outputLimit: 200000,
},
'anthropic/claude-sonnet-4': {
maxTokens: 1000000,
contextWindow: 1000000,
inputLimit: 800000,
outputLimit: 200000,
},
'anthropic/claude-3.5-haiku': {
maxTokens: 200000,
contextWindow: 200000,
inputLimit: 180000,
outputLimit: 200000,
},
'openai/gpt-5': {
maxTokens: 400000,
contextWindow: 400000,
inputLimit: 272000,
outputLimit: 400000,
},
'openai/gpt-5-mini': {
maxTokens: 400000,
contextWindow: 400000,
inputLimit: 272000,
outputLimit: 400000,
},
'openai/gpt-5-nano': {
maxTokens: 400000,
contextWindow: 400000,
inputLimit: 272000,
outputLimit: 400000,
},
'gpt-5': {
maxTokens: 400000,
contextWindow: 400000,
inputLimit: 272000,
outputLimit: 128000,
},
'claude-sonnet-4': {
maxTokens: 1000000,
contextWindow: 1000000,
inputLimit: 800000,
outputLimit: 200000,
},
'claude-3-5-haiku': {
maxTokens: 200000,
contextWindow: 200000,
inputLimit: 180000,
outputLimit: 8000,
},
default: {
maxTokens: 128000,
contextWindow: 128000,
inputLimit: 100000,
outputLimit: 4000,
},
};
Comment on lines +12 to +79
Copy link

@coderabbitai coderabbitai bot Sep 19, 2025

Choose a reason for hiding this comment

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

⚠️ Potential issue

🧩 Analysis chain

Verify model limit accuracy and consider maintenance.

The model limits contain specific values for various AI models. These may become outdated as providers update their offerings.

Please verify:

  1. The accuracy of these model limits, especially for newer models like GPT-5
  2. Consider documenting the source of these limits
  3. Plan for regular updates as model providers change their limits

🌐 Web query:

What are the current context window sizes for GPT-5 and Claude Sonnet 4?

💡 Result:

  • GPT-5 (API): total context 400,000 tokens (up to 272,000 input + up to 128,000 output). [1] (OpenAI, Aug 7, 2025)
  • Claude Sonnet 4: default/standard context 200,000 tokens; Sonnet 4 also supports a 1,000,000-token context window (public beta / tiered availability). [2][3] (Anthropic, Aug 2025)

Fix token-limit inaccuracies and add provenance

packages/ai/src/tokens/index.ts (lines 12–79):

  • openai/gpt-5, openai/gpt-5-mini, openai/gpt-5-nano — outputLimit is 400000; change to 128000 and keep inputLimit=272000, contextWindow/maxTokens=400000 (OpenAI: total 400,000 = up to 272,000 input + up to 128,000 output; Aug 7, 2025).
  • gpt-5 (alias) already uses outputLimit=128000 — consolidate duplicates to a single canonical entry to avoid conflicting values.
  • claude-sonnet-4* variants — currently set to 1,000,000; Anthropic lists default/standard 200,000 with 1,000,000 available only as a tier/beta (Aug 2025). Set default to 200000 or add explicit tiered entries + availability notes.
  • Add inline comments citing provider docs (with dates), and add a validation test that flags configs where outputLimit > provider-declared output limit or > (contextWindow - inputLimit). Add a short README/update cadence for keeping limits current.
🤖 Prompt for AI Agents 
In packages/ai/src/tokens/index.ts around lines 12–79, several model token
limits are incorrect or duplicated: adjust openai/gpt-5, openai/gpt-5-mini, and
openai/gpt-5-nano outputLimit from 400000 to 128000 while keeping
inputLimit=272000 and contextWindow/maxTokens=400000; consolidate the duplicate
"gpt-5" alias into a single canonical entry (remove conflicting duplicate);
change claude-sonnet-4* default entries from 1,000,000 to 200,000 (or add
explicit tiered entries with availability notes) and add short inline comments
citing provider docs with dates for each provider; then add a unit/validation
test that fails if any config has outputLimit greater than the provider-declared
output limit or greater than contextWindow - inputLimit, and include a brief
README note describing update cadence for keeping limits current.


export function getModelLimits(modelId: string): ModelLimits {
return MODEL_LIMITS[modelId] || MODEL_LIMITS['default']!;
}

export async function countTokensWithRoles(messages: ChatMessage[]): Promise<number> {
const perMessageExtra = 4; // ~role + metadata tokens (OpenAI chat format)
const perReplyExtra = 2; // for assistant reply priming
const perMessageExtra = 4;
const perReplyExtra = 2;
let total = 0;
for (const m of messages) {
const content = m.parts
.map((p) => {
if (p.type === 'text') {
return (p as TextUIPart).text;
} else if (p.type.startsWith('tool-')) {
return JSON.stringify((p as ToolUIPart).input); // TODO: check if this is correct
return JSON.stringify((p as ToolUIPart).input);
}
return '';
})
Expand All @@ -21,3 +101,41 @@ export async function countTokensWithRoles(messages: ChatMessage[]): Promise<num
}
return total + perReplyExtra;
}

export interface TokenUsage {
inputTokens: number;
outputTokens: number;
totalTokens: number;
}

export function countTokensInString(text: string): number {
return encode(text).length;
}

export async function getContextUsage(
messages: ChatMessage[],
modelId: string = 'openai:gpt-4',
): Promise<{
usage: TokenUsage;
limits: ModelLimits;
percentage: number;
}> {
const totalTokens = await countTokensWithRoles(messages);
const limits = getModelLimits(modelId);
const inputTokens = Math.floor(totalTokens * 0.8);
Copy link
Contributor

@ellipsis-dev ellipsis-dev bot Sep 19, 2025

Choose a reason for hiding this comment

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

The fixed 80/20 split for input/output tokens may not accurately reflect actual token usage. Consider clarifying or revisiting this logic.

const outputTokens = totalTokens - inputTokens;
Comment on lines +125 to +126
Copy link
Contributor

@graphite-app graphite-app bot Sep 19, 2025

Choose a reason for hiding this comment

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

Logic error in token calculation: The code arbitrarily assigns 80% of total tokens as input tokens and 20% as output tokens (lines 125-126). This is incorrect because it's calculating token distribution from existing messages, not predicting future usage. For existing messages, all tokens should be considered input tokens since they're already part of the conversation context. The output tokens should represent available space for the model's response, not a portion of existing messages.

Suggested change
const inputTokens = Math.floor(totalTokens * 0.8);
const outputTokens = totalTokens - inputTokens;
const inputTokens = totalTokens;
const outputTokens = maxTokens - inputTokens;

Spotted by Diamond

Fix in Graphite

Is this helpful? React 👍 or 👎 to let us know.


const usage: TokenUsage = {
inputTokens,
outputTokens,
totalTokens,
};

const percentage = (totalTokens / limits.contextWindow) * 100;

return {
usage,
limits,
percentage,
};
}
1 change: 1 addition & 0 deletions packages/ui/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
"./hooks": "./src/hooks/index.ts"
},
"dependencies": {
"@ai-sdk/ui-utils": "^1.2.11",
"@emotion/react": "^11.13.3",
"@emotion/styled": "^11.13.0",
"@hookform/resolvers": "^5.0.1",
Expand Down