File size: 6,027 Bytes
de99ac3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | import chromadb
import re
import time
from mastodon import Mastodon
from datetime import datetime
from pathlib import Path
import json
# --- 1. CONFIGURE YOUR SETTINGS ---
# !! Fill these in !!
MASTODON_SERVER_URL = r"https://mastodon.social" # e.g., "https://mastodon.social"
MASTODON_ACCESS_TOKEN = json.load(open(r"D:\cs452\TopicBuzz\secrets.json"))["token"] # The token you just generated
USER_LIST_FILE = "usersAll.txt"
DB_PATH = "./my_mastodon_db"
COLLECTION_NAME = "mastodon_posts"
POSTS_TO_FETCH = 100
# -----------------------------------
# Simple regex to strip HTML tags from post content
def strip_html(text):
return re.sub(r'<[^<]+?>', '', text)
# Helper function to sleep and respect rate limits
def polite_request(api_call, *args, **kwargs):
"""
Makes an API call and then waits 1 second.
This respects the default 300 requests / 5 min (1 req/sec) limit.
"""
response = api_call(*args, **kwargs)
time.sleep(1)
return response
def main():
print("--- Mastodon Post Fetcher ---")
# --- 2. Initialize Clients ---
try:
# Initialize Mastodon client
mastodon = Mastodon(
access_token=MASTODON_ACCESS_TOKEN,
api_base_url=MASTODON_SERVER_URL
)
print(f"Mastodon client connected to {MASTODON_SERVER_URL}")
# Initialize ChromaDB client
client = chromadb.PersistentClient(path=DB_PATH)
collection = client.get_or_create_collection(name=COLLECTION_NAME)
print(f"ChromaDB client connected to {DB_PATH}")
except Exception as e:
print(f"Error during initialization: {e}")
return
# --- 3. Read User List ---
user_file = Path(USER_LIST_FILE)
if not user_file.exists():
print(f"Error: User file not found at {USER_LIST_FILE}")
return
with open(user_file, 'r') as f:
usernames = [line.strip() for line in f if line.strip()]
print(f"Found {len(usernames)} users to track.")
# --- 4. Main Fetch Loop ---
for username in usernames:
print(f"\nFetching posts for user: {username}")
try:
# First, find the user's account ID
account = polite_request(mastodon.account_lookup, username)
if not account:
print(f"Could not find user {username}. Skipping.")
continue
user_id = account['id']
# Mastodon's API is paginated, max 40 posts per request
# We need to make multiple requests to get 100
# Fetch Page 1 (posts 1-40)
page1 = polite_request(mastodon.account_statuses, id=user_id, limit=40)
# Fetch Page 2 (posts 41-80)
page2 = polite_request(mastodon.fetch_next, page1) if page1 else []
# Fetch Page 3 (posts 81-120)
page3 = polite_request(mastodon.fetch_next, page2) if page2 else []
# Combine pages and get the most recent 100
# (or fewer if the user doesn't have 100)
all_posts = (page1 or []) + (page2 or []) + (page3 or [])
recent_posts = all_posts[:POSTS_TO_FETCH]
if not recent_posts:
print(f"No posts found for {username}. Skipping.")
continue
# --- 5. Prepare Data for ChromaDB ---
ids_batch = []
docs_batch = []
meta_batch = []
for post in recent_posts:
doc_content = ""
doc_id = str(post['id']) # Use the status ID (reblog or original) as the unique ID
doc_meta = {}
if post['reblog']:
# This is a REBLOG
original_post = post['reblog']
# Use the ORIGINAL post's content for the vector
doc_content = strip_html(original_post['content'])
doc_meta = {
"type": "reblog",
"reblogger_user_id": str(post['account']['id']), # The user who reblogged
"reblog_timestamp": int(post['created_at'].timestamp()), # Time of the reblog
"original_post_id": str(original_post['id']),
"original_author_id": str(original_post['account']['id']),
"language": str(original_post.get('language'))
}
else:
# This is an ORIGINAL post
doc_content = strip_html(post['content'])
doc_meta = {
"type": "original",
"author_user_id": str(post['account']['id']), # The user who posted
"author_handle": str(post['account']['acct']), # FIX FOR FUTURE USE WITH USERNAME
"post_timestamp": int(post['created_at'].timestamp()), # Time of the post
"language": str(post.get('language'))
}
# Don't add empty posts (e.g., if content stripping failed)
if doc_content:
ids_batch.append(doc_id)
docs_batch.append(doc_content)
meta_batch.append(doc_meta)
# --- 6. Store in ChromaDB ---
if ids_batch:
collection.upsert(
ids=ids_batch,
documents=docs_batch,
metadatas=meta_batch
)
print(f"Successfully added/updated {len(ids_batch)} posts for {username}.")
else:
print(f"No new original posts found for {username}.")
except Exception as e:
print(f"An error occurred while processing {username}: {e}")
# Continue to the next user even if one fails
pass
print("\n--- All users processed. ---")
if __name__ == "__main__":
main() |