51 lines
No EOL
1.9 KiB
Python
51 lines
No EOL
1.9 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
|
|
|
|
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```" |