35 lines
1.3 KiB
Python
35 lines
1.3 KiB
Python
|
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```"
|