feat: Refactor LLM workflow and Streamlit interface for Lang2SQL

- Enhance query generation process with refined input and table info retrieval
- Update Streamlit app to support database environment input
- Modify graph and chain logic to improve query generation
- Add DataHub server configuration to README

Author: ehddnr301

README.md

@@ -62,6 +62,7 @@ LANGCHAIN_TRACING_V2=true
 LANGCHAIN_PROJECT=autosql
 LANGCHAIN_ENDPOINT=https://api.smith.langchain.com
 LANGCHAIN_API_KEY=your-langchain-api-key
+DATAHUB_SERVER=http://localhost:8080
 ```
 
 ---

interface/streamlit_app.py

@@ -3,14 +3,19 @@
 from llm_utils.graph import builder
 
 # Streamlit 앱 제목
-st.title("AutoSQL")
+st.title("Lang2SQL")
 
 # 사용자 입력 받기
 user_query = st.text_area(
     "쿼리를 입력하세요:",
     value="고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리",
 )
 
+user_database_env = st.text_area(
+    "db 환경정보를 입력하세요:",
+    value="duckdb",
+)
+
 
 # Token usage 집계 함수 정의
 def summarize_total_tokens(data):
@@ -25,10 +30,18 @@ def summarize_total_tokens(data):
 if st.button("쿼리 실행"):
     # 그래프 컴파일 및 쿼리 실행
     graph = builder.compile()
-    human_message = HumanMessage(content=user_query)
-    res = graph.invoke(input=human_message)
-    total_tokens = summarize_total_tokens(res)
+
+    res = graph.invoke(
+        input={
+            "messages": [HumanMessage(content=user_query)],
+            "user_database_env": user_database_env,
+            "best_practice_query": "",
+        }
+    )
+    total_tokens = summarize_total_tokens(res["messages"])
 
     # 결과 출력
     st.write("총 토큰 사용량:", total_tokens)
-    st.write("결과:", res[-1].content)
+    st.write("결과:", res["generated_query"].content)
+    st.write("AI가 재해석한 사용자 질문:\n", res["refined_input"].content)
+    st.write("참고한 테이블 목록:", res["searched_tables"])

llm_utils/chains.py

@@ -1,96 +1,120 @@
 import os
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 
-from .tools import get_table_info, get_column_info
 from .llm_factory import get_llm
 
 llm = get_llm(
-    model_type="openai", model_name="gpt-4o", openai_api_key=os.getenv("OPENAI_API_KEY")
+    model_type="openai",
+    model_name="gpt-4o-mini",
+    openai_api_key=os.getenv("OPENAI_API_KEY"),
 )
 
 
-# ToolChoiceChain
-def create_tool_choice_chain(llm):
+def create_query_refiner_chain(llm):
     tool_choice_prompt = ChatPromptTemplate.from_messages(
         [
             (
                 "system",
                 """
-                너는 User Input에 대해 관련된 테이블, 컬럼을 찾는 Assistance이다.
+                당신은 데이터 분석 전문가(데이터 분석가 페르소나)입니다.
+                현재 subscription_activities, contract_activities, marketing_activities, 
+                sales_activities, success_activities, support_activities, trial_activities 데이터를 
+                보유하고 있으며, 사용자의 질문이 모호할 경우에도 우리가 가진 데이터를 기반으로 
+                충분히 답변 가능한 형태로 질문을 구체화해 주세요.
+
+                주의:
+                - 사용자에게 추가 정보를 요구하는 ‘재질문(추가 질문)’을 하지 마세요.
+                - 질문에 포함해야 할 요소(예: 특정 기간, 대상 유저 그룹, 분석 대상 로그 종류 등)가
+                  불충분하더라도, 합리적으로 추론해 가정한 뒤
+                  답변에 충분한 질문 형태로 완성해 주세요.
+
+                예시:
+                사용자가 "유저 이탈 원인이 궁금해요"라고 했다면,
+                재질문 형식이 아니라
+                "최근 1개월 간의 접속·결제 로그를 기준으로,
+                주로 어떤 사용자가 어떤 과정을 거쳐 이탈하는지를 분석해야 한다"처럼
+                분석 방향이 명확해진 질문 한 문장(또는 한 문단)으로 정리해 주세요.
+
+                최종 출력 형식 예시:
+                ------------------------------
+                구체화된 질문:
+                "최근 1개월 동안 고액 결제 경험이 있는 유저가 
+                행동 로그에서 이탈 전 어떤 패턴을 보였는지 분석"
+
+                가정한 조건:
+                - 최근 1개월치 행동 로그와 결제 로그 중심
+                - 고액 결제자(월 결제액 10만 원 이상) 그룹 대상으로 한정
+                ------------------------------
                 """,
             ),
             MessagesPlaceholder(variable_name="user_input"),
             (
                 "system",
                 """
-                위의 질의와 관련된 테이블을 찾아주세요
-                다음 tool을 사용할 수 있습니다:
-                get_table_info - 전체 table_name과 table_description을 가져옵니다.
-                get_column_info - table_name을 input으로 넣으면 column_name과 column description을 가져옵니다.
-                아래 툴을 사용해주세요
+                위 사용자의 입력을 바탕으로
+                분석 관점에서 **충분히 답변 가능한 형태**로
+                "구체화된 질문"을 작성하고,
+                필요한 경우 가정이나 전제 조건을 함께 제시해 주세요.
                 """,
             ),
-            MessagesPlaceholder(variable_name="tool_choice"),
         ]
     )
 
-    tools = [get_table_info, get_column_info]
+    return tool_choice_prompt | llm
 
-    return tool_choice_prompt | llm.bind_tools(tools)
 
-
-# TableFilterChain
-def create_table_filter_chain(llm):
-    table_filter_prompt = ChatPromptTemplate.from_messages(
+# QueryMakerChain
+def create_query_maker_chain(llm):
+    query_maker_prompt = ChatPromptTemplate.from_messages(
         [
-            MessagesPlaceholder(variable_name="user_input"),
             (
                 "system",
                 """
-                너는 위의 User Input에 대해 관련된 테이블을 찾는 Assistance이다.
-                아래 테이블 목록을 참고해서 관련된 테이블을 찾아주세요.
-                참고사항은 
-                    dim_~: 테이블은 metadata 테이블임
-                    fact_~: 테이블은 실제 데이터가 저장된 테이블임
-                    stg_~: 테이블은 데이터 적재 테이블임
-                응답형태는 'table_name - table_description' 형태로 출력해주세요.
-                최소 2개 이상의 테이블을 출력해주세요.
-                테이블 목록은 다음과 같습니다:
+                당신은 데이터 분석 전문가(데이터 분석가 페르소나)입니다.
+                사용자의 질문을 기반으로, 주어진 테이블과 컬럼 정보를 활용하여 적절한 SQL 쿼리를 생성하세요.
+                
+                주의사항:
+                - 사용자의 질문이 다소 모호하더라도, 주어진 데이터를 참고하여 합리적인 가정을 통해 SQL 쿼리를 완성하세요.
+                - 불필요한 재질문 없이, 가능한 가장 명확한 분석 쿼리를 만들어 주세요.
+                - 최종 출력 형식은 반드시 아래와 같아야 합니다.
+                
+                최종 형태 예시:
+                
+                <SQL>
+                ```sql
+                    SELECT COUNT(DISTINCT user_id) 
+                    FROM stg_users 
+                ```
+                
+                <해석>
+                ```plaintext (max_length_per_line=100)
+                    이 쿼리는 stg_users 테이블에서 고유한 사용자의 수를 계산합니다.
+                    사용자는 유니크한 user_id를 가지고 있으며
+                    중복을 제거하기 위해 COUNT(DISTINCT user_id)를 사용했습니다.
+                ```
+
                 """,
             ),
-            MessagesPlaceholder(variable_name="searched_tables"),
-        ]
-    )
-    return table_filter_prompt | llm
-
-
-# QueryMakerChain
-def create_query_maker_chain(llm):
-    query_maker_prompt = ChatPromptTemplate.from_messages(
-        [
+            (
+                "system",
+                "아래는 사용자의 질문 및 구체화된 질문입니다:",
+            ),
             MessagesPlaceholder(variable_name="user_input"),
-            ("system", "너는 위의 User Input에 대해 쿼리를 작성하는 Assistance이다."),
-            ("system", "다음 테이블과 컬럼을 참고해서 쿼리를 작성해주세요."),
-            ("system", "테이블 목록은 다음과 같습니다:"),
+            MessagesPlaceholder(variable_name="refined_input"),
+            (
+                "system",
+                "다음은 사용자의 db 환경정보와 사용 가능한 테이블 및 컬럼 정보입니다:",
+            ),
+            MessagesPlaceholder(variable_name="user_database_env"),
             MessagesPlaceholder(variable_name="searched_tables"),
-            ("system", "컬럼 목록은 다음과 같습니다:"),
-            MessagesPlaceholder(variable_name="searched_columns"),
             (
                 "system",
-                """최종 형태는 반드시 아래와 같아야합니다.
-                최종 쿼리:
-                    SELECT COUNT(DISTINCT user_id) FROM stg_users WHERE user_id = 1
-                참고한 테이블 목록:
-                    stg_users, dim_users
-                참고한 컬럼 목록:
-                    stg_users.user_id, dim_users.user_id
-                """,
+                "위 정보를 바탕으로 사용자 질문에 대한 최적의 SQL 쿼리를 최종 형태 예시와 같은 형태로 생성하세요.",
             ),
         ]
     )
     return query_maker_prompt | llm
 
 
-tool_choice_chain = create_tool_choice_chain(llm)
-table_filter_chain = create_table_filter_chain(llm)
+query_refiner_chain = create_query_refiner_chain(llm)
 query_maker_chain = create_query_maker_chain(llm)