mojentic 0.6.2__py3-none-any.whl → 0.7.2__py3-none-any.whl

mojentic/async_dispatcher.py

@@ -0,0 +1,134 @@
+import asyncio
+import logging
+from collections import deque
+from typing import Optional, Type
+from uuid import uuid4
+
+import structlog
+
+from mojentic.event import TerminateEvent
+
+logger = structlog.get_logger()
+
+
+class AsyncDispatcher:
+    """
+    AsyncDispatcher class is an asynchronous version of the Dispatcher class.
+    It uses asyncio and deque for event processing.
+    """
+    def __init__(self, router, shared_working_memory=None, batch_size=5, tracer=None):
+        """
+        Initialize the AsyncDispatcher.
+
+        Parameters
+        ----------
+        router : Router
+            The router to use for routing events to agents
+        shared_working_memory : SharedWorkingMemory, optional
+            The shared working memory to use
+        batch_size : int, optional
+            The number of events to process in each batch
+        tracer : Tracer, optional
+            The tracer to use for tracing events
+        """
+        self.router = router
+        self.batch_size = batch_size
+        self.event_queue = deque()
+        self._stop_event = asyncio.Event()
+        self._task = None
+
+        # Use null_tracer if no tracer is provided
+        from mojentic.tracer import null_tracer
+        self.tracer = tracer or null_tracer
+
+    async def start(self):
+        """
+        Start the event dispatch task.
+        """
+        logger.debug("Starting event dispatch task")
+        self._task = asyncio.create_task(self._dispatch_events())
+        return self
+
+    async def stop(self):
+        """
+        Stop the event dispatch task.
+        """
+        self._stop_event.set()
+        if self._task:
+            await self._task
+
+    async def wait_for_empty_queue(self, timeout=None):
+        """
+        Wait for the event queue to be empty.
+
+        Parameters
+        ----------
+        timeout : float, optional
+            The timeout in seconds
+
+        Returns
+        -------
+        bool
+            True if the queue is empty, False if the timeout was reached
+        """
+        start_time = asyncio.get_event_loop().time()
+        while len(self.event_queue) > 0:
+            if timeout is not None and asyncio.get_event_loop().time() - start_time > timeout:
+                return False
+            await asyncio.sleep(0.1)
+        return True
+
+    def dispatch(self, event):
+        """
+        Dispatch an event to the event queue.
+
+        Parameters
+        ----------
+        event : Event
+            The event to dispatch
+        """
+        logger.log(logging.DEBUG, f"Dispatching event: {event}")
+        if event.correlation_id is None:
+            event.correlation_id = str(uuid4())
+        self.event_queue.append(event)
+
+    async def _dispatch_events(self):
+        """
+        Dispatch events from the event queue to agents.
+        """
+        while not self._stop_event.is_set():
+            for _ in range(self.batch_size):
+                logger.debug("Checking for events")
+                if len(self.event_queue) > 0:
+                    logger.debug(f"{len(self.event_queue)} events in queue")
+                    event = self.event_queue.popleft()
+                    logger.debug(f"Processing event: {event}")
+                    agents = self.router.get_agents(event)
+                    logger.debug(f"Found {len(agents)} agents for event type {type(event)}")
+                    events = []
+                    for agent in agents:
+                        logger.debug(f"Sending event to agent {agent}")
+
+                        # Record agent interaction in tracer system
+                        self.tracer.record_agent_interaction(
+                            from_agent=str(event.source),
+                            to_agent=str(type(agent)),
+                            event_type=str(type(event).__name__),
+                            event_id=event.correlation_id,
+                            source=type(self)
+                        )
+
+                        # Process the event through the agent
+                        # If the agent is an async agent, await its receive_event method
+                        if hasattr(agent, 'receive_event_async'):
+                            received_events = await agent.receive_event_async(event)
+                        else:
+                            received_events = agent.receive_event(event)
+
+                        logger.debug(f"Agent {agent} returned {len(received_events)} events")
+                        events.extend(received_events)
+                    for fe in events:
+                        if type(fe) is TerminateEvent:
+                            self._stop_event.set()
+                        self.dispatch(fe)
+            await asyncio.sleep(0.1)  # Use asyncio.sleep instead of time.sleep

