server.py

from __future__ import annotations

from typing import Any
from typing import cast

from fastmcp import FastMCP

from schemas.output.rag import AddCrawlResult
from schemas.output.rag import AddDirectoryResult
from schemas.output.rag import AddDocumentResult
from schemas.output.rag import AddURLResult
from schemas.output.rag import DocumentListResult
from schemas.output.rag import RAGAnswer
from schemas.output.rag import RebuildIndexResult
from schemas.output.rag import RemoveDocumentResult
from schemas.output.rag import SearchResult
from schemas.output.rag import StatusInfo
from tools.news.api import news_search
from tools.rag.documents import add_directory
from tools.rag.documents import add_document
from tools.rag.documents import list_documents
from tools.rag.documents import rag_status
from tools.rag.documents import rebuild_index
from tools.rag.documents import remove_document
from tools.rag.search import doc_search
from tools.rag.search import rag_answer
from tools.rag.web_scraping import add_crawl
from tools.rag.web_scraping import add_url
from tools.stocks.api import stock_quote
from tools.stocks.api import stock_quotes
from tools.time.api import local_time
from tools.time.api import local_time_by_place
from tools.weather.api import weather_daily
from tools.weather.api import weather_now
from tools.web.api import crawl_site
from tools.web.api import fetch_url
from tools.web.api import fetch_url_readable
from tools.web.api import web_search

app: FastMCP = FastMCP()


@app.tool()
def web_search_tool(
    query: str,
    max_results: int,
) -> dict[str, object]:
    """Search the web and return the top N results.

    Always returns unified ``sources`` list so downstream agents can
    extract citations uniformly.
    """
    try:
        result = web_search(
            query=query,
            max_results=max_results,
        )
        return {
            'query': result.get('query', query),
            'results': result.get('results', []),
            'sources': result.get('sources', []),
            'error': result.get('error'),
        }
    except Exception as e:
        return {
            'query': query,
            'results': [],
            'sources': [],
            'error': str(e),
        }


@app.tool()
def news_search_tool(
    query: str,
    max_results: int,
) -> dict[str, object]:
    """Search for news articles using DuckDuckGo News."""
    try:
        result = news_search(
            query=query,
            max_results=max_results,
        )
        return {
            'query': result.get('query', query),
            'results': result.get('results', []),
            'sources': result.get('sources', []),
            'error': result.get('error'),
        }
    except Exception as e:
        return {
            'query': query,
            'results': [],
            'sources': [],
            'error': str(e),
        }


@app.tool()
def fetch_url_tool(url: str, max_length: int = 10_000) -> dict[str, object]:
    """Fetch a URL and return a lightweight representation."""
    try:
        page = fetch_url(url, max_length=max_length)
        return {
            **page,
            'sources': [
                {
                    'title': page.get('title') or url,
                    'url': url,
                    'type': 'web',
                },
            ],
        }
    except Exception as e:
        return {
            'url': url,
            'title': None,
            'text': None,
            'links': [],
            'sources': [],
            'error': str(e),
        }


@app.tool()
def fetch_url_readable_tool(
    url: str,
    max_length: int = 20_000,
) -> dict[str, object]:
    """Fetch a URL and extract readability-style main content."""
    try:
        page = fetch_url_readable(url, max_length=max_length)
        canonical = page.get('canonical') or url
        return {
            **page,
            'sources': [
                {
                    'title': page.get('title') or canonical,
                    'url': canonical,
                    'type': 'web',
                },
            ],
        }
    except Exception as e:
        return {
            'url': url,
            'title': None,
            'canonical': url,
            'meta_title': None,
            'content': None,
            'links': [],
            'sources': [],
            'error': str(e),
        }


@app.tool()
def crawl_site_tool(
    url: str,
    max_pages: int = 5,
    same_domain_only: bool = True,
    readable: bool = True,
) -> dict[str, object]:
    """Perform a small crawl rooted at the given URL."""
    try:
        if readable:
            crawl_result = cast(
                dict[str, Any],
                crawl_site(
                    url,
                    max_pages=max_pages,
                    same_domain_only=same_domain_only,
                    readable=True,
                ),
            )
        else:
            crawl_result = cast(
                dict[str, Any],
                crawl_site(
                    url,
                    max_pages=max_pages,
                    same_domain_only=same_domain_only,
                    readable=False,
                ),
            )
    except Exception as e:
        return {
            'seed': url,
            'count': 0,
            'pages': [],
            'sources': [],
            'error': str(e),
        }

    pages = cast(list[dict[str, Any]], crawl_result.get('pages', []) or [])

    sources: list[dict[str, str]] = []
    for p in pages:
        src_url = str(p.get('url') or url)
        title = str(p.get('title') or src_url)
        sources.append({'title': title, 'url': src_url, 'type': 'web'})

    return {
        'seed': crawl_result.get('seed', url),
        'count': crawl_result.get('count', len(pages)),
        'pages': pages,
        'sources': sources,
    }


