-
Notifications
You must be signed in to change notification settings - Fork 9
/
agent.ts
108 lines (97 loc) · 3.19 KB
/
agent.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import {
OpenAI,
FunctionTool,
OpenAIAgent,
Settings,
LlamaParseReader,
HuggingFaceEmbedding,
VectorStoreIndex,
QueryEngineTool
} from "llamaindex"
import 'dotenv/config'
async function main() {
// set LLM and the embedding model
Settings.llm = new OpenAI({
apiKey: process.env.OPENAI_API_KEY,
model: "gpt-4o",
})
Settings.embedModel = new HuggingFaceEmbedding({
modelType: "BAAI/bge-small-en-v1.5",
quantized: false
})
/*
Set up logging so we can see the work in progress.
Available events:
llm-start
llm-end
agent-start
agent-end
llm-tool-call
llm-tool-result
*/
Settings.callbackManager.on("llm-tool-call", (event) => {
console.log(event.detail.payload)
})
Settings.callbackManager.on("llm-tool-result", (event) => {
console.log(event.detail.payload)
})
// load our data and create a query engine
const reader = new LlamaParseReader({ resultType: "text" });
const documents = await reader.loadData("../data/sf_budget_2023_2024.pdf");
const index = await VectorStoreIndex.fromDocuments(documents)
const retriever = await index.asRetriever()
retriever.similarityTopK = 10
const queryEngine = await index.asQueryEngine({
retriever
})
// define a function to sum up numbers
const sumNumbers = ({ a, b }) => {
return `${a + b}`;
}
// define the query engine as a tool
const tools = [
new QueryEngineTool({
queryEngine: queryEngine,
metadata: {
name: "san_francisco_budget_tool",
description: `This tool can answer detailed questions about the individual components of the budget of San Francisco in 2023-2024.`,
},
}),
FunctionTool.from(
sumNumbers,
{
name: "sumNumbers",
description: "Use this function to sum two numbers",
parameters: {
type: "object",
properties: {
a: {
type: "number",
description: "First number to sum"
},
b: {
type: "number",
description: "Second number to sum"
},
},
required: ["a", "b"]
}
}
)
]
// create the agent
const agent = new OpenAIAgent({ tools })
let response = await agent.chat({
message: "What's the budget of San Francisco for community health in 2023-24?",
})
console.log(response)
let response2 = await agent.chat({
message: "What's the budget of San Francisco for public protection in 2023-24?",
})
console.log(response2)
let response3 = await agent.chat({
message: "What's the combined budget of San Francisco for community health and public protection in 2023-24?",
})
console.log(response3)
}
main().catch(console.error);