Augment

This module contains functionality related to the the augment script.

Augment

This script is used to handle chat interactions using the ChainLit library and a chat engine. Actions are observed by Langfuse. To make it work vector storage should be filled with the embeddings of the documents. To run the script execute the following command from the root directory of the project:

python src/chat.py

app_shutdown() async

Clean up resources on application shutdown. Stops the scheduler if it is running.

Source code in src/augment.py 
105
106
107
108
109
110
111
112
113
114
115
@cl.on_app_shutdown
async def app_shutdown() -> None:
    """
    Clean up resources on application shutdown.
    Stops the scheduler if it is running.
    """
    try:
        initializer.get_scheduler().stop()
        logger.info("Scheduler stopped successfully")
    except Exception as e:
        logger.warning(f"Failed to stop scheduler: {e}")

app_startup() async

Initialize the application on startup. Sets up the augmentation initializer and configuration, and starts the scheduler.

Source code in src/augment.py 
29
30
31
32
33
34
35
36
37
38
@cl.on_app_startup
async def app_startup() -> None:
    """
    Initialize the application on startup.
    Sets up the augmentation initializer and configuration, and starts the scheduler.
    """
    global initializer, configuration
    initializer = AugmentationInitializer()
    configuration = initializer.get_configuration()
    initializer.get_scheduler().start()

get_data_layer()

Initialize Chainlit's data layer with the custom service.

Note: This function may be called before app_startup(), so we need to initialize configuration lazily if it's not already set.

Returns:
  • ChainlitService( ChainlitService ) –

    The custom service for data layer.
Source code in src/augment.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@cl.data_layer
def get_data_layer() -> ChainlitService:
    """
    Initialize Chainlit's data layer with the custom service.

    Note: This function may be called before app_startup(), so we need to
    initialize configuration lazily if it's not already set.

    Returns:
        ChainlitService: The custom service for data layer.
    """
    global initializer, configuration
    if configuration is None:
        initializer = AugmentationInitializer()
        configuration = initializer.get_configuration()
    return ChainlitServiceFactory.create(configuration.augmentation)

main(user_message) async

Process user messages and generate responses.

Parameters:
  • user_message (Message) –

    Message received from user
Source code in src/augment.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
@cl.on_message
async def main(user_message: cl.Message) -> None:
    """
    Process user messages and generate responses.

    Args:
        user_message: Message received from user
    """
    try:
        chat_engine = cl.user_session.get("chat_engine")
        assistant_message = cl.Message(content="", author="Assistant")
        response = await cl.make_async(chat_engine.stream_chat)(
            message=user_message.content,
            chainlit_message_id=assistant_message.parent_id,
        )
        for token in response.response_gen:
            await assistant_message.stream_token(token)

        utils = ChainlitUtilsFactory.create(configuration.augmentation.chainlit)
        utils.add_references(assistant_message, response)
        await assistant_message.send()
    except Exception as e:
        # It is imprecise to catch all exceptions, but llamaindex doesn't provide unified RateLimitError
        logger.error(f"Error in main: {e}", exc_info=True)
        await cl.ErrorMessage(
            content="You have reached the request rate limit. Please try again later.",
        ).send()

start() async

Initialize chat session with chat engine. Sets up session-specific chat engine and displays welcome message.

Source code in src/augment.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
@cl.on_chat_start
async def start() -> None:
    """
    Initialize chat session with chat engine.
    Sets up session-specific chat engine and displays welcome message.
    """
    chat_engine = ChatEngineRegistry.get(
        configuration.augmentation.chat_engine.name
    ).create(configuration)
    chat_engine.set_session_id(cl.user_session.get("id"))
    cl.user_session.set("chat_engine", chat_engine)

    utils = ChainlitUtilsFactory.create(configuration.augmentation.chainlit)
    await utils.get_disclaimer_message().send()
    await utils.get_welcome_message().send()