Refactor AI function handling by introducing AIFunctionCall and AIFunctionResponse classes, updating function definitions, and enhancing conversation content management. Remove unused code and improve service registration for AIFunctionService.

This commit is contained in:
Andrew Beal 2025-10-06 23:01:20 +01:00
parent 35b531d9bf
commit 62104f0abd
18 changed files with 321 additions and 122 deletions

View file

@ -0,0 +1,19 @@
// platform agnostic function call class used to execute the requested function
export class AIFunctionCall {
public readonly name: string;
public readonly arguments: Record<string, any>;
constructor(name: string, args: Record<string, any>) {
this.name = name;
this.arguments = args;
}
public toConversationString(): string {
return JSON.stringify({
functionCall: {
name: this.name,
args: this.arguments
}
});
}
}

View file

@ -0,0 +1,8 @@
import type { IAIFunctionDefinition } from "./IAIFunctionDefinition";
import { ListVaultFiles } from "./ListVaultFiles";
export class AIFunctionDefinitions {
public getQueryActions(): IAIFunctionDefinition[] {
return [ListVaultFiles];
}
}

View file

@ -0,0 +1,20 @@
// Platform agnostic class for function responses
// Used by AI providers to format function execution results for API calls
export class AIFunctionResponse {
public readonly name: string;
public readonly response: object;
constructor(name: string, response: object) {
this.name = name;
this.response = response;
}
public toConversationString(): string {
return JSON.stringify({
functionResponse: {
name: this.name,
response: { result: this.response }
}
});
}
}

View file

