modmex-ai 0.1.0__py3-none-any.whl

modmex_ai/__init__.py

@@ -0,0 +1,155 @@
+"""Lightweight AI flows, agents, tools, and model adapters."""
+
+from modmex_ai.agents import (
+    Agent,
+    AgentResult,
+    AgentStreamEvent,
+    AgentStreamEventType,
+    Handoff,
+    RECOMMENDED_PROMPT_PREFIX,
+    RunContext,
+    prompt_with_handoff_instructions,
+)
+from modmex_ai.models import (
+    AsyncModelClient,
+    FakeModel,
+    FallbackModel,
+    ModelClient,
+    ModelProfile,
+    ModelRequest,
+    ModelResponse,
+    ModelStreamEvent,
+    ModelStreamEventType,
+    ModelSettings,
+    ProviderState,
+    ToolCall,
+    Usage,
+)
+from modmex_ai.sessions import InMemorySession, Session, SessionItem, SessionSnapshot
+from modmex_ai.sessions import DurableSessionStore, InMemoryDurableSessionStore, SessionConflictError
+from modmex_ai.approvals import ApprovalDecision, ApprovalDecisionType, ApprovalPolicy, ApprovalRequest
+from modmex_ai.observability import ObservabilityEvent, ObservabilityObserver
+from modmex_ai.evals import EvalCase, EvalResult, EvalRunner
+from modmex_ai.tools import Tool, ToolExecutor, ToolResult, as_tool, tool
+from modmex_ai.flows import Flow, FlowContinuation, FlowResult, FlowStreamEvent, FlowStreamEventType, FlowStateConflictError, FlowStateStatus, FlowStateStore, FlowSuspended, FlowSuspension, InMemoryFlowStateStore, PersistedFlowState
+from modmex_ai.guardrails import (
+    Guardrail,
+    GuardrailResult,
+    InputGuardrail,
+    OutputGuardrail,
+    ToolInputGuardrail,
+    ToolOutputGuardrail,
+)
+from modmex_ai.realtime import RealtimeEvent, RealtimeSession, RealtimeTransport
+from modmex_ai.voice import (
+    ChainedVoiceSession,
+    CallableSpeechToTextProvider,
+    CallableTextToSpeechProvider,
+    LiveSpeechToTextProvider,
+    LiveSpeechToTextSession,
+    SpeechToTextProvider,
+    StreamingSpeechToTextProvider,
+    StreamingTextToSpeechProvider,
+    TextToSpeechProvider,
+    Transcription,
+    VoiceInputEvent,
+    VoiceInputEventType,
+    VoiceAgentSession,
+    VoiceSessionEvent,
+    VoiceSessionEventType,
+    VoiceContinuation,
+    VoiceTerminationReason,
+    VoiceTurnOptions,
+    VoiceTurnMetrics,
+    VoiceTurnOutcome,
+    VoiceTurnReason,
+    VoiceTurnObserver,
+)
+
+__all__ = [
+    "Agent",
+    "AsyncModelClient",
+    "AgentResult",
+    "AgentStreamEvent",
+    "AgentStreamEventType",
+    "FakeModel",
+    "FallbackModel",
+    "InMemorySession",
+    "DurableSessionStore",
+    "InMemoryDurableSessionStore",
+    "SessionConflictError",
+    "ModelClient",
+    "ModelProfile",
+    "ModelRequest",
+    "ModelResponse",
+    "ModelStreamEvent",
+    "ModelStreamEventType",
+    "ModelSettings",
+    "ProviderState",
+    "RECOMMENDED_PROMPT_PREFIX",
+    "RunContext",
+    "Session",
+    "SessionItem",
+    "SessionSnapshot",
+    "ApprovalDecision",
+    "ApprovalDecisionType",
+    "ApprovalPolicy",
+    "ApprovalRequest",
+    "ObservabilityEvent",
+    "ObservabilityObserver",
+    "EvalCase",
+    "EvalResult",
+    "EvalRunner",
+    "Tool",
+    "ToolCall",
+    "ToolExecutor",
+    "ToolResult",
+    "Flow",
+    "FlowContinuation",
+    "FlowResult",
+    "FlowStreamEvent",
+    "FlowStreamEventType",
+    "FlowStateConflictError",
+    "FlowStateStatus",
+    "FlowStateStore",
+    "FlowSuspended",
+    "FlowSuspension",
+    "InMemoryFlowStateStore",
+    "PersistedFlowState",
+    "Handoff",
+    "Guardrail",
+    "GuardrailResult",
+    "InputGuardrail",
+    "OutputGuardrail",
+    "ToolInputGuardrail",
+    "ToolOutputGuardrail",
+    "Usage",
+    "ChainedVoiceSession",
+    "CallableSpeechToTextProvider",
+    "CallableTextToSpeechProvider",
+    "SpeechToTextProvider",
+    "LiveSpeechToTextProvider",
+    "LiveSpeechToTextSession",
+    "StreamingSpeechToTextProvider",
+    "StreamingTextToSpeechProvider",
+    "TextToSpeechProvider",
+    "Transcription",
+    "VoiceInputEvent",
+    "VoiceInputEventType",
+    "VoiceAgentSession",
+    "VoiceContinuation",
+    "VoiceSessionEvent",
+    "VoiceSessionEventType",
+    "VoiceTerminationReason",
+    "VoiceTurnOptions",
+    "VoiceTurnMetrics",
+    "VoiceTurnOutcome",
+    "VoiceTurnReason",
+    "VoiceTurnObserver",
+    "RealtimeEvent",
+    "RealtimeSession",
+    "RealtimeTransport",
+    "as_tool",
+    "prompt_with_handoff_instructions",
+    "tool",
+]

