Skip to content

Stream Example

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
from dotenv import load_dotenv

from llmfy import (
    LLMfy,
    Message,
    Role,
    BedrockConfig,
    BedrockModel,
    GenerationResponse,
    llmfy_usage_tracker,
)

load_dotenv()


def stream_example():
    info = """
    Irufano adalah seorang sofware engineer.
    Dia berasal dari Indonesia.
    """

    # llm
    # model="anthropic.claude-3-haiku-20240307-v1:0",
    # model="us.anthropic.claude-3-5-haiku-20241022-v1:0",
    # model="amazon.nova-lite-v1:0",

    llm = BedrockModel(
        model="amazon.nova-lite-v1:0",
        config=BedrockConfig(temperature=0.7),
    )

    SYSTEM_PROMPT = """Answer any user questions based solely on the data below:
    <data>
    {{info}}
    </data>

    DO NOT response outside context."""

    # Initialize framework
    chat = LLMfy(llm, system_message=SYSTEM_PROMPT, input_variables=["info"])

    try:
        # Example conversation with tool use
        messages = [Message(role=Role.USER, content="apa ibukota jakarta?")]
        # with openai_usage_tracker() as usage:
        with llmfy_usage_tracker() as usage:
            stream = chat.chat_stream(messages, info=info)
            full_content = ""
            for chunk in stream:
                if isinstance(chunk, GenerationResponse):
                    if chunk.result.content:
                        content = chunk.result.content
                        full_content += content
                        print(content, end="", flush=True)

            print(usage)

    except Exception as e:
        raise e


if __name__ == "__main__":
    stream_example()