from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
# 1.加载业务文档文本
biz_doc = """
业务审批规则：金额小于1万由部门主管审批，1万到10万需要经理审批，大于10万提交总监审批。
审批提交时间工作日9点‑18点，非工作时间不予受理。
"""
# 2.文档切片
splitter = RecursiveCharacterTextSplitter(
    chunk_size=300,
    chunk_overlap=50
)
chunks = splitter.split_text(biz_doc)
# 3.向量化存入向量库
embedding = OpenAIEmbeddings()
vector_db = Chroma.from_texts(chunks, embedding, persist_directory="./chroma_biz")
# 4.检索，拿到参考知识片段
query = "5万金额审批找谁"
docs = vector_db.similarity_search(query, k=2)
ref_context = "\n".join([d.page_content for d in docs])
# 组装送入大模型prompt
prompt = f"""参考下面资料回答问题，禁止编造内容：
【参考资料】
{ref_context}
【用户问题】{query}
"""
print(prompt)

import json
from typing import List
# 模拟Redis会话存储，生产替换真实Redis
session_store = {}
class SessionMemory:
    def __init__(self, session_id: str, max_round:int=6):
        self.session_id = session_id
        self.max_round = max_round
        if session_id not in session_store:
            session_store[session_id] = []
    
    def add_msg(self, role:str, content:str):
        msgs = session_store[self.session_id]
        msgs.append({"role":role, "content":content})
        # 裁剪，只保留最近max_round轮，防止上下文膨胀
        if len(msgs) > self.max_round:
            session_store[self.session_id] = msgs[-self.max_round:]
    
    def get_history(self) -> List[dict]:
        return session_store.get(self.session_id, [])
# 使用示例
mem = SessionMemory(session_id="sess_001", max_round=4)
mem.add_msg("user","我要提交5万的审批")
mem.add_msg("assistant","5万需要经理审批")
mem.add_msg("user","经理不在怎么办")
history = mem.get_history()
print(json.dumps(history, ensure_ascii=False, indent=2))

import requests
from pydantic import BaseModel, Field
# 工具入参模型，做参数校验
class SubmitApprovalTool(BaseModel):
    amount: float = Field(gt=0, description="审批金额，大于0")
    reason: str = Field(min_length=2, description="审批原因")
    user_id: str
# 模拟应用网关：校验+鉴权+幂等，模型输出的参数先过这一层
def tool_gateway(tool_name:str, raw_params:dict, login_user:str, request_id:str):
    # 1.参数校验
    if tool_name == "submit_approval":
        valid_param = SubmitApprovalTool(**raw_params)
        # 2.鉴权，最小权限校验，示例：只能操作自己的单据
        if valid_param.user_id != login_user:
            raise PermissionError("无权操作他人审批单据")
        # 3.幂等：request_id作为幂等键，避免重复提交
        headers = {"Idempotency‑Key": request_id}
        resp = requests.post(
            url="http://biz‑api/submit_approval",
            json=valid_param.model_dump(),
            headers=headers,
            timeout=10
        )
        return resp.json()
    else:
        raise ValueError("不存在该工具")
# 模拟模型输出原始参数
model_output_params = {"amount":50000,"reason":"项目采购","user_id":"u1001"}
try:
    result = tool_gateway(
        tool_name="submit_approval",
        raw_params=model_output_params,
        login_user="u1001",
        request_id="req_abc123"
    )
    print("工具调用结果：", result)
except Exception as e:
    print("调用拦截：", str(e))

import uuid
import json
import logging
from typing import Optional, Dict, Any, List
from pydantic import BaseModel, Field
# --------------------------
# 0. 初始化日志（审计日志，自动脱敏）
# --------------------------
audit_logger = logging.getLogger("llm_audit")
audit_logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
audit_logger.addHandler(handler)
def audit_record(request_id: str,
                 user_id: str,
                 session_id: str,
                 user_query: str,
                 rag_refs: List[str],
                 prompt_snippet: str,
                 tool_call_raw: Optional[Dict],
                 tool_valid_params: Optional[Dict],
                 tool_result: Optional[Dict],
                 llm_reply: str,
                 is_blocked: bool = False,
                 block_reason: str = "",
                 error: str = ""):
    """审计日志：敏感字段简单脱敏"""
    record = {
        "request_id": request_id,
        "user_id": user_id,
        "session_id": session_id,
        "user_query": user_query,
        "rag_refs": rag_refs,
        "prompt_snippet": prompt_snippet[:500],  # 截断，不存超长上下文
        "tool_call_raw": tool_call_raw,
        "tool_valid_params": tool_valid_params,
        "tool_result": tool_result,
        "llm_reply": llm_reply,
        "is_blocked": is_blocked,
        "block_reason": block_reason,
        "error": error
    }
    audit_logger.info(json.dumps(record, ensure_ascii=False))
