lucidicai 2.1.0__py3-none-any.whl → 2.1.2__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.
lucidicai/__init__.py CHANGED
@@ -14,6 +14,10 @@ from .sdk.init import (
14
14
  init as _init,
15
15
  get_session_id as _get_session_id,
16
16
  clear_state as _clear_state,
17
+ # Thread-local session management (advanced users)
18
+ set_thread_session,
19
+ clear_thread_session,
20
+ get_thread_session,
17
21
  )
18
22
 
19
23
  from .sdk.event import (
@@ -32,6 +36,7 @@ from .sdk.context import (
32
36
  session_async,
33
37
  run_session,
34
38
  run_in_session,
39
+ thread_worker_with_session, # Thread isolation helper
35
40
  current_session_id,
36
41
  current_parent_event_id,
37
42
  )
@@ -55,12 +60,15 @@ def _update_session(
55
60
  session_eval=None,
56
61
  session_eval_reason=None,
57
62
  is_successful=None,
58
- is_successful_reason=None
63
+ is_successful_reason=None,
64
+ session_id=None # Accept explicit session_id
59
65
  ):
60
66
  """Update the current session."""
61
67
  from .sdk.init import get_resources, get_session_id
62
-
63
- session_id = get_session_id()
68
+
69
+ # Use provided session_id or fall back to context
70
+ if not session_id:
71
+ session_id = get_session_id()
64
72
  if not session_id:
65
73
  return
66
74
 
@@ -87,12 +95,16 @@ def _end_session(
87
95
  session_eval_reason=None,
88
96
  is_successful=None,
89
97
  is_successful_reason=None,
90
- wait_for_flush=True
98
+ wait_for_flush=True,
99
+ session_id=None # Accept explicit session_id
91
100
  ):
92
101
  """End the current session."""
93
102
  from .sdk.init import get_resources, get_session_id, get_event_queue
94
-
95
- session_id = get_session_id()
103
+ from .sdk.shutdown_manager import get_shutdown_manager
104
+
105
+ # Use provided session_id or fall back to context
106
+ if not session_id:
107
+ session_id = get_session_id()
96
108
  if not session_id:
97
109
  return
98
110
 
@@ -114,6 +126,9 @@ def _end_session(
114
126
  # Clear session context
115
127
  clear_active_session()
116
128
 
129
+ # unregister from shutdown manager
130
+ get_shutdown_manager().unregister_session(session_id)
131
+
117
132
 
118
133
  def _get_session():
119
134
  """Get the current session object."""
@@ -282,7 +297,7 @@ get_error_history = error_boundary.get_error_history
282
297
  clear_error_history = error_boundary.clear_error_history
283
298
 
284
299
  # Version
285
- __version__ = "2.0.0"
300
+ __version__ = "2.1.2"
286
301
 
287
302
  # Apply error boundary wrapping to all SDK functions
288
303
  from .sdk.error_boundary import wrap_sdk_function
@@ -371,8 +386,14 @@ __all__ = [
371
386
  'session_async',
372
387
  'run_session',
373
388
  'run_in_session',
389
+ 'thread_worker_with_session',
374
390
  'current_session_id',
375
391
  'current_parent_event_id',
392
+
393
+ # Thread-local session management (advanced)
394
+ 'set_thread_session',
395
+ 'clear_thread_session',
396
+ 'get_thread_session',
376
397
 
377
398
  # Error types
378
399
  'LucidicError',
lucidicai/core/config.py CHANGED
@@ -19,7 +19,7 @@ class Environment(Enum):
19
19
  @dataclass
20
20
  class NetworkConfig:
21
21
  """Network and connection settings"""
22
- base_url: str = "https://api.lucidic.ai/api"
22
+ base_url: str = "https://backend.lucidic.ai/api"
23
23
  timeout: int = 30
24
24
  max_retries: int = 3
25
25
  backoff_factor: float = 0.5
@@ -31,7 +31,7 @@ class NetworkConfig:
31
31
  """Load network configuration from environment variables"""
32
32
  debug = os.getenv("LUCIDIC_DEBUG", "False").lower() == "true"
33
33
  return cls(
34
- base_url="http://localhost:8000/api" if debug else "https://api.lucidic.ai/api",
34
+ base_url="http://localhost:8000/api" if debug else "https://backend.lucidic.ai/api",
35
35
  timeout=int(os.getenv("LUCIDIC_TIMEOUT", "30")),
36
36
  max_retries=int(os.getenv("LUCIDIC_MAX_RETRIES", "3")),
37
37
  backoff_factor=float(os.getenv("LUCIDIC_BACKOFF_FACTOR", "0.5")),
lucidicai/sdk/context.py CHANGED
@@ -1,4 +1,4 @@
1
- """Async-safe context helpers for session (and step, extensible).
1
+ """Async-safe and thread-safe context helpers for session (and step, extensible).
2
2
 
3
3
  This module exposes context variables and helpers to bind a Lucidic
4
4
  session to the current execution context (threads/async tasks), so
@@ -11,6 +11,7 @@ import contextvars
11
11
  from typing import Optional, Iterator, AsyncIterator, Callable, Any, Dict
12
12
  import logging
13
13
  import os
14
+ import threading
14
15
 
15
16
 
16
17
  # Context variable for the active Lucidic session id
@@ -26,32 +27,71 @@ current_parent_event_id: contextvars.ContextVar[Optional[str]] = contextvars.Con
26
27
 
27
28
 
28
29
  def set_active_session(session_id: Optional[str]) -> None:
29
- """Bind the given session id to the current execution context."""
30
+ """Bind the given session id to the current execution context.
31
+
32
+ Sets both contextvar and thread-local storage when in a thread.
33
+ """
34
+ from .init import set_thread_session, is_main_thread
35
+
30
36
  current_session_id.set(session_id)
31
37
 
38
+ # Also set thread-local storage if we're in a non-main thread
39
+ if session_id and not is_main_thread():
40
+ set_thread_session(session_id)
41
+
32
42
 
33
43
  def clear_active_session() -> None:
34
- """Clear any active session binding in the current execution context."""
44
+ """Clear any active session binding in the current execution context.
45
+
46
+ Clears both contextvar and thread-local storage when in a thread.
47
+ """
48
+ from .init import clear_thread_session, is_main_thread
49
+
35
50
  current_session_id.set(None)
36
51
 
52
+ # Also clear thread-local storage if we're in a non-main thread
53
+ if not is_main_thread():
54
+ clear_thread_session()
55
+
37
56
 
38
57
  @contextmanager
39
58
  def bind_session(session_id: str) -> Iterator[None]:
40
- """Context manager to temporarily bind an active session id."""
59
+ """Context manager to temporarily bind an active session id.
60
+
61
+ Handles both thread-local and context variable storage for proper isolation.
62
+ """
63
+ from .init import set_thread_session, clear_thread_session, is_main_thread
64
+
41
65
  token = current_session_id.set(session_id)
66
+
67
+ # If we're in a non-main thread, also set thread-local storage
68
+ thread_local_set = False
69
+ if not is_main_thread():
70
+ set_thread_session(session_id)
71
+ thread_local_set = True
72
+
42
73
  try:
43
74
  yield
44
75
  finally:
76
+ if thread_local_set:
77
+ clear_thread_session()
45
78
  current_session_id.reset(token)
46
79
 
47
80
 
48
81
  @asynccontextmanager
49
82
  async def bind_session_async(session_id: str) -> AsyncIterator[None]:
50
83
  """Async context manager to temporarily bind an active session id."""
84
+ from .init import set_task_session, clear_task_session
85
+
51
86
  token = current_session_id.set(session_id)
87
+
88
+ # Also set task-local for async isolation
89
+ set_task_session(session_id)
90
+
52
91
  try:
53
92
  yield
54
93
  finally:
94
+ clear_task_session()
55
95
  current_session_id.reset(token)
56
96
 
57
97
 
@@ -81,9 +121,11 @@ def session(**init_params) -> Iterator[None]:
81
121
  Notes:
82
122
  - Ignores any provided auto_end parameter and ends the session on context exit.
83
123
  - If LUCIDIC_DEBUG is true, logs a warning about ignoring auto_end.
124
+ - Handles thread-local storage for proper thread isolation.
84
125
  """
85
126
  # Lazy import to avoid circular imports
86
127
  import lucidicai as lai # type: ignore
128
+ from .init import set_thread_session, clear_thread_session, is_main_thread
87
129
 
88
130
  # Force auto_end to False inside a context manager to control explicit end
89
131
  user_auto_end = init_params.get('auto_end', None)
@@ -95,12 +137,43 @@ def session(**init_params) -> Iterator[None]:
95
137
 
96
138
  session_id = lai.init(**init_params)
97
139
  token = current_session_id.set(session_id)
140
+
141
+ # If we're in a non-main thread, also set thread-local storage
142
+ thread_local_set = False
143
+ if not is_main_thread():
144
+ set_thread_session(session_id)
145
+ thread_local_set = True
146
+
98
147
  try:
99
148
  yield
100
149
  finally:
150
+ if thread_local_set:
151
+ clear_thread_session()
101
152
  current_session_id.reset(token)
102
153
  try:
103
- lai.end_session()
154
+ # Force flush OpenTelemetry spans before ending session
155
+ from .init import get_tracer_provider
156
+ from ..utils.logger import debug, info
157
+ import time
158
+
159
+ tracer_provider = get_tracer_provider()
160
+ if tracer_provider:
161
+ debug(f"[Session] Force flushing OpenTelemetry spans for session {session_id}")
162
+ try:
163
+ # Force flush with 5 second timeout to ensure all spans are exported
164
+ flush_result = tracer_provider.force_flush(timeout_millis=5000)
165
+ debug(f"[Session] Tracer provider force_flush returned: {flush_result}")
166
+
167
+ # Give a small additional delay to ensure the exporter processes everything
168
+ # This is necessary because force_flush on the provider flushes the processors,
169
+ # but the exporter might still be processing the spans
170
+ time.sleep(0.5)
171
+ debug(f"[Session] Successfully flushed spans for session {session_id}")
172
+ except Exception as e:
173
+ debug(f"[Session] Error flushing spans: {e}")
174
+
175
+ # Pass session_id explicitly to avoid context issues
176
+ lai.end_session(session_id=session_id)
104
177
  except Exception:
105
178
  # Avoid masking the original exception from the with-block
106
179
  pass
@@ -110,6 +183,7 @@ def session(**init_params) -> Iterator[None]:
110
183
  async def session_async(**init_params) -> AsyncIterator[None]:
111
184
  """Async counterpart of session(...)."""
112
185
  import lucidicai as lai # type: ignore
186
+ from .init import set_task_session, clear_task_session
113
187
 
114
188
  user_auto_end = init_params.get('auto_end', None)
115
189
  init_params = dict(init_params)
@@ -120,12 +194,40 @@ async def session_async(**init_params) -> AsyncIterator[None]:
120
194
 
121
195
  session_id = lai.init(**init_params)
122
196
  token = current_session_id.set(session_id)
197
+
198
+ # Set task-local session for true isolation in async
199
+ set_task_session(session_id)
200
+
123
201
  try:
124
202
  yield
125
203
  finally:
204
+ # Clear task-local session first
205
+ clear_task_session()
126
206
  current_session_id.reset(token)
127
207
  try:
128
- lai.end_session()
208
+ # Force flush OpenTelemetry spans before ending session
209
+ from .init import get_tracer_provider
210
+ from ..utils.logger import debug, info
211
+ import asyncio
212
+
213
+ tracer_provider = get_tracer_provider()
214
+ if tracer_provider:
215
+ debug(f"[Session] Force flushing OpenTelemetry spans for async session {session_id}")
216
+ try:
217
+ # Force flush with 5 second timeout to ensure all spans are exported
218
+ flush_result = tracer_provider.force_flush(timeout_millis=5000)
219
+ debug(f"[Session] Tracer provider force_flush returned: {flush_result}")
220
+
221
+ # Give a small additional delay to ensure the exporter processes everything
222
+ # This is necessary because force_flush on the provider flushes the processors,
223
+ # but the exporter might still be processing the spans
224
+ await asyncio.sleep(0.5)
225
+ debug(f"[Session] Successfully flushed spans for async session {session_id}")
226
+ except Exception as e:
227
+ debug(f"[Session] Error flushing spans: {e}")
228
+
229
+ # Pass session_id explicitly to avoid context issues in async
230
+ lai.end_session(session_id=session_id)
129
231
  except Exception:
130
232
  pass
131
233
 
@@ -142,3 +244,30 @@ def run_in_session(session_id: str, fn: Callable[..., Any], *fn_args: Any, **fn_
142
244
  return fn(*fn_args, **fn_kwargs)
143
245
 
144
246
 
247
+ def thread_worker_with_session(session_id: str, target: Callable[..., Any], *args, **kwargs) -> Any:
248
+ """Wrapper for thread worker functions that ensures proper session isolation.
249
+
250
+ Use this as the target function for threads to ensure each thread gets
251
+ its own session context without bleeding from the parent thread.
252
+
253
+ Example:
254
+ thread = Thread(
255
+ target=thread_worker_with_session,
256
+ args=(session_id, actual_worker_function, arg1, arg2),
257
+ kwargs={'key': 'value'}
258
+ )
259
+ """
260
+ from .init import set_thread_session, clear_thread_session
261
+
262
+ # Set thread-local session immediately
263
+ set_thread_session(session_id)
264
+
265
+ try:
266
+ # Also bind to contextvar for compatibility
267
+ with bind_session(session_id):
268
+ return target(*args, **kwargs)
269
+ finally:
270
+ # Clean up thread-local storage
271
+ clear_thread_session()
272
+
273
+
lucidicai/sdk/event.py CHANGED
@@ -12,23 +12,27 @@ from ..utils.logger import debug, truncate_id
12
12
  def create_event(
13
13
  type: str = "generic",
14
14
  event_id: Optional[str] = None,
15
+ session_id: Optional[str] = None, # accept explicit session_id
15
16
  **kwargs
16
17
  ) -> str:
17
18
  """Create a new event.
18
-
19
+
19
20
  Args:
20
21
  type: Event type (llm_generation, function_call, error_traceback, generic)
21
22
  event_id: Optional client event ID (will generate if not provided)
23
+ session_id: Optional session ID (will use context if not provided)
22
24
  **kwargs: Event-specific fields
23
-
25
+
24
26
  Returns:
25
27
  Event ID (client-generated or provided UUID)
26
28
  """
27
29
  # Import here to avoid circular dependency
28
30
  from ..sdk.init import get_session_id, get_event_queue
29
-
30
- # Get current session
31
- session_id = get_session_id()
31
+
32
+ # Use provided session_id or fall back to context
33
+ if not session_id:
34
+ session_id = get_session_id()
35
+
32
36
  if not session_id:
33
37
  # No active session, return dummy ID
34
38
  debug("[Event] No active session, returning dummy event ID")
lucidicai/sdk/init.py CHANGED
@@ -4,6 +4,9 @@ This module handles SDK initialization, separating concerns from the main __init
4
4
  """
5
5
  import uuid
6
6
  from typing import List, Optional
7
+ import asyncio
8
+ import threading
9
+ from weakref import WeakKeyDictionary
7
10
 
8
11
  from ..api.client import HttpClient
9
12
  from ..api.resources.event import EventResource
@@ -21,14 +24,18 @@ from opentelemetry.sdk.trace import TracerProvider
21
24
 
22
25
  class SDKState:
23
26
  """Container for SDK runtime state."""
24
-
27
+
25
28
  def __init__(self):
26
29
  self.http: Optional[HttpClient] = None
27
30
  self.event_queue: Optional[EventQueue] = None
28
31
  self.session_id: Optional[str] = None
29
32
  self.tracer_provider: Optional[TracerProvider] = None
30
33
  self.resources = {}
31
-
34
+ # Task-local storage for async task isolation
35
+ self.task_sessions: WeakKeyDictionary = WeakKeyDictionary()
36
+ # Thread-local storage for thread isolation
37
+ self.thread_local = threading.local()
38
+
32
39
  def reset(self):
33
40
  """Reset SDK state."""
34
41
  # Shutdown telemetry first to ensure all spans are exported
@@ -42,17 +49,21 @@ class SDKState:
42
49
  debug("[SDK] TracerProvider shutdown complete")
43
50
  except Exception as e:
44
51
  error(f"[SDK] Error shutting down TracerProvider: {e}")
45
-
52
+
46
53
  if self.event_queue:
47
54
  self.event_queue.shutdown()
48
55
  if self.http:
49
56
  self.http.close()
50
-
57
+
51
58
  self.http = None
52
59
  self.event_queue = None
53
60
  self.session_id = None
54
61
  self.tracer_provider = None
55
62
  self.resources = {}
63
+ self.task_sessions.clear()
64
+ # Clear thread-local storage for current thread
65
+ if hasattr(self.thread_local, 'session_id'):
66
+ delattr(self.thread_local, 'session_id')
56
67
 
57
68
 
58
69
  # Global SDK state
@@ -243,8 +254,87 @@ def _initialize_telemetry(providers: List[str]) -> None:
243
254
  info(f"[Telemetry] Initialized for providers: {providers}")
244
255
 
245
256
 
257
+ def set_task_session(session_id: str) -> None:
258
+ """Set session ID for current async task (if in async context)."""
259
+ try:
260
+ if task := asyncio.current_task():
261
+ _sdk_state.task_sessions[task] = session_id
262
+ debug(f"[SDK] Set task-local session {truncate_id(session_id)} for task {task.get_name()}")
263
+ except RuntimeError:
264
+ # Not in async context, ignore
265
+ pass
266
+
267
+
268
+ def clear_task_session() -> None:
269
+ """Clear session ID for current async task (if in async context)."""
270
+ try:
271
+ if task := asyncio.current_task():
272
+ _sdk_state.task_sessions.pop(task, None)
273
+ debug(f"[SDK] Cleared task-local session for task {task.get_name()}")
274
+ except RuntimeError:
275
+ # Not in async context, ignore
276
+ pass
277
+
278
+
279
+ def set_thread_session(session_id: str) -> None:
280
+ """Set session ID for current thread.
281
+
282
+ This provides true thread-local storage that doesn't inherit from parent thread.
283
+ """
284
+ _sdk_state.thread_local.session_id = session_id
285
+ current_thread = threading.current_thread()
286
+ debug(f"[SDK] Set thread-local session {truncate_id(session_id)} for thread {current_thread.name}")
287
+
288
+
289
+ def clear_thread_session() -> None:
290
+ """Clear session ID for current thread."""
291
+ if hasattr(_sdk_state.thread_local, 'session_id'):
292
+ delattr(_sdk_state.thread_local, 'session_id')
293
+ current_thread = threading.current_thread()
294
+ debug(f"[SDK] Cleared thread-local session for thread {current_thread.name}")
295
+
296
+
297
+ def get_thread_session() -> Optional[str]:
298
+ """Get session ID from thread-local storage."""
299
+ return getattr(_sdk_state.thread_local, 'session_id', None)
300
+
301
+
302
+ def is_main_thread() -> bool:
303
+ """Check if we're running in the main thread."""
304
+ return threading.current_thread() is threading.main_thread()
305
+
306
+
246
307
  def get_session_id() -> Optional[str]:
247
- """Get the current session ID."""
308
+ """Get the current session ID.
309
+
310
+ Priority:
311
+ 1. Task-local session (for async tasks)
312
+ 2. Thread-local session (for threads) - NO FALLBACK for threads
313
+ 3. SDK state session (for main thread)
314
+ 4. Context variable session (fallback for main thread only)
315
+ """
316
+ # First check task-local storage for async isolation
317
+ try:
318
+ if task := asyncio.current_task():
319
+ if task_session := _sdk_state.task_sessions.get(task):
320
+ debug(f"[SDK] Using task-local session {truncate_id(task_session)}")
321
+ return task_session
322
+ except RuntimeError:
323
+ # Not in async context
324
+ pass
325
+
326
+ # Check if we're in a thread
327
+ if not is_main_thread():
328
+ # For threads, ONLY use thread-local storage - no fallback!
329
+ # This prevents inheriting the parent thread's session
330
+ thread_session = get_thread_session()
331
+ if thread_session:
332
+ debug(f"[SDK] Using thread-local session {truncate_id(thread_session)}")
333
+ else:
334
+ debug(f"[SDK] Thread {threading.current_thread().name} has no thread-local session")
335
+ return thread_session # Return None if not set - don't fall back!
336
+
337
+ # For main thread only: fall back to SDK state or context variable
248
338
  return _sdk_state.session_id or current_session_id.get()
249
339
 
250
340
 
@@ -263,6 +353,11 @@ def get_resources() -> dict:
263
353
  return _sdk_state.resources
264
354
 
265
355
 
356
+ def get_tracer_provider() -> Optional[TracerProvider]:
357
+ """Get the tracer provider instance."""
358
+ return _sdk_state.tracer_provider
359
+
360
+
266
361
  def clear_state() -> None:
267
362
  """Clear SDK state (for testing)."""
268
363
  global _sdk_state
@@ -149,6 +149,7 @@ class LucidicLiteLLMCallback(CustomLogger):
149
149
  # Create event with correct field names
150
150
  create_event(
151
151
  type="llm_generation",
152
+ session_id=session_id, # Pass session_id explicitly
152
153
  provider=provider,
153
154
  model=model,
154
155
  messages=messages,
@@ -210,6 +211,7 @@ class LucidicLiteLLMCallback(CustomLogger):
210
211
 
211
212
  create_event(
212
213
  type="error_traceback",
214
+ session_id=session_id, # Pass session_id explicitly
213
215
  error=error_msg,
214
216
  traceback="",
215
217
  parent_event_id=parent_id, # This will be normalized by EventBuilder
@@ -41,11 +41,26 @@ class LucidicSpanExporter(SpanExporter):
41
41
  if not detect_is_llm_span(span):
42
42
  verbose(f"[Telemetry] Skipping non-LLM span: {span.name}")
43
43
  return
44
-
44
+
45
45
  debug(f"[Telemetry] Processing LLM span: {span.name}")
46
46
 
47
47
  attributes = dict(span.attributes or {})
48
48
 
49
+ # Debug: Check what attributes we have for responses.create
50
+ if span.name == "openai.responses.create":
51
+ debug(f"[Telemetry] responses.create span has {len(attributes)} attributes")
52
+ # Check for specific attributes we're interested in
53
+ has_prompts = any(k.startswith('gen_ai.prompt') for k in attributes.keys())
54
+ has_completions = any(k.startswith('gen_ai.completion') for k in attributes.keys())
55
+ debug(f"[Telemetry] Has prompt attrs: {has_prompts}, Has completion attrs: {has_completions}")
56
+
57
+ # Skip spans that are likely duplicates or incomplete
58
+ # Check if this is a responses.parse span that was already handled
59
+ if span.name == "openai.responses.create" and not attributes.get("lucidic.instrumented"):
60
+ # This might be from incorrect standard instrumentation
61
+ verbose(f"[Telemetry] Skipping potentially duplicate responses span without our marker")
62
+ return
63
+
49
64
  # Resolve session id
50
65
  target_session_id = attributes.get('lucidic.session_id')
51
66
  if not target_session_id:
@@ -84,7 +99,23 @@ class LucidicSpanExporter(SpanExporter):
84
99
  provider = self._detect_provider_name(attributes)
85
100
  messages = extract_prompts(attributes) or []
86
101
  params = self._extract_params(attributes)
87
- output_text = extract_completions(span, attributes) or "Response received"
102
+ output_text = extract_completions(span, attributes)
103
+
104
+ # Debug for responses.create
105
+ if span.name == "openai.responses.create":
106
+ debug(f"[Telemetry] Extracted messages: {messages}")
107
+ debug(f"[Telemetry] Extracted output: {output_text}")
108
+
109
+ # Skip spans with no meaningful output (likely incomplete or duplicate instrumentation)
110
+ if not output_text or output_text == "Response received":
111
+ # Only use "Response received" if we have other meaningful data
112
+ if not messages and not attributes.get("lucidic.instrumented"):
113
+ verbose(f"[Telemetry] Skipping span {span.name} with no meaningful content")
114
+ return
115
+ # Use a more descriptive default if we must
116
+ if not output_text:
117
+ output_text = "Response received"
118
+
88
119
  input_tokens = self._extract_prompt_tokens(attributes)
89
120
  output_tokens = self._extract_completion_tokens(attributes)
90
121
  cost = self._calculate_cost(attributes)
@@ -99,9 +130,10 @@ class LucidicSpanExporter(SpanExporter):
99
130
 
100
131
  try:
101
132
  # Create immutable event via non-blocking queue
102
- debug(f"[Telemetry] Creating LLM event with parent_id: {truncate_id(parent_id)}")
133
+ debug(f"[Telemetry] Creating LLM event with parent_id: {truncate_id(parent_id)}, session_id: {truncate_id(target_session_id)}")
103
134
  event_id = create_event(
104
135
  type="llm_generation",
136
+ session_id=target_session_id, # Pass the session_id explicitly
105
137
  occurred_at=occurred_at,
106
138
  duration=duration_seconds,
107
139
  provider=provider,
@@ -155,14 +187,15 @@ class LucidicSpanExporter(SpanExporter):
155
187
 
156
188
  # Create event
157
189
  event_kwargs = {
190
+ 'session_id': target_session_id, # Pass session_id explicitly
158
191
  'description': description,
159
192
  'result': "Processing...", # Will be updated when span ends
160
193
  'model': model
161
194
  }
162
-
195
+
163
196
  if images:
164
197
  event_kwargs['screenshots'] = images
165
-
198
+
166
199
  return create_event(**event_kwargs)
167
200
 
168
201
  except Exception as e:
@@ -225,31 +258,35 @@ class LucidicSpanExporter(SpanExporter):
225
258
  }
226
259
 
227
260
  def _extract_prompt_tokens(self, attributes: Dict[str, Any]) -> int:
228
- return (
229
- attributes.get(SpanAttributes.LLM_USAGE_PROMPT_TOKENS) or
230
- attributes.get('gen_ai.usage.prompt_tokens') or
231
- attributes.get('gen_ai.usage.input_tokens') or 0
232
- )
261
+ # Check each attribute and return the first non-None value
262
+ value = attributes.get(SpanAttributes.LLM_USAGE_PROMPT_TOKENS)
263
+ if value is not None:
264
+ return value
265
+ value = attributes.get('gen_ai.usage.prompt_tokens')
266
+ if value is not None:
267
+ return value
268
+ value = attributes.get('gen_ai.usage.input_tokens')
269
+ if value is not None:
270
+ return value
271
+ return 0
233
272
 
234
273
  def _extract_completion_tokens(self, attributes: Dict[str, Any]) -> int:
235
- return (
236
- attributes.get(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS) or
237
- attributes.get('gen_ai.usage.completion_tokens') or
238
- attributes.get('gen_ai.usage.output_tokens') or 0
239
- )
274
+ # Check each attribute and return the first non-None value
275
+ value = attributes.get(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS)
276
+ if value is not None:
277
+ return value
278
+ value = attributes.get('gen_ai.usage.completion_tokens')
279
+ if value is not None:
280
+ return value
281
+ value = attributes.get('gen_ai.usage.output_tokens')
282
+ if value is not None:
283
+ return value
284
+ return 0
240
285
 
241
286
  def _calculate_cost(self, attributes: Dict[str, Any]) -> Optional[float]:
242
- prompt_tokens = (
243
- attributes.get(SpanAttributes.LLM_USAGE_PROMPT_TOKENS) or
244
- attributes.get('gen_ai.usage.prompt_tokens') or
245
- attributes.get('gen_ai.usage.input_tokens') or 0
246
- )
247
- completion_tokens = (
248
- attributes.get(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS) or
249
- attributes.get('gen_ai.usage.completion_tokens') or
250
- attributes.get('gen_ai.usage.output_tokens') or 0
251
- )
252
- total_tokens = (prompt_tokens or 0) + (completion_tokens or 0)
287
+ prompt_tokens = self._extract_prompt_tokens(attributes)
288
+ completion_tokens = self._extract_completion_tokens(attributes)
289
+ total_tokens = prompt_tokens + completion_tokens
253
290
  if total_tokens > 0:
254
291
  model = (
255
292
  attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) or