import json
from typing import List

from langchain.pydantic_v1 import BaseModel, Field
from langchain.schema import BaseRetriever, Document
from langchain.tools import Tool

from backend.chat_bot.json_decoder import CustomJSONEncoder


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


def create_retriever_tool(
        retriever: BaseRetriever,
        tool_name: str,
        description: str
) -> Tool:
    """Create a tool to do retrieval of documents.

    Args:
        retriever: The retriever to use for the retrieval
        tool_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
    """
    def wrap(func):
        def wrapped_retrieve(*args, **kwargs):
            docs: List[Document] = func(*args, **kwargs)
            return json.dumps([d.dict() for d in docs], cls=CustomJSONEncoder)

        return wrapped_retrieve

    return Tool(
        name=tool_name,
        description=description,
        func=wrap(retriever.get_relevant_documents),
        coroutine=retriever.aget_relevant_documents,
        args_schema=RetrieverInput,
    )