Tool retrieval 2

[WonderPG]

Benchmark

Query speed (from sending to stop streaming)

This PR

• Hi: ~4-5s
• Find one morphology in the thalamus, plot it, get its assets and download the swc one.: ~37s
• Show me papers about neuron morphologies in the thalamus of rodents: ~16s

Main

• Hi: ~8s
• Find one morphology in the thalamus, plot it, get its assets and download the swc one.: ~42s (the download link was wrong because way shorter than it should have been, tried twice)
• Show me papers about neuron morphologies in the thalamus of rodents: ~20s

In general the latency is significantly reduced with this approach.

Token count

This PR

Currently the selection LLM has a system prompt with 7,545 tokens. After selection, the main LLM has roughly 10,000 tokens worth of tool description when selecting 10 tools.
This indeed depends on the selected tools and the number of tools. Add to that the main LLM's system prompt worth 1,164 tokens for a total of ~19,000 tokens. Bear in mind that the main LLM only has ~11k tokens injected per request since the tokens of the selection model don't make it to the main model, i.e. 7k tokens are cheaper and run faster while the remaining 11-12k are more expensive and a bit slower.

Main

Currently we have 88,022 tokens from the tools and 1,164 tokens coming from the system prompt. In total we have a minimum of 89,186 tokens injected in every query.

In general the cost goes down by a factor 8.

The following points are also to be considered:

• I observed that the LLM follows better instructions with reduced amount of token per request.
• This PR introduces a risk of not having the relevant tools make it to the main LLM. I rarely had that happen but it is a possibility. Ideally we should experiment and strike the right balance by selecting: The best intelligence/latency model for the task and the minimum amount of tools that the selection model should output everytime.
• I am currently putting the tool name and tool description into the system prompt of the selection model. More things could be added for improved performance, feel free to challenge this.

[WonderPG]

This could be turned into an env var but we have enough already. Please let me know if you would prefer an env var.

[BoBer78]

One comment, the selection LLM has some trouble to select the exa crawling tool. Might want to give it a better description.

[jankrepl]

One comment, the selection LLM has some trouble to select the exa crawling tool. Might want to give it a better description.

Great point. Not sure if directly applicable in this PR but I added a comment here: #415

[WonderPG]

Is it okay if we address this in #415 ?

[jankrepl]

Works perfectly:) Thank you @WonderPG