Skip to content

Commit a9013be

Browse files
vmelikyanvmelikyan
authored
Improve system prompts (#100)
1 parent c6b48b7 commit a9013be

31 files changed

Lines changed: 904 additions & 141 deletions

‎src/app/api/v2/ai/chat/[buildUuid]/route.ts‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,9 +115,15 @@ const postHandler = async (req: NextRequest, { params }: { params: { buildUuid:
115115
const { readable, writable } = new TransformStream();
116116
const writer = writable.getWriter();
117117
const encoder = new TextEncoder();
118+
let writerClosed = false;
118119

119120
const sendEvent = (data: AIChatSSEEvent | SSEErrorEvent) => {
120-
writer.write(encoder.encode(`data: ${JSON.stringify(data)}\n\n`)).catch(() => {});
121+
if (writerClosed) return;
122+
try {
123+
writer.write(encoder.encode(`data: ${JSON.stringify(data)}\n\n`)).catch(() => {});
124+
} catch {
125+
writerClosed = true;
126+
}
121127
};
122128

123129
const MAX_STORED_JSON_LENGTH = 8000;
@@ -132,6 +138,7 @@ const postHandler = async (req: NextRequest, { params }: { params: { buildUuid:
132138
};
133139

134140
req.signal.addEventListener('abort', () => {
141+
writerClosed = true;
135142
try {
136143
writer.close();
137144
} catch {
@@ -334,6 +341,10 @@ const postHandler = async (req: NextRequest, { params }: { params: { buildUuid:
334341
iterations: e.iterations,
335342
totalToolCalls: e.totalToolCalls,
336343
totalDurationMs: e.totalDurationMs,
344+
inputTokens: e.inputTokens,
345+
outputTokens: e.outputTokens,
346+
inputCostPerMillion: e.inputCostPerMillion,
347+
outputCostPerMillion: e.outputCostPerMillion,
337348
};
338349
}
339350
} catch {
@@ -464,6 +475,7 @@ const postHandler = async (req: NextRequest, { params }: { params: { buildUuid:
464475
modelName: 'AI model',
465476
});
466477
} finally {
478+
writerClosed = true;
467479
try {
468480
await writer.close();
469481
} catch {

‎src/components/ai-agent/ChatContainer.tsx‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import { ChatInput } from './ChatInput';
2525
import { SuggestedPrompts } from './SuggestedPrompts';
2626
import { MessageList } from './MessageList';
2727
import { globalStyles } from './styles';
28-
import { getFollowUpSuggestions } from './utils';
28+
import { getFollowUpSuggestions, computeCost, formatCost } from './utils';
2929
import type { ServiceInvestigationResult } from './types';
3030

3131
interface ChatContainerProps {
@@ -82,6 +82,26 @@ export function ChatContainer({ buildUuid }: ChatContainerProps) {
8282
[messages, loading, streaming]
8383
);
8484

85+
const sessionTotalCost = useMemo(() => {
86+
let total = 0;
87+
let hasCost = false;
88+
for (const msg of messages) {
89+
if (msg.debugMetrics?.inputTokens != null) {
90+
const cost = computeCost(
91+
msg.debugMetrics.inputTokens,
92+
msg.debugMetrics.outputTokens || 0,
93+
msg.debugMetrics.inputCostPerMillion,
94+
msg.debugMetrics.outputCostPerMillion
95+
);
96+
if (cost != null) {
97+
total += cost;
98+
hasCost = true;
99+
}
100+
}
101+
}
102+
return hasCost ? formatCost(total) : null;
103+
}, [messages]);
104+
85105
if (!mounted) {
86106
return (
87107
<div className="flex items-center justify-center h-screen bg-white">
@@ -102,6 +122,7 @@ export function ChatContainer({ buildUuid }: ChatContainerProps) {
102122
hasMessages={messages.length > 0}
103123
onLabelClick={handleLabelClick}
104124
xrayMode={xrayMode}
125+
sessionTotalCost={sessionTotalCost}
105126
/>
106127
</CardHeader>
107128

‎src/components/ai-agent/ModelSelector.tsx‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ interface ModelSelectorProps {
2727
hasMessages: boolean;
2828
onLabelClick?: () => void;
2929
xrayMode?: boolean;
30+
sessionTotalCost?: string | null;
3031
}
3132

3233
export function ModelSelector({
@@ -38,6 +39,7 @@ export function ModelSelector({
3839
hasMessages,
3940
onLabelClick,
4041
xrayMode,
42+
sessionTotalCost,
4143
}: ModelSelectorProps) {
4244
return (
4345
<>
@@ -57,7 +59,7 @@ export function ModelSelector({
5759
}
5860
labelPlacement="outside-left"
5961
size="sm"
60-
className="w-60"
62+
className="w-72"
6163
variant="bordered"
6264
selectedKeys={selectedModel ? [`${selectedModel.provider}:${selectedModel.modelId}`] : []}
6365
onSelectionChange={(keys) => {
@@ -66,8 +68,10 @@ export function ModelSelector({
6668
}}
6769
isDisabled={loading}
6870
classNames={{
69-
label: 'text-xs text-gray-600 font-semibold',
71+
label: 'text-xs text-gray-600 font-semibold whitespace-nowrap',
7072
trigger: 'border-gray-200 hover:border-gray-300',
73+
value: 'text-ellipsis overflow-visible',
74+
innerWrapper: 'overflow-visible',
7175
}}
7276
>
7377
{availableModels.map((model) => (
@@ -77,6 +81,11 @@ export function ModelSelector({
7781
)}
7882
</div>
7983
<div className="flex items-center gap-2">
84+
{xrayMode && sessionTotalCost && (
85+
<span className="text-xs font-mono text-gray-500">
86+
Session: {sessionTotalCost}
87+
</span>
88+
)}
8089
{hasMessages && (
8190
<Button
8291
onClick={onClearHistory}

‎src/components/ai-agent/ServiceCard.tsx‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,7 @@ export function ServiceCard({
222222
{service.errorSourceDetail}
223223
</p>
224224
)}
225-
{service.filePath && (
225+
{service.filePath && !(service.files && service.files.some((f) => f.path === service.filePath)) && (
226226
<div>
227227
{renderFileLink(service.filePath, repository, service.lineNumber, service.lineNumberEnd)}
228228
</div>
@@ -233,14 +233,16 @@ export function ServiceCard({
233233
<div className="space-y-3">
234234
<div className="flex items-center justify-between">
235235
<h5 className="text-sm font-semibold text-gray-500 uppercase tracking-wider">Error Output</h5>
236-
<Button
237-
size="sm"
238-
variant="light"
239-
onClick={() => setShowDetails(!showDetails)}
240-
className="text-xs"
241-
>
242-
{showDetails ? 'Hide' : 'Show'} raw output
243-
</Button>
236+
{service.keyError.length > 200 && (
237+
<Button
238+
size="sm"
239+
variant="light"
240+
onClick={() => setShowDetails(!showDetails)}
241+
className="text-xs"
242+
>
243+
{showDetails ? 'Hide' : 'Show'} raw output
244+
</Button>
245+
)}
244246
</div>
245247
<div className="bg-danger-50 border-l-4 border-danger p-4 rounded-lg">
246248
<pre className="text-sm text-danger-900 overflow-x-auto m-0 leading-relaxed font-mono whitespace-pre-wrap">

‎src/components/ai-agent/XrayPanel.tsx‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
import React from 'react';
1818
import { Chip } from '@heroui/react';
19-
import { formatDuration } from './utils';
19+
import { formatDuration, computeCost, formatCost } from './utils';
2020
import type { DebugContextData, DebugMetrics } from './types';
2121

2222
interface XrayContextPanelProps {
@@ -63,6 +63,18 @@ export function XrayContextPanel({ debugContext, debugMetrics }: XrayContextPane
6363
<span>iterations: {debugMetrics.iterations}</span>
6464
<span>tool calls: {debugMetrics.totalToolCalls}</span>
6565
<span>duration: {formatDuration(debugMetrics.totalDurationMs)}</span>
66+
{debugMetrics.inputTokens != null && (
67+
<span>tokens: {(debugMetrics.inputTokens + (debugMetrics.outputTokens || 0)).toLocaleString()} ({debugMetrics.inputTokens.toLocaleString()} in / {(debugMetrics.outputTokens || 0).toLocaleString()} out)</span>
68+
)}
69+
{debugMetrics.inputTokens != null && (() => {
70+
const cost = computeCost(
71+
debugMetrics.inputTokens!,
72+
debugMetrics.outputTokens || 0,
73+
debugMetrics.inputCostPerMillion,
74+
debugMetrics.outputCostPerMillion
75+
);
76+
return cost != null ? <span>cost: {formatCost(cost)}</span> : null;
77+
})()}
6678
</div>
6779
</div>
6880
)}

‎src/components/ai-agent/hooks/useChat.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,10 @@ export function useChat({ buildUuid, selectedModel }: UseChatOptions) {
251251
iterations: data.iterations,
252252
totalToolCalls: data.totalToolCalls,
253253
totalDurationMs: data.totalDurationMs,
254+
inputTokens: data.inputTokens,
255+
outputTokens: data.outputTokens,
256+
inputCostPerMillion: data.inputCostPerMillion,
257+
outputCostPerMillion: data.outputCostPerMillion,
254258
};
255259
setDebugMetrics(localDebugMetrics);
256260
} else if (data.type === 'chunk') {

‎src/components/ai-agent/types.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,8 @@ export interface ModelOption {
103103
displayName: string;
104104
default: boolean;
105105
maxTokens: number;
106+
inputCostPerMillion?: number;
107+
outputCostPerMillion?: number;
106108
}
107109

108110
export interface DebugToolData {
@@ -129,4 +131,8 @@ export interface DebugMetrics {
129131
iterations: number;
130132
totalToolCalls: number;
131133
totalDurationMs: number;
134+
inputTokens?: number;
135+
outputTokens?: number;
136+
inputCostPerMillion?: number;
137+
outputCostPerMillion?: number;
132138
}

‎src/components/ai-agent/utils.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,21 @@
1616

1717
import type { StructuredDebugResponse, EvidenceItem, ServiceInvestigationResult } from './types';
1818

19+
export function computeCost(
20+
inputTokens: number,
21+
outputTokens: number,
22+
inputCostPerMillion?: number,
23+
outputCostPerMillion?: number
24+
): number | null {
25+
if (inputCostPerMillion == null || outputCostPerMillion == null) return null;
26+
return (inputTokens * inputCostPerMillion + outputTokens * outputCostPerMillion) / 1_000_000;
27+
}
28+
29+
export function formatCost(cost: number): string {
30+
if (cost < 0.01) return `$${cost.toFixed(4)}`;
31+
return `$${cost.toFixed(2)}`;
32+
}
33+
1934
export function formatDuration(durationMs?: number): string {
2035
if (durationMs === undefined || durationMs === null) return '';
2136

‎src/server/lib/validation/aiAgentConfigSchemas.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ export const aiAgentConfigSchema = {
3636
enabled: { type: 'boolean' },
3737
default: { type: 'boolean' },
3838
maxTokens: { type: 'integer', minimum: 1 },
39+
inputCostPerMillion: { type: 'number', minimum: 0 },
40+
outputCostPerMillion: { type: 'number', minimum: 0 },
3941
},
4042
required: ['id', 'displayName', 'enabled', 'default', 'maxTokens'],
4143
},

‎src/server/services/ai/context/gatherer.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@ export default class AIAgentContextService extends BaseService {
171171
latestCommit: build.pullRequest.latestCommit,
172172
fullName: build.pullRequest.fullName,
173173
commentId: build.pullRequest.commentId,
174+
labels: build.pullRequest.labels || [],
174175
},
175176
environment: {
176177
id: build.environment.id,

0 commit comments

Comments
 (0)