mcp_servers/product_user_database.py

# mcp_server/product_user_database.py
import os
import pickle
import numpy as np
from typing import Dict, Any, List
from dotenv import load_dotenv
from tqdm import tqdm
from itertools import combinations

from scipy import sparse
from sklearn.metrics.pairwise import cosine_similarity
from mcp.server.fastmcp import FastMCP
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.routing import Route, Mount
import uvicorn
import pandas as pd
import torch
from transformers import CLIPProcessor, CLIPModel
from openai import AsyncOpenAI

# Load environment variables
load_dotenv()
FASHION_DATA_ROOT = os.getenv("FASHION_DATA_ROOT", "/mnt/d/PostDoc/fifth paper/code/FashionVLM/datasets/FashionRec")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_API_BASE = os.getenv("OPENAI_API_BASE")
openai = AsyncOpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_API_BASE)

###################################
#########Loading Model#############
###################################
# Load CLIP model and processor
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32", local_files_only=True)
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", local_files_only=True)
clip_model.eval()


# Load item metadata
items_df = pd.read_parquet(f"{FASHION_DATA_ROOT}/meta/items_lite.parquet").set_index("item_id")
outfits_df = pd.read_parquet(f"{FASHION_DATA_ROOT}/meta/outfits_lite.parquet").set_index("outfit_id")
users_df = pd.read_parquet(f"{FASHION_DATA_ROOT}/meta/users_lite.parquet").set_index("user_id")
image_paths = items_df["path"].to_dict()