mojentic/async_dispatcher_spec.py

@@ -0,0 +1,244 @@
+import asyncio
+import pytest
+import pytest_asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+from mojentic.async_dispatcher import AsyncDispatcher
+from mojentic.event import Event, TerminateEvent
+from mojentic.agents.base_async_agent import BaseAsyncAgent
+from mojentic.router import Router
+
+
+class SampleEvent(Event):
+    """A simple event for testing."""
+    message: str
+
+
+class SampleResponseEvent(Event):
+    """A response event for testing."""
+    response: str
+
+
+class AsyncTestAgent(BaseAsyncAgent):
+    """A test async agent that returns a SampleResponseEvent."""
+
+    async def receive_event_async(self, event):
+        if isinstance(event, SampleEvent):
+            return [SampleResponseEvent(
+                source=type(self),
+                correlation_id=event.correlation_id,
+                response=f"Processed: {event.message}"
+            )]
+        return []
+
+
+class SyncTestAgent:
+    """A test sync agent that returns a SampleResponseEvent."""
+
+    def __init__(self):
+        # Create a mock for tracking calls
+        self._mock = MagicMock()
+
+    def receive_event(self, event):
+        # Track the call
+        self._mock(event)
+
+        if isinstance(event, SampleEvent):
+            return [SampleResponseEvent(
+                source=type(self),
+                correlation_id=event.correlation_id,
+                response=f"Processed sync: {event.message}"
+            )]
+        return []
+
+    def assert_called_once(self):
+        # Delegate to the mock
+        self._mock.assert_called_once()
+
+
+@pytest.fixture
+def router():
+    """Create a router for testing."""
+    return Router()
+
+
+@pytest.fixture
+def async_agent():
+    """Create an async agent for testing."""
+    return AsyncTestAgent()
+
+
+@pytest.fixture
+def sync_agent():
+    """Create a sync agent for testing."""
+    return SyncTestAgent()
+
+
+@pytest_asyncio.fixture
+async def dispatcher(router):
+    """Create and start an AsyncDispatcher for testing."""
+    dispatcher = AsyncDispatcher(router)
+    await dispatcher.start()
+    yield dispatcher
+    await dispatcher.stop()
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_init(router):
+    """Test that the AsyncDispatcher initializes correctly."""
+    dispatcher = AsyncDispatcher(router)
+
+    assert dispatcher.router == router
+    assert dispatcher.batch_size == 5
+    assert len(dispatcher.event_queue) == 0
+    assert dispatcher._task is None
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_start_stop(router):
+    """Test that the AsyncDispatcher starts and stops correctly."""
+    dispatcher = AsyncDispatcher(router)
+
+    # Start the dispatcher
+    await dispatcher.start()
+    assert dispatcher._task is not None
+    assert not dispatcher._stop_event.is_set()
+
+    # Stop the dispatcher
+    await dispatcher.stop()
+    assert dispatcher._stop_event.is_set()
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_dispatch(dispatcher, router, async_agent):
+    """Test that the AsyncDispatcher dispatches events correctly."""
+    # Register the agent
+    router.add_route(SampleEvent, async_agent)
+
+    # Create and dispatch an event
+    event = SampleEvent(source=str, message="Hello")
+    dispatcher.dispatch(event)
+
+    # Wait for the event to be processed
+    await dispatcher.wait_for_empty_queue(timeout=1)
+
+    # The event should have been processed and removed from the queue
+    assert len(dispatcher.event_queue) == 0
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_with_async_agent(dispatcher, router, async_agent):
+    """Test that the AsyncDispatcher works with async agents."""
+    # Register the agent
+    router.add_route(SampleEvent, async_agent)
+
+    # Create and dispatch an event
+    event = SampleEvent(source=str, message="Hello")
+    dispatcher.dispatch(event)
+
+    # Wait for the event to be processed
+    await dispatcher.wait_for_empty_queue(timeout=1)
+
+    # The event should have been processed and removed from the queue
+    assert len(dispatcher.event_queue) == 0
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_with_sync_agent(dispatcher, router, sync_agent):
+    """Test that the AsyncDispatcher works with sync agents."""
+    # Register the agent
+    router.add_route(SampleEvent, sync_agent)
+
+    # Create and dispatch an event
+    event = SampleEvent(source=str, message="Hello")
+    dispatcher.dispatch(event)
+
+    # Wait for the event to be processed
+    await dispatcher.wait_for_empty_queue(timeout=1)
+
+    # The event should have been processed and removed from the queue
+    assert len(dispatcher.event_queue) == 0
+
+    # Verify that the sync agent's receive_event method was called
+    sync_agent.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_terminate_event(dispatcher, router):
+    """Test that the AsyncDispatcher handles TerminateEvent correctly."""
+    # Create a mock agent that returns a TerminateEvent
+    class TerminateAgent(BaseAsyncAgent):
+        async def receive_event_async(self, event):
+            return [TerminateEvent(source=type(self))]
+
+    terminate_agent = TerminateAgent()
+
+    # Register the agent
+    router.add_route(SampleEvent, terminate_agent)
+
+    # Create and dispatch an event
+    event = SampleEvent(source=str, message="Hello")
+    dispatcher.dispatch(event)
+
+    # Wait a moment for the event to be processed
+    await asyncio.sleep(0.5)
+
+    # The dispatcher should have stopped
+    assert dispatcher._stop_event.is_set()
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_wait_for_empty_queue(dispatcher):
+    """Test that wait_for_empty_queue works correctly."""
+    # Queue is initially empty
+    assert await dispatcher.wait_for_empty_queue() is True
+
+    # Add multiple events to the queue to test batch processing
+    for i in range(10):
+        event = SampleEvent(source=str, message=f"Event {i}")
+        dispatcher.dispatch(event)
+
+    # Queue is not empty
+    assert len(dispatcher.event_queue) > 0
+
+    # Wait for the queue to be empty with a sufficient timeout
+    # The dispatcher should process all events
+    assert await dispatcher.wait_for_empty_queue(timeout=1) is True
+
+    # Verify the queue is empty
+    assert len(dispatcher.event_queue) == 0
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_batch_processing(dispatcher, router, async_agent):
+    """Test that the AsyncDispatcher processes events in batches."""
+    # Register the agent
+    router.add_route(SampleEvent, async_agent)
+
+    # Set a small batch size
+    dispatcher.batch_size = 2
+
+    # Create and dispatch multiple events
+    for i in range(5):
+        event = SampleEvent(source=str, message=f"Event {i}")
+        dispatcher.dispatch(event)
+
+    # Wait for all events to be processed
+    await dispatcher.wait_for_empty_queue(timeout=2)
+
+    # All events should have been processed
+    assert len(dispatcher.event_queue) == 0
+
+
+@pytest.mark.asyncio
+async def test_async_dispatcher_correlation_id(dispatcher):
+    """Test that the AsyncDispatcher assigns correlation_id if not provided."""
+    # Create an event without a correlation_id
+    event = SampleEvent(source=str, message="Hello")
+    assert event.correlation_id is None
+
+    # Dispatch the event
+    dispatcher.dispatch(event)
+
+    # The event should now have a correlation_id
+    assert event.correlation_id is not None

