Add advanced formatting

modules/ask/__init__.py

@@ -1,12 +1,18 @@
 from decimal import Decimal
+import re
+from PIL import Image as PILImage
+import io
+import asyncio
 
 from langchain.callbacks import get_openai_callback
 
-from core.builtins import Bot
+from core.builtins import Bot, Plain, Image
 from core.component import module
 from core.dirty_check import check_bool
 from core.exceptions import NoReportException
-from modules.ask.agent import agent_executor
+
+from .agent import agent_executor
+from .formatting import generate_latex, generate_code_snippet
 
 ONE_K = Decimal('1000')
 # https://openai.com/pricing
@@ -44,6 +50,41 @@ async def _(msg: Bot.MessageSession):
     price = tokens / ONE_K * PRICE_PER_1K_TOKEN
     petal = price * USD_TO_CNY * CNY_TO_PETAL
     msg.data.modify_petal(-int(petal))
+
+    blocks = parse_markdown(res)
+
+    chain = []
+    for block in blocks:
+        if block['type'] == 'text':
+            chain.append(Plain(block['content']))
+        elif block['type'] == 'latex':
+            chain.append(Image(PILImage.open(io.BytesIO(await generate_latex(block['content'])))))
+        elif block['type'] == 'code':
+            chain.append(Image(PILImage.open(io.BytesIO(await generate_code_snippet(block['content']['code'], block['content']['language'])))))
+
     if await check_bool(res):
         raise NoReportException('https://wdf.ink/6OUp')
-    await msg.finish(res)
+    await msg.finish(chain)
+
+def parse_markdown(md: str):
+    regex = r'(```[\s\S]*?\n```|\$\$[\s\S]*?\$\$|[^\n]+)(?:\n|$)'
+
+    blocks = []
+    for match in re.finditer(regex, md):
+        content = match.group(1)
+        print(content)
+        if content.startswith('```'):
+            block = 'code'
+            try:
+                language, code = re.match(r'```(.*)\n([\s\S]*?)\n```', content).groups()
+            except AttributeError:
+                raise ValueError('Code block is missing language or code')
+            content = {'language': language, 'code': code}
+        elif content.startswith('$$'):
+            block = 'latex'
+            content = content[2:-2].strip()
+        else:
+            block = 'text'
+        blocks.append({'type': block, 'content': content})
+
+    return blocks

modules/ask/formatting.py

@@ -0,0 +1,24 @@
+import ujson as json
+import aiohttp
+
+async def generate_latex(formula: str):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url='https://wikimedia.org/api/rest_v1/media/math/check/inline-tex', data=json.dumps({
+            'q': formula
+        }), headers={'content-type': 'application/json'}) as req:
+            headers = req.headers
+            location = headers.get('x-resource-location')
+
+        async with session.get(url=f'https://wikimedia.org/api/rest_v1/media/math/render/png/{location}') as img:
+            return await img.read()
+
+async def generate_code_snippet(code: str, language: str):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url='https://sourcecodeshots.com/api/image', data=json.dumps({
+            'code': code,
+            'settings': {
+                'language': language,
+                'theme': 'night-owl',
+            }
+        }), headers={'content-type': 'application/json'}) as req:
+            return await req.read()

modules/ask/prompt.py

@@ -26,7 +26,7 @@ The User will ask a `Question`. Answer the `Question` as best you can. You can u
 
 Provide informative, logical, positive, interesting, intelligent, and engaging answers with details to cover multiple aspects of the question. You can generate articles and other forms of content, but do not rely on tools when doing so. Use emojis to make your answers more interesting.
 
-Use Markdown code block syntax when outputting code. Use LaTeX to output mathematical expressions, and surround the expression with dollar signs `$`, e.g. to output the mass-energy equivalence, always use $E=mc^2$. You can output multiple lines of strings.
+Use Markdown code block syntax when outputting code. Use LaTeX to output mathematical expressions, and surround the expression with dollar signs `$$`, e.g. to output the mass-energy equivalence, always use $$E=mc^2$$. You can output multiple lines of strings.
 
 Use the original question's language. For example, if I ask "什么是质能方程？", you should output your `Thought` and `Action` in Chinese like this:
 
@@ -38,7 +38,7 @@ Thought: Wolfram Alpha 不能回答这个问题。我应该使用 Search 工具
 Action: Search[质能方程是什么？]
 Observation: E = mc²，即质能等价（mass-energy equivalence）、质能守恒、质能互换，亦称为质能转换公式、质能方程，是一种阐述能量（E）与质量（m）间相互关系的理论物理学公式，公式中的 c 是物理学中代表光速的常数。
 Thought: 我现在知道问题的答案了。
-Action: Answer[质能方程是 $E=mc^2$，其中 E 代表能量，m 代表质量，c 代表光速。这意味着任何物体的质量和能量之间都有一个固定的数量关系。]
+Action: Answer[质能方程是 $$E=mc^2$$，其中 E 代表能量，m 代表质量，c 代表光速。这意味着任何物体的质量和能量之间都有一个固定的数量关系。]
 """
 
 `Action`s aren't required to be always taken.