[Backend] LangChain 답변 생성 이후 토큰 사용량 계산

기존 QA 코드

from crud.langchain import *
from langchain.prompts import PromptTemplate
from langchain_service.embeddings.get_vector import text_to_vector
from langchain_service.llms.setup import get_llm
from langchain_service.memory.relevant_message import get_relevant_messages

def qa_chain(db: Session, session_id, project_id, user_email, conversation, provider="openai", model=None):
    llm = get_llm(provider, model)  # LangChain의 get LLM
    vector = text_to_vector(conversation)  # LangChain의 get_embedding 기반으로 만든 함수 text_to_vector
    print(f"vector1: {vector}")

    # 벡터 검색을 통한 유사 대화 검색
    relevant_messages = get_relevant_messages(db, session_id, vector, top_n=5)

    print(" 검색된 유사 대화 기록:")
    for idx, msg in enumerate(relevant_messages, 1):
        print(f"{idx}. [{msg['message_role'].capitalize()}] {msg['conversation']}")

    # 관련 대화만 컨텍스트로 활용
    formatted_history = "\n".join(
        [f"{msg['message_role'].capitalize()}: {msg['conversation']}" for msg in relevant_messages]
    )

    prompt = PromptTemplate(
        input_variables=["history", "input"],
        template="{history}\nHuman: {input}\nAI:"
    )
    chain = prompt | llm
    response = chain.invoke({"history": formatted_history, "input": conversation})

    print(f"response: {response.content}")
    vector2 = text_to_vector(response.content)  # 응답도 벡터로 변환
    print(f"vector2: {vector2}")

    # 대화 저장 (유저 메시지 + AI 응답)
    add_message(db=db, session_id=session_id, project_id=project_id, user_email=user_email,
                message_role='user', conversation=conversation, vector_memory=vector)

    add_message(db=db, session_id=session_id, project_id=project_id, user_email=user_email,
                message_role='assistant', conversation=response.content, vector_memory=vector2)

    return response.content

프로젝트 내의 지식베이스 내용과 세션 내 대화 이력을 바탕으로 질문에 대하여 답변을 생성해주는 것이 기존 QA Chain 코드가 가지고 있던 기능입니다.

기존 대화 이력 및 지식베이스의 데이터는 전부 한 번에 LLM에게 전달하는 것이 아니라,

벡터 기반 검색으로 DB 내에서 관련 정보만 추출해서 전달하는 방식으로 답변을 생성합니다.

또한 새로 추가된 대화도 실시간으로 DB에 저장합니다.

신규 QA 코드

from langchain_service.embeddings.get_vector import text_to_vector
from langchain_service.llms.setup import get_llm
from langchain_service.memory.relevant_message import get_relevant_messages
from langchain_community.callbacks.manager import get_openai_callback
from langchain_core.output_parsers import StrOutputParser
from crud.langchain import *
from langchain.prompts import PromptTemplate

def qa_chain(db: Session, session_id, project_id, user_email, conversation, provider="openai", model=None):
    llm = get_llm(provider, model)
    vector = text_to_vector(conversation)

    relevant_messages = get_relevant_messages(db, session_id, vector, top_n=5)

    formatted_history = "\n".join(
        [f"{msg['message_role'].capitalize()}: {msg['conversation']}" for msg in relevant_messages]
    )

    prompt = PromptTemplate(
        input_variables=["history", "input"],
        template="{history}\nHuman: {input}\nAI:"
    )

    chain = prompt | llm | StrOutputParser()

    with get_openai_callback() as cb:
        response_text = chain.invoke({"history": formatted_history, "input": conversation})

        # 사용량 로그 출력
        print(f"[LLM 사용량] prompt: {cb.prompt_tokens} / completion: {cb.completion_tokens} / total: {cb.total_tokens} / cost: ${cb.total_cost:.6f}")


    # 응답 벡터 저장
    vector2 = text_to_vector(response_text)

    add_message(db=db, session_id=session_id, project_id=project_id, user_email=user_email,
                message_role='user', conversation=conversation, vector_memory=vector)

    add_message(db=db, session_id=session_id, project_id=project_id, user_email=user_email,
                message_role='assistant', conversation=response_text, vector_memory=vector2)

    return response_text