class InteractionDataManager:
    def __init__(self, users_df, outfits_df, items_df):
        """
        初始化类，加载数据并设置基本参数

        参数:
        - users_file: 用户数据文件路径 (parquet)
        - outfits_file: Outfit 数据文件路径 (parquet)
        - items_file: 单品数据文件路径 (parquet)
        """
        self.users_df = users_df
        self.outfits_df = outfits_df
        self.items_df = items_df

        # 创建映射
        self.item_id_to_index = {item_id: index for index, item_id in enumerate(self.items_df.index)}
        self.index_to_item_id = {index: item_id for index, item_id in enumerate(self.items_df.index)}
        self.user_id_to_index = {user_id: index for index, user_id in enumerate(self.users_df.index)}
        self.index_to_user_id = {index: user_id for index, user_id in enumerate(self.users_df.index)}
        self.outfit_ids_dict = self.outfits_df['item_ids'].to_dict()  # get outfit's item ids from outfit id
        self.item_category_dict = self.items_df['category'].to_dict()  # get item's category from item id
        self.item_subcategory_dict = self.items_df['subcategory'].to_dict()  # get item's subcategory from item id
        self.n_items = len(self.items_df)
        self.n_users = len(self.users_df)

        self.user_outfit_pairs = []
        outfit_set = set(self.outfits_df.index)
        for uid, user in self.users_df.iterrows():
            oids = user.outfit_ids.split(",")
            self.user_outfit_pairs.extend([(uid, oid) for oid in oids if oid in outfit_set])

        # 预处理类别到物品ID的映射（使用groupby）
        self.subcategory_to_items = self.items_df.groupby('subcategory').apply(lambda x: set(x.index)).to_dict()

        # 预处理类别到物品索引的映射（优化查找效率）
        self.subcategory_to_indices = {}
        for subcategory, item_ids in self.subcategory_to_items.items():
            self.subcategory_to_indices[subcategory] = set([self.item_id_to_index[item_id]
                                                            for item_id in item_ids
                                                            if item_id in self.item_id_to_index])

        item_interaction_matrix_path = f'{FASHION_DATA_ROOT}/data/personalized_recommendation/temp_matrix/item_matrix.npz'
        try:
            self.load_matrix('item', item_interaction_matrix_path)
        except FileNotFoundError:
            self.build_item_interaction_matrix()
            self.save_matrix('item', item_interaction_matrix_path)

        user_item_interaction_matrix_path = f'{FASHION_DATA_ROOT}/data/personalized_recommendation/temp_matrix/user_item_matrix.npz'
        try:
            self.load_matrix('user_item', user_item_interaction_matrix_path)
        except FileNotFoundError:
            self.build_user_item_interaction_matrix()
            self.save_matrix('user_item', user_item_interaction_matrix_path)

        # 加载item clip features
        with open(f"{FASHION_DATA_ROOT}/meta/clip_features.pkl", "rb") as f:
            print("Loading Fashion Features...")
            self.clip_features = pickle.load(f)
            print("Loading Fashion Features Successfully")

        # Prepare embeddings and item IDs
        self.item_ids = list(self.clip_features.keys())
        self.image_embeddings = np.array([self.clip_features[item_id]["image_embeds"] for item_id in item_ids])

    def save_matrix(self, matrix_type, filepath):
        """
        保存矩阵到文件

        参数:
        - matrix_type: 'item' 或 'user_item'，指定保存的矩阵类型
        - filepath: 保存路径 (例如 'temp/item_matrix.npz')
        """
        if matrix_type == 'item':
            matrix = self.item_interaction_matrix
        elif matrix_type == 'user_item':
            matrix = self.user_item_interaction_matrix
        else:
            raise ValueError("matrix_type must be 'item' or 'user_item'")

        if matrix is None:
            raise ValueError(f"{matrix_type} matrix has not been built yet.")

        sparse.save_npz(filepath, matrix)
        print(f"Saved {matrix_type} matrix to {filepath}")

    def load_matrix(self, matrix_type, filepath):
        """
        从文件加载矩阵

        参数:
        - matrix_type: 'item' 或 'user_item'，指定加载的矩阵类型
        - filepath: 加载路径 (例如 'temp/item_matrix.npz')
        """
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File {filepath} does not exist.")

        matrix = sparse.load_npz(filepath)
        if matrix_type == 'item':
            self.item_interaction_matrix = matrix
        elif matrix_type == 'user_item':
            self.user_item_interaction_matrix = matrix
        else:
            raise ValueError("matrix_type must be 'item' or 'user_item'")

        print(f"Loaded {matrix_type} matrix from {filepath}")
        return matrix

    def build_item_interaction_matrix(self):
        """构建 Item-Item 交互矩阵"""
        # 初始化单品交互矩阵
        self.item_interaction_matrix = sparse.lil_matrix((self.n_items, self.n_items), dtype=int)

        for index, outfit in tqdm(self.outfits_df.iterrows(), total=len(self.outfits_df)):
            item_ids = outfit['item_ids'].split(',')
            # 记录 item 对的共现
            for item_id1, item_id2 in combinations(item_ids, r=2):
                if item_id1 in self.item_id_to_index and item_id2 in self.item_id_to_index:
                    idx1 = self.item_id_to_index[item_id1]
                    idx2 = self.item_id_to_index[item_id2]
                    self.item_interaction_matrix[idx1, idx2] += 1
                    self.item_interaction_matrix[idx2, idx1] += 1  # 无序对称

        # 转换为 CSR 格式
        self.item_interaction_matrix = self.item_interaction_matrix.tocsr()
        return self.item_interaction_matrix

    def build_user_item_interaction_matrix(self):
        """构建 User-Item 交互矩阵"""
        # 初始化用户-单品交互矩阵
        self.user_item_interaction_matrix = sparse.lil_matrix((self.n_users, self.n_items), dtype=int)

        for uid, user in tqdm(self.users_df.iterrows(), total=len(self.users_df)):
            oids = user["outfit_ids"].split(",")
            outfits = self.outfits_df.loc[self.outfits_df.index.isin(oids)]
            for oid, outfit in outfits.iterrows():
                item_ids = outfit['item_ids'].split(',')
                # 记录 user-item 对的出现
                for iid in item_ids:
                    if iid in self.item_id_to_index:
                        uidx = self.user_id_to_index[uid]
                        iidx = self.item_id_to_index[iid]
                        self.user_item_interaction_matrix[uidx, iidx] += 1

        # 转换为 CSR 格式
        self.user_item_interaction_matrix = self.user_item_interaction_matrix.tocsr()
        return self.user_item_interaction_matrix

    def _process_interactions_for_category(
            self,
            matrix,
            given_id,
            category_indices,
            id_to_index
    ):
        """
        处理单个实体与目标类别的交互

        参数:
        - matrix: 交互矩阵
        - given_id: 给定的实体ID（用户或物品）
        - category_indices: 目标类别的物品索引集合

        返回:
        - 交互列表，每个元素为一个包含item_id、interaction_count和score的字典
        """
        interactions = []

        given_index = id_to_index[given_id]
        row = matrix[given_index]

        # 提取该行的非零元素
        row_start = row.indptr[0]
        row_end = row.indptr[1]
        col_indices = row.indices[row_start:row_end]
        data_values = row.data[row_start:row_end]

        # 筛选出属于目标类别的物品
        for col_idx, value in zip(col_indices, data_values):
            # 检查是否为目标类别的物品
            if col_idx in category_indices:
                # 获取物品ID
                output_id = self.index_to_item_id[col_idx]
                interactions.append({
                    'item_id': output_id,
                    'interaction_count': int(value),
                    'score': 0.0
                })

        return interactions

    def get_item_category_interactions(
        self,
        target_category: str,
        given_ids: List[str],
        query_type='item',  # item or user
        top_k=None,
    ):
        """
        获取指定实体（用户或单品）与目标类别的所有交互情况

        参数:
        - target_category: 待查询的subcategory
        - given_ids: List of 目标类别
        - query_type: 查询的类别， item或user
        - top_k: 返回交互次数最多的前k个物品, 如果是None直接全部返回

        返回:
        - 列表，包含与目标类别的交互统计信息，按交互次数排序
        """
        if query_type == 'item':
            matrix = self.item_interaction_matrix
            id_to_index = self.item_id_to_index
        elif query_type == 'user':
            matrix = self.user_item_interaction_matrix
            id_to_index = self.user_id_to_index
        else:
            print(f'query_type must be either item or user but got {query_type}')
            return []

        # 收集所有交互记录
        all_interactions = []
        category = target_category
        category_indices = self.subcategory_to_indices.get(category, set())  # 获取该类别的所有物品索引

        # 获取该实体的所有交互
        for given_id in given_ids:
            interactions = self._process_interactions_for_category(
                matrix, given_id, category_indices, id_to_index
            )
            # 将交互添加到结果列表
            all_interactions.extend(interactions)

        # 合并相同物品的交互次数
        item_interactions = {}
        for interaction in all_interactions:
            item_id = interaction['item_id']
            count = interaction['interaction_count']

            if item_id in item_interactions:
                item_interactions[item_id] += count
            else:
                item_interactions[item_id] = count

        # 转换为结果格式
        merged_interactions = [
            {'item_id': item_id, 'interaction_count': count, 'score': 0.0}
            for item_id, count in item_interactions.items()
        ]

        # 排序
        if merged_interactions:
            merged_interactions.sort(key=lambda x: x['interaction_count'], reverse=True)

        # 截取top-k
        if top_k and merged_interactions:
            merged_interactions = merged_interactions[:top_k]

        # 存储结果
        return merged_interactions

    def rank_by_similarity(self, item_interactions, user_interactions, beta=2.0):
        """
        计算用户交互项与商品交互项的相似度并排序
        """

        def get_combined_features(feature_dict):
            return (feature_dict['image_embeds'] + feature_dict['text_embeds']) / 2

        item_feature_list = []
        for item in item_interactions:
            item_id = item['item_id']
            if item_id not in self.clip_features:
                raise ValueError(f"Didn't find clip feature of item with id: {item_id}")

            item_features = get_combined_features(self.clip_features[item_id])
            item_feature_list.append(item_features)

        weights = np.array([x['interaction_count'] for x in item_interactions], dtype=np.float32)
        weights = weights / np.sum(weights)
        item_feature = np.sum(np.stack(item_feature_list, axis=0) * weights[:, np.newaxis], axis=0).reshape(1, -1)

        max_count = max((user_item.get('interaction_count', 1) for user_item in user_interactions), default=1)
        for user_item in user_interactions:
            user_item_id = user_item['item_id']
            if user_item_id not in self.clip_features:
                raise ValueError(f"Didn't find clip feature of item with id: {user_item_id}")

            user_item_features = get_combined_features(self.clip_features[user_item_id]).reshape(1, -1)
            similarity = cosine_similarity(user_item_features, item_feature).item()
            interaction_count = user_item['interaction_count']
            count_factor = (interaction_count / max_count) * beta + 1
            user_item['score'] = float(similarity) * count_factor

        user_interactions.sort(key=lambda x: x.get('score', 0), reverse=True)
        return user_interactions


