Custom LLM Integration Guide

Custom LLM Integration Guide

Local Deep Research now supports seamless integration with custom LangChain LLMs, allowing you to use your own language models, specialized wrappers, or third-party LLM providers alongside the built-in options.

Overview

Similar to the custom retriever support, LDR allows you to register any LangChain-compatible LLM and use it throughout the system. This enables:

• Using proprietary or fine-tuned models
• Implementing custom retry logic or preprocessing
• Integrating with LLM providers not built into LDR
• Testing with mock LLMs
• Creating specialized model configurations

Quick Start

from local_deep_research.api import quick_summary

# Option 1: Pass an LLM instance
result = quick_summary(
    query="Your research question",
    llms={"my_model": your_llm_instance},
    provider="my_model"  # Use your custom LLM
)

# Option 2: Pass a factory function
def create_llm(model_name=None, temperature=0.7, **kwargs):
    return YourCustomLLM(model=model_name, temp=temperature)

result = quick_summary(
    query="Your research question",
    llms={"custom": create_llm},
    provider="custom",
    model_name="gpt-turbo",  # Passed to factory
    temperature=0.5
)

Requirements

Your custom LLM must:

1. Inherit from langchain_core.language_models.BaseChatModel
2. Implement the required methods (_generate, _llm_type)
3. Handle the standard LangChain message formats

Example Implementation

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage
from langchain_core.outputs import ChatResult, ChatGeneration
from typing import List, Optional, Any

class CustomLLM(BaseChatModel):
    """Example custom LLM implementation."""

    def __init__(self, api_key: str, model_name: str = "custom-v1", **kwargs):
        super().__init__(**kwargs)
        self.api_key = api_key
        self.model_name = model_name

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[Any] = None,
        **kwargs: Any
    ) -> ChatResult:
        """Generate a response from your model."""
        # Call your API/model here
        response = your_api_call(messages, self.model_name, self.api_key)

        # Convert to LangChain format
        message = AIMessage(content=response.text)
        generation = ChatGeneration(message=message)

        return ChatResult(generations=[generation])

    @property
    def _llm_type(self) -> str:
        """Return identifier for this LLM."""
        return "custom"

Using with Different Research Modes

Quick Summary

from local_deep_research.api import quick_summary

result = quick_summary(
    query="Explain quantum computing",
    llms={"quantum_expert": quantum_llm},
    provider="quantum_expert",
    search_tool="arxiv"  # Search scientific papers
)

Detailed Research

from local_deep_research.api import detailed_research

result = detailed_research(
    query="Climate change impacts",
    llms={"climate_model": climate_specialized_llm},
    provider="climate_model",
    iterations=3
)

Report Generation

from local_deep_research.api import generate_report

report = generate_report(
    query="AI in healthcare",
    llms={"medical_ai": medical_llm},
    provider="medical_ai",
    output_file="healthcare_ai_report.md"
)

Advanced Usage

Multiple Custom LLMs

Register multiple LLMs for different purposes:

llms = {
    "technical": TechnicalWriterLLM(temperature=0.2),
    "creative": CreativeWriterLLM(temperature=0.9),
    "fact_checker": FactCheckingLLM(temperature=0.0)
}

# Use technical LLM for precise analysis
result = quick_summary(
    query="How do transformers work?",
    llms=llms,
    provider="technical"
)

Factory Functions with Configuration

def create_configured_llm(model_name=None, temperature=0.7, max_retries=3, **kwargs):
    """Factory that creates LLM with retry logic."""
    base_llm = YourLLM(model=model_name, temperature=temperature)
    return RetryWrapper(base_llm, max_retries=max_retries)

result = quick_summary(
    query="Your question",
    llms={"retry_llm": create_configured_llm},
    provider="retry_llm",
    model_name="your-model-v2",
    max_retries=5  # Custom parameter
)

Combining Custom LLMs and Retrievers

result = quick_summary(
    query="Internal policy on remote work",
    llms={"company_llm": company_fine_tuned_llm},
    retrievers={"company_docs": company_retriever},
    provider="company_llm",
    search_tool="company_docs"
)

Implementation Details

How It Works

1. Registration: When you pass LLMs via the llms parameter, they are registered in a global registry
2. Provider Check: When creating an LLM, the system first checks if the provider name matches a registered custom LLM
3. Factory Support: If the registered LLM is callable, it's treated as a factory and called with the provided parameters
4. Wrapping: All LLMs (custom and built-in) are wrapped with think-tag removal and token counting

Thread Safety

The LLM registry is thread-safe, allowing concurrent usage in multi-threaded applications.

Scope

Registered LLMs are available globally within the Python process. They persist until explicitly unregistered or the process ends.

Best Practices

1. Consistent Naming: Use clear, descriptive names for your custom LLMs
2. Error Handling: Implement proper error handling in your LLM's _generate method
3. Token Counting: If your LLM supports token counting, implement the appropriate methods
4. Temperature Handling: Respect the temperature parameter for consistency
5. Async Support: Implement async methods if your LLM supports asynchronous operation

Common Use Cases

Fine-tuned Models

# Use your fine-tuned model for domain-specific research
fine_tuned_llm = CustomLLM(
    model_path="/path/to/fine-tuned-model",
    domain="medical"
)

result = quick_summary(
    query="Latest treatments for condition X",
    llms={"medical_expert": fine_tuned_llm},
    provider="medical_expert"
)

Mock LLMs for Testing

class MockLLM(BaseChatModel):
    """Returns predefined responses for testing."""

    def _generate(self, messages, **kwargs):
        # Return test data
        return ChatResult(generations=[
            ChatGeneration(message=AIMessage(content="Test response"))
        ])

# Use in tests
result = quick_summary(
    query="Test query",
    llms={"mock": MockLLM()},
    provider="mock",
    search_tool="none"  # Disable search for pure testing
)

Rate-Limited Wrapper

class RateLimitedLLM(BaseChatModel):
    """Adds rate limiting to any LLM."""

    def __init__(self, base_llm, requests_per_minute=10):
        super().__init__()
        self.base_llm = base_llm
        self.rate_limiter = RateLimiter(requests_per_minute)

    def _generate(self, messages, **kwargs):
        self.rate_limiter.wait_if_needed()
        return self.base_llm._generate(messages, **kwargs)

Troubleshooting

LLM Not Found

If you get "Invalid provider" errors:

• Ensure you're passing the llms parameter to the API function
• Check that the provider name matches exactly (case-insensitive)
• Verify your LLM instance is properly initialized

Parameter Passing

When using factory functions:

• Standard parameters (model_name, temperature) are passed automatically
• Custom parameters can be passed via kwargs
• The factory receives all parameters from the API call

Compatibility Issues

Ensure your LLM:

• Inherits from BaseChatModel
• Returns proper ChatResult objects
• Handles the LangChain message format

Related Documentation

• API Documentation
• Configuration Guide
• LangChain Retriever Integration