Newer
Older
abysiuscodium / extensions / abysius-ai / src / api.ts
import * as vscode from 'vscode';

interface AbysiusConfig {
    apiKey: string;
    chatEndpoint: string;
    inlineEndpoint: string;
    model: string;
}

export interface InlineRequest {
    prompt: string;
    suffix: string;
    language: string;
    filename: string;
    max_tokens?: number;
}

export interface InlineResponse {
    completion: string;
    finish_reason: string;
}

export interface ChatMessage {
    role: 'user' | 'assistant' | 'system';
    content: string;
}

export interface ChatRequest {
    messages: ChatMessage[];
    model?: string;
    stream?: boolean;
    temperature?: number;
    max_tokens?: number;
}

export interface ChatResponse {
    message: ChatMessage;
    finish_reason: string;
    usage?: {
        prompt_tokens: number;
        completion_tokens: number;
        total_tokens: number;
    };
}

export class AbysiusApi {
    private config: AbysiusConfig;

    constructor(config: AbysiusConfig) {
        this.config = config;
    }

    updateConfig(config: Partial<AbysiusConfig>) {
        this.config = { ...this.config, ...config };
    }

    private getHeaders(): Record<string, string> {
        const headers: Record<string, string> = {
            'Content-Type': 'application/json',
            'User-Agent': 'AbysiusCodium/0.1.0'
        };
        if (this.config.apiKey) {
            headers['Authorization'] = `Bearer ${this.config.apiKey}`;
        }
        return headers;
    }

    async getInlineCompletion(request: InlineRequest, signal?: AbortSignal): Promise<InlineResponse | null> {
        try {
            const config = vscode.workspace.getConfiguration('abysius');
            const maxLength = config.get<number>('inlineCompletionMaxLength', 200);

            const response = await fetch(this.config.inlineEndpoint, {
                method: 'POST',
                headers: this.getHeaders(),
                body: JSON.stringify({
                    ...request,
                    model: this.config.model,
                    max_tokens: Math.min(request.max_tokens || maxLength, maxLength)
                }),
                signal
            });

            if (!response.ok) {
                const error = await response.text();
                console.error('[Abysius API] Inline completion error:', error);
                return null;
            }

            return await response.json() as InlineResponse;
        } catch (err) {
            if (err instanceof Error && err.name === 'AbortError') {
                return null;
            }
            console.error('[Abysius API] Inline completion failed:', err);
            return null;
        }
    }

    async *streamChat(request: ChatRequest, signal?: AbortSignal): AsyncGenerator<string, ChatResponse | null, unknown> {
        try {
            const response = await fetch(this.config.chatEndpoint, {
                method: 'POST',
                headers: this.getHeaders(),
                body: JSON.stringify({
                    ...request,
                    model: this.config.model,
                    stream: true
                }),
                signal
            });

            if (!response.ok) {
                const error = await response.text();
                console.error('[Abysius API] Chat error:', error);
                return null;
            }

            const reader = response.body?.getReader();
            if (!reader) return null;

            const decoder = new TextDecoder();
            let buffer = '';

            try {
                while (true) {
                    const { done, value } = await reader.read();
                    if (done) break;

                    buffer += decoder.decode(value, { stream: true });
                    const lines = buffer.split('\n');
                    buffer = lines.pop() || '';

                    for (const line of lines) {
                        const trimmed = line.trim();
                        if (!trimmed || !trimmed.startsWith('data: ')) continue;
                        
                        const data = trimmed.slice(6);
                        if (data === '[DONE]') continue;

                        try {
                            const chunk = JSON.parse(data);
                            const delta = chunk.choices?.[0]?.delta?.content;
                            if (delta) {
                                yield delta;
                            }
                        } catch {
                            // ignore malformed JSON
                        }
                    }
                }
            } finally {
                reader.releaseLock();
            }

            return null;
        } catch (err) {
            if (err instanceof Error && err.name === 'AbortError') {
                return null;
            }
            console.error('[Abysius API] Chat stream failed:', err);
            return null;
        }
    }

    async sendChat(request: ChatRequest, signal?: AbortSignal): Promise<ChatResponse | null> {
        try {
            const response = await fetch(this.config.chatEndpoint, {
                method: 'POST',
                headers: this.getHeaders(),
                body: JSON.stringify({
                    ...request,
                    model: this.config.model,
                    stream: false
                }),
                signal
            });

            if (!response.ok) {
                const error = await response.text();
                console.error('[Abysius API] Chat error:', error);
                return null;
            }

            return await response.json() as ChatResponse;
        } catch (err) {
            if (err instanceof Error && err.name === 'AbortError') {
                return null;
            }
            console.error('[Abysius API] Chat failed:', err);
            return null;
        }
    }
}