mojentic/context/__init__.py

@@ -1 +1,5 @@
+"""
+Mojentic context module for managing shared working memory and context.
+"""
+
 from .shared_working_memory import SharedWorkingMemory

mojentic/llm/__init__.py

@@ -1,4 +1,16 @@
+"""
+Mojentic LLM module for interacting with Large Language Models.
+"""
+
+# Main LLM components
 from .llm_broker import LLMBroker
-from .registry.llm_registry import LLMRegistry
 from .chat_session import ChatSession
-from .message_composers import MessageBuilder, FileTypeSensor
+from .message_composers import MessageBuilder, FileTypeSensor
+from .registry.llm_registry import LLMRegistry
+
+# Re-export gateway components at the LLM level
+from .gateways.models import (
+    LLMMessage,
+    LLMGatewayResponse,
+    MessageRole
+)

mojentic/llm/gateways/__init__.py

@@ -1,3 +1,25 @@
+"""
+Mojentic LLM gateways module for connecting to various LLM providers.
+"""
+
+# Gateway implementations
 from .llm_gateway import LLMGateway
 from .ollama import OllamaGateway
 from .openai import OpenAIGateway
+from .anthropic import AnthropicGateway
+from .file_gateway import FileGateway
+from .embeddings_gateway import EmbeddingsGateway
+from .tokenizer_gateway import TokenizerGateway
+
+# Message adapters
+from .anthropic_messages_adapter import adapt_messages_to_anthropic
+from .ollama_messages_adapter import adapt_messages_to_ollama
+from .openai_messages_adapter import adapt_messages_to_openai
+
+# Common models
+from .models import (
+    LLMMessage,
+    MessageRole,
+    LLMGatewayResponse,
+    LLMToolCall
+)

