1
1
use futures:: StreamExt ;
2
2
use rig:: {
3
3
agent:: Agent ,
4
+ completion:: { AssistantContent , CompletionModel } ,
4
5
message:: Message ,
5
- streaming:: { StreamingChat , StreamingChoice , StreamingCompletionModel } ,
6
+ streaming:: StreamingChat ,
6
7
} ;
7
8
use tokio:: io:: { AsyncBufReadExt , AsyncWriteExt , BufReader , BufWriter } ;
8
9
9
- pub async fn cli_chatbot < M : StreamingCompletionModel > ( chatbot : Agent < M > ) -> anyhow:: Result < ( ) > {
10
+ pub async fn cli_chatbot < M > ( chatbot : Agent < M > ) -> anyhow:: Result < ( ) >
11
+ where
12
+ M : CompletionModel ,
13
+ {
10
14
let mut chat_log = vec ! [ ] ;
11
15
12
16
let mut output = BufWriter :: new ( tokio:: io:: stdout ( ) ) ;
@@ -32,15 +36,17 @@ pub async fn cli_chatbot<M: StreamingCompletionModel>(chatbot: Agent<M>) -> anyh
32
36
let mut message_buf = String :: new ( ) ;
33
37
while let Some ( message) = response. next ( ) . await {
34
38
match message {
35
- Ok ( StreamingChoice :: Message ( text) ) => {
36
- message_buf. push_str ( & text) ;
37
- output_agent ( text, & mut output) . await ?;
39
+ Ok ( AssistantContent :: Text ( text) ) => {
40
+ message_buf. push_str ( & text. text ) ;
41
+ output_agent ( text. text , & mut output) . await ?;
38
42
}
39
- Ok ( StreamingChoice :: ToolCall ( name, _, param) ) => {
43
+ Ok ( AssistantContent :: ToolCall ( tool_call) ) => {
44
+ let name = tool_call. function . name ;
45
+ let arguments = tool_call. function . arguments ;
40
46
chat_log. push ( Message :: assistant ( format ! (
41
- "Calling tool: {name} with args: {param }"
47
+ "Calling tool: {name} with args: {arguments }"
42
48
) ) ) ;
43
- let result = chatbot. tools . call ( & name, param . to_string ( ) ) . await ;
49
+ let result = chatbot. tools . call ( & name, arguments . to_string ( ) ) . await ;
44
50
match result {
45
51
Ok ( tool_call_result) => {
46
52
stream_output_agent_finished ( & mut output) . await ?;
0 commit comments