|
| 1 | +import json |
| 2 | +import logging |
| 3 | +from typing import Any |
| 4 | +from datetime import datetime, timedelta |
| 5 | + |
| 6 | +from bson import ObjectId |
| 7 | +from tc_hivemind_backend.db.qdrant import QdrantSingleton |
| 8 | +from tc_hivemind_backend.db.mongo import MongoSingleton |
| 9 | + |
| 10 | +from temporalio import activity, workflow |
| 11 | +from qdrant_client.models import Filter, FieldCondition, MatchValue |
| 12 | + |
| 13 | +with workflow.unsafe.imports_passed_through(): |
| 14 | + from hivemind_summarizer.schema import ( |
| 15 | + TelegramSummariesActivityInput, |
| 16 | + TelegramSummariesRangeActivityInput, |
| 17 | + TelegramGetCollectionNameInput, |
| 18 | + ) |
| 19 | + |
| 20 | + |
| 21 | +def extract_summary_text(node_content: dict[str, Any]) -> str: |
| 22 | + """ |
| 23 | + Extract the actual summary text from the node_content. |
| 24 | +
|
| 25 | + Parameters |
| 26 | + ---------- |
| 27 | + node_content : dict[str, Any] |
| 28 | + The parsed node_content object |
| 29 | +
|
| 30 | + Returns |
| 31 | + ------- |
| 32 | + str |
| 33 | + The extracted summary text |
| 34 | + """ |
| 35 | + # Based on the example provided, the text is in the "text" field |
| 36 | + if isinstance(node_content, dict) and "text" in node_content: |
| 37 | + return node_content["text"] |
| 38 | + |
| 39 | + return "Summary text not found" |
| 40 | + |
| 41 | + |
| 42 | +@activity.defn |
| 43 | +async def get_collection_name(input: TelegramGetCollectionNameInput) -> str: |
| 44 | + """ |
| 45 | + Activity that extracts collection name from MongoDB based on platform_id and community_id. |
| 46 | +
|
| 47 | + Parameters |
| 48 | + ---------- |
| 49 | + input: TelegramGetCollectionNameInput |
| 50 | + Input object containing platform_id and community_id |
| 51 | +
|
| 52 | + Returns |
| 53 | + ------- |
| 54 | + str |
| 55 | + The collection name in format [communityId]_[platformName]_summary |
| 56 | +
|
| 57 | + Raises |
| 58 | + ------ |
| 59 | + Exception |
| 60 | + If platform not found or error occurs during DB access |
| 61 | + """ |
| 62 | + platform_id = input.platform_id |
| 63 | + community_id = input.community_id |
| 64 | + |
| 65 | + logging.info( |
| 66 | + f"Getting collection name for platform_id: {platform_id}, community_id: {community_id}" |
| 67 | + ) |
| 68 | + |
| 69 | + try: |
| 70 | + # Get MongoDB client |
| 71 | + mongo_client = MongoSingleton.get_instance().get_client() |
| 72 | + |
| 73 | + # Query the platform from Core database |
| 74 | + platform = mongo_client["Core"]["platforms"].find_one( |
| 75 | + {"_id": ObjectId(platform_id)} |
| 76 | + ) |
| 77 | + |
| 78 | + if not platform: |
| 79 | + raise Exception(f"Platform with ID {platform_id} not found") |
| 80 | + |
| 81 | + # Extract platform name |
| 82 | + platform_name = platform.get("name") |
| 83 | + if not platform_name: |
| 84 | + raise Exception(f"Platform name not found for platform_id {platform_id}") |
| 85 | + |
| 86 | + # Construct collection name |
| 87 | + collection_name = f"{community_id}_{platform_name}_summary" |
| 88 | + |
| 89 | + logging.info(f"Generated collection name: {collection_name}") |
| 90 | + return collection_name |
| 91 | + |
| 92 | + except Exception as e: |
| 93 | + logging.error(f"Error getting collection name: {str(e)}") |
| 94 | + raise |
| 95 | + |
| 96 | + |
| 97 | +@activity.defn |
| 98 | +async def fetch_telegram_summaries_by_date( |
| 99 | + input: TelegramSummariesActivityInput, |
| 100 | +) -> list[dict[str, Any]] | str: |
| 101 | + """ |
| 102 | + Activity that fetches Telegram summaries for a specific date from Qdrant. |
| 103 | +
|
| 104 | + Parameters |
| 105 | + ---------- |
| 106 | + input : TelegramSummariesActivityInput |
| 107 | + Input object containing date, collection_name and extract_text_only |
| 108 | +
|
| 109 | + Returns |
| 110 | + ------- |
| 111 | + list[dict[str, Any]] | str |
| 112 | + A list of summary objects for the specified date or a string of summaries |
| 113 | + """ |
| 114 | + date = input.date |
| 115 | + extract_text_only = input.extract_text_only |
| 116 | + collection_name = input.collection_name |
| 117 | + |
| 118 | + logging.info("Started fetch_telegram_summaries_by_date!") |
| 119 | + if not collection_name: |
| 120 | + raise ValueError("Collection name is required but was not provided") |
| 121 | + |
| 122 | + logging.info( |
| 123 | + f"Fetching summaries for date: {date} from collection: {collection_name}" |
| 124 | + ) |
| 125 | + |
| 126 | + try: |
| 127 | + # Get Qdrant client |
| 128 | + qdrant_client = QdrantSingleton.get_instance().get_client() |
| 129 | + |
| 130 | + # Create filter for the specified date |
| 131 | + filter_conditions = [FieldCondition(key="date", match=MatchValue(value=date))] |
| 132 | + |
| 133 | + date_filter = Filter(must=filter_conditions) |
| 134 | + |
| 135 | + # Query Qdrant for all summaries matching the date using the provided collection name |
| 136 | + search_results = qdrant_client.search( |
| 137 | + collection_name=collection_name, |
| 138 | + query_vector=[0] * 1024, |
| 139 | + query_filter=date_filter, |
| 140 | + limit=100, |
| 141 | + with_payload=True, |
| 142 | + with_vectors=False, |
| 143 | + ) |
| 144 | + |
| 145 | + summaries = [] |
| 146 | + for point in search_results: |
| 147 | + # Extract the summary data from each point |
| 148 | + summary_data = point.payload |
| 149 | + |
| 150 | + # If _node_content is a JSON string, parse it |
| 151 | + if "_node_content" in summary_data and isinstance( |
| 152 | + summary_data["_node_content"], str |
| 153 | + ): |
| 154 | + try: |
| 155 | + node_content = json.loads(summary_data["_node_content"]) |
| 156 | + if extract_text_only: |
| 157 | + summary_data = extract_summary_text(node_content) |
| 158 | + else: |
| 159 | + summary_data["parsed_content"] = node_content |
| 160 | + summary_data["summary_text"] = extract_summary_text( |
| 161 | + node_content |
| 162 | + ) |
| 163 | + except json.JSONDecodeError: |
| 164 | + logging.warning( |
| 165 | + f"Failed to parse _node_content as JSON for point with date {date}" |
| 166 | + ) |
| 167 | + |
| 168 | + summaries.append(summary_data) |
| 169 | + |
| 170 | + logging.info( |
| 171 | + f"Found {len(summaries)} summaries for date {date} in collection {collection_name}" |
| 172 | + ) |
| 173 | + return "\n".join(summaries) if extract_text_only else summaries |
| 174 | + |
| 175 | + except Exception as e: |
| 176 | + logging.error( |
| 177 | + f"Error fetching summaries for date {date} from collection {collection_name}: {str(e)}" |
| 178 | + ) |
| 179 | + raise |
| 180 | + |
| 181 | + |
| 182 | +@activity.defn |
| 183 | +async def fetch_telegram_summaries_by_date_range( |
| 184 | + input: TelegramSummariesRangeActivityInput, |
| 185 | +) -> dict[str, list[dict[str, Any] | str]]: |
| 186 | + """ |
| 187 | + Activity that fetches Telegram summaries for a range of dates from Qdrant. |
| 188 | +
|
| 189 | + Parameters |
| 190 | + ---------- |
| 191 | + input : TelegramSummariesRangeActivityInput |
| 192 | + Input object containing start_date, end_date, collection_name and extract_text_only |
| 193 | +
|
| 194 | + Returns |
| 195 | + ------- |
| 196 | + dict[str, list[dict[str, Any] | str]] |
| 197 | + A dictionary mapping dates to lists of summary objects or a string of summaries |
| 198 | +
|
| 199 | + Raises |
| 200 | + ------ |
| 201 | + ValueError |
| 202 | + If end_date is before start_date or collection_name is not provided |
| 203 | + """ |
| 204 | + start_date = input.start_date |
| 205 | + end_date = input.end_date |
| 206 | + extract_text_only = input.extract_text_only |
| 207 | + collection_name = input.collection_name |
| 208 | + |
| 209 | + if not collection_name: |
| 210 | + raise ValueError("Collection name is required but was not provided") |
| 211 | + |
| 212 | + logging.info( |
| 213 | + f"Fetching summaries for date range: {start_date} to {end_date} from collection: {collection_name}" |
| 214 | + ) |
| 215 | + |
| 216 | + try: |
| 217 | + # Parse the date strings to datetime objects |
| 218 | + start = datetime.strptime(start_date, "%Y-%m-%d").date() |
| 219 | + end = datetime.strptime(end_date, "%Y-%m-%d").date() |
| 220 | + |
| 221 | + # Validate that end_date is not before start_date |
| 222 | + if end < start: |
| 223 | + raise ValueError("End date cannot be before start date") |
| 224 | + |
| 225 | + # Calculate all dates in the range |
| 226 | + date_range = [] |
| 227 | + current = start |
| 228 | + while current <= end: |
| 229 | + date_range.append(current.strftime("%Y-%m-%d")) |
| 230 | + current += timedelta(days=1) |
| 231 | + |
| 232 | + # Fetch summaries for each date |
| 233 | + result = {} |
| 234 | + for date in date_range: |
| 235 | + date_input = TelegramSummariesActivityInput( |
| 236 | + date=date, |
| 237 | + extract_text_only=extract_text_only, |
| 238 | + collection_name=collection_name, |
| 239 | + ) |
| 240 | + summaries = await fetch_telegram_summaries_by_date(date_input) |
| 241 | + result[date] = summaries |
| 242 | + |
| 243 | + return result |
| 244 | + |
| 245 | + except Exception as e: |
| 246 | + logging.error( |
| 247 | + f"Error fetching summaries for date range {start_date} to {end_date} from collection {collection_name}: {str(e)}" |
| 248 | + ) |
| 249 | + raise |
0 commit comments