mojentic/llm/gateways/models.py

@@ -46,7 +46,7 @@ class LLMToolCall(BaseModel):
     """
     A tool call to be made available to the LLM.
 
-    Parameters
+    Attributes
     ----------
     id : Optional[str]
         The identifier of the tool call.
@@ -64,7 +64,7 @@ class LLMMessage(BaseModel):
     """
     A message to be sent to the LLM. These would accumulate during a chat session with an LLM.
 
-    Parameters
+    Attributes
     ----------
     role : MessageRole
         The role of the message in the conversation.
@@ -89,7 +89,7 @@ class LLMGatewayResponse(BaseModel):
     """
     The response from the LLM gateway, abstracting you from the quirks of a specific LLM.
 
-    Parameters
+    Attributes
     ----------
     content : Optional[Union[str, dict[str, str]]]
         The content of the response.

mojentic/llm/gateways/ollama.py

@@ -41,8 +41,8 @@ class OllamaGateway(LLMGateway):
         """
         Complete the LLM request by delegating to the Ollama service.
 
-        Parameters
-        ----------
+        Keyword Arguments
+        ----------------
         model : str
             The name of the model to use, as appears in `ollama list`.
         messages : List[LLMMessage]
@@ -107,8 +107,8 @@ class OllamaGateway(LLMGateway):
         """
         Stream the LLM response from Ollama service.
 
-        Parameters
-        ----------
+        Keyword Arguments
+        ----------------
         model : str
             The name of the model to use, as appears in `ollama list`.
         messages : List[LLMMessage]

mojentic/llm/gateways/openai.py

@@ -31,8 +31,8 @@ class OpenAIGateway(LLMGateway):
         """
         Complete the LLM request by delegating to the OpenAI service.
 
-        Parameters
-        ----------
+        Keyword Arguments
+        ----------------
         model : str
             The name of the model to use, as appears in `ollama list`.
         messages : List[LLMMessage]
@@ -148,4 +148,4 @@ class OpenAIGateway(LLMGateway):
         tokenizer = TokenizerGateway()
         tokens = tokenizer.encode(text)
         chunks_iterator = self._batched(tokens, chunk_length)
-        yield from chunks_iterator
+        yield from chunks_iterator

mojentic/llm/gateways/openai_messages_adapter.py

@@ -48,13 +48,17 @@ def get_image_type(file_path: str) -> str:
         file_path: Path to the image file
 
     Returns:
