cortexhub 0.1.3__py3-none-any.whl → 0.1.5__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.
@@ -19,11 +19,13 @@ Architectural rules:
19
19
  - No governance logic in adapter
20
20
  """
21
21
 
22
+ import json
22
23
  import os
23
24
  from functools import wraps
24
25
  from typing import TYPE_CHECKING, Any, Callable, Awaitable
25
26
 
26
27
  import structlog
28
+ from opentelemetry.trace import SpanKind, Status, StatusCode
27
29
 
28
30
  from cortexhub.adapters.base import ToolAdapter
29
31
  from cortexhub.pipeline import govern_execution
@@ -36,6 +38,10 @@ logger = structlog.get_logger(__name__)
36
38
  # Attribute names for storing originals
37
39
  _ORIGINAL_TOOL_ATTR = "__cortexhub_original_tool__"
38
40
  _PATCHED_ATTR = "__cortexhub_patched__"
41
+ _ORIGINAL_QUERY_ATTR = "__cortexhub_original_query__"
42
+ _ORIGINAL_RECEIVE_RESPONSE_ATTR = "__cortexhub_original_receive_response__"
43
+ _ORIGINAL_CLIENT_QUERY_ATTR = "__cortexhub_original_client_query__"
44
+ _PATCHED_RUN_ATTR = "__cortexhub_run_patched__"
39
45
 
40
46
 
41
47
  class ClaudeAgentsAdapter(ToolAdapter):
@@ -64,6 +70,11 @@ class ClaudeAgentsAdapter(ToolAdapter):
64
70
  def framework_name(self) -> str:
65
71
  return "claude_agents"
66
72
 
73
+ def __init__(self, cortex_hub: Any):
74
+ super().__init__(cortex_hub)
75
+ self._hook_spans: dict[str, Any] = {}
76
+ self._active_run_session_id: str | None = None
77
+
67
78
  def _get_framework_modules(self) -> list[str]:
68
79
  return ["claude_agent_sdk"]
69
80
 
@@ -145,6 +156,8 @@ class ClaudeAgentsAdapter(ToolAdapter):
145
156
  setattr(claude_agent_sdk, _PATCHED_ATTR, True)
146
157
 
147
158
  logger.info("Claude Agent SDK @tool decorator patched successfully")
159
+
160
+ self._patch_run_completion(cortex_hub)
148
161
 
149
162
  except ImportError:
150
163
  logger.debug("Claude Agent SDK not installed, skipping adapter")
@@ -161,6 +174,23 @@ class ClaudeAgentsAdapter(ToolAdapter):
161
174
  claude_agent_sdk.tool = original
162
175
  setattr(claude_agent_sdk, _PATCHED_ATTR, False)
163
176
  logger.info("Claude Agent SDK adapter unpatched")
177
+
178
+ try:
179
+ from claude_agent_sdk import ClaudeSDKClient
180
+
181
+ if hasattr(claude_agent_sdk, _ORIGINAL_QUERY_ATTR):
182
+ claude_agent_sdk.query = getattr(claude_agent_sdk, _ORIGINAL_QUERY_ATTR)
183
+ if hasattr(ClaudeSDKClient, _ORIGINAL_CLIENT_QUERY_ATTR):
184
+ ClaudeSDKClient.query = getattr(
185
+ ClaudeSDKClient, _ORIGINAL_CLIENT_QUERY_ATTR
186
+ )
187
+ if hasattr(ClaudeSDKClient, _ORIGINAL_RECEIVE_RESPONSE_ATTR):
188
+ ClaudeSDKClient.receive_response = getattr(
189
+ ClaudeSDKClient, _ORIGINAL_RECEIVE_RESPONSE_ATTR
190
+ )
191
+ setattr(claude_agent_sdk, _PATCHED_RUN_ATTR, False)
192
+ except ImportError:
193
+ pass
164
194
  except ImportError:
165
195
  pass
166
196
 
@@ -171,6 +201,94 @@ class ClaudeAgentsAdapter(ToolAdapter):
171
201
  def _discover_tools(self) -> list[dict[str, Any]]:
172
202
  """Discover tools from Claude Agent SDK (best-effort)."""
173
203
  return []
204
+
205
+ def _patch_run_completion(self, cortex_hub) -> None:
206
+ """Patch Claude Agent SDK runs to emit run completion."""
207
+ try:
208
+ import claude_agent_sdk
209
+ from claude_agent_sdk import ClaudeSDKClient, ResultMessage
210
+ adapter = self
211
+
212
+ if getattr(claude_agent_sdk, _PATCHED_RUN_ATTR, False):
213
+ return
214
+
215
+ if not hasattr(claude_agent_sdk, _ORIGINAL_QUERY_ATTR):
216
+ setattr(claude_agent_sdk, _ORIGINAL_QUERY_ATTR, claude_agent_sdk.query)
217
+ original_query = getattr(claude_agent_sdk, _ORIGINAL_QUERY_ATTR)
218
+
219
+ async def patched_query(*args, **kwargs):
220
+ status = None
221
+ failed = False
222
+ cortex_hub.start_run(framework="claude_agents")
223
+ adapter._active_run_session_id = cortex_hub.session_id
224
+ try:
225
+ async for message in original_query(*args, **kwargs):
226
+ if isinstance(message, ResultMessage):
227
+ status = "failed" if message.is_error else "completed"
228
+ yield message
229
+ except Exception:
230
+ failed = True
231
+ raise
232
+ finally:
233
+ if status is None and failed:
234
+ status = "failed"
235
+ if status:
236
+ cortex_hub.finish_run(framework="claude_agents", status=status)
237
+ adapter._active_run_session_id = None
238
+
239
+ claude_agent_sdk.query = patched_query
240
+
241
+ if not hasattr(ClaudeSDKClient, _ORIGINAL_RECEIVE_RESPONSE_ATTR):
242
+ setattr(
243
+ ClaudeSDKClient,
244
+ _ORIGINAL_RECEIVE_RESPONSE_ATTR,
245
+ ClaudeSDKClient.receive_response,
246
+ )
247
+ original_receive_response = getattr(ClaudeSDKClient, _ORIGINAL_RECEIVE_RESPONSE_ATTR)
248
+ if not hasattr(ClaudeSDKClient, _ORIGINAL_CLIENT_QUERY_ATTR):
249
+ setattr(
250
+ ClaudeSDKClient,
251
+ _ORIGINAL_CLIENT_QUERY_ATTR,
252
+ ClaudeSDKClient.query,
253
+ )
254
+ original_client_query = getattr(ClaudeSDKClient, _ORIGINAL_CLIENT_QUERY_ATTR)
255
+
256
+ async def patched_client_query(self, *args, **kwargs):
257
+ cortex_hub.start_run(framework="claude_agents")
258
+ adapter._active_run_session_id = cortex_hub.session_id
259
+ try:
260
+ return await original_client_query(self, *args, **kwargs)
261
+ except Exception:
262
+ cortex_hub.finish_run(framework="claude_agents", status="failed")
263
+ adapter._active_run_session_id = None
264
+ raise
265
+
266
+ async def patched_receive_response(self, *args, **kwargs):
267
+ status = None
268
+ failed = False
269
+ try:
270
+ async for message in original_receive_response(self, *args, **kwargs):
271
+ if isinstance(message, ResultMessage):
272
+ status = "failed" if message.is_error else "completed"
273
+ yield message
274
+ except Exception:
275
+ failed = True
276
+ raise
277
+ finally:
278
+ if status is None and failed:
279
+ status = "failed"
280
+ if status:
281
+ cortex_hub.finish_run(framework="claude_agents", status=status)
282
+ adapter._active_run_session_id = None
283
+
284
+ ClaudeSDKClient.query = patched_client_query
285
+ ClaudeSDKClient.receive_response = patched_receive_response
286
+ setattr(claude_agent_sdk, _PATCHED_RUN_ATTR, True)
287
+ logger.info("Claude Agent SDK run completion patched successfully")
288
+ except ImportError:
289
+ logger.debug("Claude Agent SDK run completion patch skipped")
290
+ except Exception as e:
291
+ logger.debug("Claude Agent SDK run completion patch failed", reason=str(e))
174
292
 
175
293
  def create_governance_hooks(self) -> dict[str, list]:
176
294
  """Create CortexHub governance hooks for Claude Agent SDK.
