learningcircuit--local-deep-research/examples/llm_integration/basic_custom_llm.py

"""
Example of using custom LangChain LLMs with Local Deep Research.

This example shows how to integrate your own LLM implementations or wrappers
with LDR's research functions.
"""

from typing import Any, List, Optional

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatResult

from local_deep_research.api import (
    create_settings_snapshot,
    detailed_research,
    quick_summary,
)


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

    model_name: str = "custom-model"
    temperature: float = 0.7

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[Any] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate a response. This is where you'd call your custom model."""
        # This is a mock implementation - replace with your actual model call
        response_text = f"This is a response from {self.model_name} to: {messages[-1].content}"

        # Create the response
        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"


def custom_llm_factory(
    model_name: str = "factory-model", temperature: float = 0.5, **kwargs
):
    """Factory function that creates a custom LLM instance."""
    return CustomLLM(model_name=model_name, temperature=temperature)


def main():
    # Example 1: Using a custom LLM instance
    custom_llm = CustomLLM(model_name="my-custom-model", temperature=0.8)

    snapshot = create_settings_snapshot(
        provider="my_custom",
        overrides={"search.tool": "wikipedia"},
    )
    result = quick_summary(
        query="What are the latest advances in quantum computing?",
        llms={"my_custom": custom_llm},
        settings_snapshot=snapshot,
    )

    print("Summary with custom LLM instance:")
    print(result["summary"])
    print("-" * 80)

    # Example 2: Using a factory function
    snapshot = create_settings_snapshot(
        provider="factory_llm",
        temperature=0.3,
        overrides={"search.tool": "wikipedia"},
    )
    result = quick_summary(
        query="Explain the benefits of renewable energy",
        llms={"factory_llm": custom_llm_factory},
        model_name="renewable-expert",  # This gets passed to the factory
        settings_snapshot=snapshot,
    )

    print("\nSummary with factory-created LLM:")
    print(result["summary"])
    print("-" * 80)

    # Example 3: Multiple custom LLMs
    llms = {
        "technical": CustomLLM(model_name="technical-writer", temperature=0.2),
        "creative": CustomLLM(model_name="creative-writer", temperature=0.9),
    }

    # Technical analysis
    snapshot = create_settings_snapshot(
        provider="technical",
        overrides={"search.tool": "arxiv"},
    )
    technical_result = detailed_research(
        query="How do neural networks work?",
        llms=llms,
        settings_snapshot=snapshot,
    )

    print("\nTechnical analysis:")
    print(technical_result["summary"])
    print("-" * 80)

    # Creative exploration
    snapshot = create_settings_snapshot(
        provider="creative",
        overrides={"search.tool": "wikipedia"},
    )
    creative_result = quick_summary(
        query="What are the philosophical implications of AI?",
        llms=llms,
        settings_snapshot=snapshot,
    )

    print("\nCreative exploration:")
    print(creative_result["summary"])


if __name__ == "__main__":
    main()