# 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

```python
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

```python
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
```python
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
```python
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
```python
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:

```python
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

```python
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

```python
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
```python
# 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
```python
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
```python
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](api-quickstart.md)
- [Configuration Guide](env_configuration.md)
- [Full Configuration Reference](CONFIGURATION.md)
- [LangChain Retriever Integration](LANGCHAIN_RETRIEVER_INTEGRATION.md)