modmex_ai/agents/__init__.py

@@ -0,0 +1,16 @@
+from modmex_ai.agents.agent import Agent
+from modmex_ai.agents.context import RunContext
+from modmex_ai.agents.handoff import Handoff, RECOMMENDED_PROMPT_PREFIX, prompt_with_handoff_instructions
+from modmex_ai.agents.result import AgentResult
+from modmex_ai.agents.stream import AgentStreamEvent, AgentStreamEventType
+
+__all__ = [
+    "Agent",
+    "AgentResult",
+    "AgentStreamEvent",
+    "AgentStreamEventType",
+    "Handoff",
+    "RECOMMENDED_PROMPT_PREFIX",
+    "RunContext",
+    "prompt_with_handoff_instructions",
+]

modmex_ai/agents/agent.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any
+
+from modmex_ai.agents.context import RunContext
+from modmex_ai.agents.execution import AgentExecution
+from modmex_ai.agents.handoff import Handoff, normalize_handoffs
+from modmex_ai.agents.result import AgentResult
+from modmex_ai.agents.stream import AgentStreamEvent, AgentStreamEventType
+from modmex_ai.errors import OutputGuardrailTriggered, OutputValidationError
+from modmex_ai.guardrails import (
+    InputGuardrail,
+    OutputGuardrail,
+    ToolInputGuardrail,
+    ToolOutputGuardrail,
+    enforce_input_guardrails,
+    enforce_output_guardrails,
+)
+from modmex_ai.messages import Message
+from modmex_ai.models import (
+    ModelClient,
+    ModelResponse,
+    ModelSettings,
+    ModelStreamEvent,
+    ModelStreamEventType,
+    ProviderState,
+    Usage,
+)
+from modmex_ai.schemas import dumps, schema_for_model, validate_model
+from modmex_ai.sessions import SessionItem
+from modmex_ai.tools import Tool, as_tool, tool as create_tool
+
+
+_NESTED_USAGE_STATE_KEY = "__modmex_ai_nested_usage"
+
+
+class Agent:
+    def __init__(
+        self,
+        *,
+        name: str,
+        instructions: str,
+        output_type: type[Any] | None = None,
+        output_strict: bool = True,
+        tools: list[Tool | Any] | None = None,
+        handoffs: list[str | Handoff] | None = None,
+        model: ModelClient | None = None,
+        settings: ModelSettings | None = None,
+        max_tool_calls: int = 8,
+        input_guardrails: list[InputGuardrail] | None = None,
+        output_guardrails: list[OutputGuardrail] | None = None,
+        max_output_guardrail_retries: int = 0,
+    ) -> None:
+        self.name = name
+        self.instructions = instructions
+        self.output_type = output_type
+        self.output_strict = output_strict
+        self.tools = [as_tool(value) for value in (tools or [])]
+        self.handoffs = normalize_handoffs(handoffs)
+        self.model = model
+        self.settings = settings
+        self.max_tool_calls = max_tool_calls
+        self.input_guardrails = input_guardrails or []
+        self.output_guardrails = output_guardrails or []
+        if max_output_guardrail_retries < 0:
+            raise ValueError("max_output_guardrail_retries must be non-negative")
+        self.max_output_guardrail_retries = max_output_guardrail_retries
+
+    def add_tool(self, value: Tool | Any) -> Tool:
+        registered = as_tool(value)
+        self.tools.append(registered)
+        return registered
+
+    def tool(
+        self,
+        func: Any = None,
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        input_guardrails: list[ToolInputGuardrail] | None = None,
+        output_guardrails: list[ToolOutputGuardrail] | None = None,
+        requires_approval: bool = False,
+        approval_reason: str | None = None,
+    ):
+        def decorator(inner):
+            return self.add_tool(create_tool(
+                inner,
+                name=name,
+                description=description,
+                input_guardrails=input_guardrails,
+                output_guardrails=output_guardrails,
+                requires_approval=requires_approval,
+                approval_reason=approval_reason,
+            ))
+        return decorator if func is None else decorator(func)
+
+    def as_tool(self, *, tool_name: str | None = None, tool_description: str | None = None) -> Tool:
+        def run_agent(input: str, ctx=None) -> Any:
+            result = self.run(input, context=ctx)
+            if isinstance(ctx, RunContext):
+                nested_usage = ctx.state.setdefault(_NESTED_USAGE_STATE_KEY, Usage())
+                nested_usage.add(result.usage)
+            return result.output
+        return create_tool(run_agent, name=tool_name or self.name, description=tool_description or self.instructions)
+
+    def run(self, input: Any, **kwargs: Any) -> AgentResult:
+        return self._run_execution(self._execution(input, **kwargs))
+
+    async def run_async(self, input: Any, **kwargs: Any) -> AgentResult:
+        execution = self._execution(input, **kwargs)
+        acomplete = getattr(execution.model, "acomplete", None)
+        if not callable(acomplete):
+            return await asyncio.to_thread(self._run_execution, execution)
+        while True:
+            step = execution.accept_response(await acomplete(execution.next_request()))
+            if step.result is not None:
+                return step.result
+            for tool_call in step.tool_calls or []:
+                execution.begin_tool_call(tool_call)
+                execution.accept_tool_result(await execution.executor.execute_async(tool_call, context=execution.context))
+
+    def run_stream(self, input: Any, **kwargs: Any):
+        execution = self._execution(input, **kwargs)
+        while True:
+            completed: ModelResponse | None = None
+            for raw_event in execution.model.stream(execution.next_request()):
+                event = self._model_stream_event(raw_event)
+                if event.type == ModelStreamEventType.TEXT_DELTA:
+                    yield AgentStreamEvent(type=AgentStreamEventType.TEXT_DELTA, text_delta=event.text_delta or "")
+                elif event.type == ModelStreamEventType.TOOL_CALL_DELTA:
+                    yield AgentStreamEvent(type=AgentStreamEventType.TOOL_CALL_DELTA, tool_call=event.tool_call)
+                elif event.type == ModelStreamEventType.COMPLETED:
+                    completed = event.response
+            if completed is None:
+                raise RuntimeError("Model stream ended without a completed response")
+            step = execution.accept_response(completed)
+            if step.result is not None:
+                yield AgentStreamEvent(type=AgentStreamEventType.HANDOFF if step.result.handoff_target else AgentStreamEventType.COMPLETED, result=step.result)
+                return
+            for tool_call in step.tool_calls or []:
+                execution.begin_tool_call(tool_call)
+                yield AgentStreamEvent(type=AgentStreamEventType.TOOL_CALL_DELTA, tool_call=tool_call)
+                tool_result = execution.executor.execute(tool_call, context=execution.context)
+                execution.accept_tool_result(tool_result)
+                yield AgentStreamEvent(type=AgentStreamEventType.TOOL_FINISHED, tool_call=tool_call, data={"output": tool_result.output})
+
+    async def arun_stream(self, input: Any, **kwargs: Any):
+        """Async counterpart of ``run_stream`` for models with native streaming."""
+        execution = self._execution(input, **kwargs)
+        astream = getattr(execution.model, "astream", None)
+        if not callable(astream):
+            raise NotImplementedError("Model does not implement native async streaming")
+        while True:
+            completed: ModelResponse | None = None
+            async for raw_event in astream(execution.next_request()):
+                event = self._model_stream_event(raw_event)
+                if event.type == ModelStreamEventType.TEXT_DELTA:
+                    yield AgentStreamEvent(type=AgentStreamEventType.TEXT_DELTA, text_delta=event.text_delta or "")
+                elif event.type == ModelStreamEventType.TOOL_CALL_DELTA:
+                    yield AgentStreamEvent(type=AgentStreamEventType.TOOL_CALL_DELTA, tool_call=event.tool_call)
+                elif event.type == ModelStreamEventType.COMPLETED:
+                    completed = event.response
+            if completed is None:
+                raise RuntimeError("Model stream ended without a completed response")
+            step = execution.accept_response(completed)
+            if step.result is not None:
+                yield AgentStreamEvent(type=AgentStreamEventType.HANDOFF if step.result.handoff_target else AgentStreamEventType.COMPLETED, result=step.result)
+                return
+            for tool_call in step.tool_calls or []:
+                execution.begin_tool_call(tool_call)
+                yield AgentStreamEvent(type=AgentStreamEventType.TOOL_CALL_DELTA, tool_call=tool_call)
+                tool_result = await execution.executor.execute_async(tool_call, context=execution.context)
+                execution.accept_tool_result(tool_result)
+                yield AgentStreamEvent(type=AgentStreamEventType.TOOL_FINISHED, tool_call=tool_call, data={"output": tool_result.output})
+
+    def _run_execution(self, execution: AgentExecution) -> AgentResult:
+        while True:
+            step = execution.accept_response(execution.model.complete(execution.next_request()))
+            if step.result is not None:
+                return step.result
+            for tool_call in step.tool_calls or []:
+                execution.begin_tool_call(tool_call)
+                execution.accept_tool_result(execution.executor.execute(tool_call, context=execution.context))
+
+    def _execution(
+        self,
+        input: Any,
+        *,
+        model: ModelClient | None = None,
+        context: RunContext | Any = None,
+        provider_state: ProviderState | None = None,
+        run_input_guardrails: bool = True,
+        run_output_guardrails: bool = True,
+    ) -> AgentExecution:
+        active_model = model or self.model
+        if active_model is None:
+            raise ValueError(f"Agent {self.name!r} has no model")
+        run_context = context if isinstance(context, RunContext) else RunContext(input=input, context=context)
+        if run_input_guardrails:
+            enforce_input_guardrails(input, self.input_guardrails, context=run_context)
+        return AgentExecution(agent=self, model=active_model, input=input, context=run_context, provider_state=provider_state, run_output_guardrails=run_output_guardrails)
+
+    @staticmethod
+    def _model_stream_event(value: ModelStreamEvent | ModelResponse) -> ModelStreamEvent:
+        if isinstance(value, ModelResponse):
+            return ModelStreamEvent.completed(value)
+        if isinstance(value, ModelStreamEvent):
+            return value
+        raise TypeError("Model stream must yield ModelStreamEvent or ModelResponse")
+
+    def _output_schema(self) -> dict[str, Any] | None:
+        return schema_for_model(self.output_type) if self.output_type else None
+
+    def _should_retry_output(self, *, output: Any, response_text: str | None, messages: list[Message | SessionItem], context: RunContext, retry_count: int) -> bool:
+        try:
+            enforce_output_guardrails(output, self.output_guardrails, context=context)
+        except OutputGuardrailTriggered as error:
+            if retry_count >= self.max_output_guardrail_retries:
+                raise
+            context.trace.add("guardrail", f"{self.name}_output_retry", attempt=retry_count + 1, reason=str(error))
+            messages.extend([
+                Message(role="assistant", content=response_text or ""),
+                Message(role="developer", content="The previous candidate output was rejected by an output guardrail: " f"{error}. Produce a corrected replacement that fully satisfies the output schema and the guardrail rule."),
+            ])
+            return True
+        return False
+
+    def _initial_messages(self, input: Any, *, include_output_schema_prompt: bool = True) -> list[Message | SessionItem]:
+        instructions = self.instructions
+        if self.output_type and include_output_schema_prompt:
+            instructions += "\n\nReturn only JSON that matches this schema:\n" + json.dumps(schema_for_model(self.output_type), separators=(",", ":"))
+        return [Message(role="system", content=instructions), *self._input_messages(input)]
+
+    def _input_messages(self, input: Any) -> list[Message | SessionItem]:
+        if isinstance(input, list) and all(isinstance(item, (Message, SessionItem, dict)) for item in input):
+            return [item if isinstance(item, (Message, SessionItem)) else SessionItem(**item) if item.get("type") else Message(**item) for item in input]
+        return [Message(role="user", content=input if isinstance(input, str) else dumps(input))]
+
+    def _parse_output(self, text: str | None) -> Any:
+        if self.output_type is None:
+            return text or ""
+        if text is None:
+            raise OutputValidationError(f"Agent {self.name!r} returned no output")
+        return validate_model(json.loads(text), self.output_type)
+
+    def _handoff_schemas(self) -> list[dict[str, Any]]:
+        return [handoff.schema() for handoff in self.handoffs]

