cli_app.py

import os

from dotenv import load_dotenv
from langchain_core.messages import AIMessageChunk, AIMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import InMemorySaver

from src.main.agents.chain_flow import ChainFlow
from src.main.common.logging import enable_logging
from src.main.common.util import get_llm_chat_model
from src.main.common.util import get_message_chunk
from src.main.datasources.database import Database
from src.main.persona import Persona
from src.main.constants import DRIVER_TABLE_FILTER, DRIVER_EXAMPLES


def main():
    driver_persona = Persona(table_filter=DRIVER_TABLE_FILTER, examples=DRIVER_EXAMPLES)
    model_name = os.environ.get('model', default="")
    chat_model = get_llm_chat_model(model_name)
    database = Database()
    memory = InMemorySaver()
    app = ChainFlow(persona=driver_persona, llm=chat_model, datasource=database, checkpointer=memory,
                    output_format="rounded_grid")

    app.draw_state_graph("./chain_flow_state.png")

    while True:
        question = input("\nQuestion: >>")
        config: RunnableConfig = {"configurable": {"thread_id": "123"}}
        stream = app.generate_response(question, config)
        for retries, chunk in stream:
            content = get_message_chunk(chunk)
            message_type = type(chunk[0])
            end = "" if message_type is AIMessageChunk else "\n"
            print(content, end=end, flush=True)


if __name__ == '__main__':
    load_dotenv()
    enable_logging()
    main()