Create app.py

app.py

@@ -0,0 +1,435 @@
+import os
+import time
+import json
+import hashlib
+import sqlite3
+import logging
+from datetime import datetime
+from typing import List, Dict, Any
+
+import streamlit as st
+import boto3
+from botocore.exceptions import ClientError
+from botocore.config import Config
+
+# ========== 配置和常量 ==========
+DB_FILE = 'audit_logs.db'
+MODEL_ID = os.getenv("BEDROCK_MODEL_ID", "anthropic.claude-3-sonnet-20240229-v1:0")
+DEFAULT_REGION = os.getenv("AWS_REGION", "us-east-1")
+
+# ========== 日志配置 ==========
+logger = logging.getLogger("audit_logger")
+logger.setLevel(logging.INFO)
+if not logger.handlers:
+    console = logging.StreamHandler()
+    console.setLevel(logging.INFO)
+    fmt = logging.Formatter("%(asctime)s %(levelname)-8s %(message)s",
+                            datefmt="%Y-%m-%d %H:%M:%S")
+    console.setFormatter(fmt)
+    logger.addHandler(console)
+
+# ========== 数据库操作 ==========
+def init_db() -> None:
+    """初始化或迁移 audit_logs 表"""
+    conn = sqlite3.connect(DB_FILE)
+    cur = conn.cursor()
+
+    # 创建基础表
+    cur.execute('''
+        CREATE TABLE IF NOT EXISTS audit_logs (
+            id        INTEGER PRIMARY KEY AUTOINCREMENT,
+            timestamp TEXT    NOT NULL,
+            user_id   TEXT    NOT NULL,
+            prompt    TEXT    NOT NULL,
+            response  TEXT    NOT NULL
+        )
+    ''')
+
+    # 检查现有列
+    cur.execute("PRAGMA table_info(audit_logs)")
+    cols = {row[1] for row in cur.fetchall()}
+
+    # 添加缺失的列
+    if 'latency_ms' not in cols:
+        cur.execute("ALTER TABLE audit_logs ADD COLUMN latency_ms REAL NOT NULL DEFAULT 0")
+    if 'success' not in cols:
+        cur.execute("ALTER TABLE audit_logs ADD COLUMN success   INTEGER NOT NULL DEFAULT 1")
+
+    conn.commit()
+    conn.close()
+
+def record_log(
+    user_id: str,
+    prompt: str,
+    response: str,
+    latency_ms: float,
+    success: bool
+) -> None:
+    """存储审计条目并打印到控制台"""
+    timestamp = datetime.utcnow().isoformat()
+    success_flag = 1 if success else 0
+
+    # 1) 持久化到 SQLite
+    with sqlite3.connect(DB_FILE) as conn:
+        conn.execute('''
+            INSERT INTO audit_logs (
+                timestamp, user_id, prompt, response, latency_ms, success
+            ) VALUES (?, ?, ?, ?, ?, ?)
+        ''', (timestamp, user_id, prompt, response, latency_ms, success_flag))
+        conn.commit()
+
+    # 2) 打印到控制台
+    logger.info(
+        f"[Audit] user={user_id!r} prompt={prompt!r} "
+        f"latency={latency_ms:.1f}ms success={bool(success_flag)}"
+    )
+
+def get_logs(limit: int = 100) -> List[Dict[str, Any]]:
+    """获取最近的审计日志条目"""
+    try:
+        with sqlite3.connect(DB_FILE) as conn:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT user_id, prompt, response, latency_ms, success, timestamp
+                FROM audit_logs
+                ORDER BY id DESC
+                LIMIT ?
+            ''', (limit,))
+            rows = cursor.fetchall()
+    except sqlite3.Error as e:
+        logger.error(f"DB error: {e}")
+        return []
+
+    return [
+        {
+            "user_id": row[0],
+            "prompt": row[1],
+            "response": row[2],
+            "latency_ms": row[3],
+            "success": bool(row[4]),
+            "timestamp": row[5]
+        }
+        for row in rows
+    ]
+
+# ========== AWS 认证 ==========
+def authenticate_aws(access_key: str, secret_key: str, region: str = 'us-east-1') -> bool:
+    """验证 AWS 凭证"""
+    try:
+        client = boto3.client(
+            'bedrock-runtime',
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+            region_name=region
+        )
+        return True
+    except ClientError:
+        return False
+
+# ========== Bedrock 客户端 ==========
+def invoke_chat_model(
+    access_key: str,
+    secret_key: str,
+    region: str,
+    model_id: str,
+    messages: list,
+    anthropic_version: str = "2023-05-31",
+    max_tokens_to_sample: int = 512,
+    temperature: float = 0.7
+) -> str:
+    """调用 Anthropic 聊天模型"""
+    # 分离系统提示并准备消息列表
+    system_prompt = None
+    chat_messages = []
+    for m in messages:
+        role = m.get("role")
+        content = m.get("content")
+        if role == "system":
+            system_prompt = content
+        else:
+            chat_messages.append({"role": role, "content": content})
+
+    # 确保第一条消息来自用户
+    while chat_messages and chat_messages[0].get("role") != "user":
+        chat_messages.pop(0)
+
+    if not chat_messages or chat_messages[0].get("role") != "user":
+        raise RuntimeError("No user message found in the conversation.")
+
+    # 构建 Bedrock 请求体
+    ver = anthropic_version
+    if not ver.startswith("bedrock-"):
+        ver = f"bedrock-{ver}"
+
+    body = {
+        "anthropic_version": ver,
+        "max_tokens": max_tokens_to_sample,
+        "temperature": temperature,
+        "messages": chat_messages,
+    }
+    if system_prompt is not None:
+        body["system"] = system_prompt
+
+    body_bytes = json.dumps(body).encode("utf-8")
+
+    # 创建 Bedrock Runtime 客户端
+    try:
+        client = boto3.client(
+            "bedrock-runtime",
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+            region_name=region,
+            config=Config(
+                retries={"max_attempts": 1},
+                read_timeout=60,
+                connect_timeout=5
+            )
+        )
+    except ClientError as e:
+        raise RuntimeError(f"Anthropic invocation failed: {e}")
+
+    try:
+        # 调用模型
+        resp = client.invoke_model(
+            modelId=model_id,
+            contentType="application/json",
+            accept="application/json",
+            body=body_bytes
+        )
+
+        # 读取和解析响应
+        raw = resp["body"].read().decode("utf-8")
+        data = json.loads(raw)
+
+        # Legacy completions API
+        if "completions" in data:
+            return data["completions"][0]["completion"]
+        if "completion" in data:
+            return data["completion"]
+
+        # Messages API → top-level "content" blocks
+        if "content" in data and isinstance(data["content"], list):
+            return "".join(
+                block.get("text", "")
+                for block in data["content"]
+                if block.get("type") == "text"
+            )
+
+        # Alternative: top-level "messages" list
+        if "messages" in data and isinstance(data["messages"], list):
+            for msg in data["messages"]:
+                if msg.get("role") == "assistant":
+                    c = msg.get("content")
+                    if isinstance(c, list):
+                        return "".join(
+                            block.get("text", "")
+                            for block in c
+                            if block.get("type") == "text"
+                        )
+                    elif isinstance(c, str):
+                        return c
+            first = data["messages"][0].get("content")
+            return first if isinstance(first, str) else ""
+
+        raise RuntimeError(f"Unrecognized Anthropic response shape: {data}")
+
+    except ClientError as e:
+        raise RuntimeError(f"Anthropic invocation failed: {e}")
+
+# ========== 聊天功能 ==========
+def chat_with_bedrock(
+    user_id: str,
+    access_key: str,
+    secret_key: str,
+    region: str,
+    messages: list
+) -> Dict[str, Any]:
+    """与 Bedrock 聊天并记录日志"""
+    # 验证凭证
+    if not authenticate_aws(access_key, secret_key, region):
+        return {"success": False, "response": "Invalid AWS credentials"}
+
+    # 调用模型并测量性能
+    start = time.perf_counter()
+    try:
+        reply = invoke_chat_model(
+            access_key,
+            secret_key,
+            region,
+            MODEL_ID,
+            messages
+        )
+        success = True
+    except Exception as e:
+        reply = str(e)
+        success = False
+    latency_ms = (time.perf_counter() - start) * 1000
+
+    # 审计日志
+    last_prompt = messages[-1].get("content", "") if messages else ""
+    record_log(
+        user_id=user_id,
+        prompt=last_prompt,
+        response=reply,
+        latency_ms=latency_ms,
+        success=success
+    )
+
+    return {"success": success, "response": reply}
+
+# ========== 前端功能 ==========
+def get_user_id(access_key: str) -> str:
+    """通过哈希 AWS access key 计算稳定的用户标识符"""
+    return hashlib.sha256(access_key.encode("utf-8")).hexdigest()
+
+def login_sidebar() -> None:
+    """渲染侧边栏登录表单"""
+    st.sidebar.header("🔐 AWS 认证")
+    access_key = st.sidebar.text_input("AWS Access Key ID", type="password")
+    secret_key = st.sidebar.text_input("AWS Secret Access Key", type="password")
+    region = st.sidebar.text_input(
+        "AWS Region", value=DEFAULT_REGION
+    )
+
+    if st.sidebar.button("登录"):
+        if not access_key or not secret_key:
+            st.sidebar.error("请输入 AWS 凭证")
+            return
+
+        # 准备测试请求以验证凭证
+        payload = {
+            "user_id": get_user_id(access_key),
+            "access_key": access_key,
+            "secret_key": secret_key,
+            "region": region,
+            "messages": [
+                {"role": "system", "content": "Authenticate"},
+                {"role": "user", "content": "Ping"}
+            ]
+        }
+        
+        result = chat_with_bedrock(**payload)
+        if result["success"]:
+            # 存储凭证供聊天使用
+            st.session_state.authenticated = True
+            st.session_state.access_key = access_key
+            st.session_state.secret_key = secret_key
+            st.session_state.region = region
+            st.sidebar.success("登录成功！")
+            st.rerun()
+        else:
+            st.sidebar.error("认证失败：请检查您的密钥")
+
+def send_message_and_get_reply(user_text: str) -> str:
+    """发送消息并获取回复"""
+    # 添加新的用户消息
+    st.session_state.messages.append({"role": "user", "content": user_text})
+
+    # 构建包含所有消息的负载
+    payload = {
+        "user_id": get_user_id(st.session_state.access_key),
+        "access_key": st.session_state.access_key,
+        "secret_key": st.session_state.secret_key,
+        "region": st.session_state.region,
+        "messages": st.session_state.messages,
+    }
+
+    # 调用后端
+    result = chat_with_bedrock(**payload)
+    if not result["success"]:
+        raise Exception(result["response"])
+    
+    return result["response"]
+
+def render_chat_interface() -> None:
+    """显示聊天消息并处理新的用户输入"""
+    # 如果是首次加载，初始化消息历史
+    if "messages" not in st.session_state:
+        st.session_state.messages = [{"role": "assistant", "content": "欢迎！我是基于 Amazon Bedrock 的聊天助手。请告诉我您需要什么帮助。"}]
+
+    # 渲染所有现有消息
+    for msg in st.session_state.messages:
+        st.chat_message(msg["role"]).write(msg["content"])
+
+    # 处理用户输入
+    if user_text := st.chat_input("您的消息..."):
+        # 立即渲染用户气泡
+        st.chat_message("user").write(user_text)
+
+        # 发送到后端并获取回复
+        try:
+            reply = send_message_and_get_reply(user_text)
+        except Exception as e:
+            st.error(f"聊天服务错误: {e}")
+            return
+
+        # 渲染助手响应
+        st.session_state.messages.append({"role": "assistant", "content": reply})
+        st.chat_message("assistant").write(reply)
+
+def render_dashboard() -> None:
+    """渲染仪表板"""
+    st.subheader("📊 聊天记录和统计")
+    
+    logs = get_logs(limit=50)
+    if not logs:
+        st.info("暂无聊天记录")
+        return
+
+    # 转换为 DataFrame 用于显示
+    import pandas as pd
+    df = pd.DataFrame(logs)
+    
+    # 显示统计信息
+    col1, col2, col3, col4 = st.columns(4)
+    with col1:
+        st.metric("总对话数", len(df))
+    with col2:
+        st.metric("成功率", f"{df['success'].mean()*100:.1f}%")
+    with col3:
+        st.metric("平均延迟", f"{df['latency_ms'].mean():.1f}ms")
+    with col4:
+        st.metric("唯一用户", df['user_id'].nunique())
+
+    # 显示最近的聊天记录
+    st.subheader("最近的聊天记录")
+    for log in logs[:10]:  # 显示最近10条
+        with st.expander(f"用户 {log['user_id'][:8]}... - {log['timestamp']}"):
+            st.write(f"**用户**: {log['prompt']}")
+            st.write(f"**助手**: {log['response']}")
+            st.write(f"**延迟**: {log['latency_ms']:.1f}ms")
+            st.write(f"**状态**: {'✅ 成功' if log['success'] else '❌ 失败'}")
+
+# ========== 主应用 ==========
+def main() -> None:
+    """主入口点"""
+    # 初始化数据库
+    init_db()
+    
+    # 页面配置和标题
+    st.set_page_config(
+        page_title="企业级 Bedrock 聊天机器人", 
+        page_icon="💼",
+        layout="wide"
+    )
+    
+    # 侧边栏登录
+    login_sidebar()
+    
+    # 主标题
+    st.title("💼 企业级 Bedrock 聊天机器人")
+    
+    # 如果已认证，显示功能选择
+    if st.session_state.get("authenticated"):
+        tab1, tab2 = st.tabs(["💬 聊天", "📊 仪表板"])
+        
+        with tab1:
+            render_chat_interface()
+            
+        with tab2:
+            render_dashboard()
+    else:
+        st.info("请在侧边栏使用 AWS 凭证登录以开始聊天。")
+
+if __name__ == "__main__":
+    main()