[LLM] LangChain과 Chainlit을 활용한 챗봇 구축 예제 Code

langchain을 통해서 구축한 RAG 기반의 추천 시스템을 Chainlit을 활용해서 챗봇 UI까지 연결해 보는 과정까지를 작성해보고자 한다.

 

데이터 전처리, vector DB에 chunk 단위로 저장하고, query input과 유사도를 측정해서 ouput을 뽑아낼 수 있는 RAG를 적용한 시스템을 구축하는 프로젝트를 진행했다. (agent까진 가지 못했고, 조금은 휴리스틱 하지만ㅜㅜ) 

 

회사 프로젝트로 함께 진행한거라서 자세한 코드나 내용에 대해서는 언급하지 못하기도 하고, 초기 기획에는 멀티턴이 가능한 챗봇으로 개발을 진행하려고 했지만, 프로젝트의 목적을 생각했을 때 좋지 못한 UI가 나올 수도 있어서 해당 부분은 추후에 고도화하는 방향으로 변경을 했다..

 

따라서, 멀티턴 챗봇까지는 아니지만 내부적으로 쿼리문을 계속해서 써보면서 디버깅도 하고 공부도 할 겸 chainlit을 활용해서 한번 간단하게 챗봇 형식을 구현해 보았다.

 

그렇기 때문에, 기본적으로 간단한 부분들만 설명을 진행해보고자 한다.

(해당 부분 이전에 RAG 부분에 대해서는 다른 포스팅을 해보려고 한당..)

 

1. chainlit 기본

본격적으로 설명을 진행하기 이전에 chainlit이란, python 기반으로 채팅 구조의 애플리케이션 개발을 지원하는 도구이다.

설치는 간단하게 pip install chainlit으로 할 수 있다.

chainlit을 활용하면 비교적으로 예쁜 UI도 간단하게 만들 수 있다고 한다.

 

간단하게 사용할 수 있는 예시는 다음과 같다.

 

@cl. on_message : 사용자가 메시지를 보낼 때마다 호출될 메시지 핸들러 함수를 정의 (사용자가 메시지를 보낼 때마다 실행)

async def main(message: cl.Message) : 사용자가 입력한 메시지를 받아 사용자 정의 로직을 수행하고 응답을 보냄

await cl.Message(content="응답 내용"). send() : 사용자에게 응답 메시지를 보내는 역할 (UI에서 사용자 또는 봇이 보내는 단일 메시지)

 

import chainlit as cl

@cl.on_message
async def main(message: cl.Message):
    # 사용자가 보낸 메시지 내용을 받아옴
    user_input = message.content
    
    # 간단한 응답을 구성하고 사용자에게 전송
    response_text = f"당신이 말한 내용: '{user_input}'에 대한 응답입니다."
    
    await cl.Message(
        content=response_text
    ).send()

 

우선 테스트를 위해서 이런 식으로 코드를 작성한 후에 터미널에서 "chainlit run app.py -w"라고 입력하면 브라우저가 켜지면서 chainlit 화면이 localhost:8000으로 켜지는 것을 확인할 수 있다.

 

 

이제 chainlit을 조금 더 커스터마이징 해보자.

 

2. chainlit을 활용한 개발 - 필요한 리소스 로드

from RetrieverRunnable_Chat import *

import pickle, os
from dotenv import load_dotenv
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from huggingface_hub import login
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_ollama import OllamaLLM
import chainlit as cl

# ============================================
# 전역 변수: 서버 시작 시 한 번만 초기화
# ============================================
RETRIEVER_RUNNABLE = None
RESPONSE_LLM = None
INITIALIZATION_ERROR = None