@@ -191,6 +309,65 @@ class ClaudeAgentsAdapter(ToolAdapter):
191
309
  )
192
310
  """
193
311
  cortex_hub = self.cortex_hub
312
+ span_store = self._hook_spans
313
+ adapter = self
314
+
315
+ def _current_session_id() -> str:
316
+ return adapter._active_run_session_id or cortex_hub.session_id
317
+
318
+ def _start_tool_span(
319
+ *,
320
+ tool_name: str,
321
+ tool_description: str,
322
+ policy_args: dict[str, Any],
323
+ raw_args: dict[str, Any],
324
+ tool_use_id: str | None,
325
+ ):
326
+ span = cortex_hub._tracer.start_span(
327
+ name="tool.invoke",
328
+ kind=SpanKind.INTERNAL,
329
+ )
330
+ span.set_attribute("cortexhub.session.id", _current_session_id())
331
+ span.set_attribute("cortexhub.agent.id", cortex_hub.agent_id)
332
+ span.set_attribute("cortexhub.tool.name", tool_name)
333
+ span.set_attribute("cortexhub.tool.framework", "claude_agents")
334
+ span.set_attribute("cortexhub.tool.description", tool_description)
335
+
336
+ if tool_use_id:
337
+ span.set_attribute("cortexhub.tool.use_id", tool_use_id)
338
+
339
+ if policy_args:
340
+ arg_names = list(policy_args.keys())
341
+ if arg_names:
342
+ span.set_attribute("cortexhub.tool.arg_names", arg_names)
343
+ arg_schema = cortex_hub._infer_arg_schema(policy_args)
344
+ if arg_schema:
345
+ span.set_attribute(
346
+ "cortexhub.tool.arg_schema",
347
+ json.dumps(arg_schema),
348
+ )
349
+
350
+ if not cortex_hub.privacy and raw_args:
351
+ span.set_attribute("cortexhub.raw.args", json.dumps(raw_args, default=str))
352
+
353
+ return span
354
+
355
+ def _finish_tool_span(
356
+ span,
357
+ *,
358
+ success: bool,
359
+ error_message: str | None = None,
360
+ result: Any | None = None,
361
+ ) -> None:
362
+ span.set_attribute("cortexhub.result.success", success)
363
+ if error_message:
364
+ span.set_attribute("cortexhub.error.message", error_message)
365
+ span.set_status(Status(StatusCode.ERROR, error_message))
366
+ else:
367
+ span.set_status(Status(StatusCode.OK))
368
+ if result is not None and not cortex_hub.privacy:
369
+ span.set_attribute("cortexhub.raw.result", json.dumps(result, default=str))
370
+ span.end()
194
371
 
195
372
  async def pre_tool_governance(
196
373
  input_data: dict[str, Any],
@@ -204,12 +381,18 @@ class ClaudeAgentsAdapter(ToolAdapter):
204
381
  """
