from typing import Any, Awaitable, Callable

from google.adk.tools.retrieval import BaseRetrievalTool
from google.adk.tools.tool_context import ToolContext
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from typing_extensions import override


class LangchainRetrievalTool(BaseRetrievalTool):
    """A retrieval tool that uses Langchain for document retrieval."""

    def __init__(
        self,
        name: str,
        description: str,
        retriever: BaseRetriever,
        post_process: Callable[[list[Document]], Awaitable[str]] | None = None,
    ):
        super().__init__(name=name, description=description)
        self._retriever = retriever
        self._post_process = post_process

    @override
    async def run_async(self, *, args: dict[str, Any], tool_context: ToolContext) -> Any:
        query = args["query"]
        documents = await self._retriever.ainvoke(query)
        if self._post_process:
            return await self._post_process(documents)
        return documents[0].page_content
