matrix-reportbot/classes/wolframalpha.py

import wolframalpha
import requests

from .logging import Logger

from typing import Dict, List, Tuple, Generator, Optional

class WolframAlpha:
    api_key: str
    logger: Logger
    client: wolframalpha.Client

    def __init__(self, api_key: str, logger: Optional[Logger] = None):
        self.api_key: str = api_key
        self.logger: Logger = logger or Logger()
        self.client = wolframalpha.Client(self.api_key)

    def generate_calculation_response(self, query: str, text: Optional[bool] = False, results_only: Optional[bool] = False) -> Generator[str | bytes, None, None]:
        self.logger.log(f"Querying WolframAlpha for {query}")

        response: wolframalpha.Result = self.client.query(query)
        
        for pod in response.pods if not results_only else response.results:
            self.logger.log(pod.keys(), "debug")
            if pod.title:
                yield "*" + pod.title + "*"
            for subpod in pod.subpods:
                self.logger.log(subpod.keys(), "debug")
                if subpod.title:
                    yield "*" + subpod.title + "*"
                if subpod.img and not text:
                    image = requests.get(subpod.img.src).content
                    yield image
                elif subpod.plaintext:
                    yield "```\n" + subpod.plaintext + "\n```"