mojentic 0.8.4__py3-none-any.whl → 1.0.0__py3-none-any.whl

mojentic/llm/gateways/openai_model_registry_spec.py

@@ -2,7 +2,6 @@
 Tests for the OpenAI Model Registry system.
 """
 
-import pytest
 from mojentic.llm.gateways.openai_model_registry import (
     OpenAIModelRegistry,
     ModelType,
@@ -178,4 +177,4 @@ class DescribeOpenAIModelRegistry:
         assert ModelType.REASONING.value == "reasoning"
         assert ModelType.CHAT.value == "chat"
         assert ModelType.EMBEDDING.value == "embedding"
-        assert ModelType.MODERATION.value == "moderation"
+        assert ModelType.MODERATION.value == "moderation"

mojentic/llm/gateways/openai_temperature_handling_spec.py

@@ -3,7 +3,7 @@ from unittest.mock import MagicMock
 
 from mojentic.llm.gateways.openai import OpenAIGateway
 from mojentic.llm.gateways.openai_model_registry import get_model_registry
-from mojentic.llm.gateways.models import LLMMessage, MessageRole, LLMGatewayResponse
+from mojentic.llm.gateways.models import LLMMessage, MessageRole
 
 
 @pytest.fixture
@@ -242,4 +242,4 @@ class DescribeModelCapabilitiesTemperatureRestrictions:
             capabilities = registry.get_model_capabilities(model)
             assert capabilities.supports_temperature(1.0) is True
             assert capabilities.supports_temperature(0.1) is False
-            assert capabilities.supported_temperatures == [1.0]
+            assert capabilities.supported_temperatures == [1.0]

mojentic/llm/llm_broker.py

@@ -1,12 +1,12 @@
 import json
 import time
-from typing import List, Optional, Type
+from typing import List, Optional, Type, Iterator
 
 import structlog
 from pydantic import BaseModel
 
 from mojentic.llm.gateways.llm_gateway import LLMGateway
-from mojentic.llm.gateways.models import MessageRole, LLMMessage, LLMGatewayResponse
+from mojentic.llm.gateways.models import MessageRole, LLMMessage, LLMGatewayResponse, LLMToolCall
 from mojentic.llm.gateways.ollama import OllamaGateway
 from mojentic.llm.gateways.tokenizer_gateway import TokenizerGateway
 from mojentic.tracer.tracer_system import TracerSystem
@@ -182,6 +182,166 @@ class LLMBroker():
 
         return result.content
 
+    def generate_stream(self, messages: List[LLMMessage], tools=None, temperature=1.0, num_ctx=32768,
+                        num_predict=-1, max_tokens=16384,
+                        correlation_id: str = None) -> Iterator[str]:
+        """
+        Generate a streaming text response from the LLM.
+
+        This method mirrors generate() but yields content chunks as they arrive from the LLM,
+        providing a better user experience for long-running requests. When tool calls are
+        detected, tools are executed and the LLM is called recursively, with the new response
+        also being streamed.
+
+        Parameters
+        ----------
+        messages : List[LLMMessage]
+            A list of messages to send to the LLM.
+        tools : List[Tool]
+            A list of tools to use with the LLM. If a tool call is requested, the tool will be
+            called and the output will be included in the response.
+        temperature : float
+            The temperature to use for the response. Defaults to 1.0
+        num_ctx : int
+            The number of context tokens to use. Defaults to 32768.
+        num_predict : int
+            The number of tokens to predict. Defaults to no limit.
+        max_tokens : int
+            The maximum number of tokens to generate. Defaults to 16384.
+        correlation_id : str
+            UUID string that is copied from cause-to-affect for tracing events.
+
+        Yields
+        ------
+        str
+            Content chunks as they arrive from the LLM.
+        """
+        # Check if gateway supports streaming
+        if not hasattr(self.adapter, 'complete_stream'):
+            raise NotImplementedError(f"Gateway {type(self.adapter).__name__} does not support streaming")
+
+        approximate_tokens = len(self.tokenizer.encode(self._content_to_count(messages)))
+        logger.info(f"Requesting streaming llm response with approx {approximate_tokens} tokens")
+
+        # Convert messages to serializable dict for audit
+        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
+        self.tracer.record_llm_call(
+            self.model,
+            messages_for_tracer,
+            temperature,
+            tools=tools_for_tracer,
+            source=type(self),
+            correlation_id=correlation_id
+        )
+
+        # Measure call duration for audit
+        start_time = time.time()
+
+        # Accumulate content and tool calls from stream
+        accumulated_content = ""
+        accumulated_tool_calls = []
+
+        stream = self.adapter.complete_stream(
+            model=self.model,
+            messages=messages,
+            tools=tools,
+            temperature=temperature,
+            num_ctx=num_ctx,
+            num_predict=num_predict,
+            max_tokens=max_tokens
+        )
+
+        for chunk in stream:
+            # Handle content chunks
+            if hasattr(chunk, 'content') and chunk.content:
+                accumulated_content += chunk.content
+                yield chunk.content
+
+            # Handle tool calls if present
+            if hasattr(chunk, 'tool_calls') and chunk.tool_calls:
+                accumulated_tool_calls.extend(chunk.tool_calls)
+
+        call_duration_ms = (time.time() - start_time) * 1000
+
+        # Record LLM response in tracer
+        tool_calls_for_tracer = [tc.model_dump() if hasattr(tc, 'model_dump') else tc for tc in
+                                 accumulated_tool_calls] if accumulated_tool_calls else None
+        self.tracer.record_llm_response(
+            self.model,
+            accumulated_content,
+            tool_calls=tool_calls_for_tracer,
+            call_duration_ms=call_duration_ms,
+            source=type(self),
+            correlation_id=correlation_id
+        )
+
+        # Process tool calls if any were accumulated
+        if accumulated_tool_calls and tools is not None:
+            logger.info("Tool call requested in streaming response")
+            for tool_call in accumulated_tool_calls:
+                # Handle both LLMToolCall objects and raw tool call data
+                if hasattr(tool_call, 'name'):
+                    tool_name = tool_call.name
+                    tool_arguments = tool_call.arguments
+                else:
+                    # Handle ollama's tool call format
+                    tool_name = tool_call.function.name
+                    tool_arguments = tool_call.function.arguments
+
+                if tool := next((t for t in tools if t.matches(tool_name)), None):
+                    logger.info('Calling function', function=tool_name)
+                    logger.info('Arguments:', arguments=tool_arguments)
+
+                    # Measure tool execution time
+                    tool_start_time = time.time()
+
+                    # Call the tool
+                    output = tool.run(**tool_arguments)
+
+                    tool_duration_ms = (time.time() - tool_start_time) * 1000
+
+                    # Record tool call in tracer
+                    self.tracer.record_tool_call(
+                        tool_name,
+                        tool_arguments,
+                        output,
+                        caller="LLMBroker.generate_stream",
+                        call_duration_ms=tool_duration_ms,
+                        source=type(self),
+                        correlation_id=correlation_id
+                    )
+
+                    logger.info('Function output', output=output)
+
+                    # Convert to LLMToolCall if needed, preserving the ID if it exists
+                    if not isinstance(tool_call, LLMToolCall):
+                        # Extract ID if available from the tool_call object
+                        tool_call_id = None
+                        if hasattr(tool_call, 'id'):
+                            tool_call_id = tool_call.id
+                        elif hasattr(tool_call, 'function') and hasattr(tool_call.function, 'id'):
+                            tool_call_id = tool_call.function.id
+
+                        tool_call = LLMToolCall(id=tool_call_id, name=tool_name, arguments=tool_arguments)
+
+                    messages.append(LLMMessage(role=MessageRole.Assistant, tool_calls=[tool_call]))
+                    messages.append(
+                        LLMMessage(role=MessageRole.Tool, content=json.dumps(output),
+                                   tool_calls=[tool_call]))
+
+                    # Recursively stream the response after tool execution
+                    yield from self.generate_stream(
+                        messages, tools, temperature, num_ctx, num_predict,
+                        max_tokens, correlation_id=correlation_id
+                    )
+                    return  # Exit after recursive call
+                else:
+                    logger.warn('Function not found', function=tool_name)
+
     def _content_to_count(self, messages: List[LLMMessage]):
         content = ""
         for message in messages:

mojentic/llm/llm_broker_spec.py

@@ -1,4 +1,3 @@
-from unittest.mock import MagicMock
 
 import pytest
 from pydantic import BaseModel
@@ -11,10 +10,12 @@ class SimpleModel(BaseModel):
     text: str
     number: int
 
+
 class NestedModel(BaseModel):
     title: str
     details: SimpleModel
 
+
 class ComplexModel(BaseModel):
     name: str
     items: list[SimpleModel]
@@ -120,7 +121,11 @@ class DescribeLLMBroker:
                 metadata={"key1": "value1", "key2": "value2"}
             )
             mock_gateway.complete.return_value = LLMGatewayResponse(
-                content='{"name": "test", "items": [{"text": "item1", "number": 1}, {"text": "item2", "number": 2}], "metadata": {"key1": "value1", "key2": "value2"}}',
+                content=(
+                    '{"name": "test", "items": [{"text": "item1", "number": 1}, '
+                    '{"text": "item2", "number": 2}], '
+                    '"metadata": {"key1": "value1", "key2": "value2"}}'
+                ),
                 object=mock_object,
                 tool_calls=[]
             )
@@ -135,3 +140,72 @@ class DescribeLLMBroker:
             assert result.items[1].number == 2
             assert result.metadata == {"key1": "value1", "key2": "value2"}
             mock_gateway.complete.assert_called_once()
+
+    class DescribeStreamingGeneration:
+
+        def should_stream_simple_response(self, llm_broker, mock_gateway, mocker):
+            from mojentic.llm.gateways.ollama import StreamingResponse
+
+            messages = [LLMMessage(role=MessageRole.User, content="Tell me a story")]
+
+            # Mock the complete_stream method to yield chunks
+            mock_gateway.complete_stream = mocker.MagicMock()
+            mock_gateway.complete_stream.return_value = iter([
+                StreamingResponse(content="Once "),
+                StreamingResponse(content="upon "),
+                StreamingResponse(content="a "),
+                StreamingResponse(content="time...")
+            ])
+
+            result_chunks = list(llm_broker.generate_stream(messages))
+
+            assert result_chunks == ["Once ", "upon ", "a ", "time..."]
+            mock_gateway.complete_stream.assert_called_once()
+
+        def should_handle_tool_calls_during_streaming(self, llm_broker, mock_gateway, mocker):
+            from mojentic.llm.gateways.ollama import StreamingResponse
+
+            messages = [LLMMessage(role=MessageRole.User, content="What is the date on Friday?")]
+            tool_call = mocker.create_autospec(LLMToolCall, instance=True)
+            tool_call.name = "resolve_date"
+            tool_call.arguments = {"date": "Friday"}
+
+            # First stream has tool call, second stream has the response after tool execution
+            mock_gateway.complete_stream = mocker.MagicMock()
+            mock_gateway.complete_stream.side_effect = [
+                iter([
+                    StreamingResponse(content="Let "),
+                    StreamingResponse(content="me "),
+                    StreamingResponse(content="check..."),
+                    StreamingResponse(tool_calls=[tool_call])
+                ]),
+                iter([
+                    StreamingResponse(content="The "),
+                    StreamingResponse(content="date "),
+                    StreamingResponse(content="is "),
+                    StreamingResponse(content="2024-11-15")
+                ])
+            ]
+
+            mock_tool = mocker.MagicMock()
+            mock_tool.matches.return_value = True
+            mock_tool.run.return_value = {"resolved_date": "2024-11-15"}
+
+            result_chunks = list(llm_broker.generate_stream(messages, tools=[mock_tool]))
+
+            # Should get chunks from first response, then chunks from second response after tool execution
+            assert result_chunks == ["Let ", "me ", "check...", "The ", "date ", "is ", "2024-11-15"]
+            assert mock_gateway.complete_stream.call_count == 2
+            mock_tool.run.assert_called_once_with(date="Friday")
+
+        def should_raise_error_if_gateway_does_not_support_streaming(self, llm_broker, mock_gateway):
+            messages = [LLMMessage(role=MessageRole.User, content="Hello")]
+
+            # Remove complete_stream method to simulate unsupported gateway
+            if hasattr(mock_gateway, 'complete_stream'):
+                delattr(mock_gateway, 'complete_stream')
+
+            with pytest.raises(NotImplementedError) as exc_info:
+                list(llm_broker.generate_stream(messages))
+
+            assert "does not support streaming" in str(exc_info.value)

mojentic/llm/message_composers.py

@@ -12,7 +12,7 @@ class FileTypeSensor:
         """
         Initialize the TypeSensor with a default mapping of file extensions to language declarations.
 
