Spaces:
Sleeping
Sleeping
File size: 2,164 Bytes
8e0dd55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | /**
* WebSocket client for chat completions
* This replaces the HTTP streaming endpoint with a WebSocket connection
*/
// Get the server base URL from environment or use default
const SERVER_BASE_URL = process.env.SERVER_BASE_URL || 'http://localhost:8001';
// Convert HTTP URL to WebSocket URL
const getWebSocketUrl = () => {
const baseUrl = SERVER_BASE_URL;
// Replace http:// with ws:// or https:// with wss://
const wsBaseUrl = baseUrl.replace(/^http/, 'ws');
return `${wsBaseUrl}/ws/chat`;
};
export interface ChatMessage {
role: 'user' | 'assistant' | 'system';
content: string;
}
export interface ChatCompletionRequest {
repo_url: string;
messages: ChatMessage[];
filePath?: string;
token?: string;
type?: string;
provider?: string;
model?: string;
language?: string;
excluded_dirs?: string;
excluded_files?: string;
}
/**
* Creates a WebSocket connection for chat completions
* @param request The chat completion request
* @param onMessage Callback for received messages
* @param onError Callback for errors
* @param onClose Callback for when the connection closes
* @returns The WebSocket connection
*/
export const createChatWebSocket = (
request: ChatCompletionRequest,
onMessage: (message: string) => void,
onError: (error: Event) => void,
onClose: () => void
): WebSocket => {
// Create WebSocket connection
const ws = new WebSocket(getWebSocketUrl());
// Set up event handlers
ws.onopen = () => {
console.log('WebSocket connection established');
// Send the request as JSON
ws.send(JSON.stringify(request));
};
ws.onmessage = (event) => {
// Call the message handler with the received text
onMessage(event.data);
};
ws.onerror = (error) => {
console.error('WebSocket error:', error);
onError(error);
};
ws.onclose = () => {
console.log('WebSocket connection closed');
onClose();
};
return ws;
};
/**
* Closes a WebSocket connection
* @param ws The WebSocket connection to close
*/
export const closeWebSocket = (ws: WebSocket | null): void => {
if (ws && ws.readyState === WebSocket.OPEN) {
ws.close();
}
};
|