from langflow.custom import CustomComponent
from langflow.io import Input, Output, MessageTextInput
from langchain_openai import ChatOpenAI

class OpenAIComponent(CustomComponent):
    display_name = "OpenAI"
    description = "A standard OpenAI model"

    # 定义 UI 输入
    inputs = [
        MessageTextInput(name="model_name", value="gpt-4"),
        Input(name="temperature", value=0.7),
    ]

    # 定义输出
    outputs = [
        Output(name="model", display_name="LLM", method="build_model"),
    ]

    # 核心构建逻辑
    def build_model(self) -> ChatOpenAI:
        return ChatOpenAI(
            model=self.model_name,
            temperature=self.temperature
        )

# 伪代码示例
from fastapi_limiter.depends import RateLimiter

async def get_rate_limit(deployment_id: str):
    # 从数据库获取该部署的限流配置
    config = await db.get_deployment_config(deployment_id)
    return config.limit_string # "10/second"

@router.post("/agent/{endpoint_slug}", dependencies=[Depends(RateLimiter(key_func=get_remote_address))])
async def invoke_agent(...):
    ...

# src/backend/base/langflow/graph/vertex/base.py 的扩展逻辑

from langfuse.callback import CallbackHandler

class Vertex:
    # ... 原有代码 ...

    def _build_results(self, custom_component, target_handle_name):
        result = custom_component.build(...)
        
        # [新增] 生产环境注入 Langfuse
        if self.graph.run_config.get("enable_tracing"):
            langfuse_handler = CallbackHandler(
                public_key=..., 
                secret_key=..., 
                host=...,
                session_id=self.graph.session_id,
                trace_name=self.graph.flow_name
            )
            
            # 如果 result 是 LangChain Runnable
            if hasattr(result, "with_config"):
                result = result.with_config(callbacks=[langfuse_handler])
            # 如果是 AgentExecutor
            elif hasattr(result, "callbacks"):
                result.callbacks.append(langfuse_handler)
                
        return result

from datetime import datetime
from typing import Optional, Dict, Any
from uuid import UUID, uuid4
from sqlmodel import Field, SQLModel, JSON, Column

class DeploymentBase(SQLModel):
    name: str = Field(index=True)
    description: Optional[str] = Field(default=None)
    endpoint_slug: str = Field(unique=True, index=True, description="The URL path for external access")
    rate_limit: Optional[str] = Field(default="60/minute", description="Rate limit string e.g., '10/second'")
    is_active: bool = Field(default=True)
    
class Deployment(DeploymentBase, table=True):
    id: UUID = Field(default_factory=uuid4, primary_key=True)
    flow_id: UUID = Field(foreign_key="flow.id")
    # 核心：存储发布时的 Flow JSON 快照，确保生产环境不受开发环境修改的影响
    graph_snapshot: Dict[str, Any] = Field(default={}, sa_column=Column(JSON))
    # 存储生产环境特定的配置（如覆盖 API Key）
    production_config: Dict[str, Any] = Field(default={}, sa_column=Column(JSON))
    
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)

class DeploymentRead(DeploymentBase):
    id: UUID
    flow_id: UUID
    created_at: datetime

class DeploymentCreate(DeploymentBase):
    flow_id: UUID
    production_config: Optional[Dict[str, Any]] = None

import os
from typing import Optional
from langfuse.callback import CallbackHandler

def get_langfuse_callback(
    user_id: Optional[str] = None, 
    session_id: Optional[str] = None, 
    trace_name: Optional[str] = None,
    tags: list = []
) -> Optional[CallbackHandler]:
    """
    工厂函数：根据环境变量安全地创建 Langfuse Handler
    """
    public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
    secret_key = os.getenv("LANGFUSE_SECRET_KEY")
    host = os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com")

    if not public_key or not secret_key:
        # 如果未配置，返回 None，避免生产环境报错
        return None

    return CallbackHandler(
        public_key=public_key,
        secret_key=secret_key,
        host=host,
        user_id=user_id,
        session_id=session_id,
        trace_name=trace_name,
        tags=["agent-server"] + tags
    )

