@illuma-ai/agents 1.1.20 → 1.1.22

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.
Files changed (246) hide show
  1. package/dist/cjs/graphs/Graph.cjs +12 -1
  2. package/dist/cjs/graphs/Graph.cjs.map +1 -1
  3. package/dist/cjs/graphs/MultiAgentGraph.cjs +85 -1
  4. package/dist/cjs/graphs/MultiAgentGraph.cjs.map +1 -1
  5. package/dist/cjs/llm/bedrock/index.cjs +14 -0
  6. package/dist/cjs/llm/bedrock/index.cjs.map +1 -1
  7. package/dist/cjs/run.cjs +20 -9
  8. package/dist/cjs/run.cjs.map +1 -1
  9. package/dist/esm/graphs/Graph.mjs +12 -1
  10. package/dist/esm/graphs/Graph.mjs.map +1 -1
  11. package/dist/esm/graphs/MultiAgentGraph.mjs +85 -1
  12. package/dist/esm/graphs/MultiAgentGraph.mjs.map +1 -1
  13. package/dist/esm/llm/bedrock/index.mjs +14 -0
  14. package/dist/esm/llm/bedrock/index.mjs.map +1 -1
  15. package/dist/esm/run.mjs +20 -9
  16. package/dist/esm/run.mjs.map +1 -1
  17. package/dist/types/graphs/MultiAgentGraph.d.ts +17 -0
  18. package/package.json +1 -1
  19. package/src/graphs/Graph.ts +12 -1
  20. package/src/graphs/MultiAgentGraph.ts +105 -1
  21. package/src/graphs/__tests__/multi-agent-delegate.test.ts +191 -0
  22. package/src/llm/bedrock/index.ts +17 -0
  23. package/src/run.ts +20 -11
  24. package/src/scripts/test-bedrock-handoff-autonomous.ts +231 -0
  25. package/src/agents/AgentContext.js +0 -782
  26. package/src/agents/AgentContext.test.js +0 -421
  27. package/src/agents/__tests__/AgentContext.test.js +0 -678
  28. package/src/agents/__tests__/resolveStructuredOutputMode.test.js +0 -117
  29. package/src/common/enum.js +0 -192
  30. package/src/common/index.js +0 -3
  31. package/src/events.js +0 -166
  32. package/src/graphs/Graph.js +0 -1857
  33. package/src/graphs/MultiAgentGraph.js +0 -1092
  34. package/src/graphs/__tests__/structured-output.integration.test.js +0 -624
  35. package/src/graphs/__tests__/structured-output.test.js +0 -144
  36. package/src/graphs/contextManagement.e2e.test.js +0 -718
  37. package/src/graphs/contextManagement.test.js +0 -485
  38. package/src/graphs/handoffValidation.test.js +0 -276
  39. package/src/graphs/index.js +0 -3
  40. package/src/index.js +0 -28
  41. package/src/instrumentation.js +0 -21
  42. package/src/llm/anthropic/index.js +0 -319
  43. package/src/llm/anthropic/types.js +0 -46
  44. package/src/llm/anthropic/utils/message_inputs.js +0 -627
  45. package/src/llm/anthropic/utils/message_outputs.js +0 -290
  46. package/src/llm/anthropic/utils/output_parsers.js +0 -89
  47. package/src/llm/anthropic/utils/tools.js +0 -25
  48. package/src/llm/bedrock/__tests__/bedrock-caching.test.js +0 -392
  49. package/src/llm/bedrock/index.js +0 -303
  50. package/src/llm/bedrock/types.js +0 -2
  51. package/src/llm/bedrock/utils/index.js +0 -6
  52. package/src/llm/bedrock/utils/message_inputs.js +0 -463
  53. package/src/llm/bedrock/utils/message_outputs.js +0 -269
  54. package/src/llm/fake.js +0 -92
  55. package/src/llm/google/index.js +0 -215
  56. package/src/llm/google/types.js +0 -12
  57. package/src/llm/google/utils/common.js +0 -670
  58. package/src/llm/google/utils/tools.js +0 -111
  59. package/src/llm/google/utils/zod_to_genai_parameters.js +0 -47
  60. package/src/llm/openai/index.js +0 -1033
  61. package/src/llm/openai/types.js +0 -2
  62. package/src/llm/openai/utils/index.js +0 -756
  63. package/src/llm/openai/utils/isReasoningModel.test.js +0 -79
  64. package/src/llm/openrouter/index.js +0 -261
  65. package/src/llm/openrouter/reasoning.test.js +0 -181
  66. package/src/llm/providers.js +0 -36
  67. package/src/llm/text.js +0 -65
  68. package/src/llm/vertexai/index.js +0 -402
  69. package/src/messages/__tests__/tools.test.js +0 -392
  70. package/src/messages/cache.js +0 -404
  71. package/src/messages/cache.test.js +0 -1167
  72. package/src/messages/content.js +0 -48
  73. package/src/messages/content.test.js +0 -314
  74. package/src/messages/core.js +0 -359
  75. package/src/messages/ensureThinkingBlock.test.js +0 -997
  76. package/src/messages/format.js +0 -973
  77. package/src/messages/formatAgentMessages.test.js +0 -2278
  78. package/src/messages/formatAgentMessages.tools.test.js +0 -362
  79. package/src/messages/formatMessage.test.js +0 -608
  80. package/src/messages/ids.js +0 -18
  81. package/src/messages/index.js +0 -9
  82. package/src/messages/labelContentByAgent.test.js +0 -725
  83. package/src/messages/prune.js +0 -438
  84. package/src/messages/reducer.js +0 -60
  85. package/src/messages/shiftIndexTokenCountMap.test.js +0 -63
  86. package/src/messages/summarize.js +0 -146
  87. package/src/messages/summarize.test.js +0 -332
  88. package/src/messages/tools.js +0 -90
  89. package/src/mockStream.js +0 -81
  90. package/src/prompts/collab.js +0 -7
  91. package/src/prompts/index.js +0 -3
  92. package/src/prompts/taskmanager.js +0 -58
  93. package/src/run.js +0 -427
  94. package/src/schemas/index.js +0 -3
  95. package/src/schemas/schema-preparation.test.js +0 -370
  96. package/src/schemas/validate.js +0 -314
  97. package/src/schemas/validate.test.js +0 -264
  98. package/src/scripts/abort.js +0 -127
  99. package/src/scripts/ant_web_search.js +0 -130
  100. package/src/scripts/ant_web_search_edge_case.js +0 -133
  101. package/src/scripts/ant_web_search_error_edge_case.js +0 -119
  102. package/src/scripts/args.js +0 -41
  103. package/src/scripts/bedrock-cache-debug.js +0 -186
  104. package/src/scripts/bedrock-content-aggregation-test.js +0 -195
  105. package/src/scripts/bedrock-merge-test.js +0 -80
  106. package/src/scripts/bedrock-parallel-tools-test.js +0 -150
  107. package/src/scripts/caching.js +0 -106
  108. package/src/scripts/cli.js +0 -152
  109. package/src/scripts/cli2.js +0 -119
  110. package/src/scripts/cli3.js +0 -163
  111. package/src/scripts/cli4.js +0 -165
  112. package/src/scripts/cli5.js +0 -165
  113. package/src/scripts/code_exec.js +0 -171
  114. package/src/scripts/code_exec_files.js +0 -180
  115. package/src/scripts/code_exec_multi_session.js +0 -185
  116. package/src/scripts/code_exec_ptc.js +0 -265
  117. package/src/scripts/code_exec_session.js +0 -217
  118. package/src/scripts/code_exec_simple.js +0 -120
  119. package/src/scripts/content.js +0 -111
  120. package/src/scripts/empty_input.js +0 -125
  121. package/src/scripts/handoff-test.js +0 -96
  122. package/src/scripts/image.js +0 -138
  123. package/src/scripts/memory.js +0 -83
  124. package/src/scripts/multi-agent-chain.js +0 -271
  125. package/src/scripts/multi-agent-conditional.js +0 -185
  126. package/src/scripts/multi-agent-document-review-chain.js +0 -171
  127. package/src/scripts/multi-agent-hybrid-flow.js +0 -264
  128. package/src/scripts/multi-agent-parallel-start.js +0 -214
  129. package/src/scripts/multi-agent-parallel.js +0 -346
  130. package/src/scripts/multi-agent-sequence.js +0 -184
  131. package/src/scripts/multi-agent-supervisor.js +0 -324
  132. package/src/scripts/multi-agent-test.js +0 -147
  133. package/src/scripts/parallel-asymmetric-tools-test.js +0 -202
  134. package/src/scripts/parallel-full-metadata-test.js +0 -176
  135. package/src/scripts/parallel-tools-test.js +0 -256
  136. package/src/scripts/programmatic_exec.js +0 -277
  137. package/src/scripts/programmatic_exec_agent.js +0 -168
  138. package/src/scripts/search.js +0 -118
  139. package/src/scripts/sequential-full-metadata-test.js +0 -143
  140. package/src/scripts/simple.js +0 -174
  141. package/src/scripts/single-agent-metadata-test.js +0 -152
  142. package/src/scripts/stream.js +0 -113
  143. package/src/scripts/test-custom-prompt-key.js +0 -132
  144. package/src/scripts/test-handoff-input.js +0 -143
  145. package/src/scripts/test-handoff-preamble.js +0 -227
  146. package/src/scripts/test-handoff-steering.js +0 -353
  147. package/src/scripts/test-multi-agent-list-handoff.js +0 -318
  148. package/src/scripts/test-parallel-agent-labeling.js +0 -253
  149. package/src/scripts/test-parallel-handoffs.js +0 -229
  150. package/src/scripts/test-thinking-handoff-bedrock.js +0 -132
  151. package/src/scripts/test-thinking-handoff.js +0 -132
  152. package/src/scripts/test-thinking-to-thinking-handoff-bedrock.js +0 -140
  153. package/src/scripts/test-tool-before-handoff-role-order.js +0 -223
  154. package/src/scripts/test-tools-before-handoff.js +0 -187
  155. package/src/scripts/test_code_api.js +0 -263
  156. package/src/scripts/thinking-bedrock.js +0 -128
  157. package/src/scripts/thinking-vertexai.js +0 -130
  158. package/src/scripts/thinking.js +0 -134
  159. package/src/scripts/tool_search.js +0 -114
  160. package/src/scripts/tools.js +0 -125
  161. package/src/specs/agent-handoffs-bedrock.integration.test.js +0 -280
  162. package/src/specs/agent-handoffs.test.js +0 -924
  163. package/src/specs/anthropic.simple.test.js +0 -287
  164. package/src/specs/azure.simple.test.js +0 -381
  165. package/src/specs/cache.simple.test.js +0 -282
  166. package/src/specs/custom-event-await.test.js +0 -148
  167. package/src/specs/deepseek.simple.test.js +0 -189
  168. package/src/specs/emergency-prune.test.js +0 -308
  169. package/src/specs/moonshot.simple.test.js +0 -237
  170. package/src/specs/observability.integration.test.js +0 -1337
  171. package/src/specs/openai.simple.test.js +0 -233
  172. package/src/specs/openrouter.simple.test.js +0 -202
  173. package/src/specs/prune.test.js +0 -733
  174. package/src/specs/reasoning.test.js +0 -144
  175. package/src/specs/spec.utils.js +0 -4
  176. package/src/specs/thinking-handoff.test.js +0 -486
  177. package/src/specs/thinking-prune.test.js +0 -600
  178. package/src/specs/token-distribution-edge-case.test.js +0 -246
  179. package/src/specs/token-memoization.test.js +0 -32
  180. package/src/specs/tokens.test.js +0 -49
  181. package/src/specs/tool-error.test.js +0 -139
  182. package/src/splitStream.js +0 -204
  183. package/src/splitStream.test.js +0 -504
  184. package/src/stream.js +0 -650
  185. package/src/stream.test.js +0 -225
  186. package/src/test/mockTools.js +0 -340
  187. package/src/tools/BrowserTools.js +0 -245
  188. package/src/tools/Calculator.js +0 -38
  189. package/src/tools/Calculator.test.js +0 -225
  190. package/src/tools/CodeExecutor.js +0 -233
  191. package/src/tools/ProgrammaticToolCalling.js +0 -602
  192. package/src/tools/StreamingToolCallBuffer.js +0 -179
  193. package/src/tools/ToolNode.js +0 -930
  194. package/src/tools/ToolSearch.js +0 -904
  195. package/src/tools/__tests__/BrowserTools.test.js +0 -306
  196. package/src/tools/__tests__/ProgrammaticToolCalling.integration.test.js +0 -276
  197. package/src/tools/__tests__/ProgrammaticToolCalling.test.js +0 -807
  198. package/src/tools/__tests__/StreamingToolCallBuffer.test.js +0 -175
  199. package/src/tools/__tests__/ToolApproval.test.js +0 -675
  200. package/src/tools/__tests__/ToolNode.recovery.test.js +0 -200
  201. package/src/tools/__tests__/ToolNode.session.test.js +0 -319
  202. package/src/tools/__tests__/ToolSearch.integration.test.js +0 -125
  203. package/src/tools/__tests__/ToolSearch.test.js +0 -812
  204. package/src/tools/__tests__/handlers.test.js +0 -799
  205. package/src/tools/__tests__/truncation-recovery.integration.test.js +0 -362
  206. package/src/tools/handlers.js +0 -306
  207. package/src/tools/schema.js +0 -25
  208. package/src/tools/search/anthropic.js +0 -34
  209. package/src/tools/search/content.js +0 -116
  210. package/src/tools/search/content.test.js +0 -133
  211. package/src/tools/search/firecrawl.js +0 -173
  212. package/src/tools/search/format.js +0 -198
  213. package/src/tools/search/highlights.js +0 -241
  214. package/src/tools/search/index.js +0 -3
  215. package/src/tools/search/jina-reranker.test.js +0 -106
  216. package/src/tools/search/rerankers.js +0 -165
  217. package/src/tools/search/schema.js +0 -102
  218. package/src/tools/search/search.js +0 -561
  219. package/src/tools/search/serper-scraper.js +0 -126
  220. package/src/tools/search/test.js +0 -129
  221. package/src/tools/search/tool.js +0 -453
  222. package/src/tools/search/types.js +0 -2
  223. package/src/tools/search/utils.js +0 -59
  224. package/src/types/graph.js +0 -24
  225. package/src/types/graph.test.js +0 -192
  226. package/src/types/index.js +0 -7
  227. package/src/types/llm.js +0 -2
  228. package/src/types/messages.js +0 -2
  229. package/src/types/run.js +0 -2
  230. package/src/types/stream.js +0 -2
  231. package/src/types/tools.js +0 -2
  232. package/src/utils/contextAnalytics.js +0 -79
  233. package/src/utils/contextAnalytics.test.js +0 -166
  234. package/src/utils/events.js +0 -26
  235. package/src/utils/graph.js +0 -11
  236. package/src/utils/handlers.js +0 -65
  237. package/src/utils/index.js +0 -10
  238. package/src/utils/llm.js +0 -21
  239. package/src/utils/llmConfig.js +0 -205
  240. package/src/utils/logging.js +0 -37
  241. package/src/utils/misc.js +0 -51
  242. package/src/utils/run.js +0 -69
  243. package/src/utils/schema.js +0 -21
  244. package/src/utils/title.js +0 -119
  245. package/src/utils/tokens.js +0 -92
  246. package/src/utils/toonFormat.js +0 -379
