mojentic 0.6.1__py3-none-any.whl → 0.7.1__py3-none-any.whl

mojentic/agents/async_aggregator_agent_spec.py

@@ -0,0 +1,227 @@
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+
+from mojentic.agents.async_aggregator_agent import AsyncAggregatorAgent
+from mojentic.event import Event
+
+
+class TestEvent1(Event):
+    """A test event type 1."""
+    message: str
+
+
+class TestEvent2(Event):
+    """A test event type 2."""
+    data: str
+
+
+class TestEvent3(Event):
+    """A test event type 3."""
+    value: int
+
+
+class TestResultEvent(Event):
+    """A test result event."""
+    result: str
+
+
+class TestAsyncAggregator(AsyncAggregatorAgent):
+    """A test implementation of AsyncAggregatorAgent."""
+    
+    async def process_events(self, events):
+        """Process the events and return a result event."""
+        # Extract events by type
+        event1 = next((e for e in events if isinstance(e, TestEvent1)), None)
+        event2 = next((e for e in events if isinstance(e, TestEvent2)), None)
+        
+        if event1 and event2:
+            # Create a result combining the events
+            return [TestResultEvent(
+                source=type(self),
+                correlation_id=event1.correlation_id,
+                result=f"{event1.message} - {event2.data}"
+            )]
+        return []
+
+
+@pytest.fixture
+def async_aggregator():
+    """Create an AsyncAggregatorAgent for testing."""
+    return AsyncAggregatorAgent(event_types_needed=[TestEvent1, TestEvent2])
+
+
+@pytest.fixture
+def test_async_aggregator():
+    """Create a TestAsyncAggregator for testing."""
+    return TestAsyncAggregator(event_types_needed=[TestEvent1, TestEvent2])
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_init():
+    """Test that the AsyncAggregatorAgent initializes correctly."""
+    agent = AsyncAggregatorAgent(event_types_needed=[TestEvent1, TestEvent2])
+    
+    assert agent.event_types_needed == [TestEvent1, TestEvent2]
+    assert agent.results == {}
+    assert agent.futures == {}
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_capture_results(async_aggregator):
+    """Test that the AsyncAggregatorAgent captures results correctly."""
+    event = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    
+    await async_aggregator._capture_results_if_needed(event)
+    
+    assert "test-id" in async_aggregator.results
+    assert len(async_aggregator.results["test-id"]) == 1
+    assert async_aggregator.results["test-id"][0] == event
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_has_all_needed(async_aggregator):
+    """Test that the AsyncAggregatorAgent checks if all needed events are captured."""
+    event1 = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    event2 = TestEvent2(source=str, correlation_id="test-id", data="World")
+    
+    # Initially, we don't have all needed events
+    assert not await async_aggregator._has_all_needed(event1)
+    
+    # Capture the first event
+    await async_aggregator._capture_results_if_needed(event1)
+    assert not await async_aggregator._has_all_needed(event1)
+    
+    # Capture the second event
+    await async_aggregator._capture_results_if_needed(event2)
+    assert await async_aggregator._has_all_needed(event2)
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_get_and_reset_results(async_aggregator):
+    """Test that the AsyncAggregatorAgent gets and resets results correctly."""
+    event1 = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    event2 = TestEvent2(source=str, correlation_id="test-id", data="World")
+    
+    # Capture the events
+    await async_aggregator._capture_results_if_needed(event1)
+    await async_aggregator._capture_results_if_needed(event2)
+    
+    # Get and reset the results
+    results = await async_aggregator._get_and_reset_results(event1)
+    
+    assert len(results) == 2
+    assert results[0] == event1
+    assert results[1] == event2
+    assert async_aggregator.results["test-id"] is None
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_wait_for_events(async_aggregator):
+    """Test that the AsyncAggregatorAgent waits for events correctly."""
+    event1 = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    event2 = TestEvent2(source=str, correlation_id="test-id", data="World")
+    
+    # Start waiting for events in a separate task
+    wait_task = asyncio.create_task(async_aggregator.wait_for_events("test-id", timeout=1))
+    
+    # Capture the events
+    await async_aggregator._capture_results_if_needed(event1)
+    await async_aggregator._capture_results_if_needed(event2)
+    
+    # Wait for the task to complete
+    results = await wait_task
+    
+    assert len(results) == 2
+    assert results[0] == event1
+    assert results[1] == event2
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_wait_for_events_timeout(async_aggregator):
+    """Test that the AsyncAggregatorAgent handles timeouts correctly."""
+    event1 = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    
+    # Capture only one event
+    await async_aggregator._capture_results_if_needed(event1)
+    
+    # Wait for events with a short timeout
+    results = await async_aggregator.wait_for_events("test-id", timeout=0.1)
+    
+    # We should get the partial results
+    assert len(results) == 1
+    assert results[0] == event1
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_receive_event_async(test_async_aggregator):
+    """Test that the AsyncAggregatorAgent processes events correctly."""
+    event1 = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    event2 = TestEvent2(source=str, correlation_id="test-id", data="World")
+    
+    # Receive the first event - should not process yet
+    result1 = await test_async_aggregator.receive_event_async(event1)
+    assert result1 == []
+    
+    # Receive the second event - should process both events
+    result2 = await test_async_aggregator.receive_event_async(event2)
+    
+    assert len(result2) == 1
+    assert isinstance(result2[0], TestResultEvent)
+    assert result2[0].result == "Hello - World"
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_receive_event_async_wrong_order(test_async_aggregator):
+    """Test that the AsyncAggregatorAgent processes events correctly regardless of order."""
+    event1 = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    event2 = TestEvent2(source=str, correlation_id="test-id", data="World")
+    
+    # Receive the second event first - should not process yet
+    result1 = await test_async_aggregator.receive_event_async(event2)
+    assert result1 == []
+    
+    # Receive the first event - should process both events
+    result2 = await test_async_aggregator.receive_event_async(event1)
+    
+    assert len(result2) == 1
+    assert isinstance(result2[0], TestResultEvent)
+    assert result2[0].result == "Hello - World"
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_receive_event_async_different_correlation_ids(test_async_aggregator):
+    """Test that the AsyncAggregatorAgent handles different correlation_ids correctly."""
+    event1_id1 = TestEvent1(source=str, correlation_id="id1", message="Hello")
+    event2_id1 = TestEvent2(source=str, correlation_id="id1", data="World")
+    
+    event1_id2 = TestEvent1(source=str, correlation_id="id2", message="Goodbye")
+    event2_id2 = TestEvent2(source=str, correlation_id="id2", data="Universe")
+    
+    # Receive events for id1
+    await test_async_aggregator.receive_event_async(event1_id1)
+    result1 = await test_async_aggregator.receive_event_async(event2_id1)
+    
+    # Receive events for id2
+    await test_async_aggregator.receive_event_async(event1_id2)
+    result2 = await test_async_aggregator.receive_event_async(event2_id2)
+    
+    # Check results for id1
+    assert len(result1) == 1
+    assert result1[0].result == "Hello - World"
+    
+    # Check results for id2
+    assert len(result2) == 1
+    assert result2[0].result == "Goodbye - Universe"
+
+
+@pytest.mark.asyncio
+async def test_async_aggregator_process_events_base_implementation(async_aggregator):
+    """Test that the base process_events implementation returns an empty list."""
+    event1 = TestEvent1(source=str, correlation_id="test-id", message="Hello")
+    event2 = TestEvent2(source=str, correlation_id="test-id", data="World")
+    
+    events = [event1, event2]
+    result = await async_aggregator.process_events(events)
+    
+    assert result == []