# --- Documents / RAG ---
@app.tool()
def docs_add_document(
    path: str,
    user_id: str,
    chat_id: str,
) -> AddDocumentResult:
    return add_document(path, user_id=user_id, chat_id=chat_id)


@app.tool()
def docs_add_directory(
    path: str,
    user_id: str,
    chat_id: str,
    recursive: bool = True,
) -> AddDirectoryResult:
    return add_directory(
        path,
        recursive=recursive,
        user_id=user_id,
        chat_id=chat_id,
    )


@app.tool()
def docs_list(user_id: str, chat_id: str) -> DocumentListResult:
    return list_documents(user_id=user_id, chat_id=chat_id)


@app.tool()
def docs_remove(
    doc_id: str,
    user_id: str,
    chat_id: str,
) -> RemoveDocumentResult:
    return remove_document(doc_id, user_id=user_id, chat_id=chat_id)


@app.tool()
def docs_search(
    query: str,
    user_id: str,
    chat_id: str,
    file_names: list[str],
    k: int,
) -> SearchResult:
    return doc_search(
        query,
        k=k,
        file_names=file_names,
        user_id=user_id,
        chat_id=chat_id,
    )


@app.tool()
def rag_answer_tool(
    question: str,
    user_id: str,
    chat_id: str,
    file_names: list[str],
    k: int,
) -> RAGAnswer:
    return rag_answer(
        question,
        k=k,
        file_names=file_names,
        user_id=user_id,
        chat_id=chat_id,
    )


@app.tool()
def rag_status_tool(user_id: str, chat_id: str) -> StatusInfo:
    return rag_status(user_id=user_id, chat_id=chat_id)


@app.tool()
def rag_rebuild(user_id: str, chat_id: str) -> RebuildIndexResult:
    return rebuild_index(user_id=user_id, chat_id=chat_id)


@app.tool()
def rag_add_url(url: str, user_id: str, chat_id: str) -> AddURLResult:
    return add_url(url, user_id=user_id, chat_id=chat_id)


@app.tool()
def rag_add_crawl(
    url: str,
    user_id: str,
    chat_id: str,
    max_pages: int = 5,
    same_domain_only: bool = True,
) -> AddCrawlResult:
    return add_crawl(
        url,
        max_pages=max_pages,
        same_domain_only=same_domain_only,
        user_id=user_id,
        chat_id=chat_id,
    )


@app.tool()
def time_local(tz: str) -> dict[str, object]:
    try:
        return cast(dict[str, object], local_time(tz))
    except Exception as e:
        return {'error': str(e)}


@app.tool()
def time_local_by_place_tool(place: str) -> dict[str, object]:
    try:
        return cast(dict[str, object], local_time_by_place(place))
    except Exception as e:
        return {'error': str(e), 'place': place}


@app.tool()
def weather_now_tool(place: str) -> dict[str, object]:
    try:
        return cast(dict[str, object], weather_now(place))
    except Exception as e:
        return {'error': str(e), 'place': place}


@app.tool()
def weather_daily_tool(place: str, days: int = 7) -> dict[str, object]:
    try:
        return cast(dict[str, object], weather_daily(place, days))
    except Exception as e:
        return {'error': str(e), 'place': place, 'days': days}


@app.tool()
def stock_quote_tool(symbol: str) -> dict[str, object]:
    try:
        return cast(dict[str, object], stock_quote(symbol))
    except Exception as e:
        return {'error': str(e), 'symbol': symbol}


@app.tool()
def stock_quotes_tool(symbols: list[str]) -> dict[str, object]:
    try:
        return cast(dict[str, object], stock_quotes(symbols))
    except Exception as e:
        return {'error': str(e), 'symbols': symbols}


def main() -> None:
    app.run(
        transport='http',
        host='127.0.0.1',
        port=8090,
        path='/mcp',
        stateless_http=True,
    )


if __name__ == '__main__':
    main()