agenta 0.30.0a2__py3-none-any.whl → 0.30.0a3__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of agenta might be problematic. Click here for more details.
- agenta/__init__.py +1 -0
- agenta/client/backend/__init__.py +32 -3
- agenta/client/backend/access_control/__init__.py +1 -0
- agenta/client/backend/access_control/client.py +167 -0
- agenta/client/backend/apps/client.py +70 -10
- agenta/client/backend/client.py +61 -45
- agenta/client/backend/configs/client.py +6 -0
- agenta/client/backend/containers/client.py +6 -0
- agenta/client/backend/core/file.py +13 -8
- agenta/client/backend/environments/client.py +6 -0
- agenta/client/backend/evaluations/client.py +14 -1
- agenta/client/backend/evaluators/client.py +24 -0
- agenta/client/backend/observability/client.py +22 -16
- agenta/client/backend/observability_v_1/__init__.py +2 -2
- agenta/client/backend/observability_v_1/client.py +203 -0
- agenta/client/backend/observability_v_1/types/__init__.py +2 -1
- agenta/client/backend/observability_v_1/types/format.py +1 -1
- agenta/client/backend/observability_v_1/types/query_analytics_response.py +7 -0
- agenta/client/backend/scopes/__init__.py +1 -0
- agenta/client/backend/scopes/client.py +114 -0
- agenta/client/backend/testsets/client.py +305 -121
- agenta/client/backend/types/__init__.py +24 -2
- agenta/client/backend/types/analytics_response.py +24 -0
- agenta/client/backend/types/app.py +2 -1
- agenta/client/backend/types/body_import_testset.py +0 -1
- agenta/client/backend/types/bucket_dto.py +26 -0
- agenta/client/backend/types/header_dto.py +22 -0
- agenta/client/backend/types/legacy_analytics_response.py +29 -0
- agenta/client/backend/types/legacy_data_point.py +27 -0
- agenta/client/backend/types/metrics_dto.py +24 -0
- agenta/client/backend/types/permission.py +1 -0
- agenta/client/backend/types/projects_response.py +28 -0
- agenta/client/backend/types/provider_key_dto.py +23 -0
- agenta/client/backend/types/provider_kind.py +21 -0
- agenta/client/backend/types/secret_dto.py +24 -0
- agenta/client/backend/types/secret_kind.py +5 -0
- agenta/client/backend/types/secret_response_dto.py +27 -0
- agenta/client/backend/variants/client.py +66 -0
- agenta/client/backend/vault/__init__.py +1 -0
- agenta/client/backend/vault/client.py +685 -0
- agenta/client/client.py +1 -1
- agenta/sdk/__init__.py +1 -0
- agenta/sdk/agenta_init.py +47 -118
- agenta/sdk/assets.py +57 -46
- agenta/sdk/context/exporting.py +25 -0
- agenta/sdk/context/routing.py +12 -12
- agenta/sdk/context/tracing.py +26 -1
- agenta/sdk/decorators/routing.py +279 -243
- agenta/sdk/decorators/tracing.py +53 -31
- agenta/sdk/managers/config.py +8 -118
- agenta/sdk/managers/secrets.py +38 -0
- agenta/sdk/middleware/auth.py +128 -93
- agenta/sdk/middleware/cache.py +4 -0
- agenta/sdk/middleware/config.py +254 -0
- agenta/sdk/middleware/cors.py +27 -0
- agenta/sdk/middleware/otel.py +40 -0
- agenta/sdk/middleware/vault.py +158 -0
- agenta/sdk/tracing/exporters.py +40 -2
- agenta/sdk/tracing/inline.py +2 -2
- agenta/sdk/tracing/processors.py +11 -3
- agenta/sdk/tracing/tracing.py +14 -12
- agenta/sdk/utils/constants.py +1 -0
- agenta/sdk/utils/exceptions.py +20 -19
- agenta/sdk/utils/globals.py +4 -8
- agenta/sdk/utils/timing.py +58 -0
- {agenta-0.30.0a2.dist-info → agenta-0.30.0a3.dist-info}/METADATA +3 -2
- {agenta-0.30.0a2.dist-info → agenta-0.30.0a3.dist-info}/RECORD +69 -44
- {agenta-0.30.0a2.dist-info → agenta-0.30.0a3.dist-info}/WHEEL +1 -1
- agenta/client/backend/types/lm_providers_enum.py +0 -21
- agenta/sdk/tracing/context.py +0 -24
- {agenta-0.30.0a2.dist-info → agenta-0.30.0a3.dist-info}/entry_points.txt +0 -0
agenta/sdk/agenta_init.py
CHANGED
|
@@ -6,8 +6,9 @@ from importlib.metadata import version
|
|
|
6
6
|
from agenta.sdk.utils.logging import log
|
|
7
7
|
from agenta.sdk.utils.globals import set_global
|
|
8
8
|
from agenta.client.backend.client import AgentaApi, AsyncAgentaApi
|
|
9
|
+
|
|
9
10
|
from agenta.sdk.tracing import Tracing
|
|
10
|
-
from agenta.
|
|
11
|
+
from agenta.sdk.context.routing import routing_context
|
|
11
12
|
|
|
12
13
|
|
|
13
14
|
class AgentaSingleton:
|
|
@@ -59,9 +60,7 @@ class AgentaSingleton:
|
|
|
59
60
|
ValueError: If `app_id` is not specified either as an argument, in the config file, or in the environment variables.
|
|
60
61
|
"""
|
|
61
62
|
|
|
62
|
-
log.info("
|
|
63
|
-
log.info("Agenta SDK - using version: %s", version("agenta"))
|
|
64
|
-
log.info("---------------------------")
|
|
63
|
+
log.info("Agenta - SDK version: %s", version("agenta"))
|
|
65
64
|
|
|
66
65
|
config = {}
|
|
67
66
|
if config_fname:
|
|
@@ -86,6 +85,13 @@ class AgentaSingleton:
|
|
|
86
85
|
|
|
87
86
|
self.api_key = api_key or getenv("AGENTA_API_KEY") or config.get("api_key")
|
|
88
87
|
|
|
88
|
+
self.base_id = getenv("AGENTA_BASE_ID")
|
|
89
|
+
|
|
90
|
+
self.service_id = getenv("AGENTA_SERVICE_ID") or self.base_id
|
|
91
|
+
|
|
92
|
+
log.info("Agenta - Service ID: %s", self.service_id)
|
|
93
|
+
log.info("Agenta - Application ID: %s", self.app_id)
|
|
94
|
+
|
|
89
95
|
self.tracing = Tracing(
|
|
90
96
|
url=f"{self.host}/api/observability/v1/otlp/traces", # type: ignore
|
|
91
97
|
redact=redact,
|
|
@@ -94,6 +100,7 @@ class AgentaSingleton:
|
|
|
94
100
|
|
|
95
101
|
self.tracing.configure(
|
|
96
102
|
api_key=self.api_key,
|
|
103
|
+
service_id=self.service_id,
|
|
97
104
|
# DEPRECATING
|
|
98
105
|
app_id=self.app_id,
|
|
99
106
|
)
|
|
@@ -108,8 +115,6 @@ class AgentaSingleton:
|
|
|
108
115
|
api_key=self.api_key if self.api_key else "",
|
|
109
116
|
)
|
|
110
117
|
|
|
111
|
-
self.base_id = getenv("AGENTA_BASE_ID")
|
|
112
|
-
|
|
113
118
|
self.config = Config(
|
|
114
119
|
host=self.host,
|
|
115
120
|
base_id=self.base_id,
|
|
@@ -120,28 +125,43 @@ class AgentaSingleton:
|
|
|
120
125
|
class Config:
|
|
121
126
|
def __init__(
|
|
122
127
|
self,
|
|
123
|
-
|
|
128
|
+
# LEGACY
|
|
129
|
+
host: Optional[str] = None,
|
|
124
130
|
base_id: Optional[str] = None,
|
|
125
|
-
api_key: Optional[str] =
|
|
131
|
+
api_key: Optional[str] = None,
|
|
132
|
+
# LEGACY
|
|
133
|
+
**kwargs,
|
|
126
134
|
):
|
|
127
|
-
self.
|
|
135
|
+
self.default_parameters = {**kwargs}
|
|
136
|
+
|
|
137
|
+
def set_default(self, **kwargs):
|
|
138
|
+
self.default_parameters.update(kwargs)
|
|
139
|
+
|
|
140
|
+
def get_default(self):
|
|
141
|
+
return self.default_parameters
|
|
142
|
+
|
|
143
|
+
def __getattr__(self, key):
|
|
144
|
+
context = routing_context.get()
|
|
145
|
+
|
|
146
|
+
parameters = context.parameters
|
|
147
|
+
|
|
148
|
+
if not parameters:
|
|
149
|
+
return None
|
|
150
|
+
|
|
151
|
+
if key in parameters:
|
|
152
|
+
value = parameters[key]
|
|
153
|
+
|
|
154
|
+
if isinstance(value, dict):
|
|
155
|
+
nested_config = Config()
|
|
156
|
+
nested_config.set_default(**value)
|
|
128
157
|
|
|
129
|
-
|
|
158
|
+
return nested_config
|
|
130
159
|
|
|
131
|
-
|
|
132
|
-
# print(
|
|
133
|
-
# "Warning: Your configuration will not be saved permanently since base_id is not provided.\n"
|
|
134
|
-
# )
|
|
135
|
-
pass
|
|
160
|
+
return value
|
|
136
161
|
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
self.persist = True
|
|
141
|
-
self.client = AgentaApi(
|
|
142
|
-
base_url=self.host + "/api",
|
|
143
|
-
api_key=api_key if api_key else "",
|
|
144
|
-
)
|
|
162
|
+
return None
|
|
163
|
+
|
|
164
|
+
### --- LEGACY --- ###
|
|
145
165
|
|
|
146
166
|
def register_default(self, overwrite=False, **kwargs):
|
|
147
167
|
"""alias for default"""
|
|
@@ -153,104 +173,13 @@ class Config:
|
|
|
153
173
|
overwrite: Whether to overwrite the existing configuration or not
|
|
154
174
|
**kwargs: A dict containing the parameters
|
|
155
175
|
"""
|
|
156
|
-
self.set(
|
|
157
|
-
**kwargs
|
|
158
|
-
) # In case there is no connectivity, we still can use the default values
|
|
159
|
-
try:
|
|
160
|
-
self.push(config_name="default", overwrite=overwrite, **kwargs)
|
|
161
|
-
except Exception as ex:
|
|
162
|
-
log.warning(
|
|
163
|
-
"Unable to push the default configuration to the server. %s", str(ex)
|
|
164
|
-
)
|
|
165
|
-
|
|
166
|
-
def push(self, config_name: str, overwrite=True, **kwargs):
|
|
167
|
-
"""Pushes the parameters for the app variant to the server
|
|
168
|
-
Args:
|
|
169
|
-
config_name: Name of the configuration to push to
|
|
170
|
-
overwrite: Whether to overwrite the existing configuration or not
|
|
171
|
-
**kwargs: A dict containing the parameters
|
|
172
|
-
"""
|
|
173
|
-
if not self.persist:
|
|
174
|
-
return
|
|
175
|
-
try:
|
|
176
|
-
self.client.configs.save_config(
|
|
177
|
-
base_id=self.base_id,
|
|
178
|
-
config_name=config_name,
|
|
179
|
-
parameters=kwargs,
|
|
180
|
-
overwrite=overwrite,
|
|
181
|
-
)
|
|
182
|
-
except Exception as ex:
|
|
183
|
-
log.warning(
|
|
184
|
-
"Failed to push the configuration to the server with error: %s", ex
|
|
185
|
-
)
|
|
186
|
-
|
|
187
|
-
def pull(
|
|
188
|
-
self, config_name: str = "default", environment_name: Optional[str] = None
|
|
189
|
-
):
|
|
190
|
-
"""Pulls the parameters for the app variant from the server and sets them to the config"""
|
|
191
|
-
if not self.persist and (
|
|
192
|
-
config_name != "default" or environment_name is not None
|
|
193
|
-
):
|
|
194
|
-
raise ValueError(
|
|
195
|
-
"Cannot pull the configuration from the server since the app_name and base_name are not provided."
|
|
196
|
-
)
|
|
197
|
-
if self.persist:
|
|
198
|
-
try:
|
|
199
|
-
if environment_name:
|
|
200
|
-
config = self.client.configs.get_config(
|
|
201
|
-
base_id=self.base_id, environment_name=environment_name
|
|
202
|
-
)
|
|
203
|
-
|
|
204
|
-
else:
|
|
205
|
-
config = self.client.configs.get_config(
|
|
206
|
-
base_id=self.base_id,
|
|
207
|
-
config_name=config_name,
|
|
208
|
-
)
|
|
209
|
-
except Exception as ex:
|
|
210
|
-
log.warning(
|
|
211
|
-
"Failed to pull the configuration from the server with error: %s",
|
|
212
|
-
str(ex),
|
|
213
|
-
)
|
|
214
|
-
try:
|
|
215
|
-
self.set(**{"current_version": config.current_version, **config.parameters})
|
|
216
|
-
except Exception as ex:
|
|
217
|
-
log.warning("Failed to set the configuration with error: %s", str(ex))
|
|
176
|
+
self.set(**kwargs)
|
|
218
177
|
|
|
219
|
-
def all(self):
|
|
220
|
-
"""Returns all the parameters for the app variant"""
|
|
221
|
-
return {
|
|
222
|
-
k: v
|
|
223
|
-
for k, v in self.__dict__.items()
|
|
224
|
-
if k
|
|
225
|
-
not in [
|
|
226
|
-
"app_name",
|
|
227
|
-
"base_name",
|
|
228
|
-
"host",
|
|
229
|
-
"base_id",
|
|
230
|
-
"api_key",
|
|
231
|
-
"persist",
|
|
232
|
-
"client",
|
|
233
|
-
]
|
|
234
|
-
}
|
|
235
|
-
|
|
236
|
-
# function to set the parameters for the app variant
|
|
237
178
|
def set(self, **kwargs):
|
|
238
|
-
|
|
239
|
-
|
|
240
|
-
Args:
|
|
241
|
-
**kwargs: A dict containing the parameters
|
|
242
|
-
"""
|
|
243
|
-
for key, value in kwargs.items():
|
|
244
|
-
setattr(self, key, value)
|
|
245
|
-
|
|
246
|
-
def dump(self):
|
|
247
|
-
"""Returns all the information about the current version in the configuration.
|
|
179
|
+
self.set_default(**kwargs)
|
|
248
180
|
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
"""
|
|
252
|
-
|
|
253
|
-
raise NotImplementedError()
|
|
181
|
+
def all(self):
|
|
182
|
+
return self.default_parameters
|
|
254
183
|
|
|
255
184
|
|
|
256
185
|
def init(
|
agenta/sdk/assets.py
CHANGED
|
@@ -1,23 +1,9 @@
|
|
|
1
1
|
supported_llm_models = {
|
|
2
|
-
"
|
|
3
|
-
"
|
|
4
|
-
"
|
|
5
|
-
"
|
|
6
|
-
"
|
|
7
|
-
],
|
|
8
|
-
"Open AI": [
|
|
9
|
-
"gpt-3.5-turbo-1106",
|
|
10
|
-
"gpt-3.5-turbo",
|
|
11
|
-
"gpt-4",
|
|
12
|
-
"gpt-4o",
|
|
13
|
-
"gpt-4o-mini",
|
|
14
|
-
"gpt-4-1106-preview",
|
|
15
|
-
],
|
|
16
|
-
"Gemini": ["gemini/gemini-1.5-pro-latest", "gemini/gemini-1.5-flash"],
|
|
17
|
-
"Cohere": [
|
|
18
|
-
"cohere/command-light",
|
|
19
|
-
"cohere/command-r-plus",
|
|
20
|
-
"cohere/command-nightly",
|
|
2
|
+
"Aleph Alpha": [
|
|
3
|
+
"luminous-base",
|
|
4
|
+
"luminous-base-control",
|
|
5
|
+
"luminous-extended-control",
|
|
6
|
+
"luminous-supreme",
|
|
21
7
|
],
|
|
22
8
|
"Anthropic": [
|
|
23
9
|
"anthropic/claude-3-5-sonnet-20240620",
|
|
@@ -33,11 +19,10 @@ supported_llm_models = {
|
|
|
33
19
|
"anyscale/meta-llama/Llama-2-13b-chat-hf",
|
|
34
20
|
"anyscale/meta-llama/Llama-2-70b-chat-hf",
|
|
35
21
|
],
|
|
36
|
-
"
|
|
37
|
-
"
|
|
38
|
-
"
|
|
39
|
-
"
|
|
40
|
-
"perplexity/pplx-70b-online",
|
|
22
|
+
"Cohere": [
|
|
23
|
+
"cohere/command-light",
|
|
24
|
+
"cohere/command-r-plus",
|
|
25
|
+
"cohere/command-nightly",
|
|
41
26
|
],
|
|
42
27
|
"DeepInfra": [
|
|
43
28
|
"deepinfra/meta-llama/Llama-2-70b-chat-hf",
|
|
@@ -46,6 +31,46 @@ supported_llm_models = {
|
|
|
46
31
|
"deepinfra/mistralai/Mistral-7B-Instruct-v0.1",
|
|
47
32
|
"deepinfra/jondurbin/airoboros-l2-70b-gpt4-1.4.1",
|
|
48
33
|
],
|
|
34
|
+
"Gemini": [
|
|
35
|
+
"gemini/gemini-1.5-pro-latest",
|
|
36
|
+
"gemini/gemini-1.5-flash",
|
|
37
|
+
],
|
|
38
|
+
"Groq": [
|
|
39
|
+
"groq/llama3-8b-8192",
|
|
40
|
+
"groq/llama3-70b-8192",
|
|
41
|
+
"groq/llama2-70b-4096",
|
|
42
|
+
"groq/mixtral-8x7b-32768",
|
|
43
|
+
"groq/gemma-7b-it",
|
|
44
|
+
],
|
|
45
|
+
"Mistral AI": [
|
|
46
|
+
"mistral/mistral-tiny",
|
|
47
|
+
"mistral/mistral-small",
|
|
48
|
+
"mistral/mistral-medium",
|
|
49
|
+
"mistral/mistral-large-latest",
|
|
50
|
+
],
|
|
51
|
+
"Open AI": [
|
|
52
|
+
"gpt-3.5-turbo-1106",
|
|
53
|
+
"gpt-3.5-turbo",
|
|
54
|
+
"gpt-4",
|
|
55
|
+
"gpt-4o",
|
|
56
|
+
"gpt-4o-mini",
|
|
57
|
+
"gpt-4-1106-preview",
|
|
58
|
+
],
|
|
59
|
+
"OpenRouter": [
|
|
60
|
+
"openrouter/openai/gpt-3.5-turbo",
|
|
61
|
+
"openrouter/openai/gpt-3.5-turbo-16k",
|
|
62
|
+
"openrouter/anthropic/claude-instant-v1",
|
|
63
|
+
"openrouter/google/palm-2-chat-bison",
|
|
64
|
+
"openrouter/google/palm-2-codechat-bison",
|
|
65
|
+
"openrouter/meta-llama/llama-2-13b-chat",
|
|
66
|
+
"openrouter/meta-llama/llama-2-70b-chat",
|
|
67
|
+
],
|
|
68
|
+
"Perplexity AI": [
|
|
69
|
+
"perplexity/pplx-7b-chat",
|
|
70
|
+
"perplexity/pplx-70b-chat",
|
|
71
|
+
"perplexity/pplx-7b-online",
|
|
72
|
+
"perplexity/pplx-70b-online",
|
|
73
|
+
],
|
|
49
74
|
"Together AI": [
|
|
50
75
|
"together_ai/togethercomputer/llama-2-70b-chat",
|
|
51
76
|
"together_ai/togethercomputer/llama-2-70b",
|
|
@@ -59,26 +84,12 @@ supported_llm_models = {
|
|
|
59
84
|
"together_ai/NousResearch/Nous-Hermes-Llama2-13b",
|
|
60
85
|
"together_ai/Austism/chronos-hermes-13b",
|
|
61
86
|
],
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
"openrouter/openai/gpt-3.5-turbo-16k",
|
|
71
|
-
"openrouter/anthropic/claude-instant-v1",
|
|
72
|
-
"openrouter/google/palm-2-chat-bison",
|
|
73
|
-
"openrouter/google/palm-2-codechat-bison",
|
|
74
|
-
"openrouter/meta-llama/llama-2-13b-chat",
|
|
75
|
-
"openrouter/meta-llama/llama-2-70b-chat",
|
|
76
|
-
],
|
|
77
|
-
"Groq": [
|
|
78
|
-
"groq/llama3-8b-8192",
|
|
79
|
-
"groq/llama3-70b-8192",
|
|
80
|
-
"groq/llama2-70b-4096",
|
|
81
|
-
"groq/mixtral-8x7b-32768",
|
|
82
|
-
"groq/gemma-7b-it",
|
|
83
|
-
],
|
|
87
|
+
}
|
|
88
|
+
|
|
89
|
+
providers_list = list(supported_llm_models.keys())
|
|
90
|
+
|
|
91
|
+
model_to_provider_mapping = {
|
|
92
|
+
model: provider
|
|
93
|
+
for provider, models in supported_llm_models.items()
|
|
94
|
+
for model in models
|
|
84
95
|
}
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
from typing import Optional
|
|
2
|
+
|
|
3
|
+
from contextlib import contextmanager
|
|
4
|
+
from contextvars import ContextVar
|
|
5
|
+
|
|
6
|
+
from pydantic import BaseModel
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class ExportingContext(BaseModel):
|
|
10
|
+
credentials: Optional[str] = None
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
exporting_context = ContextVar("exporting_context", default=ExportingContext())
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@contextmanager
|
|
17
|
+
def exporting_context_manager(
|
|
18
|
+
*,
|
|
19
|
+
context: Optional[ExportingContext] = None,
|
|
20
|
+
):
|
|
21
|
+
token = exporting_context.set(context)
|
|
22
|
+
try:
|
|
23
|
+
yield
|
|
24
|
+
finally:
|
|
25
|
+
exporting_context.reset(token)
|
agenta/sdk/context/routing.py
CHANGED
|
@@ -1,24 +1,24 @@
|
|
|
1
|
+
from typing import Any, Dict, List, Optional
|
|
2
|
+
|
|
1
3
|
from contextlib import contextmanager
|
|
2
4
|
from contextvars import ContextVar
|
|
3
|
-
from typing import Any, Dict, Optional
|
|
4
5
|
|
|
5
|
-
|
|
6
|
+
from pydantic import BaseModel
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class RoutingContext(BaseModel):
|
|
10
|
+
parameters: Optional[Dict[str, Any]] = None
|
|
11
|
+
secrets: Optional[List[Any]] = None
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
routing_context = ContextVar("routing_context", default=RoutingContext())
|
|
6
15
|
|
|
7
16
|
|
|
8
17
|
@contextmanager
|
|
9
18
|
def routing_context_manager(
|
|
10
19
|
*,
|
|
11
|
-
|
|
12
|
-
application: Optional[Dict[str, Any]] = None,
|
|
13
|
-
variant: Optional[Dict[str, Any]] = None,
|
|
14
|
-
environment: Optional[Dict[str, Any]] = None,
|
|
20
|
+
context: Optional[RoutingContext] = None,
|
|
15
21
|
):
|
|
16
|
-
context = {
|
|
17
|
-
"config": config,
|
|
18
|
-
"application": application,
|
|
19
|
-
"variant": variant,
|
|
20
|
-
"environment": environment,
|
|
21
|
-
}
|
|
22
22
|
token = routing_context.set(context)
|
|
23
23
|
try:
|
|
24
24
|
yield
|
agenta/sdk/context/tracing.py
CHANGED
|
@@ -1,3 +1,28 @@
|
|
|
1
|
+
from typing import Any, Dict, Optional
|
|
2
|
+
|
|
3
|
+
from contextlib import contextmanager
|
|
1
4
|
from contextvars import ContextVar
|
|
2
5
|
|
|
3
|
-
|
|
6
|
+
from pydantic import BaseModel
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class TracingContext(BaseModel):
|
|
10
|
+
credentials: Optional[str] = None
|
|
11
|
+
parameters: Optional[Dict[str, Any]] = None
|
|
12
|
+
references: Optional[Dict[str, Any]] = None
|
|
13
|
+
link: Optional[Dict[str, Any]] = None
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
tracing_context = ContextVar("tracing_context", default=TracingContext())
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@contextmanager
|
|
20
|
+
def tracing_context_manager(
|
|
21
|
+
*,
|
|
22
|
+
context: Optional[TracingContext] = None,
|
|
23
|
+
):
|
|
24
|
+
token = tracing_context.set(context)
|
|
25
|
+
try:
|
|
26
|
+
yield
|
|
27
|
+
finally:
|
|
28
|
+
tracing_context.reset(token)
|