Chat by LangChain

chat_by_langchain.py

"""
pip install langchain-ollama langchain-openai dotenv

Bila menggunakan ChatGPT maka buat file .env berisi:

OPENAI_API_KEY=api-key-milikmu
"""
import os
import sys
from time import time
from datetime import datetime
from langchain_openai import ChatOpenAI
from langchain_ollama import ChatOllama
from langchain_core.messages import (
    SystemMessage,
    HumanMessage,
    AIMessage,
    )
from dotenv import load_dotenv
import ollama


DAY_NAMES = ["Senin", "Selasa", "Rabu", "Kamis", "Jumat", "Sabtu", "Minggu"]


def get_day_name(dt) -> str:
    day_id = dt.weekday()
    return DAY_NAMES[day_id]


def get_token(meta):
    if "token_usage" in meta:  # OpenAI invoke
        usage = meta["token_usage"]
        token_input = usage["prompt_tokens"]
        token_output = usage["completion_tokens"]
    elif "input_tokens" in meta:  # OpenAI stream
        token_input = meta["input_tokens"]
        token_output = meta["output_tokens"]
    else:  # Ollama
        token_input = meta["prompt_eval_count"]
        token_output = meta["eval_count"]
    return dict(input=token_input, output=token_output)


load_dotenv()
is_openai = os.getenv("OPENAI_API_KEY")
llm_options = dict(temperature=0)
if not is_openai:  # Ollama by default
    llm_options["base_url"] = os.getenv("URL", "http://localhost:11434")
    llm_options["model"] = os.getenv("LLM_MODEL", "gemma4:e4b")
    client = ollama.Client(host=llm_options["base_url"])
    info = client.show(llm_options["model"])
    for key, val in info.modelinfo.items():
        if key.find("context_length") > -1:
            # Maksimalkan kapasitas context LLM
            llm_options["num_ctx"] = val
            break
    cls = ChatOllama
else:
    llm_options["model"] = os.getenv("LLM_MODEL", "gpt-5.1")
    llm_options["stream_usage"] = True
    cls = ChatOpenAI
print(f"LLM options: {llm_options}")
llm = cls(**llm_options)

prompt_from_arg = prompt = sys.argv[1:] and sys.argv[1]
messages = []
if not prompt:
    print("Silakan bertanya, contoh: 'Siapa pembuatmu?'")
while True:
    if not prompt:
        prompt = input(">>> ")
        if prompt == "exit":
            break
    dt = datetime.now()
    timestamp = dt.strftime("%Y-%m-%d %H:%M:%S")
    day_name = get_day_name(dt)
    system_prompt = f"Waktu saat ini: {day_name}, {timestamp}"
    messages.append(SystemMessage(system_prompt))
    messages.append(HumanMessage(prompt))
    begin_time = time()
    if prompt_from_arg:
        response = llm.invoke(messages)
        answer = response.content
        print(answer)
        meta = response.response_metadata
    else:
        answer = ""
        for chunk in llm.stream(messages):
            print(chunk.content, end="", flush=True)
            answer += chunk.content
            if chunk.usage_metadata:
                meta = chunk.usage_metadata
        print()
    duration = time() - begin_time
    duration = int(duration)
    messages.append(AIMessage(answer))
    token = get_token(meta)
    print(
        f"Input {token['input']} token, Output {token['output']} token, "
        f"{duration} detik")
    if prompt_from_arg:
        break
    prompt = ""