290 lines
9.2 KiB
Python
290 lines
9.2 KiB
Python
"""
|
|
PostgreSQL Database Agent - 基于 LangChain 的 PostgreSQL 数据库查询代理
|
|
"""
|
|
import os
|
|
import sys
|
|
import logging
|
|
from typing import Optional
|
|
from datetime import datetime
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel, Field
|
|
import uvicorn
|
|
|
|
# 添加 common 模块路径
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# 导入回调工具
|
|
try:
|
|
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
|
|
CALLBACK_ENABLED = True
|
|
except ImportError:
|
|
CALLBACK_ENABLED = False
|
|
AgentCallbackHandler = None
|
|
CallbackContextManager = None
|
|
|
|
# 配置日志
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 环境变量
|
|
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
|
|
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8000"))
|
|
POD_NAME = os.getenv("POD_NAME", "postgresql-agent")
|
|
USER_ID = os.getenv("USER_ID", "")
|
|
|
|
# PostgreSQL 配置
|
|
POSTGRES_HOST = os.getenv("POSTGRES_HOST", "")
|
|
POSTGRES_PORT = int(os.getenv("POSTGRES_PORT", "5432"))
|
|
POSTGRES_USER = os.getenv("POSTGRES_USER", "")
|
|
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "")
|
|
POSTGRES_DATABASE = os.getenv("POSTGRES_DATABASE", "")
|
|
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
|
|
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
|
|
|
|
# FastAPI 应用
|
|
app = FastAPI(
|
|
title="PostgreSQL Database Agent",
|
|
description="基于 LangChain 的 PostgreSQL 数据库查询代理",
|
|
version="1.0.0"
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
# ==================== 请求/响应模型 ====================
|
|
|
|
class QueryRequest(BaseModel):
|
|
"""查询请求"""
|
|
query: str = Field(..., description="自然语言查询")
|
|
user_id: Optional[str] = Field(None, description="用户ID（用于计费回调）")
|
|
openai_api_key: Optional[str] = Field(None, description="OpenAI API 密钥（可选，覆盖环境变量）")
|
|
|
|
|
|
class QueryResponse(BaseModel):
|
|
"""查询响应"""
|
|
query: str
|
|
result: str
|
|
sql: Optional[str] = None
|
|
timestamp: str
|
|
|
|
|
|
class ConnectRequest(BaseModel):
|
|
"""连接请求"""
|
|
host: str = Field(..., description="PostgreSQL 主机地址")
|
|
port: int = Field(default=5432, description="PostgreSQL 端口")
|
|
user: str = Field(..., description="PostgreSQL 用户名")
|
|
password: str = Field(..., description="PostgreSQL 密码")
|
|
database: str = Field(..., description="数据库名")
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
"""健康检查响应"""
|
|
status: str
|
|
pod_name: str
|
|
connected: bool
|
|
callback_enabled: bool
|
|
database: Optional[str] = None
|
|
timestamp: str
|
|
|
|
|
|
# ==================== 全局变量 ====================
|
|
|
|
db_agent = None
|
|
db_connection_info = None
|
|
callback_handler: Optional[AgentCallbackHandler] = None
|
|
|
|
|
|
# ==================== 生命周期 ====================
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
"""应用启动时初始化"""
|
|
global callback_handler
|
|
|
|
if CALLBACK_ENABLED:
|
|
callback_handler = AgentCallbackHandler(
|
|
agent_name=POD_NAME,
|
|
user_id=USER_ID
|
|
)
|
|
logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
|
|
else:
|
|
logger.warning("回调模块未加载，计费回调功能不可用")
|
|
|
|
# 尝试自动连接
|
|
if all([POSTGRES_HOST, POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_DATABASE, OPENAI_API_KEY]):
|
|
try:
|
|
create_db_agent(
|
|
host=POSTGRES_HOST,
|
|
port=POSTGRES_PORT,
|
|
user=POSTGRES_USER,
|
|
password=POSTGRES_PASSWORD,
|
|
database=POSTGRES_DATABASE,
|
|
api_key=OPENAI_API_KEY
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"自动连接失败: {e}")
|
|
|
|
|
|
# ==================== 辅助函数 ====================
|
|
|
|
def create_db_agent(host: str, port: int, user: str, password: str, database: str, api_key: str):
|
|
"""创建数据库 Agent"""
|
|
global db_agent, db_connection_info
|
|
|
|
try:
|
|
from langchain_community.utilities import SQLDatabase
|
|
from langchain_community.agent_toolkits import create_sql_agent
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
# 创建数据库连接
|
|
db_uri = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
|
|
db = SQLDatabase.from_uri(db_uri)
|
|
|
|
# 创建 LLM
|
|
llm = ChatOpenAI(
|
|
model=MODEL_NAME,
|
|
temperature=0,
|
|
openai_api_key=api_key
|
|
)
|
|
|
|
# 创建 SQL Agent
|
|
db_agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)
|
|
db_connection_info = {"host": host, "port": port, "database": database}
|
|
|
|
logger.info(f"PostgreSQL Agent 连接成功: {host}:{port}/{database}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"创建 PostgreSQL Agent 失败: {e}")
|
|
raise
|
|
|
|
|
|
# ==================== API 端点 ====================
|
|
|
|
@app.get("/health", response_model=HealthResponse)
|
|
@app.get("/", response_model=HealthResponse)
|
|
async def health_check():
|
|
"""健康检查"""
|
|
return HealthResponse(
|
|
status="healthy",
|
|
pod_name=POD_NAME,
|
|
connected=db_agent is not None,
|
|
callback_enabled=CALLBACK_ENABLED,
|
|
database=db_connection_info.get("database") if db_connection_info else None,
|
|
timestamp=datetime.utcnow().isoformat()
|
|
)
|
|
|
|
|
|
@app.post("/connect")
|
|
async def connect_database(request: ConnectRequest):
|
|
"""连接数据库"""
|
|
api_key = OPENAI_API_KEY
|
|
if not api_key:
|
|
raise HTTPException(status_code=400, detail="OPENAI_API_KEY 未设置")
|
|
|
|
try:
|
|
create_db_agent(
|
|
host=request.host,
|
|
port=request.port,
|
|
user=request.user,
|
|
password=request.password,
|
|
database=request.database,
|
|
api_key=api_key
|
|
)
|
|
return {
|
|
"status": "connected",
|
|
"database": request.database,
|
|
"timestamp": datetime.utcnow().isoformat()
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/query", response_model=QueryResponse)
|
|
async def query_database(request: QueryRequest):
|
|
"""执行自然语言查询"""
|
|
global db_agent
|
|
|
|
# 获取 API key - 优先使用请求中的
|
|
api_key = request.openai_api_key or OPENAI_API_KEY
|
|
if not api_key:
|
|
raise HTTPException(status_code=400, detail="OPENAI_API_KEY 未设置，请在请求中传入 openai_api_key 或设置环境变量")
|
|
|
|
# 如果未连接，尝试使用环境变量连接
|
|
if db_agent is None:
|
|
if not all([POSTGRES_HOST, POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_DATABASE]):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="数据库未连接。请先调用 /connect 或设置环境变量"
|
|
)
|
|
|
|
create_db_agent(
|
|
host=POSTGRES_HOST,
|
|
port=POSTGRES_PORT,
|
|
user=POSTGRES_USER,
|
|
password=POSTGRES_PASSWORD,
|
|
database=POSTGRES_DATABASE,
|
|
api_key=api_key
|
|
)
|
|
elif request.openai_api_key:
|
|
# 如果请求中提供了新的 API key，重新创建 agent
|
|
logger.info(f"使用请求中的 OpenAI API Key 重新初始化 Agent...")
|
|
create_db_agent(
|
|
host=db_connection_info["host"],
|
|
port=db_connection_info["port"],
|
|
user=POSTGRES_USER,
|
|
password=POSTGRES_PASSWORD,
|
|
database=db_connection_info["database"],
|
|
api_key=api_key
|
|
)
|
|
|
|
try:
|
|
# 使用回调上下文管理器
|
|
if CALLBACK_ENABLED and callback_handler and request.user_id:
|
|
with CallbackContextManager(
|
|
handler=callback_handler,
|
|
user_id=request.user_id,
|
|
request_id=f"postgresql-query-{int(datetime.utcnow().timestamp())}"
|
|
) as ctx:
|
|
ctx.add_tool("postgresql_query")
|
|
ctx.add_tool("sql_agent")
|
|
|
|
result = db_agent.invoke({"input": request.query})
|
|
|
|
return QueryResponse(
|
|
query=request.query,
|
|
result=result.get("output", str(result)),
|
|
timestamp=datetime.utcnow().isoformat()
|
|
)
|
|
else:
|
|
result = db_agent.invoke({"input": request.query})
|
|
return QueryResponse(
|
|
query=request.query,
|
|
result=result.get("output", str(result)),
|
|
timestamp=datetime.utcnow().isoformat()
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"查询失败: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# ==================== 主入口 ====================
|
|
|
|
def main():
|
|
"""主函数"""
|
|
logger.info(f"启动 PostgreSQL Agent - {POD_NAME}")
|
|
logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
|
|
|
|
uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|