from functools import partial
from typing import Optional

from langchain_core.callbacks.manager import (
    Callbacks,
)
from langchain_core.prompts import BasePromptTemplate, PromptTemplate, format_document
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.retrievers import BaseRetriever

from langchain.tools import Tool


class RetrieverInput(BaseModel):
    """Input to the retriever."""

    query: str = Field(description="query to look up in retriever")


def _get_relevant_documents(
    query: str,
    retriever: BaseRetriever,
    document_prompt: BasePromptTemplate,
    document_separator: str,
    callbacks: Callbacks = None,
) -> str:
    docs = retriever.get_relevant_documents(query, callbacks=callbacks)
    return document_separator.join(
        format_document(doc, document_prompt) for doc in docs
    )


async def _aget_relevant_documents(
    query: str,
    retriever: BaseRetriever,
    document_prompt: BasePromptTemplate,
    document_separator: str,
    callbacks: Callbacks = None,
) -> str:
    docs = await retriever.aget_relevant_documents(query, callbacks=callbacks)
    return document_separator.join(
        format_document(doc, document_prompt) for doc in docs
    )


def create_retriever_tool(
    retriever: BaseRetriever,
    name: str,
    description: str,
    *,
    document_prompt: Optional[BasePromptTemplate] = None,
    document_separator: str = "\n\n",
) -> Tool:
    """Create a tool to do retrieval of documents.

    Args:
        retriever: The retriever to use for the retrieval
        name: The name for the tool. This will be passed to the language model,
            so should be unique and somewhat descriptive.
        description: The description for the tool. This will be passed to the language
            model, so should be descriptive.

    Returns:
        Tool class to pass to an agent
    """
    document_prompt = document_prompt or PromptTemplate.from_template("{page_content}")
    func = partial(
        _get_relevant_documents,
        retriever=retriever,
        document_prompt=document_prompt,
        document_separator=document_separator,
    )
    afunc = partial(
        _aget_relevant_documents,
        retriever=retriever,
        document_prompt=document_prompt,
        document_separator=document_separator,
    )
    return Tool(
        name=name,
        description=description,
        func=func,
        coroutine=afunc,
        args_schema=RetrieverInput,
    )

langchain_core.pydantic_v1.Field langchain_core.prompts.format_document langchain.tools.Tool langchain_core.prompts.PromptTemplate.from_template
[((439, 489), 'langchain_core.pydantic_v1.Field', 'Field', ([], {'description': '"""query to look up in retriever"""'}), "(description='query to look up in retriever')\n", (444, 489), False, 'from langchain_core.pydantic_v1 import BaseModel, Field\n'), ((1996, 2126), 'functools.partial', 'partial', (['_get_relevant_documents'], {'retriever': 'retriever', 'document_prompt': 'document_prompt', 'document_separator': 'document_separator'}), '(_get_relevant_documents, retriever=retriever, document_prompt=\n    document_prompt, document_separator=document_separator)\n', (2003, 2126), False, 'from functools import partial\n'), ((2173, 2304), 'functools.partial', 'partial', (['_aget_relevant_documents'], {'retriever': 'retriever', 'document_prompt': 'document_prompt', 'document_separator': 'document_separator'}), '(_aget_relevant_documents, retriever=retriever, document_prompt=\n    document_prompt, document_separator=document_separator)\n', (2180, 2304), False, 'from functools import partial\n'), ((2350, 2450), 'langchain.tools.Tool', 'Tool', ([], {'name': 'name', 'description': 'description', 'func': 'func', 'coroutine': 'afunc', 'args_schema': 'RetrieverInput'}), '(name=name, description=description, func=func, coroutine=afunc,\n    args_schema=RetrieverInput)\n', (2354, 2450), False, 'from langchain.tools import Tool\n'), ((1938, 1984), 'langchain_core.prompts.PromptTemplate.from_template', 'PromptTemplate.from_template', (['"""{page_content}"""'], {}), "('{page_content}')\n", (1966, 1984), False, 'from langchain_core.prompts import BasePromptTemplate, PromptTemplate, format_document\n'), ((796, 833), 'langchain_core.prompts.format_document', 'format_document', (['doc', 'document_prompt'], {}), '(doc, document_prompt)\n', (811, 833), False, 'from langchain_core.prompts import BasePromptTemplate, PromptTemplate, format_document\n'), ((1176, 1213), 'langchain_core.prompts.format_document', 'format_document', (['doc', 'document_prompt'], {}), '(doc, document_prompt)\n', (1191, 1213), False, 'from langchain_core.prompts import BasePromptTemplate, PromptTemplate, format_document\n')]