mojentic/agents/async_llm_agent.py

@@ -0,0 +1,197 @@
+import json
+from typing import Annotated, List, Optional, Type
+
+from pydantic import BaseModel, Field
+
+from mojentic.agents.base_async_agent import BaseAsyncAgent
+from mojentic.context.shared_working_memory import SharedWorkingMemory
+from mojentic.event import Event
+from mojentic.llm.gateways.models import LLMMessage, MessageRole
+from mojentic.llm.llm_broker import LLMBroker
+from mojentic.llm.tools.llm_tool import LLMTool
+
+
+class BaseAsyncLLMAgent(BaseAsyncAgent):
+    """
+    BaseAsyncLLMAgent is an asynchronous version of the BaseLLMAgent.
+    It uses an LLM to generate responses asynchronously.
+    """
+    llm: LLMBroker
+    behaviour: Annotated[str, "The personality and behavioural traits of the agent."]
+
+    def __init__(self, llm: LLMBroker, behaviour: str = "You are a helpful assistant.",
+                 tools: Optional[List[LLMTool]] = None, response_model: Optional[Type[BaseModel]] = None):
+        """
+        Initialize the BaseAsyncLLMAgent.
+
+        Parameters
+        ----------
+        llm : LLMBroker
+            The LLM broker to use for generating responses
+        behaviour : str, optional
+            The personality and behavioural traits of the agent
+        tools : List[LLMTool], optional
+            The tools available to the agent
+        response_model : Type[BaseModel], optional
+            The model to use for responses
+        """
+        super().__init__()
+        self.llm = llm
+        self.behaviour = behaviour
+        self.response_model = response_model
+        self.tools = tools or []
+
+    def _create_initial_messages(self):
+        """
+        Create the initial messages for the LLM.
+
+        Returns
+        -------
+        list
+            The initial messages for the LLM
+        """
+        return [
+            LLMMessage(role=MessageRole.System, content=self.behaviour),
+        ]
+
+    def add_tool(self, tool):
+        """
+        Add a tool to the agent.
+
+        Parameters
+        ----------
+        tool : LLMTool
+            The tool to add
+        """
+        self.tools.append(tool)
+
+    async def generate_response(self, content):
+        """
+        Generate a response using the LLM asynchronously.
+
+        Parameters
+        ----------
+        content : str
+            The content to generate a response for
+
+        Returns
+        -------
+        str or BaseModel
+            The generated response
+        """
+        messages = self._create_initial_messages()
+        messages.append(LLMMessage(content=content))
+
+        if self.response_model is not None:
+            # Use asyncio.to_thread to run the synchronous generate_object method in a separate thread
+            import asyncio
+            response = await asyncio.to_thread(self.llm.generate_object, messages, object_model=self.response_model)
+        else:
+            # Use asyncio.to_thread to run the synchronous generate method in a separate thread
+            import asyncio
+            response = await asyncio.to_thread(self.llm.generate, messages, tools=self.tools)
+
+        return response
+
+    async def receive_event_async(self, event: Event) -> List[Event]:
+        """
+        Receive an event and process it asynchronously.
+        This method should be overridden by subclasses.
+
+        Parameters
+        ----------
+        event : Event
+            The event to process
+
+        Returns
+        -------
+        List[Event]
+            The events to be processed next
+        """
+        return []
+
+
+class BaseAsyncLLMAgentWithMemory(BaseAsyncLLMAgent):
+    """
+    BaseAsyncLLMAgentWithMemory is an asynchronous version of the BaseLLMAgentWithMemory.
+    It uses an LLM to generate responses asynchronously and maintains a shared working memory.
+    """
+    instructions: Annotated[str, "The instructions for the agent to follow when receiving events."]
+
+    def __init__(self, llm: LLMBroker, memory: SharedWorkingMemory, behaviour: str, instructions: str,
+                 response_model: BaseModel):
+        """
+        Initialize the BaseAsyncLLMAgentWithMemory.
+
+        Parameters
+        ----------
+        llm : LLMBroker
+            The LLM broker to use for generating responses
+        memory : SharedWorkingMemory
+            The shared working memory to use
+        behaviour : str
+            The personality and behavioural traits of the agent
+        instructions : str
+            The instructions for the agent to follow when receiving events
+        response_model : BaseModel
+            The model to use for responses
+        """
+        super().__init__(llm, behaviour)
+        self.instructions = instructions
+        self.memory = memory
+        self.response_model = response_model
+
+    def _create_initial_messages(self):
+        """
+        Create the initial messages for the LLM.
+
+        Returns
+        -------
+        list
+            The initial messages for the LLM
+        """
+        messages = super()._create_initial_messages()
+        messages.extend([
+            LLMMessage(content=f"This is what you remember:\n{json.dumps(self.memory.get_working_memory(), indent=2)}"
+                               f"\n\nRemember anything new you learn by storing it to your working memory in your response."),
+            LLMMessage(role=MessageRole.User, content=self.instructions),
+        ])
+        return messages
+
+    async def generate_response(self, content):
+        """
+        Generate a response using the LLM asynchronously.
+
+        Parameters
+        ----------
+        content : str
+            The content to generate a response for
+
+        Returns
+        -------
+        BaseModel
+            The generated response
+        """
+        class ResponseWithMemory(self.response_model):
+            memory: dict = Field(self.memory.get_working_memory(),
+                                 description="Add anything new that you have learned here.")
+
+        messages = self._create_initial_messages()
+        messages.extend([
+            LLMMessage(content=content),
+        ])
+        
+        # Use asyncio.to_thread to run the synchronous generate_object method in a separate thread
+        import asyncio
+        response = await asyncio.to_thread(
+            self.llm.generate_object,
+            messages=messages,
+            object_model=ResponseWithMemory
+        )
+        
+        self.memory.merge_to_working_memory(response.memory)
+
+        d = response.model_dump()
+        del d["memory"]
+
+        return self.response_model.model_validate(d)

