챗봇 tool 로 search_database_tables 추가

Author: ehddnr301

utils/llm/chatbot.py

@@ -4,11 +4,16 @@
 """
 
 from langchain_openai import ChatOpenAI
+from langchain_core.messages import SystemMessage
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import START, MessagesState, StateGraph
 from langgraph.prebuilt import ToolNode
 
-from utils.llm.tools import get_weather, get_famous_opensource
+from utils.llm.tools import (
+    get_weather,
+    get_famous_opensource,
+    search_database_tables,
+)
 
 
 class ChatBot:
@@ -36,7 +41,12 @@ def __init__(self, openai_api_key: str, model_name: str = "gpt-4o-mini"):
         """
         self.openai_api_key = openai_api_key
         self.model_name = model_name
-        self.tools = [get_weather, get_famous_opensource]  # 사용 가능한 tool 목록
+        # SQL 생성을 위한 데이터베이스 메타데이터 조회 도구
+        self.tools = [
+            search_database_tables,  # 데이터베이스 테이블 정보 검색
+            get_weather,  # 테스트용 도구 (추후 제거 가능)
+            get_famous_opensource,  # 테스트용 도구 (추후 제거 가능)
+        ]
         self.llm = self._setup_llm()  # LLM 인스턴스 설정
         self.app = self._setup_workflow()  # LangGraph 워크플로우 설정
 
@@ -49,7 +59,7 @@ def _setup_llm(self):
             ChatOpenAI: Tool이 바인딩된 LLM 인스턴스
         """
         llm = ChatOpenAI(
-            temperature=0.1,  # 응답의 일관성을 위해 낮은 temperature 설정
+            temperature=0.0,  # SQL 생성은 정확성이 중요하므로 0으로 설정
             openai_api_key=self.openai_api_key,
             model_name=self.model_name,
         )
@@ -79,8 +89,33 @@ def call_model(state: MessagesState):
             Returns:
                 dict: LLM 응답이 포함된 상태 업데이트
             """
-            # sys_msg = SystemMessage(content="You are a helpful assistant ")
-            response = self.llm.invoke(state["messages"])
+            # SQL 생성 전문 어시스턴트 시스템 메시지
+            sys_msg = SystemMessage(
+                content="""# 역할
+당신은 사용자가 SQL 쿼리를 생성하도록 돕는 전문 AI 어시스턴트입니다.
+
+# 주요 임무
+- 사용자의 자연어 질문을 이해하고 SQL 쿼리 생성에 필요한 정보를 파악합니다
+- 필요한 경우 데이터베이스 스키마나 메타데이터를 확인하기 위해 도구를 활용합니다
+- 단계별로 사용자와 대화하며 명확한 SQL 쿼리를 만들어갑니다
+- 생성된 SQL에 대해 이해하기 쉽게 설명합니다
+
+# 작업 프로세스
+1. 사용자의 의도를 명확히 파악
+2. 필요한 테이블/컬럼 정보 확인 (도구 사용)
+3. 사용자의 질문을 바탕으로 정보를 추출할 수 있는 명확한 질문으로 변환합니다
+
+# 주의사항
+- 항상 친절하고 명확하게 대화합니다
+- 이전 대화 맥락을 고려하여 일관성 있게 응답합니다
+- 사용자가 SQL을 이해할 수 있도록 단계별로 설명합니다
+
+---
+다음은 사용자와의 대화입니다:"""
+            )
+            # 시스템 메시지를 대화의 맨 앞에 추가
+            messages = [sys_msg] + state["messages"]
+            response = self.llm.invoke(messages)
             return {"messages": response}
 
         def route_model_output(state: MessagesState):

utils/llm/tools/__init__.py

@@ -4,12 +4,17 @@
     set_gms_server,
 )
 
-from utils.llm.tools.test import get_weather, get_famous_opensource
+from utils.llm.tools.chatbot_tool import (
+    get_weather,
+    get_famous_opensource,
+    search_database_tables,
+)
 
 __all__ = [
     "set_gms_server",
     "get_info_from_db",
     "get_metadata_from_db",
     "get_weather",
     "get_famous_opensource",
+    "search_database_tables",
 ]

utils/llm/tools/chatbot_tool.py

