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

_examples/async_dispatcher_example.py

@@ -0,0 +1,241 @@
+"""
+Example script demonstrating how to use the AsyncDispatcher, BaseAsyncAgent, and AsyncAggregatorAgent.
+
+This script shows how to create and use asynchronous agents with the AsyncDispatcher.
+"""
+
+import asyncio
+from pathlib import Path
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from mojentic.agents.async_aggregator_agent import AsyncAggregatorAgent
+from mojentic.agents.async_llm_agent import BaseAsyncLLMAgent
+from mojentic.agents.base_async_agent import BaseAsyncAgent
+from mojentic.async_dispatcher import AsyncDispatcher
+from mojentic.event import Event, TerminateEvent
+from mojentic.llm import LLMBroker
+from mojentic.llm.gateways.models import LLMMessage, MessageRole
+from mojentic.router import Router
+
+
+# Define some example events
+class TextEvent(Event):
+    text: str = Field(..., description="The text content of the event")
+
+
+class AnalysisEvent(Event):
+    analysis: str = Field(..., description="The analysis of the text")
+
+
+class SummaryEvent(Event):
+    summary: str = Field(..., description="The summary of the text")
+
+
+class CombinedResultEvent(Event):
+    text: str = Field(..., description="The original text")
+    analysis: str = Field(..., description="The analysis of the text")
+    summary: str = Field(..., description="The summary of the text")
+    combined: str = Field(..., description="The combined result")
+
+
+# Define response models for LLM agents
+class AnalysisResponse(BaseModel):
+    analysis: str = Field(..., description="The analysis of the text")
+
+
+class SummaryResponse(BaseModel):
+    summary: str = Field(..., description="The summary of the text")
+
+
+class CombinedResponse(BaseModel):
+    combined: str = Field(..., description="The combined result of the analysis and summary")
+
+
+# Define some example agents
+class TextAnalyzerAgent(BaseAsyncLLMAgent):
+    """
+    An agent that analyzes text and produces an analysis.
+    """
+
+    def __init__(self, llm: LLMBroker):
+        super().__init__(
+            llm=llm,
+            behaviour="You are a text analysis assistant. Your job is to provide a detailed analysis of the given text, including key themes, structure, and notable elements.",
+            response_model=AnalysisResponse
+        )
+
+    async def receive_event_async(self, event: Event) -> List[Event]:
+        if isinstance(event, TextEvent):
+            prompt = f"""
+Please analyze the following text in detail. Consider:
+- Main themes and topics
+- Structure and organization
+- Key points and arguments
+- Style and tone
+- Intended audience
+- Any notable or unique elements
+
+Text to analyze:
+{event.text[:1000]}... (text truncated for brevity)
+"""
+            response = await self.generate_response(prompt)
+            return [AnalysisEvent(source=type(self), correlation_id=event.correlation_id, analysis=response.analysis)]
+        return []
+
+
+class TextSummarizerAgent(BaseAsyncLLMAgent):
+    """
+    An agent that summarizes text.
+    """
+
+    def __init__(self, llm: LLMBroker):
+        super().__init__(
+            llm=llm,
+            behaviour="You are a text summarization assistant. Your job is to provide concise, accurate summaries of texts while preserving the key information and main points.",
+            response_model=SummaryResponse
+        )
+
+    async def receive_event_async(self, event: Event) -> List[Event]:
+        if isinstance(event, TextEvent):
+            prompt = f"""
+Please provide a concise summary of the following text. The summary should:
+- Capture the main points and key information
+- Be significantly shorter than the original text
+- Maintain the original meaning and intent
+- Be clear and coherent
+- Exclude unnecessary details
+
+Text to summarize:
+{event.text[:1000]}... (text truncated for brevity)
+"""
+            response = await self.generate_response(prompt)
+            return [SummaryEvent(source=type(self), correlation_id=event.correlation_id, summary=response.summary)]
+        return []
+
+
+class ResultCombinerAgent(AsyncAggregatorAgent):
+    """
+    An agent that combines the analysis and summary of a text using LLM.
+    """
+
+    def __init__(self, llm: LLMBroker):
+        super().__init__(event_types_needed=[TextEvent, AnalysisEvent, SummaryEvent])
+        self.llm = llm
+        self.behaviour = "You are an assistant that combines text analysis and summaries into a comprehensive report."
+        self.response_model = CombinedResponse
+
+    async def process_events(self, events):
+        # Extract the events
+        text_event = next((e for e in events if isinstance(e, TextEvent)), None)
+        analysis_event = next((e for e in events if isinstance(e, AnalysisEvent)), None)
+        summary_event = next((e for e in events if isinstance(e, SummaryEvent)), None)
+
+        if text_event and analysis_event and summary_event:
+            # Use LLM to create a sophisticated combination
+            prompt = f"""
+I have analyzed and summarized a text. Please combine these into a comprehensive report.
+
+Original Text (excerpt): {text_event.text[:300]}... (text truncated for brevity)
+
+Analysis: {analysis_event.analysis}
+
+Summary: {summary_event.summary}
+
+Please create a well-structured, insightful report that integrates the analysis and summary, 
+highlighting the most important aspects of the text. The report should provide a comprehensive 
+understanding of the text's content, structure, and significance.
+"""
+            # Create a temporary LLM agent to generate the response
+            messages = [
+                LLMMessage(role=MessageRole.System, content=self.behaviour),
+                LLMMessage(role=MessageRole.User, content=prompt)
+            ]
+
+            # Generate the response
+            import asyncio
+            response_json = await asyncio.to_thread(
+                self.llm.generate_object,
+                messages=messages,
+                object_model=self.response_model
+            )
+
+            combined = response_json.combined
+
+            return [CombinedResultEvent(
+                source=type(self),
+                correlation_id=text_event.correlation_id,
+                text=text_event.text,
+                analysis=analysis_event.analysis,
+                summary=summary_event.summary,
+                combined=combined
+            )]
+        return []
+
+
+class ResultOutputAgent(BaseAsyncAgent):
+    """
+    An agent that receives the CombinedResultEvent, outputs the result to the user,
+    and emits a TerminateEvent to exit the event loop.
+    """
+
+    async def receive_event_async(self, event: Event) -> List[Event]:
+        if isinstance(event, CombinedResultEvent):
+            # Output the result to the user
+            print("\n=== FINAL ANSWER ===")
+            print(event.combined)
+            print("===================\n")
+
+            # Emit a TerminateEvent to exit the event loop
+            return [TerminateEvent(source=type(self), correlation_id=event.correlation_id)]
+        return []
+
+
+async def main():
+    """
+    Main function that demonstrates the usage of AsyncDispatcher and async agents.
+    """
+    # Initialize the LLM broker with the same model as in async_llm_example
+    llm = LLMBroker(model="qwen3:30b-a3b-q4_K_M")
+
+    # Create a router and register agents
+    router = Router()
+
+    # Create agents with LLM
+    analyzer = TextAnalyzerAgent(llm)
+    summarizer = TextSummarizerAgent(llm)
+    combiner = ResultCombinerAgent(llm)
+    output_agent = ResultOutputAgent()
+
+    # Register agents with the router
+    router.add_route(TextEvent, analyzer)
+    router.add_route(TextEvent, summarizer)
+    router.add_route(TextEvent, combiner)
+    router.add_route(AnalysisEvent, combiner)
+    router.add_route(SummaryEvent, combiner)
+    router.add_route(CombinedResultEvent, output_agent)
+
+    # Create and start the dispatcher
+    dispatcher = await AsyncDispatcher(router).start()
+
+    # Create a text event
+    with open(Path.cwd().parent.parent / "README.md", "r") as f:
+        text = f.read()
+    event = TextEvent(source=type("ExampleSource", (), {}), text=text)
+
+    # Dispatch the event
+    dispatcher.dispatch(event)
+
+    # Wait for all events in the queue to be processed
+    print("Waiting for all events to be processed...")
+    queue_empty = await dispatcher.wait_for_empty_queue(timeout=10)
+    if not queue_empty:
+        print("Warning: Not all events were processed within the timeout period.")
+
+    # Stop the dispatcher
+    await dispatcher.stop()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

