Views
No views yet

[!NOTE] For evaluation, please note that all models were evaluated with minimal reasoning to ensure routing remains efficient.
1import json
2import torch
3
4from transformers import AutoTokenizer, AutoModelForCausalLM
5
6
7ORCHESTRATION_PROMPT = (
8 "You are a helpful assistant that selects the most suitable routes based on user intent.\n"
9 "You are provided with a list of available routes enclosed within <routes></routes> XML tags:\n"
10 "<routes>\n{routes}\n</routes>\n\n"
11 "You are also given the conversation context enclosed within <conversation></conversation> XML tags:\n"
12 "<conversation>\n{conversation}\n</conversation>\n\n"
13 "## Instructions\n"
14 "1. Analyze the latest user intent from the conversation.\n"
15 "2. Compare it against the available routes to find which routes can help fulfill the request.\n"
16 "3. Respond only with the exact route names from <routes>.\n"
17 "4. If no routes can help or the intent is already fulfilled, return an empty list.\n\n"
18 "## Response Format\n"
19 "Return your answer strictly in JSON as follows:\n"
20 '{{"route": ["route_name_1", "route_name_2", "..."]}}\n'
21 "If no routes are needed, return an empty list for `route`."
22)
23
24def convert_agents_to_routes(agents):
25 tools = [
26 {
27 "name": agent["name"],
28 "description": agent["description"],
29 }
30 for agent in agents
31 ]
32 return "\n".join([json.dumps(tool, ensure_ascii=False) for tool in tools])
33
34def build_messages(available_agents, conversation):
35 routes = convert_agents_to_routes(available_agents)
36 conversation_str = json.dumps(conversation, indent=4, ensure_ascii=False)
37 prompt = ORCHESTRATION_PROMPT.format(routes=routes, conversation=conversation_str)
38 return [{"role": "user", "content": prompt}]
39
40# Load model
41model_name = "katanemo/Plano-Orchestrator-4B"
42tokenizer = AutoTokenizer.from_pretrained(model_name)
43model = AutoModelForCausalLM.from_pretrained(
44 model_name,
45 torch_dtype=torch.float16,
46 device_map="auto"
47)
48
49# Define available agents
50available_agents = [
51 {"name": "WeatherAgent", "description": "Provides weather forecasts and current conditions for any location"},
52 {"name": "CodeAgent", "description": "Generates, debugs, explains, and reviews code in multiple programming languages"}
53]
54
55# Conversation history
56conversation = [
57 {"role": "user", "content": "What's the weather like today?"},
58 {"role": "assistant", "content": "I can help you with that. Could you tell me your location?"},
59 {"role": "user", "content": "San Francisco"},
60]
61
62# Build messages and generate
63model_inputs = tokenizer.apply_chat_template(
64 messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
65).to(model.device)
66
67generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
68generated_ids = [
69 output_ids[len(input_ids) :]
70 for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
71]
72
73response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
74print(response)
75# Output: {"route": ["WeatherAgent"]}