From fa3df199cc1666ff5a3636708ec0c07eb8ca81e3 Mon Sep 17 00:00:00 2001 From: Aksh Parekh Date: Tue, 4 Nov 2025 19:33:39 -0800 Subject: [PATCH] [CLEANUP] --- example.py | 61 ++++++++++++++----- .../aop_examples/auth/simple_auth_client.py | 5 +- .../aop_examples/auth/simple_auth_server.py | 11 ++-- test.py | 59 ++++++++++++++++++ 4 files changed, 113 insertions(+), 23 deletions(-) create mode 100644 test.py diff --git a/example.py b/example.py index 0a27c20c..a00fbeea 100644 --- a/example.py +++ b/example.py @@ -1,23 +1,56 @@ -import json +""" +Simple AOP Server with Custom Authentication Callback + +The auth_callback function determines ALL authentication logic. +If you provide auth_callback, authentication is enabled. +If you don't provide it, no authentication is required. +""" from swarms import Agent +from swarms.structs.aop import AOP + + +# This function governs ALL security +def custom_auth(token: str) -> bool: + """ + Your custom authentication logic goes here. + Return True to allow access, False to deny. + + This function determines everything: + - What tokens are valid + - Token format (API key, JWT, whatever) + - Any additional validation logic + """ + # Simple example: check against valid tokens + valid_tokens = { + "mytoken123", + "anothertoken456", + } + return token in valid_tokens -# Initialize the agent + +# Create agents agent = Agent( - agent_name="Quantitative-Trading-Agent", - agent_description="Advanced quantitative trading and algorithmic analysis agent", - model_name="gpt-4.1", - dynamic_temperature_enabled=True, + agent_name="Research-Agent", + model_name="claude-sonnet-4-5-20250929", max_loops=1, - dynamic_context_window=True, - streaming_on=False, - top_p=None, - output_type="dict", + system_prompt="You are a helpful research assistant.", + temperature=0.7, + top_p=None, # Can't use both temperature and top_p with Claude ) -out = agent.run( - task="What are the top five best energy stocks across nuclear, solar, gas, and other energy sources?", - n=1, +# Create server with auth callback +# If auth_callback is provided, auth is automatically enabled +server = AOP( + server_name="SimpleAuthServer", + port=5932, + auth_callback=custom_auth, # This enables and governs auth ) -print(json.dumps(out, indent=4)) +server.add_agent(agent) + +print("\nšŸš€ Server starting on port 5932") +print("šŸ” Authentication: ENABLED") +print("āœ… Valid tokens: mytoken123, anothertoken456\n") + +server.run() diff --git a/examples/aop_examples/auth/simple_auth_client.py b/examples/aop_examples/auth/simple_auth_client.py index 41ebc49a..adf29708 100644 --- a/examples/aop_examples/auth/simple_auth_client.py +++ b/examples/aop_examples/auth/simple_auth_client.py @@ -33,7 +33,7 @@ async def call_server(): result = await session.call_tool( name="discover_agents", arguments={ - "auth_token": "mytoken123" # That's it! + "auth_token": "mytoken123" }, ) @@ -43,12 +43,11 @@ async def call_server(): print("Calling Research-Agent...") print("=" * 60 + "\n") - # Same for any tool result = await session.call_tool( name="Research-Agent", arguments={ "task": "What is Python?", - "auth_token": "mytoken123" # Just include it + "auth_token": "mytoken123" }, ) diff --git a/examples/aop_examples/auth/simple_auth_server.py b/examples/aop_examples/auth/simple_auth_server.py index a00fbeea..31f61ebd 100644 --- a/examples/aop_examples/auth/simple_auth_server.py +++ b/examples/aop_examples/auth/simple_auth_server.py @@ -10,7 +10,7 @@ from swarms import Agent from swarms.structs.aop import AOP -# This function governs ALL security +# EXAMPLE: This function governs ALL security def custom_auth(token: str) -> bool: """ Your custom authentication logic goes here. @@ -21,7 +21,7 @@ def custom_auth(token: str) -> bool: - Token format (API key, JWT, whatever) - Any additional validation logic """ - # Simple example: check against valid tokens + valid_tokens = { "mytoken123", "anothertoken456", @@ -36,15 +36,14 @@ agent = Agent( max_loops=1, system_prompt="You are a helpful research assistant.", temperature=0.7, - top_p=None, # Can't use both temperature and top_p with Claude + top_p=None, ) -# Create server with auth callback -# If auth_callback is provided, auth is automatically enabled + server = AOP( server_name="SimpleAuthServer", port=5932, - auth_callback=custom_auth, # This enables and governs auth + auth_callback=custom_auth, ) server.add_agent(agent) diff --git a/test.py b/test.py new file mode 100644 index 00000000..ee7dbd1d --- /dev/null +++ b/test.py @@ -0,0 +1,59 @@ +""" +Simple AOP Client with Authentication + +Just pass your token when calling tools. That's it. +The server's auth_callback determines if it's valid. +""" + +import json +import asyncio +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + + +async def call_server(): + """Call the AOP server with authentication.""" + + url = "http://localhost:5932/mcp" + + async with streamablehttp_client(url, timeout=10) as ctx: + if len(ctx) == 2: + read, write = ctx + else: + read, write, *_ = ctx + + async with ClientSession(read, write) as session: + await session.initialize() + + print("\n" + "=" * 60) + print("Calling discover_agents...") + print("=" * 60 + "\n") + + result = await session.call_tool( + name="discover_agents", + arguments={ + "auth_token": "mytoken1234" + }, + ) + + print(json.dumps(result.model_dump(), indent=2)) + + print("\n" + "=" * 60) + print("Calling Research-Agent...") + print("=" * 60 + "\n") + + result = await session.call_tool( + name="Research-Agent", + arguments={ + "task": "What is Python?", + "auth_token": "mytoken123" + }, + ) + + print(json.dumps(result.model_dump(), indent=2)) + + +if __name__ == "__main__": + print("\nšŸ” Simple Auth Client") + print("Token: mytoken123\n") + asyncio.run(call_server())