data_manager = InteractionDataManager(users_df, outfits_df, items_df)
mcp = FastMCP('image-retrieval-server')


@mcp.tool()
async def summary_user_history(user_id: str, target_category: str, list_of_items: List[str]) -> str:
    """Summary user's buying history of specific fashion category given user_id, target_category, list_of_items
    After we collect all buying history of this user, we will summarize descriptions of these historical items through LLM.
    So we will return user's preference about target_category in sentences.

    Args:
        user_id (str): User id. Will be provided through prompt
        target_category (str): We care about user's buying history of this specific category.
        list_of_items: List of item ids for history filtering. Will be provided through prompt
    """
    # We need to find the most appropriate item to become the target item
    # It should have enough relationship with user and other items
    # Specifically, item_interaction larger than 3, history larger than 10
    item_interaction_result = data_manager.get_item_category_interactions(
        target_category, list_of_items, query_type='item'
    )
    user_interaction_result = data_manager.get_item_category_interactions(
        target_category, [user_id], query_type='user'
    )

    def get_description(item_id: str) -> str:
        return data_manager.items_df.loc[item_id].gen_description

    descriptions_for_summary = []
    if len(item_interaction_result) == 0:
        descriptions_for_summary = [get_description(x['item_id']) for x in user_interaction_result]
    else:
        if len(user_interaction_result) >= 0:
            user_interaction_result = data_manager.rank_by_similarity(
                item_interaction_result,
                user_interaction_result
            )
            descriptions_for_summary = [get_description(x['item_id']) for x in user_interaction_result[:5]]

    if descriptions_for_summary:
        user_message = f"Summary user's preference of {target_category} based on following descriptions of fashion items that user brought previously:"
        for x in descriptions_for_summary:
            user_message += f"\n{x}"
        # Get summary using OpenAI API call
        response = await openai.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": f"You are a user preference summary assistant. Your response is limited in one sentence, staring at 'I prefer ...'"},
                {"role": "user", "content": user_message}
            ],
            max_tokens=1000,
        )
        return response.choices[0].message.content
    else:
        return ""


