def recall_at_k(ranked_ids, relevant_ids, k=10):
    top_k = set(ranked_ids[:k])
    relevant = set(relevant_ids)
    if not relevant:
        raise ValueError("Recall@K 只适用于至少有一个相关文档的查询")
    return len(top_k & relevant) / len(relevant)


def mean_recall_at_k(results, labels, k=10):
    answerable_query_ids = [
        query_id for query_id, relevant_ids in labels.items() if relevant_ids
    ]
    if not answerable_query_ids:
        raise ValueError("评测集中没有可计算 Recall@K 的查询")
    scores = [
        recall_at_k(results[query_id], labels[query_id], k)
        for query_id in answerable_query_ids
    ]
    return sum(scores) / len(scores)


def no_answer_false_accept_rate(scored_results, labels, threshold):
    """scored_results[qid] 是按分数降序排列的 {id, score} 列表。"""
    no_answer_query_ids = [
        query_id for query_id, relevant_ids in labels.items() if not relevant_ids
    ]
    if not no_answer_query_ids:
        raise ValueError("评测集中没有无答案查询")
    false_accepts = sum(
        bool(scored_results.get(query_id))
        and scored_results[query_id][0]["score"] >= threshold
        for query_id in no_answer_query_ids
    )
    return false_accepts / len(no_answer_query_ids)

from pathlib import Path

import open_clip
import torch
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"
model, _, preprocess = open_clip.create_model_and_transforms(
    "ViT-B-32",
    pretrained="laion2b_s34b_b79k",
)
tokenizer = open_clip.get_tokenizer("ViT-B-32")
model = model.to(device).eval()


def preprocess_image(path: Path) -> torch.Tensor:
    with Image.open(path) as image:
        return preprocess(image.convert("RGB"))


@torch.inference_mode()
def encode_images(
    paths: list[Path],
    batch_size: int = 32,
) -> list[list[float]]:
    encoded = []
    for start in range(0, len(paths), batch_size):
        batch_paths = paths[start : start + batch_size]
        batch = torch.stack(
            [preprocess_image(path) for path in batch_paths]
        ).to(device)
        vectors = model.encode_image(batch)
        vectors = vectors / vectors.norm(dim=-1, keepdim=True)
        encoded.extend(vectors.cpu().tolist())
    return encoded


@torch.inference_mode()
def encode_texts(texts: list[str]) -> list[list[float]]:
    tokens = tokenizer(texts).to(device)
    vectors = model.encode_text(tokens)
    vectors = vectors / vectors.norm(dim=-1, keepdim=True)
    return vectors.cpu().tolist()

import json

from pymilvus import MilvusClient

client = MilvusClient(uri="./multimodal.db")
collection = "catalog_images"

manifest_path = Path("images/manifest.jsonl")
if not manifest_path.exists():
    raise FileNotFoundError("请先创建 images/manifest.jsonl")

records = [
    json.loads(line)
    for line in manifest_path.read_text(encoding="utf-8").splitlines()
    if line.strip()
]
if not records:
    raise ValueError("manifest 为空，至少需要一张授权图片")

sample_paths = [Path(record["path"]) for record in records]
missing_paths = [str(path) for path in sample_paths if not path.is_file()]
if missing_paths:
    raise FileNotFoundError(f"manifest 中的图片不存在：{missing_paths}")

vectors = encode_images(sample_paths)
dimension = len(vectors[0])

if client.has_collection(collection):
    client.drop_collection(collection)

client.create_collection(
    collection_name=collection,
    dimension=dimension,
    metric_type="COSINE",
    consistency_level="Strong",
)

client.insert(
    collection_name=collection,
    data=[
        {
            "id": index,
            "vector": vector,
            "image_id": record["image_id"],
            "uri": str(path),
            "caption": record["caption"],
            "category": record["category"],
            "source": record["source"],
            "license": record["license"],
        }
        for index, (record, path, vector) in enumerate(
            zip(records, sample_paths, vectors)
        )
    ],
)

import json


def search_images(query: str, limit: int = 8, category: str | None = None):
    query_vector = encode_texts([query])[0]
    filter_expression = (
        f"category == {json.dumps(category, ensure_ascii=False)}"
        if category
        else ""
    )
    results = client.search(
        collection_name=collection,
        data=[query_vector],
        filter=filter_expression,
        limit=limit,
        output_fields=["image_id", "uri", "caption", "category", "source"],
        search_params={"metric_type": "COSINE", "params": {}},
    )[0]
    return [
        {
            "score": hit["distance"],
            **hit["entity"],
        }
        for hit in results
    ]

query_vector = encode_images([Path("query.jpg")])[0]
results = client.search(
    collection_name=collection,
    data=[query_vector],
    limit=8,
    output_fields=["image_id", "uri", "category"],
    search_params={"metric_type": "COSINE", "params": {}},
)

results = search_images("red running shoes", limit=8, category="shoes")

import base64
import mimetypes
import os

from openai import OpenAI

vision_client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ.get("OPENAI_BASE_URL"),
)


def image_data_url(path: Path) -> str:
    mime_type = mimetypes.guess_type(path.name)[0] or "image/jpeg"
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:{mime_type};base64,{encoded}"