-        The TypeSensor is used to determine the appropriate language syntax highlighting 
+        The TypeSensor is used to determine the appropriate language syntax highlighting
         for code blocks in markdown based on file extensions.
         """
         self.extension_map: Dict[str, str] = {
@@ -129,7 +129,6 @@ class MessageBuilder():
                 f"{content.strip()}\n"
                 f"```\n")
 
-
     def add_image(self, image_path: Union[str, Path]) -> "MessageBuilder":
         """
         Add a single image to the message.
@@ -253,7 +252,11 @@ class MessageBuilder():
 
         return self
 
-    def load_content(self, file_path: Union[str, Path], template_values: Optional[Dict[str, Union[str, Path]]] = None) -> "MessageBuilder":
+    def load_content(
+        self,
+        file_path: Union[str, Path],
+        template_values: Optional[Dict[str, Union[str, Path]]] = None
+    ) -> "MessageBuilder":
         """
         Load content from a file into the content field of the MessageBuilder.
 

mojentic/llm/message_composers_spec.py

@@ -14,10 +14,12 @@ def file_gateway(mocker):
     file_gateway.is_binary.return_value = False
     return file_gateway
 
+
 @pytest.fixture
 def file_path():
     return Path("/path/to/file.txt")
 
+
 @pytest.fixture
 def whitespace_file_content():
     return "\n\n  \n  test file content with whitespace  \n\n  \n"
@@ -122,7 +124,9 @@ class DescribeMessageBuilder:
             assert "test file content" in result
             assert "```" in result
 
-        def should_strip_whitespace_from_file_content(self, message_builder, file_gateway, file_path, whitespace_file_content, mocker):
+        def should_strip_whitespace_from_file_content(
+            self, message_builder, file_gateway, file_path, whitespace_file_content, mocker
+        ):
             # Use the fixtures instead of creating file path and content directly
             file_gateway.read.return_value = whitespace_file_content
             mocker.patch.object(message_builder.type_sensor, 'get_language', return_value='text')

mojentic/llm/registry/__init__.py

@@ -1,6 +1,3 @@
 """
 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

