import wikipedia
from google.adk.agents import LlmAgent
from langchain_community.retrievers import WikipediaRetriever
from langchain_core.documents import Document

from adk_retrieval import LangchainRetrievalTool


# an example post-processing function that combines multiple documents into a single string
async def post_process(docs: list[Document]) -> str:
    sources = []
    for index, doc in enumerate(docs):
        sources.append(f"Source {index + 1}:\n{doc.page_content}\n")

    return "\n\n".join(sources)


root_agent = LlmAgent(
    name="wiki_sophisticated",
    description="Agent for summarizing Wikipedia articles.",
    instruction="""You are a specialized search & summarizer agent for Wikipedia articles.
      """,
    model="gemini-2.5-flash",
    tools=[
        LangchainRetrievalTool(
            name="wikipedia_retrieval_tool",
            description="Tool for retrieving Wikipedia articles based on user queries.",
            retriever=WikipediaRetriever(wiki_client=wikipedia, top_k_results=3),
            post_process=post_process,
        )
    ],
)