@@ -0,0 +1,120 @@
+"""
+LangGraph ChatBot에서 사용하는 도구(Tool) 함수들
+"""
+
+from langchain_core.tools import tool
+from utils.llm.retrieval import search_tables
+
+
+@tool
+def get_weather(city: str) -> str:
+    """
+    특정 도시의 현재 날씨 정보를 조회합니다.
+
+    이 함수는 도시 이름을 입력받아 해당 도시의 날씨 정보를 반환합니다.
+    사용자가 날씨, 기상, weather 등의 키워드와 함께 도시 이름을 언급하면 이 도구를 사용하세요.
+
+    Args:
+        city (str): 날씨를 확인하고 싶은 도시의 이름입니다.
+                   예: "Seoul", "New York", "Tokyo", "서울", "부산" 등
+                   영문과 한글 도시명을 모두 지원합니다.
+
+    Returns:
+        str: 해당 도시의 날씨 정보를 담은 문자열입니다.
+             현재는 항상 맑은 날씨를 반환합니다.
+
+    Examples:
+        >>> get_weather("Seoul")
+        'Seoul is sunny'
+
+        >>> get_weather("서울")
+        '서울 is sunny'
+
+    Note:
+        이 도구는 다음과 같은 경우에 사용하세요:
+        - "서울 날씨 어때?"
+        - "What's the weather in New York?"
+        - "도쿄의 날씨를 알려줘"
+        - "부산 날씨 확인해줘"
+    """
+    return f"{city} is sunny"
+
+
+@tool
+def get_famous_opensource() -> str:
+    """
+    가장 유명한 오픈소스 프로젝트를 조회합니다.
+
+    이 함수는 현재 가장 유명한 오픈소스 프로젝트의 이름을 반환합니다.
+    사용자가 유명한 오픈소스, 인기있는 오픈소스, 최고의 오픈소스 등을 물어보면 이 도구를 사용하세요.
+
+    Returns:
+        str: 가장 유명한 오픈소스 프로젝트 이름
+
+    Examples:
+        >>> get_famous_opensource()
+        'Lang2SQL'
+
+    Note:
+        이 도구는 다음과 같은 경우에 사용하세요:
+        - "제일 유명한 오픈소스가 뭐야?"
+        - "가장 인기있는 오픈소스는?"
+        - "최고의 오픈소스 프로젝트 알려줘"
+        - "유명한 오픈소스 추천해줘"
+    """
+    return "Lang2SQL"
+
+
+@tool
+def search_database_tables(
+    query: str, retriever_name: str = "기본", top_n: int = 5, device: str = "cpu"
+) -> dict:
+    """
+    사용자의 자연어 쿼리를 기반으로 관련된 데이터베이스 테이블 정보를 검색합니다.
+
+    이 함수는 SQL 쿼리 생성을 위해 필요한 테이블과 컬럼 정보를 찾아줍니다.
+    사용자가 어떤 테이블을 사용해야 할지, 어떤 컬럼이 있는지 물어보거나,
+    SQL 쿼리를 만들기 위한 스키마 정보가 필요할 때 이 도구를 사용하세요.
+
+    Args:
+        query (str): 검색하려는 자연어 질문입니다.
+                    예: "고객 정보를 조회하려면?", "주문 관련 테이블"
+        retriever_name (str, optional): 검색기 유형입니다.
+                                       "기본" 또는 "Reranker" 중 선택. 기본값은 "기본"
+        top_n (int, optional): 검색할 테이블 개수입니다. 기본값은 5개
+        device (str, optional): 모델 실행 장치입니다. "cpu" 또는 "cuda". 기본값은 "cpu"
+
+    Returns:
+        dict: 테이블 정보가 담긴 딕셔너리입니다.
+              각 테이블은 키로 저장되며, 값으로 테이블 설명과 컬럼 정보를 포함합니다.
+
+              예시 형태:
+              {
+                  "customers": {
+                      "table_description": "고객 정보 테이블",
+                      "customer_id": "고객 고유 ID",
+                      "name": "고객 이름",
+                      "email": "고객 이메일"
+                  },
+                  "orders": {
+                      "table_description": "주문 정보 테이블",
+                      "order_id": "주문 ID",
+                      "customer_id": "고객 ID (외래키)"
+                  }
+              }
+
+    Examples:
+        >>> search_database_tables("고객 정보가 필요해")
+        {'customers': {'table_description': '고객 정보 테이블', ...}}
+
+    Note:
+        이 도구는 다음과 같은 경우에 사용하세요:
+        - "어떤 테이블을 사용해야 해?"
+        - "고객 관련 테이블 정보를 알려줘"
+        - "주문 데이터는 어디에 있어?"
+        - "사용 가능한 컬럼을 보여줘"
+        - SQL 쿼리를 생성하기 전에 스키마 정보가 필요할 때
+    """
+    return search_tables(
+        query=query, retriever_name=retriever_name, top_n=top_n, device=device
+    )