205
382
  tool_name = input_data.get("tool_name", "unknown")
206
383
  tool_input = input_data.get("tool_input", {})
384
+ if not isinstance(tool_input, dict):
385
+ tool_input = {"_raw": tool_input}
386
+ policy_args = cortex_hub._sanitize_policy_args(tool_input)
207
387
 
208
- tool_metadata = {
209
- "name": tool_name,
210
- "description": f"Claude Agent SDK built-in tool: {tool_name}",
211
- "framework": "claude_agents",
212
- }
388
+ tool_description = f"Claude Agent SDK built-in tool: {tool_name}"
389
+ span = _start_tool_span(
390
+ tool_name=tool_name,
391
+ tool_description=tool_description,
392
+ policy_args=policy_args,
393
+ raw_args=tool_input,
394
+ tool_use_id=tool_use_id,
395
+ )
213
396
 
214
397
  # Build authorization request and evaluate
215
398
  from cortexhub.policy.models import (
@@ -222,7 +405,7 @@ class ClaudeAgentsAdapter(ToolAdapter):
222
405
  principal=Principal(type="Agent", id=cortex_hub.agent_id),
223
406
  action=Action(type="tool.invoke", name=tool_name),
224
407
  resource=PolicyResource(type="Tool", id=tool_name),
225
- args=tool_input,
408
+ args=policy_args,
226
409
  framework="claude_agents",
227
410
  )