from fastapi import APIRouter, Depends, HTTPException, status
from sqlmodel import Session, select
from uuid import UUID

from langflow.services.database.models.flow import Flow
from langflow.services.get_service import get_session
# 引入上面定义的模型
from langflow.services.database.models.deployment import Deployment, DeploymentCreate, DeploymentRead

router = APIRouter(prefix="/deployments", tags=["Deployments"])

@router.post("/publish", response_model=DeploymentRead)
def publish_flow(
    deploy_data: DeploymentCreate,
    session: Session = Depends(get_session),
):
    """
    将当前的 Flow 状态发布为生产 Deployment
    """
    # 1. 获取原始 Flow
    flow = session.get(Flow, deploy_data.flow_id)
    if not flow:
        raise HTTPException(status_code=404, detail="Flow not found")
    
    # 2. 检查 Slug 是否重复
    existing = session.exec(select(Deployment).where(Deployment.endpoint_slug == deploy_data.endpoint_slug)).first()
    if existing:
        # 简单策略：如果存在则更新快照 (也可以设计为报错)
        existing.graph_snapshot = flow.data
        existing.name = deploy_data.name
        existing.rate_limit = deploy_data.rate_limit
        existing.updated_at = datetime.utcnow()
        session.add(existing)
        session.commit()
        session.refresh(existing)
        return existing

    # 3. 创建新 Deployment (Snapshot)
    deployment = Deployment(
        name=deploy_data.name,
        description=deploy_data.description,
        flow_id=deploy_data.flow_id,
        endpoint_slug=deploy_data.endpoint_slug,
        rate_limit=deploy_data.rate_limit,
        # 核心：冻结当前 Flow 的 data
        graph_snapshot=flow.data, 
        production_config=deploy_data.production_config or {}
    )
    
    session.add(deployment)
    session.commit()
    session.refresh(deployment)
    return deployment

@router.get("/", response_model=list[DeploymentRead])
def list_deployments(session: Session = Depends(get_session)):
    deployments = session.exec(select(Deployment)).all()
    return deployments

import logging
from typing import Any, Dict, Optional
from fastapi import APIRouter, Depends, HTTPException, Request, Body
from fastapi_limiter.depends import RateLimiter
from sqlmodel import Session, select

from langflow.services.get_service import get_session
from langflow.services.database.models.deployment import Deployment
from langflow.graph.graph.base import Graph
from langflow.processing.process import process_graph_cached # 假设 langflow 内部有类似的构建函数
from langflow.core.tracing import get_langfuse_callback

# 初始化日志
logger = logging.getLogger("agent_server")

router = APIRouter(prefix="/agent", tags=["Agent Server"])

# --- 动态限流器 ---
async def get_rate_limit_key(request: Request):
    """自定义限流 Key：基于 Slug + IP"""
    # 假设 slug 在 path param 中，需要从 request.path_params 获取
    slug = request.path_params.get("slug")
    return f"{slug}:{request.client.host}"

async def custom_limiter(request: Request, session: Session = Depends(get_session)):
    """从数据库动态获取限流规则"""
    slug = request.path_params.get("slug")
    deployment = session.exec(select(Deployment).where(Deployment.endpoint_slug == slug)).first()
    
    if not deployment:
        # 如果找不到部署，暂时放行（后续逻辑会处理 404）
        return "100/second"
    
    return deployment.rate_limit or "60/minute"

# --- 运行接口 ---