mojentic/agents/async_llm_agent_spec.py

@@ -0,0 +1,166 @@
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from pydantic import BaseModel, Field
+
+from mojentic.agents.async_llm_agent import BaseAsyncLLMAgent
+from mojentic.event import Event
+from mojentic.llm.llm_broker import LLMBroker
+from mojentic.llm.gateways.models import LLMMessage, MessageRole
+
+
+class TestEvent(Event):
+    """A simple event for testing."""
+    message: str
+
+
+class TestResponse(BaseModel):
+    """A simple response model for testing."""
+    answer: str = Field(..., description="The answer to the question")
+
+
+@pytest.fixture
+def mock_llm_broker():
+    """Create a mock LLM broker for testing."""
+    mock_broker = MagicMock(spec=LLMBroker)
+    mock_broker.generate.return_value = "Test response"
+    mock_broker.generate_object.return_value = TestResponse(answer="Test answer")
+    return mock_broker
+
+
+@pytest.fixture
+def async_llm_agent(mock_llm_broker):
+    """Create a BaseAsyncLLMAgent for testing."""
+    return BaseAsyncLLMAgent(
+        llm=mock_llm_broker,
+        behaviour="You are a test assistant.",
+        response_model=TestResponse
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_llm_agent_init(mock_llm_broker):
+    """Test that the BaseAsyncLLMAgent initializes correctly."""
+    agent = BaseAsyncLLMAgent(
+        llm=mock_llm_broker,
+        behaviour="You are a test assistant.",
+        response_model=TestResponse
+    )
+    
+    assert agent.llm == mock_llm_broker
+    assert agent.behaviour == "You are a test assistant."
+    assert agent.response_model == TestResponse
+    assert agent.tools == []
+
+
+@pytest.mark.asyncio
+async def test_async_llm_agent_create_initial_messages(async_llm_agent):
+    """Test that the BaseAsyncLLMAgent creates initial messages correctly."""
+    messages = async_llm_agent._create_initial_messages()
+    
+    assert len(messages) == 1
+    assert messages[0].role == MessageRole.System
+    assert messages[0].content == "You are a test assistant."
+
+
+@pytest.mark.asyncio
+async def test_async_llm_agent_add_tool(async_llm_agent):
+    """Test that the BaseAsyncLLMAgent can add tools."""
+    mock_tool = MagicMock()
+    async_llm_agent.add_tool(mock_tool)
+    
+    assert mock_tool in async_llm_agent.tools
+
+
+@pytest.mark.asyncio
+async def test_async_llm_agent_generate_response_with_model(async_llm_agent, mock_llm_broker):
+    """Test that the BaseAsyncLLMAgent generates responses with a model."""
+    response = await async_llm_agent.generate_response("Test question")
+    
+    # Verify that generate_object was called
+    mock_llm_broker.generate_object.assert_called_once()
+    
+    # Verify the response
+    assert isinstance(response, TestResponse)
+    assert response.answer == "Test answer"
+
+
+@pytest.mark.asyncio
+async def test_async_llm_agent_generate_response_without_model(mock_llm_broker):
+    """Test that the BaseAsyncLLMAgent generates responses without a model."""
+    agent = BaseAsyncLLMAgent(
+        llm=mock_llm_broker,
+        behaviour="You are a test assistant."
+    )
+    
+    response = await agent.generate_response("Test question")
+    
+    # Verify that generate was called
+    mock_llm_broker.generate.assert_called_once()
+    
+    # Verify the response
+    assert response == "Test response"
+
+
+@pytest.mark.asyncio
+async def test_async_llm_agent_generate_response_with_tools(mock_llm_broker):
+    """Test that the BaseAsyncLLMAgent generates responses with tools."""
+    mock_tool = MagicMock()
+    
+    agent = BaseAsyncLLMAgent(
+        llm=mock_llm_broker,
+        behaviour="You are a test assistant.",
+        tools=[mock_tool]
+    )
+    
+    response = await agent.generate_response("Test question")
+    
+    # Verify that generate was called with tools
+    mock_llm_broker.generate.assert_called_once()
+    args, kwargs = mock_llm_broker.generate.call_args
+    assert kwargs.get('tools') == [mock_tool]
+
+
+@pytest.mark.asyncio
+async def test_async_llm_agent_receive_event_async(async_llm_agent):
+    """Test that the BaseAsyncLLMAgent's receive_event_async method works."""
+    event = TestEvent(source=str, message="Test message")
+    
+    # The base implementation should return an empty list
+    result = await async_llm_agent.receive_event_async(event)
+    
+    assert result == []
+
+
+# Create a subclass for testing the receive_event_async method
+class TestAsyncLLMAgent(BaseAsyncLLMAgent):
+    """A test async LLM agent that implements receive_event_async."""
+    
+    async def receive_event_async(self, event):
+        if isinstance(event, TestEvent):
+            response = await self.generate_response(event.message)
+            return [TestEvent(
+                source=type(self),
+                correlation_id=event.correlation_id,
+                message=f"Response: {response.answer}"
+            )]
+        return []
+
+
+@pytest.mark.asyncio
+async def test_subclass_async_llm_agent_receive_event_async(mock_llm_broker):
+    """Test that a subclass of BaseAsyncLLMAgent can implement receive_event_async."""
+    agent = TestAsyncLLMAgent(
+        llm=mock_llm_broker,
+        behaviour="You are a test assistant.",
+        response_model=TestResponse
+    )
+    
+    event = TestEvent(source=str, message="Test message")
+    
+    result = await agent.receive_event_async(event)
+    
+    assert len(result) == 1
+    assert isinstance(result[0], TestEvent)
+    assert result[0].message == "Response: Test answer"

mojentic/agents/base_async_agent.py

@@ -0,0 +1,27 @@
+from typing import List
+
+from mojentic.event import Event
+
+
+class BaseAsyncAgent:
+    """
+    BaseAsyncAgent class is the base class for all asynchronous agents.
+    It provides an async receive method for event processing.
+    """
+
+    async def receive_event_async(self, event: Event) -> List[Event]:
+        """
+        receive_event_async is the method that all async agents must implement. It takes an event as input and returns a list of
+        events as output.
+
+        In this way, you can perform work based on the event, and generate whatever subsequent events may need to be
+        processed next.
+
+        This keeps the agent decoupled from the specifics of the event routing and processing.
+
+        Events are subclasses of the Event class.
+
+        :param event: The event to process
+        :return: A list of events to be processed next
+        """
+        return []