agent_management/agent_templates/agents/postgresql_agent/postgresql_agent.py

"""
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()