def answer_with_images(query: str, hits: list[dict]) -> dict:
    if not hits:
        return {"answer": "没有检索到可用图片。", "image_ids": []}
    content = [
        {
            "type": "text",
            "text": (
                "只根据随后提供的图片和元数据回答。无法确认的细节要明确说明，"
                "每个结论标注 [图片: image_id]。\n\n问题：" + query
            ),
        }
    ]
    for hit in hits[:4]:
        content.append(
            {
                "type": "text",
                "text": f'图片 ID：{hit["image_id"]}；描述：{hit["caption"]}',
            }
        )
        content.append(
            {
                "type": "image_url",
                "image_url": {"url": image_data_url(Path(hit["uri"]))},
            }
        )
    response = vision_client.chat.completions.create(
        model=os.environ["VISION_MODEL"],
        temperature=0,
        messages=[{"role": "user", "content": content}],
    )
    return {
        "answer": response.choices[0].message.content or "",
        "image_ids": [hit["image_id"] for hit in hits[:4]],
    }

import re


def validate_image_citations(answer: str, allowed_image_ids: set[str]) -> dict:
    cited = set(re.findall(r"\[图片:\s*([^\]]+)\]", answer))
    return {
        "cited_image_ids": sorted(cited),
        "unknown_image_ids": sorted(cited - allowed_image_ids),
        "has_valid_citation": bool(cited & allowed_image_ids),
    }

def recall_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    if not relevant_ids:
        raise ValueError("Recall@K 不适用于没有相关图片的查询")
    return len(set(retrieved_ids[:k]) & relevant_ids) / len(relevant_ids)


def evaluate_text_to_image(test_cases: list[dict], k: int = 5) -> dict:
    scores = []
    failures = []
    for case in test_cases:
        hits = search_images(case["query"], limit=k)
        retrieved_ids = [hit["image_id"] for hit in hits]
        score = recall_at_k(retrieved_ids, set(case["relevant_ids"]), k)
        scores.append(score)
        if score < 1:
            failures.append({**case, "retrieved_ids": retrieved_ids})
    return {
        f"recall_at_{k}": sum(scores) / len(scores),
        "failures": failures,
    }

import hashlib
import re
from dataclasses import dataclass


@dataclass
class Chunk:
    document_id: str
    chunk_id: str
    title: str
    text: str
    source: str
    updated_at: str
    access_scope: str


def split_markdown(
    source: str,
    content: str,
    updated_at: str,
    access_scope: str,
    document_id: str | None = None,
    max_chars: int = 1800,
):
    # 生产系统应从内容管理系统传入永久 document_id；这里只用来源地址派生。
    document_id = document_id or hashlib.sha256(source.encode()).hexdigest()[:16]
    sections = re.split(r"(?=^##\s+)", content, flags=re.MULTILINE)
    chunks = []
    chunk_occurrences = {}
    for section in sections:
        section = section.strip()
        if not section:
            continue
        first_line, *_ = section.splitlines()
        title = first_line.removeprefix("## ").strip()
        for window_index, start in enumerate(range(0, len(section), max_chars)):
            text = section[start : start + max_chars].strip()
            content_hash = hashlib.sha256(
                f"{title}\n{window_index}\n{text}".encode()
            ).hexdigest()[:16]
            occurrence = chunk_occurrences.get(content_hash, 0)
            chunk_occurrences[content_hash] = occurrence + 1
            chunk_hash = (
                content_hash
                if occurrence == 0
                else hashlib.sha256(
                    f"{content_hash}:{occurrence}".encode()
                ).hexdigest()[:16]
            )
            chunks.append(
                Chunk(
                    document_id=document_id,
                    chunk_id=f"{document_id}:{chunk_hash}",
                    title=title,
                    text=text,
                    source=source,
                    updated_at=updated_at,
                    access_scope=access_scope,
                )
            )
    return chunks

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")


def encode_documents(texts: list[str]) -> list[list[float]]:
    vectors = model.encode(
        texts,
        normalize_embeddings=True,
        batch_size=16,
        show_progress_bar=True,
    )
    return vectors.tolist()


def encode_query(query: str) -> list[float]:
    vector = model.encode(
        [query],
        prompt_name="query",
        normalize_embeddings=True,
    )[0]
    return vector.tolist()

from pymilvus import MilvusClient

client = MilvusClient(uri="./knowledge_base.db")
collection = "knowledge_chunks"
dimension = model.get_sentence_embedding_dimension()

if client.has_collection(collection):
    client.drop_collection(collection)

client.create_collection(
    collection_name=collection,
    dimension=dimension,
    metric_type="COSINE",
    consistency_level="Strong",
)

with open("handbook.md", encoding="utf-8") as source_file:
    chunks = split_markdown(
        source="handbook.md",
        content=source_file.read(),
        updated_at="2026-07-28T00:00:00Z",
        access_scope="employees",
        document_id="employee-handbook",
    )
vectors = encode_documents([chunk.text for chunk in chunks])

