|
| 1 | +# Code bellow written following examples here: https://llama-stack.readthedocs.io/en/latest/building_applications |
| 2 | +from llama_stack_client.lib.agents.agent import Agent |
| 3 | +from llama_stack_client.lib.agents.event_logger import EventLogger |
| 4 | +from llama_stack_client import LlamaStackClient |
| 5 | +import logging |
| 6 | +import os |
| 7 | +from dotenv import load_dotenv |
| 8 | +import json |
| 9 | + |
| 10 | + |
| 11 | + |
| 12 | + |
| 13 | +def test_mcp_server(file_path, REMOTE_MCP_URL, model, mcp_server_name): |
| 14 | + |
| 15 | + # Open the JSON file and load its content |
| 16 | + with open(file_path, 'r') as file: |
| 17 | + data = json.load(file) |
| 18 | + |
| 19 | + queries = data['queries'] |
| 20 | + user_prompts = [query['query'] for query in queries] |
| 21 | + |
| 22 | + logger = logging.getLogger(__name__) |
| 23 | + if not logger.hasHandlers(): |
| 24 | + logger.setLevel(logging.INFO) |
| 25 | + stream_handler = logging.StreamHandler() |
| 26 | + stream_handler.setLevel(logging.INFO) |
| 27 | + formatter = logging.Formatter('%(message)s') |
| 28 | + stream_handler.setFormatter(formatter) |
| 29 | + logger.addHandler(stream_handler) |
| 30 | + |
| 31 | + logger.info(f"Testing {mcp_server_name} toolgroup ") |
| 32 | + logger.info(f'Model is {model} \n \n') |
| 33 | + |
| 34 | + |
| 35 | + # Connect to a llama stack server |
| 36 | + base_url = os.getenv('REMOTE_BASE_URL') |
| 37 | + client = LlamaStackClient(base_url=base_url) |
| 38 | + logger.info(f"Connected to Llama Stack server @ {base_url} \n") |
| 39 | + |
| 40 | + # Get tool info and register tools |
| 41 | + registered_tools = client.tools.list() |
| 42 | + registered_tools_identifiers = [t.identifier for t in registered_tools] |
| 43 | + registered_toolgroups = [t.toolgroup_id for t in registered_tools] |
| 44 | + |
| 45 | + if mcp_server_name not in registered_toolgroups: |
| 46 | + # Register MCP tools |
| 47 | + client.toolgroups.register( |
| 48 | + toolgroup_id=mcp_server_name, |
| 49 | + provider_id="model-context-protocol", |
| 50 | + mcp_endpoint={"uri":REMOTE_MCP_URL}, |
| 51 | + ) |
| 52 | + |
| 53 | + logger.info(f"""Your Server has access the the following toolgroups: |
| 54 | + {set(registered_toolgroups)} |
| 55 | + """) |
| 56 | + |
| 57 | + for prompt in user_prompts: |
| 58 | + # Create simple agent with tools |
| 59 | + agent = Agent( |
| 60 | + client, |
| 61 | + model=model, |
| 62 | + instructions = """You are a helpful assistant. You have access to a number of tools. |
| 63 | + Whenever a tool is called, be sure return the Response in a friendly and helpful tone. |
| 64 | + When you are asked to search the web you must use a tool. keep answer concise |
| 65 | + """ , |
| 66 | + tools=[mcp_server_name], |
| 67 | + tool_config={"tool_choice":"auto"}, |
| 68 | + sampling_params={"max_tokens":4096} |
| 69 | + ) |
| 70 | + session_id = agent.create_session(session_name="Conversation_demo") |
| 71 | + turn_response = agent.create_turn( |
| 72 | + messages=[ |
| 73 | + { |
| 74 | + "role":"user", |
| 75 | + "content": prompt |
| 76 | + } |
| 77 | + ], |
| 78 | + session_id=session_id |
| 79 | + ) |
| 80 | + for log in EventLogger().log(turn_response): |
| 81 | + log.print() |
| 82 | + |
| 83 | + |
| 84 | +if __name__ == "__main__": |
| 85 | + load_dotenv() |
| 86 | + |
| 87 | + ansible_mcp_server_url = os.getenv('ANSIBLE_MCP_SERVER_URL') |
| 88 | + github_mcp_server_url = os.getenv('GITHUB_MCP_SERVER_URL') |
| 89 | + ocp_mcp_server_url = os.getenv('OCP_MCP_SERVER_URL') |
| 90 | + custom_mcp_server_url = os.getenv('CUSTOM_MCP_SERVER_URL') |
| 91 | + # Call the function with the file path |
| 92 | + ansilbe_file_path = './queries/anisble_queries.json' |
| 93 | + custom_file_path = './queries/custom_queries.json' |
| 94 | + github_file_path = './queries/github_queries.json' |
| 95 | + openshift_file_path = './queries/ocp_queries.json' |
| 96 | + |
| 97 | + test_mcp_server(ansilbe_file_path,ansible_mcp_server_url,"meta-llama/Llama-3.2-3B-Instruct","mcp::ansible") |
| 98 | + test_mcp_server(github_file_path,github_mcp_server_url,"ibm-granite/granite-3.2-8b-instruct","mcp::github") |
| 99 | + test_mcp_server(openshift_file_path,ocp_mcp_server_url,"ibm-granite/granite-3.2-8b-instruct","mcp::openshift") |
| 100 | + test_mcp_server(custom_file_path,custom_mcp_server_url,"meta-llama/Llama-3.2-3B-Instruct","mcp::custom_tool") |
0 commit comments