用 Amazon Nova 多模态嵌入打造航空制造文档智能检索系统

# 安装依赖
pip install boto3 Pillow pdf2image

# 确保 AWS 凭证已配置（ ~/.aws/credentials 或环境变量）
# 需要的 IAM 权限：
#   - bedrock:InvokeModel  (Nova Multimodal Embeddings)
#   - s3:PutObject / s3:GetObject
#   - s3-vector:PutVectorIndex / s3-vector:QueryVectorIndex

import boto3
import json
import base64
from pathlib import Path

bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")

MODEL_ID = "amazon.nova-multimodal-embedding-v1"
EMBED_DIM = 1024  # 支持维度：256 / 512 / 1024，制造场景建议用 1024

def encode_image_to_base64(image_path: str) -> str:
    """将本地图片转为 base64 字符串，供 Bedrock 调用"""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def get_multimodal_embedding(text: str, image_path: str = None) -> list[float]:
    """
    调用 Nova Multimodal Embeddings，同时编码文字和图片。
    如果没有图片，只传文字即可。
    """
    content = [{"text": text}]

    if image_path:
        content.append({
            "image": {
                "format": Path(image_path).suffix.lstrip("."),  # jpg / png
                "source": {"bytes": encode_image_to_base64(image_path)}
            }
        })

    body = {
        "input": {"content": content},
        "configuration": {"outputDimensionality": EMBED_DIM}
    }

    response = bedrock.invoke_model(
        modelId=MODEL_ID,
        body=json.dumps(body),
        accept="application/json",
        contentType="application/json"
    )

    result = json.loads(response["body"].read())
    return result["embedding"]


# ---- 示例：编码一页制造文档 ----
page_text = "焊接顺序：先焊内环缝（G1），再焊外环缝（G2），层间温度控制在150°C以下。"
page_image = "weld_sequence_diagram.png"  # 本地图片路径

embedding = get_multimodal_embedding(page_text, page_image)
print(f"向量维度: {len(embedding)}")  # 1024

import boto3

s3v = boto3.client("s3-vector", region_name="us-east-1")

VECTOR_BUCKET = "aero-manufacturing-vectors"
INDEX_NAME = "doc-pages-index"

# 1. 创建向量索引（首次运行）
s3v.create_vector_index(
    vectorBucketName=VECTOR_BUCKET,
    indexName=INDEX_NAME,
    dimension=EMBED_DIM,
    distanceMetric="cosine"
)

# 2. 写入向量
def upsert_page_vector(page_id: str, embedding: list[float], metadata: dict):
    s3v.put_vectors(
        vectorBucketName=VECTOR_BUCKET,
        indexName=INDEX_NAME,
        vectors=[
            {
                "key": {"id": page_id},
                "values": embedding,
                "metadata": metadata  # 如 {"doc": "WS-2024-037", "page": 5}
            }
        ]
    )

upsert_page_vector(
    page_id="WS-2024-037_p5",
    embedding=embedding,
    metadata={"doc": "WS-2024-037", "page": 5, "type": "weld_procedure"}
)

# 3. 查询检索
def search_pages(query_text: str, top_k: int = 5) -> list[dict]:
    query_emb = get_multimodal_embedding(query_text)  # 纯文本查询
    result = s3v.query_vector_index(
        vectorBucketName=VECTOR_BUCKET,
        indexName=INDEX_NAME,
        queryVector={"values": query_emb},
        topK=top_k,
        returnMetadata=True
    )
    return result["results"]

# 实际查询
hits = search_pages("内环缝焊接的层间温度要求是多少？")
for hit in hits:
    print(f"文档: {hit['metadata']['doc']}, 页码: {hit['metadata']['page']}, 距离: {hit['distance']}")

from anthropic import AnthropicBedrock

anthropic = AnthropicBedrock()

def generate_answer(query: str, top_k: int = 5) -> str:
    hits = search_pages(query, top_k)
    # 组装上下文：取回页面的文字部分（metadata 中可扩展存储原文）
    context_blocks = []
    for hit in hits:
        context_blocks.append({
            "type": "text",
            "text": f"[文档 {hit['metadata']['doc']} 第{hit['metadata']['page']}页]\n"
                    + hit["metadata"].get("page_text", "")
        })
        # 如果检索到的页面有图片，也可以把图片直接传给 LLM
        img_path = hit["metadata"].get("page_image_path")
        if img_path:
            context_blocks.append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": encode_image_to_base64(img_path)
                }
            })

    response = anthropic.messages.create(
        model="anthropic.claude-3-5-sonnet-20241022-v2:0",
        max_tokens=1024,
        system="你是航空制造工艺专家，根据提供的文档内容准确回答问题。如果文档中没有相关信息，明确说明。",
        messages=[{
            "role": "user",
            "content": context_blocks + [{"type": "text", "text": query}]
        }]
    )
    return response.content[0].text

answer = generate_answer("内环缝焊接的层间温度要求是多少？")
print(answer)