279 lines
9.2 KiB
Python
279 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
|
"""NL2SQL 전용 워커 프로세스
|
|
|
|
Usage: python nl2sql_worker.py <port>
|
|
|
|
담당 도구:
|
|
run_sql, query_pv_history, get_tag_metadata, list_drawings, query_with_nl
|
|
|
|
특징:
|
|
- PostgreSQL 직접 연결
|
|
- LLM SQL 생성 + DB 실행 분리
|
|
- 메모리: ~1GB (SQL 생성용 LLM)
|
|
- 생명주기: 메인 서버 종료 시까지 유지
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
import sys
|
|
import os
|
|
|
|
# mcp-server 디렉토리를 Python 경로에 추가
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import logging
|
|
import asyncio
|
|
from functools import lru_cache
|
|
|
|
from fastapi import FastAPI, Request
|
|
import uvicorn
|
|
import httpx
|
|
|
|
# ── 설정 ─────────────────────────────────────────────────────────────────────
|
|
|
|
DB_CONNECTION_STRING = "postgresql://postgres:postgres@localhost:5432/iiot_platform"
|
|
DB_TIMEOUT = 10
|
|
|
|
VLLM_BASE_URL = "http://localhost:8000/v1"
|
|
VLLM_MODEL = "Qwen/Qwen3-Coder-Next-FP8"
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
stream=sys.stderr,
|
|
format="%(asctime)s [nl2sql_worker] %(levelname)s %(message)s",
|
|
)
|
|
|
|
app = FastAPI()
|
|
|
|
# ── DB 연결 풀 ───────────────────────────────────────────────────────────────
|
|
|
|
def _get_db_connection():
|
|
import psycopg
|
|
return psycopg.connect(DB_CONNECTION_STRING, connect_timeout=DB_TIMEOUT)
|
|
|
|
# ── LLM 클라이언트 ───────────────────────────────────────────────────────────
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _llm_client():
|
|
from openai import AsyncOpenAI
|
|
return AsyncOpenAI(base_url=VLLM_BASE_URL, api_key="dummy")
|
|
|
|
async def _generate_sql(natural_language: str) -> str:
|
|
"""자연어를 SQL로 변환."""
|
|
client = _llm_client()
|
|
|
|
prompt = f"""다음 자연어 질문을 PostgreSQL SQL 쿼리로 변환하세요.
|
|
데이터베이스는 TimescaleDB 기반의 IIoT 플랫폼입니다.
|
|
|
|
질문:
|
|
{natural_language}
|
|
|
|
SQL 쿼리 (SELECT 문만, 설명 없이):"""
|
|
|
|
response = await client.chat.completions.create(
|
|
model=VLLM_MODEL,
|
|
messages=[
|
|
{"role": "system", "content": "You are a PostgreSQL SQL expert. Generate only valid SQL queries."},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
max_tokens=1024,
|
|
temperature=0.1,
|
|
)
|
|
return response.choices[0].message.content.strip()
|
|
|
|
# ── NL2SQL 도구 구현 ─────────────────────────────────────────────────────────
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
"""워커 헬스체크."""
|
|
return {"status": "ok"}
|
|
|
|
@app.post("/execute")
|
|
async def execute(request: Request):
|
|
"""HTTP 요청을 MCP 도구 호출로 변환."""
|
|
body = await request.json()
|
|
tool = body["tool"]
|
|
params = body["params"]
|
|
|
|
try:
|
|
if tool == "run_sql":
|
|
result = await _run_sql(**params)
|
|
elif tool == "query_pv_history":
|
|
result = await _query_pv_history(**params)
|
|
elif tool == "get_tag_metadata":
|
|
result = await _get_tag_metadata(**params)
|
|
elif tool == "list_drawings":
|
|
result = await _list_drawings(**params)
|
|
elif tool == "query_with_nl":
|
|
result = await _query_with_nl(**params)
|
|
else:
|
|
return {"success": False, "error": f"Unknown tool: {tool}"}
|
|
|
|
return result
|
|
except Exception as e:
|
|
logging.error(f"Error executing {tool}: {e}")
|
|
return {"success": False, "error": str(e)}
|
|
|
|
async def _run_sql(sql: str) -> str:
|
|
"""SQL 실행."""
|
|
conn = _get_db_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute(sql)
|
|
if cur.description:
|
|
columns = [desc[0] for desc in cur.description]
|
|
rows = cur.fetchall()
|
|
data = [dict(zip(columns, row)) for row in rows]
|
|
return {
|
|
"success": True,
|
|
"columns": columns,
|
|
"count": len(data),
|
|
"data": data,
|
|
}
|
|
else:
|
|
conn.commit()
|
|
return {
|
|
"success": True,
|
|
"message": f"Query executed successfully. {cur.rowcount} rows affected.",
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
async def _query_pv_history(tag_names: list[str], time_from: str, time_to: str, limit: int = 100) -> str:
|
|
"""과거 값(PV) 히스토리 조회."""
|
|
if not tag_names:
|
|
return {"success": False, "error": "tag_names is required"}
|
|
|
|
conn = _get_db_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
# TimescaleDB의 time_bucket 함수 사용
|
|
cur.execute(
|
|
"""
|
|
SELECT time_bucket('1 min', ts) AS time, tag_name, value
|
|
FROM realtime_table
|
|
WHERE tag_name = ANY(%s)
|
|
AND ts >= %s
|
|
AND ts <= %s
|
|
ORDER BY time DESC
|
|
LIMIT %s
|
|
""",
|
|
(tag_names, time_from, time_to, limit),
|
|
)
|
|
columns = ["time", "tag_name", "value"]
|
|
rows = cur.fetchall()
|
|
data = [dict(zip(columns, row)) for row in rows]
|
|
return {
|
|
"success": True,
|
|
"tag_names": tag_names,
|
|
"time_range": {"from": time_from, "to": time_to},
|
|
"limit": limit,
|
|
"count": len(data),
|
|
"data": data,
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
async def _get_tag_metadata(query: str, limit: int = 10) -> str:
|
|
"""태그 메타데이터 검색."""
|
|
conn = _get_db_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT DISTINCT tag_name, unit, description
|
|
FROM realtime_table
|
|
WHERE tag_name ILIKE %s
|
|
ORDER BY tag_name
|
|
LIMIT %s
|
|
""",
|
|
(f"%{query}%", limit),
|
|
)
|
|
columns = ["tag_name", "unit", "description"]
|
|
rows = cur.fetchall()
|
|
data = [dict(zip(columns, row)) for row in rows]
|
|
return {
|
|
"success": True,
|
|
"query": query,
|
|
"count": len(data),
|
|
"tags": data,
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
async def _list_drawings(unit_no: str = None) -> str:
|
|
"""단위별 도면 목록 조회."""
|
|
conn = _get_db_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
if unit_no:
|
|
cur.execute(
|
|
"""
|
|
SELECT DISTINCT name
|
|
FROM node_map_master
|
|
WHERE name LIKE %s
|
|
ORDER BY name
|
|
""",
|
|
(f"{unit_no}%",),
|
|
)
|
|
else:
|
|
cur.execute(
|
|
"""
|
|
SELECT DISTINCT name
|
|
FROM node_map_master
|
|
ORDER BY name
|
|
"""
|
|
)
|
|
columns = ["name"]
|
|
rows = cur.fetchall()
|
|
data = [dict(zip(columns, row[0])) for row in rows]
|
|
return {
|
|
"success": True,
|
|
"unit_no": unit_no,
|
|
"count": len(data),
|
|
"names": [d["name"] for d in data],
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
async def _query_with_nl(question: str) -> str:
|
|
"""자연어로 SQL 쿼리 실행."""
|
|
sql = await _generate_sql(question)
|
|
|
|
conn = _get_db_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute(sql)
|
|
if cur.description:
|
|
columns = [desc[0] for desc in cur.description]
|
|
rows = cur.fetchall()
|
|
data = [dict(zip(columns, row)) for row in rows]
|
|
return {
|
|
"success": True,
|
|
"sql": sql,
|
|
"columns": columns,
|
|
"count": len(data),
|
|
"data": data,
|
|
}
|
|
else:
|
|
conn.commit()
|
|
return {
|
|
"success": True,
|
|
"sql": sql,
|
|
"message": f"Query executed successfully. {cur.rowcount} rows affected.",
|
|
}
|
|
except Exception as db_error:
|
|
return {
|
|
"success": False,
|
|
"sql": sql,
|
|
"error": str(db_error),
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
# ── 메인 ─────────────────────────────────────────────────────────────────────
|
|
|
|
if __name__ == "__main__":
|
|
port = int(sys.argv[1]) if len(sys.argv) > 1 else 5003
|
|
logging.info(f"Starting NL2SQL worker on port {port}")
|
|
uvicorn.run(app, host="0.0.0.0", port=port)
|