# --------------------------
# 1. 数据模型：工具入参定义(Pydantic校验)
# --------------------------
class SubmitApprovalTool(BaseModel):
    """提交审批工具参数"""
    amount: float = Field(gt=0, description="审批金额必须大于0")
    reason: str = Field(min_length=2, max_length=200, description="审批原因")
    target_user_id: str = Field(description="操作单据所属用户ID")
# --------------------------
# 2. 模拟存储层：会话记忆、长期记忆、幂等记录表
# --------------------------
mock_session_store: Dict[str, List[Dict[str, str]]] = {}   # session_id -> message list
mock_long_memory: Dict[str, List[str]] = {}                 # user_id -> 用户长期记忆事实
mock_idempotent_table: Dict[str, bool] = {}                 # request_id幂等表
mock_biz_api_db: List[Dict] = []                            # 模拟业务后端存储
# --------------------------
# 3. 模拟底层能力：记忆、RAG、大模型、MCP/API调用
# --------------------------
def load_session_memory(session_id: str) -> List[Dict[str, str]]:
    """读取会话短期记忆"""
    return mock_session_store.get(session_id, [])
def load_long_memory(user_id: str) -> List[str]:
    """加载用户长期记忆事实"""
    return mock_long_memory.get(user_id, [])
def mock_rag_retrieve(query: str) -> List[str]:
    """模拟RAG检索，返回参考知识片段"""
    if "审批" in query:
        return [
            "业务审批规则：金额小于1万部门主管审批；1万‑10万经理审批；大于10万总监审批。",
            "审批仅工作日9‑18点受理。"
        ]
    return []
def truncate_context(messages: List[Dict], max_msg_count: int = 8) -> List[Dict]:
    """上下文截断：只保留最近N轮消息，防止token爆炸"""
    if len(messages) <= max_msg_count:
        return messages
    return messages[-max_msg_count:]
def mock_llm_infer(prompt_messages: List[Dict]) -> Dict[str, Any]:
    """模拟大模型推理：两种返回，普通回答 / tool_call工具调用"""
    last_user_msg = prompt_messages[-1]["content"]
    if "提交审批" in last_user_msg:
        # 模拟模型输出工具调用
        return {
            "reply_content": "",
            "tool_call": {
                "tool_name": "submit_approval",
                "raw_params": {"amount": 50000.0, "reason": "项目物料采购", "target_user_id": "u_001"}
            }
        }
    else:
        return {
            "reply_content": "已收到你的咨询，审批相关规则已参考制度说明。",
            "tool_call": None
        }
def mock_biz_api_submit_approval(params: Dict, idempotency_key: str) -> Dict:
    """模拟业务API，自带幂等判断"""
    global mock_biz_api_db
    if idempotency_key in mock_idempotent_table:
        return {"code": 0, "msg": "幂等命中，已处理过", "data": None}
    mock_idempotent_table[idempotency_key] = True
    mock_biz_api_db.append(params)
    return {"code":0, "msg":"审批单据提交成功", "data":{"approval_id":"ap_10086"}}
# --------------------------
# 4. 核心网关：鉴权、最小权限校验
# --------------------------
def check_permission(user_id: str, tool_name: str, parsed_params: BaseModel) -> tuple[bool, str]:
    """
    鉴权&最小权限校验
    return (是否通过, 拒绝原因)
    """
    if tool_name == "submit_approval":
        p = parsed_params.model_dump()
        # 业务规则：用户只能操作属于自己的单据 target_user_id == 当前登录user_id
        if p["target_user_id"] != user_id:
            return False, f"权限拒绝：用户{user_id}不允许操作target_user_id={p['target_user_id']}的单据"
        return True, ""
    return False, f"不支持工具:{tool_name}"
