11use futures:: StreamExt ;
22use rig:: {
33 agent:: Agent ,
4+ completion:: { AssistantContent , CompletionModel } ,
45 message:: Message ,
5- streaming:: { StreamingChat , StreamingChoice , StreamingCompletionModel } ,
6+ streaming:: StreamingChat ,
67} ;
78use tokio:: io:: { AsyncBufReadExt , AsyncWriteExt , BufReader , BufWriter } ;
89
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+ {
1014 let mut chat_log = vec ! [ ] ;
1115
1216 let mut output = BufWriter :: new ( tokio:: io:: stdout ( ) ) ;
@@ -32,15 +36,17 @@ pub async fn cli_chatbot<M: StreamingCompletionModel>(chatbot: Agent<M>) -> anyh
3236 let mut message_buf = String :: new ( ) ;
3337 while let Some ( message) = response. next ( ) . await {
3438 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 ?;
3842 }
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 ;
4046 chat_log. push ( Message :: assistant ( format ! (
41- "Calling tool: {name} with args: {param }"
47+ "Calling tool: {name} with args: {arguments }"
4248 ) ) ) ;
43- let result = chatbot. tools . call ( & name, param . to_string ( ) ) . await ;
49+ let result = chatbot. tools . call ( & name, arguments . to_string ( ) ) . await ;
4450 match result {
4551 Ok ( tool_call_result) => {
4652 stream_output_agent_finished ( & mut output) . await ?;
0 commit comments