_examples/async_llm_example.py

@@ -0,0 +1,236 @@
+"""
+Example script demonstrating how to use the AsyncDispatcher with BaseAsyncLLMAgent.
+
+This script shows how to create and use asynchronous LLM agents with the AsyncDispatcher.
+"""
+
+import asyncio
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+from mojentic.agents.async_aggregator_agent import AsyncAggregatorAgent
+from mojentic.agents.async_llm_agent import BaseAsyncLLMAgent
+from mojentic.async_dispatcher import AsyncDispatcher
+from mojentic.context.shared_working_memory import SharedWorkingMemory
+from mojentic.event import Event
+from mojentic.llm import LLMBroker
+from mojentic.router import Router
+
+
+# Define some example events
+class QuestionEvent(Event):
+    question: str = Field(..., description="The question to answer")
+
+
+class FactCheckEvent(Event):
+    question: str = Field(..., description="The original question")
+    facts: List[str] = Field(..., description="The facts related to the question")
+
+
+class AnswerEvent(Event):
+    question: str = Field(..., description="The original question")
+    answer: str = Field(..., description="The answer to the question")
+    confidence: float = Field(..., description="The confidence level of the answer (0-1)")
+
+
+class FinalAnswerEvent(Event):
+    question: str = Field(..., description="The original question")
+    answer: str = Field(..., description="The final answer to the question")
+    facts: List[str] = Field(..., description="The facts used to answer the question")
+    confidence: float = Field(..., description="The confidence level of the answer (0-1)")
+
+
+# Define response models for LLM agents
+class FactCheckResponse(BaseModel):
+    facts: List[str] = Field(..., description="The facts related to the question")
+
+
+class AnswerResponse(BaseModel):
+    answer: str = Field(..., description="The answer to the question")
+    confidence: float = Field(..., description="The confidence level of the answer (0-1)")
+
+
+# Define some example agents
+class FactCheckerAgent(BaseAsyncLLMAgent):
+    """
+    An agent that checks facts related to a question.
+    """
+    def __init__(self, llm: LLMBroker):
+        super().__init__(
+            llm=llm,
+            behaviour="You are a fact-checking assistant. Your job is to provide relevant facts about a question.",
+            response_model=FactCheckResponse
+        )
+
+    async def receive_event_async(self, event: Event) -> List[Event]:
+        if isinstance(event, QuestionEvent):
+            prompt = f"Please provide relevant facts about the following question: {event.question}"
+            response = await self.generate_response(prompt)
+            return [FactCheckEvent(
+                source=type(self),
+                correlation_id=event.correlation_id,
+                question=event.question,
+                facts=response.facts
+            )]
+        return []
+
+
+class AnswerGeneratorAgent(BaseAsyncLLMAgent):
+    """
+    An agent that generates an answer to a question.
+    """
+    def __init__(self, llm: LLMBroker):
+        super().__init__(
+            llm=llm,
+            behaviour="You are a question-answering assistant. Your job is to provide accurate answers to questions.",
+            response_model=AnswerResponse
+        )
+
+    async def receive_event_async(self, event: Event) -> List[Event]:
+        if isinstance(event, QuestionEvent):
+            prompt = f"Please answer the following question: {event.question}"
+            response = await self.generate_response(prompt)
+            return [AnswerEvent(
+                source=type(self),
+                correlation_id=event.correlation_id,
+                question=event.question,
+                answer=response.answer,
+                confidence=response.confidence
+            )]
+        return []
+
+
+class FinalAnswerAgent(AsyncAggregatorAgent):
+    """
+    An agent that combines facts and answers to produce a final answer.
+    """
+    def __init__(self, llm: LLMBroker):
+        super().__init__(event_types_needed=[FactCheckEvent, AnswerEvent])
+        self.llm = llm
+        self.final_answer_event = None
+
+    async def receive_event_async(self, event: Event) -> list:
+        print(f"FinalAnswerAgent received event: {type(event).__name__}")
+        result = await super().receive_event_async(event)
+        # Store any FinalAnswerEvent created
+        for e in result:
+            if isinstance(e, FinalAnswerEvent):
+                self.final_answer_event = e
+        return result
+
+    async def process_events(self, events):
+        print(f"FinalAnswerAgent processing events: {[type(e).__name__ for e in events]}")
+        # Extract the events
+        fact_check_event = next((e for e in events if isinstance(e, FactCheckEvent)), None)
+        answer_event = next((e for e in events if isinstance(e, AnswerEvent)), None)
+
+        if fact_check_event and answer_event:
+            print("FinalAnswerAgent has both FactCheckEvent and AnswerEvent")
+            # In a real implementation, we might use the LLM to refine the answer based on the facts
+            # For this example, we'll just combine them
+
+            # Adjust confidence based on facts
+            confidence = answer_event.confidence
+            if len(fact_check_event.facts) > 0:
+                # Increase confidence if we have facts
+                confidence = min(1.0, confidence + 0.1)
+
+            final_answer_event = FinalAnswerEvent(
+                source=type(self),
+                correlation_id=fact_check_event.correlation_id,
+                question=fact_check_event.question,
+                answer=answer_event.answer,
+                facts=fact_check_event.facts,
+                confidence=confidence
+            )
+            print(f"FinalAnswerAgent created FinalAnswerEvent: {final_answer_event}")
+            self.final_answer_event = final_answer_event
+            return [final_answer_event]
+        print("FinalAnswerAgent missing either FactCheckEvent or AnswerEvent")
+        return []
+
+    async def get_final_answer(self, correlation_id, timeout=30):
+        """
+        Get the final answer for a specific correlation_id.
+
+        Parameters
+        ----------
+        correlation_id : str
+            The correlation_id to get the final answer for
+        timeout : float, optional
+            The timeout in seconds
+
+        Returns
+        -------
+        FinalAnswerEvent or None
+            The final answer event, or None if not found
+        """
+        # First wait for all needed events
+        await self.wait_for_events(correlation_id, timeout)
+
+        # Then check if we have a final answer
+        if self.final_answer_event and self.final_answer_event.correlation_id == correlation_id:
+            return self.final_answer_event
+
+        return None
+
+
+async def main():
+    """
+    Main function that demonstrates the usage of AsyncDispatcher with async LLM agents.
+    """
+    # Initialize the LLM broker with your preferred model
+    llm = LLMBroker(model="qwen3:30b-a3b-q4_K_M")
+
+    # Create a router and register agents
+    router = Router()
+
+    # Create agents
+    fact_checker = FactCheckerAgent(llm)
+    answer_generator = AnswerGeneratorAgent(llm)
+    final_answer_agent = FinalAnswerAgent(llm)
+
+    # Register agents with the router
+    router.add_route(QuestionEvent, fact_checker)
+    router.add_route(QuestionEvent, answer_generator)
+    router.add_route(QuestionEvent, final_answer_agent)
+    router.add_route(FactCheckEvent, final_answer_agent)
+    router.add_route(AnswerEvent, final_answer_agent)
+
+    # Create and start the dispatcher
+    dispatcher = await AsyncDispatcher(router).start()
+
+    # Create a question event
+    question = "What is the capital of France?"
+    event = QuestionEvent(source=type("ExampleSource", (), {}), question=question)
+
+    # Dispatch the event
+    print("Dispatching question event")
+    dispatcher.dispatch(event)
+
+    # Give the dispatcher a moment to start processing the event
+    print("Waiting for dispatcher to start processing")
+    await asyncio.sleep(0.1)
+
+    # Wait for the final answer from the FinalAnswerAgent
+    print("Waiting for final answer from FinalAnswerAgent")
+    final_answer_event = await final_answer_agent.get_final_answer(event.correlation_id, timeout=30)
+
+    # Print the final answer
+    if final_answer_event:
+        print(f"Question: {final_answer_event.question}")
+        print(f"Answer: {final_answer_event.answer}")
+        print(f"Confidence: {final_answer_event.confidence}")
+        print("Facts:")
+        for fact in final_answer_event.facts:
+            print(f"  - {fact}")
+    else:
+        print("No FinalAnswerEvent found")
+
+    # Stop the dispatcher
+    await dispatcher.stop()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