def initialize_global_resources():
    """서버 시작 시 한 번만 실행되는 초기화 함수"""
    global RETRIEVER_RUNNABLE, RESPONSE_LLM, INITIALIZATION_ERROR
    
    try:
        print("=" * 60)
        print("서버 리소스 초기화 시작")
        print("=" * 60)
        
        # 1. 환경 변수 로드
        print("✅ 환경 변수 로딩")
        load_dotenv("config/.env")
        huggingface_Key = os.environ["huggingface_Key"]
        login(huggingface_Key)
        
        # 2. JSON 쿼리 LLM 초기화
        print("✅ JSON Query LLM 초기화")
        json_query_llm = OllamaLLM(model="llama3.1:8b", temperature=0)
        
        # 3. Embedding 모델 로드
        print("✅ Embedding 모델 로딩")
        embedding = HuggingFaceEmbeddings(
            model_name="snunlp/KR-SBERT-V40K-klueNLI-augSTS",
            model_kwargs={"device": "cuda"}
        )
        
        # 4. Chroma DB 로드
        print("✅ Vector DB (Chroma) 로딩")
        post_chroma_path = "vector_db/"
        post_chroma_index = Chroma(
            persist_directory=post_chroma_path,
            embedding_function=embedding,
            collection_name="post_chroma_index"
        )
        
        # 5. BM25 Retriever 로드
        print("✅ BM25 Retriever 로딩")
        with open("bm25_retriever.pkl", "rb") as f:
            post_bm25_retriever = pickle.load(f)
        
        # 6. Retriever Runnable 생성
        print("✅ Retriever Runnable 생성")
        RETRIEVER_RUNNABLE = RetrieverRunnable(
            post_bm25_retriever, 
            post_chroma_index, 
            json_query_llm
        )
        
        # 7. 응답 생성 LLM 초기화
        print("✅ Response LLM 초기화")
        RESPONSE_LLM = OllamaLLM(
            model="gemma3:12b", 
            temperature=0,
            top_p=0.9, 
            num_predict=2040, 
            repeat_penalty=1.1
        )
        
        print("=" * 60)
        print("✅ 서버 리소스 초기화 완료!")
        print("=" * 60)
        
    except Exception as e:
        INITIALIZATION_ERROR = str(e)
        print("=" * 60)
        print(f"❌ 서버 리소스 초기화 실패: {e}")
        print("=" * 60)
        import traceback
        traceback.print_exc()

 

우선 내가 진행한 프로젝트의 경우는 기본적으로 vectorDB에 저장해 둔 정보를 기반으로 쿼리문이 들어오면 그 조건에 해당하는 답변을 출력하는 일종의 추천 시스템이다. 

맨 처음에 로딩한 모듈 중 RetrieverRunnable_Chat이 해당 추천 과정을 진행할 수 있도록 만든 RAG 모듈이다. 

 

그렇다면 chainlit을 본격적으로 세팅하기 이전에, 애플리케이션 전체에서 이러한 공유할 무거운 리소스들을 미리미리 로드를 해주는 과정이 필요하다. 

 

1. 환경 변수 로딩

embedding model 로딩을 위해서 허깅페이스에 로그인을 행하기 때문에, 나의 key값이 들어있는. env 파일을 로딩한 후에 이를 이용해서 임베딩 모델을 불러온다. 

2. json 쿼리 llm 초기화

또, 나의 경우에는 주석에 보면 2번으로 json 쿼리 llm 초기화라는 부분이 있는데, 해당 부분의 역할은 user의 input인 프롬프트로 받은 쿼리문 자체를 한번 더 rewriting 해서 원하는 결과를 json으로 파싱 하는 역할을 하는 별도의 llm이다. 

단순히 RAG의 성능을 올리기 위해서 쿼리문을 rewriting 하기 위함이라기보다는 vectorDB에서 search를 진행할 때, 메타데이터를 잘 활용할 수 있게 하기 위해서 쿼리문을 분석해서 그에 맞는 메타데이터를 추출해 낼 수 있도록 하려고 진행했다. 

해당 부분은 실제 디코딩하는 llm과는 다른 모델을 사용했다.

(쿼리문을 sql-like 형태로 파싱 하거나 json으로 출력하는 데 있어서는 llama 모델이 더 성능이 좋다는 GPT의 정보를 활용...)

3. embedding model 로드

내가 사용했던 embedding model은 "snunlp/KR-SBERT-V40K-klueNLI-augSTS" 모델이다. 

나의 경우 정제한 데이터 셋 자체가 SNS 데이터들이었기 때문에, 다양한 오픈소스 한국어 임베딩 모델들을 활용해 봤지만 유사도 수치도 그렇고 가장 좋았던 모델이다. (qwen을 포함해서 진행해 봤지만, 이게 젤 잘 되는 것 같았다..ㅜ)

4. chroma DB 로드

그리고 임베딩 모델을 통해서 청킹한 문서들을 document로 만들고, 이를 chroma DB를 활용해서 DB를 만들었다.

로컬에 저장해 둔 DB를 불러오는 작업을 진행해 준다.

5. BM25 retriever 로드

많은 경우에도 retriever를 단순히 하나만 사용하지 않고 하이브리드로 많이 활용하고 있다. 