@router.post("/{slug}", dependencies=[Depends(RateLimiter(key_func=get_rate_limit_key, limit_func=custom_limiter))])
async def run_agent(
    slug: str,
    inputs: Dict[str, Any] = Body(..., description="Input values for the agent"),
    session_id: Optional[str] = Body(None, description="Session ID for conversation history"),
    db_session: Session = Depends(get_session)
):
    """
    生产级 Agent 运行接口
    """
    # 1. 加载 Deployment
    deployment = db_session.exec(select(Deployment).where(Deployment.endpoint_slug == slug, Deployment.is_active == True)).first()
    if not deployment:
        raise HTTPException(status_code=404, detail="Agent deployment not found or inactive")

    try:
        # 2. 构建 Graph (从快照)
        # 这里的关键是使用 langflow 的 Graph 类将 JSON 转换为 Python 对象
        # 注意：这里我们根据 Langflow 的源码逻辑，使用 Deployment 中的 graph_snapshot
        graph = Graph.from_payload(deployment.graph_snapshot)
        
        # 3. 准备 Langfuse Callback
        langfuse_cb = get_langfuse_callback(
            session_id=session_id,
            trace_name=f"agent-{deployment.name}",
            tags=["production", slug]
        )
        
        callbacks = [langfuse_cb] if langfuse_cb else []

        # 4. 执行 Graph
        # Langflow 的 Graph.process 或 build 方法通常返回 artifacts 或 runnable
        # 我们这里模拟构建 Runnable 并运行的过程
        
        # A. 构建 Runnable
        # process=True 会尝试编译 Graph，如果有 LangGraph 组件，这里会返回 CompiledGraph
        artifacts = graph.build() 
        
        # 找到入口节点或 LangChain 对象进行调用
        # 注意：Langflow 的运行逻辑比较复杂，通常是找到 outputs 并 process
        # 这里简化为：使用 graph 自身的 process 方法，并注入 callback
        
        # *深度修改*: 为了注入 Callback，我们需要 Hack Graph 的执行过程
        # 如果是 LangChain/LangGraph 组件，通常在 .invoke() 时传入 config={'callbacks': ...}
        
        result = await graph.process(
            inputs, 
            session_id=session_id,
            callbacks=callbacks # 这需要修改 Langflow 的 Graph.process 源码来支持 callbacks 参数
        )
        
        return {
            "status": "success",
            "outputs": result,
            "trace_id": langfuse_cb.get_trace_id() if langfuse_cb else None
        }

    except Exception as e:
        logger.error(f"Error running agent {slug}: {str(e)}", exc_info=True)
        # 生产环境不返回具体堆栈
        raise HTTPException(status_code=500, detail=f"Internal Agent Error: {str(e)}")

# 修改前
# async def process(self, inputs, ...):

# 修改后
async def process(
    self, 
    inputs: Dict[str, Any], 
    session_id: str = None, 
    callbacks: list = [] # [新增]
):
    # ... 原有逻辑 ...
    
    for vertex in self.vertices:
        # 在构建每个节点时，如果是 LangChain 组件，尝试注入 callbacks
        # 这通常需要在 vertex.build() 内部处理，或者在这里处理
        
        # 方案：如果是 Runnable，使用 with_config
        if hasattr(vertex._built_result, "with_config") and callbacks:
            vertex._built_result = vertex._built_result.with_config(callbacks=callbacks)
            
    # ... 继续原有执行逻辑 ...

import redis.asyncio as redis
from fastapi_limiter import FastAPILimiter
from langflow.api.v1.deployments import router as deployments_router
from langflow.api.v1.agent_server import router as agent_server_router

# ... 在 create_app 函数中 ...

@app.on_event("startup")
async def startup():
    # 连接 Redis，用于限流
    redis_url = os.getenv("REDIS_URL", "redis://localhost:6379")
    r = redis.from_url(redis_url, encoding="utf-8", decode_responses=True)
    await FastAPILimiter.init(r)

# ... 注册路由 ...
app.include_router(deployments_router)
app.include_router(agent_server_router)

from datetime import datetime
from typing import Optional, Dict, Any
from uuid import UUID, uuid4
from sqlmodel import Field, SQLModel, JSON, Column

class DeploymentBase(SQLModel):
    name: str = Field(index=True)
    description: Optional[str] = Field(default=None)
    endpoint_slug: str = Field(unique=True, index=True)
    rate_limit: Optional[str] = Field(default="60/minute")
    is_active: bool = Field(default=True)

