import gradio as gr
from string import Template

from rag import search


async def rag_search(query: str) -> str:
    """
    Search for information based on a query.

    Args:
        query (str): The search query

    Returns:
        str: A message indicating the search result
    """
    results = await search(query, top_k=5)

    source_template = Template(
        '''
        ### Source
        $source
        
        ### Content
        $content
        '''
    )

    return f"## Results for query {query}\n" + "\n".join(
        source_template.substitute(
            source=result['metadata']['source'],
            content=result['content']
        ) for result in results
    ) if results else "No results found."


# Create a Gradio interface for the RAG search functionality
with gr.Blocks() as demo:
    gr.Interface(rag_search, gr.Textbox(), gr.Textbox(), api_name="search_content", title="RAG Search")

    with open("README.md") as f:
        readme_content = f.read()

        readme_content = readme_content[readme_content.index("#"):]

        readme = gr.Markdown(readme_content)

if __name__ == "__main__":
    demo.launch(mcp_server=True)