🔍 Code Extractor

A mock implementation of an LLM client designed for testing extractor components without making actual API calls to language models.

Purpose

MockLLMClient simulates the behavior of a real LLM client for unit testing purposes. It allows tests to define predefined responses based on prompt keywords, tracks all calls made to it, and returns mock JSON responses. This enables testing of extractor logic without incurring API costs or dealing with network latency and variability in real LLM responses.

Source Code

class MockLLMClient:
    """Mock LLMClient for testing extractors."""
    
    def __init__(self, responses=None):
        self.responses = responses or {}
        self.calls = []
    
    def generate(self, prompt, system_message=None):
        """Mock generating text with an LLM."""
        self.calls.append((prompt, system_message))
        
        # Return predefined response based on key words in the prompt
        for key, response in self.responses.items():
            if key in prompt:
                return response
        
        # Default response
        return json.dumps({"result": "mock response"})

Parameters

Parameter Details

responses: Optional dictionary mapping keyword strings to response strings. When the generate() method is called, if any key from this dictionary is found in the prompt, the corresponding response value is returned. If None or not provided, defaults to an empty dictionary. This allows tests to control what the mock returns based on prompt content.

Return Value

Instantiation returns a MockLLMClient object. The generate() method returns a string, typically JSON-formatted. If a keyword match is found in the responses dictionary, it returns the corresponding predefined response; otherwise, it returns a default JSON string: '{"result": "mock response"}'.

Class Interface

Methods

Attributes

Dependencies

• json

Required Imports

import json

Usage Example

# Basic instantiation with no predefined responses
mock_client = MockLLMClient()
response = mock_client.generate('Extract data from document')
print(response)  # '{"result": "mock response"}'

# Instantiation with predefined responses
responses = {
    'invoice': json.dumps({'invoice_number': '12345', 'total': 100.50}),
    'receipt': json.dumps({'receipt_id': 'REC-001', 'amount': 25.00})
}
mock_client = MockLLMClient(responses=responses)

# Generate with keyword match
response = mock_client.generate('Extract invoice data')
print(response)  # '{"invoice_number": "12345", "total": 100.5}'

# Check call history
print(len(mock_client.calls))  # 1
print(mock_client.calls[0])  # ('Extract invoice data', None)

# Generate with system message
response = mock_client.generate('Process receipt', system_message='You are a data extractor')
print(mock_client.calls[-1])  # ('Process receipt', 'You are a data extractor')

Best Practices

• Always instantiate with a responses dictionary when you need specific outputs for different prompt types in your tests
• Use the calls attribute to verify that the correct prompts and system messages were passed during testing
• Clear or reset the calls list between test cases if you're reusing the same mock instance
• Ensure response values in the responses dictionary are properly formatted (typically JSON strings) to match what a real LLM client would return
• Use keyword matching strategically - choose unique keywords that won't accidentally match unintended prompts
• This mock is stateless except for call tracking, making it safe to reuse across multiple test methods if responses don't need to change

Tags

Similar Components

AI-powered semantic similarity - components with related functionality:

• class TestLLMClient 80.0% similar

  Unit test class for testing the LLMClient class, which provides comprehensive test coverage for initialization, text generation, structured data extraction, and error handling across multiple LLM providers (OpenAI, Anthropic, Azure, local).

  From: /tf/active/vicechatdev/invoice_extraction/tests/test_utils.py

• function test_llm_client 72.8% similar

  Tests the LLM client functionality by analyzing a sample contract text and verifying the extraction of key contract metadata such as third parties, dates, and status.

  From: /tf/active/vicechatdev/contract_validity_analyzer/test_implementation.py

• class LLMClient_v1 68.4% similar

  A client class for interacting with Large Language Models (LLMs), specifically designed to work with OpenAI's chat completion API.

  From: /tf/active/vicechatdev/QA_updater/core/llm_client.py

• class LLMClient_v2 67.1% similar

  Client for interacting with LLM providers (OpenAI, Anthropic, Azure, etc.)

  From: /tf/active/vicechatdev/contract_validity_analyzer/utils/llm_client.py

• class LLMClient 66.4% similar

  A singleton client class for interacting with multiple LLM providers (OpenAI, Anthropic, Azure OpenAI, and local models) with unified interface for text generation and structured data extraction.

  From: /tf/active/vicechatdev/invoice_extraction/utils/llm_client.py