diff --git a/botspot/components/new/contact_manager.py b/botspot/components/new/contact_manager.py index a8a44ab..2c50951 100644 --- a/botspot/components/new/contact_manager.py +++ b/botspot/components/new/contact_manager.py @@ -1,13 +1,32 @@ -from typing import TYPE_CHECKING +"""Contact Manager component for managing user contacts.""" +from datetime import date, datetime, timezone +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from aiogram import Dispatcher, Router +from aiogram.filters import Command +from aiogram.types import Message +from loguru import logger +from pydantic import BaseModel, Field from pydantic_settings import BaseSettings -if TYPE_CHECKING: - from motor.motor_asyncio import AsyncIOMotorCollection # noqa: F401 +from botspot.utils.internal import get_logger +from botspot import commands_menu +from botspot.components.new.queue_manager import QueueItem, get_queue_manager + +logger = get_logger() +# --------------------------------------------- +# region Settings and Model Definitions +# --------------------------------------------- + class ContactManagerSettings(BaseSettings): enabled: bool = False + collection: str = "contacts" + message_parser_enabled: bool = True + random_contact_enabled: bool = True + allow_everyone: bool = False # If False, only friends and admins can use class Config: env_prefix = "BOTSPOT_CONTACT_MANAGER_" @@ -16,20 +35,628 @@ class Config: extra = "ignore" +if TYPE_CHECKING: + from motor.motor_asyncio import AsyncIOMotorCollection # noqa: F401 + from botspot.commands_menu import Visibility + from botspot.components.new.queue_manager import QueueManager + + +class ContactItem(QueueItem): + """Contact information model.""" + + name: str + user_id: Optional[int] = None # Telegram user ID if applicable + phone: Optional[str] = None # Phone number + email: Optional[str] = None # Email address + telegram: Optional[str] = None # Telegram username + birthday: Optional[date] = None # Birthday + notes: Optional[str] = None # Any additional notes + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + owner_id: Optional[int] = None # Who added this contact + + @property + def display_info(self) -> str: + """Return formatted contact information.""" + parts = [f"📇 **{self.name}**"] + + if self.phone: + parts.append(f"📱 Phone: {self.phone}") + if self.email: + parts.append(f"📧 Email: {self.email}") + if self.telegram: + username = self.telegram.strip('@') + parts.append(f"📨 Telegram: @{username}") + if self.birthday: + parts.append(f"🎂 Birthday: {self.birthday.strftime('%d %B')}") + if self.notes: + parts.append(f"📝 Notes: {self.notes}") + if self.user_id: + parts.append(f"👤 User ID: {self.user_id}") + if self.owner_id: + parts.append(f"👤 Owner ID: {self.owner_id}") + + return "\n".join(parts) + + +# --------------------------------------------- +# endregion Settings and Model Definitions +# --------------------------------------------- + + +# --------------------------------------------- +# region Contact Manager Implementation +# --------------------------------------------- + +# TODO: change the way this works, make a collection of Queues, one queue per Owner_id. +# Currently its only viable for personal use(i.e. only one user per bot) class ContactManager: - pass + """Manager for user contacts.""" + + def __init__( + self, + settings: ContactManagerSettings, + queue_manager: Optional["QueueManager"] = None, + ): + """Initialize contact manager with settings.""" + self.settings = settings + self._queue_manager = queue_manager + self.queue = None + + @classmethod + async def create( + cls, + settings: ContactManagerSettings, + queue_manager: Optional["QueueManager"] = None, + ) -> "ContactManager": + """Create a new instance of ContactManager with async initialization.""" + instance = cls(settings, queue_manager) + instance._queue_manager = queue_manager or await get_queue_manager() + instance.queue = instance._queue_manager.create_queue( + key=instance.settings.collection, + item_model=ContactItem, + use_timestamp=True, + use_priority=False, + use_done=False, + ) + return instance + + async def add_contact(self, contact: ContactItem, owner_id: Optional[int] = None) -> bool: + """Add a new contact to the queue.""" + try: + if owner_id: + contact.owner_id = owner_id + await self.queue.add_item(contact, user_id=owner_id) + logger.info(f"Added contact: {contact.name}") + return True + except Exception as e: + logger.error(f"Error adding contact: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to add contact") from e + + async def update_contact( + self, contact_id: str, data: Dict[str, Any], owner_id: Optional[int] = None + ) -> bool: + """Update an existing contact.""" + try: + # Get the current contact + contact = await self.get_contact_by_id(contact_id, owner_id) + if not contact: + from botspot.core.errors import ContactDataError + raise ContactDataError("Contact not found") + + # Update fields + for key, value in data.items(): + if hasattr(contact, key): + setattr(contact, key, value) + contact.updated_at = datetime.now(timezone.utc) + + # Update in queue + result = await self.queue.update_record(contact_id, contact.model_dump()) + if not result: + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to update contact") + return True + except Exception as e: + logger.error(f"Error updating contact: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to update contact") from e + + async def delete_contact(self, contact_id: str, owner_id: Optional[int] = None) -> bool: + """Delete a contact from the queue.""" + try: + result = await self.queue.delete_record(contact_id, user_id=owner_id) + if not result: + from botspot.core.errors import ContactDataError + raise ContactDataError("Contact not found or deletion failed") + return True + except Exception as e: + logger.error(f"Error deleting contact: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to delete contact") from e + + async def get_contact_by_id( + self, contact_id: str, owner_id: Optional[int] = None + ) -> Optional[ContactItem]: + """Get a contact by its ID.""" + try: + contact_data = await self.queue.find({"_id": contact_id}, user_id=owner_id) + if not contact_data: + return None + return ContactItem(**contact_data) + except Exception as e: + logger.error(f"Error getting contact: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to get contact") from e + + async def find_contacts( + self, + query: Dict[str, Any], + limit: int = 10, + owner_id: Optional[int] = None, + ) -> List[ContactItem]: + """Find contacts matching the given query.""" + try: + contacts_data = await self.queue.find_many( + query, + limit=limit, + user_id=owner_id + ) + return [ContactItem(**contact) for contact in contacts_data] + except Exception as e: + logger.error(f"Error finding contacts: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to find contacts") from e + + async def search_contacts( + self, + text: str, + limit: int = 10, + owner_id: Optional[int] = None, + ) -> List[ContactItem]: + """Search for contacts by name, email, etc.""" + try: + # Create a query with $or for text matching across multiple fields + query: Dict[str, Any] = { + "$or": [ + {"name": {"$regex": text, "$options": "i"}}, + {"email": {"$regex": text, "$options": "i"}}, + {"phone": {"$regex": text, "$options": "i"}}, + {"telegram": {"$regex": text, "$options": "i"}}, + {"notes": {"$regex": text, "$options": "i"}} + ] + } + return await self.find_contacts( + query, + limit=limit, + owner_id=owner_id + ) + except Exception as e: + logger.error(f"Error searching contacts: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to search contacts") from e + + async def get_random_contact( + self, owner_id: Optional[int] = None + ) -> Optional[ContactItem]: + """Get a random contact from the queue.""" + try: + contact_data = await self.queue.get_random_record(user_id=owner_id) + if not contact_data: + return None + return ContactItem(**contact_data) + except Exception as e: + logger.error(f"Error getting random contact: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to get random contact") from e + + async def parse_contact_with_llm( + self, text: str, owner_id: Optional[int] = None + ) -> Optional[ContactItem]: + """Parse contact information from text using LLM.""" + try: + from botspot.utils.deps_getters import get_llm_provider + + llm = get_llm_provider() + if not llm: + from botspot.core.errors import ContactDataError + raise ContactDataError("LLM provider not available") + + # System message for the LLM + system_message = """ + You are a contact information parser. Extract contact details from the user's message. + Return ONLY a valid JSON object with the following structure: + { + "name": "Person's name", + "phone": "Phone number (if found)", + "email": "Email (if found)", + "telegram": "Telegram username (if found)", + "birthday": "Birthday in YYYY-MM-DD format (if found)", + "notes": "Any additional information", + "owner_id": null + } + + If a field is not found, set it to null. + The name field is required - make your best guess if not explicit. + """ + + # Create a temporary model for parsing + class TempContact(BaseModel): + name: str + phone: Optional[str] = None + email: Optional[str] = None + telegram: Optional[str] = None + birthday: Optional[date] = None + notes: Optional[str] = None + owner_id: Optional[int] = None + + # Use structured output for better parsing + temp_contact = await llm.aquery_llm_structured( + prompt=text, + output_schema=TempContact, + user=owner_id, + system_message=system_message, + temperature=0.2, # Lower temperature for more deterministic output + ) + + # Create a new dictionary with owner_id if provided + contact_data = temp_contact.model_dump() + if owner_id is not None: + contact_data["owner_id"] = owner_id + + # Convert to ContactItem + return ContactItem.model_validate(contact_data) + except Exception as e: + logger.error(f"Error parsing contact with LLM: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to parse contact information") from e + + async def get_contacts_by_owner( + self, owner_id: int, limit: int = 10 + ) -> List[ContactItem]: + """Get all contacts owned by a specific user.""" + try: + query = {"owner_id": owner_id} + return await self.find_contacts( + query, + limit=limit, + owner_id=owner_id + ) + except Exception as e: + logger.error(f"Error getting contacts by owner: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to get contacts by owner") from e + + async def get_contacts_by_telegram( + self, telegram: str, limit: int = 10 + ) -> List[ContactItem]: + """Get contacts by Telegram username.""" + try: + query = {"telegram": telegram} + return await self.find_contacts( + query, + limit=limit, # Already validated as int + ) + except Exception as e: + logger.error(f"Error getting contacts by telegram: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to get contacts by telegram") from e + + async def get_contacts_by_phone( + self, phone: str, limit: int = 10 + ) -> List[ContactItem]: + """Get contacts by phone number.""" + try: + query = {"phone": phone} + return await self.find_contacts( + query, + limit=limit, # Already validated as int + ) + except Exception as e: + logger.error(f"Error getting contacts by phone: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to get contacts by phone") from e + + async def get_contacts_by_email( + self, email: str, limit: int = 10 + ) -> List[ContactItem]: + """Get contacts by email address.""" + try: + query = {"email": email} + return await self.find_contacts( + query, + limit=limit, # Already validated as int + ) + except Exception as e: + logger.error(f"Error getting contacts by email: {e}") + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to get contacts by email") from e + + +# --------------------------------------------- +# endregion Contact Manager Implementation +# --------------------------------------------- + + +# --------------------------------------------- +# region Command Handlers +# --------------------------------------------- + +def setup_command_handlers(dp: Dispatcher, manager: ContactManager, settings: ContactManagerSettings): + """Set up command handlers for contact management.""" + from botspot.commands_menu import Visibility + + router = Router(name="contact_manager") + + # Add command to add contact + @commands_menu.add_command("add_contact", "Add a new contact", Visibility.PUBLIC) + @router.message(Command("add_contact")) + async def add_contact_cmd(message: Message): + """Add a new contact.""" + # Extract contact info after the command + if len(message.text.split()) <= 1: + await message.reply( + "Please provide contact information after the command. Example:\n" + "/add_contact John Doe, phone: 555-1234, email: john@example.com" + ) + return + + # Extract contact text (everything after the command) + contact_text = message.text.split(maxsplit=1)[1] + + # Parse contact with LLM + contact = await manager.parse_contact_with_llm(contact_text, message.from_user.id) + + if not contact: + await message.reply("I couldn't parse the contact information. Please try again with more details.") + return + + # Check if any required fields are missing + missing_fields = [] + if not contact.phone and not contact.email and not contact.telegram: + missing_fields.append("a way to contact (phone, email, or telegram)") + + # If missing fields, ask user to provide them + if missing_fields: + missing_info = ", ".join(missing_fields) + await message.reply( + f"I need {missing_info} for this contact. Please add this information and try again." + ) + return + + # Add contact to database + success = await manager.add_contact(contact, message.from_user.id) + + if success: + await message.reply( + f"✅ Contact added successfully!\n\n{contact.display_info}", + parse_mode="Markdown" + ) + else: + await message.reply("❌ Failed to add contact. Please try again.") + + # Add command to find contact + @commands_menu.add_command("find_contact", "Find contacts by name, phone, etc.", Visibility.PUBLIC) + @router.message(Command("find_contact")) + async def find_contact_cmd(message: Message): + """Find contacts matching a search term.""" + if len(message.text.split()) <= 1: + await message.reply( + "Please provide a search term after the command. Example:\n" + "/find_contact John" + ) + return + + # Extract search text + search_text = message.text.split(maxsplit=1)[1] + # Search for contacts + contacts = await manager.search_contacts(search_text, message.from_user.id) + + if not contacts: + await message.reply(f"No contacts found matching '{search_text}'.") + return + + # Format results + if len(contacts) == 1: + await message.reply( + f"Found 1 contact:\n\n{contacts[0].display_info}", + parse_mode="Markdown" + ) + else: + results = "\n\n".join([c.display_info for c in contacts[:5]]) + count_msg = f"Found {len(contacts)} contacts" + (", showing first 5:" if len(contacts) > 5 else ":") + await message.reply( + f"{count_msg}\n\n{results}", + parse_mode="Markdown" + ) + + # Add command for random contact + if settings.random_contact_enabled: + @commands_menu.add_command("random_contact", "Get a random contact", Visibility.PUBLIC) + @router.message(Command("random_contact")) + async def random_contact_cmd(message: Message): + """Get a random contact.""" + contact = await manager.get_random_contact(message.from_user.id) + + if not contact: + await message.reply("You don't have any contacts saved yet.") + return + + await message.reply( + f"Random contact:\n\n{contact.display_info}", + parse_mode="Markdown" + ) + + # Handle message parsing if enabled + if settings.message_parser_enabled: + @router.message() + async def parse_contact_message(message: Message): + """Parse potential contact information from regular messages.""" + # Skip commands + if message.text and message.text.startswith('/'): + return + + # Simple heuristic: Look for some indicators of contact info + # This is intentionally simple as the LLM will do the heavy lifting + if not message.text or not any(kw in message.text.lower() for kw in [ + 'phone', 'email', 'telegram', '@', 'contact', 'reach', 'name' + ]): + return + + # Try to parse as contact + contact = await manager.parse_contact_with_llm(message.text, message.from_user.id) + + # If we have a name and at least one contact method, offer to save + if contact and contact.name and (contact.phone or contact.email or contact.telegram): + await message.reply( + f"I detected contact information. Would you like to save this contact?\n\n" + f"{contact.display_info}\n\n" + f"Reply with /add_contact to confirm.", + parse_mode="Markdown" + ) + + # Include router in dispatcher + dp.include_router(router) + + +# --------------------------------------------- +# endregion Command Handlers +# --------------------------------------------- + + +# --------------------------------------------- +# region Initialization and Setup +# --------------------------------------------- + +def setup_dispatcher(dp: Dispatcher, **kwargs): + """Setup contact manager component in the dispatcher.""" + settings = ContactManagerSettings(**kwargs) + + if not settings.enabled: + logger.info("Contact Manager component is disabled") + return dp + + # Make sure dependencies are available + from botspot.core.dependency_manager import get_dependency_manager + + deps = get_dependency_manager() + + if deps.contact_manager is None: + logger.warning("Contact Manager component is not initialized") + return dp + + logger.info("Setting up Contact Manager component") + + # Set up command handlers + setup_command_handlers(dp, deps.contact_manager, settings) -def setup_dispatcher(dp): return dp -def initialize(settings: ContactManagerSettings) -> ContactManager: - pass +async def initialize(settings: ContactManagerSettings) -> ContactManager: + """Initialize the Contact Manager component.""" + if not settings.enabled: + logger.info("Contact Manager component is disabled") + return ContactManager(settings) # Return empty manager instead of None + logger.info("Initializing Contact Manager component") + manager = await ContactManager.create(settings) + return manager -def get_contact_manager(): + +def get_contact_manager() -> ContactManager: + """Get the Contact Manager from dependency manager.""" from botspot.core.dependency_manager import get_dependency_manager deps = get_dependency_manager() + + if deps.contact_manager is None: + raise RuntimeError("Contact Manager is not initialized") + return deps.contact_manager + + +# --------------------------------------------- +# endregion Initialization and Setup +# --------------------------------------------- + + +# --------------------------------------------- +# region Utils +# --------------------------------------------- + + +async def add_contact(contact: ContactItem, owner_id: Optional[int] = None) -> bool: + """Add a new contact to the database.""" + manager = get_contact_manager() + return await manager.add_contact(contact, owner_id) + + +async def update_contact(contact_id: str, data: Dict[str, Any], owner_id: Optional[int] = None) -> bool: + """Update an existing contact.""" + manager = get_contact_manager() + return await manager.update_contact(contact_id, data, owner_id) + + +async def delete_contact(contact_id: str, owner_id: Optional[int] = None) -> bool: + """Delete a contact by ID.""" + manager = get_contact_manager() + return await manager.delete_contact(contact_id, owner_id) + + +async def get_contact_by_id(contact_id: str, owner_id: Optional[int] = None) -> Optional[ContactItem]: + """Get a contact by ID.""" + manager = get_contact_manager() + return await manager.get_contact_by_id(contact_id, owner_id) + + +async def find_contacts(query: Dict[str, Any], owner_id: Optional[int] = None, limit: int = 10) -> List[ContactItem]: + """Find contacts matching the query.""" + manager = get_contact_manager() + return await manager.find_contacts(query, limit, owner_id) + + +async def search_contacts(text: str, owner_id: Optional[int] = None, limit: int = 10) -> List[ContactItem]: + """Search for contacts by name, email, etc.""" + manager = get_contact_manager() + return await manager.search_contacts(text, limit, owner_id) + + +async def get_random_contact(owner_id: Optional[int] = None) -> Optional[ContactItem]: + """Get a random contact.""" + manager = get_contact_manager() + return await manager.get_random_contact(owner_id) + + +async def parse_contact_with_llm(text: str, owner_id: Optional[int] = None) -> Optional[ContactItem]: + """Parse contact information from text using LLM.""" + manager = get_contact_manager() + return await manager.parse_contact_with_llm(text, owner_id) + + +async def get_contacts_by_owner(owner_id: int, limit: int = 10) -> List[ContactItem]: + """Get all contacts owned by a specific user.""" + manager = get_contact_manager() + return await manager.get_contacts_by_owner(owner_id, limit) + + +async def get_contacts_by_telegram(telegram: str, limit: int = 10) -> List[ContactItem]: + """Get contacts by Telegram username.""" + manager = get_contact_manager() + return await manager.get_contacts_by_telegram(telegram, limit) + + +async def get_contacts_by_phone(phone: str, limit: int = 10) -> List[ContactItem]: + """Get contacts by phone number.""" + manager = get_contact_manager() + return await manager.get_contacts_by_phone(phone, limit) + + +async def get_contacts_by_email(email: str, limit: int = 10) -> List[ContactItem]: + """Get contacts by email address.""" + manager = get_contact_manager() + return await manager.get_contacts_by_email(email, limit) + +# --------------------------------------------- +# endregion Utils diff --git a/botspot/components/new/queue_manager.py b/botspot/components/new/queue_manager.py index 6b7b1a0..65a0a26 100644 --- a/botspot/components/new/queue_manager.py +++ b/botspot/components/new/queue_manager.py @@ -88,6 +88,110 @@ async def add_item(self, item: T, user_id: Optional[int] = None): doc = self.enrich_item(item, user_id) await self.collection.insert_one(doc) + async def update_record(self, record_id: Any, data: Dict[str, Any]) -> bool: + """ + Update a record by its ID. + + Args: + record_id: The record ID to update + data: Dictionary containing the fields to update + + Returns: + True if update was successful, False otherwise + """ + try: + # Update timestamp if enabled + if self.use_timestamp: + data["timestamp"] = datetime.now() + + # Update the document + result = await self.collection.update_one( + {"_id": record_id}, + {"$set": data} + ) + return result.modified_count > 0 + except Exception as e: + logger.error(f"Error updating record: {e}") + return False + + async def find(self, query: Dict[str, Any], user_id: Optional[int] = None) -> Optional[Dict[str, Any]]: + """ + Find a single record matching the given query. + + Args: + query: MongoDB query dictionary + user_id: Optional user ID to filter by + + Returns: + The record as a dictionary if found, None otherwise + """ + try: + if not self.single_user_mode and user_id is None: + from botspot.core.errors import QueuePermissionError + raise QueuePermissionError("user_id is required unless single_user_mode is enabled") + + if user_id is not None: + query["user_id"] = user_id + + return await self.collection.find_one(query) + except Exception as e: + logger.error(f"Error finding record: {e}") + return None + + async def find_many(self, query: Dict[str, Any], user_id: Optional[int] = None, limit: Optional[int] = None) -> List[Dict[str, Any]]: + """ + Find multiple records matching the given query. + + Args: + query: MongoDB query dictionary + user_id: Optional user ID to filter by + limit: Maximum number of records to return + + Returns: + List of records as dictionaries + """ + try: + if not self.single_user_mode and user_id is None: + from botspot.core.errors import QueuePermissionError + raise QueuePermissionError("user_id is required unless single_user_mode is enabled") + + if user_id is not None: + query["user_id"] = user_id + + cursor = self.collection.find(query) + if limit is not None: + cursor = cursor.limit(limit) + return await cursor.to_list(length=limit) + except Exception as e: + logger.error(f"Error finding records: {e}") + return [] + + async def delete_record(self, record_id: Any, user_id: Optional[int] = None) -> bool: + """ + Delete a record by its ID. + + Args: + record_id: The record ID to delete + user_id: Optional user ID to verify ownership + + Returns: + True if deletion was successful, False otherwise + """ + try: + if not self.single_user_mode and user_id is None: + from botspot.core.errors import QueuePermissionError + raise QueuePermissionError("user_id is required unless single_user_mode is enabled") + + query = {"_id": record_id} + if user_id is not None: + query["user_id"] = user_id + + result = await self.collection.delete_one(query) + return result.deleted_count > 0 + except Exception as e: + logger.error(f"Error deleting record: {e}") + return False + async def get_items( self, user_id: Optional[int] = None, @@ -123,6 +227,41 @@ async def get_records( cursor = cursor.limit(limit) return await cursor.to_list(length=limit) + async def get_random_record(self, user_id: Optional[int] = None) -> Optional[Dict[str, Any]]: + """ + Get a random record from the queue. + + Args: + user_id: Optional user ID to filter by + + Returns: + A random record as a dictionary, or None if no records found + """ + try: + if not self.single_user_mode and user_id is None: + from botspot.core.errors import QueuePermissionError + raise QueuePermissionError("user_id is required unless single_user_mode is enabled") + + # Build the pipeline + pipeline = [] + + # Add user filter if needed + if user_id is not None: + pipeline.append({"$match": {"user_id": user_id}}) + + # Add sample stage to get a random document + pipeline.append({"$sample": {"size": 1}}) + + # Execute the aggregation + cursor = self.collection.aggregate(pipeline) + async for doc in cursor: + return doc + + return None + except Exception as e: + logger.error(f"Error getting random record: {e}") + return None + class QueueManager: def __init__(self, settings: QueueManagerSettings, single_user_mode: Optional[bool] = None): diff --git a/botspot/contact_manager.py b/botspot/contact_manager.py new file mode 100644 index 0000000..1dbf0cc --- /dev/null +++ b/botspot/contact_manager.py @@ -0,0 +1,30 @@ +from botspot.components.new.contact_manager import ( + add_contact, + delete_contact, + find_contacts, + get_contact_by_id, + get_random_contact, + parse_contact_with_llm, + search_contacts, + update_contact, + get_contacts_by_owner, + get_contacts_by_telegram, + get_contacts_by_phone, + get_contacts_by_email, +) + +__all__ = [ + "add_contact", + "update_contact", + "delete_contact", + "get_contact_by_id", + "find_contacts", + "search_contacts", + "get_random_contact", + "parse_contact_with_llm", + "get_contacts_by_owner", + "get_contacts_by_telegram", + "get_contacts_by_phone", + "get_contacts_by_email", +] + diff --git a/botspot/core/bot_manager.py b/botspot/core/bot_manager.py index 65f6318..103db9a 100644 --- a/botspot/core/bot_manager.py +++ b/botspot/core/bot_manager.py @@ -17,6 +17,7 @@ auto_archive, chat_binder, chat_fetcher, + contact_manager, llm_provider, queue_manager, ) @@ -66,6 +67,9 @@ def __init__( if self.settings.llm_provider.enabled: self.deps.llm_provider = llm_provider.initialize(self.settings.llm_provider) + + if self.settings.contact_manager.enabled: + self.deps.contact_manager = contact_manager.initialize(self.settings.contact_manager) if self.settings.queue_manager.enabled: self.deps.queue_manager = queue_manager.initialize(self.settings.queue_manager) @@ -120,6 +124,9 @@ def setup_dispatcher(self, dp: Dispatcher): if self.settings.llm_provider.enabled: llm_provider.setup_dispatcher(dp) + + if self.settings.contact_manager.enabled: + contact_manager.setup_dispatcher(dp) if self.settings.queue_manager.enabled: queue_manager.setup_dispatcher(dp) diff --git a/botspot/core/botspot_settings.py b/botspot/core/botspot_settings.py index 83af6ca..d2c8388 100644 --- a/botspot/core/botspot_settings.py +++ b/botspot/core/botspot_settings.py @@ -15,6 +15,7 @@ from botspot.components.new.auto_archive import AutoArchiveSettings from botspot.components.new.chat_binder import ChatBinderSettings from botspot.components.new.chat_fetcher import ChatFetcherSettings +from botspot.components.new.contact_manager import ContactManagerSettings from botspot.components.new.llm_provider import LLMProviderSettings from botspot.components.new.queue_manager import QueueManagerSettings from botspot.components.qol.bot_commands_menu import BotCommandsMenuSettings @@ -63,6 +64,7 @@ def friends(self) -> List[str]: chat_binder: ChatBinderSettings = ChatBinderSettings() chat_fetcher: ChatFetcherSettings = ChatFetcherSettings() llm_provider: LLMProviderSettings = LLMProviderSettings() + contact_manager: ContactManagerSettings = ContactManagerSettings() queue_manager: QueueManagerSettings = QueueManagerSettings() auto_archive: AutoArchiveSettings = AutoArchiveSettings() diff --git a/botspot/core/dependency_manager.py b/botspot/core/dependency_manager.py index af17119..1cfaba9 100644 --- a/botspot/core/dependency_manager.py +++ b/botspot/core/dependency_manager.py @@ -15,6 +15,7 @@ from botspot.components.new.auto_archive import AutoArchive from botspot.components.new.chat_binder import ChatBinder from botspot.components.new.chat_fetcher import ChatFetcher + from botspot.components.new.contact_manager import ContactManager from botspot.components.new.llm_provider import LLMProvider from botspot.components.new.queue_manager import QueueManager @@ -39,6 +40,7 @@ def __init__( self._user_manager = None self._chat_binder = None self._llm_provider = None + self._contact_manager = None self._queue_manager = None self._chat_fetcher = None self._auto_archive = None @@ -141,6 +143,16 @@ def llm_provider(self) -> "LLMProvider": @llm_provider.setter def llm_provider(self, value): self._llm_provider = value + + @property + def contact_manager(self) -> "ContactManager": + if self._contact_manager is None: + raise RuntimeError("Contact Manager is not initialized") + return self._contact_manager + + @contact_manager.setter + def contact_manager(self, value): + self._contact_manager = value @property def queue_manager(self) -> "QueueManager": diff --git a/botspot/utils/deps_getters.py b/botspot/utils/deps_getters.py index 37b7ea2..e2401f5 100644 --- a/botspot/utils/deps_getters.py +++ b/botspot/utils/deps_getters.py @@ -14,6 +14,7 @@ from botspot.components.main.telethon_manager import get_telethon_manager from botspot.components.new.chat_binder import get_chat_binder from botspot.components.new.chat_fetcher import get_chat_fetcher +from botspot.components.new.contact_manager import get_contact_manager from botspot.components.new.llm_provider import get_llm_provider from botspot.components.new.queue_manager import get_queue_manager @@ -78,6 +79,7 @@ async def get_telethon_client( "get_mongo_client", "get_chat_binder", "get_chat_fetcher", + "get_contact_manager", "get_queue_manager", "get_llm_provider", ] diff --git a/dev/workalong_contact_manager/rework_protocol.md b/dev/workalong_contact_manager/rework_protocol.md new file mode 100644 index 0000000..e28f6e4 --- /dev/null +++ b/dev/workalong_contact_manager/rework_protocol.md @@ -0,0 +1,58 @@ +This is a protocol for refactoring / reworking code components with claude-code or other +ai tools. +The main goal of this protocol is to ensure the codebase doesn't bloat and grow riddled +with duplicated code during AI refactoring / reworks. + +# Pre-rework + +- Commit. Make sure everything is committed + - Optional: Create a new branch for the rework +- Optional: Check tests, if present. + +# Simple protocol: + +If the task is clear, go directly to coding phase + +# Complex protocol: + +If there's a complex task - first, read spec.md or ask user for a list of requirements + +- After that, evaluate the current solution on the match with specs + - Which additional featuers were implemented that are not present in the spec + - What is the likely reason were they added + - + - Which components / features explicitly listed in the spec are missing + - How difficult it is to add this + - write to workalong.md + - proceed to coding + +## Coding: + +- Before coding, lay out a plan to the user in clear terms. + - Which components / features will be added + - Which modified + - Which moved / removed + - Make an explicit enumeration for user to specify which steps are approved and + which are declined + - Each item should be formulated in as simple terms as possible, 1 line maximum per + item, not longer than a few words +- Always remove duplicate code or alternative / previous implementations of the same + feature + - After making a change, call git diff and make sure file is in a desired target + state and the changes you planned are correctly reflected +- proceed with implementing each item one by one and track progress with checkboxes in + workalong.md + - [x] item 1 - keep to original item descriptions, DO NOT ADD SUB-ITEMS. List which + files were affected for this feature. + +AI Issue resolution + +1) Failed deletions + Sometimes AI applier fails to delete code according to instructions + This results in really confusing situations with duplicate functions / methods + present in multiple places across the codebase + To mitigate that + +- Check file diff after each modifications and make sure it reflects the changes you've + made + diff --git a/dev/workalong_contact_manager/spec.md b/dev/workalong_contact_manager/spec.md new file mode 100644 index 0000000..21b8264 --- /dev/null +++ b/dev/workalong_contact_manager/spec.md @@ -0,0 +1,21 @@ +# Contact Manager (for People Bot) + +- **Main Feature**: Manages user contacts with a basic schema (name, phone, email, telegram, birthday) stored in MongoDB. +- **Demo Bot**: Send a message (e.g., "Jane, 555-5678"); it auto-parses and saves as a contact, asking for missing fields if needed. +- **Useful Bot**: **People Bot** - Auto-ingests contacts from chat messages, parses with LLM, stores them, and lets you retrieve random contacts. + +## Developer Notes + +0) need to decide whether to implement contact manager on top of queue manager or as a completely separate + +1) for contacts - people should never get out of the queue + +2) AI loves to invent features here. Forbid it. Make sure it is minimal possible implementation to start with - and then extend that + +3) need some way to easily ingest contacts: from telegram, from iphone, from gmail, from my old notion / remnote / obsidian. + +4) I guess we do this on top of queue manager. Just need data model. The question is how will we add fields later? + +5) how to make this extensible into a full-fledged crm? + +6) bring over features to send messages to people - from ef random coffee bot. /Users/petrlavrov/work/experiments/ef-community-bot-rc-notifier diff --git a/dev/workalong_contact_manager/workalong.md b/dev/workalong_contact_manager/workalong.md new file mode 100644 index 0000000..cc954d9 --- /dev/null +++ b/dev/workalong_contact_manager/workalong.md @@ -0,0 +1,50 @@ +# Contact Manager Implementation Analysis + +## Features/Methods/Fields Added Beyond the Spec + +1. **Settings beyond basic enablement:** + - `message_parser_enabled` - Toggle for auto-parsing contacts from messages + - `random_contact_enabled` - Toggle for random contact feature + - `allow_everyone` - Access control setting + - `collection` - MongoDB collection name setting + +2. **Contact Model Fields beyond basic schema:** + - `user_id` - Telegram user ID if applicable + - `notes` - Additional notes field + - `created_at` - Timestamp when contact was created + - `updated_at` - Timestamp when contact was last updated + - `owner_id` - Who added this contact + +3. **Additional Methods in ContactManager:** + - `update_contact()` - Update existing contacts + - `delete_contact()` - Delete contacts + - `get_contact_by_id()` - Get specific contact by ID + - `find_contacts()` - Generic method to find contacts by any criteria + - `search_contacts()` - Specialized search method with text matching + +4. **Commands and User Interface:** + - `/find_contact` command - Search functionality + - Detailed response formatting with emoji and markdown + +5. **Other Features:** + - Visibility control through command menu + - Contact ownership and per-user contacts + +## Features Mentioned in Spec but Not Implemented + +1. **Queue Manager Integration:** + - Notes mention "implement contact manager on top of queue manager" (line 9) + - "for contacts - people should never get out of the queue" (line 11) + - Not implemented as a queue-based system + +2. **External Data Import:** + - "need some way to easily ingest contacts: from telegram, from iphone, from gmail, from my old notion / remnote / obsidian" (line 15) + - No import functionality from external sources implemented + +3. **Message Feature:** + - "bring over features to send messages to people - from ef random coffee bot" (line 21) + - No functionality to send messages to contacts + +4. **Extensibility for Full CRM:** + - "how to make this extensible into a full-fledged crm?" (line 19) + - While the implementation is somewhat extensible, no explicit CRM-focused extensibility mechanisms \ No newline at end of file diff --git a/examples/components_examples/contact_manager_demo/README.md b/examples/components_examples/contact_manager_demo/README.md new file mode 100644 index 0000000..5278b13 --- /dev/null +++ b/examples/components_examples/contact_manager_demo/README.md @@ -0,0 +1,53 @@ +# Contact Manager Demo + +This example demonstrates how to use the Contact Manager component to create a simple contact management bot. + +## Features + +- Add contacts with a basic schema (name, phone, email, telegram, birthday) +- Auto-parse contact information from regular messages using LLM +- Search contacts by name, email, phone, or other fields +- Get random contacts from your saved list + +## Setup + +1. Copy `sample.env` to `.env` and update with your settings: + ```bash + cp sample.env .env + ``` + +2. Edit `.env` file to add your Telegram bot token and other settings + +3. Make sure MongoDB is running (required for this component) + +4. Install dependencies: + ```bash + poetry install + ``` + +5. Run the bot: + ```bash + python bot.py + ``` + +## Usage + +- `/add_contact [contact info]` - Add a new contact + Example: `/add_contact John Doe, phone: 555-1234, email: john@example.com` + +- `/find_contact [search term]` - Find contacts matching search term + Example: `/find_contact John` + +- `/random_contact` - Get a random contact from your saved list + +- The bot will also automatically detect potential contact information in regular messages and offer to save it. + +## Behind the Scenes + +This demo uses: +- **MongoDB** for storing contacts +- **LLM Provider** for parsing contact details from text +- **User Data** for managing user access +- **Bot Commands Menu** for registering commands + +The Contact Manager component is designed to be extensible for more advanced CRM-like functionality in the future. \ No newline at end of file diff --git a/examples/components_examples/contact_manager_demo/bot.py b/examples/components_examples/contact_manager_demo/bot.py new file mode 100644 index 0000000..7cb07a0 --- /dev/null +++ b/examples/components_examples/contact_manager_demo/bot.py @@ -0,0 +1,58 @@ +import asyncio +import logging +import os + +from aiogram import Bot, Dispatcher +from aiogram.enums import ParseMode +from aiogram.filters import Command +from aiogram.types import Message +from dotenv import load_dotenv + +from botspot.core.bot_manager import BotManager +from botspot.utils.deps_getters import get_contact_manager + + +# Configure logging +logging.basicConfig(level=logging.INFO) + +# Load environment variables +load_dotenv() + + +async def main(): + # Initialize Bot and Dispatcher + bot = Bot(token=os.getenv("BOT_TOKEN"), parse_mode=ParseMode.MARKDOWN) + dp = Dispatcher() + + # Initialize BotManager with our bot and dispatcher + manager = BotManager(bot=bot, dispatcher=dp) + + # Set up the dispatcher with all enabled components + manager.setup_dispatcher(dp) + + # Add custom command handler (optional, Contact Manager already adds its own handlers) + @dp.message(Command("contacts_help")) + async def contacts_help_cmd(message: Message): + help_text = """ +📇 **Contact Manager Help** + +This bot helps you manage your contacts. You can: + +- `/add_contact [info]` - Add a new contact + Example: `/add_contact John Doe, phone: 555-1234, email: john@example.com` + +- `/find_contact [search]` - Find contacts + Example: `/find_contact John` + +- `/random_contact` - Get a random contact + +You can also simply send a message with contact info, and I'll offer to save it! + """ + await message.reply(help_text) + + # Start the bot + await dp.start_polling(bot) + + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/examples/components_examples/contact_manager_demo/sample.env b/examples/components_examples/contact_manager_demo/sample.env new file mode 100644 index 0000000..77b444b --- /dev/null +++ b/examples/components_examples/contact_manager_demo/sample.env @@ -0,0 +1,37 @@ +# Bot token (required) +BOT_TOKEN=your_bot_token_here + +# MongoDB settings (required for Contact Manager) +BOTSPOT_MONGO_DATABASE_ENABLED=true +BOTSPOT_MONGO_DATABASE_CONN_STR=mongodb://localhost:27017 +BOTSPOT_MONGO_DATABASE_DATABASE=botspot_contacts + +# User data settings (required for Contact Manager) +BOTSPOT_USER_DATA_ENABLED=true +BOTSPOT_USER_DATA_COLLECTION=users + +# Contact Manager settings +BOTSPOT_CONTACT_MANAGER_ENABLED=true +BOTSPOT_CONTACT_MANAGER_COLLECTION=contacts +BOTSPOT_CONTACT_MANAGER_MESSAGE_PARSER_ENABLED=true +BOTSPOT_CONTACT_MANAGER_RANDOM_CONTACT_ENABLED=true +BOTSPOT_CONTACT_MANAGER_ALLOW_EVERYONE=true + +# LLM Provider settings (required for Contact Manager) +BOTSPOT_LLM_PROVIDER_ENABLED=true +BOTSPOT_LLM_PROVIDER_DEFAULT_MODEL=claude-3.7 +BOTSPOT_LLM_PROVIDER_ALLOW_EVERYONE=true + +# Admin and friends settings +BOTSPOT_ADMINS_STR=@your_username +BOTSPOT_FRIENDS_STR= + +# Bot Commands Menu +BOTSPOT_BOT_COMMANDS_MENU_ENABLED=true + +# QoL Components +BOTSPOT_BOT_INFO_ENABLED=true +BOTSPOT_PRINT_BOT_URL_ENABLED=true + +# Error Handler +BOTSPOT_ERROR_HANDLING_ENABLED=true \ No newline at end of file diff --git a/tests/components/new/test_contact_manager.py b/tests/components/new/test_contact_manager.py new file mode 100644 index 0000000..f521624 --- /dev/null +++ b/tests/components/new/test_contact_manager.py @@ -0,0 +1,388 @@ +import asyncio +import datetime +import pytest +import pytest_asyncio +from unittest.mock import AsyncMock, MagicMock, patch, create_autospec, Mock + +from botspot.components.new.contact_manager import ContactItem, ContactManager, ContactManagerSettings, initialize, setup_command_handlers +from botspot.core.errors import ContactDataError + + +@pytest.fixture +def contact_manager_settings(): + """Fixture to create contact manager settings""" + return ContactManagerSettings( + enabled=True, + collection="test_contacts", + message_parser_enabled=True, + random_contact_enabled=True + ) + + +@pytest.fixture +def mock_queue(): + """Create a mock queue for testing.""" + queue = AsyncMock() + queue.add_item = AsyncMock(return_value=True) + queue.update_record = AsyncMock(return_value=True) + queue.delete_record = AsyncMock(return_value=True) + queue.find = AsyncMock(return_value=None) + queue.find_many = AsyncMock(return_value=[]) + queue.get_random_record = AsyncMock(return_value=None) + return queue + + +@pytest.fixture +def mock_queue_manager(mock_queue): + """Create a mock queue manager for testing.""" + manager = Mock() + manager.create_queue = Mock(return_value=mock_queue) + return manager + + +@pytest.fixture +def mock_dispatcher(): + """Fixture to create a mock dispatcher""" + return MagicMock() + + +@pytest_asyncio.fixture +async def contact_manager(mock_queue_manager): + """Create a contact manager instance for testing.""" + settings = ContactManagerSettings(enabled=True) + return await ContactManager.create(settings, queue_manager=mock_queue_manager) + + +@pytest.fixture +def sample_contact(): + """Create a sample contact for testing.""" + return ContactItem( + data="Test Contact", + name="Test Contact", + phone="555-1234", + email="test@example.com", + telegram="@testuser", + owner_id=9876, + created_at=datetime.datetime(2025, 4, 5, 14, 5, 26, 882324, tzinfo=datetime.timezone.utc), + updated_at=datetime.datetime(2025, 4, 5, 14, 5, 26, 882324, tzinfo=datetime.timezone.utc), + ) + + +# Basic CRUD tests +@pytest.mark.asyncio +async def test_add_contact(contact_manager, sample_contact): + """Test adding a contact""" + result = await contact_manager.add_contact(sample_contact) + assert result is True + contact_manager.queue.add_item.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_contact(contact_manager, sample_contact): + """Test updating a contact""" + contact_manager.queue.find = AsyncMock(return_value=sample_contact.model_dump()) + result = await contact_manager.update_contact("test_id", {"name": "Updated Name"}) + assert result is True + contact_manager.queue.update_record.assert_called_once() + + +@pytest.mark.asyncio +async def test_delete_contact(contact_manager): + """Test deleting a contact""" + result = await contact_manager.delete_contact("test_id") + assert result is True + contact_manager.queue.delete_record.assert_called_once() + + +@pytest.mark.asyncio +async def test_get_contact_by_id(contact_manager, sample_contact): + """Test retrieving a contact by ID""" + contact_dict = sample_contact.model_dump() + contact_manager.queue.find = AsyncMock(return_value=contact_dict) + result = await contact_manager.get_contact_by_id("test_id") + assert isinstance(result, ContactItem) + assert result.name == sample_contact.name + contact_manager.queue.find.assert_called_once() + + +@pytest.mark.asyncio +async def test_find_contacts(contact_manager, sample_contact): + """Test finding contacts by criteria""" + contact_dicts = [sample_contact.model_dump()] + contact_manager.queue.find_many = AsyncMock(return_value=contact_dicts) + result = await contact_manager.find_contacts({"name": "Test"}) + assert len(result) == 1 + assert isinstance(result[0], ContactItem) + assert result[0].name == sample_contact.name + contact_manager.queue.find_many.assert_called_once() + + +@pytest.mark.asyncio +async def test_search_contacts(contact_manager, sample_contact): + """Test searching contacts by text""" + contact_dicts = [sample_contact.model_dump()] + contact_manager.queue.find_many = AsyncMock(return_value=contact_dicts) + result = await contact_manager.search_contacts("Test") + assert len(result) == 1 + assert isinstance(result[0], ContactItem) + assert result[0].name == sample_contact.name + contact_manager.queue.find_many.assert_called_once() + + +@pytest.mark.asyncio +async def test_get_random_contact(contact_manager, sample_contact): + """Test getting a random contact""" + contact_manager.queue.get_random_record = AsyncMock(return_value=sample_contact.model_dump()) + result = await contact_manager.get_random_contact() + assert isinstance(result, ContactItem) + assert result.name == sample_contact.name + contact_manager.queue.get_random_record.assert_called_once() + + +# User Story Tests +@pytest.mark.asyncio +async def test_parse_contact_with_llm_success(): + """Test parsing contact info from text using LLM - successful case""" + # Create mock LLM provider + mock_llm = AsyncMock() + + # Mock the LLM response - create an actual ContactItem + mock_contact = ContactItem( + data="Jane Smith", + name="Jane Smith", + phone="555-5678", + email="jane@example.com", + telegram="@janesmith" + ) + mock_llm.aquery_llm_structured.return_value = mock_contact + + # Create mock queue and queue manager + mock_queue = AsyncMock() + mock_queue_manager = AsyncMock() + mock_queue_manager.create_queue.return_value = mock_queue + + # Patch dependencies - mock both queue manager and llm provider + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.get_queue_manager", return_value=mock_queue_manager): + + # Create manager + settings = ContactManagerSettings(collection="test_contacts") + manager = ContactManager(settings) + + # Test parsing + contact = await manager.parse_contact_with_llm("Jane Smith, 555-5678, jane@example.com, @janesmith", owner_id=1234) + + # Assert + assert contact is not None + assert contact.name == "Jane Smith" + assert contact.phone == "555-5678" + assert contact.email == "jane@example.com" + assert contact.telegram == "@janesmith" + assert contact.owner_id == 1234 + + # Verify LLM was called with appropriate system message + mock_llm.aquery_llm_structured.assert_called_once() + + +@pytest.mark.asyncio +async def test_parse_contact_with_llm_missing_fields(): + """Test parsing contact with missing fields""" + # Create mock LLM provider + mock_llm = AsyncMock() + + # Mock the LLM response with minimal info + mock_contact = ContactItem( + data="Jane Smith", + name="Jane Smith", + notes="Not much info provided" + ) + mock_llm.aquery_llm_structured.return_value = mock_contact + + # Create mock queue and queue manager + mock_queue = AsyncMock() + mock_queue_manager = AsyncMock() + mock_queue_manager.create_queue.return_value = mock_queue + + # Patch dependencies + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.get_queue_manager", return_value=mock_queue_manager): + + # Create manager + settings = ContactManagerSettings(collection="test_contacts") + manager = ContactManager(settings) + + # Test parsing + contact = await manager.parse_contact_with_llm("Jane Smith", owner_id=1234) + + # Assert + assert contact is not None + assert contact.name == "Jane Smith" + assert contact.notes == "Not much info provided" + assert contact.owner_id == 1234 + assert contact.phone is None + assert contact.email is None + assert contact.telegram is None + + +@pytest.mark.asyncio +async def test_parse_contact_with_llm_failure(): + """Test failed contact parsing""" + # Create mock LLM provider + mock_llm = AsyncMock() + + # Mock the LLM to raise an exception on structured output + mock_llm.aquery_llm_structured.side_effect = Exception("Invalid output format") + + # Create mock queue and queue manager + mock_queue = AsyncMock() + mock_queue_manager = AsyncMock() + mock_queue_manager.create_queue.return_value = mock_queue + + # Patch dependencies + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.get_queue_manager", return_value=mock_queue_manager): + + # Create manager + settings = ContactManagerSettings(collection="test_contacts") + manager = ContactManager(settings) + + # Test parsing - should raise ContactDataError + with pytest.raises(Exception): + await manager.parse_contact_with_llm("Invalid contact info", owner_id=1234) + + +@pytest.mark.asyncio +async def test_parse_contact_llm_exception(): + """Test exception handling during contact parsing""" + # Create mock LLM provider + mock_llm = AsyncMock() + + # Mock the LLM to raise an exception + mock_llm.aquery_llm_structured.side_effect = Exception("LLM service unavailable") + + # Create mock queue and queue manager + mock_queue = AsyncMock() + mock_queue_manager = AsyncMock() + mock_queue_manager.create_queue.return_value = mock_queue + + # Patch dependencies + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.get_queue_manager", return_value=mock_queue_manager): + + # Create manager + settings = ContactManagerSettings(collection="test_contacts") + manager = ContactManager(settings) + + # Test parsing - should raise ContactDataError + with pytest.raises(Exception): + await manager.parse_contact_with_llm("Some contact info", owner_id=1234) + + +@pytest.mark.asyncio +async def test_initialization(): + """Test the initialization function""" + settings = ContactManagerSettings(enabled=True, collection="test_contacts") + mock_queue = AsyncMock() + mock_queue_manager = Mock() + mock_queue_manager.create_queue = Mock(return_value=mock_queue) + + async def mock_get_queue_manager(): + return mock_queue_manager + + with patch("botspot.components.new.contact_manager.get_queue_manager", mock_get_queue_manager): + manager = await initialize(settings) + assert isinstance(manager, ContactManager) + assert manager.settings == settings + assert manager.queue == mock_queue + + +@pytest.mark.asyncio +async def test_get_contacts_by_owner(contact_manager, sample_contact): + """Test getting contacts by owner ID""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[sample_contact.model_dump()]) + + # Test getting contacts by owner + contacts = await contact_manager.get_contacts_by_owner(9876) + assert len(contacts) == 1 + assert contacts[0].name == "Test Contact" + assert contacts[0].owner_id == 9876 + + +@pytest.mark.asyncio +async def test_get_contacts_by_telegram(contact_manager, sample_contact): + """Test getting contacts by Telegram username""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[sample_contact.model_dump()]) + + # Test getting contacts by Telegram + contacts = await contact_manager.get_contacts_by_telegram("@testuser") + assert len(contacts) == 1 + assert contacts[0].name == "Test Contact" + assert contacts[0].telegram == "@testuser" + + +@pytest.mark.asyncio +async def test_get_contacts_by_phone(contact_manager, sample_contact): + """Test getting contacts by phone number""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[sample_contact.model_dump()]) + + # Test getting contacts by phone + contacts = await contact_manager.get_contacts_by_phone("555-1234") + assert len(contacts) == 1 + assert contacts[0].name == "Test Contact" + assert contacts[0].phone == "555-1234" + + +@pytest.mark.asyncio +async def test_get_contacts_by_email(contact_manager, sample_contact): + """Test getting contacts by email""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[sample_contact.model_dump()]) + + # Test getting contacts by email + contacts = await contact_manager.get_contacts_by_email("test@example.com") + assert len(contacts) == 1 + assert contacts[0].name == "Test Contact" + assert contacts[0].email == "test@example.com" + + +@pytest.mark.asyncio +async def test_get_contacts_by_owner_empty(contact_manager): + """Test getting contacts by owner when none exist""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[]) + + contacts = await contact_manager.get_contacts_by_owner(9876) + assert len(contacts) == 0 + + +@pytest.mark.asyncio +async def test_get_contacts_by_telegram_empty(contact_manager): + """Test getting contacts by Telegram when none exist""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[]) + + contacts = await contact_manager.get_contacts_by_telegram("@testuser") + assert len(contacts) == 0 + + +@pytest.mark.asyncio +async def test_get_contacts_by_phone_empty(contact_manager): + """Test getting contacts by phone when none exist""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[]) + + contacts = await contact_manager.get_contacts_by_phone("555-1234") + assert len(contacts) == 0 + + +@pytest.mark.asyncio +async def test_get_contacts_by_email_empty(contact_manager): + """Test getting contacts by email when none exist""" + # Set up mock return values + contact_manager.queue.find_many = AsyncMock(return_value=[]) + + contacts = await contact_manager.get_contacts_by_email("test@example.com") + assert len(contacts) == 0 \ No newline at end of file