user_id = "115"
# 根据类别和given outfit找到这个用户的历史交互
partial_outfit = ["25479e5dacebbfaed18a7dc4830bd5cd19114486", "becc7b46236e9abb6f6760e7a1569b06bbc236c1",
                  "180c32b5c8c164f3c632f3e73d6002ccfa6fea57"]
target_category = "Skirts"
summary_user_history(user_id, target_category, partial_outfit)


async def compute_text_embedding(text: str) -> np.ndarray:
    inputs = clip_processor(text=text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        text_embedding = clip_model.get_text_features(**inputs).numpy()
    return text_embedding / np.linalg.norm(text_embedding, axis=1, keepdims=True)


async def find_most_similar_image(text_embedding: np.ndarray) -> Dict[str, Any]:
    similarities = np.dot(data_manager.image_embeddings, text_embedding.T).flatten()
    most_similar_idx = np.argmax(similarities)
    most_similar_item_id = data_manager.item_ids[most_similar_idx]
    return {
        "image_path": image_paths[most_similar_item_id],
        "similarity": float(similarities[most_similar_idx])
    }


@mcp.tool()
async def retrieve_image(text: str) -> Dict[str, Any]:
    """Search for the most similar fashion image based on a text description.

    Args:
        text (str): Text description of the fashion item to search.
    """
    print(f"Searching for {text}")
    text_embedding = await compute_text_embedding(text)
    return await find_most_similar_image(text_embedding)


mcp_server = mcp._mcp_server  # 获取内部 Server 对象
sse_transport = SseServerTransport("/messages/")


async def handle_sse(request):
    print("Handling SSE connection")
    async with sse_transport.connect_sse(request.scope, request.receive, request._send) as streams:
        read_stream, write_stream = streams
        await mcp_server.run(
            read_stream,
            write_stream,
            mcp_server.create_initialization_options(),
        )

# 定义路由
routes = [
    Route("/sse", endpoint=handle_sse),
    Mount("/messages/", app=sse_transport.handle_post_message),
]

# 创建 Starlette 应用
starlette_app = Starlette(routes=routes)


if __name__ == "__main__":
    print("Starting Image Retrieval server with HTTP and SSE...")
    uvicorn.run(starlette_app, host="0.0.0.0", port=8001)  # 使用 8001 端口，避免与 FashionVLM 冲突