dnsagent.py

import asyncio
import json

from agno.agent.agent import Agent
from aiodns import DNSResolver


async def query_address_records(hostname: str) -> str:
    """
    Use this function to find the A and AAAA records for a given hostname.

    Args:
        domain (str): The domain to fetch the A and AAAA records for.

    Returns:
        A list of A (IPv4) and AAAA (IPv6) records. It is a list of strings.
        If an error occurs then a dict with an error field is returned.
    """

    try:
        addresses = []

        resolver = DNSResolver()
        a_task = resolver.query(hostname, "A")
        aaaa_task = resolver.query(hostname, "AAAA")

        a_result, aaaa_result = await asyncio.gather(a_task, aaaa_task, return_exceptions=True)

        if not isinstance(a_result, Exception):
            addresses += [record.host for record in a_result]
        if not isinstance(aaaa_result, Exception):
            addresses += [record.host for record in aaaa_result]

        return json.dumps(addresses)
    except Exception as e:
        error = {"error": f"Error looking up MX records for {domain}: {e}"}
        return json.dumps(error)


async def query_mx_records(domain: str) -> str:
    """
    Use this function to find the MX records for a given domain name.

    Args:
        domain (str): The domain to fetch the MX records for.

    Returns:
        A list of MX records. Each record will contain the hostname and preference.
        If an error occurs then a dict with an error field is returned.
    """

    try:
        resolver = DNSResolver()
        response = await resolver.query(domain, "MX")
        records = [{"host": mx.host, "priority": mx.priority} for mx in response]
        return json.dumps(records)
    except Exception as e:
        error = {"error": f"Error looking up MX records for {domain}: {e}"}
        return json.dumps(error)


if __name__ == "__main__":
    agent = Agent(
        tools=[query_mx_records, query_address_records],
        show_tool_calls=True,
    )

    # asyncio.run(agent.aprint_response("What is the email provider for wavelo.com, sateh.com and thestar.com", stream=True))
    asyncio.run(agent.aprint_response("Does mozilla.com have any IPv6 MX hosts?", stream=True))