|
|
|
import { GoogleGenAI, Chat, HarmCategory, HarmBlockThreshold, GenerateContentResponse, Part, Content } from "@google/genai"; |
|
import { MODEL_NAME } from '../constants'; |
|
import type { GroundingChunk, UploadedFile } from '../types'; |
|
|
|
const API_KEY = process.env.API_KEY; |
|
|
|
if (!API_KEY) { |
|
console.error("API_KEY for Gemini is not set. Please set the API_KEY environment variable."); |
|
} |
|
|
|
const ai = new GoogleGenAI({ apiKey: API_KEY! }); |
|
|
|
const generationConfigValues = { |
|
temperature: 0.7, |
|
topK: 1, |
|
topP: 1, |
|
maxOutputTokens: 4096, |
|
}; |
|
|
|
const safetySettingsList = [ |
|
{ category: HarmCategory.HARM_CATEGORY_HARASSMENT, threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE }, |
|
{ category: HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE }, |
|
{ category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE }, |
|
{ category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE }, |
|
]; |
|
|
|
async function createChatSessionWithHistory(systemInstructionText: string, history: Content[]): Promise<Chat> { |
|
const filteredHistory = history.filter(content => content.role !== "system"); |
|
|
|
return ai.chats.create({ |
|
model: MODEL_NAME, |
|
history: filteredHistory, |
|
config: { |
|
systemInstruction: systemInstructionText, |
|
temperature: generationConfigValues.temperature, |
|
topK: generationConfigValues.topK, |
|
topP: generationConfigValues.topP, |
|
maxOutputTokens: generationConfigValues.maxOutputTokens, |
|
safetySettings: safetySettingsList, |
|
tools: [{ googleSearch: {} }], |
|
|
|
} |
|
}); |
|
} |
|
|
|
async function* sendMessageStream( |
|
chat: Chat, |
|
messageText: string, |
|
imageFile?: UploadedFile |
|
): AsyncGenerator<{ text: string; groundingChunks?: GroundingChunk[] }, void, undefined> { |
|
try { |
|
const parts: Part[] = []; |
|
const userProvidedText = messageText.trim(); |
|
|
|
if (imageFile) { |
|
if (imageFile.dataUrl && imageFile.type.startsWith('image/')) { |
|
const base64Data = imageFile.dataUrl.split(',')[1]; |
|
if (base64Data) { |
|
if (!userProvidedText) { |
|
parts.push({ text: "Describe this image" }); |
|
} else { |
|
parts.push({ text: userProvidedText }); |
|
} |
|
parts.push({ |
|
inlineData: { |
|
mimeType: imageFile.type, |
|
data: base64Data, |
|
}, |
|
}); |
|
} |
|
} else { |
|
const docInfo = `(User uploaded a document: ${imageFile.name})`; |
|
if (userProvidedText) { |
|
parts.push({ text: `${userProvidedText}\n${docInfo}` }); |
|
} else { |
|
parts.push({ text: docInfo }); |
|
} |
|
} |
|
} else { |
|
if (userProvidedText) { |
|
parts.push({ text: userProvidedText }); |
|
} |
|
} |
|
|
|
if (parts.length === 0) { |
|
parts.push({ text: " " }); |
|
} |
|
|
|
const result = await chat.sendMessageStream({ message: parts }); |
|
for await (const chunk of result) { |
|
const text = chunk.text; |
|
const groundingMetadata = chunk.candidates?.[0]?.groundingMetadata; |
|
let groundingChunks: GroundingChunk[] | undefined = undefined; |
|
if (groundingMetadata?.groundingChunks && groundingMetadata.groundingChunks.length > 0) { |
|
groundingChunks = groundingMetadata.groundingChunks |
|
.filter(gc => gc.web?.uri && gc.web?.title) |
|
.map(gc => ({ web: { uri: gc.web!.uri, title: gc.web!.title } })); |
|
} |
|
yield { text, groundingChunks }; |
|
} |
|
} catch (error) { |
|
console.error("Error in sendMessageStream:", error); |
|
if (error instanceof Error) { |
|
yield { text: `\n\n[AI Error: ${error.message}]` }; |
|
} else { |
|
yield { text: `\n\n[An unexpected AI error occurred]` }; |
|
} |
|
throw error; |
|
} |
|
} |
|
|
|
|
|
|
|
export const geminiService = { |
|
createChatSessionWithHistory, |
|
sendMessageStream, |
|
|
|
}; |
|
|