81 lines
2.1 KiB
TypeScript
81 lines
2.1 KiB
TypeScript
import { useState, useEffect, useCallback } from 'react';
|
|
|
|
export interface AgentEvent {
|
|
id: string;
|
|
timestamp: string;
|
|
agent: string;
|
|
tier: 'quick' | 'mid' | 'deep';
|
|
type: 'thought' | 'tool' | 'result' | 'system';
|
|
message: string;
|
|
details?: {
|
|
model_used: string;
|
|
latency_ms: number;
|
|
input_tokens: number;
|
|
output_tokens: number;
|
|
raw_json_response: string;
|
|
};
|
|
}
|
|
|
|
export const useAgentStream = (runId: string | null) => {
|
|
const [events, setEvents] = useState<AgentEvent[]>([]);
|
|
const [status, setStatus] = useState<'idle' | 'connecting' | 'streaming' | 'completed' | 'error'>('idle');
|
|
const [error, setError] = useState<string | null>(null);
|
|
|
|
const connect = useCallback(() => {
|
|
if (!runId) return;
|
|
|
|
setStatus('connecting');
|
|
setError(null);
|
|
|
|
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
|
const host = window.location.host; // Change this to your backend host if different
|
|
const socket = new WebSocket(`${protocol}//${host}/ws/stream/${runId}`);
|
|
|
|
socket.onopen = () => {
|
|
setStatus('streaming');
|
|
console.log(`Connected to run: ${runId}`);
|
|
};
|
|
|
|
socket.onmessage = (event) => {
|
|
const data = JSON.parse(event.data);
|
|
|
|
if (data.type === 'system' && data.message === 'Run completed.') {
|
|
setStatus('completed');
|
|
} else if (data.type === 'system' && data.message.startsWith('Error:')) {
|
|
setStatus('error');
|
|
setError(data.message);
|
|
} else {
|
|
setEvents((prev) => [...prev, data as AgentEvent]);
|
|
}
|
|
};
|
|
|
|
socket.onclose = () => {
|
|
if (status !== 'completed' && status !== 'error') {
|
|
setStatus('idle');
|
|
}
|
|
console.log(`Disconnected from run: ${runId}`);
|
|
};
|
|
|
|
socket.onerror = (err) => {
|
|
setStatus('error');
|
|
setError('WebSocket error occurred');
|
|
console.error(err);
|
|
};
|
|
|
|
return () => {
|
|
socket.close();
|
|
};
|
|
}, [runId, status]);
|
|
|
|
useEffect(() => {
|
|
if (runId) {
|
|
const cleanup = connect();
|
|
return cleanup;
|
|
}
|
|
}, [runId, connect]);
|
|
|
|
const clearEvents = () => setEvents([]);
|
|
|
|
return { events, status, error, clearEvents };
|
|
};
|