app.py

import os
import logging
import pandas as pd
from tqdm import tqdm
from pathlib import Path
import autosim
from autosim.agent import Agent, LongTermMemory, WorkingMemory
from autosim.base import Movie
from autosim.dataset.movie_lens import MovieLensDataset
from autosim.model_manager import ModelManager
from autosim.model_evaluator import ModelEvaluator
from pages import MovieWebsite
from prompts import LONG_TERM_MEMORY_PROMPT, USER_SYSTEM_PROMPT
from tools import DataStore, MovieSearchTool
from utils import format_user_profile
from actions import *

autosim.set_env()

os.makedirs("output", exist_ok=True)

logging.basicConfig(
    filename="output/simulation.txt",
    level=logging.INFO,
    format="%(message)s",
    encoding="latin-1",
)


DF_TASTE = pd.read_csv("data/user_taste.csv")
MOVIE_DATA_STORE = DataStore(data_dir="data")


def create_user(user_id, user_profile, user_taste, activity_trait):
    fatigue_map = {"low": 20, "medium": 30, "high": 40}
    max_fatigue = fatigue_map.get(activity_trait, fatigue_map["medium"])
    user_system_prompt = USER_SYSTEM_PROMPT.format(
        user_profile=user_profile, movie_taste=user_taste
    )
    long_term_memory = LongTermMemory(
        system_prompt=LONG_TERM_MEMORY_PROMPT,
        tools=[
            MovieSearchTool(user_id=user_id, data_store=MOVIE_DATA_STORE),
        ],
    )
    working_memory = WorkingMemory(
        system_prompt=user_system_prompt,
        max_fatigue=max_fatigue,
    )
    return Agent(
        name="User",
        system_prompt=user_system_prompt,
        long_term_memory=long_term_memory,
        working_memory=working_memory,
        actions=[
            click_movie,
            watch_and_rate_movie,
            next_page,
            previous_page,
            back_action,
            exit_action,
        ],
    )


def create_movie_website(user, movie_ids, dataset, image_dir=None):
    movies = []
    for movie_id in movie_ids:
        item_info = dataset.get_item_info(movie_id)
        movie_data = item_info.copy()

        movie_data.update(
            {
                "id": movie_id,
                "rating": float(f"{item_info['rating'] / 2:.1f}"),
                "genres": item_info["genres"].strip().split(","),
            }
        )

        if image_dir is not None and "poster_path" in item_info:
            movie_data["poster_path"] = str(Path(image_dir) / item_info["poster_path"])
        else:
            movie_data["poster_path"] = None

        movies.append(Movie(**movie_data))
    return MovieWebsite(user, movies=movies)


output_folder = "output"
manager = ModelManager(model_folder="models", output_folder=output_folder)
evaluator = ModelEvaluator(output_folder)
dataset = MovieLensDataset(data_dir="dataset/ml-1m")
user_activity_trait_df = pd.read_parquet("data/user_activity_trait.parquet")
IMAGE_DIR = "dataset/ml-1m/posters"

for user_id in tqdm(manager.get_user_ids(), desc="Running simulation"):
    logging.info(f"User ID: {user_id}")

    # 检查是否已处理过该用户
    if all(
        user_id in manager.get_processed_user_ids(model_name)
        for model_name in manager.models
    ):
        logging.info(f"User ID {user_id} already processed. Skipping.")
        continue

    try:
        user_data = dataset.get_user_detail(user_id, 978133367)
        user_profile = format_user_profile(user_data["user_info"])
        user_taste = DF_TASTE[DF_TASTE["user_id"] == user_id]["user_taste"].values[0]
        user_activity_trait = user_activity_trait_df.loc[user_id]["trait"]
    except Exception as e:
        logging.error(f"Error preparing user data for User ID {user_id}: {e}")
        print(f"Error preparing user data for User ID {user_id}: {e}")
        continue

    for model_name, _ in manager:
        try:
            if user_id in manager.get_processed_user_ids(model_name):
                logging.info(
                    f"User ID {user_id} already processed for model {model_name}. Skipping."
                )
                continue

            logging.info(f"Model: {model_name}")
            user = create_user(user_id, user_profile, user_taste, user_activity_trait)
            user.print_system_message()
            recommendations = manager.get_recommendations(model_name, user_id)
            simulator = create_movie_website(user, recommendations, dataset, IMAGE_DIR)
            simulator.run()
            manager.log(model_name, user_id, simulator.logs)
        except Exception as e:
            logging.error(
                f"Error processing User ID {user_id} with model {model_name}: {e}"
            )
            print(f"Error processing User ID {user_id} with model {model_name}: {e}")
            continue
    evaluator.evaluate()