# --------------------------
# 5. 主流程函数：完整复现你伪代码逻辑
# --------------------------
def llm_agent_entry(user_id: str, session_id: str, user_query: str) -> Dict[str, Any]:
    """
    完整链路：
    接收用户请求
    生成全局request_id
    读取会话记忆、加载长期记忆
    执行RAG检索，获取参考知识片段
    组装Prompt上下文，做长度截断
    传给大模型推理
    if 模型返回工具调用：
        拿到模型输出raw_params
        【参数校验】Pydantic校验，失败返回错误
        【鉴权最小权限】校验用户是否允许调用该工具，资源是否归属用户，不通过直接拦截
        【幂等】带上request_id幂等键
        执行工具调用（API/MCP）
    记录完整审计日志（脱敏）
    把工具返回结果放回上下文
    大模型整合结果输出回答
    """
    request_id = str(uuid.uuid4())
    # 1.读取记忆
    session_msgs = load_session_memory(session_id)
    long_facts = load_long_memory(user_id)
    # 2.RAG检索
    rag_refs = mock_rag_retrieve(user_query)
    # 3.组装prompt上下文
    system_parts = []
    if long_facts:
        system_parts.append(f"【用户记忆】：{'；'.join(long_facts)}")
    if rag_refs:
        system_parts.append(f"【参考知识库】：{'\n'.join(rag_refs)}")
    system_msg = {"role":"system", "content":"\n".join(system_parts)}
    prompt_messages = [system_msg] + session_msgs + [{"role":"user", "content": user_query}]
    prompt_messages = truncate_context(prompt_messages)
    prompt_snippet = json.dumps(prompt_messages, ensure_ascii=False)
    # 4.调用大模型推理
    llm_out = mock_llm_infer(prompt_messages)
    tool_call = llm_out.get("tool_call")
    tool_raw_params: Optional[Dict] = None
    valid_parsed_obj: Optional[BaseModel] = None
    tool_response: Optional[Dict] = None
    blocked = False
    block_msg = ""
    error_msg = ""
    # ========== 工具调用分支 ==========
    if tool_call is not None:
        tool_name = tool_call["tool_name"]
        tool_raw_params = tool_call["raw_params"]
        try:
            # 【参数校验 Pydantic】
            if tool_name == "submit_approval":
                valid_parsed_obj = SubmitApprovalTool(**tool_raw_params)
            else:
                raise Exception(f"未知工具 {tool_name}")
            # 【鉴权、最小权限校验】
            permit_ok, reason = check_permission(user_id, tool_name, valid_parsed_obj)
            if not permit_ok:
                blocked = True
                block_msg = reason
            else:
                # 【幂等：request_id作为幂等key】执行业务调用
                if tool_name == "submit_approval":
                    tool_response = mock_biz_api_submit_approval(
                        params=valid_parsed_obj.model_dump(),
                        idempotency_key=request_id
                    )
                # 将工具结果追加到上下文，给大模型二次整合
                prompt_messages.append({
                    "role": "tool",
                    "content": json.dumps(tool_response, ensure_ascii=False)
                })
                # 模拟大模型拿到工具结果后再生成最终回复
                final_llm = mock_llm_infer(prompt_messages)
                llm_out["reply_content"] = final_llm["reply_content"] or json.dumps(tool_response)
        except Exception as e:
            blocked = True
            error_msg = str(e)
    final_answer = llm_out.get("reply_content", "")
    # 记录审计日志（脱敏）
    audit_record(
        request_id=request_id,
        user_id=user_id,
        session_id=session_id,
        user_query=user_query,
        rag_refs=rag_refs,
        prompt_snippet=prompt_snippet,
        tool_call_raw=tool_raw_params,
        tool_valid_params=valid_parsed_obj.model_dump() if valid_parsed_obj else None,
        tool_result=tool_response,
        llm_reply=final_answer,
        is_blocked=blocked,
        block_reason=block_msg,
        error=error_msg
    )
    # 更新会话记忆
    if session_id not in mock_session_store:
        mock_session_store[session_id] = []
    mock_session_store[session_id].append({"role":"user", "content": user_query})
    mock_session_store[session_id].append({"role":"assistant", "content": final_answer})
    return {
        "request_id": request_id,
        "final_answer": final_answer,
        "blocked": blocked,
        "block_reason": block_msg,
        "error": error_msg,
        "tool_response": tool_response
    }
# --------------------------
# 测试运行示例
# --------------------------
if __name__ == "__main__":
    # 测试用例1：正常提交审批，身份匹配
    res1 = llm_agent_entry(user_id="u_001", session_id="sess_001", user_query="帮我提交审批，50000，项目物料采购")
    print(json.dumps(res1, ensure_ascii=False, indent=2))
    print("-"*60)
    # 测试用例2：越权场景，user_id u_002尝试操作u_001单据，会被鉴权拦截
    res2 = llm_agent_entry(user_id="u_002", session_id="sess_002", user_query="帮我提交审批，50000，项目物料采购")
    print(json.dumps(res2, ensure_ascii=False, indent=2))