@ -1,4 +1,5 @@
export interface FunctionDefinition {
// platform agnostic function definition used to present function calls in an API call
export interface IAIFunctionDefinition {
name: string;
description: string;
parameters: {

View file

@ -1,7 +1,7 @@
import { AIFunction } from "Enums/AIFunction";
import type { FunctionDefinition } from "./AIFunctionDefinition";
import type { IAIFunctionDefinition } from "./IAIFunctionDefinition";
export const ListVaultFiles: FunctionDefinition = {
export const ListVaultFiles: IAIFunctionDefinition = {
name: AIFunction.ListVaultFiles,
description: `Returns complete list of vault files with metadata (names, paths, sizes).
Call this whenever you need to know what files exist in the vault to answer questions,
@ -9,6 +9,11 @@ export const ListVaultFiles: FunctionDefinition = {
when vault contents would inform your response.`,
parameters: {
type: "object",
properties: {}
properties: {
user_message: {
type: "string",
description: "A short message to be displayed to the user that explains the action being taken"
}
}
}
}

View file

@ -6,76 +6,146 @@ import { StreamingService, type StreamChunk } from "Services/StreamingService";
import type { Conversation } from "Conversations/Conversation";
import { Role } from "Enums/Role";
import { AIProviderURL } from "Enums/ApiProvider";
import { AIFunctionCall } from "AIClasses/AIFunctionCall";
import type { AIFunctionDefinitions } from "AIClasses/FunctionDefinitions/AIFunctionDefinitions";
import type { IAIFunctionDefinition } from "AIClasses/FunctionDefinitions/IAIFunctionDefinition";
export class Gemini implements IAIClass {
private readonly REQUEST_WEB_SEARCH: string = "request_web_search";
private readonly apiKey: string;
private readonly aiPrompt: IPrompt;
private readonly streamingService: StreamingService;
private readonly aiFunctionDefinitions: AIFunctionDefinitions;
private accumulatedFunctionName: string | null = null;
private accumulatedFunctionArgs: Record<string, any> = {};
public constructor(apiKey: string) {
this.apiKey = apiKey;
this.aiPrompt = Resolve(Services.IPrompt);
this.streamingService = Resolve(Services.StreamingService);
this.aiFunctionDefinitions = Resolve(Services.AIFunctionDefinitions);
}
public async* streamRequest(conversation: Conversation): AsyncGenerator<StreamChunk, void, unknown> {
// next request should use web search only (gemini api doesn't support custom tooling and grounding at the same time)
let requestWebSearch = this.accumulatedFunctionName == this.REQUEST_WEB_SEARCH;
// Reset function call accumulation state for new request
this.accumulatedFunctionName = null;
this.accumulatedFunctionArgs = {};
const contents = conversation.contents.map(content => ({
role: content.role === Role.User ? "user" : "model",
parts: [
{
text: content.content
}
]
parts: (content.isFunctionCall || content.isFunctionCallResponse)
? [JSON.parse(content.content)] : [{ text: content.content }]
}));
const tools = requestWebSearch ? { google_search: {} } :
{
functionDeclarations: [
{
name: "request_web_search",
description: `Use this function when you need to search the web for current
information, recent events, news, or facts that may have changed.
After calling this, you will be able to perform web searches.`,
parameters: {
type: "object",
properties: {
reasoning: {
type: "string",
description: "Brief explanation of why web search is needed"
}
},
required: ["reasoning"]
}
},
...this.mapFunctionDefinitions(this.aiFunctionDefinitions.getQueryActions()),
]
}
const requestBody = {
system_instruction: {
parts: [
{
text: this.aiPrompt.systemInstruction()
},
{
text: `IMPORTANT: When you need current information from the web (recent events, news, current prices, weather, etc.), you should:
1. First call the 'request_web_search' function to indicate you need web access
2. After that, you'll be given access to Google Search
3. Once you have the information from the search, you can answer the user's question
4. Subsequent communication will return to providing custom function calls`
},
{
text: await this.aiPrompt.userInstruction()
}
]
},
contents: contents,
tools: [
{
google_search: {},
functionDeclarations: [],
},
]
tools: [tools]
};
console.log(requestBody);
yield* this.streamingService.streamRequest(
AIProviderURL.Gemini.replace("API_KEY", this.apiKey),
requestBody,
this.parseStreamChunk
this.parseStreamChunk.bind(this)
);
}
private parseStreamChunk(chunk: string): StreamChunk {
try {
const data = JSON.parse(chunk);
let text = "";
let functionCall: AIFunctionCall | undefined = undefined;
const candidate = data.candidates?.[0];
if (candidate) {
if (candidate.content?.parts?.[0]?.text) {
text = candidate.content.parts[0].text;
} else if (candidate.text) {
text = candidate.text;
// Check for text content
if (candidate.content?.parts?.[0]?.text) {
text = candidate.content.parts[0].text;
} else if (candidate.text) {
text = candidate.text;
}
// Check for function call and accumulate
const parts = candidate.content?.parts || [];
for (const part of parts) {
if (part.functionCall) {
// Accumulate function name
if (part.functionCall.name) {
this.accumulatedFunctionName = part.functionCall.name;
}
// Accumulate function arguments (merge with existing)
if (part.functionCall.args) {
this.accumulatedFunctionArgs = {
...this.accumulatedFunctionArgs,
...part.functionCall.args
};
}
break; // Only handle first function call per chunk
}
}
}
const isComplete = !!candidate?.finishReason;
// If streaming is complete and we have accumulated a function call, return it
if (isComplete && this.accumulatedFunctionName) {
functionCall = new AIFunctionCall(
this.accumulatedFunctionName,
this.accumulatedFunctionArgs
);
}
return {
content: text,
isComplete: isComplete,
functionCall: functionCall,
};
} catch (error) {
const message = error instanceof Error ? error.message : "Unknown parsing error";
@ -83,4 +153,12 @@ export class Gemini implements IAIClass {
return { content: "", isComplete: false, error: `Failed to parse chunk: ${message}` };
}
}
private mapFunctionDefinitions(aiFunctionDefinitions: IAIFunctionDefinition[]): object[] {
return aiFunctionDefinitions.map((functionDefinition) => ({
name: functionDefinition.name,
description: functionDefinition.description,
parameters: functionDefinition.parameters
}));
}
}

View file

@ -1,6 +0,0 @@
interface CreateFileRequest {
file_path: string;
file_content: string;
}

View file

@ -54,6 +54,7 @@ You are also capable of helping with:
- Simple conversational exchanges
**Best practices:**
- Always provide a user friendly description when using a tool in addition to the function object
- Use tools proactively when vault content would enhance your answer
- Offer to search rather than assuming - e.g., "Would you like me to check your notes for related information?"
- Don't over-use tools for simple questions that don't require vault access

View file

@ -163,28 +163,30 @@
<div class="chat-area" bind:this={chatContainer} on:scroll={handleScroll}>
{#each messages as message, messageIndex (`${message.role}-${messageIndex}`)}
{#if message.role === Role.User}
<div class="message-container user">
<div class="message-bubble user">
<p class="message-text-user fade-in-fast">{message.content}</p>
</div>
</div>
{:else}
<div class="message-container assistant">
<div class="message-bubble assistant">
<div class="markdown-content fade-in-fast" class:streaming={isStreaming && messageIndex === messages.length - 1}>
<!-- Streaming message: use action for initialization -->
{#if isStreaming && messageIndex === messages.length - 1}
<div use:streamingAction={`${message.role}-${messageIndex}`} class="streaming-content"></div>
<StreamingIndicator/>
<ChatAreaThought/>
{:else}
<!-- Static message: use traditional rendering -->
{@html getStaticHTML(message, messageIndex)}
{/if}
{#if !message.isFunctionCall && !message.isFunctionCallResponse}
{#if message.role === Role.User}
<div class="message-container user">
<div class="message-bubble user">
<p class="message-text-user fade-in-fast">{message.content}</p>
</div>
</div>
</div>
{:else}
<div class="message-container assistant">
<div class="message-bubble assistant">
<div class="markdown-content fade-in-fast" class:streaming={isStreaming && messageIndex === messages.length - 1}>
<!-- Streaming message: use action for initialization -->
{#if isStreaming && messageIndex === messages.length - 1}
<div use:streamingAction={`${message.role}-${messageIndex}`} class="streaming-content"></div>
<StreamingIndicator/>
<ChatAreaThought/>
{:else}
<!-- Static message: use traditional rendering -->
{@html getStaticHTML(message, messageIndex)}
{/if}
</div>
</div>
</div>
{/if}
{/if}
{/each}

View file

@ -11,9 +11,13 @@
import { Role } from "Enums/Role";
import { Conversation } from "Conversations/Conversation";
import { ConversationContent } from "Conversations/ConversationContent";
import type { AIFunctionCall } from "AIClasses/AIFunctionCall";
import type { AIFunctionService } from "Services/AIFunctionService";
import type { AIFunctionResponse } from "AIClasses/FunctionDefinitions/AIFunctionResponse";
let ai: IAIClass = Resolve(Services.IAIClass);
let conversationService: ConversationFileSystemService = Resolve(Services.ConversationFileSystemService);
let aiFunctionService: AIFunctionService = Resolve(Services.AIFunctionService);
let semaphore: Semaphore = new Semaphore(1, false);
let textareaElement: HTMLTextAreaElement;
@ -27,76 +31,103 @@
let conversation = new Conversation();
async function handleSubmit() {
if (userRequest.trim() === "" || isSubmitting) {
return;
}
isSubmitting = true;
const requestToSend = userRequest;
userRequest = "";
textareaElement.value = "";
autoResize();
if (!await semaphore.wait()) {
return;
}
// Add user message to chat
conversation.contents = [...conversation.contents, new ConversationContent(Role.User, requestToSend)];
await conversationService.saveConversation(conversation);
scrollToBottom();
try {
// Create AI message placeholder
const aiMessageIndex = conversation.contents.length;
conversation.contents = [...conversation.contents, new ConversationContent(Role.Assistant, "")];
isStreaming = true;
if (userRequest.trim() === "" || isSubmitting) {
return;
}
isSubmitting = true;
// Stream the response
let accumulatedContent = "";
conversation.contents = [...conversation.contents, new ConversationContent(Role.User, userRequest)];
await conversationService.saveConversation(conversation);
for await (const chunk of ai.streamRequest(conversation)) {
if (chunk.error) {
console.error("Streaming error:", chunk.error);
// Update message with error
conversation.contents = conversation.contents.map((msg, messageIndex) =>
messageIndex === aiMessageIndex
? { ...msg, content: "Error: " + chunk.error }
: msg
);
isStreaming = false;
break;
}
textareaElement.value = "";
userRequest = "";
autoResize();
if (chunk.content) {
accumulatedContent += chunk.content;
// Update the message with accumulated content
conversation.contents = conversation.contents.map((msg, messageIndex) =>
messageIndex === aiMessageIndex
? { ...msg, content: accumulatedContent }
: msg
);
}
scrollToBottom();
if (chunk.isComplete) {
// Mark streaming as complete
isStreaming = false;
conversation.contents = conversation.contents.map((msg, messageIndex) =>
messageIndex === aiMessageIndex
? { ...msg, content: accumulatedContent }
: msg
);
let functionCall: AIFunctionCall | null = await streamRequestResponse();
while (functionCall) {
if ('user_message' in functionCall.arguments) {
conversation.contents = [...conversation.contents, new ConversationContent(
Role.Assistant, functionCall.arguments.user_message)];
await conversationService.saveConversation(conversation);
}
conversation.contents = [...conversation.contents, new ConversationContent(
Role.Assistant, functionCall.toConversationString(), new Date(), true)];
await conversationService.saveConversation(conversation);
const functionResponse: AIFunctionResponse = await aiFunctionService.performAIFunction(functionCall);
conversation.contents = [...conversation.contents, new ConversationContent(
Role.User, functionResponse.toConversationString(), new Date(), false, true)];
await conversationService.saveConversation(conversation);
functionCall = await streamRequestResponse();
}
} finally {
semaphore.release();
isSubmitting = false;
tick().then(() => {
textareaElement?.focus();
});
}
}
async function streamRequestResponse(): Promise<AIFunctionCall | null> {
// Create AI message placeholder
const aiMessageIndex = conversation.contents.length;
conversation.contents = [...conversation.contents, new ConversationContent(Role.Assistant, "")];
isStreaming = true;
let accumulatedContent = "";
let capturedFunctionCall: AIFunctionCall | null = null;
for await (const chunk of ai.streamRequest(conversation)) {
if (chunk.error) {
console.error("Streaming error:", chunk.error);
conversation.contents = conversation.contents.map((msg, messageIndex) =>
messageIndex === aiMessageIndex
? { ...msg, content: "Error: " + chunk.error }
: msg
);
isStreaming = false;
break;
}
if (chunk.content) {
accumulatedContent += chunk.content;
conversation.contents = conversation.contents.map((msg, messageIndex) =>
messageIndex === aiMessageIndex
? { ...msg, content: accumulatedContent }
: msg
);
}
if (chunk.functionCall) {
capturedFunctionCall = chunk.functionCall;
}
if (chunk.isComplete) {
isStreaming = false;
conversation.contents = conversation.contents.map((msg, messageIndex) =>
messageIndex === aiMessageIndex
? { ...msg, content: accumulatedContent }
: msg
);
await conversationService.saveConversation(conversation);
}
}
return capturedFunctionCall;
}
function handleKeydown(e: KeyboardEvent) {
if (e.key === 'Enter') {
if (e.shiftKey) {
@ -161,10 +192,10 @@
rows="1">
</textarea>
<button
<button
id="submit"
bind:this={submitButton}
on:click={handleSubmit}
on:click={() => { handleSubmit() }}
disabled={isSubmitting || userRequest.trim() === ""}
aria-label="Send Message">
</button>
@ -213,6 +244,12 @@
resize: none;
overflow-y: auto;
color: var(--font-interface-theme);
transition: border-color 0.3s ease-out;
}
#input:focus {
border-color: var(--color-accent);
transition: border-color 0.3s ease-out;
}
#submit {

View file

@ -2,6 +2,8 @@ export class ConversationContent {
role: string;
content: string
timestamp: Date;
isFunctionCall: boolean;
isFunctionCallResponse: boolean;
public static isConversationContentData(data: unknown): data is { role: string; content: string; timestamp: string } {
return (
@ -10,15 +12,21 @@ export class ConversationContent {
'role' in data &&
'content' in data &&
'timestamp' in data &&
'isFunctionCall' in data &&
'isFunctionCallResponse' in data &&
typeof data.role === 'string' &&
typeof data.content === 'string' &&
typeof data.timestamp === 'string'
typeof data.timestamp === 'string' &&
typeof data.isFunctionCall == 'boolean' &&
typeof data.isFunctionCallResponse == 'boolean'
);
}
constructor(role: string, content: string) {
constructor(role: string, content: string, timestamp: Date = new Date(), isFunctionCall = false, isFunctionCallResponse = false) {
this.role = role;
this.content = content;
this.timestamp = new Date();
this.timestamp = timestamp;
this.isFunctionCall = isFunctionCall;
this.isFunctionCallResponse = isFunctionCallResponse;
}
}

View file

@ -33,7 +33,9 @@ export class ConversationHistoryModal extends Modal {
override async open() {
this.conversations = await this.conversationFileSystemService.getAllConversations();
this.items = this.conversations.map((conversation, index) => ({
this.items = this.conversations
.sort((a, b) => b.created.getTime() - a.created.getTime())
.map((conversation, index) => ({
id: index.toString(),
date: dateToString(conversation.created, false),
title: conversation.title,

View file

@ -1,6 +1,6 @@
<script lang="ts">
import { setIcon } from "obsidian";
import { fade, slide } from "svelte/transition";
import { fade } from "svelte/transition";
export let items: Array<{id: string, date: string, title: string, selected: boolean}>;
export let onClose: () => void;
@ -18,7 +18,7 @@
}
let selectedItems = new Set<string>();
let searchQuery = '';
let searchQuery = "";
$: filteredItems = items.filter(item =>
item.title.toLowerCase().includes(searchQuery.toLowerCase())
@ -112,6 +112,7 @@
display: grid;
grid-template-rows: var(--size-4-3) auto var(--size-4-3) 1fr var(--size-4-3);
grid-template-columns: var(--size-4-3) 1fr var(--size-4-3);
max-height: 60vh;
}
.conversation-history-modal-top-bar {
@ -210,7 +211,8 @@
#delete-button {
grid-row: 1;
grid-column: 2;
transition: width 300ms ease-out, opacity 300ms ease-out, padding 300ms ease-out;
margin-right: var(--size-4-2);
transition: width 300ms ease-out, opacity 300ms ease-out, padding 300ms ease-out, margin 300ms ease-out;
}
#delete-button.hidden {
@ -218,6 +220,7 @@
min-width: 0;
padding: 0;
opacity: 0;
margin: 0;
overflow: hidden;
pointer-events: none;
}
@ -225,7 +228,6 @@
#close-button {
grid-row: 1;
grid-column: 5;
margin-right: var(--size-4-2);
}
#search-input {

View file

@ -2,12 +2,32 @@ import { Resolve } from "./DependencyService";
import { Services } from "./Services";
import type { FileSystemService } from "./FileSystemService";
import { TFile } from "obsidian";
import { AIFunction } from "Enums/AIFunction";
import { AIFunctionResponse } from "AIClasses/FunctionDefinitions/AIFunctionResponse";
import type { AIFunctionCall } from "AIClasses/AIFunctionCall";
export class AIFunctionService {
private fileSystemService: FileSystemService = Resolve(Services.FileSystemService);
public async listVaultFiles(): Promise<object[]> {
public async performAIFunction(functionCall: AIFunctionCall): Promise<AIFunctionResponse> {
switch (functionCall.name) {
case AIFunction.ListVaultFiles:
return new AIFunctionResponse(
functionCall.name,
await this.listVaultFiles()
);
default:
const error = `Unknown function request ${functionCall.name}`
console.error(error);
return new AIFunctionResponse(
functionCall.name,
{ error: error }
);
}
}
private async listVaultFiles(): Promise<object> {
const files: TFile[] = await this.fileSystemService.listFilesInDirectory("/");
return files.map((file) => ({
name: file.basename,

View file

@ -18,11 +18,6 @@ export class ConversationFileSystemService {
return `${Path.Conversations}/${conversation.title}.json`;
}
// public async loadConversation(filePath: string): Promise<{ messages:
// Array<{ role: string, content: string, timestamp: Date }>, created: Date, title: string } | null> {
// const conversation: TFile | null = this.fileSystemService.readObjectFromFile(filePath) as unknown as TFile;
// }
public async saveConversation(conversation: Conversation): Promise<string> {
if (!this.currentConversationPath) {
this.currentConversationPath = this.generateConversationPath(conversation);
@ -34,7 +29,9 @@ export class ConversationFileSystemService {
contents: conversation.contents.map(content => ({
role: content.role,
content: content.content,
timestamp: content.timestamp.toISOString()
timestamp: content.timestamp.toISOString(),
isFunctionCall: content.isFunctionCall,
isFunctionCallResponse: content.isFunctionCallResponse
}))
};
@ -79,9 +76,8 @@ export class ConversationFileSystemService {
conversation.title = data.title;
conversation.created = new Date(data.created);
conversation.contents = data.contents.map(content => {
const conversationContent = new ConversationContent(content.role, content.content);
conversationContent.timestamp = new Date(content.timestamp);
return conversationContent;
return new ConversationContent(
content.role, content.content, new Date(content.timestamp), content.isFunctionCall);
});
conversations.push(conversation);
}

View file

@ -14,6 +14,7 @@ import { ConversationHistoryModal } from "Modals/ConversationHistoryModal";
import type { App } from "obsidian";
import { AIFunctionService } from "./AIFunctionService";
import { StreamingService } from "./StreamingService";
import { AIFunctionDefinitions } from "AIClasses/FunctionDefinitions/AIFunctionDefinitions";
export function RegisterDependencies(plugin: AIAgentPlugin) {
RegisterSingleton(Services.MessageService, new MessageService());
@ -23,6 +24,7 @@ export function RegisterDependencies(plugin: AIAgentPlugin) {
RegisterSingleton(Services.ConversationFileSystemService, new ConversationFileSystemService());
RegisterSingleton<IPrompt>(Services.IPrompt, new AIPrompt());
RegisterSingleton<AIFunctionDefinitions>(Services.AIFunctionDefinitions, new AIFunctionDefinitions());
RegisterSingleton<AIFunctionService>(Services.AIFunctionService, new AIFunctionService());
RegisterSingleton<StreamingService>(Services.StreamingService, new StreamingService());

View file

@ -7,6 +7,7 @@ export class Services {
static StreamingService = Symbol("StreamingService");
static MarkdownService = Symbol("MarkdownService");
static StreamingMarkdownService = Symbol("StreamingMarkdownService");
static AIFunctionDefinitions = Symbol("AIFunctionDefinitions");
static AIFunctionService = Symbol("AIFunctionService");
// interfaces

View file

@ -1,7 +1,10 @@
import type { AIFunctionCall } from "AIClasses/AIFunctionCall";
export interface StreamChunk {
content: string;
isComplete: boolean;
error?: string;
functionCall?: AIFunctionCall;
}
export class StreamingService {