Add function calling

src/services/ai/ai.ts

@@ -1,28 +1,123 @@
 import { GoogleGenAI } from "@google/genai";
 import { config } from "../../config.js";
-import type { AIResponseImage, AIResponseText } from "./types.js";
+import type {
+    AIResponseImage,
+    AIResponseText,
+    AIToolMatrixData,
+} from "./types.js";
 import type { IAIInstructions } from "../../modules/ai/types.js";
+import { FunctionCallingConfigMode } from "@google/genai";
+import { toolFunctions, tools } from "./tools.js";
+import type { FunctionResponse } from "@google/genai";
+import type { Content } from "@google/genai";
 
 const googleAI = new GoogleGenAI({
     apiKey: config.app.ai.api.key,
 });
 
 const getTextGemini = async (
+    matrixData: AIToolMatrixData,
     instructions: IAIInstructions,
     input: string,
     oldInput?: string,
 ): Promise<AIResponseText> => {
+    const inputContent: Content = {
+        role: "user",
+        parts: [
+            {
+                text: input,
+            },
+        ],
+    };
+    const oldInputContent: Content = {
+        role: "user",
+        parts: [
+            {
+                text: oldInput ?? "",
+            },
+        ],
+    };
+
+    const contents: Content[] = oldInput
+        ? [oldInputContent, inputContent]
+        : [inputContent];
+
     const response = await googleAI.models.generateContent({
         model: "gemini-3-flash-preview",
-        contents: oldInput ? [oldInput, input] : input,
+        contents: contents,
         config: {
             systemInstruction: JSON.stringify(instructions),
+            toolConfig: {
+                functionCallingConfig: {
+                    mode: FunctionCallingConfigMode.AUTO,
+                },
+            },
+            tools: [{ functionDeclarations: tools }],
+        },
+    });
+
+    let text = response.text ?? "AI Error";
+    let token = response.usageMetadata?.totalTokenCount ?? 0;
+
+    const content = response.candidates?.at(0)?.content;
+    const functionCall = content?.parts?.at(0)?.functionCall;
+
+    if (response.text || !content || !functionCall) {
+        return {
+            text: text,
+            tokens: token,
+        };
+    }
+
+    text = `Calling function ${functionCall.name}`;
+
+    const func = toolFunctions.find(
+        (func) => func.name === functionCall.name,
+    )?.function;
+    if (!func) {
+        return {
+            text: text,
+            tokens: token,
+        };
+    }
+
+    const output = func(matrixData, functionCall.args);
+    const functionResponse: FunctionResponse = {
+        id: functionCall.id ?? "",
+        name: functionCall.name ?? "",
+        response: {
+            output: JSON.stringify(output),
+        },
+    };
+
+    const responseTool = await googleAI.models.generateContent({
+        model: "gemini-3-flash-preview",
+        contents: [
+            ...contents,
+            content,
+            {
+                role: "tool",
+                parts: [
+                    {
+                        functionResponse: functionResponse,
+                    },
+                ],
+            },
+        ],
+        config: {
+            systemInstruction: JSON.stringify(instructions),
+            toolConfig: {
+                functionCallingConfig: {
+                    mode: FunctionCallingConfigMode.AUTO,
+                },
+            },
+            tools: [{ functionDeclarations: tools }],
         },
     });
 
     return {
-        text: response.text ?? "AI Error",
-        tokens: response.usageMetadata?.totalTokenCount ?? 0,
+        text: responseTool.text ?? "AI Error",
+        tokens: token + (responseTool.usageMetadata?.totalTokenCount ?? 0),
     };
 };
 

src/services/ai/index.ts

@@ -1,2 +1,3 @@
 export * from "./ai.js";
 export * from "./types.js";
+export * from "./tools.js";

src/services/ai/tools.ts

@@ -0,0 +1,139 @@
+import type { FunctionDeclaration } from "@google/genai";
+import { config, packageConfig } from "../../config.js";
+import type { AIToolFunction } from "./types.js";
+import { save, state } from "../../store/store.js";
+
+const tools: FunctionDeclaration[] = [
+    {
+        name: "getBotInfo",
+        description: "Gets information about this AI bot",
+        parametersJsonSchema: {
+            type: "object",
+            properties: {},
+            required: [],
+        },
+    },
+    {
+        name: "shutdownBot",
+        description: "Shuts/Turns this AI bot down",
+        parametersJsonSchema: {
+            type: "object",
+            properties: {
+                password: {
+                    type: "string",
+                    description: "shutdown password provided by the user",
+                },
+            },
+            required: ["password"],
+        },
+    },
+    {
+        name: "remember",
+        description: "Gets information about this AI bot",
+        parametersJsonSchema: {
+            type: "object",
+            properties: {
+                knowledge: {
+                    type: "string",
+                    description:
+                        "new knowledge to remember and add to bot total knowledge",
+                },
+            },
+            required: ["knowledge"],
+        },
+    },
+    {
+        name: "changeUsername",
+        description: "Changes the username of this AI bot",
+        parametersJsonSchema: {
+            type: "object",
+            properties: {
+                username: {
+                    type: "string",
+                    description: "new username",
+                },
+            },
+            required: ["username"],
+        },
+    },
+    {
+        name: "getFullUserData",
+        description: "Gets the full user data about all users",
+        parametersJsonSchema: {
+            type: "object",
+            properties: {},
+            required: [],
+        },
+    },
+];
+
+const toolFunctions: AIToolFunction[] = [
+    {
+        name: "getBotInfo",
+        function: () => {
+            return {
+                botInfo: `Aslobot version ${packageConfig.version}`,
+            };
+        },
+    },
+    {
+        name: "shutdownBot",
+        function: async (matrix, args) => {
+            if (args.password === config.app.ai.password) {
+                await matrix.client.sendTextMessage(
+                    matrix.roomId,
+                    "Saving data...",
+                );
+                save();
+
+                await matrix.client.sendTextMessage(
+                    matrix.roomId,
+                    "Shutting down...",
+                );
+                process.exit(0);
+            }
+            return {
+                message: "incorrect password",
+            };
+        },
+    },
+    {
+        name: "remember",
+        function: (_matrix, args) => {
+            if (!state.aiMemory) {
+                state.aiMemory = [];
+            }
+            if (args.knowledge) {
+                state.aiMemory.push(args.knowledge);
+            }
+
+            return {
+                message: "success",
+            };
+        },
+    },
+    {
+        name: "changeUsername",
+        function: (matrix, args) => {
+            if (args.username) {
+                state.aiMemory.push(args.knowledge);
+            }
+
+            matrix.client.setDisplayName(args.username);
+
+            return {
+                message: "success",
+            };
+        },
+    },
+    {
+        name: "getFullUserData",
+        function: (_matrix, _args) => {
+            return {
+                users: state.users,
+            };
+        },
+    },
+];
+
+export { tools, toolFunctions };

src/services/ai/types.ts

@@ -1,3 +1,5 @@
+import type { MatrixClient } from "matrix-js-sdk";
+
 interface AIResponseText {
     text: string;
     tokens: number;
@@ -8,4 +10,20 @@ interface AIResponseImage {
     tokens: number;
 }
 
-export { type AIResponseText, type AIResponseImage };
+interface AIToolFunction {
+    name: string;
+    function: (matrix: AIToolMatrixData, args: any) => object;
+}
+
+interface AIToolMatrixData {
+    client: MatrixClient;
+    roomId: string;
+    sender: string;
+}
+
+export {
+    type AIResponseText,
+    type AIResponseImage,
+    type AIToolFunction,
+    type AIToolMatrixData,
+};