ragaai-catalyst 2.2.5b3__py3-none-any.whl → 2.2.5b5__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.
- ragaai_catalyst/tracers/agentic_tracing/upload/session_manager.py +156 -0
- ragaai_catalyst/tracers/agentic_tracing/upload/trace_uploader.py +125 -34
- ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py +74 -36
- ragaai_catalyst/tracers/agentic_tracing/upload/upload_code.py +93 -54
- ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py +72 -26
- ragaai_catalyst/tracers/exporters/ragaai_trace_exporter.py +4 -0
- ragaai_catalyst/tracers/tracer.py +1 -1
- {ragaai_catalyst-2.2.5b3.dist-info → ragaai_catalyst-2.2.5b5.dist-info}/METADATA +1 -1
- {ragaai_catalyst-2.2.5b3.dist-info → ragaai_catalyst-2.2.5b5.dist-info}/RECORD +12 -11
- {ragaai_catalyst-2.2.5b3.dist-info → ragaai_catalyst-2.2.5b5.dist-info}/WHEEL +0 -0
- {ragaai_catalyst-2.2.5b3.dist-info → ragaai_catalyst-2.2.5b5.dist-info}/licenses/LICENSE +0 -0
- {ragaai_catalyst-2.2.5b3.dist-info → ragaai_catalyst-2.2.5b5.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,156 @@
|
|
1
|
+
import logging
|
2
|
+
import os
|
3
|
+
import threading
|
4
|
+
|
5
|
+
from requests.adapters import HTTPAdapter
|
6
|
+
from urllib3.util.retry import Retry
|
7
|
+
from urllib3.exceptions import PoolError, MaxRetryError, NewConnectionError
|
8
|
+
from requests.exceptions import ConnectionError, Timeout
|
9
|
+
from http.client import RemoteDisconnected
|
10
|
+
from ragaai_catalyst import RagaAICatalyst
|
11
|
+
import requests
|
12
|
+
|
13
|
+
logger = logging.getLogger(__name__)
|
14
|
+
|
15
|
+
|
16
|
+
class SessionManager:
|
17
|
+
"""Shared session manager with connection pooling for HTTP requests"""
|
18
|
+
_instance = None
|
19
|
+
_session = None
|
20
|
+
_lock = threading.Lock()
|
21
|
+
|
22
|
+
def __new__(cls):
|
23
|
+
if cls._instance is None:
|
24
|
+
with cls._lock: # Thread-safe singleton
|
25
|
+
if cls._instance is None: # Double-check locking
|
26
|
+
logger.info("Creating new SessionManager singleton instance")
|
27
|
+
cls._instance = super(SessionManager, cls).__new__(cls)
|
28
|
+
cls._instance._initialize_session()
|
29
|
+
else:
|
30
|
+
logger.debug("SessionManager instance already exists, returning existing instance")
|
31
|
+
else:
|
32
|
+
logger.debug("SessionManager instance exists, returning existing instance")
|
33
|
+
return cls._instance
|
34
|
+
|
35
|
+
def _initialize_session(self):
|
36
|
+
"""Initialize session with connection pooling and retry strategy"""
|
37
|
+
logger.info("Initializing HTTP session with connection pooling and retry strategy")
|
38
|
+
self._session = requests.Session()
|
39
|
+
|
40
|
+
retry_strategy = Retry(
|
41
|
+
total=3, # number of retries
|
42
|
+
connect=3, # number of retries for connection-related errors
|
43
|
+
read=3, # number of retries for read-related errors
|
44
|
+
backoff_factor=0.5, # wait 0.5, 1, 2... seconds between retries
|
45
|
+
status_forcelist=[500, 502, 503, 504] # HTTP status codes to retry on
|
46
|
+
)
|
47
|
+
|
48
|
+
adapter = HTTPAdapter(
|
49
|
+
max_retries=retry_strategy,
|
50
|
+
pool_connections=5, # number of connection pools to cache (per host)
|
51
|
+
pool_maxsize=50, # maximum number of connections in each pool
|
52
|
+
pool_block=True # Block/wait when pool is full rather than raising error
|
53
|
+
)
|
54
|
+
|
55
|
+
self._session.mount("http://", adapter)
|
56
|
+
self._session.mount("https://", adapter)
|
57
|
+
|
58
|
+
# Set session-level configuration to handle connection issues
|
59
|
+
self._session.headers.update({
|
60
|
+
'Connection': 'keep-alive',
|
61
|
+
'User-Agent': 'RagaAI-Catalyst/1.0'
|
62
|
+
})
|
63
|
+
|
64
|
+
logger.info("HTTP session initialized successfully with adapters mounted for http:// and https://")
|
65
|
+
|
66
|
+
# Warm up connection pool using RagaAICatalyst.BASE_URL
|
67
|
+
if os.getenv("RAGAAI_CATALYST_BASE_URL") is not None:
|
68
|
+
base_url = os.getenv("RAGAAI_CATALYST_BASE_URL")
|
69
|
+
logger.info(f"Warming up connection pool using RagaAICatalyst.BASE_URL: {base_url}")
|
70
|
+
self.warm_up_connections(base_url)
|
71
|
+
else:
|
72
|
+
logger.warning(f"RAGAAI_CATALYST_BASE_URL not available, skipping connection warmup")
|
73
|
+
|
74
|
+
@property
|
75
|
+
def session(self):
|
76
|
+
if self._session is None:
|
77
|
+
logger.warning("Session accessed but not initialized, reinitializing...")
|
78
|
+
self._initialize_session()
|
79
|
+
return self._session
|
80
|
+
|
81
|
+
def warm_up_connections(self, base_url, num_connections=3):
|
82
|
+
"""
|
83
|
+
Warm up the connection pool by making lightweight requests to healthcheck endpoint.
|
84
|
+
This can help prevent RemoteDisconnected errors on initial requests.
|
85
|
+
"""
|
86
|
+
if not self._session:
|
87
|
+
return
|
88
|
+
|
89
|
+
# Construct healthcheck URL
|
90
|
+
healthcheck_url = f"{base_url.rstrip('/')}/healthcheck"
|
91
|
+
logger.info(f"Warming up connection pool with {num_connections} connections to {healthcheck_url}")
|
92
|
+
|
93
|
+
for i in range(num_connections):
|
94
|
+
try:
|
95
|
+
# Make a lightweight HEAD request to the healthcheck endpoint to warm up the connection
|
96
|
+
response = self._session.head(healthcheck_url, timeout=10)
|
97
|
+
logger.info(f"Warmup connection {i+1}: Status {response.status_code}")
|
98
|
+
except Exception as e:
|
99
|
+
logger.warning(f"Warmup connection {i+1} failed (this may be normal): {e}")
|
100
|
+
# Ignore other failures during warmup as they're expected
|
101
|
+
continue
|
102
|
+
|
103
|
+
logger.info("Connection pool warmup completed")
|
104
|
+
|
105
|
+
def close(self):
|
106
|
+
"""Close the session"""
|
107
|
+
if self._session:
|
108
|
+
logger.info("Closing HTTP session")
|
109
|
+
self._session.close()
|
110
|
+
self._session = None
|
111
|
+
logger.info("HTTP session closed successfully")
|
112
|
+
else:
|
113
|
+
logger.debug("Close called but session was already None")
|
114
|
+
|
115
|
+
def handle_request_exceptions(self, e, operation_name):
|
116
|
+
"""Handle common request exceptions with appropriate logging"""
|
117
|
+
logger.error(f"Exception occurred during {operation_name}")
|
118
|
+
if isinstance(e, (PoolError, MaxRetryError)):
|
119
|
+
logger.error(f"Connection pool exhausted during {operation_name}: {e}")
|
120
|
+
elif isinstance(e, NewConnectionError):
|
121
|
+
logger.error(f"Failed to establish new connection during {operation_name}: {e}")
|
122
|
+
elif isinstance(e, RemoteDisconnected):
|
123
|
+
logger.error(f"Remote connection closed unexpectedly during {operation_name}: {e}")
|
124
|
+
elif isinstance(e, ConnectionError):
|
125
|
+
logger.error(f"Connection error during {operation_name}: {e}")
|
126
|
+
elif isinstance(e, Timeout):
|
127
|
+
logger.error(f"Request timeout during {operation_name}: {e}")
|
128
|
+
else:
|
129
|
+
logger.error(f"Unexpected error during {operation_name}: {e}")
|
130
|
+
|
131
|
+
def make_request_with_retry(self, method, url, **kwargs):
|
132
|
+
"""
|
133
|
+
Make HTTP request with additional retry logic for RemoteDisconnected errors
|
134
|
+
that may not be caught by urllib3's retry mechanism.
|
135
|
+
"""
|
136
|
+
max_retries = 3
|
137
|
+
for attempt in range(max_retries):
|
138
|
+
try:
|
139
|
+
response = self._session.request(method, url, **kwargs)
|
140
|
+
return response
|
141
|
+
except (RemoteDisconnected, ConnectionError) as e:
|
142
|
+
logger.warning(f"Connection error on attempt {attempt + 1}/{max_retries}: {e}")
|
143
|
+
if attempt == max_retries - 1:
|
144
|
+
# Re-raise the exception on the last attempt
|
145
|
+
raise
|
146
|
+
# Wait before retrying (exponential backoff)
|
147
|
+
import time
|
148
|
+
wait_time = 2 ** attempt
|
149
|
+
logger.info(f"Retrying in {wait_time} seconds...")
|
150
|
+
time.sleep(wait_time)
|
151
|
+
|
152
|
+
|
153
|
+
# Global session manager instance
|
154
|
+
logger.info("Creating global SessionManager instance")
|
155
|
+
session_manager = SessionManager()
|
156
|
+
logger.info(f"Global SessionManager instance created with ID: {id(session_manager)}")
|
@@ -22,6 +22,7 @@ from typing import Dict, Any, Optional
|
|
22
22
|
import threading
|
23
23
|
import uuid
|
24
24
|
|
25
|
+
|
25
26
|
# Set up logging
|
26
27
|
log_dir = os.path.join(tempfile.gettempdir(), "ragaai_logs")
|
27
28
|
os.makedirs(log_dir, exist_ok=True)
|
@@ -49,11 +50,13 @@ try:
|
|
49
50
|
from ragaai_catalyst.tracers.agentic_tracing.upload.upload_code import upload_code
|
50
51
|
# from ragaai_catalyst.tracers.agentic_tracing.upload.upload_trace_metric import upload_trace_metric
|
51
52
|
from ragaai_catalyst.tracers.agentic_tracing.utils.create_dataset_schema import create_dataset_schema_with_trace
|
53
|
+
from ragaai_catalyst.tracers.agentic_tracing.upload.session_manager import session_manager
|
52
54
|
from ragaai_catalyst import RagaAICatalyst
|
53
55
|
IMPORTS_AVAILABLE = True
|
54
56
|
except ImportError:
|
55
57
|
logger.warning("RagaAI Catalyst imports not available - running in test mode")
|
56
58
|
IMPORTS_AVAILABLE = False
|
59
|
+
session_manager = None
|
57
60
|
|
58
61
|
# Define task queue directory
|
59
62
|
QUEUE_DIR = os.path.join(tempfile.gettempdir(), "ragaai_tasks")
|
@@ -72,6 +75,10 @@ _executor_lock = threading.Lock()
|
|
72
75
|
_futures: Dict[str, Any] = {}
|
73
76
|
_futures_lock = threading.Lock()
|
74
77
|
|
78
|
+
# Dataset creation cache to avoid redundant API calls
|
79
|
+
_dataset_cache: Dict[str, Dict[str, Any]] = {}
|
80
|
+
_dataset_cache_lock = threading.Lock()
|
81
|
+
DATASET_CACHE_DURATION = 600 # 10 minutes in seconds
|
75
82
|
|
76
83
|
_cleanup_lock = threading.Lock()
|
77
84
|
_last_cleanup = 0
|
@@ -88,7 +95,7 @@ def get_executor(max_workers=None):
|
|
88
95
|
if _executor is None:
|
89
96
|
# Calculate optimal worker count
|
90
97
|
if max_workers is None:
|
91
|
-
max_workers = min(
|
98
|
+
max_workers = min(8, (os.cpu_count() or 1) * 4)
|
92
99
|
|
93
100
|
logger.info(f"Creating ThreadPoolExecutor with {max_workers} workers")
|
94
101
|
_executor = concurrent.futures.ThreadPoolExecutor(
|
@@ -110,9 +117,57 @@ def generate_unique_task_id():
|
|
110
117
|
unique_id = str(uuid.uuid4())[:8] # Short UUID
|
111
118
|
return f"task_{int(time.time())}_{os.getpid()}_{counter}_{unique_id}"
|
112
119
|
|
120
|
+
def _generate_dataset_cache_key(dataset_name: str, project_name: str, base_url: str) -> str:
|
121
|
+
"""Generate a unique cache key for dataset creation"""
|
122
|
+
return f"{dataset_name}#{project_name}#{base_url}"
|
123
|
+
|
124
|
+
def _is_dataset_cached(cache_key: str) -> bool:
|
125
|
+
"""Check if dataset creation is cached and still valid"""
|
126
|
+
with _dataset_cache_lock:
|
127
|
+
if cache_key not in _dataset_cache:
|
128
|
+
return False
|
129
|
+
|
130
|
+
cache_entry = _dataset_cache[cache_key]
|
131
|
+
cache_time = cache_entry.get('timestamp', 0)
|
132
|
+
current_time = time.time()
|
133
|
+
|
134
|
+
# Check if cache is still valid (within 10 minutes)
|
135
|
+
if current_time - cache_time <= DATASET_CACHE_DURATION:
|
136
|
+
logger.info(f"Dataset creation cache hit for key: {cache_key}")
|
137
|
+
return True
|
138
|
+
else:
|
139
|
+
# Cache expired, remove it
|
140
|
+
logger.info(f"Dataset creation cache expired for key: {cache_key}")
|
141
|
+
del _dataset_cache[cache_key]
|
142
|
+
return False
|
143
|
+
|
144
|
+
def _cache_dataset_creation(cache_key: str, response: Any) -> None:
|
145
|
+
"""Cache successful dataset creation"""
|
146
|
+
with _dataset_cache_lock:
|
147
|
+
_dataset_cache[cache_key] = {
|
148
|
+
'timestamp': time.time(),
|
149
|
+
'response': response
|
150
|
+
}
|
151
|
+
|
152
|
+
def _cleanup_expired_cache_entries() -> None:
|
153
|
+
"""Remove expired cache entries"""
|
154
|
+
current_time = time.time()
|
155
|
+
with _dataset_cache_lock:
|
156
|
+
expired_keys = []
|
157
|
+
for cache_key, cache_entry in _dataset_cache.items():
|
158
|
+
cache_time = cache_entry.get('timestamp', 0)
|
159
|
+
if current_time - cache_time > DATASET_CACHE_DURATION:
|
160
|
+
expired_keys.append(cache_key)
|
161
|
+
|
162
|
+
for key in expired_keys:
|
163
|
+
del _dataset_cache[key]
|
164
|
+
|
165
|
+
if expired_keys:
|
166
|
+
logger.info(f"Cleaned up {len(expired_keys)} expired dataset cache entries")
|
167
|
+
|
113
168
|
def process_upload(task_id: str, filepath: str, hash_id: str, zip_path: str,
|
114
169
|
project_name: str, project_id: str, dataset_name: str,
|
115
|
-
user_details: Dict[str, Any], base_url: str, timeout=120, fail_on_trace_error=True) -> Dict[str, Any]:
|
170
|
+
user_details: Dict[str, Any], base_url: str, tracer_type, timeout=120, fail_on_trace_error=True) -> Dict[str, Any]:
|
116
171
|
"""
|
117
172
|
Process a single upload task
|
118
173
|
|
@@ -165,20 +220,39 @@ def process_upload(task_id: str, filepath: str, hash_id: str, zip_path: str,
|
|
165
220
|
save_task_status(result)
|
166
221
|
return result
|
167
222
|
|
168
|
-
# Step 1: Create dataset schema
|
223
|
+
# Step 1: Create dataset schema (with caching)
|
169
224
|
logger.info(f"Creating dataset schema for {dataset_name} with base_url: {base_url} and timeout: {timeout}")
|
170
|
-
|
171
|
-
|
172
|
-
|
173
|
-
|
174
|
-
|
175
|
-
|
176
|
-
|
177
|
-
|
178
|
-
|
179
|
-
|
180
|
-
|
181
|
-
|
225
|
+
|
226
|
+
# Generate cache key and check if dataset creation is already cached
|
227
|
+
cache_key = _generate_dataset_cache_key(dataset_name, project_name, base_url)
|
228
|
+
|
229
|
+
if _is_dataset_cached(cache_key):
|
230
|
+
logger.info(f"Dataset schema creation skipped (cached) for {dataset_name}")
|
231
|
+
else:
|
232
|
+
try:
|
233
|
+
# Clean up expired cache entries periodically
|
234
|
+
# _cleanup_expired_cache_entries()
|
235
|
+
|
236
|
+
response = create_dataset_schema_with_trace(
|
237
|
+
dataset_name=dataset_name,
|
238
|
+
project_name=project_name,
|
239
|
+
base_url=base_url,
|
240
|
+
user_details=user_details,
|
241
|
+
timeout=timeout
|
242
|
+
)
|
243
|
+
|
244
|
+
if response is None:
|
245
|
+
logger.error(f"Dataset schema creation failed for {dataset_name} - received None response")
|
246
|
+
elif hasattr(response, 'status_code') and response.status_code in [200, 201]:
|
247
|
+
logger.info(f"Dataset schema created successfully: {response.status_code}")
|
248
|
+
_cache_dataset_creation(cache_key, response)
|
249
|
+
logger.info(f"Response cached successfully for dataset: {dataset_name} and key: {cache_key}")
|
250
|
+
else:
|
251
|
+
logger.warning(f"Dataset schema creation returned unexpected response: {response}")
|
252
|
+
|
253
|
+
except Exception as e:
|
254
|
+
logger.error(f"Error creating dataset schema: {e}")
|
255
|
+
# Continue with other steps
|
182
256
|
|
183
257
|
# Step 2: Upload trace metrics
|
184
258
|
# if filepath and os.path.exists(filepath):
|
@@ -238,28 +312,34 @@ def process_upload(task_id: str, filepath: str, hash_id: str, zip_path: str,
|
|
238
312
|
logger.error(error_msg)
|
239
313
|
|
240
314
|
# Step 4: Upload code hash
|
241
|
-
if
|
242
|
-
logger.info(f"
|
243
|
-
|
244
|
-
|
245
|
-
|
246
|
-
|
247
|
-
|
248
|
-
|
249
|
-
|
250
|
-
|
251
|
-
|
252
|
-
|
253
|
-
|
254
|
-
|
255
|
-
|
256
|
-
|
257
|
-
|
315
|
+
if tracer_type.startswith("agentic/"):
|
316
|
+
logger.info(f"Tracer type '{tracer_type}' matches agentic pattern, proceeding with code upload")
|
317
|
+
if hash_id and zip_path and os.path.exists(zip_path):
|
318
|
+
logger.info(f"Uploading code hash {hash_id} with base_url: {base_url} and timeout: {timeout}")
|
319
|
+
try:
|
320
|
+
response = upload_code(
|
321
|
+
hash_id=hash_id,
|
322
|
+
zip_path=zip_path,
|
323
|
+
project_name=project_name,
|
324
|
+
dataset_name=dataset_name,
|
325
|
+
base_url=base_url,
|
326
|
+
timeout=timeout
|
327
|
+
)
|
328
|
+
if response is None:
|
329
|
+
error_msg = "Code hash not uploaded"
|
330
|
+
logger.error(error_msg)
|
331
|
+
else:
|
332
|
+
logger.info(f"Code hash uploaded successfully: {response}")
|
333
|
+
except Exception as e:
|
334
|
+
logger.error(f"Error uploading code hash: {e}")
|
335
|
+
else:
|
336
|
+
logger.warning(f"Code zip {zip_path} not found, skipping code upload")
|
337
|
+
|
258
338
|
# Mark task as completed
|
259
339
|
result["status"] = STATUS_COMPLETED
|
260
340
|
result["end_time"] = datetime.now().isoformat()
|
261
341
|
logger.info(f"Task {task_id} completed successfully")
|
262
|
-
|
342
|
+
|
263
343
|
except Exception as e:
|
264
344
|
logger.error(f"Error processing task {task_id}: {e}")
|
265
345
|
result["status"] = STATUS_FAILED
|
@@ -302,7 +382,8 @@ def save_task_status(task_status: Dict[str, Any]):
|
|
302
382
|
with open(status_path, "w") as f:
|
303
383
|
json.dump(task_status, f, indent=2)
|
304
384
|
|
305
|
-
def submit_upload_task(filepath, hash_id, zip_path, project_name, project_id, dataset_name, user_details, base_url,
|
385
|
+
def submit_upload_task(filepath, hash_id, zip_path, project_name, project_id, dataset_name, user_details, base_url,
|
386
|
+
tracer_type, timeout=120):
|
306
387
|
"""
|
307
388
|
Submit a new upload task using futures.
|
308
389
|
|
@@ -349,6 +430,7 @@ def submit_upload_task(filepath, hash_id, zip_path, project_name, project_id, da
|
|
349
430
|
dataset_name=dataset_name,
|
350
431
|
user_details=user_details,
|
351
432
|
base_url=base_url,
|
433
|
+
tracer_type = tracer_type,
|
352
434
|
timeout=timeout,
|
353
435
|
fail_on_trace_error=True
|
354
436
|
)
|
@@ -379,6 +461,7 @@ def submit_upload_task(filepath, hash_id, zip_path, project_name, project_id, da
|
|
379
461
|
dataset_name=dataset_name,
|
380
462
|
user_details=user_details,
|
381
463
|
base_url=base_url,
|
464
|
+
tracer_type=tracer_type,
|
382
465
|
timeout=timeout,
|
383
466
|
fail_on_trace_error=True
|
384
467
|
)
|
@@ -550,6 +633,14 @@ def shutdown(timeout=120):
|
|
550
633
|
|
551
634
|
_executor = None
|
552
635
|
|
636
|
+
# Close the session manager to clean up HTTP connections
|
637
|
+
if session_manager is not None:
|
638
|
+
try:
|
639
|
+
session_manager.close()
|
640
|
+
logger.info("Session manager closed successfully")
|
641
|
+
except Exception as e:
|
642
|
+
logger.error(f"Error closing session manager: {e}")
|
643
|
+
|
553
644
|
# Register shutdown handler
|
554
645
|
atexit.register(shutdown)
|
555
646
|
|
@@ -4,8 +4,10 @@ import os
|
|
4
4
|
import re
|
5
5
|
import time
|
6
6
|
from urllib.parse import urlparse, urlunparse
|
7
|
-
|
8
|
-
import
|
7
|
+
from urllib3.exceptions import PoolError, MaxRetryError, NewConnectionError
|
8
|
+
from requests.exceptions import ConnectionError, Timeout, RequestException
|
9
|
+
from http.client import RemoteDisconnected
|
10
|
+
from .session_manager import session_manager
|
9
11
|
|
10
12
|
logger = logging.getLogger(__name__)
|
11
13
|
|
@@ -48,7 +50,7 @@ class UploadAgenticTraces:
|
|
48
50
|
start_time = time.time()
|
49
51
|
endpoint = f"{self.base_url}/v1/llm/presigned-url"
|
50
52
|
# Changed to POST from GET
|
51
|
-
response =
|
53
|
+
response = session_manager.make_request_with_retry(
|
52
54
|
"POST", endpoint, headers=headers, data=payload, timeout=self.timeout
|
53
55
|
)
|
54
56
|
elapsed_ms = (time.time() - start_time) * 1000
|
@@ -60,9 +62,35 @@ class UploadAgenticTraces:
|
|
60
62
|
presignedURLs = response.json()["data"]["presignedUrls"][0]
|
61
63
|
presignedurl = self.update_presigned_url(presignedURLs, self.base_url)
|
62
64
|
return presignedurl
|
65
|
+
elif response.status_code == 401:
|
66
|
+
logger.warning("Received 401 error while getting presign url. Attempting to refresh token.")
|
67
|
+
token = RagaAICatalyst.get_token(force_refresh=True)
|
68
|
+
headers = {
|
69
|
+
"Content-Type": "application/json",
|
70
|
+
"Authorization": f"Bearer {token}",
|
71
|
+
"X-Project-Name": self.project_name,
|
72
|
+
}
|
73
|
+
response = session_manager.make_request_with_retry(
|
74
|
+
"POST", endpoint, headers=headers, data=payload, timeout=self.timeout
|
75
|
+
)
|
76
|
+
elapsed_ms = (time.time() - start_time) * 1000
|
77
|
+
logger.debug(
|
78
|
+
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
79
|
+
)
|
80
|
+
if response.status_code == 200:
|
81
|
+
presignedURLs = response.json()["data"]["presignedUrls"][0]
|
82
|
+
presignedurl = self.update_presigned_url(presignedURLs, self.base_url)
|
83
|
+
return presignedurl
|
84
|
+
else:
|
85
|
+
logger.error(
|
86
|
+
f"Error while getting presigned url after token refresh: {response.json()['message']}"
|
87
|
+
)
|
88
|
+
return None
|
63
89
|
else:
|
64
90
|
# If POST fails, try GET
|
65
|
-
|
91
|
+
logger.warning(
|
92
|
+
f"POST request failed for getting presign url with status{response.status_code}.Falling back to GET request.")
|
93
|
+
response = session_manager.make_request_with_retry(
|
66
94
|
"GET", endpoint, headers=headers, data=payload, timeout=self.timeout
|
67
95
|
)
|
68
96
|
elapsed_ms = (time.time() - start_time) * 1000
|
@@ -76,14 +104,14 @@ class UploadAgenticTraces:
|
|
76
104
|
)
|
77
105
|
return presignedurl
|
78
106
|
elif response.status_code == 401:
|
79
|
-
logger.warning("Received 401 error. Attempting to refresh token.")
|
107
|
+
logger.warning("Received 401 error while getting presign url. Attempting to refresh token.")
|
80
108
|
token = RagaAICatalyst.get_token(force_refresh=True)
|
81
109
|
headers = {
|
82
110
|
"Content-Type": "application/json",
|
83
111
|
"Authorization": f"Bearer {token}",
|
84
112
|
"X-Project-Name": self.project_name,
|
85
113
|
}
|
86
|
-
response =
|
114
|
+
response = session_manager.make_request_with_retry(
|
87
115
|
"POST",
|
88
116
|
endpoint,
|
89
117
|
headers=headers,
|
@@ -102,7 +130,7 @@ class UploadAgenticTraces:
|
|
102
130
|
return presignedurl
|
103
131
|
else:
|
104
132
|
logger.error(
|
105
|
-
f"Error while getting presigned url: {response.json()['message']}"
|
133
|
+
f"Error while getting presigned url after token refresh: {response.json()['message']}"
|
106
134
|
)
|
107
135
|
return None
|
108
136
|
else:
|
@@ -110,8 +138,10 @@ class UploadAgenticTraces:
|
|
110
138
|
f"Error while getting presigned url: {response.json()['message']}"
|
111
139
|
)
|
112
140
|
return None
|
113
|
-
|
114
|
-
|
141
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
142
|
+
session_manager.handle_request_exceptions(e, "getting presigned URL")
|
143
|
+
return None
|
144
|
+
except RequestException as e:
|
115
145
|
logger.error(f"Error while getting presigned url: {e}")
|
116
146
|
return None
|
117
147
|
|
@@ -138,16 +168,16 @@ class UploadAgenticTraces:
|
|
138
168
|
|
139
169
|
if "blob.core.windows.net" in presignedUrl: # Azure
|
140
170
|
headers["x-ms-blob-type"] = "BlockBlob"
|
141
|
-
|
171
|
+
logger.info("Uploading agentic traces to presigned URL...")
|
142
172
|
try:
|
143
173
|
with open(filename) as f:
|
144
174
|
payload = f.read().replace("\n", "").replace("\r", "").encode()
|
145
175
|
except Exception as e:
|
146
|
-
|
176
|
+
logger.error(f"Error while reading file: {e}")
|
147
177
|
return False
|
148
178
|
try:
|
149
179
|
start_time = time.time()
|
150
|
-
response =
|
180
|
+
response = session_manager.make_request_with_retry(
|
151
181
|
"PUT", presignedUrl, headers=headers, data=payload, timeout=self.timeout
|
152
182
|
)
|
153
183
|
elapsed_ms = (time.time() - start_time) * 1000
|
@@ -157,8 +187,11 @@ class UploadAgenticTraces:
|
|
157
187
|
if response.status_code != 200 or response.status_code != 201:
|
158
188
|
return response, response.status_code
|
159
189
|
return True
|
160
|
-
except
|
161
|
-
|
190
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
191
|
+
session_manager.handle_request_exceptions(e, "uploading trace to presigned URL")
|
192
|
+
return False
|
193
|
+
except RequestException as e:
|
194
|
+
logger.error(f"Error while uploading trace to presigned url: {e}")
|
162
195
|
return False
|
163
196
|
|
164
197
|
def insert_traces(self, presignedUrl):
|
@@ -177,16 +210,16 @@ class UploadAgenticTraces:
|
|
177
210
|
try:
|
178
211
|
start_time = time.time()
|
179
212
|
endpoint = f"{self.base_url}/v1/llm/insert/trace"
|
180
|
-
response =
|
213
|
+
response = session_manager.make_request_with_retry(
|
181
214
|
"POST", endpoint, headers=headers, data=payload, timeout=self.timeout
|
182
215
|
)
|
183
216
|
elapsed_ms = (time.time() - start_time) * 1000
|
184
217
|
logger.debug(
|
185
218
|
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
186
219
|
)
|
187
|
-
if response.status_code
|
188
|
-
|
189
|
-
return
|
220
|
+
if response.status_code in [200, 201]:
|
221
|
+
logger.info(f"Traces inserted successfully: {response.json()['message']}")
|
222
|
+
return True
|
190
223
|
elif response.status_code == 401:
|
191
224
|
logger.warning("Received 401 error. Attempting to refresh token.")
|
192
225
|
token = RagaAICatalyst.get_token(force_refresh=True)
|
@@ -195,7 +228,7 @@ class UploadAgenticTraces:
|
|
195
228
|
"Content-Type": "application/json",
|
196
229
|
"X-Project-Name": self.project_name,
|
197
230
|
}
|
198
|
-
response =
|
231
|
+
response = session_manager.make_request_with_retry(
|
199
232
|
"POST",
|
200
233
|
endpoint,
|
201
234
|
headers=headers,
|
@@ -206,17 +239,21 @@ class UploadAgenticTraces:
|
|
206
239
|
logger.debug(
|
207
240
|
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
208
241
|
)
|
209
|
-
if response.status_code
|
210
|
-
|
211
|
-
return
|
242
|
+
if response.status_code in [200, 201]:
|
243
|
+
logger.info(f"Traces inserted successfully: {response.json()['message']}")
|
244
|
+
return True
|
212
245
|
else:
|
213
|
-
|
246
|
+
logger.error(f"Error while inserting traces after 401: {response.json()['message']}")
|
214
247
|
return False
|
215
248
|
else:
|
216
|
-
|
217
|
-
|
218
|
-
|
219
|
-
|
249
|
+
logger.error(f"Error while inserting traces: {response.json()['message']}")
|
250
|
+
return False
|
251
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
252
|
+
session_manager.handle_request_exceptions(e, "inserting traces")
|
253
|
+
return False
|
254
|
+
except RequestException as e:
|
255
|
+
logger.error(f"Error while inserting traces: {e}")
|
256
|
+
return False
|
220
257
|
|
221
258
|
def _get_dataset_spans(self):
|
222
259
|
try:
|
@@ -245,26 +282,26 @@ class UploadAgenticTraces:
|
|
245
282
|
continue
|
246
283
|
return dataset_spans
|
247
284
|
except Exception as e:
|
248
|
-
|
285
|
+
logger.error(f"Error while reading dataset spans: {e}")
|
249
286
|
return None
|
250
287
|
|
251
288
|
def upload_agentic_traces(self):
|
252
289
|
try:
|
253
290
|
presigned_url = self._get_presigned_url()
|
254
291
|
if presigned_url is None:
|
255
|
-
|
292
|
+
logger.warning("Warning: Failed to obtain presigned URL")
|
256
293
|
return False
|
257
294
|
|
258
295
|
# Upload the file using the presigned URL
|
259
296
|
upload_result = self._put_presigned_url(presigned_url, self.json_file_path)
|
260
297
|
if not upload_result:
|
261
|
-
|
298
|
+
logger.error("Error: Failed to upload file to presigned URL")
|
262
299
|
return False
|
263
300
|
elif isinstance(upload_result, tuple):
|
264
301
|
response, status_code = upload_result
|
265
302
|
if status_code not in [200, 201]:
|
266
|
-
|
267
|
-
f"Error:
|
303
|
+
logger.error(
|
304
|
+
f"Error: Uploading agentic traces failed with status code {status_code}: {response.text if hasattr(response, 'text') else 'Unknown error'}")
|
268
305
|
return False
|
269
306
|
# Insert trace records
|
270
307
|
insert_success = self.insert_traces(presigned_url)
|
@@ -272,13 +309,14 @@ class UploadAgenticTraces:
|
|
272
309
|
print("Error: Failed to insert trace records")
|
273
310
|
return False
|
274
311
|
|
275
|
-
|
312
|
+
logger.info("Successfully uploaded agentic traces")
|
276
313
|
return True
|
277
314
|
except FileNotFoundError:
|
278
|
-
|
315
|
+
logger.error(f"Error: Trace file not found at {self.json_file_path}")
|
279
316
|
return False
|
280
317
|
except ConnectionError as e:
|
281
|
-
|
318
|
+
logger.error(f"Error: Network connection failed while uploading traces: {e}")
|
282
319
|
return False
|
283
320
|
except Exception as e:
|
284
|
-
|
321
|
+
logger.error(f"Error while uploading agentic traces: {e}")
|
322
|
+
return False
|
@@ -1,15 +1,17 @@
|
|
1
|
-
import json
|
2
1
|
import logging
|
3
2
|
import os
|
4
3
|
import time
|
5
|
-
|
6
|
-
import
|
4
|
+
import re
|
5
|
+
import json
|
6
|
+
from urllib.parse import urlparse, urlunparse
|
7
|
+
from urllib3.exceptions import PoolError, MaxRetryError, NewConnectionError
|
8
|
+
from requests.exceptions import ConnectionError, Timeout, RequestException
|
9
|
+
from http.client import RemoteDisconnected
|
7
10
|
|
8
11
|
from ragaai_catalyst.ragaai_catalyst import RagaAICatalyst
|
12
|
+
from .session_manager import session_manager
|
9
13
|
|
10
14
|
logger = logging.getLogger(__name__)
|
11
|
-
import re
|
12
|
-
from urllib.parse import urlparse, urlunparse
|
13
15
|
|
14
16
|
|
15
17
|
def upload_code(
|
@@ -19,11 +21,26 @@ def upload_code(
|
|
19
21
|
project_name, dataset_name, base_url, timeout=timeout
|
20
22
|
)
|
21
23
|
|
24
|
+
# Handle None case during exceptions - do not proceed
|
25
|
+
if code_hashes_list is None:
|
26
|
+
logger.error("Failed to fetch existing code hashes, cannot proceed with upload")
|
27
|
+
return None
|
28
|
+
|
22
29
|
if hash_id not in code_hashes_list:
|
23
30
|
presigned_url = _fetch_presigned_url(
|
24
31
|
project_name, dataset_name, base_url, timeout=timeout
|
25
32
|
)
|
26
|
-
|
33
|
+
# Handle None case for presigned URL
|
34
|
+
if presigned_url is None:
|
35
|
+
logger.error("Failed to fetch presigned URL, cannot proceed with upload")
|
36
|
+
return None
|
37
|
+
|
38
|
+
upload_result = _put_zip_presigned_url(project_name, presigned_url, zip_path, timeout=timeout)
|
39
|
+
|
40
|
+
# Handle upload failure
|
41
|
+
if upload_result is False or (isinstance(upload_result, tuple) and upload_result[1] not in [200, 201]):
|
42
|
+
logger.error("Failed to upload zip file")
|
43
|
+
return None
|
27
44
|
|
28
45
|
response = _insert_code(
|
29
46
|
dataset_name,
|
@@ -33,6 +50,10 @@ def upload_code(
|
|
33
50
|
base_url,
|
34
51
|
timeout=timeout,
|
35
52
|
)
|
53
|
+
# Handle None response from insert_code
|
54
|
+
if response is None:
|
55
|
+
logger.error("Failed to insert code metadata")
|
56
|
+
return None
|
36
57
|
return response
|
37
58
|
else:
|
38
59
|
return "Code already exists"
|
@@ -49,7 +70,7 @@ def _fetch_dataset_code_hashes(project_name, dataset_name, base_url=None, timeou
|
|
49
70
|
url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
|
50
71
|
start_time = time.time()
|
51
72
|
endpoint = f"{url_base}/v2/llm/dataset/code?datasetName={dataset_name}"
|
52
|
-
response =
|
73
|
+
response = session_manager.make_request_with_retry(
|
53
74
|
"GET", endpoint, headers=headers, data=payload, timeout=timeout
|
54
75
|
)
|
55
76
|
elapsed_ms = (time.time() - start_time) * 1000
|
@@ -57,7 +78,7 @@ def _fetch_dataset_code_hashes(project_name, dataset_name, base_url=None, timeou
|
|
57
78
|
f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
58
79
|
)
|
59
80
|
|
60
|
-
if response.status_code
|
81
|
+
if response.status_code in [200, 201]:
|
61
82
|
return response.json()["data"]["codeHashes"]
|
62
83
|
elif response.status_code == 401:
|
63
84
|
logger.warning("Received 401 error. Attempting to refresh token.")
|
@@ -66,22 +87,25 @@ def _fetch_dataset_code_hashes(project_name, dataset_name, base_url=None, timeou
|
|
66
87
|
"Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
|
67
88
|
"X-Project-Name": project_name,
|
68
89
|
}
|
69
|
-
response =
|
90
|
+
response = session_manager.make_request_with_retry(
|
70
91
|
"GET", endpoint, headers=headers, data=payload, timeout=timeout
|
71
92
|
)
|
72
93
|
elapsed_ms = (time.time() - start_time) * 1000
|
73
|
-
logger.debug(
|
74
|
-
|
75
|
-
)
|
76
|
-
if response.status_code == 200:
|
94
|
+
logger.debug(f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
|
95
|
+
if response.status_code in [200, 201]:
|
77
96
|
return response.json()["data"]["codeHashes"]
|
78
97
|
else:
|
79
|
-
logger.error(
|
80
|
-
|
81
|
-
|
82
|
-
|
83
|
-
|
84
|
-
|
98
|
+
logger.error(f"Failed to fetch code hashes: {response.json()['message']}")
|
99
|
+
return None
|
100
|
+
else:
|
101
|
+
logger.error(f"Error while fetching dataset code hashes: {response.json()['message']}")
|
102
|
+
return None
|
103
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
104
|
+
session_manager.handle_request_exceptions(e, "fetching dataset code hashes")
|
105
|
+
return None
|
106
|
+
except RequestException as e:
|
107
|
+
logger.error(f"Failed to fetch dataset code hashes: {e}")
|
108
|
+
return None
|
85
109
|
|
86
110
|
|
87
111
|
def update_presigned_url(presigned_url, base_url):
|
@@ -116,7 +140,7 @@ def _fetch_presigned_url(project_name, dataset_name, base_url=None, timeout=120)
|
|
116
140
|
start_time = time.time()
|
117
141
|
# Changed to POST from GET
|
118
142
|
endpoint = f"{url_base}/v1/llm/presigned-url"
|
119
|
-
response =
|
143
|
+
response = session_manager.make_request_with_retry(
|
120
144
|
"POST", endpoint, headers=headers, data=payload, timeout=timeout
|
121
145
|
)
|
122
146
|
elapsed_ms = (time.time() - start_time) * 1000
|
@@ -124,20 +148,20 @@ def _fetch_presigned_url(project_name, dataset_name, base_url=None, timeout=120)
|
|
124
148
|
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
125
149
|
)
|
126
150
|
|
127
|
-
if response.status_code
|
151
|
+
if response.status_code in [200, 201]:
|
128
152
|
presigned_url = response.json()["data"]["presignedUrls"][0]
|
129
153
|
presigned_url = update_presigned_url(presigned_url, url_base)
|
130
154
|
return presigned_url
|
131
155
|
else:
|
132
156
|
# If POST fails, try GET
|
133
|
-
response =
|
157
|
+
response = session_manager.make_request_with_retry(
|
134
158
|
"POST", endpoint, headers=headers, data=payload, timeout=timeout
|
135
159
|
)
|
136
160
|
elapsed_ms = (time.time() - start_time) * 1000
|
137
161
|
logger.debug(
|
138
162
|
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
139
163
|
)
|
140
|
-
if response.status_code
|
164
|
+
if response.status_code in [200, 201]:
|
141
165
|
presigned_url = response.json()["data"]["presignedUrls"][0]
|
142
166
|
presigned_url = update_presigned_url(presigned_url, url_base)
|
143
167
|
return presigned_url
|
@@ -149,28 +173,32 @@ def _fetch_presigned_url(project_name, dataset_name, base_url=None, timeout=120)
|
|
149
173
|
"Content-Type": "application/json",
|
150
174
|
"X-Project-Name": project_name,
|
151
175
|
}
|
152
|
-
response =
|
176
|
+
response = session_manager.make_request_with_retry(
|
153
177
|
"POST", endpoint, headers=headers, data=payload, timeout=timeout
|
154
178
|
)
|
155
179
|
elapsed_ms = (time.time() - start_time) * 1000
|
156
180
|
logger.debug(
|
157
181
|
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
158
182
|
)
|
159
|
-
if response.status_code
|
183
|
+
if response.status_code in [200, 201]:
|
160
184
|
presigned_url = response.json()["data"]["presignedUrls"][0]
|
161
185
|
presigned_url = update_presigned_url(presigned_url, url_base)
|
162
186
|
return presigned_url
|
163
187
|
else:
|
164
188
|
logger.error(
|
165
|
-
f"Failed to fetch code
|
189
|
+
f"Failed to fetch presigned URL for code upload after 401: {response.json()['message']}"
|
166
190
|
)
|
167
191
|
else:
|
168
192
|
logger.error(
|
169
|
-
f"Failed to fetch code
|
193
|
+
f"Failed to fetch presigned URL for code upload: {response.json()['message']}"
|
170
194
|
)
|
171
|
-
|
172
|
-
|
173
|
-
|
195
|
+
return None
|
196
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
197
|
+
session_manager.handle_request_exceptions(e, "fetching presigned URL for code upload")
|
198
|
+
return None
|
199
|
+
except RequestException as e:
|
200
|
+
logger.error(f"Failed to fetch presigned URL for code upload: {e}")
|
201
|
+
return None
|
174
202
|
|
175
203
|
|
176
204
|
def _put_zip_presigned_url(project_name, presignedUrl, filename, timeout=120):
|
@@ -181,21 +209,28 @@ def _put_zip_presigned_url(project_name, presignedUrl, filename, timeout=120):
|
|
181
209
|
|
182
210
|
if "blob.core.windows.net" in presignedUrl: # Azure
|
183
211
|
headers["x-ms-blob-type"] = "BlockBlob"
|
184
|
-
|
185
|
-
|
186
|
-
|
187
|
-
|
188
|
-
start_time = time.time()
|
189
|
-
response = requests.request(
|
190
|
-
"PUT", presignedUrl, headers=headers, data=payload, timeout=timeout
|
191
|
-
)
|
192
|
-
elapsed_ms = (time.time() - start_time) * 1000
|
193
|
-
logger.debug(
|
194
|
-
f"API Call: [PUT] {presignedUrl} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
195
|
-
)
|
196
|
-
if response.status_code != 200 or response.status_code != 201:
|
197
|
-
return response, response.status_code
|
212
|
+
logger.info("Uploading code to presigned URL...")
|
213
|
+
try:
|
214
|
+
with open(filename, "rb") as f:
|
215
|
+
payload = f.read()
|
198
216
|
|
217
|
+
start_time = time.time()
|
218
|
+
response = session_manager.make_request_with_retry(
|
219
|
+
"PUT", presignedUrl, headers=headers, data=payload, timeout=timeout
|
220
|
+
)
|
221
|
+
elapsed_ms = (time.time() - start_time) * 1000
|
222
|
+
logger.debug(
|
223
|
+
f"API Call: [PUT] {presignedUrl} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
224
|
+
)
|
225
|
+
if response.status_code not in [200, 201]:
|
226
|
+
return response, response.status_code
|
227
|
+
return True
|
228
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
229
|
+
session_manager.handle_request_exceptions(e, "uploading zip to presigned URL")
|
230
|
+
return False
|
231
|
+
except RequestException as e:
|
232
|
+
logger.error(f"Failed to upload zip: {e}")
|
233
|
+
return False
|
199
234
|
|
200
235
|
def _insert_code(
|
201
236
|
dataset_name, hash_id, presigned_url, project_name, base_url=None, timeout=120
|
@@ -218,39 +253,43 @@ def _insert_code(
|
|
218
253
|
url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
|
219
254
|
start_time = time.time()
|
220
255
|
endpoint = f"{url_base}/v2/llm/dataset/code"
|
221
|
-
response =
|
256
|
+
response = session_manager.make_request_with_retry(
|
222
257
|
"POST", endpoint, headers=headers, data=payload, timeout=timeout
|
223
258
|
)
|
224
259
|
elapsed_ms = (time.time() - start_time) * 1000
|
225
260
|
logger.debug(
|
226
261
|
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
227
262
|
)
|
228
|
-
if response.status_code
|
263
|
+
if response.status_code in [200, 201]:
|
229
264
|
return response.json()["message"]
|
230
265
|
|
231
266
|
elif response.status_code == 401:
|
232
|
-
logger.warning("Received 401 error. Attempting to refresh token.")
|
267
|
+
logger.warning("Received 401 error during inserting code. Attempting to refresh token.")
|
233
268
|
token = RagaAICatalyst.get_token(force_refresh=True)
|
234
269
|
headers = {
|
235
270
|
"X-Project-Name": project_name,
|
236
271
|
"Content-Type": "application/json",
|
237
272
|
"Authorization": f"Bearer {token}",
|
238
273
|
}
|
239
|
-
response =
|
274
|
+
response = session_manager.make_request_with_retry(
|
240
275
|
"POST", endpoint, headers=headers, data=payload, timeout=timeout
|
241
276
|
)
|
242
277
|
elapsed_ms = (time.time() - start_time) * 1000
|
243
278
|
logger.debug(
|
244
279
|
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
245
280
|
)
|
246
|
-
if response.status_code
|
281
|
+
if response.status_code in [200, 201]:
|
282
|
+
logger.info(f"Code inserted successfully after 401: {response.json()['message']}")
|
247
283
|
return response.json()["message"]
|
248
284
|
else:
|
249
|
-
logger.error(f"Failed to insert code: {response.json()['message']}")
|
250
|
-
|
285
|
+
logger.error(f"Failed to insert code after 401: {response.json()['message']}")
|
286
|
+
return None
|
251
287
|
else:
|
252
288
|
logger.error(f"Failed to insert code: {response.json()['message']}")
|
253
|
-
|
254
|
-
except
|
289
|
+
return None
|
290
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
291
|
+
session_manager.handle_request_exceptions(e, "inserting code")
|
292
|
+
return None
|
293
|
+
except RequestException as e:
|
255
294
|
logger.error(f"Failed to insert code: {e}")
|
256
|
-
|
295
|
+
return None
|
@@ -1,8 +1,16 @@
|
|
1
1
|
import os
|
2
2
|
import json
|
3
3
|
import re
|
4
|
-
import
|
4
|
+
import logging
|
5
|
+
import time
|
6
|
+
from urllib3.exceptions import PoolError, MaxRetryError, NewConnectionError
|
7
|
+
from requests.exceptions import ConnectionError, Timeout, RequestException
|
8
|
+
from http.client import RemoteDisconnected
|
9
|
+
|
5
10
|
from ragaai_catalyst import RagaAICatalyst
|
11
|
+
from ragaai_catalyst.tracers.agentic_tracing.upload.session_manager import session_manager
|
12
|
+
|
13
|
+
logger = logging.getLogger(__name__)
|
6
14
|
|
7
15
|
def create_dataset_schema_with_trace(project_name, dataset_name, base_url=None, user_details=None, timeout=120):
|
8
16
|
SCHEMA_MAPPING = {}
|
@@ -13,31 +21,69 @@ def create_dataset_schema_with_trace(project_name, dataset_name, base_url=None,
|
|
13
21
|
continue
|
14
22
|
SCHEMA_MAPPING[key] = {"columnType": "metadata"}
|
15
23
|
|
16
|
-
|
17
|
-
|
18
|
-
|
19
|
-
|
20
|
-
|
21
|
-
|
22
|
-
|
23
|
-
|
24
|
-
|
25
|
-
|
26
|
-
|
27
|
-
|
28
|
-
|
29
|
-
|
30
|
-
|
31
|
-
|
32
|
-
|
24
|
+
headers = {
|
25
|
+
"Content-Type": "application/json",
|
26
|
+
"Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
|
27
|
+
"X-Project-Name": project_name,
|
28
|
+
}
|
29
|
+
|
30
|
+
if SCHEMA_MAPPING:
|
31
|
+
payload = json.dumps({
|
32
|
+
"datasetName": dataset_name,
|
33
|
+
"traceFolderUrl": None,
|
34
|
+
"schemaMapping": SCHEMA_MAPPING
|
35
|
+
})
|
36
|
+
else:
|
37
|
+
payload = json.dumps({
|
38
|
+
"datasetName": dataset_name,
|
39
|
+
"traceFolderUrl": None,
|
40
|
+
})
|
41
|
+
|
42
|
+
try:
|
33
43
|
# Use provided base_url or fall back to default
|
34
44
|
url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
|
35
|
-
|
36
|
-
|
37
|
-
|
38
|
-
|
39
|
-
timeout=timeout
|
45
|
+
start_time = time.time()
|
46
|
+
endpoint = f"{url_base}/v1/llm/dataset/logs"
|
47
|
+
|
48
|
+
response = session_manager.make_request_with_retry(
|
49
|
+
"POST", endpoint, headers=headers, data=payload, timeout=timeout
|
50
|
+
)
|
51
|
+
|
52
|
+
elapsed_ms = (time.time() - start_time) * 1000
|
53
|
+
logger.debug(
|
54
|
+
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
40
55
|
)
|
41
|
-
|
42
|
-
|
43
|
-
|
56
|
+
|
57
|
+
if response.status_code in [200, 201]:
|
58
|
+
logger.info(f"Dataset schema created successfully: {response.status_code}")
|
59
|
+
return response
|
60
|
+
elif response.status_code == 401:
|
61
|
+
logger.warning("Received 401 error during dataset schema creation. Attempting to refresh token.")
|
62
|
+
RagaAICatalyst.get_token(force_refresh=True)
|
63
|
+
headers = {
|
64
|
+
"Content-Type": "application/json",
|
65
|
+
"Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
|
66
|
+
"X-Project-Name": project_name,
|
67
|
+
}
|
68
|
+
response = session_manager.make_request_with_retry(
|
69
|
+
"POST", endpoint, headers=headers, data=payload, timeout=timeout
|
70
|
+
)
|
71
|
+
elapsed_ms = (time.time() - start_time) * 1000
|
72
|
+
logger.debug(
|
73
|
+
f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
|
74
|
+
)
|
75
|
+
if response.status_code in [200, 201]:
|
76
|
+
logger.info(f"Dataset schema created successfully after 401: {response.status_code}")
|
77
|
+
return response
|
78
|
+
else:
|
79
|
+
logger.error(f"Failed to create dataset schema after 401: {response.status_code}")
|
80
|
+
return None
|
81
|
+
else:
|
82
|
+
logger.error(f"Failed to create dataset schema: {response.status_code}")
|
83
|
+
return None
|
84
|
+
except (PoolError, MaxRetryError, NewConnectionError, ConnectionError, Timeout, RemoteDisconnected) as e:
|
85
|
+
session_manager.handle_request_exceptions(e, "creating dataset schema")
|
86
|
+
return None
|
87
|
+
except RequestException as e:
|
88
|
+
logger.error(f"Failed to create dataset schema: {e}")
|
89
|
+
return None
|
@@ -77,6 +77,9 @@ class RAGATraceExporter(SpanExporter):
|
|
77
77
|
|
78
78
|
if trace_id not in self.trace_spans:
|
79
79
|
self.trace_spans[trace_id] = list()
|
80
|
+
|
81
|
+
if span_json.get("attributes").get("openinference.span.kind", None) is None:
|
82
|
+
span_json["attributes"]["openinference.span.kind"] = "UNKNOWN"
|
80
83
|
|
81
84
|
self.trace_spans[trace_id].append(span_json)
|
82
85
|
|
@@ -225,6 +228,7 @@ class RAGATraceExporter(SpanExporter):
|
|
225
228
|
dataset_name=self.dataset_name,
|
226
229
|
user_details=self.user_details,
|
227
230
|
base_url=self.base_url,
|
231
|
+
tracer_type=self.tracer_type,
|
228
232
|
timeout=self.timeout
|
229
233
|
)
|
230
234
|
|
@@ -170,7 +170,7 @@ class Tracer(AgenticTracing):
|
|
170
170
|
logger.error(f"Failed to retrieve projects list: {e}")
|
171
171
|
|
172
172
|
# Handle agentic tracers
|
173
|
-
if tracer_type == "agentic" or tracer_type.startswith("agentic/") or tracer_type == "langchain" or tracer_type == "google-adk":
|
173
|
+
if tracer_type == "agentic" or tracer_type.startswith("agentic/") or tracer_type == "langchain" or tracer_type == "llamaindex" or tracer_type == "google-adk":
|
174
174
|
# Setup instrumentors based on tracer type
|
175
175
|
instrumentors = []
|
176
176
|
|
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: ragaai_catalyst
|
3
|
-
Version: 2.2.
|
3
|
+
Version: 2.2.5b5
|
4
4
|
Summary: RAGA AI CATALYST
|
5
5
|
Author-email: Kiran Scaria <kiran.scaria@raga.ai>, Kedar Gaikwad <kedar.gaikwad@raga.ai>, Dushyant Mahajan <dushyant.mahajan@raga.ai>, Siddhartha Kosti <siddhartha.kosti@raga.ai>, Ritika Goel <ritika.goel@raga.ai>, Vijay Chaurasia <vijay.chaurasia@raga.ai>, Tushar Kumar <tushar.kumar@raga.ai>, Rishabh Pandey <rishabh.pandey@raga.ai>, Jyotsana C G <jyotsana@raga.ai>
|
6
6
|
Requires-Python: <=3.13.2,>=3.10
|
@@ -27,7 +27,7 @@ ragaai_catalyst/redteaming/utils/issue_description.py,sha256=iB0XbeOjdqHTPrikCKS
|
|
27
27
|
ragaai_catalyst/redteaming/utils/rt.png,sha256=HzVC8bz_4UgwafKXuMe8RJVI6CyK_UmSgo53ceAOQK8,282154
|
28
28
|
ragaai_catalyst/tracers/__init__.py,sha256=LfgTes-nHpazssbGKnn8kyLZNr49kIPrlkrqqoTFTfc,301
|
29
29
|
ragaai_catalyst/tracers/distributed.py,sha256=CGPuOh4CsgEk428PPibieLaAG2Tt3BVygF6ZlmbXxg4,10009
|
30
|
-
ragaai_catalyst/tracers/tracer.py,sha256=
|
30
|
+
ragaai_catalyst/tracers/tracer.py,sha256=5jM-AJozLoc95OIR-K3K2fs0OJ2Oo47Lgtca_6ToFzk,33649
|
31
31
|
ragaai_catalyst/tracers/agentic_tracing/README.md,sha256=X4QwLb7-Jg7GQMIXj-SerZIgDETfw-7VgYlczOR8ZeQ,4508
|
32
32
|
ragaai_catalyst/tracers/agentic_tracing/__init__.py,sha256=yf6SKvOPSpH-9LiKaoLKXwqj5sez8F_5wkOb91yp0oE,260
|
33
33
|
ragaai_catalyst/tracers/agentic_tracing/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
@@ -35,11 +35,12 @@ ragaai_catalyst/tracers/agentic_tracing/data/data_structure.py,sha256=icAtNzKN_I
|
|
35
35
|
ragaai_catalyst/tracers/agentic_tracing/tracers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
36
36
|
ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py,sha256=Wq4LFclPlLy47LyXvbaLeYiSMQABj7VYS3J87xyea_E,4159
|
37
37
|
ragaai_catalyst/tracers/agentic_tracing/upload/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
38
|
-
ragaai_catalyst/tracers/agentic_tracing/upload/
|
39
|
-
ragaai_catalyst/tracers/agentic_tracing/upload/
|
40
|
-
ragaai_catalyst/tracers/agentic_tracing/upload/
|
38
|
+
ragaai_catalyst/tracers/agentic_tracing/upload/session_manager.py,sha256=sOlxeIYIP8tycaTtZC9xkZosi6EDJUxvDw0_rc_NLI8,6823
|
39
|
+
ragaai_catalyst/tracers/agentic_tracing/upload/trace_uploader.py,sha256=NuetePZOdDmwRgN3aAsQrDIytXFicNylaAqORdz8C2o,25051
|
40
|
+
ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py,sha256=yLzYiyNk_XtwwlUFbq5uaRahDKKF-eFBoI6xqHUNedw,14526
|
41
|
+
ragaai_catalyst/tracers/agentic_tracing/upload/upload_code.py,sha256=IAhNFS-nbV_ImNz8Xp98qU4r-2naj49qg9q08x53TFE,12521
|
41
42
|
ragaai_catalyst/tracers/agentic_tracing/utils/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
|
42
|
-
ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py,sha256=
|
43
|
+
ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py,sha256=1c6HA_Alm5yofF-eifqwdwHboxr25mvW3xxVYBktfjA,3667
|
43
44
|
ragaai_catalyst/tracers/agentic_tracing/utils/file_name_tracker.py,sha256=YG601l1a29ov9VPu9Vl4RXxgL7l16k54_WWnoTNoG58,2064
|
44
45
|
ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py,sha256=PiyXvEj_qu0EnJFjk4GfGyWFZbwlvQQh0hdQ_lm0p8E,22976
|
45
46
|
ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json,sha256=2tzGw_cKCTPcfjEm7iGvFE6pTw7gMTPzeBov_MTaXNY,321336
|
@@ -50,15 +51,15 @@ ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py,sha256
|
|
50
51
|
ragaai_catalyst/tracers/exporters/__init__.py,sha256=qA3vx7z9CQ5kTGCn9LIDtIFvW9fJHQLkvF9-xBQUm94,237
|
51
52
|
ragaai_catalyst/tracers/exporters/dynamic_trace_exporter.py,sha256=Rm-QaLv1qMAKpHKcFOcK_HWaKHwFBoUH45_4QYipE-g,6843
|
52
53
|
ragaai_catalyst/tracers/exporters/file_span_exporter.py,sha256=_icciSCktK6c86KB2HV3GZMFHvUitgKJ8x_IdPmgi1M,6363
|
53
|
-
ragaai_catalyst/tracers/exporters/ragaai_trace_exporter.py,sha256=
|
54
|
+
ragaai_catalyst/tracers/exporters/ragaai_trace_exporter.py,sha256=M3Bj311o_MpjY1gO9SGpfIIT2O0yz76uCgw6mKUwHdo,10032
|
54
55
|
ragaai_catalyst/tracers/instrumentators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
55
56
|
ragaai_catalyst/tracers/utils/__init__.py,sha256=KeMaZtYaTojilpLv65qH08QmpYclfpacDA0U3wg6Ybw,64
|
56
57
|
ragaai_catalyst/tracers/utils/model_prices_and_context_window_backup.json,sha256=WlZCZeOQ54aMVjYS8BAeka2uaFC3ftBTMZ8zzzA8TAI,495947
|
57
58
|
ragaai_catalyst/tracers/utils/rag_extraction_logic_final.py,sha256=3ygkRT__lLDRflRttjzPu28tIA8cTCiGQVMQjqMItqQ,11309
|
58
59
|
ragaai_catalyst/tracers/utils/trace_json_converter.py,sha256=-HZVmijeUFLO7e9OAvi1RJdWVTxPRUHPd1MkKQlCD54,11785
|
59
60
|
ragaai_catalyst/tracers/utils/utils.py,sha256=o-p9n2ZuophdrV0wrixu-BqRHCkovup_klc3mS8mU8g,2374
|
60
|
-
ragaai_catalyst-2.2.
|
61
|
-
ragaai_catalyst-2.2.
|
62
|
-
ragaai_catalyst-2.2.
|
63
|
-
ragaai_catalyst-2.2.
|
64
|
-
ragaai_catalyst-2.2.
|
61
|
+
ragaai_catalyst-2.2.5b5.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
|
62
|
+
ragaai_catalyst-2.2.5b5.dist-info/METADATA,sha256=Oho91NQ2Otn-agTiKESdvBL1iZwLyQLeBtWTgu5nfc8,17735
|
63
|
+
ragaai_catalyst-2.2.5b5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
64
|
+
ragaai_catalyst-2.2.5b5.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
|
65
|
+
ragaai_catalyst-2.2.5b5.dist-info/RECORD,,
|
File without changes
|
File without changes
|
File without changes
|