matrix-gptbot/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

    api_code: str = "wolfram"
    calculation_api: str = "alpha"

    operator: str = "Wolfram ([https://www.wolfram.com](https://www.wolfram.com))"

    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, user: Optional[str] = None) -> Generator[str | bytes, None, None]:
        self.logger.log(f"Querying WolframAlpha for {query}")

        response: wolframalpha.Result = self.client.query(query)

        if not response.success:
            yield "Could not process your query."

            if response.didyoumeans:
                yield "Did you mean: " + response.didyoumeans["didyoumean"][0]["#text"]

            return

        if response.error:
            self.logger.log("Error in query to WolframAlpha: " + response.error, "error")

        for pod in response.pods if not results_only else (response.results if len(list(response.results)) else response.pods):
            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```"