mojentic/agents/async_aggregator_agent.py

@@ -0,0 +1,162 @@
+import asyncio
+import structlog
+
+from mojentic.agents.base_async_agent import BaseAsyncAgent
+from mojentic.event import Event
+
+logger = structlog.get_logger()
+
+
+class AsyncAggregatorAgent(BaseAsyncAgent):
+    """
+    AsyncAggregatorAgent is an asynchronous version of the BaseAggregatingAgent.
+    It aggregates events based on their correlation_id and processes them when all required events are available.
+    """
+    def __init__(self, event_types_needed=None):
+        """
+        Initialize the AsyncAggregatorAgent.
+
+        Parameters
+        ----------
+        event_types_needed : list, optional
+            List of event types that need to be captured before processing
+        """
+        super().__init__()
+        self.results = {}
+        self.event_types_needed = event_types_needed or []
+        self.futures = {}  # Maps correlation_id to Future objects
+
+    async def _get_and_reset_results(self, event):
+        """
+        Get and reset the results for a specific correlation_id.
+
+        Parameters
+        ----------
+        event : Event
+            The event to get results for
+
+        Returns
+        -------
+        list
+            The results for the event
+        """
+        results = self.results[event.correlation_id]
+        self.results[event.correlation_id] = None
+        return results
+
+    async def _capture_results_if_needed(self, event):
+        """
+        Capture results for a specific correlation_id.
+
+        Parameters
+        ----------
+        event : Event
+            The event to capture results for
+        """
+        results = self.results.get(event.correlation_id, [])
+        results.append(event)
+        self.results[event.correlation_id] = results
+
+        # Check if we have all needed events and set the future if we do
+        event_types_captured = [type(e) for e in self.results.get(event.correlation_id, [])]
+        finished = all([event_type in event_types_captured for event_type in self.event_types_needed])
+
+        if finished and event.correlation_id in self.futures:
+            future = self.futures[event.correlation_id]
+            if not future.done():
+                future.set_result(self.results[event.correlation_id])
+
+    async def _has_all_needed(self, event):
+        """
+        Check if all needed event types have been captured for a specific correlation_id.
+
+        Parameters
+        ----------
+        event : Event
+            The event to check
+
+        Returns
+        -------
+        bool
+            True if all needed event types have been captured, False otherwise
+        """
+        event_types_captured = [type(e) for e in self.results.get(event.correlation_id, [])]
+        finished = all([event_type in event_types_captured for event_type in self.event_types_needed])
+        logger.debug(f"Captured: {event_types_captured}, Needed: {self.event_types_needed}, Finished: {finished}")
+        return finished
+
+    async def wait_for_events(self, correlation_id, timeout=None):
+        """
+        Wait for all needed events for a specific correlation_id.
+
+        Parameters
+        ----------
+        correlation_id : str
+            The correlation_id to wait for
+        timeout : float, optional
+            The timeout in seconds
+
+        Returns
+        -------
+        list
+            The events for the correlation_id
+        """
+        if correlation_id not in self.futures:
+            self.futures[correlation_id] = asyncio.Future()
+
+        # If we already have all needed events, return them
+        if correlation_id in self.results:
+            event_types_captured = [type(e) for e in self.results.get(correlation_id, [])]
+            if all([event_type in event_types_captured for event_type in self.event_types_needed]):
+                return self.results[correlation_id]
+
+        # Otherwise, wait for the future to be set
+        try:
+            return await asyncio.wait_for(self.futures[correlation_id], timeout)
+        except asyncio.TimeoutError:
+            logger.warning(f"Timeout waiting for events for correlation_id {correlation_id}")
+            return self.results.get(correlation_id, [])
+
+    async def receive_event_async(self, event: Event) -> list:
+        """
+        Receive an event and process it if all needed events are available.
+
+        Parameters
+        ----------
+        event : Event
+            The event to process
+
+        Returns
+        -------
+        list
+            The events to be processed next
+        """
+        # First capture the event
+        await self._capture_results_if_needed(event)
+
+        # Then check if we have all needed events
+        event_types_captured = [type(e) for e in self.results.get(event.correlation_id, [])]
+        finished = all([event_type in event_types_captured for event_type in self.event_types_needed])
+
+        # If we have all needed events, process them
+        if finished:
+            return await self.process_events(await self._get_and_reset_results(event))
+
+        return []
+
+    async def process_events(self, events):
+        """
+        Process a list of events.
+        This method should be overridden by subclasses.
+
+        Parameters
+        ----------
+        events : list
+            The events to process
+
+        Returns
+        -------
+        list
+            The events to be processed next
+        """
+        return []