上下文工程

# !pip install ipynbname

import os
import uuid
import sqlite3

from typing import Callable
from dotenv import load_dotenv
from dataclasses import dataclass
from langchain_openai import ChatOpenAI
from langchain.tools import tool, ToolRuntime
from langchain.agents import create_agent
from langchain.agents.middleware import dynamic_prompt, wrap_model_call, ModelRequest, ModelResponse, SummarizationMiddleware
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.store.memory import InMemoryStore
from langgraph.store.sqlite import SqliteStore

# 加载模型配置
_ = load_dotenv()

# 加载模型
llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url=os.getenv("DASHSCOPE_BASE_URL"),
    model="qwen3-coder-plus",
    temperature=0.7,
)

@dynamic_prompt
def state_aware_prompt(request: ModelRequest) -> str:
    # request.messages is a shortcut for request.state["messages"]
    message_count = len(request.messages)

    base = "You are a helpful assistant."

    if message_count > 6:
        base += "\nThis is a long conversation - be extra concise."

    # 临时打印base看效果
    print(base)

    return base

agent = create_agent(
    model=llm,
    middleware=[state_aware_prompt]
)

result = agent.invoke(
    {"messages": [
        {"role": "user", "content": "广州今天的天气怎么样？"},
        {"role": "assistant", "content": "广州天气很好"},
        {"role": "user", "content": "吃点什么好呢"},
        {"role": "assistant", "content": "要不要吃香茅鳗鱼煲"},
        {"role": "user", "content": "香茅是什么"},
        {"role": "assistant", "content": "香茅又名柠檬草，常见于泰式冬阴功汤、越南烤肉"},
        {"role": "user", "content": "auv 那还等什么，咱吃去吧"},
    ]},
)

for message in result['messages']:
    message.pretty_print()

@dataclass
class Context:
    user_id: str

@dynamic_prompt
def store_aware_prompt(request: ModelRequest) -> str:
    user_id = request.runtime.context.user_id

    # Read from Store: get user preferences
    store = request.runtime.store
    user_prefs = store.get(("preferences",), user_id)

    base = "You are a helpful assistant."

    if user_prefs:
        style = user_prefs.value.get("communication_style", "balanced")
        base += f"\nUser prefers {style} responses."

    return base

store = InMemoryStore()

agent = create_agent(
    model=llm,
    middleware=[store_aware_prompt],
    context_schema=Context,
    store=store,
)

# 预置两条偏好信息
store.put(("preferences",), "user_1", {"communication_style": "Chinese"})
store.put(("preferences",), "user_2", {"communication_style": "Korean"})

# 用户1喜欢中文回复
result = agent.invoke(
    {"messages": [
        {"role": "system", "content": "You are a helpful assistant. Please be extra concise."},
        {"role": "user", "content": 'What is a "hold short line"?'}
    ]},
    context=Context(user_id="user_1"),
)

for message in result['messages']:
    message.pretty_print()

# 用户2喜欢韩文回复
result = agent.invoke(
    {"messages": [
        {"role": "system", "content": "You are a helpful assistant. Please be extra concise."},
        {"role": "user", "content": 'What is a "hold short line"?'}
    ]},
    context=Context(user_id="user_2"),
)

for message in result['messages']:
    message.pretty_print()

@dataclass
class Context:
    user_role: str
    deployment_env: str

@dynamic_prompt
def context_aware_prompt(request: ModelRequest) -> str:
    # Read from Runtime Context: user role and environment
    user_role = request.runtime.context.user_role
    env = request.runtime.context.deployment_env

    base = "You are a helpful assistant."

    if user_role == "admin":
        base += "\nYou can use the get_weather tool."
    else:
        base += "\nYou are prohibited from using the get_weather tool."

    if env == "production":
        base += "\nBe extra careful with any data modifications."

    return base

@tool
def get_weather(city: str) -> str:
    """Get weather for a given city."""
    return f"It's always sunny in {city}!"

agent = create_agent(
    model=llm,
    tools=[get_weather],
    middleware=[context_aware_prompt],
    context_schema=Context,
    checkpointer=InMemorySaver(),
)

# 利用 Runtime 中的两个变量，动态控制 System prompt
# 将 user_role 设为 admin，允许使用天气查询工具
config = {'configurable': {'thread_id': str(uuid.uuid4())}}
result = agent.invoke(
    {"messages": [{"role": "user", "content": "广州今天的天气怎么样？"}]},
    context=Context(user_role="admin", deployment_env="production"),
    config=config,
)

for message in result['messages']:
    message.pretty_print()

# 若将 user_role 改为 viewer，则无法使用天气查询工具
config = {'configurable': {'thread_id': str(uuid.uuid4())}}
result = agent.invoke(
    {"messages": [{"role": "user", "content": "广州今天的天气怎么样？"}]},
    context=Context(user_role="viewer", deployment_env="production"),
    config=config,
)

for message in result['messages']:
    message.pretty_print()

result['messages']

@dataclass
class FileContext:
    uploaded_files: list[dict]

