mantisdk 0.1.0__py3-none-any.whl

mantisdk/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+__version__ = "0.1.0"
+
+from .adapter import *
+from .algorithm import *
+from .client import MantisdkClient, DevTaskLoader  # deprecated  # type: ignore
+from .config import *
+from .emitter import *
+from .env_var import *
+from .execution import *
+from .litagent import *
+from .llm_proxy import *
+from .logging import configure_logger  # deprecated  # type: ignore
+from .logging import setup as setup_logging  # type: ignore
+from .logging import setup_module as setup_module_logging  # type: ignore
+from .runner import *
+from .server import MantisdkServer  # deprecated  # type: ignore
+from .store import *
+from .tracer import *
+from .trainer import *
+from .types import *

mantisdk/adapter/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import Adapter, OtelTraceAdapter, TraceAdapter
+from .messages import TraceToMessages
+from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase
+
+__all__ = [
+    "TraceAdapter",
+    "OtelTraceAdapter",
+    "Adapter",
+    "TraceToTripletBase",
+    "TracerTraceToTriplet",
+    "LlmProxyTraceToTriplet",
+    "TraceToMessages",
+]

mantisdk/adapter/base.py

@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Generic, Sequence, TypeVar
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from mantisdk.types import Span
+
+T_from = TypeVar("T_from")
+T_to = TypeVar("T_to")
+
+
+class Adapter(Generic[T_from, T_to]):
+    """Base class for synchronous adapters that convert data from one format to another.
+
+    The class defines a minimal protocol so that adapters can be treated like callables while
+    still allowing subclasses to supply the concrete transformation logic.
+
+    !!! note
+        Subclasses must override [`adapt()`][mantisdk.Adapter.adapt] to provide
+        the actual conversion.
+
+    Type Variables:
+
+        T_from: Source data type supplied to the adapter.
+
+        T_to: Target data type produced by the adapter.
+
+    Examples:
+        >>> class IntToStrAdapter(Adapter[int, str]):
+        ...     def adapt(self, source: int) -> str:
+        ...         return str(source)
+        ...
+        >>> adapter = IntToStrAdapter()
+        >>> adapter(42)
+        '42'
+    """
+
+    def __call__(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        This method delegates to [`adapt()`][mantisdk.Adapter.adapt] so that an
+        instance of [`Adapter`][mantisdk.Adapter] can be used like a standard
+        function.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        return self.adapt(source)
+
+    def adapt(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        Subclasses must override this method with the concrete transformation logic. The base
+        implementation raises `NotImplementedError` to make the requirement explicit.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        raise NotImplementedError("Adapter.adapt() is not implemented")
+
+
+class OtelTraceAdapter(Adapter[Sequence[ReadableSpan], T_to], Generic[T_to]):
+    """Base class for adapters that convert OpenTelemetry trace spans into other formats.
+
+    This specialization of [`Adapter`][mantisdk.Adapter] expects a list of
+    `opentelemetry.sdk.trace.ReadableSpan` instances and produces any target format, such as
+    reinforcement learning trajectories, structured logs, or analytics-ready payloads.
+
+    Examples:
+        >>> class TraceToDictAdapter(OtelTraceAdapter[dict]):
+        ...     def adapt(self, spans: List[ReadableSpan]) -> dict:
+        ...         return {"count": len(spans)}
+        ...
+        >>> adapter = TraceToDictAdapter()
+        >>> adapter([span1, span2])
+        {'count': 2}
+    """
+
+
+class TraceAdapter(Adapter[Sequence[Span], T_to], Generic[T_to]):
+    """Base class for adapters that convert trace spans into other formats.
+
+    This class specializes [`Adapter`][mantisdk.Adapter] for working with
+    [`Span`][mantisdk.Span] instances emitted by Mantisdk instrumentation.
+    Subclasses receive entire trace slices and return a format suited for the downstream consumer,
+    for example reinforcement learning training data or observability metrics.
+    """

mantisdk/adapter/messages.py

@@ -0,0 +1,270 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, Generator, Iterable, List, Optional, Sequence, TypedDict, Union, cast
+
+from pydantic import TypeAdapter
+
+from mantisdk.types import Span
+
+from .base import TraceAdapter
+
+if TYPE_CHECKING:
+    from openai.types.chat import (
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+
+
+class OpenAIMessages(TypedDict):
+    """OpenAI-style chat messages with optional tool definitions.
+
+    Attributes:
+        messages: Ordered chat messages that describe the conversation.
+        tools: Tool specifications available to the assistant, if any.
+    """
+
+    messages: List[ChatCompletionMessageParam]
+    tools: Optional[List[ChatCompletionFunctionToolParam]]
+
+
+class _RawSpanInfo(TypedDict):
+    """Intermediate representation parsed from a span.
+
+    Attributes:
+        prompt: Prompt messages reconstructed from span attributes.
+        completion: Assistant completions following tool invocations.
+        request: Request payload recorded in the trace.
+        response: Response payload recorded in the trace.
+        tools: Tool call metadata extracted from child spans.
+    """
+
+    prompt: List[Dict[str, Any]]
+    completion: List[Dict[str, Any]]
+    request: Dict[str, Any]
+    response: Dict[str, Any]
+    tools: List[Dict[str, Any]]
+
+
+def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
+    """Convert flattened trace attributes into nested structures.
+
+    Attributes emitted by the tracing pipeline often arrive as dotted paths (for example
+    `gen_ai.prompt.0.role`). This helper groups those keys into nested dictionaries or lists so that
+    downstream processing can operate on structured data.
+
+    Args:
+        data: Flat dictionary whose keys are dotted paths.
+        prefix: Top-level key (for example `gen_ai.prompt`) that determines which attributes are
+            grouped.
+
+    Returns:
+        A nested dictionary (no numeric index detected) or list (numeric indices detected) containing
+        the grouped values.
+    """
+    result: Union[Dict[str, Any], List[Any]] = {}
+
+    # Collect keys that match the prefix
+    relevant = {k[len(prefix) + 1 :]: v for k, v in data.items() if k.startswith(prefix + ".")}
+
+    # Detect if we have numeric indices (-> list) or not (-> dict)
+    indexed = any(part.split(".")[0].isdigit() for part in relevant.keys())
+
+    if indexed:
+        # Group by index
+        grouped: Dict[int, Dict[str, Any]] = defaultdict(dict)
+        for k, v in relevant.items():
+            parts = k.split(".")
+            if not parts[0].isdigit():
+                continue
+            idx, rest = int(parts[0]), ".".join(parts[1:])
+            grouped[idx][rest] = v
+        # Recursively build
+        result = []
+        for i in sorted(grouped.keys()):
+            result.append(group_genai_dict({f"{prefix}.{rest}": val for rest, val in grouped[i].items()}, prefix))
+    else:
+        # No indices: build dict
+        nested: Dict[str, Any] = defaultdict(dict)
+        for k, v in relevant.items():
+            if "." in k:
+                head, _tail = k.split(".", 1)
+                nested[head][f"{prefix}.{k}"] = v
+            else:
+                result[k] = v
+        # Recurse into nested dicts
+        for head, subdict in nested.items():
+            result[head] = group_genai_dict(subdict, prefix + "." + head)
+
+    return result
+
+
+def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]:
+    """Convert raw trace payloads into OpenAI-style chat messages.
+
+    The function consumes an iterable produced by
+    [`TraceToMessages.adapt()`][mantisdk.TraceToMessages.adapt] and yields
+    structures that match the OpenAI fine-tuning JSONL schema, including tool definitions.
+
+    Args:
+        prompt_completion_list: Raw prompt/completion/tool payloads extracted from a trace.
+
+    Returns:
+        A generator that yields [`OpenAIMessages`][mantisdk.adapter.messages.OpenAIMessages]
+        entries compatible with the OpenAI Functions fine-tuning format.
+    """
+
+    # Import locally to avoid legacy OpenAI version type import errors
+    from openai.types.chat import (
+        ChatCompletionAssistantMessageParam,
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+
+    for pc_entry in prompt_completion_list:
+        messages: List[ChatCompletionMessageParam] = []
+
+        # Extract messages
+        for msg in pc_entry["prompt"]:
+            role = msg["role"]
+
+            if role == "assistant" and "tool_calls" in msg:
+                # Use the tool_calls directly
+                # This branch is usually not used in the wild.
+                tool_calls: List[ChatCompletionMessageFunctionToolCallParam] = [
+                    ChatCompletionMessageFunctionToolCallParam(
+                        id=call["id"],
+                        type="function",
+                        function={"name": call["name"], "arguments": call["arguments"]},
+                    )
+                    for call in msg["tool_calls"]
+                ]
+                messages.append(
+                    ChatCompletionAssistantMessageParam(role="assistant", content=None, tool_calls=tool_calls)
+                )
+            else:
+                # Normal user/system/tool content
+                message = cast(
+                    ChatCompletionMessageParam,
+                    TypeAdapter(ChatCompletionMessageParam).validate_python(
+                        dict(role=role, content=msg.get("content", ""), tool_call_id=msg.get("tool_call_id", None))
+                    ),
+                )
+                messages.append(message)
+
+        # Extract completions (assistant outputs after tool responses)
+        for comp in pc_entry["completion"]:
+            if comp.get("role") == "assistant":
+                content = comp.get("content")
+                if pc_entry["tools"]:
+                    tool_calls = [
+                        ChatCompletionMessageFunctionToolCallParam(
+                            id=tool["call"]["id"],
+                            type=tool["call"]["type"],
+                            function={"name": tool["name"], "arguments": tool["parameters"]},
+                        )
+                        for tool in pc_entry["tools"]
+                    ]
+                    messages.append(
+                        ChatCompletionAssistantMessageParam(role="assistant", content=content, tool_calls=tool_calls)
+                    )
+                else:
+                    messages.append(ChatCompletionAssistantMessageParam(role="assistant", content=content))
+
+        # Build tools definitions (if available)
+        if "functions" in pc_entry["request"]:
+            tools = [
+                ChatCompletionFunctionToolParam(
+                    type="function",
+                    function={
+                        "name": fn["name"],
+                        "description": fn.get("description", ""),
+                        "parameters": (
+                            json.loads(fn["parameters"]) if isinstance(fn["parameters"], str) else fn["parameters"]
+                        ),
+                    },
+                )
+                for fn in pc_entry["request"]["functions"]
+            ]
+            yield OpenAIMessages(messages=messages, tools=tools)
+        else:
+            yield OpenAIMessages(messages=messages, tools=None)
+
+
+class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
+    """Convert trace spans into OpenAI-compatible conversation messages.
+
+    The adapter reconstructs prompts, completions, tool calls, and function definitions from
+    `gen_ai.*` span attributes. The resulting objects match the JSONL structure expected by the
+    OpenAI fine-tuning pipeline.
+
+    !!! warning
+        The adapter assumes all spans share a common trace and that tool call spans are direct
+        children of the associated completion span.
+    """
+
+    def get_tool_calls(self, completion: Span, all_spans: Sequence[Span], /) -> Iterable[Dict[str, Any]]:
+        """Yield tool call payloads for a completion span.
+
+        Args:
+            completion: The completion span whose descendants should be inspected.
+            all_spans: The complete span list belonging to the trace.
+
+        Yields:
+            Dictionaries describing tool calls with identifiers, names, and arguments.
+
+        Raises:
+            ValueError: If a candidate tool span cannot be converted into a dictionary.
+        """
+        # Get all the spans that are children of the completion span
+        children = [span for span in all_spans if span.parent_id == completion.span_id]
+        # Get the tool calls from the children
+        for maybe_tool_call in children:
+            tool_call = group_genai_dict(maybe_tool_call.attributes, "tool")
+            if not isinstance(tool_call, dict):
+                raise ValueError(f"Extracted tool call from trace is not a dict: {tool_call}")
+            if tool_call:
+                yield tool_call
+
+    def adapt(self, source: Sequence[Span], /) -> List[OpenAIMessages]:
+        """Transform trace spans into OpenAI chat payloads.
+
+        Args:
+            source: Spans containing `gen_ai.*` attributes emitted by the tracing pipeline.
+
+        Returns:
+            A list of [`OpenAIMessages`][mantisdk.adapter.messages.OpenAIMessages] entries that
+            capture prompts, completions, tools, and metadata.
+        """
+        raw_prompt_completions: List[_RawSpanInfo] = []
+
+        for span in source:
+            attributes = {k: v for k, v in span.attributes.items()}
+
+            # Get all related information from the trace span
+            prompt = group_genai_dict(attributes, "gen_ai.prompt") or []
+            completion = group_genai_dict(attributes, "gen_ai.completion") or []
+            request = group_genai_dict(attributes, "gen_ai.request") or {}
+            response = group_genai_dict(attributes, "gen_ai.response") or {}
+            if not isinstance(prompt, list):
+                raise ValueError(f"Extracted prompt from trace is not a list: {prompt}")
+            if not isinstance(completion, list):
+                raise ValueError(f"Extracted completion from trace is not a list: {completion}")
+            if not isinstance(request, dict):
+                raise ValueError(f"Extracted request from trace is not a dict: {request}")
+            if not isinstance(response, dict):
+                raise ValueError(f"Extracted response from trace is not a dict: {response}")
+            if prompt or completion or request or response:
+                tools = list(self.get_tool_calls(span, source)) or []
+                raw_prompt_completions.append(
+                    _RawSpanInfo(
+                        prompt=prompt or [], completion=completion, request=request, response=response, tools=tools
+                    )
+                )
+
+        return list(convert_to_openai_messages(raw_prompt_completions))