Skip to content

Commit 3c0c5bb

Browse files
cuericleecuericlee
authored
Merge pull request #11 from loveyana/main
feat(cloud): add OAuth2/JWT authentication support for CloudStrategy
2 parents 57144db + d898089 commit 3c0c5bb

File tree

102 files changed

+7660
-4886
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

102 files changed

+7660
-4886
lines changed

‎agentkit/apps/a2a_app/telemetry.py‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646

4747
logger = logging.getLogger("agentkit." + __name__)
4848

49+
4950
class Telemetry:
5051
def __init__(self):
5152
self.tracer = get_tracer("agentkit.a2a_app")
@@ -57,7 +58,6 @@ def __init__(self):
5758
explicit_bucket_boundaries_advisory=_GEN_AI_CLIENT_OPERATION_DURATION_BUCKETS,
5859
)
5960

60-
6161
def trace_a2a_agent(
6262
self,
6363
func: Callable,
@@ -69,7 +69,9 @@ def trace_a2a_agent(
6969
"""Get current span and set required attributes."""
7070
trace_id = span.get_span_context().trace_id
7171
span_id = span.get_span_context().span_id
72-
logger.debug(f"Set attributes for span with trace_id={trace_id}, span_id={span_id}")
72+
logger.debug(
73+
f"Set attributes for span with trace_id={trace_id}, span_id={span_id}"
74+
)
7375

7476
# ===============================
7577
# Set attributes for current span
@@ -87,12 +89,18 @@ def trace_a2a_agent(
8789
if user_id:
8890
span.set_attribute(key="gen_ai.user.id", value=user_id)
8991

90-
span.set_attribute(key="gen_ai.input", value=safe_serialize_to_json_string(request.message.parts))
92+
span.set_attribute(
93+
key="gen_ai.input",
94+
value=safe_serialize_to_json_string(request.message.parts),
95+
)
9196

9297
span.set_attribute(key="gen_ai.span.kind", value="a2a_agent")
9398
span.set_attribute(key="gen_ai.operation.name", value="invoke_agent")
9499
span.set_attribute(key="gen_ai.operation.type", value="a2a_agent")
95-
attributes={"gen_ai_operation_name": "invoke_agent", "gen_ai_operation_type": "a2a_agent"}
100+
attributes = {
101+
"gen_ai_operation_name": "invoke_agent",
102+
"gen_ai_operation_type": "a2a_agent",
103+
}
96104

97105
if exception:
98106
self.handle_exception(span, exception)
@@ -116,4 +124,3 @@ def handle_exception(span: trace.Span, exception: Exception) -> None:
116124

117125

118126
telemetry = Telemetry()
119-

‎agentkit/apps/mcp_app/telemetry.py‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444

4545
logger = logging.getLogger("agentkit." + __name__)
4646

47+
4748
class Telemetry:
4849
def __init__(self):
4950
self.tracer = get_tracer("agentkit.mcp_app")
@@ -56,20 +57,22 @@ def __init__(self):
5657
)
5758

5859
def trace_tool(
59-
self,
60-
func: Callable,
61-
span: Span,
62-
args: dict,
63-
func_result: any,
64-
operation_type: str,
65-
exception: Exception,
60+
self,
61+
func: Callable,
62+
span: Span,
63+
args: dict,
64+
func_result: any,
65+
operation_type: str,
66+
exception: Exception,
6667
) -> None:
6768
"""Get current span and set required attributes."""
6869

6970
trace_id = span.get_span_context().trace_id
7071
span_id = span.get_span_context().span_id
7172

72-
logger.debug(f"Set attributes for span with trace_id={trace_id}, span_id={span_id}")
73+
logger.debug(
74+
f"Set attributes for span with trace_id={trace_id}, span_id={span_id}"
75+
)
7376

7477
# ===============================
7578
# Set attributes for current span
@@ -79,14 +82,19 @@ def trace_tool(
7982

8083
span.set_attribute(key="gen_ai.func_name", value=func.__name__)
8184

82-
span.set_attribute(key="gen_ai.input", value=safe_serialize_to_json_string(args))
85+
span.set_attribute(
86+
key="gen_ai.input", value=safe_serialize_to_json_string(args)
87+
)
8388
safe_result = safe_serialize_to_json_string(func_result)
8489
span.set_attribute(key="gen_ai.output", value=safe_result)
8590

8691
span.set_attribute(key="gen_ai.span.kind", value="tool")
8792
span.set_attribute(key="gen_ai.operation.name", value="tool")
8893
span.set_attribute(key="gen_ai.operation.type", value=operation_type)
89-
attributes = {"gen_ai_operation_name": "tool", "gen_ai_operation_type": operation_type}
94+
attributes = {
95+
"gen_ai_operation_name": "tool",
96+
"gen_ai_operation_type": operation_type,
97+
}
9098

9199
if exception:
92100
self.handle_exception(span, exception)
@@ -110,4 +118,3 @@ def handle_exception(span: trace.Span, exception: Exception) -> None:
110118

111119

112120
telemetry = Telemetry()
113-

‎agentkit/apps/simple_app/simple_app_handlers.py‎

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -73,14 +73,10 @@ async def handle(self, request: Request) -> Response | StreamingResponse:
7373

7474
logger.info("Returning non-streaming response")
7575
safe_json_string = safe_serialize_to_json_string(result)
76-
telemetry.trace_agent_finish(
77-
safe_json_string, None
78-
)
76+
telemetry.trace_agent_finish(safe_json_string, None)
7977
except Exception as e:
8078
logger.error("Invoke handler function failed: %s", e)
81-
telemetry.trace_agent_finish(
82-
"", e
83-
)
79+
telemetry.trace_agent_finish("", e)
8480
raise e
8581

8682
return Response(safe_json_string, media_type="application/json")
@@ -160,9 +156,7 @@ def _sync_stream_with_error_handling(self, generator):
160156
finally:
161157
# finish trace span and record metrics
162158
result = safe_serialize_to_json_string(last_value)
163-
telemetry.trace_agent_finish(
164-
result, exception
165-
)
159+
telemetry.trace_agent_finish(result, exception)
166160

167161
async def _stream_with_error_handling(self, generator):
168162
"""Wrap async generator to handle errors and convert to SSE format."""
@@ -184,9 +178,7 @@ async def _stream_with_error_handling(self, generator):
184178
finally:
185179
# finish trace span and record metrics
186180
result = safe_serialize_to_json_string(last_value)
187-
telemetry.trace_agent_finish(
188-
result, exception
189-
)
181+
telemetry.trace_agent_finish(result, exception)
190182

191183

192184
class PingHandler(BaseHandler):

‎agentkit/apps/simple_app/telemetry.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,9 @@ def trace_agent(
9292
trace_id = span.get_span_context().trace_id
9393
span_id = span.get_span_context().span_id
9494

95-
logger.debug(f"Set attributes for span with trace_id={trace_id}, span_id={span_id}")
95+
logger.debug(
96+
f"Set attributes for span with trace_id={trace_id}, span_id={span_id}"
97+
)
9698

9799
# ===============================
98100
# Set attributes for current span
@@ -112,7 +114,9 @@ def trace_agent(
112114
if user_id:
113115
span.set_attribute(key="gen_ai.user.id", value=user_id)
114116

115-
span.set_attribute(key="gen_ai.input", value=safe_serialize_to_json_string(payload))
117+
span.set_attribute(
118+
key="gen_ai.input", value=safe_serialize_to_json_string(payload)
119+
)
116120

117121
span.set_attribute(key="gen_ai.span.kind", value="workflow")
118122
span.set_attribute(key="gen_ai.operation.name", value="invoke_agent")
@@ -128,7 +132,10 @@ def trace_agent_finish(
128132

129133
if span and span.is_recording():
130134
span.set_attribute(key="gen_ai.output", value=func_result)
131-
attributes={"gen_ai_operation_name": "invoke_agent", "gen_ai_operation_type": "agent"}
135+
attributes = {
136+
"gen_ai_operation_name": "invoke_agent",
137+
"gen_ai_operation_type": "agent",
138+
}
132139
if exception:
133140
self.handle_exception(span, exception)
134141
attributes["error_type"] = exception.__class__.__name__

‎agentkit/client/__init__.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,8 @@
2323
from agentkit.client.base_iam_client import BaseIAMClient
2424

2525
__all__ = [
26-
'BaseServiceClient',
27-
'BaseAgentkitClient',
28-
'BaseIAMClient',
29-
'ApiConfig',
26+
"BaseServiceClient",
27+
"BaseAgentkitClient",
28+
"BaseIAMClient",
29+
"ApiConfig",
3030
]

‎agentkit/client/base_agentkit_client.py‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,20 +26,20 @@
2626
class BaseAgentkitClient(BaseServiceClient):
2727
"""
2828
Base client for all AgentKit services.
29-
29+
3030
This class provides:
3131
1. Common credential initialization
3232
2. Unified API invocation logic with error handling
3333
3. Automatic ApiInfo generation with flexible configuration
34-
34+
3535
Subclasses should override API_ACTIONS with either:
3636
- Simple dict mapping: {"ActionName": "ActionName"}
3737
- Detailed ApiConfig: {"ActionName": ApiConfig(action="ActionName", method="GET", path="/custom")}
3838
"""
39-
39+
4040
# Subclasses should override this with their API action configurations
4141
API_ACTIONS: Dict[str, Union[str, ApiConfig]] = {}
42-
42+
4343
def __init__(
4444
self,
4545
access_key: str = "",
@@ -50,7 +50,7 @@ def __init__(
5050
) -> None:
5151
"""
5252
Initialize the AgentKit client.
53-
53+
5454
Args:
5555
access_key: Volcengine access key
5656
secret_key: Volcengine secret key
@@ -64,23 +64,23 @@ def __init__(
6464
region=region,
6565
session_token=session_token,
6666
service_name=service_name,
67-
credential_env_prefix='AGENTKIT',
67+
credential_env_prefix="AGENTKIT",
6868
)
69-
69+
7070
def _get_service_config(self) -> Dict[str, str]:
7171
"""
7272
Get AgentKit service configuration.
73-
73+
7474
Returns:
7575
Dictionary with host, api_version, and service
7676
"""
7777
host, api_version, service = get_volc_agentkit_host_info()
7878
return {
79-
'host': host,
80-
'api_version': api_version,
81-
'service': service,
79+
"host": host,
80+
"api_version": api_version,
81+
"service": service,
8282
}
83-
83+
8484
def _get(self, api_action: str, params: Dict[str, Any] = None) -> str:
8585
"""Legacy method for GET requests."""
8686
try:

‎agentkit/client/base_iam_client.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -25,25 +25,25 @@
2525
class BaseIAMClient(BaseServiceClient):
2626
"""
2727
Base client for IAM services.
28-
28+
2929
This class provides the same interface as BaseAgentkitClient but for IAM services:
3030
1. Common credential initialization
3131
2. Unified API invocation logic with error handling
3232
3. Automatic ApiInfo generation with flexible configuration
33-
33+
3434
Subclasses should override API_ACTIONS with either:
3535
- Simple dict mapping: {"ActionName": "ActionName"}
3636
- Detailed ApiConfig: {"ActionName": ApiConfig(action="ActionName", method="GET", path="/custom")}
3737
"""
38-
38+
3939
# Subclasses should override this with their API action configurations
4040
API_ACTIONS: Dict[str, Union[str, ApiConfig]] = {}
41-
41+
4242
# IAM service specific configuration
4343
IAM_API_VERSION = "2018-01-01"
4444
IAM_SERVICE_CODE = "iam"
45-
IAM_HOST ="open.volcengineapi.com"
46-
45+
IAM_HOST = "open.volcengineapi.com"
46+
4747
def __init__(
4848
self,
4949
access_key: str = "",
@@ -54,7 +54,7 @@ def __init__(
5454
) -> None:
5555
"""
5656
Initialize the IAM client.
57-
57+
5858
Args:
5959
access_key: Volcengine access key
6060
secret_key: Volcengine secret key
@@ -68,18 +68,18 @@ def __init__(
6868
region=region,
6969
session_token=session_token,
7070
service_name=service_name,
71-
credential_env_prefix='IAM',
71+
credential_env_prefix="IAM",
7272
)
73-
73+
7474
def _get_service_config(self) -> Dict[str, str]:
7575
"""
7676
Get IAM service configuration.
77-
77+
7878
Returns:
7979
Dictionary with host, api_version, and service
8080
"""
8181
return {
82-
'host': self.IAM_HOST,
83-
'api_version': self.IAM_API_VERSION,
84-
'service': self.IAM_SERVICE_CODE,
82+
"host": self.IAM_HOST,
83+
"api_version": self.IAM_API_VERSION,
84+
"service": self.IAM_SERVICE_CODE,
8585
}

0 commit comments

Comments
 (0)