client.insert(
    collection_name=collection,
    data=[
        {
            "id": index,
            "vector": vector,
            "document_id": chunk.document_id,
            "chunk_id": chunk.chunk_id,
            "title": chunk.title,
            "source": chunk.source,
            "updated_at": chunk.updated_at,
            "access_scope": chunk.access_scope,
            "text": chunk.text,
        }
        for index, (chunk, vector) in enumerate(zip(chunks, vectors))
    ],
)

import json


def retrieve(query: str, access_scope: str, limit: int = 20):
    if not access_scope:
        raise ValueError("access_scope 不能为空")
    access_scope_literal = json.dumps(access_scope, ensure_ascii=False)
    result = client.search(
        collection_name=collection,
        data=[encode_query(query)],
        filter=f"access_scope == {access_scope_literal}",
        limit=limit,
        output_fields=[
            "document_id",
            "chunk_id",
            "title",
            "source",
            "updated_at",
            "access_scope",
            "text",
        ],
        search_params={"metric_type": "COSINE", "params": {}},
    )[0]
    return [
        {
            "score": hit["distance"],
            **hit["entity"],
        }
        for hit in result
    ]

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

reranker_name = "Qwen/Qwen3-Reranker-0.6B"
reranker_tokenizer = AutoTokenizer.from_pretrained(
    reranker_name,
    padding_side="left",
)
reranker = AutoModelForCausalLM.from_pretrained(reranker_name).eval()
reranker_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
reranker = reranker.to(reranker_device)

prefix = (
    "<|im_start|>system\nJudge whether the Document meets the requirements "
    "based on the Query and the Instruct provided. Note that the answer can "
    'only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
)
suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
prefix_tokens = reranker_tokenizer.encode(prefix, add_special_tokens=False)
suffix_tokens = reranker_tokenizer.encode(suffix, add_special_tokens=False)
yes_token_id = reranker_tokenizer.convert_tokens_to_ids("yes")
no_token_id = reranker_tokenizer.convert_tokens_to_ids("no")
max_length = 8192


@torch.inference_mode()
def score_pairs(pairs: list[tuple[str, str]]) -> list[float]:
    task = "根据用户问题，判断文档是否包含能够支持回答的相关信息"
    texts = [
        f"<Instruct>: {task}\n<Query>: {query}\n<Document>: {document}"
        for query, document in pairs
    ]
    encoded = reranker_tokenizer(
        texts,
        padding=False,
        truncation="longest_first",
        max_length=max_length - len(prefix_tokens) - len(suffix_tokens),
        return_attention_mask=False,
    )
    for index, token_ids in enumerate(encoded["input_ids"]):
        encoded["input_ids"][index] = prefix_tokens + token_ids + suffix_tokens
    inputs = reranker_tokenizer.pad(encoded, padding=True, return_tensors="pt")
    inputs = {name: value.to(reranker_device) for name, value in inputs.items()}
    next_token_logits = reranker(**inputs).logits[:, -1, :]
    binary_logits = torch.stack(
        [next_token_logits[:, no_token_id], next_token_logits[:, yes_token_id]],
        dim=1,
    )
    return binary_logits.softmax(dim=1)[:, 1].cpu().tolist()


def rerank(query: str, candidates: list[dict]) -> list[dict]:
    if not candidates:
        return []
    scores = score_pairs([(query, item["text"]) for item in candidates])
    ranked = [
        {**candidate, "rerank_score": float(score)}
        for candidate, score in zip(candidates, scores)
    ]
    return sorted(ranked, key=lambda item: item["rerank_score"], reverse=True)

import os
import re

from openai import OpenAI

llm = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ.get("OPENAI_BASE_URL"),
)


def answer(query: str, access_scope: str, top_k: int = 5) -> dict:
    candidates = rerank(
        query,
        retrieve(query, access_scope=access_scope, limit=20),
    )[:top_k]
    if not candidates:
        return {
            "answer": "现有权限范围内没有检索到可用资料。",
            "sources": [],
        }
    context = "\n\n".join(
        f'[来源: {item["chunk_id"]}]\n{item["text"]}' for item in candidates
    )
    response = llm.chat.completions.create(
        model=os.environ["CHAT_MODEL"],
        temperature=0,
        messages=[
            {
                "role": "system",
                "content": (
                    "你是企业知识库助手。只能根据提供的上下文回答。"
                    "证据不足时回答‘现有资料无法确认’。"
                    "每个关键结论后必须标注 [来源: chunk_id]。"
                ),
            },
            {
                "role": "user",
                "content": f"问题：{query}\n\n上下文：\n{context}",
            },
        ],
    )
    answer_text = response.choices[0].message.content or ""
    allowed_sources = {item["chunk_id"] for item in candidates}
    cited_sources = set(re.findall(r"\[来源:\s*([^\]]+)\]", answer_text))
    unknown_sources = cited_sources - allowed_sources
    if unknown_sources:
        raise ValueError(f"回答引用了未提供的来源：{sorted(unknown_sources)}")
    return {
        "answer": answer_text,
        "sources": [
            {"chunk_id": item["chunk_id"], "source": item["source"]}
            for item in candidates
        ],
    }