From 6d68fa83f5358469765aa393492a5824751b80d2 Mon Sep 17 00:00:00 2001 From: Petr Lavrov Date: Fri, 21 Mar 2025 03:35:45 +0300 Subject: [PATCH 1/7] update CLAUDE.md --- CLAUDE.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 3dfad56..41de419 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,3 +1,25 @@ +# 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 + # Botspot Development Guide ## Build & Test Commands From 317cda80555b913d07cda04dfeac70c0fdf5b241 Mon Sep 17 00:00:00 2001 From: Petr Lavrov Date: Mon, 24 Mar 2025 01:36:48 +0300 Subject: [PATCH 2/7] contact manager: Spec and workalong --- CLAUDE.md | 21 ------- .../rework_protocol.md | 58 +++++++++++++++++++ dev/workalong_contact_manager/spec.md | 21 +++++++ dev/workalong_contact_manager/workalong.md | 50 ++++++++++++++++ 4 files changed, 129 insertions(+), 21 deletions(-) create mode 100644 dev/workalong_contact_manager/rework_protocol.md create mode 100644 dev/workalong_contact_manager/spec.md create mode 100644 dev/workalong_contact_manager/workalong.md diff --git a/CLAUDE.md b/CLAUDE.md index 41de419..327d0bc 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,24 +1,3 @@ -# 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 # Botspot Development Guide 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 From 2d0be6e5f55b4237735adf8adc11aa0485c2c83c Mon Sep 17 00:00:00 2001 From: Petr Lavrov Date: Mon, 24 Mar 2025 01:39:46 +0300 Subject: [PATCH 3/7] contact manager: dump changes --- botspot/components/new/contact_manager.py | 444 +++++++++++++++++- botspot/core/bot_manager.py | 8 +- botspot/core/botspot_settings.py | 2 + botspot/core/dependency_manager.py | 12 + botspot/utils/deps_getters.py | 4 + .../contact_manager_demo/README.md | 53 +++ .../contact_manager_demo/bot.py | 58 +++ .../contact_manager_demo/sample.env | 37 ++ 8 files changed, 611 insertions(+), 7 deletions(-) create mode 100644 examples/components_examples/contact_manager_demo/README.md create mode 100644 examples/components_examples/contact_manager_demo/bot.py create mode 100644 examples/components_examples/contact_manager_demo/sample.env diff --git a/botspot/components/new/contact_manager.py b/botspot/components/new/contact_manager.py index a8a44ab..7335624 100644 --- a/botspot/components/new/contact_manager.py +++ b/botspot/components/new/contact_manager.py @@ -1,13 +1,35 @@ -from typing import TYPE_CHECKING +from datetime import date, datetime, timezone +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from aiogram import Dispatcher, Router +from aiogram.filters import Command +from aiogram.types import Message +from pydantic import BaseModel, Field from pydantic_settings import BaseSettings +from botspot import commands_menu +from botspot.commands_menu import Visibility +from botspot.utils.admin_filter import AdminFilter +from botspot.utils.deps_getters import get_database +from botspot.utils.internal import get_logger +from botspot.utils.user_ops import UserLike, get_chat_id + if TYPE_CHECKING: - from motor.motor_asyncio import AsyncIOMotorCollection # noqa: F401 + from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase # noqa: F401 + +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 +38,430 @@ class Config: extra = "ignore" +class Contact(BaseModel): + """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 (user_id) + + @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}") + + return "\n".join(parts) + + +# --------------------------------------------- +# endregion Settings and Model Definitions +# --------------------------------------------- + + +# --------------------------------------------- +# region Contact Manager Implementation +# --------------------------------------------- + class ContactManager: - pass + """Manager for user contacts.""" + + def __init__(self, db: "AsyncIOMotorDatabase", collection: str): + """Initialize contact manager with database connection.""" + self.db = db + self.collection = collection + + @property + def contacts_collection(self) -> "AsyncIOMotorCollection": + """Get the contacts collection.""" + return self.db[self.collection] + + async def add_contact(self, contact: Contact) -> bool: + """Add a new contact to the database.""" + try: + contact.updated_at = datetime.now(timezone.utc) + result = await self.contacts_collection.insert_one(contact.model_dump()) + logger.info(f"Added contact: {contact.name} with ID {result.inserted_id}") + return True + except Exception as e: + logger.error(f"Error adding contact {contact.name}: {e}") + return False + + async def update_contact(self, contact_id: str, data: Dict[str, Any]) -> bool: + """Update an existing contact.""" + try: + # Add updated timestamp + data["updated_at"] = datetime.now(timezone.utc) + result = await self.contacts_collection.update_one( + {"_id": contact_id}, + {"$set": data} + ) + return result.modified_count > 0 + except Exception as e: + logger.error(f"Error updating contact {contact_id}: {e}") + return False + + async def delete_contact(self, contact_id: str) -> bool: + """Delete a contact by ID.""" + try: + result = await self.contacts_collection.delete_one({"_id": contact_id}) + return result.deleted_count > 0 + except Exception as e: + logger.error(f"Error deleting contact {contact_id}: {e}") + return False + + async def get_contact_by_id(self, contact_id: str) -> Optional[Contact]: + """Get a contact by ID.""" + try: + data = await self.contacts_collection.find_one({"_id": contact_id}) + return Contact(**data) if data else None + except Exception as e: + logger.error(f"Error getting contact {contact_id}: {e}") + return None + + async def find_contacts( + self, + query: Dict[str, Any], + limit: int = 10, + owner_id: Optional[int] = None + ) -> List[Contact]: + """Find contacts matching the query.""" + try: + # If owner_id is provided, filter by owner + if owner_id is not None: + query["owner_id"] = owner_id + + cursor = self.contacts_collection.find(query).limit(limit) + contacts = [] + + async for doc in cursor: + contacts.append(Contact(**doc)) + + return contacts + except Exception as e: + logger.error(f"Error finding contacts: {e}") + return [] + + async def search_contacts( + self, + text: str, + owner_id: Optional[int] = None, + limit: int = 10 + ) -> List[Contact]: + """Search for contacts by name, email, etc.""" + query = { + "$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, owner_id) + + async def get_random_contact(self, owner_id: Optional[int] = None) -> Optional[Contact]: + """Get a random contact.""" + try: + pipeline = [] + + # If owner_id is provided, filter by owner + if owner_id is not None: + pipeline.append({"$match": {"owner_id": owner_id}}) + + # Add sample stage to get random document + pipeline.append({"$sample": {"size": 1}}) + + cursor = self.contacts_collection.aggregate(pipeline) + async for doc in cursor: + return Contact(**doc) + + return None + except Exception as e: + logger.error(f"Error getting random contact: {e}") + return None + + async def parse_contact_with_llm( + self, + text: str, + user_id: Optional[int] = None + ) -> Optional[Contact]: + """ + Parse contact information from text using LLM. + + Args: + text: Text message to parse + user_id: User ID who sent the message (for ownership) + + Returns: + Contact object if parsing successful, else None + """ + try: + from botspot.components.new.llm_provider import get_llm_provider + from botspot.utils.deps_getters import get_llm_provider + + # Get LLM provider + llm = get_llm_provider() + + # 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" + } + + If a field is not found, set it to null. + The name field is required - make your best guess if not explicit. + """ + + # Use structured output for better parsing + result = await llm.aquery_llm_structured( + prompt=text, + output_schema=Contact, + user=user_id, + system_message=system_message, + temperature=0.2, # Lower temperature for more deterministic output + ) + + # Add owner ID + result.owner_id = user_id + + return result + except Exception as e: + logger.error(f"Error parsing contact with LLM: {e}") + return None + +# --------------------------------------------- +# endregion Contact Manager Implementation +# --------------------------------------------- -def setup_dispatcher(dp): +# --------------------------------------------- +# region Command Handlers +# --------------------------------------------- + +def setup_command_handlers(dp: Dispatcher, manager: ContactManager, settings: ContactManagerSettings): + """Set up command handlers for contact management.""" + 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) + + 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) + return dp def initialize(settings: ContactManagerSettings) -> ContactManager: - pass + """Initialize the Contact Manager component.""" + if not settings.enabled: + logger.info("Contact Manager component is disabled") + return None + + logger.info("Initializing Contact Manager component") + + # Get MongoDB database + db = get_database() + + # Create Contact Manager instance + manager = ContactManager(db, settings.collection) + + 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 +# --------------------------------------------- diff --git a/botspot/core/bot_manager.py b/botspot/core/bot_manager.py index 885fe8f..79ed7d4 100644 --- a/botspot/core/bot_manager.py +++ b/botspot/core/bot_manager.py @@ -13,7 +13,7 @@ from botspot.components.features import user_interactions from botspot.components.main import event_scheduler, single_user_mode, telethon_manager, trial_mode from botspot.components.middlewares import error_handler -from botspot.components.new import chat_binder, llm_provider +from botspot.components.new import chat_binder, contact_manager, llm_provider from botspot.components.qol import bot_commands_menu, bot_info, print_bot_url from botspot.core.botspot_settings import BotspotSettings from botspot.core.dependency_manager import DependencyManager @@ -60,6 +60,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) def setup_dispatcher(self, dp: Dispatcher): """Setup dispatcher with components""" @@ -105,3 +108,6 @@ 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) diff --git a/botspot/core/botspot_settings.py b/botspot/core/botspot_settings.py index 7c1d186..48e92e0 100644 --- a/botspot/core/botspot_settings.py +++ b/botspot/core/botspot_settings.py @@ -13,6 +13,7 @@ from botspot.components.main.trial_mode import TrialModeSettings from botspot.components.middlewares.error_handler import ErrorHandlerSettings from botspot.components.new.chat_binder import ChatBinderSettings +from botspot.components.new.contact_manager import ContactManagerSettings from botspot.components.new.llm_provider import LLMProviderSettings from botspot.components.qol.bot_commands_menu import BotCommandsMenuSettings from botspot.components.qol.bot_info import BotInfoSettings @@ -59,6 +60,7 @@ def friends(self) -> List[str]: admin_filter: AdminFilterSettings = AdminFilterSettings() chat_binder: ChatBinderSettings = ChatBinderSettings() llm_provider: LLMProviderSettings = LLMProviderSettings() + contact_manager: ContactManagerSettings = ContactManagerSettings() class Config: env_prefix = "BOTSPOT_" diff --git a/botspot/core/dependency_manager.py b/botspot/core/dependency_manager.py index 88bf491..0a8aacc 100644 --- a/botspot/core/dependency_manager.py +++ b/botspot/core/dependency_manager.py @@ -15,6 +15,7 @@ from botspot.components.data.user_data import UserManager from botspot.components.main.telethon_manager import TelethonManager + from botspot.components.new.contact_manager import ContactManager from botspot.components.new.llm_provider import LLMProvider @@ -38,6 +39,7 @@ def __init__( self._user_manager = None self._chat_binder = None self._llm_provider = None + self._contact_manager = None self.__dict__.update(kwargs) @property @@ -137,6 +139,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 @classmethod def is_initialized(cls) -> bool: diff --git a/botspot/utils/deps_getters.py b/botspot/utils/deps_getters.py index b508dc4..3fac394 100644 --- a/botspot/utils/deps_getters.py +++ b/botspot/utils/deps_getters.py @@ -12,6 +12,8 @@ from botspot.components.data.user_data import get_user_manager from botspot.components.main.event_scheduler import get_scheduler from botspot.components.main.telethon_manager import get_telethon_manager +from botspot.components.new.contact_manager import get_contact_manager +from botspot.components.new.llm_provider import get_llm_provider if TYPE_CHECKING: from aiogram import Bot, Dispatcher @@ -70,4 +72,6 @@ async def get_telethon_client( "get_telethon_manager", "get_telethon_client", "get_mongo_client", + "get_contact_manager", + "get_llm_provider", ] 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 From d8a37571fab88a3a6c01a9d9e783f0223f6401f6 Mon Sep 17 00:00:00 2001 From: GMKirkon <367307@edu.itmo.ru> Date: Sat, 5 Apr 2025 02:10:37 +0300 Subject: [PATCH 4/7] add necessary functionality to Queue class use Queue inside contact_manager --- botspot/components/new/contact_manager.py | 140 ++++---- botspot/components/new/queue_manager.py | 110 ++++++ tests/components/new/test_contact_manager.py | 337 +++++++++++++++++++ 3 files changed, 513 insertions(+), 74 deletions(-) create mode 100644 tests/components/new/test_contact_manager.py diff --git a/botspot/components/new/contact_manager.py b/botspot/components/new/contact_manager.py index 7335624..0287aa7 100644 --- a/botspot/components/new/contact_manager.py +++ b/botspot/components/new/contact_manager.py @@ -1,6 +1,11 @@ from datetime import date, datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +if TYPE_CHECKING: + from motor.motor_asyncio import AsyncIOMotorCollection # noqa: F401 + from botspot.commands_menu import Visibility + from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase # noqa: F401 + from aiogram import Dispatcher, Router from aiogram.filters import Command from aiogram.types import Message @@ -8,14 +13,11 @@ from pydantic_settings import BaseSettings from botspot import commands_menu -from botspot.commands_menu import Visibility from botspot.utils.admin_filter import AdminFilter from botspot.utils.deps_getters import get_database from botspot.utils.internal import get_logger -from botspot.utils.user_ops import UserLike, get_chat_id - -if TYPE_CHECKING: - from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase # noqa: F401 +from botspot.utils.user_ops import UserLike +from botspot.components.new.queue_manager import create_queue, get_queue logger = get_logger() @@ -38,7 +40,7 @@ class Config: extra = "ignore" -class Contact(BaseModel): +class ContactItem(BaseModel): """Contact information model.""" name: str @@ -84,22 +86,23 @@ def display_info(self) -> str: class ContactManager: """Manager for user contacts.""" - def __init__(self, db: "AsyncIOMotorDatabase", collection: str): - """Initialize contact manager with database connection.""" - self.db = db - self.collection = collection - - @property - def contacts_collection(self) -> "AsyncIOMotorCollection": - """Get the contacts collection.""" - return self.db[self.collection] + def __init__(self, settings: ContactManagerSettings): + """Initialize contact manager with settings.""" + self.settings = settings + self.queue = create_queue( + key="contacts", + item_model=ContactItem + ) - async def add_contact(self, contact: Contact) -> bool: + async def add_contact(self, contact: ContactItem) -> bool: """Add a new contact to the database.""" try: contact.updated_at = datetime.now(timezone.utc) - result = await self.contacts_collection.insert_one(contact.model_dump()) - logger.info(f"Added contact: {contact.name} with ID {result.inserted_id}") + await self.queue.add_item(contact, user_id=contact.owner_id) + logger.info(f"Added contact: {contact.name}") + # used to be + # logger.info(f"Added contact: {contact.name}, with ID {result.inserted_id}"), but QM does not provide ID, add_item returns void. + # may change QM iterface in the future return True except Exception as e: logger.error(f"Error adding contact {contact.name}: {e}") @@ -108,13 +111,11 @@ async def add_contact(self, contact: Contact) -> bool: async def update_contact(self, contact_id: str, data: Dict[str, Any]) -> bool: """Update an existing contact.""" try: - # Add updated timestamp + # Update the timestamp data["updated_at"] = datetime.now(timezone.utc) - result = await self.contacts_collection.update_one( - {"_id": contact_id}, - {"$set": data} - ) - return result.modified_count > 0 + + # Use the queue's update_record method to directly update the contact + return await self.queue.update_record(contact_id, data) except Exception as e: logger.error(f"Error updating contact {contact_id}: {e}") return False @@ -122,17 +123,16 @@ async def update_contact(self, contact_id: str, data: Dict[str, Any]) -> bool: async def delete_contact(self, contact_id: str) -> bool: """Delete a contact by ID.""" try: - result = await self.contacts_collection.delete_one({"_id": contact_id}) - return result.deleted_count > 0 + return await self.queue.delete_record(contact_id) except Exception as e: logger.error(f"Error deleting contact {contact_id}: {e}") return False - async def get_contact_by_id(self, contact_id: str) -> Optional[Contact]: + async def get_contact_by_id(self, contact_id: str) -> Optional[ContactItem]: """Get a contact by ID.""" try: - data = await self.contacts_collection.find_one({"_id": contact_id}) - return Contact(**data) if data else None + contact_dict = await self.queue.find({"_id": contact_id}) + return ContactItem(**contact_dict) if contact_dict else None except Exception as e: logger.error(f"Error getting contact {contact_id}: {e}") return None @@ -142,20 +142,16 @@ async def find_contacts( query: Dict[str, Any], limit: int = 10, owner_id: Optional[int] = None - ) -> List[Contact]: + ) -> List[ContactItem]: """Find contacts matching the query.""" try: - # If owner_id is provided, filter by owner + # If owner_id is provided, add it to the query if owner_id is not None: query["owner_id"] = owner_id - - cursor = self.contacts_collection.find(query).limit(limit) - contacts = [] - async for doc in cursor: - contacts.append(Contact(**doc)) - - return contacts + # Use the queue's find_many method + contact_dicts = await self.queue.find_many(query, limit=limit) + return [ContactItem(**doc) for doc in contact_dicts] except Exception as e: logger.error(f"Error finding contacts: {e}") return [] @@ -165,37 +161,37 @@ async def search_contacts( text: str, owner_id: Optional[int] = None, limit: int = 10 - ) -> List[Contact]: + ) -> List[ContactItem]: """Search for contacts by name, email, etc.""" - query = { - "$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, owner_id) - - async def get_random_contact(self, owner_id: Optional[int] = None) -> Optional[Contact]: - """Get a random contact.""" try: - pipeline = [] + # Create a query with $or for text matching across multiple fields + query = { + "$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"}} + ] + } - # If owner_id is provided, filter by owner + # Add owner_id constraint if provided if owner_id is not None: - pipeline.append({"$match": {"owner_id": owner_id}}) - - # Add sample stage to get random document - pipeline.append({"$sample": {"size": 1}}) + query["owner_id"] = owner_id - cursor = self.contacts_collection.aggregate(pipeline) - async for doc in cursor: - return Contact(**doc) - - return None + # Use find_many with the query + contact_dicts = await self.queue.find_many(query, limit=limit) + return [ContactItem(**doc) for doc in contact_dicts] + except Exception as e: + logger.error(f"Error searching contacts: {e}") + return [] + + async def get_random_contact(self, owner_id: Optional[int] = None) -> Optional[ContactItem]: + """Get a random contact.""" + try: + # Use the queue's get_random_record method + contact_dict = await self.queue.get_random_record(user_id=owner_id) + return ContactItem(**contact_dict) if contact_dict else None except Exception as e: logger.error(f"Error getting random contact: {e}") return None @@ -204,7 +200,7 @@ async def parse_contact_with_llm( self, text: str, user_id: Optional[int] = None - ) -> Optional[Contact]: + ) -> Optional[ContactItem]: """ Parse contact information from text using LLM. @@ -242,7 +238,7 @@ async def parse_contact_with_llm( # Use structured output for better parsing result = await llm.aquery_llm_structured( prompt=text, - output_schema=Contact, + output_schema=ContactItem, user=user_id, system_message=system_message, temperature=0.2, # Lower temperature for more deterministic output @@ -267,6 +263,8 @@ async def parse_contact_with_llm( 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 @@ -441,13 +439,7 @@ def initialize(settings: ContactManagerSettings) -> ContactManager: return None logger.info("Initializing Contact Manager component") - - # Get MongoDB database - db = get_database() - - # Create Contact Manager instance - manager = ContactManager(db, settings.collection) - + manager = ContactManager(settings) return manager @@ -464,4 +456,4 @@ def get_contact_manager() -> ContactManager: # --------------------------------------------- # endregion Initialization and Setup -# --------------------------------------------- +# --------------------------------------------- \ No newline at end of file diff --git a/botspot/components/new/queue_manager.py b/botspot/components/new/queue_manager.py index 6b7b1a0..e057827 100644 --- a/botspot/components/new/queue_manager.py +++ b/botspot/components/new/queue_manager.py @@ -88,6 +88,85 @@ 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]) -> Optional[Dict[str, Any]]: + """ + Find a single record matching the given query. + + Args: + query: MongoDB query dictionary + + Returns: + The record as a dictionary if found, None otherwise + """ + try: + 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], limit: Optional[int] = None) -> List[Dict[str, Any]]: + """ + Find multiple records matching the given query. + + Args: + query: MongoDB query dictionary + limit: Maximum number of records to return + + Returns: + List of records as dictionaries + """ + try: + 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) -> bool: + """ + Delete a record by its ID. + + Args: + record_id: The record ID to delete + + Returns: + True if deletion was successful, False otherwise + """ + try: + result = await self.collection.delete_one({"_id": record_id}) + 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 +202,37 @@ 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: + # 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/tests/components/new/test_contact_manager.py b/tests/components/new/test_contact_manager.py new file mode 100644 index 0000000..800b05f --- /dev/null +++ b/tests/components/new/test_contact_manager.py @@ -0,0 +1,337 @@ +import asyncio +import datetime +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from botspot.components.new.contact_manager import ContactItem, ContactManager, ContactManagerSettings, initialize, setup_command_handlers + + +@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(): + """Fixture to create a mock queue for the contact manager""" + queue = AsyncMock() + queue.collection = MagicMock() + queue.collection.find_one.return_value = None # Default: contact not found + queue.find_one.return_value = None + return queue + + +@pytest.fixture +def mock_dispatcher(): + """Fixture to create a mock dispatcher""" + return MagicMock() + + +@pytest.fixture +def contact_manager(contact_manager_settings, mock_queue): + """Fixture to create a contact manager with a mock queue""" + with patch("botspot.components.new.contact_manager.create_queue", return_value=mock_queue): + manager = ContactManager(contact_manager_settings) + return manager + + +@pytest.fixture +def sample_contact(): + """Fixture to create a sample contact""" + return ContactItem( + name="Test Contact", + phone="555-1234", + email="test@example.com", + telegram="@testcontact", + birthday=datetime.date(1990, 1, 1), + notes="Test notes", + owner_id=9876 + ) + + +# Basic CRUD tests +def test_add_contact(contact_manager, sample_contact): + """Test adding a contact""" + contact_manager.queue.add_item.return_value = None + + # Execute + loop = asyncio.get_event_loop() + result = loop.run_until_complete(contact_manager.add_contact(sample_contact)) + + # Assert + assert result == True + contact_manager.queue.add_item.assert_called_once() + + +def test_update_contact(contact_manager, sample_contact): + """Test updating a contact""" + contact_manager.queue.update_record.return_value = True + + # Execute + loop = asyncio.get_event_loop() + result = loop.run_until_complete(contact_manager.update_contact("contact123", {"name": "Updated Name"})) + + # Assert + assert result == True + contact_manager.queue.update_record.assert_called_once() + + +def test_delete_contact(contact_manager): + """Test deleting a contact""" + contact_manager.queue.delete_record.return_value = True + + # Execute + loop = asyncio.get_event_loop() + result = loop.run_until_complete(contact_manager.delete_contact("contact123")) + + # Assert + assert result == True + contact_manager.queue.delete_record.assert_called_once_with("contact123") + + +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.return_value = contact_dict + + # Execute + loop = asyncio.get_event_loop() + result = loop.run_until_complete(contact_manager.get_contact_by_id("contact123")) + + # Assert + assert isinstance(result, ContactItem) + assert result.name == "Test Contact" + contact_manager.queue.find.assert_called_once_with({"_id": "contact123"}) + + +def test_find_contacts(contact_manager, sample_contact): + """Test finding contacts by criteria""" + contact_dicts = [sample_contact.model_dump()] + contact_manager.queue.find_many.return_value = contact_dicts + + # Execute + loop = asyncio.get_event_loop() + result = loop.run_until_complete(contact_manager.find_contacts({"name": "Test Contact"})) + + # Assert + assert len(result) == 1 + assert isinstance(result[0], ContactItem) + assert result[0].name == "Test Contact" + contact_manager.queue.find_many.assert_called_once() + + +def test_search_contacts(contact_manager, sample_contact): + """Test searching contacts by text""" + contact_dicts = [sample_contact.model_dump()] + contact_manager.queue.find_many.return_value = contact_dicts + + # Execute + loop = asyncio.get_event_loop() + result = loop.run_until_complete(contact_manager.search_contacts("Test")) + + # Assert + assert len(result) == 1 + assert isinstance(result[0], ContactItem) + assert result[0].name == "Test Contact" + contact_manager.queue.find_many.assert_called_once() + + +def test_get_random_contact(contact_manager, sample_contact): + """Test getting a random contact""" + contact_manager.queue.get_random_record.return_value = sample_contact.model_dump() + + # Execute + loop = asyncio.get_event_loop() + result = loop.run_until_complete(contact_manager.get_random_contact()) + + # Assert + assert isinstance(result, ContactItem) + assert result.name == "Test Contact" + 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( + name="Jane Smith", + phone="555-5678", + email="jane@example.com", + telegram="@janesmith" + ) + mock_llm.aquery_llm_structured.return_value = mock_contact + + # Create mock queue + mock_queue = AsyncMock() + + # Patch dependencies - mock both queue and llm provider + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.create_queue", return_value=mock_queue): + + # Create manager directly - don't rely on queue_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", user_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() + assert "parser" in mock_llm.aquery_llm_structured.call_args[1]["system_message"].lower() + + +@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( + name="Jane Smith", + notes="Not much info provided" + ) + mock_llm.aquery_llm_structured.return_value = mock_contact + + # Create mock queue + mock_queue = AsyncMock() + + # Patch dependencies + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.create_queue", return_value=mock_queue): + + # Create manager + settings = ContactManagerSettings(collection="test_contacts") + manager = ContactManager(settings) + + # Test parsing + contact = await manager.parse_contact_with_llm("Jane Smith", user_id=1234) + + # Assert + assert contact is not None + assert contact.name == "Jane Smith" + assert contact.phone is None + assert contact.email is None + assert contact.telegram is None + assert contact.owner_id == 1234 + + +@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 + mock_queue = AsyncMock() + + # Patch dependencies + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.create_queue", return_value=mock_queue): + + # Create manager + settings = ContactManagerSettings(collection="test_contacts") + manager = ContactManager(settings) + + # Test parsing + contact = await manager.parse_contact_with_llm("Jane Smith", user_id=1234) + + # Assert + assert contact is None + + +@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 + mock_queue = AsyncMock() + + # Patch dependencies + with patch("botspot.utils.deps_getters.get_llm_provider", return_value=mock_llm), \ + patch("botspot.components.new.contact_manager.create_queue", return_value=mock_queue): + + # Create manager + settings = ContactManagerSettings(collection="test_contacts") + manager = ContactManager(settings) + + # Test parsing + contact = await manager.parse_contact_with_llm("Jane Smith", user_id=1234) + + # Assert + assert contact is None + + +@pytest.mark.asyncio +async def test_add_contact_command_handler(): + """Test the add_contact command handler""" + # Skip this test for now - will be implemented after fixing other issues + pytest.skip("Handler tests will be implemented after fixing core functionality") + + +@pytest.mark.asyncio +async def test_find_contact_command_handler(): + """Test the find_contact command handler""" + # Skip this test for now - will be implemented after fixing other issues + pytest.skip("Handler tests will be implemented after fixing core functionality") + + +@pytest.mark.asyncio +async def test_random_contact_command_handler(): + """Test the random_contact command handler""" + # Skip this test for now - will be implemented after fixing other issues + pytest.skip("Handler tests will be implemented after fixing core functionality") + + +@pytest.mark.asyncio +async def test_message_parser(): + """Test the message parser that automatically detects contact info in chat messages""" + # Skip this test for now - will be implemented after fixing other issues + pytest.skip("Handler tests will be implemented after fixing core functionality") + + +# Integration tests +@pytest.mark.asyncio +async def test_initialization(): + """Test the initialization function""" + settings = ContactManagerSettings(enabled=True, collection="test_contacts") + + with patch("botspot.components.new.contact_manager.create_queue") as mock_create_queue: + mock_queue = AsyncMock() + mock_create_queue.return_value = mock_queue + + # Call initialize + manager = initialize(settings) + + # Assert + assert isinstance(manager, ContactManager) + + # Verify create_queue was called with the correct parameters - note the 'key' named parameter + mock_create_queue.assert_called_once_with(key="contacts", item_model=ContactItem) \ No newline at end of file From 2b6bbd0809aa4e4361ab851151e0b146bbe5fd5f Mon Sep 17 00:00:00 2001 From: GMKirkon <367307@edu.itmo.ru> Date: Sat, 5 Apr 2025 14:02:10 +0300 Subject: [PATCH 5/7] permission checks in added methods --- botspot/components/new/contact_manager.py | 38 +++++++++++++------- botspot/components/new/queue_manager.py | 37 ++++++++++++++++--- tests/components/new/test_contact_manager.py | 4 +-- 3 files changed, 61 insertions(+), 18 deletions(-) diff --git a/botspot/components/new/contact_manager.py b/botspot/components/new/contact_manager.py index 0287aa7..9f0e9d7 100644 --- a/botspot/components/new/contact_manager.py +++ b/botspot/components/new/contact_manager.py @@ -128,10 +128,10 @@ async def delete_contact(self, contact_id: str) -> bool: logger.error(f"Error deleting contact {contact_id}: {e}") return False - async def get_contact_by_id(self, contact_id: str) -> Optional[ContactItem]: + async def get_contact_by_id(self, contact_id: str, owner_id: Optional[int] = None) -> Optional[ContactItem]: """Get a contact by ID.""" try: - contact_dict = await self.queue.find({"_id": contact_id}) + contact_dict = await self.queue.find({"_id": contact_id}, user_id=owner_id) return ContactItem(**contact_dict) if contact_dict else None except Exception as e: logger.error(f"Error getting contact {contact_id}: {e}") @@ -150,7 +150,7 @@ async def find_contacts( query["owner_id"] = owner_id # Use the queue's find_many method - contact_dicts = await self.queue.find_many(query, limit=limit) + contact_dicts = await self.queue.find_many(query, user_id=owner_id, limit=limit) return [ContactItem(**doc) for doc in contact_dicts] except Exception as e: logger.error(f"Error finding contacts: {e}") @@ -165,7 +165,7 @@ async def search_contacts( """Search for contacts by name, email, etc.""" try: # Create a query with $or for text matching across multiple fields - query = { + query: Dict[str, Any] = { "$or": [ {"name": {"$regex": text, "$options": "i"}}, {"email": {"$regex": text, "$options": "i"}}, @@ -180,7 +180,7 @@ async def search_contacts( query["owner_id"] = owner_id # Use find_many with the query - contact_dicts = await self.queue.find_many(query, limit=limit) + contact_dicts = await self.queue.find_many(query, user_id=owner_id, limit=limit) return [ContactItem(**doc) for doc in contact_dicts] except Exception as e: logger.error(f"Error searching contacts: {e}") @@ -228,26 +228,40 @@ async def parse_contact_with_llm( "email": "Email (if found)", "telegram": "Telegram username (if found)", "birthday": "Birthday in YYYY-MM-DD format (if found)", - "notes": "Any additional information" + "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 - result = await llm.aquery_llm_structured( + temp_contact = await llm.aquery_llm_structured( prompt=text, - output_schema=ContactItem, + output_schema=TempContact, user=user_id, system_message=system_message, temperature=0.2, # Lower temperature for more deterministic output ) - # Add owner ID - result.owner_id = user_id + # Create a new dictionary with owner_id if provided + contact_data = temp_contact.model_dump() + if user_id is not None: + contact_data["owner_id"] = user_id - return result + # Convert to ContactItem + return ContactItem.model_validate(contact_data) except Exception as e: logger.error(f"Error parsing contact with LLM: {e}") return None @@ -436,7 +450,7 @@ def initialize(settings: ContactManagerSettings) -> ContactManager: """Initialize the Contact Manager component.""" if not settings.enabled: logger.info("Contact Manager component is disabled") - return None + return ContactManager(settings) # Return empty manager instead of None logger.info("Initializing Contact Manager component") manager = ContactManager(settings) diff --git a/botspot/components/new/queue_manager.py b/botspot/components/new/queue_manager.py index e057827..65a0a26 100644 --- a/botspot/components/new/queue_manager.py +++ b/botspot/components/new/queue_manager.py @@ -114,34 +114,50 @@ async def update_record(self, record_id: Any, data: Dict[str, Any]) -> bool: logger.error(f"Error updating record: {e}") return False - async def find(self, query: Dict[str, Any]) -> Optional[Dict[str, Any]]: + 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], limit: Optional[int] = None) -> List[Dict[str, Any]]: + 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) @@ -150,18 +166,27 @@ async def find_many(self, query: Dict[str, Any], limit: Optional[int] = None) -> logger.error(f"Error finding records: {e}") return [] - async def delete_record(self, record_id: Any) -> bool: + 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: - result = await self.collection.delete_one({"_id": record_id}) + 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}") @@ -213,6 +238,10 @@ async def get_random_record(self, user_id: Optional[int] = None) -> Optional[Dic 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 = [] diff --git a/tests/components/new/test_contact_manager.py b/tests/components/new/test_contact_manager.py index 800b05f..1558448 100644 --- a/tests/components/new/test_contact_manager.py +++ b/tests/components/new/test_contact_manager.py @@ -102,12 +102,12 @@ def test_get_contact_by_id(contact_manager, sample_contact): # Execute loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.get_contact_by_id("contact123")) + result = loop.run_until_complete(contact_manager.get_contact_by_id("contact123", owner_id=9876)) # Assert assert isinstance(result, ContactItem) assert result.name == "Test Contact" - contact_manager.queue.find.assert_called_once_with({"_id": "contact123"}) + contact_manager.queue.find.assert_called_once_with({"_id": "contact123"}, user_id=9876) def test_find_contacts(contact_manager, sample_contact): From 53aa3e9b18e5e3f9157f4b92b4d1933e2bbc263a Mon Sep 17 00:00:00 2001 From: GMKirkon <367307@edu.itmo.ru> Date: Sat, 5 Apr 2025 16:18:52 +0300 Subject: [PATCH 6/7] add utils func exposion --- botspot/components/new/contact_manager.py | 237 ++++++++++++++-------- botspot/contact_manager.py | 22 ++ 2 files changed, 175 insertions(+), 84 deletions(-) create mode 100644 botspot/contact_manager.py diff --git a/botspot/components/new/contact_manager.py b/botspot/components/new/contact_manager.py index 9f0e9d7..9ccf7b5 100644 --- a/botspot/components/new/contact_manager.py +++ b/botspot/components/new/contact_manager.py @@ -1,6 +1,8 @@ from datetime import date, datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from botspot.core.errors import ContactDataError + if TYPE_CHECKING: from motor.motor_asyncio import AsyncIOMotorCollection # noqa: F401 from botspot.commands_menu import Visibility @@ -42,7 +44,7 @@ class Config: class ContactItem(BaseModel): """Contact information model.""" - + name: str user_id: Optional[int] = None # Telegram user ID if applicable phone: Optional[str] = None # Phone number @@ -53,12 +55,12 @@ class ContactItem(BaseModel): 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 (user_id) - + @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: @@ -70,7 +72,7 @@ def display_info(self) -> str: parts.append(f"🎂 Birthday: {self.birthday.strftime('%d %B')}") if self.notes: parts.append(f"📝 Notes: {self.notes}") - + return "\n".join(parts) @@ -85,7 +87,7 @@ def display_info(self) -> str: class ContactManager: """Manager for user contacts.""" - + def __init__(self, settings: ContactManagerSettings): """Initialize contact manager with settings.""" self.settings = settings @@ -93,41 +95,49 @@ def __init__(self, settings: ContactManagerSettings): key="contacts", item_model=ContactItem ) - - async def add_contact(self, contact: ContactItem) -> bool: - """Add a new contact to the database.""" + + async def add_contact(self, contact: ContactItem, user_id: Optional[int] = None) -> bool: + """Add a new contact to the database""" try: contact.updated_at = datetime.now(timezone.utc) - await self.queue.add_item(contact, user_id=contact.owner_id) + if user_id is not None: + contact.owner_id = user_id + await self.queue.add_item(contact, user_id=user_id) logger.info(f"Added contact: {contact.name}") - # used to be - # logger.info(f"Added contact: {contact.name}, with ID {result.inserted_id}"), but QM does not provide ID, add_item returns void. - # may change QM iterface in the future return True except Exception as e: logger.error(f"Error adding contact {contact.name}: {e}") - return False - - async def update_contact(self, contact_id: str, data: Dict[str, Any]) -> bool: - """Update an existing contact.""" + from botspot.core.errors import BotspotError + raise BotspotError(f"Failed to add contact: {str(e)}") + + async def update_contact(self, contact_id: str, data: Dict[str, Any], user_id: Optional[int] = None) -> bool: + """Update an existing contact""" try: # Update the timestamp data["updated_at"] = datetime.now(timezone.utc) - + # Use the queue's update_record method to directly update the contact - return await self.queue.update_record(contact_id, data) + success = await self.queue.update_record(contact_id, data) + if not success: + raise ContactDataError("Contact not found or update failed") + return True except Exception as e: logger.error(f"Error updating contact {contact_id}: {e}") - return False - - async def delete_contact(self, contact_id: str) -> bool: - """Delete a contact by ID.""" + from botspot.core.errors import BotspotError + raise BotspotError(f"Failed to update contact: {str(e)}") + + async def delete_contact(self, contact_id: str, user_id: Optional[int] = None) -> bool: + """Delete a contact by ID""" try: - return await self.queue.delete_record(contact_id) + success = await self.queue.delete_record(contact_id, user_id=user_id) + if not success: + raise ContactDataError("Contact not found or deletion failed") + return True except Exception as e: logger.error(f"Error deleting contact {contact_id}: {e}") - return False - + from botspot.core.errors import BotspotError + raise BotspotError(f"Failed to delete contact: {str(e)}") + async def get_contact_by_id(self, contact_id: str, owner_id: Optional[int] = None) -> Optional[ContactItem]: """Get a contact by ID.""" try: @@ -136,31 +146,31 @@ async def get_contact_by_id(self, contact_id: str, owner_id: Optional[int] = Non except Exception as e: logger.error(f"Error getting contact {contact_id}: {e}") return None - + async def find_contacts( - self, - query: Dict[str, Any], - limit: int = 10, - owner_id: Optional[int] = None + self, + query: Dict[str, Any], + limit: int = 10, + owner_id: Optional[int] = None ) -> List[ContactItem]: """Find contacts matching the query.""" try: # If owner_id is provided, add it to the query if owner_id is not None: query["owner_id"] = owner_id - + # Use the queue's find_many method contact_dicts = await self.queue.find_many(query, user_id=owner_id, limit=limit) return [ContactItem(**doc) for doc in contact_dicts] except Exception as e: logger.error(f"Error finding contacts: {e}") return [] - + async def search_contacts( - self, - text: str, - owner_id: Optional[int] = None, - limit: int = 10 + self, + text: str, + owner_id: Optional[int] = None, + limit: int = 10 ) -> List[ContactItem]: """Search for contacts by name, email, etc.""" try: @@ -174,18 +184,18 @@ async def search_contacts( {"notes": {"$regex": text, "$options": "i"}} ] } - + # Add owner_id constraint if provided if owner_id is not None: query["owner_id"] = owner_id - + # Use find_many with the query contact_dicts = await self.queue.find_many(query, user_id=owner_id, limit=limit) return [ContactItem(**doc) for doc in contact_dicts] except Exception as e: logger.error(f"Error searching contacts: {e}") return [] - + async def get_random_contact(self, owner_id: Optional[int] = None) -> Optional[ContactItem]: """Get a random contact.""" try: @@ -195,11 +205,11 @@ async def get_random_contact(self, owner_id: Optional[int] = None) -> Optional[C except Exception as e: logger.error(f"Error getting random contact: {e}") return None - + async def parse_contact_with_llm( - self, - text: str, - user_id: Optional[int] = None + self, + text: str, + user_id: Optional[int] = None ) -> Optional[ContactItem]: """ Parse contact information from text using LLM. @@ -214,10 +224,10 @@ async def parse_contact_with_llm( try: from botspot.components.new.llm_provider import get_llm_provider from botspot.utils.deps_getters import get_llm_provider - + # Get LLM provider llm = get_llm_provider() - + # System message for the LLM system_message = """ You are a contact information parser. Extract contact details from the user's message. @@ -235,7 +245,7 @@ async def parse_contact_with_llm( 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 @@ -245,7 +255,7 @@ class TempContact(BaseModel): 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, @@ -254,18 +264,19 @@ class TempContact(BaseModel): 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 user_id is not None: contact_data["owner_id"] = user_id - + # Convert to ContactItem return ContactItem.model_validate(contact_data) except Exception as e: logger.error(f"Error parsing contact with LLM: {e}") return None + # --------------------------------------------- # endregion Contact Manager Implementation # --------------------------------------------- @@ -278,9 +289,9 @@ class TempContact(BaseModel): 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")) @@ -293,22 +304,22 @@ async def add_contact_cmd(message: Message): "/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) @@ -316,10 +327,10 @@ async def add_contact_cmd(message: Message): 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) - + success = await manager.add_contact(contact, message.from_user.id) + if success: await message.reply( f"✅ Contact added successfully!\n\n{contact.display_info}", @@ -327,7 +338,7 @@ async def add_contact_cmd(message: Message): ) 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")) @@ -339,17 +350,17 @@ async def find_contact_cmd(message: Message): "/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( @@ -363,24 +374,24 @@ async def find_contact_cmd(message: Message): 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) + @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() @@ -389,17 +400,17 @@ async def parse_contact_message(message: Message): # 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( @@ -408,10 +419,11 @@ async def parse_contact_message(message: Message): f"Reply with /add_contact to confirm.", parse_mode="Markdown" ) - + # Include router in dispatcher dp.include_router(router) + # --------------------------------------------- # endregion Command Handlers # --------------------------------------------- @@ -424,25 +436,25 @@ async def parse_contact_message(message: Message): 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) - + return dp @@ -451,7 +463,7 @@ def initialize(settings: ContactManagerSettings) -> ContactManager: 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 = ContactManager(settings) return manager @@ -462,12 +474,69 @@ def get_contact_manager() -> ContactManager: 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 -# --------------------------------------------- \ No newline at end of file +# --------------------------------------------- + + +# --------------------------------------------- +# region Utils +# --------------------------------------------- + + +async def add_contact(contact: ContactItem, user_id: Optional[int] = None) -> bool: + """Add a new contact to the database.""" + manager = get_contact_manager() + return await manager.add_contact(contact, user_id) + + +async def update_contact(contact_id: str, data: Dict[str, Any], user_id: Optional[int] = None) -> bool: + """Update an existing contact.""" + manager = get_contact_manager() + return await manager.update_contact(contact_id, data, user_id) + + +async def delete_contact(contact_id: str, user_id: Optional[int] = None) -> bool: + """Delete a contact by ID.""" + manager = get_contact_manager() + return await manager.delete_contact(contact_id, user_id) + + +async def get_contact_by_id(contact_id: str, user_id: Optional[int] = None) -> Optional[ContactItem]: + """Get a contact by ID.""" + manager = get_contact_manager() + return await manager.get_contact_by_id(contact_id, user_id) + + +async def find_contacts(query: Dict[str, Any], user_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=limit, owner_id=user_id) + + +async def search_contacts(text: str, user_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, user_id, limit) + + +async def get_random_contact(user_id: Optional[int] = None) -> Optional[ContactItem]: + """Get a random contact.""" + manager = get_contact_manager() + return await manager.get_random_contact(user_id) + + +async def parse_contact_with_llm(text: str, user_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, user_id) + +# --------------------------------------------- +# endregion Utils diff --git a/botspot/contact_manager.py b/botspot/contact_manager.py new file mode 100644 index 0000000..9c883d8 --- /dev/null +++ b/botspot/contact_manager.py @@ -0,0 +1,22 @@ +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, +) + +__all__ = [ + "add_contact", + "update_contact", + "delete_contact", + "get_contact_by_id", + "find_contacts", + "search_contacts", + "get_random_contact", + "parse_contact_with_llm", +] + From d58676dc4eedb8598562b9f1a92ffd04e21c2e30 Mon Sep 17 00:00:00 2001 From: GMKirkon <367307@edu.itmo.ru> Date: Sat, 5 Apr 2025 17:08:47 +0300 Subject: [PATCH 7/7] refactor contact manager, better naming, new funcs for searching --- botspot/components/new/contact_manager.py | 352 +++++++++++++------ botspot/contact_manager.py | 8 + tests/components/new/test_contact_manager.py | 325 +++++++++-------- 3 files changed, 432 insertions(+), 253 deletions(-) diff --git a/botspot/components/new/contact_manager.py b/botspot/components/new/contact_manager.py index 9ccf7b5..2c50951 100644 --- a/botspot/components/new/contact_manager.py +++ b/botspot/components/new/contact_manager.py @@ -1,25 +1,18 @@ -from datetime import date, datetime, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union - -from botspot.core.errors import ContactDataError +"""Contact Manager component for managing user contacts.""" -if TYPE_CHECKING: - from motor.motor_asyncio import AsyncIOMotorCollection # noqa: F401 - from botspot.commands_menu import Visibility - from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase # noqa: F401 +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 -from botspot import commands_menu -from botspot.utils.admin_filter import AdminFilter -from botspot.utils.deps_getters import get_database from botspot.utils.internal import get_logger -from botspot.utils.user_ops import UserLike -from botspot.components.new.queue_manager import create_queue, get_queue +from botspot import commands_menu +from botspot.components.new.queue_manager import QueueItem, get_queue_manager logger = get_logger() @@ -42,7 +35,13 @@ class Config: extra = "ignore" -class ContactItem(BaseModel): +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 @@ -54,7 +53,7 @@ class ContactItem(BaseModel): 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 (user_id) + owner_id: Optional[int] = None # Who added this contact @property def display_info(self) -> str: @@ -72,6 +71,10 @@ def display_info(self) -> str: 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) @@ -85,92 +88,131 @@ def display_info(self) -> str: # 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: """Manager for user contacts.""" - def __init__(self, settings: ContactManagerSettings): + def __init__( + self, + settings: ContactManagerSettings, + queue_manager: Optional["QueueManager"] = None, + ): """Initialize contact manager with settings.""" self.settings = settings - self.queue = create_queue( - key="contacts", - item_model=ContactItem + 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, user_id: Optional[int] = None) -> bool: - """Add a new contact to the database""" + async def add_contact(self, contact: ContactItem, owner_id: Optional[int] = None) -> bool: + """Add a new contact to the queue.""" try: - contact.updated_at = datetime.now(timezone.utc) - if user_id is not None: - contact.owner_id = user_id - await self.queue.add_item(contact, user_id=user_id) + 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 {contact.name}: {e}") - from botspot.core.errors import BotspotError - raise BotspotError(f"Failed to add contact: {str(e)}") - - async def update_contact(self, contact_id: str, data: Dict[str, Any], user_id: Optional[int] = None) -> bool: - """Update an existing contact""" + 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: - # Update the timestamp - data["updated_at"] = datetime.now(timezone.utc) + # 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") - # Use the queue's update_record method to directly update the contact - success = await self.queue.update_record(contact_id, data) - if not success: - raise ContactDataError("Contact not found or update failed") + # 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 {contact_id}: {e}") - from botspot.core.errors import BotspotError - raise BotspotError(f"Failed to update contact: {str(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, user_id: Optional[int] = None) -> bool: - """Delete a contact by ID""" + async def delete_contact(self, contact_id: str, owner_id: Optional[int] = None) -> bool: + """Delete a contact from the queue.""" try: - success = await self.queue.delete_record(contact_id, user_id=user_id) - if not success: + 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 {contact_id}: {e}") - from botspot.core.errors import BotspotError - raise BotspotError(f"Failed to delete contact: {str(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 ID.""" + 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_dict = await self.queue.find({"_id": contact_id}, user_id=owner_id) - return ContactItem(**contact_dict) if contact_dict else None + 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 {contact_id}: {e}") - return None + 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 + owner_id: Optional[int] = None, ) -> List[ContactItem]: - """Find contacts matching the query.""" + """Find contacts matching the given query.""" try: - # If owner_id is provided, add it to the query - if owner_id is not None: - query["owner_id"] = owner_id - - # Use the queue's find_many method - contact_dicts = await self.queue.find_many(query, user_id=owner_id, limit=limit) - return [ContactItem(**doc) for doc in contact_dicts] + 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}") - return [] + 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, - limit: int = 10 ) -> List[ContactItem]: """Search for contacts by name, email, etc.""" try: @@ -184,49 +226,41 @@ async def search_contacts( {"notes": {"$regex": text, "$options": "i"}} ] } - - # Add owner_id constraint if provided - if owner_id is not None: - query["owner_id"] = owner_id - - # Use find_many with the query - contact_dicts = await self.queue.find_many(query, user_id=owner_id, limit=limit) - return [ContactItem(**doc) for doc in contact_dicts] + return await self.find_contacts( + query, + limit=limit, + owner_id=owner_id + ) except Exception as e: logger.error(f"Error searching contacts: {e}") - return [] + 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.""" + async def get_random_contact( + self, owner_id: Optional[int] = None + ) -> Optional[ContactItem]: + """Get a random contact from the queue.""" try: - # Use the queue's get_random_record method - contact_dict = await self.queue.get_random_record(user_id=owner_id) - return ContactItem(**contact_dict) if contact_dict else None + 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}") - return None + from botspot.core.errors import ContactDataError + raise ContactDataError("Failed to get random contact") from e async def parse_contact_with_llm( - self, - text: str, - user_id: Optional[int] = None + self, text: str, owner_id: Optional[int] = None ) -> Optional[ContactItem]: - """ - Parse contact information from text using LLM. - - Args: - text: Text message to parse - user_id: User ID who sent the message (for ownership) - - Returns: - Contact object if parsing successful, else None - """ + """Parse contact information from text using LLM.""" try: - from botspot.components.new.llm_provider import get_llm_provider from botspot.utils.deps_getters import get_llm_provider - # 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 = """ @@ -260,21 +294,83 @@ class TempContact(BaseModel): temp_contact = await llm.aquery_llm_structured( prompt=text, output_schema=TempContact, - user=user_id, + 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 user_id is not None: - contact_data["owner_id"] = user_id + 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}") - return None + 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 # --------------------------------------------- @@ -458,14 +554,14 @@ def setup_dispatcher(dp: Dispatcher, **kwargs): return dp -def initialize(settings: ContactManagerSettings) -> ContactManager: +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 = ContactManager(settings) + manager = await ContactManager.create(settings) return manager @@ -491,52 +587,76 @@ def get_contact_manager() -> ContactManager: # --------------------------------------------- -async def add_contact(contact: ContactItem, user_id: Optional[int] = None) -> bool: +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, user_id) + return await manager.add_contact(contact, owner_id) -async def update_contact(contact_id: str, data: Dict[str, Any], user_id: Optional[int] = None) -> bool: +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, user_id) + return await manager.update_contact(contact_id, data, owner_id) -async def delete_contact(contact_id: str, user_id: Optional[int] = None) -> bool: +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, user_id) + return await manager.delete_contact(contact_id, owner_id) -async def get_contact_by_id(contact_id: str, user_id: Optional[int] = None) -> Optional[ContactItem]: +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, user_id) + return await manager.get_contact_by_id(contact_id, owner_id) -async def find_contacts(query: Dict[str, Any], user_id: Optional[int] = None, limit: int = 10) -> List[ContactItem]: +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=limit, owner_id=user_id) + return await manager.find_contacts(query, limit, owner_id) -async def search_contacts(text: str, user_id: Optional[int] = None, limit: int = 10) -> List[ContactItem]: +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, user_id, limit) + return await manager.search_contacts(text, limit, owner_id) -async def get_random_contact(user_id: Optional[int] = None) -> Optional[ContactItem]: +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(user_id) + return await manager.get_random_contact(owner_id) -async def parse_contact_with_llm(text: str, user_id: Optional[int] = None) -> Optional[ContactItem]: +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, user_id) + 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/contact_manager.py b/botspot/contact_manager.py index 9c883d8..1dbf0cc 100644 --- a/botspot/contact_manager.py +++ b/botspot/contact_manager.py @@ -7,6 +7,10 @@ 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__ = [ @@ -18,5 +22,9 @@ "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/tests/components/new/test_contact_manager.py b/tests/components/new/test_contact_manager.py index 1558448..f521624 100644 --- a/tests/components/new/test_contact_manager.py +++ b/tests/components/new/test_contact_manager.py @@ -1,9 +1,11 @@ import asyncio import datetime import pytest -from unittest.mock import AsyncMock, MagicMock, patch +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 @@ -19,140 +21,121 @@ def contact_manager_settings(): @pytest.fixture def mock_queue(): - """Fixture to create a mock queue for the contact manager""" + """Create a mock queue for testing.""" queue = AsyncMock() - queue.collection = MagicMock() - queue.collection.find_one.return_value = None # Default: contact not found - queue.find_one.return_value = None + 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.fixture -def contact_manager(contact_manager_settings, mock_queue): - """Fixture to create a contact manager with a mock queue""" - with patch("botspot.components.new.contact_manager.create_queue", return_value=mock_queue): - manager = ContactManager(contact_manager_settings) - return manager +@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(): - """Fixture to create a sample contact""" + """Create a sample contact for testing.""" return ContactItem( + data="Test Contact", name="Test Contact", phone="555-1234", email="test@example.com", - telegram="@testcontact", - birthday=datetime.date(1990, 1, 1), - notes="Test notes", - owner_id=9876 + 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 -def test_add_contact(contact_manager, sample_contact): +@pytest.mark.asyncio +async def test_add_contact(contact_manager, sample_contact): """Test adding a contact""" - contact_manager.queue.add_item.return_value = None - - # Execute - loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.add_contact(sample_contact)) - - # Assert - assert result == True + result = await contact_manager.add_contact(sample_contact) + assert result is True contact_manager.queue.add_item.assert_called_once() -def test_update_contact(contact_manager, sample_contact): +@pytest.mark.asyncio +async def test_update_contact(contact_manager, sample_contact): """Test updating a contact""" - contact_manager.queue.update_record.return_value = True - - # Execute - loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.update_contact("contact123", {"name": "Updated Name"})) - - # Assert - assert result == True + 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() -def test_delete_contact(contact_manager): +@pytest.mark.asyncio +async def test_delete_contact(contact_manager): """Test deleting a contact""" - contact_manager.queue.delete_record.return_value = True - - # Execute - loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.delete_contact("contact123")) - - # Assert - assert result == True - contact_manager.queue.delete_record.assert_called_once_with("contact123") + result = await contact_manager.delete_contact("test_id") + assert result is True + contact_manager.queue.delete_record.assert_called_once() -def test_get_contact_by_id(contact_manager, sample_contact): +@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.return_value = contact_dict - - # Execute - loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.get_contact_by_id("contact123", owner_id=9876)) - - # Assert + 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 == "Test Contact" - contact_manager.queue.find.assert_called_once_with({"_id": "contact123"}, user_id=9876) + assert result.name == sample_contact.name + contact_manager.queue.find.assert_called_once() -def test_find_contacts(contact_manager, sample_contact): +@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.return_value = contact_dicts - - # Execute - loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.find_contacts({"name": "Test Contact"})) - - # Assert + 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 == "Test Contact" + assert result[0].name == sample_contact.name contact_manager.queue.find_many.assert_called_once() -def test_search_contacts(contact_manager, sample_contact): +@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.return_value = contact_dicts - - # Execute - loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.search_contacts("Test")) - - # Assert + 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 == "Test Contact" + assert result[0].name == sample_contact.name contact_manager.queue.find_many.assert_called_once() -def test_get_random_contact(contact_manager, sample_contact): +@pytest.mark.asyncio +async def test_get_random_contact(contact_manager, sample_contact): """Test getting a random contact""" - contact_manager.queue.get_random_record.return_value = sample_contact.model_dump() - - # Execute - loop = asyncio.get_event_loop() - result = loop.run_until_complete(contact_manager.get_random_contact()) - - # Assert + 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 == "Test Contact" + assert result.name == sample_contact.name contact_manager.queue.get_random_record.assert_called_once() @@ -165,6 +148,7 @@ async def test_parse_contact_with_llm_success(): # Mock the LLM response - create an actual ContactItem mock_contact = ContactItem( + data="Jane Smith", name="Jane Smith", phone="555-5678", email="jane@example.com", @@ -172,19 +156,21 @@ async def test_parse_contact_with_llm_success(): ) mock_llm.aquery_llm_structured.return_value = mock_contact - # Create mock queue + # 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 and llm provider + # 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.create_queue", return_value=mock_queue): + patch("botspot.components.new.contact_manager.get_queue_manager", return_value=mock_queue_manager): - # Create manager directly - don't rely on 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", user_id=1234) + contact = await manager.parse_contact_with_llm("Jane Smith, 555-5678, jane@example.com, @janesmith", owner_id=1234) # Assert assert contact is not None @@ -196,7 +182,6 @@ async def test_parse_contact_with_llm_success(): # Verify LLM was called with appropriate system message mock_llm.aquery_llm_structured.assert_called_once() - assert "parser" in mock_llm.aquery_llm_structured.call_args[1]["system_message"].lower() @pytest.mark.asyncio @@ -207,32 +192,36 @@ async def test_parse_contact_with_llm_missing_fields(): # 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 + # 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.create_queue", return_value=mock_queue): + 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", user_id=1234) + 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 - assert contact.owner_id == 1234 @pytest.mark.asyncio @@ -244,22 +233,22 @@ async def test_parse_contact_with_llm_failure(): # Mock the LLM to raise an exception on structured output mock_llm.aquery_llm_structured.side_effect = Exception("Invalid output format") - # Create mock queue + # 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.create_queue", return_value=mock_queue): + 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", user_id=1234) - - # Assert - assert contact is None + # Test parsing - should raise ContactDataError + with pytest.raises(Exception): + await manager.parse_contact_with_llm("Invalid contact info", owner_id=1234) @pytest.mark.asyncio @@ -271,67 +260,129 @@ async def test_parse_contact_llm_exception(): # Mock the LLM to raise an exception mock_llm.aquery_llm_structured.side_effect = Exception("LLM service unavailable") - # Create mock queue + # 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.create_queue", return_value=mock_queue): + 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", user_id=1234) - - # Assert - assert contact is None + # 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_add_contact_command_handler(): - """Test the add_contact command handler""" - # Skip this test for now - will be implemented after fixing other issues - pytest.skip("Handler tests will be implemented after fixing core functionality") +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_find_contact_command_handler(): - """Test the find_contact command handler""" - # Skip this test for now - will be implemented after fixing other issues - pytest.skip("Handler tests will be implemented after fixing core functionality") +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_random_contact_command_handler(): - """Test the random_contact command handler""" - # Skip this test for now - will be implemented after fixing other issues - pytest.skip("Handler tests will be implemented after fixing core functionality") +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_message_parser(): - """Test the message parser that automatically detects contact info in chat messages""" - # Skip this test for now - will be implemented after fixing other issues - pytest.skip("Handler tests will be implemented after fixing core functionality") +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" -# Integration tests @pytest.mark.asyncio -async def test_initialization(): - """Test the initialization function""" - settings = ContactManagerSettings(enabled=True, collection="test_contacts") +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()]) - with patch("botspot.components.new.contact_manager.create_queue") as mock_create_queue: - mock_queue = AsyncMock() - mock_create_queue.return_value = mock_queue - - # Call initialize - manager = initialize(settings) - - # Assert - assert isinstance(manager, ContactManager) - - # Verify create_queue was called with the correct parameters - note the 'key' named parameter - mock_create_queue.assert_called_once_with(key="contacts", item_model=ContactItem) \ No newline at end of file + # 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