class Deployment(DeploymentBase, table=True):
    id: UUID = Field(default_factory=uuid4, primary_key=True)
    flow_id: UUID = Field(foreign_key="flow.id")
    version: int = Field(default=1) # 新增：用于缓存失效控制
    
    # 存储 Flow 的静态快照
    graph_snapshot: Dict[str, Any] = Field(default={}, sa_column=Column(JSON))
    # 存储生产环境的覆盖配置 (Tweaks)
    production_config: Dict[str, Any] = Field(default={}, sa_column=Column(JSON))
    
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)

class DeploymentCreate(DeploymentBase):
    flow_id: UUID
    production_config: Optional[Dict[str, Any]] = None

class DeploymentRead(DeploymentBase):
    id: UUID
    flow_id: UUID
    version: int
    created_at: datetime

import logging
import os
from typing import Any, Dict, Optional, Tuple
from uuid import UUID

from async_lru import alru_cache
from langchain_core.runnables import Runnable, RunnableConfig
from langfuse.callback import CallbackHandler
from sqlmodel import Session, select

from langflow.graph.graph.base import Graph
from langflow.services.database.models.deployment import Deployment
from langflow.services.get_service import get_session

logger = logging.getLogger("agent_server")

class AgentServerService:
    """
    负责 Agent 的缓存管理、图构建和执行
    """
    
    # --- 1. 缓存层 ---
    # 使用 async_lru 缓存构建好的 Runnable 对象
    # Key: (deployment_id, version)
    # Value: (Runnable, input_key_name)
    @staticmethod
    @alru_cache(maxsize=100)
    async def get_compiled_agent(deployment_id: str, version: int, snapshot_hash: int) -> Tuple[Runnable, str]:
        """
        从数据库加载 Snapshot -> 应用 Config -> 编译为 Runnable -> 缓存
        snapshot_hash 用于确保内容一致性，作为缓存键的一部分
        """
        logger.info(f"Building graph cache for deployment {deployment_id} v{version}")
        
        # 为了在静态方法中获取 DB session，这里需要独立的上下文，或者传入数据
        # 实际生产中建议传入 snapshot 数据而非在 cache 函数内查库以保持纯函数特性
        # 但为了简化调用，这里演示逻辑：
        
        # 注意：由于 alru_cache 限制参数必须可哈希，我们通常在外层查好数据传进来，
        # 或者在这里临时创建 session (不推荐高频)。
        # 最佳实践：将此函数逻辑拆分，外层查 DB，内层只负责构建。
        raise NotImplementedError("Internal implementation detail: See load_and_build below")

    async def load_and_build_graph(self, deployment: Deployment) -> Tuple[Runnable, str]:
        """
        核心构建逻辑：将 Langflow JSON 转换为可执行的 LangChain Runnable
        """
        try:
            # 1. 初始化 Graph
            graph = Graph.from_payload(deployment.graph_snapshot)
            
            # 2. 应用生产环境配置 (Tweaks)
            # production_config 格式应为 { "NodeID": { "param": "value" } }
            if deployment.production_config:
                # Langflow 的 Graph 对象通常有 apply_tweaks 方法或类似的逻辑
                # 这里我们模拟手动覆盖：
                for node in graph.vertices:
                    if node.id in deployment.production_config:
                        for key, value in deployment.production_config[node.id].items():
                            node.params[key] = value
            
            # 3. 构建 (Build)
            # build() 会实例化所有 Python 对象
            built_artifacts = graph.build()
            
            # 4. 寻找入口和出口
            # Langflow 的图可能有多个输出，我们需要找到主要的 Runnable
            # 通常 graph.build() 后，我们可以遍历 vertices 找到 LangChain 的 Runnable
            
            # 简化策略：寻找第一个 ChatOutput 或 Output 节点，或者直接使用 graph 编译后的结果
            # 如果是 LangGraph，通常有一个 compiled graph
            
            runnable = None
            input_key = "input_value" # 默认值
            
            # 尝试从构建结果中提取 Runnable
            # 这里的逻辑依赖 Langflow 具体版本，假设 v1.0+
            for vertex in graph.vertices:
                if vertex.is_input:
                    input_key = vertex.id # 或者是 vertex.display_name
                
                # 如果节点构建出了 Runnable，且是输出节点
                if vertex.is_output and hasattr(vertex._built_result, "invoke"):
                    runnable = vertex._built_result
                    break
            
            # 如果没找到明确的输出，尝试使用 Graph 自身的运行逻辑封装
            if not runnable:
                # Fallback: 某些版本的 Langflow graph.build() 返回的就是 Runnable
                if hasattr(built_artifacts, "invoke"):
                    runnable = built_artifacts
                else:
                    raise ValueError("Could not extract a valid Runnable from the graph")

            return runnable, input_key

        except Exception as e:
            logger.error(f"Failed to build graph for {deployment.id}: {e}")
            raise e

    # --- 2. 内存缓存封装 (解决 alru_cache 参数问题) ---
    _cache = {}

    async def get_agent_runnable(self, deployment: Deployment) -> Tuple[Runnable, str]:
        """
        带有内存缓存的获取方法
        """
        cache_key = f"{deployment.id}_v{deployment.version}"
        
        if cache_key in self._cache:
            return self._cache[cache_key]
        
        # Cache Miss
        runnable, input_key = await self.load_and_build_graph(deployment)
        self._cache[cache_key] = (runnable, input_key)
        return runnable, input_key

    # --- 3. 执行层 ---
    async def execute_agent(
        self, 
        deployment: Deployment, 
        inputs: Dict[str, Any], 
        session_id: Optional[str] = None
    ) -> Dict[str, Any]:
        
        # 1. 获取缓存的可执行对象
        runnable, default_input_key = await self.get_agent_runnable(deployment)
        
        # 2. 处理输入
        # 如果用户传了 {"input_value": "hi"} 则直接用
        # 如果用户传了 {"query": "hi"} 但图需要 "input_value"，这里做简单映射 (可选)
        final_inputs = inputs
        
        # 3. 配置 Langfuse
        callbacks = []
        lf_handler = self._create_langfuse_handler(
            trace_name=f"agent-{deployment.name}",
            session_id=session_id,
            user_id=inputs.get("user_id"),
            tags=[deployment.endpoint_slug, f"v{deployment.version}"]
        )
        if lf_handler:
            callbacks.append(lf_handler)

        # 4. 执行 (Invoke)
        # 使用 LangChain 标准 API，无需修改 Langflow 源码
        config = RunnableConfig(
            callbacks=callbacks,
            configurable={"session_id": session_id} # 传给 Checkpointer
        )
        
        try:
            # 关键：直接调用 Runnable，绕过 Langflow 的 Graph.process 封装
            # 这样我们可以完全控制 callbacks
            if hasattr(runnable, "ainvoke"):
                result = await runnable.ainvoke(final_inputs, config=config)
            else:
                result = runnable.invoke(final_inputs, config=config)
                
            # 5. 格式化输出
            # 如果结果是 Message 对象，转字符串
            if hasattr(result, "content"):
                output_content = result.content
            elif isinstance(result, dict) and "output" in result:
                output_content = result["output"]
            else:
                output_content = result

            return {
                "outputs": output_content,
                "trace_id": lf_handler.get_trace_id() if lf_handler else None
            }
            
        except Exception as e:
            logger.error(f"Execution failed: {e}")
            raise e

    def _create_langfuse_handler(self, **kwargs):
        pk = os.getenv("LANGFUSE_PUBLIC_KEY")
        sk = os.getenv("LANGFUSE_SECRET_KEY")
        if pk and sk:
            return CallbackHandler(public_key=pk, secret_key=sk, host=os.getenv("LANGFUSE_HOST"), **kwargs)
        return None

