AI module - init

This commit is contained in:
Philipinho
2025-08-15 23:18:51 -07:00
parent 0f29eb8842
commit bb83d12c8b
5 changed files with 236 additions and 1 deletions

View File

@ -0,0 +1,61 @@
import { useState, useCallback, useRef } from "react";
import { useAiGenerateStreamMutation } from "@/ee/ai/queries/ai-query.ts";
import { AiGenerateDto } from "@/ee/ai/types/ai.types.ts";
export function useAiStream() {
const [content, setContent] = useState("");
const [isStreaming, setIsStreaming] = useState(false);
const abortControllerRef = useRef<AbortController | null>(null);
const mutation = useAiGenerateStreamMutation();
const startStream = useCallback(
async (data: AiGenerateDto) => {
setContent("");
setIsStreaming(true);
try {
const controller = await mutation.mutateAsync({
...data,
onChunk: (chunk) => {
setContent((prev) => prev + chunk.content);
},
onError: (error) => {
console.error("AI stream error:", error);
setIsStreaming(false);
},
onComplete: () => {
setIsStreaming(false);
},
});
abortControllerRef.current = controller;
} catch (error) {
console.error("Failed to start stream:", error);
setIsStreaming(false);
}
},
[mutation]
);
const stopStream = useCallback(() => {
if (abortControllerRef.current) {
abortControllerRef.current.abort();
abortControllerRef.current = null;
setIsStreaming(false);
}
}, []);
const resetContent = useCallback(() => {
setContent("");
}, []);
return {
content,
isStreaming,
startStream,
stopStream,
resetContent,
isLoading: mutation.isPending,
error: mutation.error,
};
}

View File

@ -0,0 +1,45 @@
import {
useMutation,
UseMutationResult,
useQuery,
UseQueryResult,
} from "@tanstack/react-query";
import {
generateAiContent,
generateAiContentStream,
getAiConfig,
} from "@/ee/ai/services/ai-service.ts";
import {
AiConfigResponse,
AiContentResponse,
AiGenerateDto,
AiStreamChunk,
AiStreamError,
} from "@/ee/ai/types/ai.types.ts";
export function useAiGenerateMutation(): UseMutationResult<
AiContentResponse,
Error,
AiGenerateDto
> {
return useMutation({
mutationFn: (data: AiGenerateDto) => generateAiContent(data),
});
}
interface StreamCallbacks {
onChunk: (chunk: AiStreamChunk) => void;
onError?: (error: AiStreamError) => void;
onComplete?: () => void;
}
export function useAiGenerateStreamMutation(): UseMutationResult<
AbortController,
Error,
AiGenerateDto & StreamCallbacks
> {
return useMutation({
mutationFn: ({ onChunk, onError, onComplete, ...data }) =>
generateAiContentStream(data, onChunk, onError, onComplete),
});
}

View File

@ -0,0 +1,89 @@
import api from "@/lib/api-client.ts";
import {
AiGenerateDto,
AiContentResponse,
AiStreamChunk,
AiStreamError,
} from "@/ee/ai/types/ai.types.ts";
export async function generateAiContent(
data: AiGenerateDto,
): Promise<AiContentResponse> {
const req = await api.post<AiContentResponse>("/ai/generate", data);
return req.data;
}
export async function generateAiContentStream(
data: AiGenerateDto,
onChunk: (chunk: AiStreamChunk) => void,
onError?: (error: AiStreamError) => void,
onComplete?: () => void,
): Promise<AbortController> {
const abortController = new AbortController();
try {
const response = await fetch("/api/ai/generate/stream", {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(data),
signal: abortController.signal,
credentials: "include", // This ensures cookies are sent, matching axios withCredentials
});
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}
const reader = response.body?.getReader();
const decoder = new TextDecoder();
if (!reader) {
throw new Error("Response body is not readable");
}
const processStream = async () => {
try {
while (true) {
const { done, value } = await reader.read();
if (done) break;
const chunk = decoder.decode(value, { stream: true });
const lines = chunk.split("\n");
for (const line of lines) {
if (line.startsWith("data: ")) {
const data = line.slice(6);
if (data === "[DONE]") {
onComplete?.();
return;
}
try {
const parsed = JSON.parse(data);
if (parsed.error) {
onError?.(parsed);
} else {
onChunk(parsed);
}
} catch (e) {
// Ignore parse errors for incomplete chunks
}
}
}
}
} catch (error) {
if (error.name !== "AbortError") {
onError?.({ error: error.message });
}
} finally {
reader.releaseLock();
}
};
processStream();
} catch (error) {
onError?.({ error: error.message });
}
return abortController;
}

View File

@ -0,0 +1,40 @@
export enum AiAction {
IMPROVE_WRITING = "improve_writing",
FIX_SPELLING_GRAMMAR = "fix_spelling_grammar",
MAKE_SHORTER = "make_shorter",
MAKE_LONGER = "make_longer",
SIMPLIFY = "simplify",
CHANGE_TONE = "change_tone",
SUMMARIZE = "summarize",
CONTINUE_WRITING = "continue_writing",
TRANSLATE = "translate",
CUSTOM = "custom",
}
export interface AiGenerateDto {
action?: AiAction;
content: string;
prompt?: string;
}
export interface AiContentResponse {
content: string;
usage?: {
promptTokens: number;
completionTokens: number;
totalTokens: number;
};
}
export interface AiConfigResponse {
configured: boolean;
availableActions: AiAction[];
}
export interface AiStreamChunk {
content: string;
}
export interface AiStreamError {
error: string;
}