@@ -1,1337 +0,0 @@
1
- /**
2
- * Integration tests for Illuma Observability SDK integration in agents.
3
- *
4
- * These tests verify that the ObservabilityCallbackHandler and IllumaSpanProcessor
5
- * correctly route LangChain/OTel events to trace payloads with proper structure,
6
- * parent-child relationships, and token usage metrics.
7
- *
8
- * Uses a mock client to capture enqueued events without hitting a real server.
9
- */
10
- import { ObservabilityCallbackHandler } from '@illuma-ai/observability-langchain';
11
- // ---------------------------------------------------------------------------
12
- // Mock Client - captures enqueued events for assertions
13
- // ---------------------------------------------------------------------------
14
- class MockObservabilityClient {
15
- events = [];
16
- flushed = false;
17
- shutdownCalled = false;
18
- enqueue(event) {
19
- this.events.push(event);
20
- }
21
- async flush() {
22
- this.flushed = true;
23
- }
24
- async shutdown() {
25
- this.shutdownCalled = true;
26
- await this.flush();
27
- }
28
- /** Get events filtered by type */
29
- getByType(type) {
30
- return this.events.filter((e) => e.type === type);
31
- }
32
- /** Get all event types in order */
33
- getEventTypes() {
34
- return this.events.map((e) => e.type);
35
- }
36
- /** Pretty-print all events for debugging */
37
- printEvents() {
38
- for (const event of this.events) {
39
- const body = event.body;
40
- const usage = body.usage;
41
- console.log(` ${event.type} | name=${body.name ?? '—'} | id=${body.id ?? '—'} | traceId=${body.traceId ?? '—'} | parentObsId=${body.parentObservationId ?? '—'}` +
42
- (usage ? ` | usage={prompt:${usage.promptTokens ?? '—'},completion:${usage.completionTokens ?? '—'},total:${usage.totalTokens ?? '—'}}` : '') +
43
- (body.model ? ` | model=${body.model}` : '') +
44
- (body.provider ? ` | provider=${body.provider}` : '') +
45
- (body.level === 'ERROR' ? ` | ERROR: ${body.statusMessage}` : ''));
46
- }
47
- }
48
- }
49
- // ---------------------------------------------------------------------------
50
- // Helpers
51
- // ---------------------------------------------------------------------------
52
- const serialized = (name) => ({
53
- lc: 1,
54
- type: 'not_implemented',
55
- id: ['langchain', name],
56
- });
57
- function makeUUID() {
58
- return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
59
- const r = (Math.random() * 16) | 0;
60
- const v = c === 'x' ? r : (r & 0x3) | 0x8;
61
- return v.toString(16);
62
- });
63
- }
64
- // ---------------------------------------------------------------------------
65
- // Tests
66
- // ---------------------------------------------------------------------------
67
- describe('ObservabilityCallbackHandler integration', () => {
68
- let client;
69
- let handler;
70
- beforeEach(() => {
71
- client = new MockObservabilityClient();
72
- handler = new ObservabilityCallbackHandler({
73
- client,
74
- traceName: 'test-agent-run',
75
- userId: 'user-42',
76
- sessionId: 'session-abc',
77
- tags: ['integration-test'],
78
- metadata: { messageId: 'msg-001' },
79
- environment: 'test',
80
- debug: false,
81
- });
82
- });
83
- // -----------------------------------------------------------------------
84
- // 1. Root trace creation
85
- // -----------------------------------------------------------------------
86
- describe('root trace creation', () => {
87
- it('should create a trace on the first chain start', async () => {
88
- const chainRunId = makeUUID();
89
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'Hello agent' }, chainRunId);
90
- const traceEvents = client.getByType('trace-create');
91
- expect(traceEvents).toHaveLength(1);
92
- const trace = traceEvents[0].body;
93
- expect(trace.name).toBe('test-agent-run');
94
- expect(trace.userId).toBe('user-42');
95
- expect(trace.sessionId).toBe('session-abc');
96
- expect(trace.tags).toEqual(['integration-test']);
97
- expect(trace.metadata).toEqual({ messageId: 'msg-001' });
98
- expect(trace.environment).toBe('test');
99
- expect(trace.input).toEqual({ input: 'Hello agent' });
100
- });
101
- it('should not create a second trace for nested chains', async () => {
102
- const outerRunId = makeUUID();
103
- const innerRunId = makeUUID();
104
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'outer' }, outerRunId);
105
- await handler.handleChainStart(serialized('ChatPromptTemplate'), { input: 'inner' }, innerRunId, outerRunId);
106
- const traceEvents = client.getByType('trace-create');
107
- expect(traceEvents).toHaveLength(1); // Only one trace, not two
108
- });
109
- it('should return the trace ID after creation', async () => {
110
- expect(handler.getTraceId()).toBeNull();
111
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, makeUUID());
112
- expect(handler.getTraceId()).toBeTruthy();
113
- expect(typeof handler.getTraceId()).toBe('string');
114
- });
115
- });
116
- // -----------------------------------------------------------------------
117
- // 2. LLM generation events
118
- // -----------------------------------------------------------------------
119
- describe('LLM generation tracing', () => {
120
- it('should create generation-create with input prompts', async () => {
121
- const chainRunId = makeUUID();
122
- const llmRunId = makeUUID();
123
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
124
- await handler.handleLLMStart(serialized('ChatAnthropic'), ['You are a helpful assistant.\n\nHuman: Hello'], llmRunId, chainRunId, undefined, undefined, { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' });
125
- const genEvents = client.getByType('generation-create');
126
- expect(genEvents).toHaveLength(1);
127
- const gen = genEvents[0].body;
128
- expect(gen.name).toBe('ChatAnthropic');
129
- expect(gen.model).toBe('claude-sonnet-4-20250514');
130
- expect(gen.provider).toBe('anthropic');
131
- expect(gen.input).toEqual(['You are a helpful assistant.\n\nHuman: Hello']);
132
- expect(gen.traceId).toBe(handler.getTraceId());
133
- // Parent should be the chain
134
- const chainEvents = client.getByType('chain-create');
135
- const chainObsId = chainEvents[0].body.id;
136
- expect(gen.parentObservationId).toBe(chainObsId);
137
- });
138
- it('should create generation-update with token usage on LLM end', async () => {
139
- const chainRunId = makeUUID();
140
- const llmRunId = makeUUID();
141
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
142
- await handler.handleLLMStart(serialized('ChatOpenAI'), ['Hello'], llmRunId, chainRunId);
143
- const llmResult = {
144
- generations: [[{ text: 'Hi there! How can I help?', generationInfo: {} }]],
145
- llmOutput: {
146
- tokenUsage: {
147
- promptTokens: 150,
148
- completionTokens: 42,
149
- totalTokens: 192,
150
- },
151
- modelName: 'gpt-4o',
152
- provider: 'openai',
153
- },
154
- };
155
- await handler.handleLLMEnd(llmResult, llmRunId);
156
- const genUpdates = client.getByType('generation-update');
157
- expect(genUpdates).toHaveLength(1);
158
- const update = genUpdates[0].body;
159
- expect(update.output).toBe('Hi there! How can I help?');
160
- expect(update.model).toBe('gpt-4o');
161
- expect(update.provider).toBe('openai');
162
- expect(update.endTime).toBeDefined();
163
- const usage = update.usage;
164
- expect(usage.promptTokens).toBe(150);
165
- expect(usage.completionTokens).toBe(42);
166
- expect(usage.totalTokens).toBe(192);
167
- });
168
- it('should handle LLM errors with ERROR level', async () => {
169
- const chainRunId = makeUUID();
170
- const llmRunId = makeUUID();
171
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
172
- await handler.handleLLMStart(serialized('ChatAnthropic'), ['Hello'], llmRunId, chainRunId);
173
- await handler.handleLLMError(new Error('Rate limit exceeded'), llmRunId);
174
- const genUpdates = client.getByType('generation-update');
175
- expect(genUpdates).toHaveLength(1);
176
- const update = genUpdates[0].body;
177
- expect(update.level).toBe('ERROR');
178
- expect(update.statusMessage).toBe('Rate limit exceeded');
179
- expect(update.endTime).toBeDefined();
180
- });
181
- });
182
- // -----------------------------------------------------------------------
183
- // 3. Chain events
184
- // -----------------------------------------------------------------------
185
- describe('chain tracing', () => {
186
- it('should create chain-create events with input', async () => {
187
- const chainRunId = makeUUID();
188
- await handler.handleChainStart(serialized('RunnableSequence'), { query: 'What is AI?' }, chainRunId);
189
- const chainEvents = client.getByType('chain-create');
190
- expect(chainEvents).toHaveLength(1);
191
- const chain = chainEvents[0].body;
192
- expect(chain.name).toBe('RunnableSequence');
193
- expect(chain.input).toEqual({ query: 'What is AI?' });
194
- expect(chain.traceId).toBe(handler.getTraceId());
195
- });
196
- it('should create span-update on chain end with output', async () => {
197
- const chainRunId = makeUUID();
198
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
199
- await handler.handleChainEnd({ output: 'AI is artificial intelligence.' }, chainRunId);
200
- const spanUpdates = client.getByType('span-update');
201
- expect(spanUpdates).toHaveLength(1);
202
- const update = spanUpdates[0].body;
203
- expect(update.output).toEqual({ output: 'AI is artificial intelligence.' });
204
- expect(update.endTime).toBeDefined();
205
- });
206
- it('should handle chain error with ERROR level', async () => {
207
- const chainRunId = makeUUID();
208
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
209
- await handler.handleChainError(new Error('Chain failed: missing tool'), chainRunId);
210
- const spanUpdates = client.getByType('span-update');
211
- const errorUpdate = spanUpdates.find((e) => e.body.level === 'ERROR');
212
- expect(errorUpdate).toBeDefined();
213
- expect(errorUpdate.body.statusMessage).toBe('Chain failed: missing tool');
214
- });
215
- });
216
- // -----------------------------------------------------------------------
217
- // 4. Tool events
218
- // -----------------------------------------------------------------------
219
- describe('tool tracing', () => {
220
- it('should create tool-create events', async () => {
221
- const chainRunId = makeUUID();
222
- const toolRunId = makeUUID();
223
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
224
- await handler.handleToolStart(serialized('web_search'), '{"query": "latest AI news"}', toolRunId, chainRunId);
225
- const toolEvents = client.getByType('tool-create');
226
- expect(toolEvents).toHaveLength(1);
227
- const tool = toolEvents[0].body;
228
- expect(tool.name).toBe('Tool: web_search');
229
- expect(tool.input).toBe('{"query": "latest AI news"}');
230
- expect(tool.traceId).toBe(handler.getTraceId());
231
- });
232
- it('should create span-update on tool end', async () => {
233
- const chainRunId = makeUUID();
234
- const toolRunId = makeUUID();
235
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
236
- await handler.handleToolStart(serialized('calculator'), '{"expression": "2+2"}', toolRunId, chainRunId);
237
- await handler.handleToolEnd('4', toolRunId);
238
- const spanUpdates = client.getByType('span-update');
239
- expect(spanUpdates).toHaveLength(1);
240
- expect(spanUpdates[0].body.output).toBe('4');
241
- });
242
- it('should handle tool error', async () => {
243
- const chainRunId = makeUUID();
244
- const toolRunId = makeUUID();
245
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
246
- await handler.handleToolStart(serialized('api_call'), '{"url": "https://example.com"}', toolRunId, chainRunId);
247
- await handler.handleToolError(new Error('Connection timeout'), toolRunId);
248
- const spanUpdates = client.getByType('span-update');
249
- const errorUpdate = spanUpdates.find((e) => e.body.level === 'ERROR');
250
- expect(errorUpdate).toBeDefined();
251
- expect(errorUpdate.body.statusMessage).toBe('Connection timeout');
252
- });
253
- });
254
- // -----------------------------------------------------------------------
255
- // 5. Retriever events
256
- // -----------------------------------------------------------------------
257
- describe('retriever tracing', () => {
258
- it('should create retriever-create events', async () => {
259
- const chainRunId = makeUUID();
260
- const retrieverRunId = makeUUID();
261
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
262
- await handler.handleRetrieverStart(serialized('VectorStoreRetriever'), 'How does photosynthesis work?', retrieverRunId, chainRunId);
263
- const retrieverEvents = client.getByType('retriever-create');
264
- expect(retrieverEvents).toHaveLength(1);
265
- const retriever = retrieverEvents[0].body;
266
- expect(retriever.name).toBe('Retriever: VectorStoreRetriever');
267
- expect(retriever.input).toBe('How does photosynthesis work?');
268
- });
269
- it('should produce span-update with documents on retriever end', async () => {
270
- const chainRunId = makeUUID();
271
- const retrieverRunId = makeUUID();
272
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'test' }, chainRunId);
273
- await handler.handleRetrieverStart(serialized('VectorStoreRetriever'), 'photosynthesis', retrieverRunId, chainRunId);
274
- const documents = [
275
- { pageContent: 'Photosynthesis is the process...', metadata: { source: 'wiki', page: 1 } },
276
- { pageContent: 'Plants use chlorophyll...', metadata: { source: 'textbook', page: 42 } },
277
- ];
278
- await handler.handleRetrieverEnd(documents, retrieverRunId);
279
- const spanUpdates = client.getByType('span-update');
280
- expect(spanUpdates).toHaveLength(1);
281
- const update = spanUpdates[0].body;
282
- const output = update.output;
283
- expect(output).toHaveLength(2);
284
- expect(output[0].pageContent).toBe('Photosynthesis is the process...');
285
- expect(output[0].metadata).toEqual({ source: 'wiki', page: 1 });
286
- const meta = update.metadata;
287
- expect(meta.documentCount).toBe(2);
288
- });
289
- });
290
- // -----------------------------------------------------------------------
291
- // 6. Parent-child hierarchy
292
- // -----------------------------------------------------------------------
293
- describe('parent-child hierarchy', () => {
294
- it('should correctly link nested observations to parent', async () => {
295
- const outerChainId = makeUUID();
296
- const innerChainId = makeUUID();
297
- const llmRunId = makeUUID();
298
- const toolRunId = makeUUID();
299
- // Outer chain starts
300
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'orchestrate' }, outerChainId);
301
- // Inner chain nested under outer
302
- await handler.handleChainStart(serialized('ChatPromptTemplate'), { input: 'format prompt' }, innerChainId, outerChainId);
303
- // LLM nested under inner chain
304
- await handler.handleLLMStart(serialized('ChatOpenAI'), ['formatted prompt'], llmRunId, innerChainId);
305
- // Tool nested under outer chain
306
- await handler.handleToolStart(serialized('calculator'), '2+2', toolRunId, outerChainId);
307
- // Verify hierarchy
308
- const chainCreates = client.getByType('chain-create');
309
- const genCreates = client.getByType('generation-create');
310
- const toolCreates = client.getByType('tool-create');
311
- // All share the same traceId
312
- const traceId = handler.getTraceId();
313
- expect(chainCreates[0].body.traceId).toBe(traceId);
314
- expect(chainCreates[1].body.traceId).toBe(traceId);
315
- expect(genCreates[0].body.traceId).toBe(traceId);
316
- expect(toolCreates[0].body.traceId).toBe(traceId);
317
- // Inner chain -> parent is outer chain's observationId
318
- const outerChainObsId = chainCreates[0].body.id;
319
- const innerChainObsId = chainCreates[1].body.id;
320
- expect(chainCreates[1].body.parentObservationId).toBe(outerChainObsId);
321
- // LLM -> parent is inner chain
322
- expect(genCreates[0].body.parentObservationId).toBe(innerChainObsId);
323
- // Tool -> parent is outer chain
324
- expect(toolCreates[0].body.parentObservationId).toBe(outerChainObsId);
325
- });
326
- });
327
- // -----------------------------------------------------------------------
328
- // 7. Full agent pipeline simulation
329
- // -----------------------------------------------------------------------
330
- describe('full agent pipeline simulation', () => {
331
- it('should produce correct event sequence for a RAG agent with tools', async () => {
332
- const chainRunId = makeUUID();
333
- const agentChainId = makeUUID();
334
- const retrieverRunId = makeUUID();
335
- const llmRunId1 = makeUUID();
336
- const toolRunId = makeUUID();
337
- const llmRunId2 = makeUUID();
338
- // 1. Root chain starts (the main graph)
339
- await handler.handleChainStart(serialized('CompiledStateGraph'), { messages: [{ role: 'user', content: 'What papers did OpenAI publish in 2024?' }] }, chainRunId);
340
- // 2. Agent chain starts (nested)
341
- await handler.handleChainStart(serialized('AgentExecutor'), { input: 'What papers did OpenAI publish in 2024?' }, agentChainId, chainRunId);
342
- // 3. First LLM call (decides to use tool)
343
- await handler.handleLLMStart(serialized('ChatAnthropic'), ['System: You are a research assistant.\nHuman: What papers...'], llmRunId1, agentChainId, undefined, undefined, { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' });
344
- await handler.handleLLMEnd({
345
- generations: [[{
346
- text: 'I\'ll search for OpenAI papers from 2024.',
347
- generationInfo: {},
348
- }]],
349
- llmOutput: {
350
- tokenUsage: { promptTokens: 200, completionTokens: 30, totalTokens: 230 },
351
- modelName: 'claude-sonnet-4-20250514',
352
- provider: 'anthropic',
353
- },
354
- }, llmRunId1);
355
- // 4. Retriever runs
356
- await handler.handleRetrieverStart(serialized('VectorStoreRetriever'), 'OpenAI papers 2024', retrieverRunId, agentChainId);
357
- await handler.handleRetrieverEnd([
358
- { pageContent: 'GPT-4o: omni-model for text, vision, audio...', metadata: { source: 'arxiv', year: 2024 } },
359
- { pageContent: 'Sora: text-to-video generation model...', metadata: { source: 'blog', year: 2024 } },
360
- ], retrieverRunId);
361
- // 5. Tool call (web search)
362
- await handler.handleToolStart(serialized('web_search'), '{"query": "OpenAI 2024 research papers list"}', toolRunId, agentChainId);
363
- await handler.handleToolEnd(JSON.stringify([
364
- { title: 'GPT-4o Technical Report', url: 'https://arxiv.org/...' },
365
- { title: 'Sora', url: 'https://openai.com/sora' },
366
- ]), toolRunId);
367
- // 6. Second LLM call (final answer)
368
- await handler.handleLLMStart(serialized('ChatAnthropic'), ['System: ...\nContext: ...\nHuman: What papers...'], llmRunId2, agentChainId, undefined, undefined, { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' });
369
- await handler.handleLLMEnd({
370
- generations: [[{
371
- text: 'OpenAI published several notable papers in 2024 including GPT-4o and Sora...',
372
- generationInfo: {},
373
- }]],
374
- llmOutput: {
375
- tokenUsage: { promptTokens: 800, completionTokens: 150, totalTokens: 950 },
376
- modelName: 'claude-sonnet-4-20250514',
377
- provider: 'anthropic',
378
- },
379
- }, llmRunId2);
380
- // 7. Chains end
381
- await handler.handleChainEnd({ output: 'OpenAI published several notable papers...' }, agentChainId);
382
- await handler.handleChainEnd({ messages: [{ role: 'assistant', content: 'OpenAI published several notable papers...' }] }, chainRunId);
383
- // Flush
384
- await handler.flushAsync();
385
- expect(client.flushed).toBe(true);
386
- // --- Assertions ---
387
- // Uncomment to debug: client.printEvents();
388
- // Verify event counts
389
- const eventTypes = client.getEventTypes();
390
- const typeCounts = eventTypes.reduce((acc, t) => { acc[t] = (acc[t] || 0) + 1; return acc; }, {});
391
- expect(typeCounts['trace-create']).toBe(1); // Root trace
392
- expect(typeCounts['chain-create']).toBe(1); // Outer chain only
393
- expect(typeCounts['agent-create']).toBe(1); // AgentExecutor detected as agent!
394
- expect(typeCounts['generation-create']).toBe(2); // Two LLM calls
395
- expect(typeCounts['generation-update']).toBe(2); // Two LLM ends
396
- expect(typeCounts['retriever-create']).toBe(1); // One retriever
397
- expect(typeCounts['tool-create']).toBe(1); // One tool
398
- expect(typeCounts['span-update']).toBe(4); // retriever end + tool end + 2 chain ends
399
- // Verify all events share the same traceId
400
- const traceId = handler.getTraceId();
401
- for (const event of client.events) {
402
- const body = event.body;
403
- if (body.traceId) {
404
- expect(body.traceId).toBe(traceId);
405
- }
406
- }
407
- // Verify token usage in generation-update events
408
- const genUpdates = client.getByType('generation-update');
409
- // First LLM call usage
410
- const firstLLMUpdate = genUpdates[0].body;
411
- const firstUsage = firstLLMUpdate.usage;
412
- expect(firstUsage.promptTokens).toBe(200);
413
- expect(firstUsage.completionTokens).toBe(30);
414
- expect(firstUsage.totalTokens).toBe(230);
415
- expect(firstLLMUpdate.model).toBe('claude-sonnet-4-20250514');
416
- expect(firstLLMUpdate.provider).toBe('anthropic');
417
- // Second LLM call usage
418
- const secondLLMUpdate = genUpdates[1].body;
419
- const secondUsage = secondLLMUpdate.usage;
420
- expect(secondUsage.promptTokens).toBe(800);
421
- expect(secondUsage.completionTokens).toBe(150);
422
- expect(secondUsage.totalTokens).toBe(950);
423
- // Verify total token computation
424
- const totalPrompt = 200 + 800;
425
- const totalCompletion = 30 + 150;
426
- const totalTokens = totalPrompt + totalCompletion;
427
- expect(totalPrompt).toBe(1000);
428
- expect(totalCompletion).toBe(180);
429
- expect(totalTokens).toBe(1180);
430
- // Total: 1180 tokens (1000 prompt + 180 completion) across 2 LLM calls
431
- });
432
- });
433
- // -----------------------------------------------------------------------
434
- // 8. Token usage edge cases
435
- // -----------------------------------------------------------------------
436
- describe('token usage edge cases', () => {
437
- it('should handle different token usage field names (prompt_tokens vs promptTokens)', async () => {
438
- const chainRunId = makeUUID();
439
- const llmRunId = makeUUID();
440
- await handler.handleChainStart(serialized('Chain'), { input: 'test' }, chainRunId);
441
- await handler.handleLLMStart(serialized('ChatOpenAI'), ['test'], llmRunId, chainRunId);
442
- // OpenAI style: uses prompt_tokens (snake_case)
443
- await handler.handleLLMEnd({
444
- generations: [[{ text: 'response', generationInfo: {} }]],
445
- llmOutput: {
446
- tokenUsage: {
447
- prompt_tokens: 100,
448
- completion_tokens: 50,
449
- total_tokens: 150,
450
- },
451
- },
452
- }, llmRunId);
453
- const updates = client.getByType('generation-update');
454
- const usage = updates[0].body.usage;
455
- // The handler should normalize both field naming conventions
456
- expect(usage.promptTokens ?? usage.prompt_tokens).toBeTruthy();
457
- });
458
- it('should handle missing token usage gracefully', async () => {
459
- const chainRunId = makeUUID();
460
- const llmRunId = makeUUID();
461
- await handler.handleChainStart(serialized('Chain'), { input: 'test' }, chainRunId);
462
- await handler.handleLLMStart(serialized('ChatOpenAI'), ['test'], llmRunId, chainRunId);
463
- // No llmOutput at all
464
- await handler.handleLLMEnd({ generations: [[{ text: 'response', generationInfo: {} }]] }, llmRunId);
465
- const updates = client.getByType('generation-update');
466
- expect(updates).toHaveLength(1);
467
- // Should not crash, usage should be present but with undefined values
468
- const usage = updates[0].body.usage;
469
- expect(usage).toBeDefined();
470
- });
471
- });
472
- // -----------------------------------------------------------------------
473
- // 9. Flush and shutdown
474
- // -----------------------------------------------------------------------
475
- describe('flush and shutdown', () => {
476
- it('flushAsync should call client.flush()', async () => {
477
- await handler.handleChainStart(serialized('Chain'), { input: 'test' }, makeUUID());
478
- await handler.flushAsync();
479
- expect(client.flushed).toBe(true);
480
- });
481
- it('shutdown should call client.flush() (not shutdown) for external client', async () => {
482
- await handler.handleChainStart(serialized('Chain'), { input: 'test' }, makeUUID());
483
- await handler.shutdown();
484
- // Since we provided an external client, it should only flush, not shutdown
485
- expect(client.flushed).toBe(true);
486
- expect(client.shutdownCalled).toBe(false);
487
- });
488
- });
489
- // -----------------------------------------------------------------------
490
- // 10. Agent detection
491
- // -----------------------------------------------------------------------
492
- describe('agent detection', () => {
493
- it('should emit agent-create for chains with langgraph_node agent= metadata', async () => {
494
- const outerChainId = makeUUID();
495
- const agentChainId = makeUUID();
496
- await handler.handleChainStart(serialized('CompiledStateGraph'), { messages: [{ role: 'user', content: 'Hello' }] }, outerChainId);
497
- // Agent chain with langgraph_node metadata
498
- // Cast to any to pass tags/metadata params (ts-jest resolves different .d.ts for BaseCallbackHandler)
499
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'agent task' }, agentChainId, outerChainId, undefined, // tags
500
- { langgraph_node: 'agent=research-agent' });
501
- const agentEvents = client.getByType('agent-create');
502
- expect(agentEvents).toHaveLength(1);
503
- const agentBody = agentEvents[0].body;
504
- expect(agentBody.name).toBe('Agent: research-agent');
505
- expect(agentBody.traceId).toBe(handler.getTraceId());
506
- const meta = agentBody.metadata;
507
- expect(meta.agentId).toBe('research-agent');
508
- expect(meta.langchainMetadata).toEqual({ langgraph_node: 'agent=research-agent' });
509
- });
510
- it('should emit agent-create for AgentExecutor chain name', async () => {
511
- const agentRunId = makeUUID();
512
- await handler.handleChainStart(serialized('AgentExecutor'), { input: 'do something' }, agentRunId);
513
- const agentEvents = client.getByType('agent-create');
514
- expect(agentEvents).toHaveLength(1);
515
- expect(agentEvents[0].body.name).toBe('AgentExecutor');
516
- // Should NOT produce a chain-create
517
- const chainEvents = client.getByType('chain-create');
518
- expect(chainEvents).toHaveLength(0);
519
- });
520
- it('should emit agent-create for chain names ending in Agent', async () => {
521
- const agentRunId = makeUUID();
522
- await handler.handleChainStart(serialized('ReactAgent'), { input: 'think and act' }, agentRunId);
523
- const agentEvents = client.getByType('agent-create');
524
- expect(agentEvents).toHaveLength(1);
525
- expect(agentEvents[0].body.name).toBe('ReactAgent');
526
- });
527
- it('should NOT emit agent-create for regular chains', async () => {
528
- const chainRunId = makeUUID();
529
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'regular chain' }, chainRunId);
530
- const agentEvents = client.getByType('agent-create');
531
- expect(agentEvents).toHaveLength(0);
532
- const chainEvents = client.getByType('chain-create');
533
- expect(chainEvents).toHaveLength(1);
534
- });
535
- it('should correctly end agent observations with span-update', async () => {
536
- const agentRunId = makeUUID();
537
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'agent task' }, agentRunId, undefined, // parentRunId
538
- undefined, // tags
539
- { langgraph_node: 'agent=planner' });
540
- await handler.handleChainEnd({ output: 'agent completed' }, agentRunId);
541
- const agentEvents = client.getByType('agent-create');
542
- expect(agentEvents).toHaveLength(1);
543
- const spanUpdates = client.getByType('span-update');
544
- expect(spanUpdates).toHaveLength(1);
545
- const update = spanUpdates[0].body;
546
- expect(update.output).toEqual({ output: 'agent completed' });
547
- expect(update.endTime).toBeDefined();
548
- // The span-update should reference the same observation ID as agent-create
549
- expect(update.id).toBe(agentEvents[0].body.id);
550
- });
551
- it('should handle agent error with ERROR level', async () => {
552
- const agentRunId = makeUUID();
553
- await handler.handleChainStart(serialized('AgentExecutor'), { input: 'fail task' }, agentRunId);
554
- await handler.handleChainError(new Error('Agent failed: tool not found'), agentRunId);
555
- const agentEvents = client.getByType('agent-create');
556
- expect(agentEvents).toHaveLength(1);
557
- const spanUpdates = client.getByType('span-update');
558
- const errorUpdate = spanUpdates.find((e) => e.body.level === 'ERROR');
559
- expect(errorUpdate).toBeDefined();
560
- expect(errorUpdate.body.statusMessage).toBe('Agent failed: tool not found');
561
- });
562
- it('should integrate agent-create in full pipeline with correct counts', async () => {
563
- const rootChainId = makeUUID();
564
- const agentChainId = makeUUID();
565
- const llmRunId = makeUUID();
566
- // Root chain
567
- await handler.handleChainStart(serialized('CompiledStateGraph'), { messages: [{ role: 'user', content: 'Hello' }] }, rootChainId);
568
- // Agent chain (detected via metadata)
569
- await handler.handleChainStart(serialized('RunnableSequence'), { input: 'agent work' }, agentChainId, rootChainId, undefined, // tags
570
- { langgraph_node: 'agent=qa-agent' });
571
- // LLM inside agent
572
- await handler.handleLLMStart(serialized('ChatAnthropic'), ['prompt'], llmRunId, agentChainId, undefined, undefined, { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' });
573
- await handler.handleLLMEnd({
574
- generations: [[{ text: 'answer', generationInfo: {} }]],
575
- llmOutput: {
576
- tokenUsage: { promptTokens: 100, completionTokens: 20, totalTokens: 120 },
577
- modelName: 'claude-sonnet-4-20250514',
578
- },
579
- }, llmRunId);
580
- await handler.handleChainEnd({ output: 'done' }, agentChainId);
581
- await handler.handleChainEnd({ output: 'done' }, rootChainId);
582
- const typeCounts = client.getEventTypes().reduce((acc, t) => { acc[t] = (acc[t] || 0) + 1; return acc; }, {});
583
- expect(typeCounts['trace-create']).toBe(1);
584
- expect(typeCounts['chain-create']).toBe(1); // Only root chain
585
- expect(typeCounts['agent-create']).toBe(1); // Agent detected!
586
- expect(typeCounts['generation-create']).toBe(1);
587
- expect(typeCounts['generation-update']).toBe(1);
588
- expect(typeCounts['span-update']).toBe(2); // agent end + root chain end
589
- // LLM parent should be the agent observation
590
- const agentObsId = client.getByType('agent-create')[0].body.id;
591
- const genCreate = client.getByType('generation-create')[0].body;
592
- expect(genCreate.parentObservationId).toBe(agentObsId);
593
- });
594
- });
595
- // -----------------------------------------------------------------------
596
- // 11. Event structure validation
597
- // -----------------------------------------------------------------------
598
- describe('event structure validation', () => {
599
- it('every event should have id, type, timestamp, and body', async () => {
600
- const chainRunId = makeUUID();
601
- const llmRunId = makeUUID();
602
- await handler.handleChainStart(serialized('Chain'), { input: 'test' }, chainRunId);
603
- await handler.handleLLMStart(serialized('LLM'), ['test'], llmRunId, chainRunId);
604
- await handler.handleLLMEnd({ generations: [[{ text: 'ok', generationInfo: {} }]], llmOutput: { tokenUsage: {} } }, llmRunId);
605
- await handler.handleChainEnd({ output: 'done' }, chainRunId);
606
- for (const event of client.events) {
607
- expect(event.id).toBeDefined();
608
- expect(typeof event.id).toBe('string');
609
- expect(event.type).toBeDefined();
610
- expect(event.timestamp).toMatch(/^\d{4}-\d{2}-\d{2}T/);
611
- expect(event.body).toBeDefined();
612
- expect(typeof event.body).toBe('object');
613
- }
614
- });
615
- it('observation events should have traceId and observationId', async () => {
616
- const chainRunId = makeUUID();
617
- await handler.handleChainStart(serialized('Chain'), { input: 'test' }, chainRunId);
618
- await handler.handleChainEnd({ output: 'done' }, chainRunId);
619
- const nonTraceEvents = client.events.filter((e) => e.type !== 'trace-create');
620
- for (const event of nonTraceEvents) {
621
- const body = event.body;
622
- expect(body.traceId).toBe(handler.getTraceId());
623
- expect(body.id).toBeDefined();
624
- }
625
- });
626
- });
627
- });
628
- // ---------------------------------------------------------------------------
629
- // IllumaSpanProcessor unit tests (OTel integration)
630
- // ---------------------------------------------------------------------------
631
- describe('IllumaSpanProcessor integration', () => {
632
- it('should be importable from @illuma-ai/observability-otel', async () => {
633
- const { IllumaSpanProcessor } = await import('@illuma-ai/observability-otel');
634
- expect(IllumaSpanProcessor).toBeDefined();
635
- expect(typeof IllumaSpanProcessor).toBe('function');
636
- });
637
- });
638
- // ---------------------------------------------------------------------------
639
- // Core SDK features: sampling, PII masking, observation types, prompt/dataset
640
- // ---------------------------------------------------------------------------
641
- describe('Core SDK features', () => {
642
- // Use require to bypass ts-jest module resolution issues with dynamic imports
643
- // eslint-disable-next-line @typescript-eslint/no-require-imports
644
- const { ObservabilityCoreClient, TraceClient, SpanClient, IngestionEventType, } = require('@illuma-ai/observability-core');
645
- // Create a concrete subclass at runtime to avoid TS abstract class issues
646
- const TestClientClass = class extends ObservabilityCoreClient {
647
- async fetchWithRetry() {
648
- return { status: 200, statusText: 'OK', ok: true, json: async () => ({}), text: async () => '' };
649
- }
650
- };
651
- function createTestClient(config = {}) {
652
- const events = [];
653
- const client = new TestClientClass({
654
- publicKey: 'pk-test',
655
- secretKey: 'sk-test',
656
- baseUrl: 'http://localhost:9999',
657
- flushInterval: 0,
658
- ...config,
659
- });
660
- // Intercept enqueue to capture events
661
- const origEnqueue = client.enqueue.bind(client);
662
- client.enqueue = (event) => {
663
- origEnqueue(event);
664
- events.push(event);
665
- };
666
- return { client, events };
667
- }
668
- // -----------------------------------------------------------------------
669
- // 12. Sampling rate
670
- // -----------------------------------------------------------------------
671
- describe('sampling rate', () => {
672
- it('should trace everything when sampleRate=1.0 (default)', () => {
673
- const { client, events } = createTestClient({ sampleRate: 1.0 });
674
- for (let i = 0; i < 20; i++) {
675
- client.trace({ name: `trace-${i}` });
676
- }
677
- // All 20 traces should be created (each produces a trace-create event queued)
678
- const traceEvents = events.filter((e) => e.type === 'trace-create');
679
- expect(traceEvents.length).toBe(20);
680
- });
681
- it('should trace nothing when sampleRate=0.0', () => {
682
- const { client, events } = createTestClient({ sampleRate: 0.0 });
683
- for (let i = 0; i < 20; i++) {
684
- const trace = client.trace({ name: `trace-${i}` });
685
- // Even child events should be silently dropped
686
- trace.generation({ name: 'gen', model: 'gpt-4o' });
687
- trace.span({ name: 'span' });
688
- }
689
- // No events should be enqueued at all
690
- expect(events.length).toBe(0);
691
- });
692
- it('should sample approximately the right percentage', () => {
693
- // Use a fixed seed approach: run many traces with 50% sample rate
694
- const { client, events } = createTestClient({ sampleRate: 0.5 });
695
- const iterations = 1000;
696
- for (let i = 0; i < iterations; i++) {
697
- client.trace({ name: `trace-${i}` });
698
- }
699
- const traceEvents = events.filter((e) => e.type === 'trace-create');
700
- // With 1000 iterations at 50%, expect roughly 400-600 traces
701
- expect(traceEvents.length).toBeGreaterThan(300);
702
- expect(traceEvents.length).toBeLessThan(700);
703
- });
704
- it('should include all children once a trace is sampled in', () => {
705
- const { client, events } = createTestClient({ sampleRate: 1.0 });
706
- const trace = client.trace({ name: 'sampled-trace' });
707
- trace.span({ name: 'child-span' });
708
- trace.generation({ name: 'child-gen', model: 'gpt-4o' });
709
- trace.agent({ name: 'child-agent' });
710
- trace.tool({ name: 'child-tool' });
711
- // 1 trace + 4 children = 5 events
712
- expect(events.length).toBe(5);
713
- });
714
- it('should clamp sampleRate to [0, 1] range', () => {
715
- // sampleRate > 1 should be clamped to 1 (trace everything)
716
- const { events: events1 } = createTestClient({ sampleRate: 5.0 });
717
- // sampleRate < 0 should be clamped to 0 (trace nothing)
718
- const { client: client2, events: events2 } = createTestClient({ sampleRate: -1 });
719
- client2.trace({ name: 'should-not-appear' });
720
- expect(events2.length).toBe(0);
721
- });
722
- });
723
- // -----------------------------------------------------------------------
724
- // 13. PII masking
725
- // -----------------------------------------------------------------------
726
- describe('PII masking', () => {
727
- it('should apply mask function to event bodies', () => {
728
- const maskFunction = (body) => {
729
- const masked = { ...body };
730
- if (typeof masked.input === 'string') {
731
- masked.input = masked.input.replace(/[\w.-]+@[\w.-]+\.\w+/g, '[EMAIL]');
732
- }
733
- return masked;
734
- };
735
- const { client } = createTestClient({ maskFunction });
736
- const trace = client.trace({ name: 'pii-test' });
737
- trace.span({
738
- name: 'user-query',
739
- input: 'My email is john@example.com and I need help',
740
- });
741
- // Check the queued events - the span should have masked input
742
- const queue = client.getQueue();
743
- const spanEvent = [...queue].find((e) => e.type === 'span-create');
744
- if (spanEvent) {
745
- expect(spanEvent.body.input).toBe('My email is [EMAIL] and I need help');
746
- expect(spanEvent.body.input).not.toContain('john@example.com');
747
- }
748
- });
749
- it('should mask multiple PII patterns', () => {
750
- const maskFunction = (body) => {
751
- const masked = { ...body };
752
- const str = JSON.stringify(masked);
753
- const redacted = str
754
- .replace(/[\w.-]+@[\w.-]+\.\w+/g, '[EMAIL]')
755
- .replace(/\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/g, '[PHONE]')
756
- .replace(/\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b/g, '[CARD]');
757
- return JSON.parse(redacted);
758
- };
759
- const { client } = createTestClient({ maskFunction });
760
- const trace = client.trace({
761
- name: 'multi-pii',
762
- metadata: { userEmail: 'test@corp.com', phone: '555-123-4567' },
763
- });
764
- const queue = client.getQueue();
765
- const traceEvent = [...queue].find((e) => e.type === 'trace-create');
766
- if (traceEvent) {
767
- const meta = traceEvent.body.metadata;
768
- expect(meta.userEmail).toBe('[EMAIL]');
769
- expect(meta.phone).toBe('[PHONE]');
770
- }
771
- });
772
- it('should gracefully handle mask function errors', () => {
773
- const maskFunction = () => {
774
- throw new Error('Mask function exploded');
775
- };
776
- const { client } = createTestClient({ maskFunction });
777
- // Should not throw — falls back to unmasked
778
- expect(() => {
779
- client.trace({ name: 'should-not-crash' });
780
- }).not.toThrow();
781
- const queue = client.getQueue();
782
- expect(queue.length).toBeGreaterThan(0);
783
- });
784
- it('should not modify events when no mask function is set', () => {
785
- const { client } = createTestClient();
786
- client.trace({
787
- name: 'no-mask',
788
- metadata: { email: 'visible@example.com' },
789
- });
790
- const queue = client.getQueue();
791
- const traceEvent = [...queue].find((e) => e.type === 'trace-create');
792
- expect(traceEvent.body.metadata.email).toBe('visible@example.com');
793
- });
794
- });
795
- // -----------------------------------------------------------------------
796
- // 14. All observation types (TraceClient + SpanClient)
797
- // -----------------------------------------------------------------------
798
- describe('observation types completeness', () => {
799
- it('TraceClient should have all observation type methods', () => {
800
- const enqueue = jest.fn();
801
- const trace = new TraceClient('trace-123', enqueue);
802
- // All typed observation methods should exist and return SpanClient or GenerationClient
803
- expect(typeof trace.span).toBe('function');
804
- expect(typeof trace.agent).toBe('function');
805
- expect(typeof trace.tool).toBe('function');
806
- expect(typeof trace.chain).toBe('function');
807
- expect(typeof trace.retriever).toBe('function');
808
- expect(typeof trace.guardrail).toBe('function');
809
- expect(typeof trace.evaluator).toBe('function');
810
- expect(typeof trace.embedding).toBe('function');
811
- expect(typeof trace.generation).toBe('function');
812
- expect(typeof trace.event).toBe('function');
813
- expect(typeof trace.score).toBe('function');
814
- });
815
- it('SpanClient should have all observation type methods', () => {
816
- const enqueue = jest.fn();
817
- const span = new SpanClient('span-123', 'trace-123', enqueue);
818
- expect(typeof span.span).toBe('function');
819
- expect(typeof span.agent).toBe('function');
820
- expect(typeof span.tool).toBe('function');
821
- expect(typeof span.chain).toBe('function');
822
- expect(typeof span.retriever).toBe('function');
823
- expect(typeof span.guardrail).toBe('function');
824
- expect(typeof span.evaluator).toBe('function');
825
- expect(typeof span.embedding).toBe('function');
826
- expect(typeof span.generation).toBe('function');
827
- expect(typeof span.event).toBe('function');
828
- expect(typeof span.score).toBe('function');
829
- });
830
- it('should emit correct event types for each observation method on TraceClient', () => {
831
- const events = [];
832
- const enqueue = (event) => events.push(event);
833
- const trace = new TraceClient('trace-456', enqueue);
834
- trace.span({ name: 'span' });
835
- trace.agent({ name: 'agent' });
836
- trace.tool({ name: 'tool' });
837
- trace.chain({ name: 'chain' });
838
- trace.retriever({ name: 'retriever' });
839
- trace.guardrail({ name: 'guardrail' });
840
- trace.evaluator({ name: 'evaluator' });
841
- trace.embedding({ name: 'embedding' });
842
- trace.generation({ name: 'generation' });
843
- trace.event({ name: 'event' });
844
- const types = events.map((e) => e.type);
845
- expect(types).toContain('span-create');
846
- expect(types).toContain('agent-create');
847
- expect(types).toContain('tool-create');
848
- expect(types).toContain('chain-create');
849
- expect(types).toContain('retriever-create');
850
- expect(types).toContain('guardrail-create');
851
- expect(types).toContain('evaluator-create');
852
- expect(types).toContain('embedding-create');
853
- expect(types).toContain('generation-create');
854
- expect(types).toContain('event-create');
855
- });
856
- it('should emit correct event types for each observation method on SpanClient', () => {
857
- const events = [];
858
- const enqueue = (event) => events.push(event);
859
- const span = new SpanClient('span-789', 'trace-789', enqueue);
860
- span.span({ name: 'child-span' });
861
- span.agent({ name: 'child-agent' });
862
- span.tool({ name: 'child-tool' });
863
- span.chain({ name: 'child-chain' });
864
- span.retriever({ name: 'child-retriever' });
865
- span.guardrail({ name: 'child-guardrail' });
866
- span.evaluator({ name: 'child-evaluator' });
867
- span.embedding({ name: 'child-embedding' });
868
- span.generation({ name: 'child-generation' });
869
- span.event({ name: 'child-event' });
870
- const types = events.map((e) => e.type);
871
- expect(types).toContain('span-create');
872
- expect(types).toContain('agent-create');
873
- expect(types).toContain('tool-create');
874
- expect(types).toContain('chain-create');
875
- expect(types).toContain('retriever-create');
876
- expect(types).toContain('guardrail-create');
877
- expect(types).toContain('evaluator-create');
878
- expect(types).toContain('embedding-create');
879
- expect(types).toContain('generation-create');
880
- expect(types).toContain('event-create');
881
- });
882
- it('SpanClient children should set parentObservationId automatically', () => {
883
- const events = [];
884
- const enqueue = (event) => events.push(event);
885
- const span = new SpanClient('parent-span', 'trace-abc', enqueue);
886
- span.guardrail({ name: 'guardrail-check' });
887
- span.evaluator({ name: 'eval-run' });
888
- span.embedding({ name: 'embed-op' });
889
- for (const event of events) {
890
- expect(event.body.parentObservationId).toBe('parent-span');
891
- expect(event.body.traceId).toBe('trace-abc');
892
- }
893
- });
894
- });
895
- // -----------------------------------------------------------------------
896
- // 15. Prompt & Dataset SDK methods
897
- // -----------------------------------------------------------------------
898
- describe('prompt and dataset SDK methods', () => {
899
- it('core client should expose getPrompt method', () => {
900
- const { client } = createTestClient();
901
- expect(typeof client.getPrompt).toBe('function');
902
- });
903
- it('core client should expose createPrompt method', () => {
904
- const { client } = createTestClient();
905
- expect(typeof client.createPrompt).toBe('function');
906
- });
907
- it('core client should expose getDataset method', () => {
908
- const { client } = createTestClient();
909
- expect(typeof client.getDataset).toBe('function');
910
- });
911
- it('core client should expose createDataset method', () => {
912
- const { client } = createTestClient();
913
- expect(typeof client.createDataset).toBe('function');
914
- });
915
- it('core client should expose createDatasetItem method', () => {
916
- const { client } = createTestClient();
917
- expect(typeof client.createDatasetItem).toBe('function');
918
- });
919
- it('getPrompt should return parsed response from server', async () => {
920
- const { client } = createTestClient();
921
- // Mock fetchWithRetry returns { ok: true, json: async () => ({}) }
922
- const result = await client.getPrompt({ name: 'my-prompt' });
923
- expect(result).toBeDefined();
924
- expect(result).toEqual({});
925
- });
926
- it('getDataset should return parsed response from server', async () => {
927
- const { client } = createTestClient();
928
- const result = await client.getDataset({ name: 'my-dataset' });
929
- expect(result).toBeDefined();
930
- expect(result).toEqual({});
931
- });
932
- it('prompt/dataset methods should return null when disabled', async () => {
933
- const { client } = createTestClient({ enabled: false });
934
- expect(await client.getPrompt({ name: 'test' })).toBeNull();
935
- expect(await client.createPrompt({ name: 'test', prompt: 'hello' })).toBeNull();
936
- expect(await client.getDataset({ name: 'test' })).toBeNull();
937
- expect(await client.createDataset({ name: 'test' })).toBeNull();
938
- expect(await client.createDatasetItem({ datasetName: 'test', input: 'x' })).toBeNull();
939
- });
940
- });
941
- // -----------------------------------------------------------------------
942
- // 16. Node SDK: sampling + PII masking propagation
943
- // -----------------------------------------------------------------------
944
- describe('node SDK config propagation', () => {
945
- // eslint-disable-next-line @typescript-eslint/no-require-imports
946
- const { Observability } = require('@illuma-ai/observability-node');
947
- it('Observability class should accept sampleRate option', async () => {
948
- expect(Observability).toBeDefined();
949
- const obs = new Observability({
950
- publicKey: 'pk-test',
951
- secretKey: 'sk-test',
952
- baseUrl: 'http://localhost:9999',
953
- sampleRate: 0.5,
954
- flushInterval: 0,
955
- });
956
- expect(obs).toBeDefined();
957
- await obs.shutdown();
958
- });
959
- it('Observability class should accept maskFunction option', async () => {
960
- const obs = new Observability({
961
- publicKey: 'pk-test',
962
- secretKey: 'sk-test',
963
- baseUrl: 'http://localhost:9999',
964
- maskFunction: (body) => ({ ...body, input: '[REDACTED]' }),
965
- flushInterval: 0,
966
- });
967
- expect(obs).toBeDefined();
968
- await obs.shutdown();
969
- });
970
- });
971
- });
972
- // ---------------------------------------------------------------------------
973
- // 17. Guardrail tracing via ObservabilityCallbackHandler
974
- // ---------------------------------------------------------------------------
975
- /** Re-use MockObservabilityClient shape for guardrail tests (same file, top-level scope) */
976
- class GuardrailMockClient {
977
- events = [];
978
- flushed = false;
979
- shutdownCalled = false;
980
- enqueue(event) { this.events.push(event); }
981
- async flush() { this.flushed = true; }
982
- async shutdown() { this.shutdownCalled = true; }
983
- getByType(type) { return this.events.filter((e) => e.type === type); }
984
- getEventTypes() { return this.events.map((e) => e.type); }
985
- printEvents() {
986
- for (const event of this.events) {
987
- const body = event.body;
988
- const usage = body.usage;
989
- console.log(` ${event.type} | name=${body.name ?? '—'} | id=${body.id ?? '—'} | traceId=${body.traceId ?? '—'} | parentObsId=${body.parentObservationId ?? '—'}` +
990
- (usage ? ` | usage={prompt:${usage.promptTokens ?? '—'},completion:${usage.completionTokens ?? '—'},total:${usage.totalTokens ?? '—'}}` : '') +
991
- (body.model ? ` | model=${body.model}` : '') +
992
- (body.level === 'ERROR' || body.level === 'WARNING' ? ` | ${body.level}: ${body.statusMessage}` : ''));
993
- }
994
- }
995
- }
996
- const grSerialized = (name) => ({
997
- lc: 1,
998
- type: 'not_implemented',
999
- id: ['langchain', name],
1000
- });
1001
- describe('Guardrail tracing', () => {
1002
- let grClient;
1003
- let grHandler; // Use any to bypass ts-jest type resolution for traceGuardrail
1004
- beforeEach(() => {
1005
- grClient = new GuardrailMockClient();
1006
- grHandler = new ObservabilityCallbackHandler({
1007
- client: grClient,
1008
- });
1009
- });
1010
- it('outcome=passed should create guardrail-create + span-update with DEFAULT level', async () => {
1011
- const chainRunId = makeUUID();
1012
- await grHandler.handleChainStart(grSerialized('RunnableSequence'), { input: 'hello' }, chainRunId);
1013
- const result = grHandler.traceGuardrail({
1014
- name: 'Output Moderation',
1015
- guardrailId: 'bedrock-guardrail-123',
1016
- guardrailVersion: '1',
1017
- outcome: 'passed',
1018
- actionApplied: false,
1019
- action: 'NONE',
1020
- reason: 'passed',
1021
- source: 'OUTPUT',
1022
- input: 'What is the weather?',
1023
- output: 'The weather is sunny.',
1024
- violations: [],
1025
- assessments: [{ contentPolicy: { filters: [] } }],
1026
- });
1027
- expect(result).not.toBeNull();
1028
- const guardrailEvents = grClient.getByType('guardrail-create');
1029
- expect(guardrailEvents).toHaveLength(1);
1030
- const grBody = guardrailEvents[0].body;
1031
- expect(grBody.name).toBe('Output Moderation');
1032
- expect(grBody.input).toBe('What is the weather?');
1033
- expect(grBody.metadata.guardrailId).toBe('bedrock-guardrail-123');
1034
- expect(grBody.metadata.outcome).toBe('passed');
1035
- expect(grBody.metadata.source).toBe('OUTPUT');
1036
- // span-update should show PASSED
1037
- const spanUpdates = grClient.getByType('span-update');
1038
- const guardrailUpdate = spanUpdates.find((e) => e.body.id === result);
1039
- expect(guardrailUpdate).toBeDefined();
1040
- expect(guardrailUpdate.body.output).toBe('The weather is sunny.');
1041
- expect(guardrailUpdate.body.level).toBe('DEFAULT');
1042
- expect(guardrailUpdate.body.statusMessage).toBe('PASSED');
1043
- });
1044
- it('outcome=blocked should mark with ERROR level when enforced', async () => {
1045
- const chainRunId = makeUUID();
1046
- await grHandler.handleChainStart(grSerialized('RunnableSequence'), { input: 'test' }, chainRunId);
1047
- grHandler.traceGuardrail({
1048
- name: 'Input Moderation',
1049
- guardrailId: 'bedrock-guardrail-456',
1050
- outcome: 'blocked',
1051
- actionApplied: true,
1052
- action: 'GUARDRAIL_INTERVENED',
1053
- reason: 'policy_violation',
1054
- source: 'INPUT',
1055
- input: 'How to hack a system?',
1056
- output: 'Sorry, I cannot help with that.',
1057
- violations: [{ type: 'CONTENT_POLICY', category: 'VIOLENCE', action: 'BLOCKED' }],
1058
- });
1059
- const spanUpdates = grClient.getByType('span-update');
1060
- const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
1061
- expect(guardrailUpdate.body.level).toBe('ERROR');
1062
- expect(guardrailUpdate.body.statusMessage).toBe('BLOCKED: policy_violation');
1063
- expect(guardrailUpdate.body.metadata.outcome).toBe('blocked');
1064
- expect(guardrailUpdate.body.metadata.actionApplied).toBe(true);
1065
- expect(guardrailUpdate.body.metadata.violations).toHaveLength(1);
1066
- });
1067
- it('outcome=blocked with actionApplied=false should note "not enforced"', async () => {
1068
- const chainRunId = makeUUID();
1069
- await grHandler.handleChainStart(grSerialized('RunnableSequence'), { input: 'test' }, chainRunId);
1070
- grHandler.traceGuardrail({
1071
- name: 'Input Moderation',
1072
- outcome: 'blocked',
1073
- actionApplied: false,
1074
- reason: 'policy_violation',
1075
- source: 'INPUT',
1076
- input: 'Some flagged content',
1077
- });
1078
- const spanUpdates = grClient.getByType('span-update');
1079
- const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
1080
- expect(guardrailUpdate.body.level).toBe('ERROR');
1081
- expect(guardrailUpdate.body.statusMessage).toContain('not enforced');
1082
- expect(guardrailUpdate.body.metadata.actionApplied).toBe(false);
1083
- });
1084
- it('outcome=anonymized should mark with WARNING and include modified content', async () => {
1085
- const chainRunId = makeUUID();
1086
- await grHandler.handleChainStart(grSerialized('RunnableSequence'), { input: 'test' }, chainRunId);
1087
- grHandler.traceGuardrail({
1088
- name: 'Output Moderation',
1089
- outcome: 'anonymized',
1090
- actionApplied: true,
1091
- action: 'GUARDRAIL_INTERVENED',
1092
- reason: 'anonymized',
1093
- source: 'OUTPUT',
1094
- input: 'My email is john@example.com and SSN is 123-45-6789',
1095
- originalContent: 'My email is john@example.com and SSN is 123-45-6789',
1096
- modifiedContent: 'My email is [EMAIL] and SSN is [SSN]',
1097
- violations: [
1098
- { type: 'PII_POLICY', category: 'EMAIL', action: 'ANONYMIZED' },
1099
- { type: 'PII_POLICY', category: 'SSN', action: 'ANONYMIZED' },
1100
- ],
1101
- });
1102
- const spanUpdates = grClient.getByType('span-update');
1103
- const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
1104
- expect(guardrailUpdate.body.level).toBe('WARNING');
1105
- expect(guardrailUpdate.body.statusMessage).toBe('ANONYMIZED: PII detected and masked');
1106
- // output should be the modified content
1107
- expect(guardrailUpdate.body.output).toBe('My email is [EMAIL] and SSN is [SSN]');
1108
- const meta = guardrailUpdate.body.metadata;
1109
- expect(meta.originalContent).toBe('My email is john@example.com and SSN is 123-45-6789');
1110
- expect(meta.modifiedContent).toBe('My email is [EMAIL] and SSN is [SSN]');
1111
- expect(meta.violations).toHaveLength(2);
1112
- });
1113
- it('outcome=intervened should mark with WARNING level', async () => {
1114
- const chainRunId = makeUUID();
1115
- await grHandler.handleChainStart(grSerialized('RunnableSequence'), { input: 'test' }, chainRunId);
1116
- grHandler.traceGuardrail({
1117
- name: 'Output Moderation',
1118
- outcome: 'intervened',
1119
- actionApplied: true,
1120
- action: 'GUARDRAIL_INTERVENED',
1121
- reason: 'intervened_passthrough',
1122
- source: 'OUTPUT',
1123
- input: 'Original response text',
1124
- modifiedContent: 'Modified response text',
1125
- });
1126
- const spanUpdates = grClient.getByType('span-update');
1127
- const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
1128
- expect(guardrailUpdate.body.level).toBe('WARNING');
1129
- expect(guardrailUpdate.body.statusMessage).toBe('INTERVENED: intervened_passthrough');
1130
- });
1131
- it('traceGuardrail() should return null when no trace exists', () => {
1132
- const freshClient = new GuardrailMockClient();
1133
- const freshHandler = new ObservabilityCallbackHandler({
1134
- client: freshClient,
1135
- });
1136
- const result = freshHandler.traceGuardrail({
1137
- name: 'Pre-trace Guardrail',
1138
- outcome: 'passed',
1139
- actionApplied: false,
1140
- source: 'INPUT',
1141
- });
1142
- expect(result).toBeNull();
1143
- expect(freshClient.getByType('guardrail-create')).toHaveLength(0);
1144
- });
1145
- it('traceGuardrail() should use custom startTime and endTime', async () => {
1146
- const chainRunId = makeUUID();
1147
- await grHandler.handleChainStart(grSerialized('Chain'), { input: 'test' }, chainRunId);
1148
- const startTime = '2024-01-01T00:00:00.000Z';
1149
- const endTime = '2024-01-01T00:00:01.500Z';
1150
- grHandler.traceGuardrail({
1151
- name: 'Timed Guardrail',
1152
- outcome: 'passed',
1153
- actionApplied: false,
1154
- source: 'OUTPUT',
1155
- startTime,
1156
- endTime,
1157
- });
1158
- const guardrailEvents = grClient.getByType('guardrail-create');
1159
- expect(guardrailEvents[0].body.startTime).toBe(startTime);
1160
- const spanUpdates = grClient.getByType('span-update');
1161
- const lastUpdate = spanUpdates[spanUpdates.length - 1];
1162
- expect(lastUpdate.body.endTime).toBe(endTime);
1163
- });
1164
- it('should trace both input and output guardrails on the same trace', async () => {
1165
- const chainRunId = makeUUID();
1166
- await grHandler.handleChainStart(grSerialized('RunnableSequence'), { input: 'test' }, chainRunId);
1167
- const inputId = grHandler.traceGuardrail({
1168
- name: 'Input Moderation',
1169
- outcome: 'passed',
1170
- actionApplied: false,
1171
- source: 'INPUT',
1172
- input: 'What is 2+2?',
1173
- });
1174
- const outputId = grHandler.traceGuardrail({
1175
- name: 'Output Moderation',
1176
- outcome: 'passed',
1177
- actionApplied: false,
1178
- source: 'OUTPUT',
1179
- input: 'The answer is 4.',
1180
- });
1181
- expect(inputId).not.toBeNull();
1182
- expect(outputId).not.toBeNull();
1183
- expect(inputId).not.toBe(outputId);
1184
- const guardrailEvents = grClient.getByType('guardrail-create');
1185
- expect(guardrailEvents).toHaveLength(2);
1186
- const traceId = guardrailEvents[0].body.traceId;
1187
- expect(guardrailEvents[1].body.traceId).toBe(traceId);
1188
- expect(guardrailEvents[0].body.metadata.source).toBe('INPUT');
1189
- expect(guardrailEvents[1].body.metadata.source).toBe('OUTPUT');
1190
- });
1191
- it('should include full AWS Bedrock assessments and violations in metadata', async () => {
1192
- const chainRunId = makeUUID();
1193
- await grHandler.handleChainStart(grSerialized('Chain'), { input: 'test' }, chainRunId);
1194
- grHandler.traceGuardrail({
1195
- name: 'Detailed Guardrail',
1196
- outcome: 'blocked',
1197
- actionApplied: true,
1198
- action: 'GUARDRAIL_INTERVENED',
1199
- reason: 'policy_violation',
1200
- source: 'INPUT',
1201
- violations: [
1202
- { type: 'CONTENT_POLICY', category: 'VIOLENCE', confidence: 'HIGH', action: 'BLOCKED' },
1203
- { type: 'PII_POLICY', category: 'SSN', action: 'BLOCKED' },
1204
- ],
1205
- assessments: [
1206
- {
1207
- contentPolicy: { filters: [{ type: 'VIOLENCE', confidence: 'HIGH', action: 'BLOCKED' }] },
1208
- sensitiveInformationPolicy: { piiEntities: [{ type: 'SSN', action: 'BLOCKED' }] },
1209
- },
1210
- ],
1211
- });
1212
- const spanUpdates = grClient.getByType('span-update');
1213
- const lastUpdate = spanUpdates[spanUpdates.length - 1];
1214
- const meta = lastUpdate.body.metadata;
1215
- expect(meta.violations).toHaveLength(2);
1216
- expect(meta.assessments).toHaveLength(1);
1217
- expect(meta.violations[0]).toEqual({
1218
- type: 'CONTENT_POLICY', category: 'VIOLENCE', confidence: 'HIGH', action: 'BLOCKED',
1219
- });
1220
- });
1221
- });
1222
- // ---------------------------------------------------------------------------
1223
- // 18. Guardrail tracing with debug logging — full pipeline with all 4 outcomes
1224
- // ---------------------------------------------------------------------------
1225
- describe('Guardrail tracing with debug output', () => {
1226
- it('should produce debug-friendly trace output for manual verification', async () => {
1227
- const debugClient = new GuardrailMockClient();
1228
- const debugHandler = new ObservabilityCallbackHandler({
1229
- client: debugClient,
1230
- traceName: 'guardrail-debug-test',
1231
- userId: 'user-123',
1232
- sessionId: 'session-456',
1233
- debug: true,
1234
- });
1235
- // 1. Chain starts (creates trace)
1236
- const chainRunId = makeUUID();
1237
- await debugHandler.handleChainStart(grSerialized('RunnableSequence'), { messages: [{ role: 'user', content: 'Explain quantum computing' }] }, chainRunId);
1238
- // 2. LLM call
1239
- const llmRunId = makeUUID();
1240
- await debugHandler.handleLLMStart(grSerialized('ChatOpenAI'), ['Explain quantum computing'], llmRunId, chainRunId);
1241
- const llmResult = {
1242
- generations: [[{ text: 'Quantum computing uses qubits...', generationInfo: {} }]],
1243
- llmOutput: {
1244
- tokenUsage: { promptTokens: 10, completionTokens: 25, totalTokens: 35 },
1245
- modelName: 'gpt-4o',
1246
- },
1247
- };
1248
- await debugHandler.handleLLMEnd(llmResult, llmRunId);
1249
- await debugHandler.handleChainEnd({ output: 'Quantum computing uses qubits...' }, chainRunId);
1250
- // 3. Input guardrail — passed
1251
- const inputGuardrailId = debugHandler.traceGuardrail({
1252
- name: 'Input Moderation (Bedrock)',
1253
- guardrailId: 'arn:aws:bedrock:us-east-1:123456:guardrail/abc123',
1254
- guardrailVersion: '3',
1255
- outcome: 'passed',
1256
- actionApplied: false,
1257
- action: 'NONE',
1258
- reason: 'passed',
1259
- source: 'INPUT',
1260
- input: 'Explain quantum computing',
1261
- startTime: '2024-06-15T10:00:00.000Z',
1262
- endTime: '2024-06-15T10:00:00.150Z',
1263
- });
1264
- // 4. Output guardrail — passed
1265
- const outputGuardrailId = debugHandler.traceGuardrail({
1266
- name: 'Output Moderation (Bedrock)',
1267
- guardrailId: 'arn:aws:bedrock:us-east-1:123456:guardrail/abc123',
1268
- guardrailVersion: '3',
1269
- outcome: 'passed',
1270
- actionApplied: false,
1271
- action: 'NONE',
1272
- reason: 'passed',
1273
- source: 'OUTPUT',
1274
- input: 'Quantum computing uses qubits...',
1275
- violations: [],
1276
- assessments: [{ topicPolicy: { topics: [] }, contentPolicy: { filters: [] } }],
1277
- startTime: '2024-06-15T10:00:01.000Z',
1278
- endTime: '2024-06-15T10:00:01.200Z',
1279
- });
1280
- expect(inputGuardrailId).not.toBeNull();
1281
- expect(outputGuardrailId).not.toBeNull();
1282
- console.log('\n=== SAMPLE GUARDRAIL TRACE OUTPUT ===');
1283
- debugClient.printEvents();
1284
- console.log('=== END TRACE ===\n');
1285
- // Verify complete event sequence
1286
- const types = debugClient.getEventTypes();
1287
- expect(types).toEqual([
1288
- 'trace-create',
1289
- 'chain-create',
1290
- 'generation-create',
1291
- 'generation-update',
1292
- 'span-update', // chain end
1293
- 'guardrail-create', // input guardrail
1294
- 'span-update', // input guardrail update
1295
- 'guardrail-create', // output guardrail
1296
- 'span-update', // output guardrail update
1297
- ]);
1298
- // Verify trace ID consistency
1299
- const traceId = debugClient.getByType('trace-create')[0].body.id;
1300
- for (const evt of debugClient.events) {
1301
- const body = evt.body;
1302
- if (evt.type !== 'trace-create') {
1303
- expect(body.traceId).toBe(traceId);
1304
- }
1305
- }
1306
- });
1307
- it('should trace all 4 guardrail outcomes with correct levels', async () => {
1308
- const debugClient = new GuardrailMockClient();
1309
- const handler = new ObservabilityCallbackHandler({
1310
- client: debugClient,
1311
- debug: true,
1312
- });
1313
- // Create trace
1314
- const chainRunId = makeUUID();
1315
- await handler.handleChainStart(grSerialized('Chain'), { input: 'test' }, chainRunId);
1316
- // All 4 outcomes
1317
- handler.traceGuardrail({ name: 'Passed', outcome: 'passed', actionApplied: false, source: 'INPUT' });
1318
- handler.traceGuardrail({ name: 'Blocked', outcome: 'blocked', actionApplied: true, reason: 'policy_violation', source: 'INPUT' });
1319
- handler.traceGuardrail({ name: 'Anonymized', outcome: 'anonymized', actionApplied: true, reason: 'anonymized', source: 'OUTPUT', modifiedContent: '[EMAIL]' });
1320
- handler.traceGuardrail({ name: 'Intervened', outcome: 'intervened', actionApplied: true, reason: 'intervened_passthrough', source: 'OUTPUT' });
1321
- console.log('\n=== ALL 4 GUARDRAIL OUTCOMES ===');
1322
- debugClient.printEvents();
1323
- console.log('=== END ===\n');
1324
- const guardrailEvents = debugClient.getByType('guardrail-create');
1325
- expect(guardrailEvents).toHaveLength(4);
1326
- const spanUpdates = debugClient.getByType('span-update');
1327
- // 4 guardrail updates (chain not ended in this test)
1328
- expect(spanUpdates).toHaveLength(4);
1329
- // Check levels: passed=DEFAULT, blocked=ERROR, anonymized=WARNING, intervened=WARNING
1330
- const guardrailUpdates = spanUpdates;
1331
- expect(guardrailUpdates[0].body.level).toBe('DEFAULT');
1332
- expect(guardrailUpdates[1].body.level).toBe('ERROR');
1333
- expect(guardrailUpdates[2].body.level).toBe('WARNING');
1334
- expect(guardrailUpdates[3].body.level).toBe('WARNING');
1335
- });
1336
- });
1337
- //# sourceMappingURL=observability.integration.test.js.map