Updates tool managers various states with new information
(&mut self)
| 927 | |
| 928 | /// Updates tool managers various states with new information |
| 929 | pub async fn update(&mut self) { |
| 930 | // A hashmap of <tool name, tool spec> |
| 931 | let mut tool_specs = HashMap::<String, ToolSpec>::new(); |
| 932 | let new_tools = { |
| 933 | let mut new_tool_specs = self.new_tool_specs.lock().await; |
| 934 | new_tool_specs.drain().fold( |
| 935 | HashMap::<ServerName, (HashMap<ModelToolName, ToolInfo>, Vec<ToolSpec>)>::new(), |
| 936 | |mut acc, (server_name, v)| { |
| 937 | acc.insert(server_name, v); |
| 938 | acc |
| 939 | }, |
| 940 | ) |
| 941 | }; |
| 942 | |
| 943 | let mut updated_servers = HashSet::<ToolOrigin>::new(); |
| 944 | let mut conflicts = HashMap::<ServerName, String>::new(); |
| 945 | for (server_name, (tool_name_map, specs)) in new_tools { |
| 946 | // First we evict the tools that were already in the tn_map |
| 947 | self.tn_map.retain(|_, tool_info| tool_info.server_name != server_name); |
| 948 | |
| 949 | // And update them with the new tools queried |
| 950 | // valid: tools that do not have conflicts in naming |
| 951 | let (valid, invalid) = tool_name_map |
| 952 | .into_iter() |
| 953 | .partition::<HashMap<ModelToolName, ToolInfo>, _>(|(model_tool_name, _)| { |
| 954 | !self.tn_map.contains_key(model_tool_name) |
| 955 | }); |
| 956 | // We reject tools that are conflicting with the existing tools by not including them |
| 957 | // in the tn_map. We would also want to report this error. |
| 958 | if !invalid.is_empty() { |
| 959 | let msg = invalid.into_iter().fold("The following tools are rejected because they conflict with existing tools in names. Avoid this via setting aliases for them: \n".to_string(), |mut acc, (model_tool_name, tool_info)| { |
| 960 | acc.push_str(&format!(" - {} from {}\n", model_tool_name, tool_info.server_name)); |
| 961 | acc |
| 962 | }); |
| 963 | conflicts.insert(server_name, msg); |
| 964 | } |
| 965 | if let Some(spec) = specs.first() { |
| 966 | updated_servers.insert(spec.tool_origin.clone()); |
| 967 | } |
| 968 | // We want to filter for specs that are valid |
| 969 | // Note that [ToolSpec::name] is a model facing name (thus you should be comparing it |
| 970 | // with the keys of a tn_map) |
| 971 | for spec in specs.into_iter().filter(|spec| valid.contains_key(&spec.name)) { |
| 972 | tool_specs.insert(spec.name.clone(), spec); |
| 973 | } |
| 974 | |
| 975 | self.tn_map.extend(valid); |
| 976 | } |
| 977 | |
| 978 | // Update schema |
| 979 | // As we are writing over the ensemble of tools in a given server, we will need to first |
| 980 | // remove everything that it has. |
| 981 | self.schema |
| 982 | .retain(|_tool_name, spec| !updated_servers.contains(&spec.tool_origin)); |
| 983 | self.schema.extend(tool_specs); |
| 984 | |
| 985 | // if block here to avoid repeatedly asking for loc |
| 986 | if !conflicts.is_empty() { |
no test coverage detected