228
411
 
@@ -230,8 +413,23 @@ class ClaudeAgentsAdapter(ToolAdapter):
230
413
  if cortex_hub.enforce and cortex_hub.evaluator:
231
414
  from cortexhub.policy.effects import Effect
232
415
  decision = cortex_hub.evaluator.evaluate(request)
416
+
417
+ span.add_event(
418
+ "policy.decision",
419
+ attributes={
420
+ "decision.effect": decision.effect.value,
421
+ "decision.policy_id": decision.policy_id or "",
422
+ "decision.reasoning": decision.reasoning,
423
+ "decision.policy_name": decision.policy_name or "",
424
+ },
425
+ )
233
426
 
234
427
  if decision.effect == Effect.DENY:
428
+ _finish_tool_span(
429
+ span,
430
+ success=False,
431
+ error_message=decision.reasoning,
432
+ )
235
433
  return {
236
434
  "hookSpecificOutput": {
237
435
  "hookEventName": "PreToolUse",
@@ -242,7 +440,7 @@ class ClaudeAgentsAdapter(ToolAdapter):
242
440
 
243
441
  if decision.effect == Effect.ESCALATE:
244
442
  try:
245
- context_hash = cortex_hub._compute_context_hash(tool_name, tool_input)
443
+ context_hash = cortex_hub._compute_context_hash(tool_name, policy_args)
246
444
  approval_response = cortex_hub.backend.create_approval(
247
445
  run_id=cortex_hub.session_id,
248
446
  trace_id=cortex_hub._get_current_trace_id(),
@@ -260,6 +458,11 @@ class ClaudeAgentsAdapter(ToolAdapter):
260
458
  approval_id = approval_response.get("approval_id", "unknown")
261
459
  except Exception as e:
262
460
  logger.error("Failed to create approval", error=str(e))
461
+ _finish_tool_span(
462
+ span,
463
+ success=False,
464
+ error_message=str(e),
465
+ )
263
466
  return {
264
467
  "hookSpecificOutput": {
265
468
  "hookEventName": "PreToolUse",
@@ -270,6 +473,20 @@ class ClaudeAgentsAdapter(ToolAdapter):
270
473
  }
271
474
  }
272
475
 
476
+ span.add_event(
477
+ "approval.created",
478
+ attributes={
479
+ "approval_id": approval_id,
480
+ "tool_name": tool_name,
481
+ "policy_id": decision.policy_id or "",
482
+ "expires_at": approval_response.get("expires_at", ""),
483
+ },
484
+ )
485
+ _finish_tool_span(
486
+ span,
487
+ success=False,
488
+ error_message="Approval required",
489
+ )
273
490
  return {
274
491
  "hookSpecificOutput": {
275
492
  "hookEventName": "PreToolUse",
@@ -281,6 +498,10 @@ class ClaudeAgentsAdapter(ToolAdapter):
281
498
  }
282
499
 
283
500
  # Allow execution
501
+ if tool_use_id:
502
+ span_store[tool_use_id] = span
503
+ else:
504
+ span.end()
284
505
  return {}
285
506
 
286
507
  async def post_tool_governance(
@@ -294,6 +515,21 @@ class ClaudeAgentsAdapter(ToolAdapter):
294
515
  """
295
516
  tool_name = input_data.get("tool_name", "unknown")
296
517
  tool_response = input_data.get("tool_response", {})
518
+ tool_input = input_data.get("tool_input", {})
519
+ if not isinstance(tool_input, dict):
520
+ tool_input = {"_raw": tool_input}
521
+ policy_args = cortex_hub._sanitize_policy_args(tool_input)
522
+
523
+ span = span_store.pop(tool_use_id, None) if tool_use_id else None
524
+ if span is None:
525
+ tool_description = f"Claude Agent SDK built-in tool: {tool_name}"
526
+ span = _start_tool_span(
527
+ tool_name=tool_name,
528
+ tool_description=tool_description,
529
+ policy_args=policy_args,
530
+ raw_args=tool_input,
531
+ tool_use_id=tool_use_id,
532
+ )
297
533
 
298
534
  # Log the tool execution
299
535
  logger.debug(
@@ -301,6 +537,8 @@ class ClaudeAgentsAdapter(ToolAdapter):
301
537
  tool=tool_name,
302
538
  framework="claude_agents",
303
539
  )
540
+
541
+ _finish_tool_span(span, success=True, result=tool_response)
304
542
 
305
543
  return {}
306
544
 
@@ -28,11 +28,21 @@ logger = structlog.get_logger(__name__)
28
28
  # Attribute names for storing originals on class
29
29
  _ORIGINAL_INVOKE_ATTR = "__cortexhub_original_invoke__"
30
30
  _ORIGINAL_RUN_ATTR = "__cortexhub_original_run__"
31
+ _ORIGINAL_RUN_METHOD_ATTR = "__cortexhub_original_run_method__"
31
32
  _PATCHED_ATTR = "__cortexhub_patched__"
32
33
  _PATCHED_TOOL_ATTR = "__cortexhub_tool_patched__"
34
+ _PATCHED_RUN_METHOD_ATTR = "__cortexhub_run_method_patched__"
35
+ _ORIGINAL_TOOL_RUN_ATTR = "__cortexhub_original_tool_run__"
36
+ _PATCHED_TOOL_RUN_ATTR = "__cortexhub_tool_run_patched__"
33
37
  _PATCHED_LLM_ATTR = "__cortexhub_llm_patched__"
34
38
  _ORIGINAL_COMPLETION_ATTR = "__cortexhub_original_completion__"
35
39
  _ORIGINAL_ACOMPLETION_ATTR = "__cortexhub_original_acompletion__"
40
+ _ORIGINAL_KICKOFF_ATTR = "__cortexhub_original_kickoff__"
41
+ _ORIGINAL_KICKOFF_ASYNC_ATTR = "__cortexhub_original_kickoff_async__"
42
+ _PATCHED_RUN_ATTR = "__cortexhub_run_patched__"
43
+ _ORIGINAL_LLM_CALL_ATTR = "__cortexhub_original_llm_call__"
44
+ _ORIGINAL_LLM_ACALL_ATTR = "__cortexhub_original_llm_acall__"
45
+ _PATCHED_LLM_CALL_ATTR = "__cortexhub_llm_call_patched__"
36
46
 
37
47
 
38
48
  class CrewAIAdapter(ToolAdapter):
@@ -105,7 +115,7 @@ class CrewAIAdapter(ToolAdapter):
105
115
 
106
116
  # Also patch BaseTool._run for direct tool.run() calls
107
117
  try:
108
- from crewai.tools.base_tool import BaseTool
118
+ from crewai.tools.base_tool import BaseTool, Tool
109
119
 
110
120
  if not getattr(BaseTool, _PATCHED_TOOL_ATTR, False):
111
121
  if not hasattr(BaseTool, _ORIGINAL_RUN_ATTR):
@@ -142,6 +152,76 @@ class CrewAIAdapter(ToolAdapter):
142
152
  BaseTool._run = patched_run
143
153
  setattr(BaseTool, _PATCHED_TOOL_ATTR, True)
144
154
  logger.info("CrewAI BaseTool._run patched")
155
+
156
+ if not getattr(BaseTool, _PATCHED_RUN_METHOD_ATTR, False):
157
+ if not hasattr(BaseTool, _ORIGINAL_RUN_METHOD_ATTR):
158
+ setattr(BaseTool, _ORIGINAL_RUN_METHOD_ATTR, BaseTool.run)
159
+
160
+ original_run_method = getattr(BaseTool, _ORIGINAL_RUN_METHOD_ATTR)
161
+
162
+ def patched_run_method(self, *args, **kwargs):
163
+ """Governed BaseTool.run execution."""
164
+ tool_name = getattr(self, "name", "unknown_tool")
165
+ tool_description = getattr(self, "description", None)
166
+
167
+ tool_metadata = {
168
+ "name": tool_name,
169
+ "description": tool_description,
170
+ "framework": "crewai",
171
+ }
172
+
173
+ governed_fn = govern_execution(
174
+ tool_fn=lambda **_kw: original_run_method(self, *args, **kwargs),
175
+ tool_metadata=tool_metadata,
176
+ cortex_hub=cortex_hub,
177
+ )
178
+
179
+ if kwargs:
180
+ return governed_fn(**kwargs)
181
+ if len(args) == 1 and isinstance(args[0], dict):
182
+ return governed_fn(**args[0])
183
+ if args:
184
+ return governed_fn(_raw=args[0])
185
+ return governed_fn()
186
+
187
+ BaseTool.run = patched_run_method
188
+ setattr(BaseTool, _PATCHED_RUN_METHOD_ATTR, True)
189
+ logger.info("CrewAI BaseTool.run patched")
190
+
191
+ if not getattr(Tool, _PATCHED_TOOL_RUN_ATTR, False):
192
+ if not hasattr(Tool, _ORIGINAL_TOOL_RUN_ATTR):
193
+ setattr(Tool, _ORIGINAL_TOOL_RUN_ATTR, Tool.run)
194
+
195
+ original_tool_run = getattr(Tool, _ORIGINAL_TOOL_RUN_ATTR)
196
+
197
+ def patched_tool_run(self, *args, **kwargs):
198
+ """Governed Tool.run execution (bypasses BaseTool.run)."""
199
+ tool_name = getattr(self, "name", "unknown_tool")
200
+ tool_description = getattr(self, "description", None)
201
+
202
+ tool_metadata = {
203
+ "name": tool_name,
204
+ "description": tool_description,
205
+ "framework": "crewai",
206
+ }
207
+
208
+ governed_fn = govern_execution(
209
+ tool_fn=lambda **_kw: original_tool_run(self, *args, **kwargs),
210
+ tool_metadata=tool_metadata,
211
+ cortex_hub=cortex_hub,
212
+ )
213
+
214
+ if kwargs:
215
+ return governed_fn(**kwargs)
216
+ if len(args) == 1 and isinstance(args[0], dict):
217
+ return governed_fn(**args[0])
218
+ if args:
219
+ return governed_fn(_raw=args[0])
220
+ return governed_fn()
221
+
222
+ Tool.run = patched_tool_run
223
+ setattr(Tool, _PATCHED_TOOL_RUN_ATTR, True)
224
+ logger.info("CrewAI Tool.run patched")
145
225
 
146
226
  except ImportError:
147
227
  logger.debug("CrewAI BaseTool not available")
@@ -150,6 +230,8 @@ class CrewAIAdapter(ToolAdapter):
150
230
 
151
231
  # Patch LiteLLM for LLM call governance (guardrails, PII)
152
232
  self._patch_litellm(cortex_hub)
233
+ self._patch_llm_call(cortex_hub)
234
+ self._patch_run_completion(cortex_hub)
153
235
 
154
236
  except ImportError:
155
237
  logger.debug("CrewAI not available, skipping adapter")
@@ -251,6 +333,183 @@ class CrewAIAdapter(ToolAdapter):
251
333
  except Exception as e:
252
334
  logger.debug("CrewAI LiteLLM interception skipped", reason=str(e))
253
335
 
336
+ def _patch_llm_call(self, cortex_hub) -> None:
337
+ """Patch CrewAI LLM provider call methods for llm.call spans."""
338
+ try:
339
+ from crewai.llms.base_llm import BaseLLM
340
+ provider_modules = [
341
+ "crewai.llms.providers.openai.completion",
342
+ "crewai.llms.providers.anthropic.completion",
343
+ "crewai.llms.providers.azure.completion",
344
+ "crewai.llms.providers.gemini.completion",
345
+ "crewai.llms.providers.bedrock.completion",
346
+ ]
347
+ for module_path in provider_modules:
348
+ try:
349
+ __import__(module_path)
350
+ except Exception:
351
+ continue
352
+
353
+ def _all_subclasses(base):
354
+ seen = set()
355
+ stack = list(base.__subclasses__())
356
+ while stack:
357
+ cls = stack.pop()
358
+ if cls in seen:
359
+ continue
360
+ seen.add(cls)
361
+ stack.extend(cls.__subclasses__())
362
+ yield cls
363
+
364
+ patched_any = False
365
+ for llm_cls in _all_subclasses(BaseLLM):
366
+ if getattr(llm_cls, _PATCHED_LLM_CALL_ATTR, False):
367
+ continue
368
+
369
+ if hasattr(llm_cls, "call"):
370
+ if not hasattr(llm_cls, _ORIGINAL_LLM_CALL_ATTR):
371
+ setattr(llm_cls, _ORIGINAL_LLM_CALL_ATTR, llm_cls.call)
372
+ original_call = getattr(llm_cls, _ORIGINAL_LLM_CALL_ATTR)
373
+
374
+ def _make_patched_call(call_impl):
375
+ def patched_call(self, messages, *args, **kwargs):
376
+ if getattr(self, "is_litellm", False):
377
+ return call_impl(self, messages, *args, **kwargs)
378
+
379
+ model = getattr(self, "model", "unknown")
380
+ prompt = messages
381
+
382
+ def call_original(prompt_override):
383
+ new_messages = (
384
+ prompt_override if prompt_override is not None else messages
385
+ )
386
+ return call_impl(self, new_messages, *args, **kwargs)
387
+
388
+ llm_metadata = {
389
+ "kind": "llm",
390
+ "framework": "crewai",
391
+ "model": model,
392
+ "prompt": prompt,
393
+ "call_original": call_original,
394
+ }
395
+
396
+ governed = govern_execution(
397
+ tool_fn=lambda *a, **kw: call_impl(
398
+ self, messages, *args, **kwargs
399
+ ),
400
+ tool_metadata=llm_metadata,
401
+ cortex_hub=cortex_hub,
402
+ )
403
+ return governed()
404
+
405
+ return patched_call
406
+
407
+ llm_cls.call = _make_patched_call(original_call)
408
+ patched_any = True
409
+
410
+ if hasattr(llm_cls, "acall"):
411
+ if not hasattr(llm_cls, _ORIGINAL_LLM_ACALL_ATTR):
412
+ setattr(llm_cls, _ORIGINAL_LLM_ACALL_ATTR, llm_cls.acall)
413
+ original_acall = getattr(llm_cls, _ORIGINAL_LLM_ACALL_ATTR)
414
+
415
+ def _make_patched_acall(call_impl):
416
+ async def patched_acall(self, messages, *args, **kwargs):
417
+ if getattr(self, "is_litellm", False):
418
+ return await call_impl(self, messages, *args, **kwargs)
419
+
420
+ model = getattr(self, "model", "unknown")
421
+ prompt = messages
422
+
423
+ async def call_original(prompt_override):
424
+ new_messages = (
425
+ prompt_override if prompt_override is not None else messages
426
+ )
427
+ return await call_impl(self, new_messages, *args, **kwargs)
428
+
429
+ llm_metadata = {
430
+ "kind": "llm",
431
+ "framework": "crewai",
432
+ "model": model,
433
+ "prompt": prompt,
434
+ "call_original": call_original,
435
+ }
436
+
437
+ governed = govern_execution(
438
+ tool_fn=lambda *a, **kw: call_impl(
439
+ self, messages, *args, **kwargs
440
+ ),
441
+ tool_metadata=llm_metadata,
442
+ cortex_hub=cortex_hub,
443
+ )
444
+ return await governed()
445
+
446
+ return patched_acall
447
+
448
+ llm_cls.acall = _make_patched_acall(original_acall)
449
+ patched_any = True
450
+
451
+ setattr(llm_cls, _PATCHED_LLM_CALL_ATTR, True)
452
+
453
+ if patched_any:
454
+ logger.info("CrewAI LLM provider calls patched successfully")
455
+ except ImportError:
456
+ logger.debug("CrewAI LLM base not available, skipping LLM call patch")
457
+ except Exception as e:
458
+ logger.debug("CrewAI LLM call patch skipped", reason=str(e))
459
+
460
+ def _patch_run_completion(self, cortex_hub) -> None:
461
+ """Patch CrewAI crew kickoff methods to emit run completion."""
462
+ try:
463
+ import crewai
464
+ Crew = getattr(crewai, "Crew", None)
465
+ if Crew is None:
466
+ from crewai.crew import Crew
467
+
468
+ if getattr(Crew, _PATCHED_RUN_ATTR, False):
469
+ return
470
+
471
+ if not hasattr(Crew, _ORIGINAL_KICKOFF_ATTR):
472
+ setattr(Crew, _ORIGINAL_KICKOFF_ATTR, Crew.kickoff)
473
+ original_kickoff = getattr(Crew, _ORIGINAL_KICKOFF_ATTR)
474
+
475
+ def patched_kickoff(self, *args, **kwargs):
476
+ status = "completed"
477
+ cortex_hub.start_run(framework="crewai")
478
+ try:
479
+ return original_kickoff(self, *args, **kwargs)
480
+ except Exception:
481
+ status = "failed"
482
+ raise
483
+ finally:
484
+ cortex_hub.finish_run(framework="crewai", status=status)
485
+
486
+ Crew.kickoff = patched_kickoff
487
+
488
+ if hasattr(Crew, "kickoff_async"):
489
+ if not hasattr(Crew, _ORIGINAL_KICKOFF_ASYNC_ATTR):
490
+ setattr(Crew, _ORIGINAL_KICKOFF_ASYNC_ATTR, Crew.kickoff_async)
491
+ original_kickoff_async = getattr(Crew, _ORIGINAL_KICKOFF_ASYNC_ATTR)
492
+
493
+ async def patched_kickoff_async(self, *args, **kwargs):
494
+ status = "completed"
495
+ cortex_hub.start_run(framework="crewai")
496
+ try:
497
+ return await original_kickoff_async(self, *args, **kwargs)
498
+ except Exception:
499
+ status = "failed"
500
+ raise
501
+ finally:
502
+ cortex_hub.finish_run(framework="crewai", status=status)
503
+
504
+ Crew.kickoff_async = patched_kickoff_async
505
+
506
+ setattr(Crew, _PATCHED_RUN_ATTR, True)
507
+ logger.info("CrewAI run completion patched successfully")
508
+ except ImportError:
509
+ logger.debug("CrewAI run completion patch skipped")
510
+ except Exception as e:
511
+ logger.debug("CrewAI run completion patch failed", reason=str(e))
512
+
254
513
  def unpatch(self) -> None:
255
514
  """Restore original CrewAI methods."""
256
515
  try:
@@ -284,6 +543,21 @@ class CrewAIAdapter(ToolAdapter):
284
543
  logger.info("CrewAI LiteLLM unpatched")
285
544
  except ImportError:
286
545
  pass
546
+
547
+ # Restore run completion patches
548
+ try:
549
+ import crewai
550
+ Crew = getattr(crewai, "Crew", None)
551
+ if Crew is None:
552
+ from crewai.crew import Crew
553
+
554
+ if hasattr(Crew, _ORIGINAL_KICKOFF_ATTR):
555
+ Crew.kickoff = getattr(Crew, _ORIGINAL_KICKOFF_ATTR)
556
+ if hasattr(Crew, _ORIGINAL_KICKOFF_ASYNC_ATTR):
557
+ Crew.kickoff_async = getattr(Crew, _ORIGINAL_KICKOFF_ASYNC_ATTR)
558
+ setattr(Crew, _PATCHED_RUN_ATTR, False)
559
+ except ImportError:
560
+ pass
287
561
 
288
562
  except ImportError:
289
563
  pass