from langchain.tools import tool
from vector_store import search_documents, add_documents


@tool
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """Semantic search in the knowledge base. Use this tool to find relevant information.

    Args:
        query: The search query string.
        max_results: Maximum number of results to return (default 5).

    Returns:
        Formatted string with search results and relevance scores.
    """
    results = search_documents(query, max_results=max_results)
    if not results:
        return "No results found in the knowledge base."
    output_lines = [f"Found {len(results)} result(s):\n"]
    for i, r in enumerate(results, 1):
        title = r["metadata"].get("title", "Unknown")
        content = r["content"]
        output_lines.append(f"[{i}] Title: {title}")
        output_lines.append(f"    {content}\n")
    return "\n".join(output_lines)


@tool
def add_to_knowledge_base(content: str, title: str) -> str:
    """Add a document to the knowledge base. Use this tool to store new information.

    Args:
        content: The full text content of the document to add.
        title: A descriptive title for the document.

    Returns:
        Confirmation message with the number of chunks stored.
    """
    num_chunks = add_documents(content=content, title=title)
    return f"Successfully added document '{title}' to the knowledge base ({num_chunks} chunk(s) stored)."