import logging

from langchain.chat_models import ChatOpenAI

from dreamsboard.dreams.builder_cosplay_code.base import StructuredDreamsStoryboard
from dreamsboard.dreams.dreams_personality_chain.base import StoryBoardDreamsGenerationChain
import langchain

from dreamsboard.engine.generate.code_generate import QueryProgramGenerator, EngineProgramGenerator, AIProgramGenerator
from dreamsboard.engine.loading import load_store_from_storage
from dreamsboard.engine.storage.storage_context import StorageContext

langchain.verbose = True
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 控制台打印
handler = logging.StreamHandler()
handler.setLevel(logging.DEBUG)

logger.addHandler(handler)


def test_structured_dreams_storyboard_store_print() -> None:
    try:

        storage_context = StorageContext.from_defaults(persist_dir="./storage")
        code_gen_builder = load_store_from_storage(storage_context)
        index_loaded = True
    except:
        index_loaded = False

    if not index_loaded:
        llm = ChatOpenAI(
            verbose=True
        )

        dreams_generation_chain = StoryBoardDreamsGenerationChain.from_dreams_personality_chain(
            llm=llm, csv_file_path="/media/checkpoint/speech_data/抖音作品/id46Bv3g/str/id46Bv3g_keyframe.csv")

        output = dreams_generation_chain.run()
        logger.info("dreams_guidance_context:" + output.get("dreams_guidance_context"))
        logger.info("dreams_personality_context:" + output.get("dreams_personality_context"))
        dreams_guidance_context = output.get("dreams_guidance_context")
        dreams_personality_context = output.get("dreams_personality_context")

        storyboard_executor = StructuredDreamsStoryboard.form_builder(llm=llm,
                                                                      builder=dreams_generation_chain.builder,
                                                                      dreams_guidance_context=dreams_guidance_context,
                                                                      dreams_personality_context=dreams_personality_context
                                                                      )
        code_gen_builder = storyboard_executor.loader_cosplay_builder()


    executor = code_gen_builder.build_executor()
    logger.info(executor.executor_code)
    assert True

langchain.chat_models.ChatOpenAI
[((545, 572), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (562, 572), False, 'import logging\n'), ((623, 646), 'logging.StreamHandler', 'logging.StreamHandler', ([], {}), '()\n', (644, 646), False, 'import logging\n'), ((806, 859), 'dreamsboard.engine.storage.storage_context.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {'persist_dir': '"""./storage"""'}), "(persist_dir='./storage')\n", (834, 859), False, 'from dreamsboard.engine.storage.storage_context import StorageContext\n'), ((887, 927), 'dreamsboard.engine.loading.load_store_from_storage', 'load_store_from_storage', (['storage_context'], {}), '(storage_context)\n', (910, 927), False, 'from dreamsboard.engine.loading import load_store_from_storage\n'), ((1037, 1061), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'verbose': '(True)'}), '(verbose=True)\n', (1047, 1061), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1119, 1285), 'dreamsboard.dreams.dreams_personality_chain.base.StoryBoardDreamsGenerationChain.from_dreams_personality_chain', 'StoryBoardDreamsGenerationChain.from_dreams_personality_chain', ([], {'llm': 'llm', 'csv_file_path': '"""/media/checkpoint/speech_data/抖音作品/id46Bv3g/str/id46Bv3g_keyframe.csv"""'}), "(llm=llm,\n    csv_file_path=\n    '/media/checkpoint/speech_data/抖音作品/id46Bv3g/str/id46Bv3g_keyframe.csv')\n", (1180, 1285), False, 'from dreamsboard.dreams.dreams_personality_chain.base import StoryBoardDreamsGenerationChain\n'), ((1701, 1909), 'dreamsboard.dreams.builder_cosplay_code.base.StructuredDreamsStoryboard.form_builder', 'StructuredDreamsStoryboard.form_builder', ([], {'llm': 'llm', 'builder': 'dreams_generation_chain.builder', 'dreams_guidance_context': 'dreams_guidance_context', 'dreams_personality_context': 'dreams_personality_context'}), '(llm=llm, builder=\n    dreams_generation_chain.builder, dreams_guidance_context=\n    dreams_guidance_context, dreams_personality_context=\n    dreams_personality_context)\n', (1740, 1909), False, 'from dreamsboard.dreams.builder_cosplay_code.base import StructuredDreamsStoryboard\n')]