나의 경우에도 마찬가지로 경우에 따라서 bm25_retriever를 추가적으로 활용해 주기 위해서 미리 저장해 둔 retriever를 불러온다.

6. retriever runnable 생성

chromaDB 즉, dense retriever와 bm25_retriever 두 가지를 활용해서 직접 커스터마이징 한 RAG 함수를 설정해 준다.

7. 디코딩 llm 초기화

마지막으로 디코딩을 진행할 모델을 로딩해 준다. 

나의 경우에는 기본적으로 플로우 자체를 langchain을 활용해서 개발을 진행했기 때문에, 

랭체인 올라마를 활용해서 모델을 로드해 줬다.

사실상 MVP 정도까지는 올라마를 통해서 서비스를 해도 좋다곤 하지만, 실제로 서비스를 상용화해주기 위해서는 vLLM이나 다른 툴들을 쓰는 게 더욱 효율적이라고 알고 있다.

(나의 경우도 완전한 서비스의 개념이기보단 현재까지는 MVP 수준이기 때문에, 서빙 전까진 이렇게 진행을 해줬다.)

 

3. chainlit을 활용한 개발 - chainlit 세션 설정

from RetrieverRunnable_Chat import *

import pickle, os
from dotenv import load_dotenv
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from huggingface_hub import login
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_ollama import OllamaLLM
import chainlit as cl

# ============================================
# 전역 변수: 서버 시작 시 한 번만 초기화
# ============================================
RETRIEVER_RUNNABLE = None
RESPONSE_LLM = None
INITIALIZATION_ERROR = None

def initialize_global_resources():
    """서버 시작 시 한 번만 실행되는 초기화 함수"""
    global RETRIEVER_RUNNABLE, RESPONSE_LLM, INITIALIZATION_ERROR
    
    try:
        print("=" * 60)
        print("서버 리소스 초기화 시작")
        print("=" * 60)
        
        # 1. 환경 변수 로드
        print("✅ 환경 변수 로딩")
        load_dotenv("config/.env")
        huggingface_Key = os.environ["huggingface_Key"]
        login(huggingface_Key)
        
        # 2. JSON 쿼리 LLM 초기화
        print("✅ JSON Query LLM 초기화")
        json_query_llm = OllamaLLM(model="llama3.1:8b", temperature=0)
        
        # 3. Embedding 모델 로드
        print("✅ Embedding 모델 로딩")
        embedding = HuggingFaceEmbeddings(
            model_name="snunlp/KR-SBERT-V40K-klueNLI-augSTS",
            model_kwargs={"device": "cuda"}
        )
        
        # 4. Chroma DB 로드
        print("✅ Vector DB (Chroma) 로딩")
        post_chroma_path = "vector_db/"
        post_chroma_index = Chroma(
            persist_directory=post_chroma_path,
            embedding_function=embedding,
            collection_name="post_chroma_index"
        )
        
        # 5. BM25 Retriever 로드
        print("✅ BM25 Retriever 로딩")
        with open("bm25_retriever.pkl", "rb") as f:
            post_bm25_retriever = pickle.load(f)
        
        # 6. Retriever Runnable 생성
        print("✅ Retriever Runnable 생성")
        RETRIEVER_RUNNABLE = RetrieverRunnable(
            post_bm25_retriever, 
            post_chroma_index, 
            json_query_llm
        )
        
        # 7. 응답 생성 LLM 초기화
        print("✅ Response LLM 초기화")
        RESPONSE_LLM = OllamaLLM(
            model="gemma3:12b", 
            temperature=0,
            top_p=0.9, 
            num_predict=2040, 
            repeat_penalty=1.1
        )
        
        print("=" * 60)
        print("✅ 서버 리소스 초기화 완료!")
        print("=" * 60)
        
    except Exception as e:
        INITIALIZATION_ERROR = str(e)
        print("=" * 60)
        print(f"❌ 서버 리소스 초기화 실패: {e}")
        print("=" * 60)
        import traceback
        traceback.print_exc()

# ============================================
# 서버 시작 시 자동 실행
# ============================================
initialize_global_resources()

