Skip to content

Commit

Permalink
Add tool calling (#112)
Browse files Browse the repository at this point in the history
- add tool calling to the api interface
- add a tool calling example
  • Loading branch information
BruceMacD authored Jul 17, 2024
1 parent 2d31f21 commit 59ef1ce
Show file tree
Hide file tree
Showing 2 changed files with 118 additions and 0 deletions.
88 changes: 88 additions & 0 deletions examples/tools/tools.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
import ollama from 'ollama';

// Simulates an API call to get flight times
// In a real application, this would fetch data from a live database or API
function getFlightTimes(departure: string, arrival: string) {
const flights = {
"NYC-LAX": { departure: "08:00 AM", arrival: "11:30 AM", duration: "5h 30m" },
"LAX-NYC": { departure: "02:00 PM", arrival: "10:30 PM", duration: "5h 30m" },
"LHR-JFK": { departure: "10:00 AM", arrival: "01:00 PM", duration: "8h 00m" },
"JFK-LHR": { departure: "09:00 PM", arrival: "09:00 AM", duration: "7h 00m" },
"CDG-DXB": { departure: "11:00 AM", arrival: "08:00 PM", duration: "6h 00m" },
"DXB-CDG": { departure: "03:00 AM", arrival: "07:30 AM", duration: "7h 30m" }
};

const key = `${departure}-${arrival}`.toUpperCase();
return JSON.stringify(flights[key] || { error: "Flight not found" });
}

async function run(model: string) {
// Initialize conversation with a user query
let messages = [{ role: 'user', content: 'What is the flight time from New York (NYC) to Los Angeles (LAX)?' }];

// First API call: Send the query and function description to the model
const response = await ollama.chat({
model: model,
messages: messages,
tools: [
{
type: 'function',
function: {
name: 'get_flight_times',
description: 'Get the flight times between two cities',
parameters: {
type: 'object',
properties: {
departure: {
type: 'string',
description: 'The departure city (airport code)',
},
arrival: {
type: 'string',
description: 'The arrival city (airport code)',
},
},
required: ['departure', 'arrival'],
},
},
},
],
})
// Add the model's response to the conversation history
messages.push(response.message);

// Check if the model decided to use the provided function
if (!response.message.tool_calls || response.message.tool_calls.length === 0) {
console.log("The model didn't use the function. Its response was:");
console.log(response.message.content);
return;
}

// Process function calls made by the model
if (response.message.tool_calls) {
const availableFunctions = {
get_flight_times: getFlightTimes,
};
for (const tool of response.message.tool_calls) {
const functionToCall = availableFunctions[tool.function.name];
const functionResponse = functionToCall(
tool.function.arguments.departure,
tool.function.arguments.arrival
);
// Add function response to the conversation
messages.push({
role: 'tool',
content: functionResponse,
});
}
}

// Second API call: Get final response from the model
const finalResponse = await ollama.chat({
model: model,
messages: messages,
});
console.log(finalResponse.message.content);
}

run('mistral').catch(error => console.error("An error occurred:", error));
30 changes: 30 additions & 0 deletions src/interfaces.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,35 @@ export interface Message {
role: string
content: string
images?: Uint8Array[] | string[]
tool_calls?: ToolCall[]
}

export interface ToolCall {
function: {
name: string;
arguments: {
[key: string]: any;
};
};
}

export interface Tool {
type: string;
function: {
name: string;
description: string;
parameters: {
type: string;
required: string[];
properties: {
[key: string]: {
type: string;
description: string;
enum?: string[];
};
};
};
};
}

export interface ChatRequest {
Expand All @@ -70,6 +99,7 @@ export interface ChatRequest {
stream?: boolean
format?: string
keep_alive?: string | number
tools?: Tool[]

options?: Partial<Options>
}
Expand Down

0 comments on commit 59ef1ce

Please sign in to comment.