modmex_ai/agents/context.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from dataclasses import field
+from typing import Any
+
+from modmex import BaseModel
+
+from modmex_ai.tracing import Trace
+
+
+class RunContext(BaseModel):
+    input: Any
+    context: Any = None
+    state: dict[str, Any] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    trace: Trace = field(default_factory=Trace)

modmex_ai/agents/execution.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from modmex_ai.agents.context import RunContext
+from modmex_ai.agents.handoff import find_handoff
+from modmex_ai.agents.result import AgentResult
+from modmex_ai.errors import MaxToolCallsExceeded
+from modmex_ai.messages import Message
+from modmex_ai.models import ModelClient, ModelRequest, ModelResponse, ProviderState, ToolCall, Usage
+from modmex_ai.sessions import SessionItem
+from modmex_ai.tools import ToolExecutor, ToolResult
+
+if TYPE_CHECKING:
+    from modmex_ai.agents.agent import Agent
+
+
+@dataclass
+class AgentExecutionStep:
+    """The next action requested by the provider-neutral agent state machine."""
+
+    result: AgentResult | None = None
+    tool_calls: list[ToolCall] | None = None
+
+
+class AgentExecution:
+    """Shared Agent state machine; sync and async drivers supply I/O only."""
+
+    def __init__(
+        self,
+        *,
+        agent: Agent,
+        model: ModelClient,
+        input: Any,
+        context: RunContext,
+        provider_state: ProviderState | None,
+        run_output_guardrails: bool,
+    ) -> None:
+        self.agent = agent
+        self.model = model
+        self.context = context
+        self.run_output_guardrails = run_output_guardrails
+        self.messages = agent._initial_messages(
+            input,
+            include_output_schema_prompt=not self.model.profile.supports_structured_output,
+        )
+        self.executor = ToolExecutor(agent.tools)
+        self.tool_schemas = [*self.executor.schemas(), *agent._handoff_schemas()]
+        self.output_schema = agent._output_schema()
+        self.tool_call_count = 0
+        self.output_guardrail_retry_count = 0
+        self.usage = Usage()
+        self.generated_items: list[SessionItem] = []
+        self.provider_state = provider_state
+
+    def next_request(self) -> ModelRequest:
+        self.context.trace.add(
+            "model_request",
+            self.agent.name,
+            tools=[tool["name"] for tool in self.tool_schemas],
+        )
+        return ModelRequest(
+            messages=self.messages,
+            tools=self.tool_schemas,
+            output_schema=self.output_schema,
+            output_strict=self.agent.output_strict,
+            settings=self.agent.settings,
+            provider_state=self.provider_state,
+        )
+
+    def accept_response(self, response: ModelResponse) -> AgentExecutionStep:
+        if response.provider_state is not None:
+            self.provider_state = response.provider_state
+        self.usage.add(response.usage)
+        self.context.trace.add(
+            "model_response",
+            self.agent.name,
+            request_id=response.request_id,
+            tool_calls=[call.name for call in response.tool_calls],
+            usage=response.usage,
+        )
+        if response.tool_calls:
+            return self._tool_step(response.tool_calls)
+
+        output = self.agent._parse_output(response.output_text)
+        if self.run_output_guardrails and self.agent._should_retry_output(
+            output=output,
+            response_text=response.output_text,
+            messages=self.messages,
+            context=self.context,
+            retry_count=self.output_guardrail_retry_count,
+        ):
+            self.output_guardrail_retry_count += 1
+            return AgentExecutionStep()
+        return AgentExecutionStep(result=self._result(output=output))
+
+    def begin_tool_call(self, tool_call: ToolCall) -> None:
+        item = SessionItem(
+            type="function_call",
+            tool_call_id=tool_call.tool_call_id,
+            name=tool_call.name,
+            arguments=tool_call.arguments,
+        )
+        self.generated_items.append(item)
+        self.messages.append(item)
+
+    def accept_tool_result(self, result: ToolResult) -> None:
+        item = SessionItem(
+            type="function_call_output",
+            tool_call_id=result.tool_call_id,
+            name=result.name,
+            output=result.output,
+        )
+        self.generated_items.append(item)
+        nested_usage = self.context.state.pop("__modmex_ai_nested_usage", None)
+        if nested_usage:
+            self.usage.add(nested_usage)
+        self.context.trace.add("tool_call", result.name, tool_call_id=result.tool_call_id)
+        self.messages.append(item)
+
+    def _tool_step(self, tool_calls: list[ToolCall]) -> AgentExecutionStep:
+        self.tool_call_count += len(tool_calls)
+        if self.tool_call_count > self.agent.max_tool_calls:
+            raise MaxToolCallsExceeded(f"Agent {self.agent.name!r} exceeded max_tool_calls")
+        for tool_call in tool_calls:
+            handoff = find_handoff(tool_call, self.agent.handoffs)
+            if handoff is None:
+                continue
+            handoff_input = handoff.invoke(self.context, tool_call.arguments)
+            self.generated_items.extend([
+                SessionItem(
+                    type="handoff_call",
+                    tool_call_id=tool_call.tool_call_id,
+                    name=handoff.name,
+                    arguments=tool_call.arguments,
+                ),
+                SessionItem(
+                    type="handoff_call_output",
+                    tool_call_id=tool_call.tool_call_id,
+                    name=handoff.name,
+                    output={"transferred": True},
+                ),
+            ])
+            self.context.trace.add(
+                "handoff_call",
+                handoff.name,
+                tool_call_id=tool_call.tool_call_id,
+                to=handoff.agent,
+            )
+            return AgentExecutionStep(result=AgentResult(
+                output=None,
+                agent=self.agent.name,
+                handoff_target=handoff.agent,
+                handoff_name=handoff.name,
+                handoff_input=handoff_input,
+                usage=self.usage,
+                items=self.generated_items,
+                provider_state=self.provider_state,
+                trace=self.context.trace,
+            ))
+        return AgentExecutionStep(tool_calls=tool_calls)
+
+    def _result(self, *, output: Any) -> AgentResult:
+        return AgentResult(
+            output=output,
+            agent=self.agent.name,
+            usage=self.usage,
+            items=self.generated_items,
+            trace=self.context.trace,
+            provider_state=self.provider_state,
+        )