@traccia2/sdk 0.0.1

package/src/__tests__/integrations-langchain.test.ts

@@ -0,0 +1,384 @@
+/**
+ * Tests for LangChain integration (TraciaCallbackHandler)
+ */
+
+import { TraciaCallbackHandler } from '../integrations/langchain-callback';
+import { startTracing, stopTracing } from '../auto';
+
+describe('TraciaCallbackHandler', () => {
+  beforeEach(() => {
+    startTracing({
+      enableTokenCounting: false,
+      enableCostTracking: false,
+    });
+  });
+
+  afterEach(() => {
+    stopTracing();
+  });
+
+  describe('handleLLMStart', () => {
+    it('should create a span for LLM invocation', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockLLM = {
+        name: 'gpt-4',
+        _modelType: 'openai',
+      };
+
+      await handler.handleLLMStart(mockLLM, ['prompt1', 'prompt2'], 'run-1');
+
+      // Verify span was created in the handler's span map
+      expect(handler['spanStack'].has('run-1')).toBe(true);
+    });
+
+    it('should capture LLM model name in attributes', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockLLM = {
+        name: 'claude-3',
+        _modelType: 'anthropic',
+      };
+
+      await handler.handleLLMStart(mockLLM, ['test prompt'], 'run-2');
+
+      const span = handler['spanStack'].get('run-2');
+      expect(span?.attributes?.model).toMatch(/claude-3|anthropic/);
+    });
+
+    it('should capture prompt count', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockLLM = { name: 'test-model' };
+      const prompts = ['p1', 'p2', 'p3'];
+
+      await handler.handleLLMStart(mockLLM, prompts, 'run-3');
+
+      const span = handler['spanStack'].get('run-3');
+      expect(span?.attributes?.prompt_count).toBe(3);
+    });
+
+    it('should handle missing model name gracefully', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockLLM = {}; // No name or _modelType
+
+      await handler.handleLLMStart(mockLLM, ['prompt'], 'run-4');
+
+      const span = handler['spanStack'].get('run-4');
+      expect(span?.attributes?.model).toBe('unknown');
+    });
+  });
+
+  describe('handleLLMEnd', () => {
+    it('should end LLM span and record token counts', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockLLM = { name: 'gpt-4' };
+
+      // Start span
+      await handler.handleLLMStart(mockLLM, ['prompt'], 'run-5');
+
+      // End span
+      const mockOutput = {
+        generations: [[{ text: 'response', generation_info: { usage: { prompt_tokens: 10, completion_tokens: 5 } } }]],
+      };
+      await handler.handleLLMEnd(mockOutput, 'run-5');
+
+      // Verify span was removed after ending
+      expect(handler['spanStack'].has('run-5')).toBe(false);
+    });
+
+    it('should record output length', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockLLM = { name: 'test' };
+
+      await handler.handleLLMStart(mockLLM, ['prompt'], 'run-6');
+
+      const output = {
+        generations: [[{ text: 'This is a test response that has some length to it' }]],
+      };
+      await handler.handleLLMEnd(output, 'run-6');
+
+      // Span should be ended
+      expect(handler['spanStack'].has('run-6')).toBe(false);
+    });
+  });
+
+  describe('handleChainStart', () => {
+    it('should create a span for chain execution', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockChain = {
+        name: 'test-chain',
+        _chainType: 'stuff',
+      };
+
+      await handler.handleChainStart(mockChain, { input: 'test' }, 'chain-1');
+
+      expect(handler['spanStack'].has('chain-1')).toBe(true);
+    });
+
+    it('should capture chain name and type', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockChain = {
+        name: 'qa-chain',
+        _chainType: 'retrieval_qa',
+      };
+
+      await handler.handleChainStart(mockChain, { query: 'test' }, 'chain-2');
+
+      const span = handler['spanStack'].get('chain-2');
+      expect(span?.attributes?.chain_name).toBe('qa-chain');
+      expect(span?.attributes?.chain_type).toBe('retrieval_qa');
+    });
+
+    it('should record input keys', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockChain = { name: 'chain' };
+      const inputs = { key1: 'value1', key2: 'value2' };
+
+      await handler.handleChainStart(mockChain, inputs, 'chain-3');
+
+      const span = handler['spanStack'].get('chain-3');
+      expect(span?.attributes?.input_keys).toContain('key1');
+      expect(span?.attributes?.input_keys).toContain('key2');
+    });
+  });
+
+  describe('handleChainEnd', () => {
+    it('should end chain span and record output', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockChain = { name: 'chain' };
+
+      await handler.handleChainStart(mockChain, { input: 'test' }, 'chain-4');
+
+      const output = { output: 'result' };
+      await handler.handleChainEnd(output, 'chain-4');
+
+      expect(handler['spanStack'].has('chain-4')).toBe(false);
+    });
+  });
+
+  describe('handleToolStart', () => {
+    it('should create a span for tool invocation', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockTool = {
+        name: 'search-api',
+      };
+
+      await handler.handleToolStart(mockTool, 'search query', 'tool-1');
+
+      expect(handler['spanStack'].has('tool-1')).toBe(true);
+    });
+
+    it('should capture tool name', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockTool = { name: 'calculator' };
+
+      await handler.handleToolStart(mockTool, '2+2', 'tool-2');
+
+      const span = handler['spanStack'].get('tool-2');
+      expect(span?.attributes?.tool_name).toBe('calculator');
+    });
+
+    it('should record input length', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockTool = { name: 'tool' };
+      const input = 'test input with some length';
+
+      await handler.handleToolStart(mockTool, input, 'tool-3');
+
+      const span = handler['spanStack'].get('tool-3');
+      expect(span?.attributes?.input_length).toBeGreaterThan(0);
+    });
+  });
+
+  describe('handleToolEnd', () => {
+    it('should end tool span', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockTool = { name: 'tool' };
+
+      await handler.handleToolStart(mockTool, 'input', 'tool-4');
+      await handler.handleToolEnd('output', 'tool-4');
+
+      expect(handler['spanStack'].has('tool-4')).toBe(false);
+    });
+
+    it('should record output length', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockTool = { name: 'tool' };
+
+      await handler.handleToolStart(mockTool, 'input', 'tool-5');
+      const output = 'This is the tool output response';
+      await handler.handleToolEnd(output, 'tool-5');
+
+      expect(handler['spanStack'].has('tool-5')).toBe(false);
+    });
+  });
+
+  describe('handleAgentAction', () => {
+    it('should update existing span on agent action', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockChain = { name: 'chain' };
+
+      // Start a span first
+      await handler.handleChainStart(mockChain, {}, 'agent-1');
+
+      // Then handle agent action
+      const mockAction = {
+        tool: 'search',
+        toolInput: 'query',
+      };
+      await handler.handleAgentAction(mockAction, 'agent-1');
+
+      // Span should still exist
+      expect(handler['spanStack'].has('agent-1')).toBe(true);
+    });
+
+    it('should record agent action details', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockChain = { name: 'chain' };
+
+      // Start a span first
+      await handler.handleChainStart(mockChain, {}, 'agent-2');
+
+      // Then handle agent action
+      const mockAction = {
+        tool: 'calculator',
+        toolInput: '42 / 7',
+      };
+      await handler.handleAgentAction(mockAction, 'agent-2');
+
+      // Span should have agent action recorded
+      const span = handler['spanStack'].get('agent-2');
+      expect(span).toBeDefined();
+    });
+  });
+
+  describe('handleAgentFinish', () => {
+    it('should end agent span on finish', async () => {
+      const handler = new TraciaCallbackHandler();
+
+      // Create a span first
+      const mockAction = { tool: 'test', toolInput: 'input' };
+      await handler.handleAgentAction(mockAction, 'agent-3');
+
+      // Finish agent
+      const finish = { output: 'final result' };
+      await handler.handleAgentFinish(finish, 'agent-3');
+
+      expect(handler['spanStack'].has('agent-3')).toBe(false);
+    });
+
+    it('should record final output', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockAction = { tool: 'test', toolInput: 'input' };
+
+      await handler.handleAgentAction(mockAction, 'agent-4');
+
+      const finish = { output: 'completed result' };
+      await handler.handleAgentFinish(finish, 'agent-4');
+
+      expect(handler['spanStack'].has('agent-4')).toBe(false);
+    });
+  });
+
+  describe('Error Handling', () => {
+    it('should handle LLM errors gracefully', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockLLM = { name: 'llm' };
+
+      await handler.handleLLMStart(mockLLM, ['prompt'], 'error-1');
+
+      const error = new Error('LLM failed');
+      await handler.handleLLMError?.(error, 'error-1');
+
+      // Span should still be cleaned up
+      expect(handler['spanStack'].has('error-1')).toBe(false);
+    });
+
+    it('should handle chain errors gracefully', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockChain = { name: 'chain' };
+
+      await handler.handleChainStart(mockChain, {}, 'error-2');
+
+      const error = new Error('Chain failed');
+      await handler.handleChainError?.(error, 'error-2');
+
+      expect(handler['spanStack'].has('error-2')).toBe(false);
+    });
+
+    it('should handle tool errors gracefully', async () => {
+      const handler = new TraciaCallbackHandler();
+      const mockTool = { name: 'tool' };
+
+      await handler.handleToolStart(mockTool, 'input', 'error-3');
+
+      const error = new Error('Tool failed');
+      await handler.handleToolError?.(error, 'error-3');
+
+      expect(handler['spanStack'].has('error-3')).toBe(false);
+    });
+  });
+
+  describe('Span Nesting', () => {
+    it('should handle nested spans (chain containing LLM)', async () => {
+      const handler = new TraciaCallbackHandler();
+
+      // Start chain
+      const chain = { name: 'chain' };
+      await handler.handleChainStart(chain, {}, 'chain-outer');
+
+      // Start LLM inside chain
+      const llm = { name: 'llm' };
+      await handler.handleLLMStart(llm, ['prompt'], 'llm-inner');
+
+      // Both should exist
+      expect(handler['spanStack'].has('chain-outer')).toBe(true);
+      expect(handler['spanStack'].has('llm-inner')).toBe(true);
+
+      // End LLM first
+      await handler.handleLLMEnd({ generations: [[{ text: 'response' }]] }, 'llm-inner');
+      expect(handler['spanStack'].has('llm-inner')).toBe(false);
+
+      // Chain should still exist
+      expect(handler['spanStack'].has('chain-outer')).toBe(true);
+
+      // End chain
+      await handler.handleChainEnd({ output: 'result' }, 'chain-outer');
+      expect(handler['spanStack'].has('chain-outer')).toBe(false);
+    });
+
+    it('should handle multiple concurrent spans', async () => {
+      const handler = new TraciaCallbackHandler();
+
+      // Start multiple spans
+      const chain1 = { name: 'chain1' };
+      const chain2 = { name: 'chain2' };
+
+      await handler.handleChainStart(chain1, {}, 'c1');
+      await handler.handleChainStart(chain2, {}, 'c2');
+
+      expect(handler['spanStack'].has('c1')).toBe(true);
+      expect(handler['spanStack'].has('c2')).toBe(true);
+
+      // End one
+      await handler.handleChainEnd({ output: 'r1' }, 'c1');
+
+      // Other should still exist
+      expect(handler['spanStack'].has('c1')).toBe(false);
+      expect(handler['spanStack'].has('c2')).toBe(true);
+    });
+  });
+
+  describe('Integration with Tracer', () => {
+    it('should use the SDK tracer', () => {
+      const handler = new TraciaCallbackHandler();
+      expect(handler['tracer']).toBeDefined();
+    });
+
+    it('should handle unavailable tracer gracefully', async () => {
+      const handler = new TraciaCallbackHandler();
+      
+      // Should not throw even if operations fail
+      const mockLLM = { name: 'llm' };
+      await expect(handler.handleLLMStart(mockLLM, ['prompt'], 'test')).resolves.not.toThrow();
+    });
+  });
+});