Merge pull request #17 from TimTProd/add_nicknames

feat: add get_sender_name function to enhance message formatting
This commit is contained in:
Eugene Evstafev 2025-08-06 17:53:53 +01:00 committed by GitHub
commit 574fe149e5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

59
main.py
View file

@ -186,6 +186,24 @@ def format_message(message) -> Dict[str, Any]:
return result return result
def get_sender_name(message) -> str:
"""Helper function to get sender name from a message."""
if not message.sender:
return "Unknown"
# Check for group/channel title first
if hasattr(message.sender, 'title') and message.sender.title:
return message.sender.title
elif hasattr(message.sender, 'first_name'):
# User sender
first_name = getattr(message.sender, 'first_name', '') or ''
last_name = getattr(message.sender, 'last_name', '') or ''
full_name = f"{first_name} {last_name}".strip()
return full_name if full_name else "Unknown"
else:
return "Unknown"
@mcp.tool() @mcp.tool()
async def get_chats(page: int = 1, page_size: int = 20) -> str: async def get_chats(page: int = 1, page_size: int = 20) -> str:
""" """
@ -229,7 +247,8 @@ async def get_messages(chat_id: int, page: int = 1, page_size: int = 20) -> str:
return "No messages found for this page." return "No messages found for this page."
lines = [] lines = []
for msg in messages: for msg in messages:
lines.append(f"ID: {msg.id} | Date: {msg.date} | Message: {msg.message}") sender_name = get_sender_name(msg)
lines.append(f"ID: {msg.id} | {sender_name} | Date: {msg.date} | Message: {msg.message}")
return "\n".join(lines) return "\n".join(lines)
except Exception as e: except Exception as e:
return log_and_format_error( return log_and_format_error(
@ -400,15 +419,10 @@ async def list_messages(
lines = [] lines = []
for msg in messages: for msg in messages:
sender = "" sender_name = get_sender_name(msg)
if msg.sender: message_text = msg.message or '[Media/No text]'
sender_name = getattr(msg.sender, "first_name", "") or getattr(
msg.sender, "title", "Unknown"
)
sender = f"{sender_name} | "
lines.append( lines.append(
f"ID: {msg.id} | {sender}Date: {msg.date} | Message: {msg.message or '[Media/No text]'}" f"ID: {msg.id} | {sender_name} | Date: {msg.date} | Message: {message_text}"
) )
return "\n".join(lines) return "\n".join(lines)
@ -730,11 +744,7 @@ async def get_message_context(chat_id: int, message_id: int, context_size: int =
all_messages.sort(key=lambda m: m.id) all_messages.sort(key=lambda m: m.id)
results = [f"Context for message {message_id} in chat {chat_id}:"] results = [f"Context for message {message_id} in chat {chat_id}:"]
for msg in all_messages: for msg in all_messages:
sender_name = "Unknown" sender_name = get_sender_name(msg)
if msg.sender:
sender_name = getattr(msg.sender, "first_name", "") or getattr(
msg.sender, "title", "Unknown"
)
highlight = " [THIS MESSAGE]" if msg.id == message_id else "" highlight = " [THIS MESSAGE]" if msg.id == message_id else ""
results.append( results.append(
f"ID: {msg.id} | {sender_name} | {msg.date}{highlight}\n{msg.message or '[Media/No text]'}\n" f"ID: {msg.id} | {sender_name} | {msg.date}{highlight}\n{msg.message or '[Media/No text]'}\n"
@ -2028,7 +2038,11 @@ async def search_messages(chat_id: int, query: str, limit: int = 20) -> str:
try: try:
entity = await client.get_entity(chat_id) entity = await client.get_entity(chat_id)
messages = await client.get_messages(entity, limit=limit, search=query) messages = await client.get_messages(entity, limit=limit, search=query)
return "\n".join([f"ID: {m.id} | {m.date} | {m.message}" for m in messages]) lines = []
for msg in messages:
sender_name = get_sender_name(msg)
lines.append(f"ID: {msg.id} | {sender_name} | Date: {msg.date} | Message: {msg.message}")
return "\n".join(lines)
except Exception as e: except Exception as e:
return log_and_format_error( return log_and_format_error(
"search_messages", e, chat_id=chat_id, query=query, limit=limit "search_messages", e, chat_id=chat_id, query=query, limit=limit
@ -2354,7 +2368,11 @@ async def get_history(chat_id: int, limit: int = 100) -> str:
try: try:
entity = await client.get_entity(chat_id) entity = await client.get_entity(chat_id)
messages = await client.get_messages(entity, limit=limit) messages = await client.get_messages(entity, limit=limit)
return "\n".join([f"ID: {m.id} | {m.date} | {m.message}" for m in messages]) lines = []
for msg in messages:
sender_name = get_sender_name(msg)
lines.append(f"ID: {msg.id} | {sender_name} | Date: {msg.date} | Message: {msg.message}")
return "\n".join(lines)
except Exception as e: except Exception as e:
return log_and_format_error("get_history", e, chat_id=chat_id, limit=limit) return log_and_format_error("get_history", e, chat_id=chat_id, limit=limit)
@ -2429,9 +2447,12 @@ async def get_pinned_messages(chat_id: int) -> str:
if not messages: if not messages:
return "No pinned messages found in this chat." return "No pinned messages found in this chat."
return "\n".join( lines = []
[f"ID: {m.id} | {m.date} | {m.message or '[Media/No text]'}" for m in messages] for msg in messages:
) sender_name = get_sender_name(msg)
lines.append(f"ID: {msg.id} | {sender_name} | Date: {msg.date} | Message: {msg.message or '[Media/No text]'}")
return "\n".join(lines)
except Exception as e: except Exception as e:
logger.exception(f"get_pinned_messages failed (chat_id={chat_id})") logger.exception(f"get_pinned_messages failed (chat_id={chat_id})")
return log_and_format_error("get_pinned_messages", e, chat_id=chat_id) return log_and_format_error("get_pinned_messages", e, chat_id=chat_id)