@@ -24,9 +24,9 @@ def register_llms_from_ollama(url: str, registry: LLMRegistry):
         #         'quantization_level': 'Q4_K_M'
         #     },
         #     'digest': '4bd6cbf2d094264457a17aab6bd6acd1ed7a72fb8f8be3cfb193f63c78dd56df',
-        #     'model': 'qwen2.5-coder:32b',
+        #     'model': 'qwen3-coder:32b',
         #     'modified_at': '2025-01-29T22:37:29.191797577-05:00',
-        #     'name': 'qwen2.5-coder:32b',
+        #     'name': 'qwen3-coder:32b',
         #     'size': 19851349856
         # }
 

mojentic/llm/tools/__init__.py

@@ -3,16 +3,7 @@ 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/ask_user_tool.py

@@ -3,8 +3,8 @@ from mojentic.llm.tools.llm_tool import LLMTool
 
 class AskUserTool(LLMTool):
     def run(self, user_request: str) -> str:
-        print(f"\n\n\nI NEED YOUR HELP!\n{user_request}")
-        return input(f"Your response: ")
+        print("\n\n\nI NEED YOUR HELP!\n{user_request}")
+        return input("Your response: ")
 
     @property
     def descriptor(self):
@@ -12,16 +12,22 @@ class AskUserTool(LLMTool):
             "type": "function",
             "function": {
                 "name": "ask_user",
-                "description": "If you do not know how to proceed, ask the user a question, or ask them for help or to do something for you.",
+                "description": (
+                    "If you do not know how to proceed, ask the user a question, or ask them for "
+                    "help or to do something for you."
+                ),
                 "parameters": {
                     "type": "object",
                     "properties": {
                         "user_request": {
                             "type": "string",
-                            "description": "The question you need the user to answer, or the task you need the user to do for you."
+                            "description": (
+                                "The question you need the user to answer, or the task you need the user to "
+                                "do for you."
+                            )
                         }
                     },
                     "required": ["user_request"]
                 },
             },
