from v1 import client

import asyncio

import os
import json


class TestDifyClient:
    async def test_chat_messages(self):
        cln = client.AsyncDifyServiceClient(api_key=os.getenv("DIFY_API_KEY"), base_url=os.getenv("DIFY_BASE_URL"))

        async for chunk in cln.chat_messages(inputs={}, query="调用工具查看现在几点？", user="test"):
            print(json.dumps(chunk, ensure_ascii=False, indent=4))

    async def test_upload_file(self):
        cln = client.AsyncDifyServiceClient(api_key=os.getenv("DIFY_API_KEY"), base_url=os.getenv("DIFY_BASE_URL"))

        file_bytes = open("img.png", "rb").read()

        print(type(file_bytes))

        file = ("img2.png", file_bytes, "image/png")

        resp = await cln.upload_file(file=file, user="test")
        print(json.dumps(resp, ensure_ascii=False, indent=4))

    async def test_workflow_run(self):
        cln = client.AsyncDifyServiceClient(api_key=os.getenv("DIFY_API_KEY"), base_url=os.getenv("DIFY_BASE_URL"))

        # resp = await cln.workflow_run(inputs={}, user="test")
        # # print(json.dumps(resp, ensure_ascii=False, indent=4))
        # print(resp)
        chunks = []

        ignored_events = ['text_chunk']
        async for chunk in cln.workflow_run(inputs={}, user="test"):
            if chunk['event'] in ignored_events:
                continue
            chunks.append(chunk)
        print(json.dumps(chunks, ensure_ascii=False, indent=4))

if __name__ == "__main__":
    asyncio.run(TestDifyClient().test_chat_messages())