3 - ⭐️[Example ready] context builder

botspot/components/new/context_builder.py

@@ -1,57 +1,28 @@
 """
+Context Builder Component
 
-1 - tokens
-a) make sure necessary libs are installed - anthropic, openai etc
-b)
-c) how do I get a client that's working?
-
-2 - utils. for calmlib
-a)
-b)
-c)
-d)
-
-3 - quotas
-a) for friend - unlimited
-b) single user mode - unlimited
-c)
-d)
-
-----
-
-1) Simple query_llm funciton
-2) support 'telegram user' arg or check for single_user_mode
-3) future: support interoperability with calmlib: make calmlib query_llm check if botspot is enabled, if so - use botspot's llm provider
-4) async (have both versions - complete and acomplete)
-5)
-6)
-
-Decisions:
-- Do I need LLMProvider class? Just for the sake of having a class corresponding to the component?
-    - to wire with db, deps etc? Well, i store all those in deps.. do i need like properties or something?
-- Do I need _llm_query_base function? e.g. to return raw response, parsed json or just text
--
-
+Collects messages forwarded together and bundles them as a single context object.
 """
+from typing import Any, Awaitable, Callable, Dict, List, Optional
+from datetime import datetime, timedelta
 
-from typing import TYPE_CHECKING
+from aiogram import Dispatcher
+from aiogram.types import Message, TelegramObject, Update
+from pydantic import Field
+from pydantic_settings import BaseSettings
 
 from botspot.utils.internal import get_logger
 
-if TYPE_CHECKING:
-    from motor.motor_asyncio import AsyncIOMotorDatabase  # noqa: F401
-
 logger = get_logger()
-from typing import TYPE_CHECKING
-
-from pydantic_settings import BaseSettings
-
-if TYPE_CHECKING:
-    from motor.motor_asyncio import AsyncIOMotorCollection  # noqa: F401
 
 
 class ContextBuilderSettings(BaseSettings):
     enabled: bool = False
+    bundle_timeout: float = 0.5  # Time window in seconds to bundle messages
+    include_replies: bool = True
+    include_forwards: bool = True
+    include_media: bool = True
+    include_text: bool = True
 
     class Config:
         env_prefix = "BOTSPOT_CONTEXT_BUILDER_"
@@ -60,20 +31,148 @@ class Config:
         extra = "ignore"
 
 
-class ContextBuilder:
-    pass
+class MessageContext:
+    """Represents a bundle of messages that form a context."""
+
+    def __init__(self):
+        self.messages: List[Message] = []
+        self.last_update_time: datetime = datetime.now()
+        self.text_content: str = ""
+        self.media_descriptions: List[str] = []
+
+    def add_message(self, message: Message) -> None:
+        """Add a message to the context bundle."""
+        self.messages.append(message)
+        self.last_update_time = datetime.now()
+
+        # Extract text content
+        if message.text:
+            if self.text_content:
+                self.text_content += f"\n\n{message.text}"
+            else:
+                self.text_content = message.text
+
+        # Extract media descriptions
+        if message.photo:
+            self.media_descriptions.append("[Image]")
+        elif message.video:
+            self.media_descriptions.append("[Video]")
+        elif message.voice:
+            self.media_descriptions.append("[Voice Message]")
+        elif message.audio:
+            self.media_descriptions.append("[Audio]")
+        elif message.document:
+            self.media_descriptions.append(f"[Document: {message.document.file_name}]")
+
+    def is_expired(self, timeout_seconds: float) -> bool:
+        """Check if the context has expired based on timeout."""
+        return (datetime.now() - self.last_update_time) > timedelta(seconds=timeout_seconds)
+
+    def get_combined_text(self) -> str:
+        """Get the combined text from all messages in the context."""
+        return self.text_content
+
+    def get_full_context(self) -> str:
+        """Get the full context including text and media descriptions."""
+        result = self.text_content
+        if self.media_descriptions:
+            media_text = "\n".join(self.media_descriptions)
+            if result:
+                result += f"\n\n{media_text}"
+            else:
+                result = media_text
+        return result
 
 
-def setup_dispatcher(dp):
+class ContextBuilder:
+    """Component that builds context from multiple messages."""
+
+    def __init__(self, settings: ContextBuilderSettings):
+        self.settings = settings
+        self.message_contexts: Dict[int, MessageContext] = {}  # chat_id -> MessageContext
+
+    async def build_context(self, message: Message) -> Optional[str]:
+        """
+        Build context from a message and its related messages.
+        Returns the current context text if available.
+        """
+        chat_id = message.chat.id
+
+        # Process message based on settings
+        context_text = None
+
+        # Get or create message context for this chat
+        if chat_id not in self.message_contexts:
+            self.message_contexts[chat_id] = MessageContext()
+
+        context = self.message_contexts[chat_id]
+
+        # If context has expired, create a new one
+        if context.is_expired(self.settings.bundle_timeout):
+            # If there were messages in the previous context, we consider it complete
+            if context.messages:
+                context_text = context.get_full_context()
+            # Start a new context
+            self.message_contexts[chat_id] = MessageContext()
+            context = self.message_contexts[chat_id]
+
+        # Add the current message to the context
+        context.add_message(message)
+
+        return context_text
+
+
+async def context_builder_middleware(
+    handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
+    event: Update,
+    data: Dict[str, Any],
+) -> Any:
+    """
+    Middleware that builds context from messages and adds it to the message metadata.
+    """
+    # Only process message events
+    if not hasattr(event, "message") or event.message is None or not isinstance(event.message, Message):
+        return await handler(event, data)
+
+    message = event.message
+
+    from botspot.core.dependency_manager import get_dependency_manager
+    deps = get_dependency_manager()
+
+    if not hasattr(deps, "context_builder"):
+        return await handler(event, data)
+
+    context_builder = deps.context_builder
+
+    # Build context from the message
+    context_text = await context_builder.build_context(message)
+
+    # Add context to message data
+    if context_text:
+        logger.debug(f"Adding context to message data: {context_text[:100]}...")
+        data["context_text"] = context_text
+
+    # Set metadata flag to indicate that this message is part of a context bundle
+    data["is_in_context_bundle"] = True
+
+    return await handler(event, data)
+
+
+def setup_dispatcher(dp: Dispatcher):
+    """Set up the context builder middleware."""
+    logger.debug("Adding context builder middleware")
+    dp.update.middleware(context_builder_middleware)
     return dp
 
 
 def initialize(settings: ContextBuilderSettings) -> ContextBuilder:
-    pass
+    """Initialize the context builder component."""
+    logger.info("Initializing context builder component")
+    return ContextBuilder(settings)
 
 
 def get_context_builder() -> ContextBuilder:
+    """Get the context builder from the dependency manager."""
     from botspot.core.dependency_manager import get_dependency_manager
-
     deps = get_dependency_manager()
-    return deps.context_builder
+    return deps.context_builder

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, queue_manager
+from botspot.components.new import chat_binder, context_builder, llm_provider, queue_manager
 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
@@ -58,9 +58,12 @@ def __init__(
         if self.settings.chat_binder.enabled:
             self.deps.chat_binder = chat_binder.initialize(self.settings.chat_binder)
 
+        if self.settings.context_builder.enabled:
+            self.deps.context_builder = context_builder.initialize(self.settings.context_builder)
+
         if self.settings.llm_provider.enabled:
             self.deps.llm_provider = llm_provider.initialize(self.settings.llm_provider)
-            
+
         if self.settings.queue_manager.enabled:
             self.deps.queue_manager = queue_manager.initialize(self.settings.queue_manager)
 
@@ -106,8 +109,11 @@ def setup_dispatcher(self, dp: Dispatcher):
         if self.settings.chat_binder.enabled:
             chat_binder.setup_dispatcher(dp)
 
+        if self.settings.context_builder.enabled:
+            context_builder.setup_dispatcher(dp)
+
         if self.settings.llm_provider.enabled:
             llm_provider.setup_dispatcher(dp)
-            
+
         if self.settings.queue_manager.enabled:
             queue_manager.setup_dispatcher(dp)

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.context_builder import ContextBuilderSettings
 from botspot.components.new.llm_provider import LLMProviderSettings
 from botspot.components.new.queue_manager import QueueManagerSettings
 from botspot.components.qol.bot_commands_menu import BotCommandsMenuSettings
@@ -59,6 +60,7 @@ def friends(self) -> List[str]:
     send_safe: SendSafeSettings = SendSafeSettings()
     admin_filter: AdminFilterSettings = AdminFilterSettings()
     chat_binder: ChatBinderSettings = ChatBinderSettings()
+    context_builder: ContextBuilderSettings = ContextBuilderSettings()
     llm_provider: LLMProviderSettings = LLMProviderSettings()
     queue_manager: QueueManagerSettings = QueueManagerSettings()
 

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.context_builder import ContextBuilder
     from botspot.components.new.chat_binder import ChatBinder
     from botspot.components.new.llm_provider import LLMProvider
     from botspot.components.new.queue_manager import QueueManager
@@ -39,6 +40,7 @@ def __init__(
         self._telethon_manager = None
         self._user_manager = None
         self._chat_binder = None
+        self._context_builder = None
         self._llm_provider = None
         self._queue_manager = None
         self.__dict__.update(kwargs)
@@ -131,6 +133,16 @@ def chat_binder(self) -> "ChatBinder":
     def chat_binder(self, value):
         self._chat_binder = value
 
+    @property
+    def context_builder(self) -> "ContextBuilder":
+        if self._context_builder is None:
+            raise RuntimeError("Context Builder is not initialized")
+        return self._context_builder
+
+    @context_builder.setter
+    def context_builder(self, value):
+        self._context_builder = value
+
     @property
     def llm_provider(self) -> "LLMProvider":
         if self._llm_provider is None:

botspot/utils/deps_getters.py

@@ -12,6 +12,7 @@
 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.context_builder import get_context_builder
 from botspot.components.new.chat_binder import get_chat_binder
 from botspot.components.new.queue_manager import get_queue_manager
 
@@ -74,4 +75,5 @@ async def get_telethon_client(
     "get_mongo_client",
     "get_chat_binder",
     "get_queue_manager",
+    "get_context_builder",
 ]

dev/workalong_context_builder/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
+