@wrap_model_call
def inject_file_context(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse]
) -> ModelResponse:
    """Inject context about files user has uploaded this session."""
    uploaded_files = request.runtime.context.uploaded_files

    try:
        base_dir = os.path.dirname(os.path.abspath(__file__))
    except Exception as e:
        import ipynbname
        import os
        notebook_path = ipynbname.path()
        base_dir = os.path.dirname(notebook_path)

    file_sections = []
    for file in uploaded_files:
        name, ftype = "", ""
        path = file.get("path")
        if path:
            base_filename = os.path.basename(path)
            stem, ext = os.path.splitext(base_filename)
            name = stem or base_filename
            ftype = (ext.lstrip(".") if ext else None)

            # 构建文件描述内容
            content_list = [f"名称: {name}"]
            if ftype:
                content_list.append(f"类型: {ftype}")

            # 解析相对路径为绝对路径
            abs_path = path if os.path.isabs(path) else os.path.join(base_dir, path)

            # 读取文件内容
            content_block = ""
            if abs_path and os.path.exists(abs_path):
                try:
                    with open(abs_path, "r", encoding="utf-8") as f:
                        content_block = f.read()
                except Exception as e:
                    content_block = f"[读取文件错误 '{abs_path}': {e}]"
            else:
                content_block = "[文件路径缺失或未找到]"

            section = (
                f"---\n"
                f"{chr(10).join(content_list)}\n\n"
                f"{content_block}\n"
                f"---"
            )
            file_sections.append(section)

        file_context = (
            "已加载的会话文件：\n"
            f"{chr(10).join(file_sections)}"
            "\n回答问题时请参考这些文件。"
        )

        # Inject file context before recent messages
        messages = [  
            *request.messages,
            {"role": "user", "content": file_context},
        ]
        request = request.override(messages=messages)  

    return handler(request)

agent = create_agent(
    model=llm,
    middleware=[inject_file_context],
    context_schema=FileContext,
)

result = agent.invoke(
    {
        "messages": [{
            "role": "user",
            "content": "关于上海地铁的无脸乘客，有什么需要注意的？",
        }],
    },
    context=FileContext(uploaded_files=[{"path": "./docs/rule_horror.md"}]),
)

for message in result['messages']:
    message.pretty_print()

# 删除SQLite数据库
if os.path.exists("user-info.db"):
    os.remove("user-info.db")

# 创建SQLite存储
conn = sqlite3.connect("user-info.db", check_same_thread=False, isolation_level=None)
conn.execute("PRAGMA journal_mode=WAL;")
conn.execute("PRAGMA busy_timeout = 30000;")

store = SqliteStore(conn)

# 预置两条用户信息
store.put(("user_info",), "柳如烟", {"description": "清冷才女，身怀绝技，为寻身世之谜踏入江湖。", "birthplace": "吴兴县"})
store.put(("user_info",), "苏慕白", {"description": "孤傲剑客，剑法超群，背负家族血仇，隐于市井追寻真相。", "birthplace": "杭县"})

@tool
def fetch_user_data(
    user_id: str,
    runtime: ToolRuntime
) -> str:
    """
    Fetch user information from the in-memory store.

    :param user_id: The unique identifier of the user.
    :param runtime: The tool runtime context injected by the framework.
    :return: The user's description string if found; an empty string otherwise.
    """
    store = runtime.store
    user_info = store.get(("user_info",), user_id)

    user_desc = ""
    if user_info:
        user_desc = user_info.value.get("description", "")

    return user_desc

agent = create_agent(
    model=llm,
    tools=[fetch_user_data],
    store=store,
)

result = agent.invoke({
    "messages": [{
        "role": "user",
        "content": "五分钟之内，我要柳如烟的全部信息"
    }]
})

for message in result['messages']:
    message.pretty_print()

@dataclass
class Context:
    key: str

@tool
def fetch_user_data(
    user_id: str,
    runtime: ToolRuntime[Context]
) -> str:
    """
    Fetch user information from the in-memory store.

    :param user_id: The unique identifier of the user.
    :param runtime: The tool runtime context injected by the framework.
    :return: The user's description string if found; an empty string otherwise.
    """
    key = runtime.context.key

    store = runtime.store
    user_info = store.get(("user_info",), user_id)

    user_desc = ""
    if user_info:
        user_desc = user_info.value.get(key, "")

    return f"{key}: {user_desc}"

agent = create_agent(
    model=llm,
    tools=[fetch_user_data],
    store=store,
)

result = agent.invoke(
    {"messages": [{"role": "user", "content": "五分钟之内，我要柳如烟的全部信息"}]},
    context=Context(key="birthplace"),
)

for message in result['messages']:
    message.pretty_print()

# 创建短期记忆
checkpointer = InMemorySaver()

# 创建带内置摘要中间件的Agent
# 为了让配置能在我们的例子里生效，这里的触发值设得很小
agent = create_agent(
    model=llm,
    middleware=[
        SummarizationMiddleware(
            model=llm,
            max_tokens_before_summary=40,  # Trigger summarization at 40 tokens
            messages_to_keep=1,  # Keep last 1 messages after summary
        ),
    ],
)

result = agent.invoke(
    {"messages": [
        {"role": "user", "content": "广州今天的天气怎么样？"},
        {"role": "assistant", "content": "广州天气很好"},
        {"role": "user", "content": "吃点什么好呢"},
        {"role": "assistant", "content": "要不要吃香茅鳗鱼煲"},
        {"role": "user", "content": "香茅是什么"},
        {"role": "assistant", "content": "香茅又名柠檬草，常见于泰式冬阴功汤、越南烤肉"},
        {"role": "user", "content": "auv 那还等什么，咱吃去吧"},
    ]},
    checkpointer=checkpointer,
)

for message in result['messages']:
    message.pretty_print()