aiqa-client 0.4.7__tar.gz → 0.5.2__tar.gz
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.
- {aiqa_client-0.4.7/aiqa_client.egg-info → aiqa_client-0.5.2}/PKG-INFO +1 -1
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/__init__.py +1 -1
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/client.py +97 -15
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/constants.py +1 -1
- aiqa_client-0.5.2/aiqa/http_utils.py +143 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/tracing.py +48 -20
- {aiqa_client-0.4.7 → aiqa_client-0.5.2/aiqa_client.egg-info}/PKG-INFO +1 -1
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa_client.egg-info/SOURCES.txt +2 -1
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/pyproject.toml +1 -1
- aiqa_client-0.5.2/tests/test_api_key.py +96 -0
- aiqa_client-0.5.2/tests/test_integration.py +285 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/tests/test_object_serialiser.py +0 -33
- aiqa_client-0.5.2/tests/test_startup_reliability.py +108 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/tests/test_tracing.py +188 -5
- aiqa_client-0.4.7/aiqa/aiqa_exporter.py +0 -772
- aiqa_client-0.4.7/aiqa/http_utils.py +0 -69
- aiqa_client-0.4.7/tests/test_startup_reliability.py +0 -249
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/LICENSE.txt +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/MANIFEST.in +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/README.md +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/experiment_runner.py +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/object_serialiser.py +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/py.typed +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa/tracing_llm_utils.py +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa_client.egg-info/dependency_links.txt +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa_client.egg-info/requires.txt +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/aiqa_client.egg-info/top_level.txt +0 -0
- {aiqa_client-0.4.7 → aiqa_client-0.5.2}/setup.cfg +0 -0
|
@@ -6,7 +6,7 @@ The client initializes automatically when WithTracing is first used.
|
|
|
6
6
|
|
|
7
7
|
Set environment variables:
|
|
8
8
|
AIQA_SERVER_URL: URL of the AIQA server
|
|
9
|
-
AIQA_API_KEY: API key for authentication
|
|
9
|
+
AIQA_API_KEY: API key for authentication (required)
|
|
10
10
|
AIQA_COMPONENT_TAG: Optional component identifier
|
|
11
11
|
AIQA_STARTUP_DELAY_SECONDS: Optional delay before first flush (default: 10s)
|
|
12
12
|
|
|
@@ -2,10 +2,12 @@
|
|
|
2
2
|
import os
|
|
3
3
|
import logging
|
|
4
4
|
from functools import lru_cache
|
|
5
|
-
from typing import Optional, TYPE_CHECKING, Any
|
|
5
|
+
from typing import Optional, TYPE_CHECKING, Any, Dict
|
|
6
6
|
from opentelemetry import trace
|
|
7
7
|
from opentelemetry.sdk.trace import TracerProvider
|
|
8
8
|
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
|
9
|
+
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
|
10
|
+
import requests
|
|
9
11
|
|
|
10
12
|
from .constants import AIQA_TRACER_NAME, LOG_TAG
|
|
11
13
|
|
|
@@ -30,7 +32,7 @@ except ImportError:
|
|
|
30
32
|
# Set to None so we can check later
|
|
31
33
|
TraceIdRatioBased = None
|
|
32
34
|
|
|
33
|
-
from .http_utils import get_server_url, get_api_key
|
|
35
|
+
from .http_utils import get_server_url, get_api_key, build_headers, format_http_error
|
|
34
36
|
|
|
35
37
|
class AIQAClient:
|
|
36
38
|
"""
|
|
@@ -100,14 +102,14 @@ class AIQAClient:
|
|
|
100
102
|
self.enabled = False
|
|
101
103
|
if self._provider:
|
|
102
104
|
self._provider.shutdown()
|
|
103
|
-
|
|
104
|
-
self._exporter.shutdown()
|
|
105
|
+
# OTLP exporter doesn't have a separate shutdown method - it's handled by the provider
|
|
105
106
|
except Exception as e:
|
|
106
107
|
logger.error(f"Error shutting down tracing: {e}")
|
|
107
108
|
# Still disable even if shutdown had errors
|
|
108
109
|
self.enabled = False
|
|
109
110
|
|
|
110
111
|
|
|
112
|
+
|
|
111
113
|
# Global singleton instance
|
|
112
114
|
client: AIQAClient = AIQAClient()
|
|
113
115
|
|
|
@@ -176,11 +178,10 @@ def _init_tracing() -> None:
|
|
|
176
178
|
server_url = get_server_url()
|
|
177
179
|
api_key = get_api_key()
|
|
178
180
|
|
|
179
|
-
if not
|
|
181
|
+
if not api_key:
|
|
180
182
|
client.enabled = False
|
|
181
|
-
missing_vars = [var for var, val in [("AIQA_SERVER_URL", server_url), ("AIQA_API_KEY", api_key)] if not val]
|
|
182
183
|
logger.warning(
|
|
183
|
-
f"AIQA tracing is disabled: missing required environment variables:
|
|
184
|
+
f"AIQA tracing is disabled: missing required environment variables: AIQA_API_KEY"
|
|
184
185
|
)
|
|
185
186
|
client._initialized = True
|
|
186
187
|
return
|
|
@@ -227,20 +228,39 @@ def _init_tracing() -> None:
|
|
|
227
228
|
|
|
228
229
|
def _attach_aiqa_processor(provider: TracerProvider) -> None:
|
|
229
230
|
"""Attach AIQA span processor to the provider. Idempotent - safe to call multiple times."""
|
|
230
|
-
from .aiqa_exporter import AIQASpanExporter
|
|
231
|
-
|
|
232
231
|
try:
|
|
233
232
|
# Check if already attached
|
|
234
233
|
for p in provider._active_span_processor._span_processors:
|
|
235
|
-
if isinstance(getattr(p, "exporter", None),
|
|
234
|
+
if isinstance(getattr(p, "exporter", None), OTLPSpanExporter):
|
|
236
235
|
logger.debug(f"AIQA span processor already attached, skipping")
|
|
237
236
|
return
|
|
238
237
|
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
|
|
238
|
+
server_url = get_server_url()
|
|
239
|
+
api_key = get_api_key()
|
|
240
|
+
|
|
241
|
+
# Build headers for authentication
|
|
242
|
+
# OTLP exporter sets its own Content-Type, so we only need Authorization
|
|
243
|
+
auth_headers = {}
|
|
244
|
+
if api_key:
|
|
245
|
+
auth_headers["Authorization"] = f"ApiKey {api_key}"
|
|
246
|
+
elif os.getenv("AIQA_API_KEY"):
|
|
247
|
+
auth_headers["Authorization"] = f"ApiKey {os.getenv('AIQA_API_KEY')}"
|
|
248
|
+
|
|
249
|
+
# OTLP HTTP exporter requires the full endpoint URL including /v1/traces
|
|
250
|
+
# Ensure server_url doesn't have trailing slash or /v1/traces, then append /v1/traces
|
|
251
|
+
base_url = server_url.rstrip('/')
|
|
252
|
+
if base_url.endswith('/v1/traces'):
|
|
253
|
+
endpoint = base_url
|
|
254
|
+
else:
|
|
255
|
+
endpoint = f"{base_url}/v1/traces"
|
|
256
|
+
|
|
257
|
+
# Create OTLP exporter with authentication headers only
|
|
258
|
+
# The exporter will set Content-Type and other headers automatically
|
|
259
|
+
exporter = OTLPSpanExporter(
|
|
260
|
+
endpoint=endpoint,
|
|
261
|
+
headers=auth_headers if auth_headers else None,
|
|
243
262
|
)
|
|
263
|
+
|
|
244
264
|
provider.add_span_processor(BatchSpanProcessor(exporter))
|
|
245
265
|
global client
|
|
246
266
|
client.exporter = exporter
|
|
@@ -266,4 +286,66 @@ def get_aiqa_tracer() -> trace.Tracer:
|
|
|
266
286
|
except Exception as e:
|
|
267
287
|
# Log issue but still return a tracer
|
|
268
288
|
logger.info(f"Issue getting AIQA tracer with version: {e}, using fallback")
|
|
269
|
-
return trace.get_tracer(AIQA_TRACER_NAME)
|
|
289
|
+
return trace.get_tracer(AIQA_TRACER_NAME)
|
|
290
|
+
|
|
291
|
+
|
|
292
|
+
def get_organisation(
|
|
293
|
+
organisation_id: str,
|
|
294
|
+
server_url: Optional[str] = None,
|
|
295
|
+
api_key: Optional[str] = None
|
|
296
|
+
) -> Dict[str, Any]:
|
|
297
|
+
"""
|
|
298
|
+
Get organisation information based on API key via an API call.
|
|
299
|
+
|
|
300
|
+
Args:
|
|
301
|
+
organisation_id: ID of the organisation to retrieve
|
|
302
|
+
server_url: Optional server URL (defaults to AIQA_SERVER_URL env var)
|
|
303
|
+
api_key: Optional API key (defaults to AIQA_API_KEY env var)
|
|
304
|
+
|
|
305
|
+
Returns:
|
|
306
|
+
Organisation object as a dictionary
|
|
307
|
+
"""
|
|
308
|
+
url = get_server_url(server_url)
|
|
309
|
+
key = get_api_key(api_key)
|
|
310
|
+
headers = build_headers(key)
|
|
311
|
+
|
|
312
|
+
response = requests.get(
|
|
313
|
+
f"{url}/organisation/{organisation_id}",
|
|
314
|
+
headers=headers,
|
|
315
|
+
)
|
|
316
|
+
|
|
317
|
+
if not response.ok:
|
|
318
|
+
raise Exception(format_http_error(response, "get organisation"))
|
|
319
|
+
|
|
320
|
+
return response.json()
|
|
321
|
+
|
|
322
|
+
|
|
323
|
+
def get_api_key_info(
|
|
324
|
+
api_key_id: str,
|
|
325
|
+
server_url: Optional[str] = None,
|
|
326
|
+
api_key: Optional[str] = None
|
|
327
|
+
) -> Dict[str, Any]:
|
|
328
|
+
"""
|
|
329
|
+
Get API key information via an API call.
|
|
330
|
+
|
|
331
|
+
Args:
|
|
332
|
+
api_key_id: ID of the API key to retrieve
|
|
333
|
+
server_url: Optional server URL (defaults to AIQA_SERVER_URL env var)
|
|
334
|
+
api_key: Optional API key (defaults to AIQA_API_KEY env var)
|
|
335
|
+
|
|
336
|
+
Returns:
|
|
337
|
+
ApiKey object as a dictionary
|
|
338
|
+
"""
|
|
339
|
+
url = get_server_url(server_url)
|
|
340
|
+
key = get_api_key(api_key)
|
|
341
|
+
headers = build_headers(key)
|
|
342
|
+
|
|
343
|
+
response = requests.get(
|
|
344
|
+
f"{url}/api-key/{api_key_id}",
|
|
345
|
+
headers=headers,
|
|
346
|
+
)
|
|
347
|
+
|
|
348
|
+
if not response.ok:
|
|
349
|
+
raise Exception(format_http_error(response, "get api key info"))
|
|
350
|
+
|
|
351
|
+
return response.json()
|
|
@@ -3,6 +3,6 @@ Constants used across the AIQA client package.
|
|
|
3
3
|
"""
|
|
4
4
|
|
|
5
5
|
AIQA_TRACER_NAME = "aiqa-tracer"
|
|
6
|
-
VERSION = "0.
|
|
6
|
+
VERSION = "0.5.2" # automatically updated by set-version-json.sh
|
|
7
7
|
|
|
8
8
|
LOG_TAG = "AIQA" # Used in all logging output to identify AIQA messages
|
|
@@ -0,0 +1,143 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Shared HTTP utilities for AIQA client.
|
|
3
|
+
Provides common functions for building headers, handling errors, and accessing environment variables.
|
|
4
|
+
Supports AIQA-specific env vars (AIQA_SERVER_URL, AIQA_API_KEY) with fallback to OTLP standard vars.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import os
|
|
8
|
+
from typing import Dict, Optional
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def build_headers(api_key: Optional[str] = None) -> Dict[str, str]:
|
|
12
|
+
"""
|
|
13
|
+
Build HTTP headers for AIQA API requests.
|
|
14
|
+
|
|
15
|
+
Checks AIQA_API_KEY first, then falls back to OTEL_EXPORTER_OTLP_HEADERS if not set.
|
|
16
|
+
|
|
17
|
+
Args:
|
|
18
|
+
api_key: Optional API key. If not provided, will try to get from AIQA_API_KEY env var,
|
|
19
|
+
then from OTEL_EXPORTER_OTLP_HEADERS.
|
|
20
|
+
|
|
21
|
+
Returns:
|
|
22
|
+
Dictionary with Content-Type, Accept-Encoding, and optionally Authorization header.
|
|
23
|
+
"""
|
|
24
|
+
headers = {
|
|
25
|
+
"Content-Type": "application/json",
|
|
26
|
+
"Accept-Encoding": "gzip, deflate, br", # Request compression (aiohttp handles decompression automatically)
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
# Check parameter first
|
|
30
|
+
if api_key:
|
|
31
|
+
headers["Authorization"] = f"ApiKey {api_key}"
|
|
32
|
+
return headers
|
|
33
|
+
|
|
34
|
+
# Check AIQA_API_KEY env var
|
|
35
|
+
aiqa_api_key = os.getenv("AIQA_API_KEY")
|
|
36
|
+
if aiqa_api_key:
|
|
37
|
+
headers["Authorization"] = f"ApiKey {aiqa_api_key}"
|
|
38
|
+
return headers
|
|
39
|
+
|
|
40
|
+
# Fallback to OTLP headers (format: "key1=value1,key2=value2")
|
|
41
|
+
otlp_headers = os.getenv("OTEL_EXPORTER_OTLP_HEADERS")
|
|
42
|
+
if otlp_headers:
|
|
43
|
+
# Parse comma-separated key=value pairs
|
|
44
|
+
for header_pair in otlp_headers.split(","):
|
|
45
|
+
header_pair = header_pair.strip()
|
|
46
|
+
if "=" in header_pair:
|
|
47
|
+
key, value = header_pair.split("=", 1)
|
|
48
|
+
key = key.strip()
|
|
49
|
+
value = value.strip()
|
|
50
|
+
if key.lower() == "authorization":
|
|
51
|
+
headers["Authorization"] = value
|
|
52
|
+
else:
|
|
53
|
+
headers[key] = value
|
|
54
|
+
|
|
55
|
+
return headers
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def get_server_url(server_url: Optional[str] = None) -> str:
|
|
59
|
+
"""
|
|
60
|
+
Get server URL from parameter or environment variable, with trailing slash removed.
|
|
61
|
+
|
|
62
|
+
Checks AIQA_SERVER_URL first, then falls back to OTEL_EXPORTER_OTLP_ENDPOINT if not set.
|
|
63
|
+
|
|
64
|
+
Args:
|
|
65
|
+
server_url: Optional server URL. If not provided, will get from AIQA_SERVER_URL env var,
|
|
66
|
+
then from OTEL_EXPORTER_OTLP_ENDPOINT.
|
|
67
|
+
|
|
68
|
+
Returns:
|
|
69
|
+
Server URL with trailing slash removed. Defaults to https://server-aiqa.winterwell.com if not set.
|
|
70
|
+
"""
|
|
71
|
+
# Check parameter first
|
|
72
|
+
if server_url:
|
|
73
|
+
return server_url.rstrip("/")
|
|
74
|
+
|
|
75
|
+
# Check AIQA_SERVER_URL env var
|
|
76
|
+
url = os.getenv("AIQA_SERVER_URL")
|
|
77
|
+
if url:
|
|
78
|
+
return url.rstrip("/")
|
|
79
|
+
|
|
80
|
+
# Fallback to OTLP endpoint
|
|
81
|
+
url = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
|
|
82
|
+
if url:
|
|
83
|
+
return url.rstrip("/")
|
|
84
|
+
|
|
85
|
+
# Default fallback
|
|
86
|
+
return "https://server-aiqa.winterwell.com"
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def get_api_key(api_key: Optional[str] = None) -> str:
|
|
90
|
+
"""
|
|
91
|
+
Get API key from parameter or environment variable.
|
|
92
|
+
|
|
93
|
+
Checks AIQA_API_KEY first, then falls back to OTEL_EXPORTER_OTLP_HEADERS if not set.
|
|
94
|
+
|
|
95
|
+
Args:
|
|
96
|
+
api_key: Optional API key. If not provided, will get from AIQA_API_KEY env var,
|
|
97
|
+
then from OTEL_EXPORTER_OTLP_HEADERS (looking for Authorization header).
|
|
98
|
+
|
|
99
|
+
Returns:
|
|
100
|
+
API key or empty string if not set.
|
|
101
|
+
"""
|
|
102
|
+
# Check parameter first
|
|
103
|
+
if api_key:
|
|
104
|
+
return api_key
|
|
105
|
+
|
|
106
|
+
# Check AIQA_API_KEY env var
|
|
107
|
+
aiqa_api_key = os.getenv("AIQA_API_KEY")
|
|
108
|
+
if aiqa_api_key:
|
|
109
|
+
return aiqa_api_key
|
|
110
|
+
|
|
111
|
+
# Fallback to OTLP headers (look for Authorization header)
|
|
112
|
+
otlp_headers = os.getenv("OTEL_EXPORTER_OTLP_HEADERS")
|
|
113
|
+
if otlp_headers:
|
|
114
|
+
for header_pair in otlp_headers.split(","):
|
|
115
|
+
header_pair = header_pair.strip()
|
|
116
|
+
if "=" in header_pair:
|
|
117
|
+
key, value = header_pair.split("=", 1)
|
|
118
|
+
key = key.strip()
|
|
119
|
+
value = value.strip()
|
|
120
|
+
if key.lower() == "authorization":
|
|
121
|
+
# Extract API key from "ApiKey <key>" or just return the value
|
|
122
|
+
if value.startswith("ApiKey "):
|
|
123
|
+
return value[7:]
|
|
124
|
+
return value
|
|
125
|
+
|
|
126
|
+
return ""
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def format_http_error(response, operation: str) -> str:
|
|
130
|
+
"""
|
|
131
|
+
Format an HTTP error message from a response object.
|
|
132
|
+
|
|
133
|
+
Args:
|
|
134
|
+
response: Response object with status_code, reason, and text attributes
|
|
135
|
+
operation: Description of the operation that failed (e.g., "fetch dataset")
|
|
136
|
+
|
|
137
|
+
Returns:
|
|
138
|
+
Formatted error message string.
|
|
139
|
+
"""
|
|
140
|
+
error_text = response.text if hasattr(response, "text") else "Unknown error"
|
|
141
|
+
status_code = getattr(response, "status_code", getattr(response, "status", "unknown"))
|
|
142
|
+
reason = getattr(response, "reason", "")
|
|
143
|
+
return f"Failed to {operation}: {status_code} {reason} - {error_text}"
|
|
@@ -8,13 +8,13 @@ import logging
|
|
|
8
8
|
import inspect
|
|
9
9
|
import os
|
|
10
10
|
import copy
|
|
11
|
+
import requests
|
|
11
12
|
from typing import Any, Callable, Optional, List
|
|
12
13
|
from functools import wraps
|
|
13
14
|
from opentelemetry import trace
|
|
14
15
|
from opentelemetry.sdk.trace import TracerProvider
|
|
15
16
|
from opentelemetry.trace import Status, StatusCode, SpanContext, TraceFlags
|
|
16
17
|
from opentelemetry.propagate import inject, extract
|
|
17
|
-
from .aiqa_exporter import AIQASpanExporter
|
|
18
18
|
from .client import get_aiqa_client, get_component_tag, set_component_tag as _set_component_tag, get_aiqa_tracer
|
|
19
19
|
from .constants import AIQA_TRACER_NAME, LOG_TAG
|
|
20
20
|
from .object_serialiser import serialize_for_span
|
|
@@ -31,13 +31,11 @@ async def flush_tracing() -> None:
|
|
|
31
31
|
if you want to flush immediately, e.g. before exiting a process.
|
|
32
32
|
A common use is if you are tracing unit tests or experiment runs.
|
|
33
33
|
|
|
34
|
-
This flushes
|
|
34
|
+
This flushes the BatchSpanProcessor (OTLP exporter doesn't have a separate flush method).
|
|
35
35
|
"""
|
|
36
36
|
client = get_aiqa_client()
|
|
37
37
|
if client.provider:
|
|
38
38
|
client.provider.force_flush() # Synchronous method
|
|
39
|
-
if client.exporter:
|
|
40
|
-
await client.exporter.flush()
|
|
41
39
|
|
|
42
40
|
|
|
43
41
|
# Export provider and exporter accessors for advanced usage
|
|
@@ -117,12 +115,10 @@ class TracingOptions:
|
|
|
117
115
|
self.filter_output = filter_output
|
|
118
116
|
|
|
119
117
|
|
|
120
|
-
def _prepare_input(args: tuple, kwargs: dict) -> Any:
|
|
118
|
+
def _prepare_input(args: tuple, kwargs: dict, sig: Optional[inspect.Signature] = None) -> Any:
|
|
121
119
|
"""Prepare input for span attributes.
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
For single-arg-dicts or kwargs-only, returns a shallow copy of the input data.
|
|
120
|
+
Converts args and kwargs into a unified dict structure using function signature when available.
|
|
121
|
+
Falls back to legacy behavior for functions without inspectable signatures.
|
|
126
122
|
|
|
127
123
|
Note: This function does NOT serialize values - it just structures the data.
|
|
128
124
|
Serialization happens later via serialize_for_span() to avoid double-encoding
|
|
@@ -130,6 +126,22 @@ def _prepare_input(args: tuple, kwargs: dict) -> Any:
|
|
|
130
126
|
"""
|
|
131
127
|
if not args and not kwargs:
|
|
132
128
|
return None
|
|
129
|
+
|
|
130
|
+
# Try to bind args to parameter names using function signature
|
|
131
|
+
if sig is not None:
|
|
132
|
+
try:
|
|
133
|
+
bound = sig.bind(*args, **kwargs)
|
|
134
|
+
bound.apply_defaults()
|
|
135
|
+
# Return dict of all arguments (positional args are now named)
|
|
136
|
+
result = bound.arguments.copy()
|
|
137
|
+
# Shallow copy to protect against mutating the input
|
|
138
|
+
return result
|
|
139
|
+
except (TypeError, ValueError):
|
|
140
|
+
# Binding failed (e.g., wrong number of args, *args/**kwargs issues)
|
|
141
|
+
# Fall through to legacy behavior
|
|
142
|
+
pass
|
|
143
|
+
|
|
144
|
+
# in case binding fails
|
|
133
145
|
if not kwargs:
|
|
134
146
|
if len(args) == 1:
|
|
135
147
|
arg0 = args[0]
|
|
@@ -150,15 +162,17 @@ def _prepare_and_filter_input(
|
|
|
150
162
|
kwargs: dict,
|
|
151
163
|
filter_input: Optional[Callable[[Any], Any]],
|
|
152
164
|
ignore_input: Optional[List[str]],
|
|
165
|
+
sig: Optional[inspect.Signature] = None,
|
|
153
166
|
) -> Any:
|
|
154
167
|
"""
|
|
155
168
|
Prepare and filter input for span attributes - applies the user's filter_input and ignore_input.
|
|
156
|
-
|
|
169
|
+
Converts all args to a dict using function signature when available.
|
|
157
170
|
"""
|
|
158
171
|
# Handle "self" in ignore_input by skipping the first argument
|
|
159
172
|
filtered_args = args
|
|
160
173
|
filtered_kwargs = kwargs.copy() if kwargs else {}
|
|
161
174
|
filtered_ignore_input = ignore_input
|
|
175
|
+
filtered_sig = sig
|
|
162
176
|
if ignore_input and "self" in ignore_input:
|
|
163
177
|
# Remove "self" from ignore_input list (we'll handle it specially)
|
|
164
178
|
filtered_ignore_input = [key for key in ignore_input if key != "self"]
|
|
@@ -168,8 +182,14 @@ def _prepare_and_filter_input(
|
|
|
168
182
|
# Also remove "self" from kwargs if present
|
|
169
183
|
if "self" in filtered_kwargs:
|
|
170
184
|
del filtered_kwargs["self"]
|
|
171
|
-
|
|
172
|
-
|
|
185
|
+
# Adjust signature to remove "self" parameter if present
|
|
186
|
+
# This is needed because we removed self from args, so signature binding will fail otherwise
|
|
187
|
+
if filtered_sig is not None:
|
|
188
|
+
params = list(filtered_sig.parameters.values())
|
|
189
|
+
if params and params[0].name == "self":
|
|
190
|
+
filtered_sig = filtered_sig.replace(parameters=params[1:])
|
|
191
|
+
# turn args, kwargs into one "nice" object (now always a dict when signature is available)
|
|
192
|
+
input_data = _prepare_input(filtered_args, filtered_kwargs, filtered_sig)
|
|
173
193
|
if filter_input and input_data is not None:
|
|
174
194
|
input_data = filter_input(input_data)
|
|
175
195
|
if filtered_ignore_input and len(filtered_ignore_input) > 0:
|
|
@@ -447,6 +467,15 @@ def WithTracing(
|
|
|
447
467
|
is_generator = inspect.isgeneratorfunction(fn)
|
|
448
468
|
is_async_generator = inspect.isasyncgenfunction(fn) if hasattr(inspect, 'isasyncgenfunction') else False
|
|
449
469
|
|
|
470
|
+
# Get function signature once at decoration time for efficient arg name resolution
|
|
471
|
+
fn_sig: Optional[inspect.Signature] = None
|
|
472
|
+
try:
|
|
473
|
+
fn_sig = inspect.signature(fn)
|
|
474
|
+
except (ValueError, TypeError):
|
|
475
|
+
# Some callables (e.g., builtins, C extensions) don't have inspectable signatures
|
|
476
|
+
# Will fall back to legacy behavior
|
|
477
|
+
pass
|
|
478
|
+
|
|
450
479
|
# Don't get tracer here - get it lazily when function is called
|
|
451
480
|
# This ensures initialization only happens when tracing is actually used
|
|
452
481
|
|
|
@@ -595,7 +624,7 @@ def WithTracing(
|
|
|
595
624
|
if is_async_generator:
|
|
596
625
|
@wraps(fn)
|
|
597
626
|
async def async_gen_traced_fn(*args, **kwargs):
|
|
598
|
-
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
|
|
627
|
+
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
|
|
599
628
|
return await _execute_generator_async(
|
|
600
629
|
lambda: fn(*args, **kwargs),
|
|
601
630
|
input_data
|
|
@@ -607,7 +636,7 @@ def WithTracing(
|
|
|
607
636
|
elif is_generator:
|
|
608
637
|
@wraps(fn)
|
|
609
638
|
def gen_traced_fn(*args, **kwargs):
|
|
610
|
-
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
|
|
639
|
+
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
|
|
611
640
|
return _execute_generator_sync(
|
|
612
641
|
lambda: fn(*args, **kwargs),
|
|
613
642
|
input_data
|
|
@@ -619,7 +648,7 @@ def WithTracing(
|
|
|
619
648
|
elif is_async:
|
|
620
649
|
@wraps(fn)
|
|
621
650
|
async def async_traced_fn(*args, **kwargs):
|
|
622
|
-
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
|
|
651
|
+
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
|
|
623
652
|
return await _execute_with_span_async(
|
|
624
653
|
lambda: fn(*args, **kwargs),
|
|
625
654
|
input_data
|
|
@@ -631,7 +660,7 @@ def WithTracing(
|
|
|
631
660
|
else:
|
|
632
661
|
@wraps(fn)
|
|
633
662
|
def sync_traced_fn(*args, **kwargs):
|
|
634
|
-
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
|
|
663
|
+
input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
|
|
635
664
|
return _execute_with_span_sync(
|
|
636
665
|
lambda: fn(*args, **kwargs),
|
|
637
666
|
input_data
|
|
@@ -678,6 +707,7 @@ def get_active_span() -> Optional[trace.Span]:
|
|
|
678
707
|
|
|
679
708
|
def set_conversation_id(conversation_id: str) -> bool:
|
|
680
709
|
"""
|
|
710
|
+
Naturally a conversation might span several traces.
|
|
681
711
|
Set the gen_ai.conversation.id attribute on the active span.
|
|
682
712
|
This allows you to group multiple traces together that are part of the same conversation.
|
|
683
713
|
See https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/ for more details.
|
|
@@ -1027,14 +1057,12 @@ def get_span(span_id: str, organisation_id: Optional[str] = None, exclude: Optio
|
|
|
1027
1057
|
print(f"Found span: {span['name']}")
|
|
1028
1058
|
my_function(**span['input'])
|
|
1029
1059
|
"""
|
|
1030
|
-
import os
|
|
1031
|
-
import requests
|
|
1032
|
-
|
|
1033
1060
|
server_url = get_server_url()
|
|
1034
1061
|
api_key = get_api_key()
|
|
1035
1062
|
org_id = organisation_id or os.getenv("AIQA_ORGANISATION_ID", "")
|
|
1036
1063
|
|
|
1037
|
-
if not
|
|
1064
|
+
# Check if server_url is the default (meaning AIQA_SERVER_URL was not set)
|
|
1065
|
+
if not os.getenv("AIQA_SERVER_URL"):
|
|
1038
1066
|
raise ValueError("AIQA_SERVER_URL is not set. Cannot retrieve span.")
|
|
1039
1067
|
if not org_id:
|
|
1040
1068
|
raise ValueError("Organisation ID is required. Provide it as parameter or set AIQA_ORGANISATION_ID environment variable.")
|
|
@@ -3,7 +3,6 @@ MANIFEST.in
|
|
|
3
3
|
README.md
|
|
4
4
|
pyproject.toml
|
|
5
5
|
aiqa/__init__.py
|
|
6
|
-
aiqa/aiqa_exporter.py
|
|
7
6
|
aiqa/client.py
|
|
8
7
|
aiqa/constants.py
|
|
9
8
|
aiqa/experiment_runner.py
|
|
@@ -17,6 +16,8 @@ aiqa_client.egg-info/SOURCES.txt
|
|
|
17
16
|
aiqa_client.egg-info/dependency_links.txt
|
|
18
17
|
aiqa_client.egg-info/requires.txt
|
|
19
18
|
aiqa_client.egg-info/top_level.txt
|
|
19
|
+
tests/test_api_key.py
|
|
20
|
+
tests/test_integration.py
|
|
20
21
|
tests/test_object_serialiser.py
|
|
21
22
|
tests/test_startup_reliability.py
|
|
22
23
|
tests/test_tracing.py
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "aiqa-client"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.5.2"
|
|
8
8
|
description = "OpenTelemetry-based Python client for tracing functions and sending traces to the AIQA server"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.8"
|
|
@@ -0,0 +1,96 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Integration test for get_api_key_info functionality.
|
|
3
|
+
|
|
4
|
+
This test verifies that get_api_key_info works correctly with environment variables
|
|
5
|
+
loaded from .env files.
|
|
6
|
+
|
|
7
|
+
Prerequisites:
|
|
8
|
+
- AIQA server must be running and accessible
|
|
9
|
+
- Set AIQA_SERVER_URL and AIQA_API_KEY environment variables in .env
|
|
10
|
+
- Server must have PostgreSQL configured
|
|
11
|
+
- The API key must be valid and associated with an organisation
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
import os
|
|
15
|
+
import pytest
|
|
16
|
+
import requests
|
|
17
|
+
from typing import Optional
|
|
18
|
+
from aiqa.client import get_api_key_info
|
|
19
|
+
from aiqa.http_utils import get_server_url, get_api_key, build_headers
|
|
20
|
+
from dotenv import load_dotenv
|
|
21
|
+
|
|
22
|
+
load_dotenv()
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def get_api_key_id_from_list() -> Optional[str]:
|
|
26
|
+
"""
|
|
27
|
+
Helper function to get an API key ID by listing API keys.
|
|
28
|
+
|
|
29
|
+
Returns:
|
|
30
|
+
First API key ID found
|
|
31
|
+
"""
|
|
32
|
+
|
|
33
|
+
server_url = get_server_url()
|
|
34
|
+
api_key = get_api_key()
|
|
35
|
+
|
|
36
|
+
url = f"{server_url}/api-key"
|
|
37
|
+
params = {}
|
|
38
|
+
headers = build_headers(api_key)
|
|
39
|
+
|
|
40
|
+
response = requests.get(url, params=params, headers=headers, timeout=5)
|
|
41
|
+
if response.status_code == 200:
|
|
42
|
+
api_keys = response.json()
|
|
43
|
+
if api_keys and len(api_keys) > 0:
|
|
44
|
+
return api_keys[0].get("id")
|
|
45
|
+
return None
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
@pytest.fixture(scope="function")
|
|
49
|
+
def check_server_available():
|
|
50
|
+
"""Fixture that skips tests if server is not available."""
|
|
51
|
+
server_url = os.getenv("AIQA_SERVER_URL")
|
|
52
|
+
api_key = os.getenv("AIQA_API_KEY")
|
|
53
|
+
|
|
54
|
+
if not server_url or not api_key:
|
|
55
|
+
pytest.skip("AIQA_SERVER_URL and AIQA_API_KEY environment variables must be set")
|
|
56
|
+
|
|
57
|
+
# Try to connect to server
|
|
58
|
+
try:
|
|
59
|
+
url = f"{server_url.rstrip('/')}/span"
|
|
60
|
+
headers = build_headers(api_key)
|
|
61
|
+
response = requests.get(
|
|
62
|
+
url,
|
|
63
|
+
params={"q": "name:nonexistent", "limit": "1"},
|
|
64
|
+
headers=headers,
|
|
65
|
+
timeout=5
|
|
66
|
+
)
|
|
67
|
+
# 200 means server is up (even if no results)
|
|
68
|
+
# 401/403 means server is up but auth failed
|
|
69
|
+
# Other errors might mean server is down
|
|
70
|
+
if response.status_code not in [200, 401, 403]:
|
|
71
|
+
pytest.skip(f"Server appears to be down (status {response.status_code})")
|
|
72
|
+
except requests.exceptions.RequestException as e:
|
|
73
|
+
pytest.skip(f"Cannot connect to server: {e}")
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def test_get_api_key_info(check_server_available):
|
|
77
|
+
"""Test that get_api_key_info works correctly with environment variables from .env."""
|
|
78
|
+
# Get an API key ID to test with
|
|
79
|
+
api_key_id = get_api_key_id_from_list()
|
|
80
|
+
|
|
81
|
+
if not api_key_id:
|
|
82
|
+
pytest.skip("AIQA_ORGANISATION_ID must be set to test get_api_key_info, or no API keys found")
|
|
83
|
+
|
|
84
|
+
# Test get_api_key_info - it should load server_url and api_key from .env
|
|
85
|
+
api_key_info = get_api_key_info(api_key_id)
|
|
86
|
+
|
|
87
|
+
# Verify the response structure
|
|
88
|
+
assert api_key_info is not None, "API key info should not be None"
|
|
89
|
+
assert "id" in api_key_info, "API key info should have 'id' field"
|
|
90
|
+
assert api_key_info["id"] == api_key_id, f"API key ID should match: expected {api_key_id}, got {api_key_info['id']}"
|
|
91
|
+
|
|
92
|
+
# Verify other expected fields
|
|
93
|
+
assert "organisation" in api_key_info, "API key info should have 'organisation' field"
|
|
94
|
+
assert "role" in api_key_info, "API key info should have 'role' field"
|
|
95
|
+
assert api_key_info["role"] in ["trace", "developer", "admin"], "API key role should be valid"
|
|
96
|
+
|