-        Image type (e.g., 'jpg', 'png')
+        Image type (e.g., 'jpeg', 'png')
     """
     _, ext = os.path.splitext(file_path)
     image_type = ext.lstrip('.').lower()
-    if image_type not in ['jpeg', 'jpg', 'png', 'gif', 'webp']:
-        image_type = 'jpeg'  # Default to jpeg if unknown extension
-    return image_type
+
+    # Convert 'jpg' to 'jpeg'
+    if image_type == 'jpg':
+        return 'jpeg'
+
+    # Use 'jpeg' for unknown extensions, otherwise use the detected type
+    return image_type if image_type in ['jpeg', 'png', 'gif', 'webp'] else 'jpeg'
 
 
 def adapt_messages_to_openai(messages: List[LLMMessage]):

mojentic/llm/llm_broker.py

@@ -88,7 +88,7 @@ class LLMBroker():
         logger.info(f"Requesting llm response with approx {approximate_tokens} tokens")
 
         # Convert messages to serializable dict for audit
-        messages_for_tracer = [m.dict() for m in messages]
+        messages_for_tracer = [m.model_dump() for m in messages]
 
         # Record LLM call in tracer
         tools_for_tracer = [{"name": t.name, "description": t.description} for t in tools] if tools else None
@@ -115,7 +115,7 @@ class LLMBroker():
         call_duration_ms = (time.time() - start_time) * 1000
 
         # Record LLM response in tracer
-        tool_calls_for_tracer = [tc.dict() for tc in result.tool_calls] if result.tool_calls else None
+        tool_calls_for_tracer = [tc.model_dump() for tc in result.tool_calls] if result.tool_calls else None
         self.tracer.record_llm_response(
             self.model,
             result.content,
@@ -199,7 +199,7 @@ class LLMBroker():
         logger.info(f"Requesting llm response with approx {approximate_tokens} tokens")
 
         # Convert messages to serializable dict for audit
-        messages_for_tracer = [m.dict() for m in messages]
+        messages_for_tracer = [m.model_dump() for m in messages]
 
         # Record LLM call in tracer
         self.tracer.record_llm_call(
@@ -221,7 +221,7 @@ class LLMBroker():
 
         # Record LLM response in tracer with object representation
         # Convert object to string for tracer
-        object_str = str(result.object.dict()) if hasattr(result.object, "dict") else str(result.object)
+        object_str = str(result.object.model_dump()) if hasattr(result.object, "model_dump") else str(result.object)
         self.tracer.record_llm_response(
             self.model,
             f"Structured response: {object_str}",

mojentic/llm/message_composers.py

@@ -216,7 +216,7 @@ class MessageBuilder():
 
         return self
 
-    def add_files(self, *file_paths: Union[str, Path]) -> "MessageBuilder":
+    def add_files(self, *file_paths: List[Union[str, Path]]) -> "MessageBuilder":
         """
         Add multiple text files to the message, ignoring binary files.
 

mojentic/llm/registry/__init__.py

@@ -0,0 +1,6 @@
+"""
+Mojentic LLM registry module for managing model registrations.
+"""
+
+from .llm_registry import LLMRegistry
+from .models import ModelInfo, Modality, Quantization

mojentic/llm/registry/populate_registry_from_ollama.py

@@ -46,15 +46,16 @@ def register_llms_from_ollama(url: str, registry: LLMRegistry):
         registry.register(entry)
 
 