-        }
+        }

mojentic/llm/tools/current_datetime.py

@@ -6,12 +6,12 @@ class CurrentDateTimeTool(LLMTool):
     def run(self, format_string: str = "%Y-%m-%d %H:%M:%S") -> dict:
         """
         Returns the current date and time.
-        
+
         Parameters
         ----------
         format_string : str, optional
             The format string for the datetime, by default "%Y-%m-%d %H:%M:%S"
-            
+
         Returns
         -------
         dict
@@ -19,13 +19,13 @@ class CurrentDateTimeTool(LLMTool):
         """
         current_time = datetime.now()
         formatted_time = current_time.strftime(format_string)
-        
+
         return {
             "current_datetime": formatted_time,
             "timestamp": current_time.timestamp(),
             "timezone": datetime.now().astimezone().tzname()
         }
-    
+
     @property
     def descriptor(self):
         return {
@@ -38,10 +38,13 @@ class CurrentDateTimeTool(LLMTool):
                     "properties": {
                         "format_string": {
                             "type": "string",
-                            "description": "Format string for the datetime (e.g., '%Y-%m-%d %H:%M:%S', '%A, %B %d, %Y'). Default is ISO format."
+                            "description": (
+                                "Format string for the datetime (e.g., '%Y-%m-%d %H:%M:%S', '%A, %B %d, %Y'). "
+                                "Default is ISO format."
+                            )
                         }
                     },
                     "required": []
                 }
             }