@cl.on_chat_start
async def on_chat_start():
    """각 사용자 세션 시작 시 실행"""
    
    # 1. 초기 메시지
    initial_message = (
        "안녕하세요!"
    )
    await cl.Message(content=initial_message).send()
    
    # 2. 초기화 에러 체크
    if INITIALIZATION_ERROR:
        await cl.Message(
            content=f"⚠️ 시스템 초기화 중 오류가 발생했습니다:\n{INITIALIZATION_ERROR}"
        ).send()
        return
    
    if not RETRIEVER_RUNNABLE or not RESPONSE_LLM:
        await cl.Message(
            content="⚠️ 시스템이 아직 초기화 중입니다. 잠시 후 다시 시도해주세요."
        ).send()
        return
    
    try:
        # 3. 프롬프트 템플릿 생성 (가벼운 작업)
        chat_prompt = ChatPromptTemplate.from_messages([
            (
                "system",
                """
                원하는 system prompt를 작성
                """
            ),
            (   
                "user",
                """
                위의 검색 결과를 정리해서 제시해주세요."""
            )
        ])
        
        chain = RETRIEVER_RUNNABLE | chat_prompt | RESPONSE_LLM
        cl.user_session.set("runnable", chain)
        
    except Exception as e:
        await cl.Message(content=f"⚠️ 세션 초기화 중 오류 발생: {str(e)}").send()
        print(f"[Session Error] {e}")
        import traceback
        traceback.print_exc()

@cl.on_message
async def on_message(message: cl.Message):
    """메시지 처리"""
    user_input = str(message.content)
    runnable = cl.user_session.get("runnable")
    
    if not runnable:
        await cl.Message(content="⚠️ 시스템이 초기화되지 않았습니다. 페이지를 새로고침해주세요.").send()
        return
    
    msg = cl.Message(content="")
    await msg.send()
    
    try:
        config = RunnableConfig(callbacks=[cl.LangchainCallbackHandler()])
        async for chunk in runnable.astream(user_input, config=config):
            await msg.stream_token(chunk)
        await msg.update()
    except Exception as e:
        await cl.Message(content=f"⚠️ 오류가 발생했습니다: {str(e)}").send()
        print(f"[Message Error] {e}")
        import traceback
        traceback.print_exc()

 

앞서 전역 변수를 설정해 준 것에 이어서 작성하자면,

이렇게 필요한 것들을 전역변수로 설정해서 서버를 처음 시작할 때, initialize_global_resources를 활용해서 초기에 한번 불러와준다.

 

1. @cl.on_chat_start

그 이후에는 chainlit을 위한 세션 설정을 진행해 준다. 맨 처음에 보여준 예제와 마찬가지로 코드를 비슷하게 작성해 주면 되는데, 

여기서 조금 다른 점이라고 한다면, 이후에 디코딩할 llm의 프롬프트를 작성해 준다는 점과 불러온 RAG 모듈인 retriever_runnable 그리고 llm을 하나의 체인으로 연결해 준다. 

이렇게 설정해 주면 사용자의 메시지가 retriever_runnable에 입력이 되고, 검색된 문서들은 chat_prompt에 전달이 되어 최종적으로 응답이 생성된다. 그리고 이 완성된 chain 객체를 cl.user_session에 저장하여 해당 세션 내의 모든 메세지 처리에서 재사용할 수 있게 한다. 

 

2. @cl.on_message

세션에 저장된 runnable 체인을 불러오고, 사용자에게 응답이 전송될 cl.Message 객체를 생성하고 미리 화면에 표시한다.

이후에 체인을 실행하면 비동기 스트리밍으로 실시간으로 응답이 생성되는 게 사용자에게 보인다.

 

또, 마찬가지로 cl.LangchainCallbackHandler()를 callbacks로 설정하여 LangChain/LCEL 체인의 **중간 단계 (Chain of Thought)**를 Chainlit UI에 자동으로 시각화하게 된다.

 

스트리밍이 완료된 후에  await msg.update()를 호출해서 메시지를 최종적으로 완료한다.

 

 

이렇게 해서 chainlit을 통해서 간단한 챗봇의 형태를 쉽게 개발해 볼 수 있다.

물론 이를 아직 챗봇형태로 chainlit을 이용해서 배포를 하지는 않았지만, 추후에 이걸 활용하게 된다면 해당 부분도 포스팅해보려고 한다.

 

사실상 이거만 보면, 내가 만든 retriever_runnable이 어떻게 작동하는지 알 수 없기 때문에 무슨 챗봇이 만들어졌고, 어떻게 답변을 하고 있는지를 볼 수 없기 때문에 이해하기 쉽지 않으리라 생각하지만, 전반적으로 chainlit을 어떻게 구축하고 어떤 순서로 진행되어야 하는지에 파악하는 데는 도움이 될 수 있으리라 생각합니다..

 

감사합니다.