Skip to content

Commit 1df5c42

Browse files
authored
Function calling in assistants (#77)
1 parent ef84b37 commit 1df5c42

File tree

14 files changed

+2378
-21
lines changed

14 files changed

+2378
-21
lines changed
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
../function_calling/assistants
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
#!/usr/bin/env python3
2+
3+
from __future__ import annotations
4+
5+
import asyncio
6+
7+
from yandex_cloud_ml_sdk import AsyncYCloudML
8+
9+
10+
def create_tools(sdk: AsyncYCloudML):
11+
# it is imported inside only because yandex-cloud-ml-sdk does not require pydantic by default
12+
# pylint: disable=import-outside-toplevel
13+
from pydantic import BaseModel, Field
14+
15+
calculator_tool = sdk.tools.function(
16+
name="calculator",
17+
description="A simple calculator that performs basic arithmetic operations.",
18+
# parameters could contain valid jsonschema with function parameters types and description
19+
parameters={
20+
"type": "object",
21+
"properties": {
22+
"expression": {
23+
"type": "string",
24+
"description": "The mathematical expression to evaluate (e.g., '2 + 3 * 4').",
25+
}
26+
},
27+
"required": ["expression"],
28+
}
29+
)
30+
31+
class Weather(BaseModel):
32+
"""Getting the weather in specified location for specified date"""
33+
34+
# It is important to describe all the arguments with a natural language
35+
location: str = Field(description="Name of the place for fetching weatcher at")
36+
date: str = Field(description="Date which a user interested in")
37+
38+
weather_tool = sdk.tools.function(Weather)
39+
40+
return [calculator_tool, weather_tool]
41+
42+
43+
def calculator(expression: str) -> str:
44+
print(f'calculator got {expression=}')
45+
return "-5"
46+
47+
48+
def weather(location: str, date: str) -> str:
49+
print(f"weather func got {location=} and {date=}")
50+
return "-10 celsius"
51+
52+
53+
def process_tool_calls(tool_calls) -> list[dict]:
54+
"""
55+
This function is an example how you could organize
56+
dispatching of function calls in general case
57+
"""
58+
59+
function_map = {
60+
'calculator': calculator,
61+
'Weather': weather
62+
}
63+
64+
result = []
65+
for tool_call in tool_calls:
66+
# only functions are available at the moment
67+
assert tool_call.function
68+
69+
function = function_map[tool_call.function.name]
70+
71+
answer = function(**tool_call.function.arguments) # type: ignore[operator]
72+
73+
result.append({'name': tool_call.function.name, 'content': answer})
74+
75+
return result
76+
77+
78+
async def run_on_thread(assistant, thread):
79+
run = await assistant.run_stream(thread)
80+
81+
event = None
82+
async for event in run:
83+
print(
84+
'Stream event from a run - '
85+
f'status={event.status.name!r}, text={event.text!r}, tool_calls={event.tool_calls}'
86+
)
87+
if event.tool_calls:
88+
tool_results = process_tool_calls(event.tool_calls)
89+
await run.submit_tool_results(tool_results)
90+
91+
# Just to show you it is the final event:
92+
assert event
93+
assert event.status.name == 'DONE'
94+
95+
# Final event have text field you could use as a result
96+
return event
97+
98+
99+
async def main() -> None:
100+
sdk = AsyncYCloudML(folder_id='b1ghsjum2v37c2un8h64')
101+
sdk.setup_default_logging()
102+
103+
model = sdk.models.completions('yandexgpt', model_version='rc')
104+
105+
assistant = await sdk.assistants.create(
106+
model,
107+
tools=create_tools(sdk)
108+
)
109+
thread = await sdk.threads.create()
110+
111+
try:
112+
for question in [
113+
"How much it would be 7 * 8?",
114+
"What is the weather like in Paris at 12 of March?"
115+
]:
116+
await thread.write(question)
117+
result = await run_on_thread(assistant, thread)
118+
print(f'Assistant response on question "{question}": {result.text}')
119+
120+
print('All of the thread messages:')
121+
messages = reversed([message async for message in thread])
122+
for message in messages:
123+
print(f' {message.author.role}: {message.text}')
124+
finally:
125+
await thread.delete()
126+
await assistant.delete()
127+
128+
129+
if __name__ == '__main__':
130+
asyncio.run(main())
Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
#!/usr/bin/env python3
2+
3+
from __future__ import annotations
4+
5+
import asyncio
6+
7+
from yandex_cloud_ml_sdk import AsyncYCloudML
8+
9+
10+
def create_tools(sdk: AsyncYCloudML):
11+
# it is imported inside only because yandex-cloud-ml-sdk does not require pydantic by default
12+
# pylint: disable=import-outside-toplevel
13+
from pydantic import BaseModel, Field
14+
15+
calculator_tool = sdk.tools.function(
16+
name="calculator",
17+
description="A simple calculator that performs basic arithmetic operations.",
18+
# parameters could contain valid jsonschema with function parameters types and description
19+
parameters={
20+
"type": "object",
21+
"properties": {
22+
"expression": {
23+
"type": "string",
24+
"description": "The mathematical expression to evaluate (e.g., '2 + 3 * 4').",
25+
}
26+
},
27+
"required": ["expression"],
28+
}
29+
)
30+
31+
class Weather(BaseModel):
32+
"""Getting the weather in specified location for specified date"""
33+
34+
# It is important to describe all the arguments with a natural language
35+
location: str = Field(description="Name of the place for fetching weatcher at")
36+
date: str = Field(description="Date which a user interested in")
37+
38+
weather_tool = sdk.tools.function(Weather)
39+
40+
return [calculator_tool, weather_tool]
41+
42+
43+
def calculator(expression: str) -> str:
44+
print(f'calculator got {expression=}')
45+
return "-5"
46+
47+
48+
def weather(location: str, date: str) -> str:
49+
print(f"weather func got {location=} and {date=}")
50+
return "-10 celsius"
51+
52+
53+
def process_tool_calls(tool_calls) -> list[dict]:
54+
"""
55+
This function is an example how you could organize
56+
dispatching of function calls in general case
57+
"""
58+
59+
function_map = {
60+
'calculator': calculator,
61+
'Weather': weather
62+
}
63+
64+
result = []
65+
for tool_call in tool_calls:
66+
# only functions are available at the moment
67+
assert tool_call.function
68+
69+
function = function_map[tool_call.function.name]
70+
71+
answer = function(**tool_call.function.arguments) # type: ignore[operator]
72+
73+
result.append({'name': tool_call.function.name, 'content': answer})
74+
75+
return result
76+
77+
78+
async def run_on_thread(assistant, thread):
79+
run = await assistant.run(thread)
80+
result = await run
81+
82+
if result.tool_calls:
83+
tool_results = process_tool_calls(result.tool_calls)
84+
await run.submit_tool_results(tool_results)
85+
86+
# run starting to work again after submitting tool results
87+
# and we need to wait it's results again
88+
result = await run
89+
90+
return result
91+
92+
93+
async def main() -> None:
94+
sdk = AsyncYCloudML(folder_id='b1ghsjum2v37c2un8h64')
95+
sdk.setup_default_logging()
96+
97+
model = sdk.models.completions('yandexgpt', model_version='rc')
98+
99+
assistant = await sdk.assistants.create(
100+
model,
101+
tools=create_tools(sdk)
102+
)
103+
thread = await sdk.threads.create()
104+
105+
try:
106+
for question in [
107+
"How much it would be 7 * 8?",
108+
"What is the weather like in Paris at 12 of March?"
109+
]:
110+
await thread.write(question)
111+
result = await run_on_thread(assistant, thread)
112+
print(f'Assistant response on question "{question}": {result.text}')
113+
114+
print('All of the thread messages:')
115+
messages = reversed([message async for message in thread])
116+
for message in messages:
117+
print(f' {message.author.role}: {message.text}')
118+
finally:
119+
await thread.delete()
120+
await assistant.delete()
121+
122+
123+
if __name__ == '__main__':
124+
asyncio.run(main())
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
../function_calling/assistants
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
#!/usr/bin/env python3
2+
3+
from __future__ import annotations
4+
5+
from yandex_cloud_ml_sdk import YCloudML
6+
7+
8+
def create_tools(sdk: YCloudML):
9+
# it is imported inside only because yandex-cloud-ml-sdk does not require pydantic by default
10+
# pylint: disable=import-outside-toplevel
11+
from pydantic import BaseModel, Field
12+
13+
calculator_tool = sdk.tools.function(
14+
name="calculator",
15+
description="A simple calculator that performs basic arithmetic operations.",
16+
# parameters could contain valid jsonschema with function parameters types and description
17+
parameters={
18+
"type": "object",
19+
"properties": {
20+
"expression": {
21+
"type": "string",
22+
"description": "The mathematical expression to evaluate (e.g., '2 + 3 * 4').",
23+
}
24+
},
25+
"required": ["expression"],
26+
}
27+
)
28+
29+
class Weather(BaseModel):
30+
"""Getting the weather in specified location for specified date"""
31+
32+
# It is important to describe all the arguments with a natural language
33+
location: str = Field(description="Name of the place for fetching weatcher at")
34+
date: str = Field(description="Date which a user interested in")
35+
36+
weather_tool = sdk.tools.function(Weather)
37+
38+
return [calculator_tool, weather_tool]
39+
40+
41+
def calculator(expression: str) -> str:
42+
print(f'calculator got {expression=}')
43+
return "-5"
44+
45+
46+
def weather(location: str, date: str) -> str:
47+
print(f"weather func got {location=} and {date=}")
48+
return "-10 celsius"
49+
50+
51+
def process_tool_calls(tool_calls) -> list[dict]:
52+
"""
53+
This function is an example how you could organize
54+
dispatching of function calls in general case
55+
"""
56+
57+
function_map = {
58+
'calculator': calculator,
59+
'Weather': weather
60+
}
61+
62+
result = []
63+
for tool_call in tool_calls:
64+
# only functions are available at the moment
65+
assert tool_call.function
66+
67+
function = function_map[tool_call.function.name]
68+
69+
answer = function(**tool_call.function.arguments) # type: ignore[operator]
70+
71+
result.append({'name': tool_call.function.name, 'content': answer})
72+
73+
return result
74+
75+
76+
def run_on_thread(assistant, thread):
77+
run = assistant.run_stream(thread)
78+
79+
event = None
80+
for event in run:
81+
print(
82+
'Stream event from a run - '
83+
f'status={event.status.name!r}, text={event.text!r}, tool_calls={event.tool_calls}'
84+
)
85+
if event.tool_calls:
86+
tool_results = process_tool_calls(event.tool_calls)
87+
run.submit_tool_results(tool_results)
88+
89+
# Just to show you it is the final event:
90+
assert event
91+
assert event.status.name == 'DONE'
92+
93+
# Final event have text field you could use as a result
94+
return event
95+
96+
97+
def main() -> None:
98+
sdk = YCloudML(folder_id='b1ghsjum2v37c2un8h64')
99+
sdk.setup_default_logging()
100+
101+
model = sdk.models.completions('yandexgpt', model_version='rc')
102+
103+
assistant = sdk.assistants.create(
104+
model,
105+
tools=create_tools(sdk)
106+
)
107+
thread = sdk.threads.create()
108+
109+
try:
110+
for question in [
111+
"How much it would be 7 * 8?",
112+
"What is the weather like in Paris at 12 of March?"
113+
]:
114+
thread.write(question)
115+
result = run_on_thread(assistant, thread)
116+
print(f'Assistant response on question "{question}": {result.text}')
117+
118+
print('All of the thread messages:')
119+
messages = reversed(list(thread))
120+
for message in messages:
121+
print(f' {message.author.role}: {message.text}')
122+
finally:
123+
thread.delete()
124+
assistant.delete()
125+
126+
127+
if __name__ == '__main__':
128+
main()

0 commit comments

Comments
 (0)