-        }
+        }

mojentic/llm/tools/date_resolver.py

@@ -7,7 +7,7 @@ from mojentic.llm.tools.llm_tool import LLMTool
 
 # Avoid circular imports with TYPE_CHECKING
 if TYPE_CHECKING:
-    from mojentic.llm.llm_broker import LLMBroker
+    pass
 
 
 class ResolveDateTool(LLMTool):
@@ -34,7 +34,10 @@ class ResolveDateTool(LLMTool):
             "type": "function",
             "function": {
                 "name": "resolve_date",
-                "description": "Take text that specifies a relative date, and output an absolute date. If no reference date is available, the current date is used.",
+                "description": (
+                    "Take text that specifies a relative date, and output an absolute date. If no "
+                    "reference date is available, the current date is used."
+                ),
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -44,8 +47,11 @@ class ResolveDateTool(LLMTool):
                         },
                         "reference_date_in_iso8601": {
                             "type": "string",
-                            "description": "The date from which the resolved date should be calculated, in YYYY-MM-DD"
-                                           " format. Do not provide if you weren't provided one, I will assume the current date."
+                            "description": (
+                                "The date from which the resolved date should be calculated, in YYYY-MM-DD "
+                                "format. Do not provide if you weren't provided one, I will assume the "
+                                "current date."
+                            )
                         }
                     },
                     "additionalProperties": False,

mojentic/llm/tools/date_resolver_spec.py

@@ -20,7 +20,6 @@ class DescribeResolveDateTool:
         assert result["relative_date"] == "next Friday"
         assert "resolved_date" in result
 
-
     def should_resolve_date_with_reference(self, date_resolver):
         """
         Given a date resolver

mojentic/llm/tools/ephemeral_task_manager/append_task_tool.py

@@ -61,7 +61,10 @@ class AppendTaskTool(LLMTool):
             "type": "function",
             "function": {
                 "name": "append_task",
-                "description": "Append a new task to the end of the task list with a description. The task will start with 'pending' status.",
+                "description": (
+                    "Append a new task to the end of the task list with a description. The task will "
+                    "start with 'pending' status."
+                ),
                 "parameters": {
                     "type": "object",
                     "properties": {

mojentic/llm/tools/ephemeral_task_manager/ephemeral_task_list.py

@@ -5,7 +5,7 @@ This module provides a class for managing a list of tasks with state transitions
 """
 
 from enum import Enum
-from typing import List, Optional
+from typing import List
 from pydantic import BaseModel
 
 

mojentic/llm/tools/ephemeral_task_manager/insert_task_after_tool.py

@@ -64,7 +64,10 @@ class InsertTaskAfterTool(LLMTool):
             "type": "function",
             "function": {
                 "name": "insert_task_after",
-                "description": "Insert a new task after an existing task in the task list. The task will start with 'pending' status.",
+                "description": (
+                    "Insert a new task after an existing task in the task list. The task will start with "
+                    "'pending' status."
+                ),
                 "parameters": {
                     "type": "object",
                     "properties": {

mojentic/llm/tools/ephemeral_task_manager/prepend_task_tool.py

@@ -61,7 +61,10 @@ class PrependTaskTool(LLMTool):
             "type": "function",
             "function": {
                 "name": "prepend_task",
-                "description": "Prepend a new task to the beginning of the task list with a description. The task will start with 'pending' status.",
+                "description": (
+                    "Prepend a new task to the beginning of the task list with a description. The task "
+                    "will start with 'pending' status."
+                ),
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -74,4 +77,4 @@ class PrependTaskTool(LLMTool):
                     "additionalProperties": False
                 }
             }
-        }
+        }