-# Example usage
-ollama_url = "http://localhost:11434/api/tags"
-registry = LLMRegistry()
-register_llms_from_ollama(ollama_url, registry)
-
-# print(f"Tool using: {registry.find_first(tools=True, structured_output=True).name}")
-print(f"Fastest model: {registry.find_fastest().name}")
-print(f"Fastest model with tools: {registry.find_fastest(tools=True).name}")
-print(f"Fastest model with structured output: {registry.find_fastest(structured_output=True).name}")
-print(f"Smartest model: {registry.find_smartest().name}")
-print(f"Smartest model with tools: {registry.find_smartest(tools=True).name}")
-print(f"Smartest model with structured output: {registry.find_smartest(structured_output=True).name}")
+if __name__ == "__main__":
+    # Example usage
+    ollama_url = "http://localhost:11434/api/tags"
+    registry = LLMRegistry()
+    register_llms_from_ollama(ollama_url, registry)
+
+    # print(f"Tool using: {registry.find_first(tools=True, structured_output=True).name}")
+    print(f"Fastest model: {registry.find_fastest().name}")
+    print(f"Fastest model with tools: {registry.find_fastest(tools=True).name}")
+    print(f"Fastest model with structured output: {registry.find_fastest(structured_output=True).name}")
+    print(f"Smartest model: {registry.find_smartest().name}")
+    print(f"Smartest model with tools: {registry.find_smartest(tools=True).name}")
+    print(f"Smartest model with structured output: {registry.find_smartest(structured_output=True).name}")

mojentic/llm/tools/__init__.py

@@ -0,0 +1,18 @@
+"""
+Mojentic LLM tools module for extending LLM capabilities.
+"""
+
+# Base tool class
+from .llm_tool import LLMTool
+from .tool_wrapper import ToolWrapper
+
+# Common tools
+from .ask_user_tool import AskUserTool
+from .current_datetime import CurrentDateTimeTool
+from .date_resolver import ResolveDateTool
+from .organic_web_search import OrganicWebSearchTool
+from .tell_user_tool import TellUserTool
+
+# Import tool modules
+from . import file_manager
+from . import ephemeral_task_manager

mojentic/llm/tools/date_resolver.py

@@ -1,11 +1,14 @@
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from parsedatetime import Calendar, VERSION_CONTEXT_STYLE
 from pytz import timezone
 
-from mojentic.llm import LLMBroker
 from mojentic.llm.tools.llm_tool import LLMTool
 
+# Avoid circular imports with TYPE_CHECKING
+if TYPE_CHECKING:
+    from mojentic.llm.llm_broker import LLMBroker
+
 
 class ResolveDateTool(LLMTool):
     def run(self, relative_date_found: str, reference_date_in_iso8601: Optional[str] = None) -> dict[str, str]:

mojentic/llm/tools/ephemeral_task_manager/__init__.py

@@ -5,14 +5,14 @@ This module provides tools for appending, prepending, inserting, starting, compl
 Tasks follow a state machine that transitions from PENDING through IN_PROGRESS to COMPLETED.
 """
 
-from mojentic.llm.tools.ephemeral_task_manager.append_task_tool import AppendTaskTool
-from mojentic.llm.tools.ephemeral_task_manager.clear_tasks_tool import ClearTasksTool
-from mojentic.llm.tools.ephemeral_task_manager.complete_task_tool import CompleteTaskTool
-from mojentic.llm.tools.ephemeral_task_manager.insert_task_after_tool import InsertTaskAfterTool
-from mojentic.llm.tools.ephemeral_task_manager.list_tasks_tool import ListTasksTool
-from mojentic.llm.tools.ephemeral_task_manager.prepend_task_tool import PrependTaskTool
-from mojentic.llm.tools.ephemeral_task_manager.start_task_tool import StartTaskTool
-from mojentic.llm.tools.ephemeral_task_manager.ephemeral_task_list import EphemeralTaskList, Task
+from .append_task_tool import AppendTaskTool
+from .clear_tasks_tool import ClearTasksTool
+from .complete_task_tool import CompleteTaskTool
+from .insert_task_after_tool import InsertTaskAfterTool
+from .list_tasks_tool import ListTasksTool
+from .prepend_task_tool import PrependTaskTool
+from .start_task_tool import StartTaskTool
+from .ephemeral_task_list import EphemeralTaskList, Task
 
 __all__ = [
     "EphemeralTaskList",