# 单例模式
agent_service = AgentServerService()

from typing import Any, Dict, Optional
from fastapi import APIRouter, Depends, HTTPException, Request, Body
from fastapi_limiter.depends import RateLimiter
from sqlmodel import Session, select

from langflow.services.get_service import get_session
from langflow.services.database.models.deployment import Deployment
# 引入上面的 Service
from langflow.services.agent_server.service import agent_service

router = APIRouter(prefix="/agent", tags=["Agent Server"])

async def get_deployment_limit(request: Request, session: Session = Depends(get_session)):
    slug = request.path_params.get("slug")
    dep = session.exec(select(Deployment).where(Deployment.endpoint_slug == slug)).first()
    return dep.rate_limit if dep else "10/minute"

async def get_rate_limit_key(request: Request):
    return f"{request.path_params.get('slug')}:{request.client.host}"

@router.post("/{slug}", dependencies=[Depends(RateLimiter(key_func=get_rate_limit_key, limit_func=get_deployment_limit))])
async def run_agent_endpoint(
    slug: str,
    inputs: Dict[str, Any] = Body(..., example={"input_value": "Hello Agent"}),
    session_id: Optional[str] = Body(None),
    db: Session = Depends(get_session)
):
    """
    生产环境 Agent 执行入口
    """
    # 1. 查库
    deployment = db.exec(select(Deployment).where(Deployment.endpoint_slug == slug, Deployment.is_active == True)).first()
    
    if not deployment:
        raise HTTPException(status_code=404, detail="Agent not found")

    try:
        # 2. 委托给 Service 执行
        # 所有的构建、缓存、Tracing 逻辑都在 Service 内部处理
        result = await agent_service.execute_agent(
            deployment=deployment,
            inputs=inputs,
            session_id=session_id
        )
        return result

    except Exception as e:
        # 生产环境隐藏具体堆栈
        raise HTTPException(status_code=500, detail=f"Agent execution error: {str(e)}")

