84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
"""
|
|
Test script for domain-specific tool retrieval functions
|
|
"""
|
|
|
|
import json
|
|
from pprint import pprint
|
|
import sys
|
|
sys.path.append("/home/ubuntu/sas0/lzy/multi_mcp_server")
|
|
from sci_mcp.core.llm_tools import get_domain_tools, get_domain_tool_schemas, get_all_tool_schemas,get_all_tools
|
|
|
|
def test_get_domain_tools():
|
|
"""Test retrieving tools from specific domains"""
|
|
print("\n=== Testing get_domain_tools ===")
|
|
|
|
# Test with material and general domains
|
|
domains = ['material', 'general']
|
|
print(f"Getting tools for domains: {domains}")
|
|
|
|
domain_tools = get_domain_tools(domains)
|
|
|
|
# Print results
|
|
for domain, tools in domain_tools.items():
|
|
print(f"\nDomain: {domain}")
|
|
print(f"Number of tools: {len(tools)}")
|
|
print("Tool names:")
|
|
for tool_name in tools.keys():
|
|
print(f" - {tool_name}")
|
|
|
|
def test_get_domain_tool_schemas():
|
|
"""Test retrieving tool schemas from specific domains"""
|
|
print("\n=== Testing get_domain_tool_schemas (OpenAI format) ===")
|
|
|
|
# Test with material and general domains
|
|
domains = ['material', 'general']
|
|
print(f"Getting tool schemas for domains: {domains}")
|
|
|
|
domain_schemas = get_domain_tool_schemas(domains)
|
|
|
|
# Print results
|
|
for domain, schemas in domain_schemas.items():
|
|
print(f"\nDomain: {domain}")
|
|
print(f"Number of schemas: {len(schemas)}")
|
|
print("Tool names in schemas:")
|
|
for schema in schemas:
|
|
print(f" - {schema['function']['name']}")
|
|
|
|
|
|
def test_all_tool_schemas():
|
|
"""Test retrieving all tool schemas in both formats"""
|
|
print("\n=== Testing get_all_tool_schemas ===")
|
|
|
|
# OpenAI format
|
|
print("Getting all tool schemas in OpenAI format")
|
|
openai_schemas = get_all_tool_schemas()
|
|
print(f"Number of schemas: {len(openai_schemas)}")
|
|
print("Tool names:")
|
|
for schema in openai_schemas:
|
|
print(f" - {schema['function']['name']}")
|
|
|
|
# MCP format
|
|
print("\nGetting all tool schemas in MCP format")
|
|
mcp_schemas = get_all_tool_schemas(use_mcp_format=True)
|
|
print(f"Number of schemas: {len(mcp_schemas)}")
|
|
print("Tool names:")
|
|
for schema in mcp_schemas:
|
|
print(f" - {schema['name']}")
|
|
|
|
def main():
|
|
"""Main function to run tests"""
|
|
print("Testing domain-specific tool retrieval functions")
|
|
|
|
#test_get_domain_tools()
|
|
test_get_domain_tool_schemas()
|
|
|
|
test_all_tool_schemas()
|
|
|
|
if __name__ == "__main__":
|
|
#print(get_all_tool_schemas())
|
|
from rich.console import Console
|
|
console = Console()
|
|
console.print("Testing domain-specific tool retrieval functions", style="bold green")
|
|
console.print(get_domain_tool_schemas(['chemistry']))
|
|
|