from datetime import datetime
from fastapi import APIRouter, Depends, HTTPException
from sqlmodel import Session, select
from langflow.services.database.models.flow import Flow
from langflow.services.database.models.deployment import Deployment, DeploymentCreate, DeploymentRead
from langflow.services.get_service import get_session

router = APIRouter(prefix="/deployments", tags=["Deployments"])

@router.post("/publish", response_model=DeploymentRead)
def publish_flow(data: DeploymentCreate, session: Session = Depends(get_session)):
    flow = session.get(Flow, data.flow_id)
    if not flow:
        raise HTTPException(status_code=404, detail="Flow not found")

    # 检查是否已存在 Slug
    existing = session.exec(select(Deployment).where(Deployment.endpoint_slug == data.endpoint_slug)).first()
    
    if existing:
        # 更新模式：版本号 + 1，这将自动触发缓存失效（因为缓存Key包含版本号）
        existing.version += 1
        existing.graph_snapshot = flow.data.copy() # Deep copy recommended
        existing.name = data.name
        existing.updated_at = datetime.utcnow()
        if data.production_config:
            existing.production_config = data.production_config
        
        session.add(existing)
        session.commit()
        session.refresh(existing)
        return existing
    
    # 新建模式
    new_deployment = Deployment(
        name=data.name,
        flow_id=data.flow_id,
        endpoint_slug=data.endpoint_slug,
        graph_snapshot=flow.data.copy(),
        production_config=data.production_config or {},
        version=1
    )
    session.add(new_deployment)
    session.commit()
    session.refresh(new_deployment)
    return new_deployment

# ... 现有 imports ...
import redis.asyncio as redis
from fastapi_limiter import FastAPILimiter
from langflow.api.v1.deployments import router as dep_router
from langflow.api.v1.agent_server import router as agent_router

# ... inside create_app() ...

    @app.on_event("startup")
    async def startup():
        # 初始化 Redis 连接池
        redis_url = os.getenv("REDIS_URL", "redis://localhost:6379")
        r = redis.from_url(redis_url, encoding="utf-8", decode_responses=True)
        await FastAPILimiter.init(r)
        
    # 注册路由
    app.include_router(dep_router)
    app.include_router(agent_router)