mwin 0.2.2__tar.gz → 0.2.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (60) hide show
  1. {mwin-0.2.2 → mwin-0.2.4}/PKG-INFO +3 -3
  2. {mwin-0.2.2 → mwin-0.2.4}/README.md +2 -2
  3. {mwin-0.2.2 → mwin-0.2.4}/mwin/__init__.py +14 -15
  4. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/schemas/request/log_request.py +2 -4
  5. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/sync_client.py +0 -8
  6. {mwin-0.2.2 → mwin-0.2.4}/mwin/context/runner.py +2 -2
  7. mwin-0.2.4/mwin/eval/__init__.py +26 -0
  8. mwin-0.2.4/mwin/eval/llm_as_judge.py +148 -0
  9. mwin-0.2.4/mwin/eval/metric.py +187 -0
  10. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/args_helper.py +4 -4
  11. mwin-0.2.4/mwin/helper/prompt_helper.py +26 -0
  12. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/serialize_helper.py +71 -71
  13. mwin-0.2.4/mwin/logger.py +39 -0
  14. {mwin-0.2.2 → mwin-0.2.4}/mwin/models/__init__.py +1 -2
  15. {mwin-0.2.2 → mwin-0.2.4}/mwin/models/common.py +2 -0
  16. {mwin-0.2.2 → mwin-0.2.4}/mwin/models/key_models.py +2 -9
  17. {mwin-0.2.2 → mwin-0.2.4}/mwin/patches/openai/async_completions.py +21 -19
  18. {mwin-0.2.2 → mwin-0.2.4}/mwin/patches/openai/completions.py +5 -12
  19. {mwin-0.2.2 → mwin-0.2.4}/mwin/track/base.py +14 -14
  20. {mwin-0.2.2 → mwin-0.2.4}/mwin/track/options.py +7 -7
  21. {mwin-0.2.2 → mwin-0.2.4}/mwin.egg-info/PKG-INFO +3 -3
  22. {mwin-0.2.2 → mwin-0.2.4}/mwin.egg-info/SOURCES.txt +4 -2
  23. {mwin-0.2.2 → mwin-0.2.4}/pyproject.toml +1 -1
  24. mwin-0.2.2/mwin/logger.py +0 -106
  25. mwin-0.2.2/mwin/prompt/__init__.py +0 -6
  26. mwin-0.2.2/mwin/prompt/wrapper.py +0 -112
  27. {mwin-0.2.2 → mwin-0.2.4}/LICENSE +0 -0
  28. {mwin-0.2.2 → mwin-0.2.4}/mwin/_client.py +0 -0
  29. {mwin-0.2.2 → mwin-0.2.4}/mwin/_exception.py +0 -0
  30. {mwin-0.2.2 → mwin-0.2.4}/mwin/cli.py +0 -0
  31. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/__init__.py +0 -0
  32. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/config.py +0 -0
  33. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/schemas/__init__.py +0 -0
  34. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/schemas/request/__init__.py +0 -0
  35. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/schemas/response/__init__.py +0 -0
  36. {mwin-0.2.2 → mwin-0.2.4}/mwin/client/schemas/response/log_response.py +0 -0
  37. {mwin-0.2.2 → mwin-0.2.4}/mwin/config/__init__.py +0 -0
  38. {mwin-0.2.2 → mwin-0.2.4}/mwin/config/configurator.py +0 -0
  39. {mwin-0.2.2 → mwin-0.2.4}/mwin/config/loader.py +0 -0
  40. {mwin-0.2.2 → mwin-0.2.4}/mwin/config/types.py +0 -0
  41. {mwin-0.2.2 → mwin-0.2.4}/mwin/context/__init__.py +0 -0
  42. {mwin-0.2.2 → mwin-0.2.4}/mwin/context/func_context.py +0 -0
  43. {mwin-0.2.2 → mwin-0.2.4}/mwin/context/storage.py +0 -0
  44. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/__init__.py +0 -0
  45. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/cli_helper.py +0 -0
  46. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/exception_helper.py +0 -0
  47. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/id_helper.py +0 -0
  48. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/inspect_helper.py +0 -0
  49. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/llm/__init__.py +0 -0
  50. {mwin-0.2.2 → mwin-0.2.4}/mwin/helper/llm/openai_helper.py +0 -0
  51. {mwin-0.2.2 → mwin-0.2.4}/mwin/patches/llm_patch_config.py +0 -0
  52. {mwin-0.2.2 → mwin-0.2.4}/mwin/patches/std.py +0 -0
  53. {mwin-0.2.2 → mwin-0.2.4}/mwin/track/__init__.py +0 -0
  54. {mwin-0.2.2 → mwin-0.2.4}/mwin/track/_types.py +0 -0
  55. {mwin-0.2.2 → mwin-0.2.4}/mwin/track/at_track.py +0 -0
  56. {mwin-0.2.2 → mwin-0.2.4}/mwin.egg-info/dependency_links.txt +0 -0
  57. {mwin-0.2.2 → mwin-0.2.4}/mwin.egg-info/entry_points.txt +0 -0
  58. {mwin-0.2.2 → mwin-0.2.4}/mwin.egg-info/requires.txt +0 -0
  59. {mwin-0.2.2 → mwin-0.2.4}/mwin.egg-info/top_level.txt +0 -0
  60. {mwin-0.2.2 → mwin-0.2.4}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mwin
3
- Version: 0.2.2
3
+ Version: 0.2.4
4
4
  Summary: Track OpenAI, Claude, Gemini and OpenAI-compatible models then give solutions to improve your agent system.
5
5
  Author-email: yanghui <dasss90ovo@gmail.com>
6
6
  Project-URL: Homepage, https://github.com/yanghui1-arch/mwin.git
@@ -73,14 +73,14 @@ It's the most recommended method to use mwin to track the trace in a project. mw
73
73
 
74
74
  ### Demo
75
75
  ```python
76
- from mwin import track, StepType, start_trace
76
+ from mwin import track, start_trace
77
77
  from openai import OpenAI
78
78
 
79
79
  openai_apikey = "<llm_api_key>"
80
80
  openai_base_url = "<llm_base_url>"
81
81
  model = "<llm_model>"
82
82
 
83
- @track(step_type=StepType.TOOL)
83
+ @track(step_type="tool")
84
84
  def execute_bash(command: str):
85
85
  # assume execute bash and get a stdout
86
86
  return "<bash_stdout>"
@@ -56,14 +56,14 @@ It's the most recommended method to use mwin to track the trace in a project. mw
56
56
 
57
57
  ### Demo
58
58
  ```python
59
- from mwin import track, StepType, start_trace
59
+ from mwin import track, start_trace
60
60
  from openai import OpenAI
61
61
 
62
62
  openai_apikey = "<llm_api_key>"
63
63
  openai_base_url = "<llm_base_url>"
64
64
  model = "<llm_model>"
65
65
 
66
- @track(step_type=StepType.TOOL)
66
+ @track(step_type="tool")
67
67
  def execute_bash(command: str):
68
68
  # assume execute bash and get a stdout
69
69
  return "<bash_stdout>"
@@ -1,16 +1,15 @@
1
- from .track import track
2
- from .models import LLMProvider, StepType, Step, Trace, Conversation
3
- from .context import start_trace, start_trace_async
4
- from .prompt.wrapper import template_prompt
5
-
6
- __all__ = [
7
- "track",
8
- "start_trace",
9
- "start_trace_async",
10
- "template_prompt",
11
- "LLMProvider",
12
- "StepType",
13
- "Step",
14
- "Trace",
15
- "Conversation",
1
+ from .track import track
2
+ from .models import LLMProvider, Step, Trace, Conversation
3
+ from .context import start_trace, start_trace_async
4
+ from . import eval
5
+
6
+ __all__ = [
7
+ "track",
8
+ "start_trace",
9
+ "start_trace_async",
10
+ "LLMProvider",
11
+ "Step",
12
+ "Trace",
13
+ "Conversation",
14
+ "eval",
16
15
  ]
@@ -62,9 +62,6 @@ class LogStepRequest(BaseModel):
62
62
  prompt_version: str | None = None
63
63
  """Prompt version"""
64
64
 
65
- system_prompt: str | None = None
66
- """Raw system prompt string captured via template_prompt() inside the step"""
67
-
68
65
  @field_serializer('input', 'output')
69
66
  def serialize_any_field(self, value: Any):
70
67
  return serialize_helper.safe_serialize(value)
@@ -93,4 +90,5 @@ class LogTraceRequest(BaseModel):
93
90
 
94
91
  @field_serializer("start_time", "last_update_timestamp")
95
92
  def serialize_datetime(self, value: datetime):
96
- return value.strftime("%Y-%m-%d %H:%M:%S.%f")
93
+ return value.strftime("%Y-%m-%d %H:%M:%S.%f")
94
+
@@ -58,7 +58,6 @@ class SyncClient:
58
58
  llm_provider: LLMProvider,
59
59
  pipeline: str | None,
60
60
  prompt_name: str | None,
61
- system_prompt: str | None,
62
61
  prompt_version: str | None,
63
62
  ) -> LogStepResponse:
64
63
  """Create a step and log it in server."""
@@ -68,12 +67,6 @@ class SyncClient:
68
67
  parent_step_id = None
69
68
 
70
69
  llm_provider = llm_provider.value
71
- # prompt_version not set -> don't pass system_prompt.
72
- # TODO: Give a warn
73
- if prompt_version is None:
74
- pipeline = None
75
- system_prompt = None
76
- prompt_name = None
77
70
 
78
71
  log_step_req = LogStepRequest(
79
72
  project_name=self._project_name,
@@ -94,7 +87,6 @@ class SyncClient:
94
87
  llm_provider=llm_provider,
95
88
  prompt_pipeline=pipeline,
96
89
  prompt_name=prompt_name,
97
- system_prompt=system_prompt,
98
90
  prompt_version=prompt_version,
99
91
  )
100
92
 
@@ -22,13 +22,13 @@ It works well at beginning of project but not recommend it for developer. it's v
22
22
  Please use following demo to develop with mwin. It's more safe, easier and more controllable.
23
23
 
24
24
  ```python
25
- from mwin import track, StepType
25
+ from mwin import track
26
26
  from mwin.runner import start_trace
27
27
  @track()
28
28
  def call_openai() -> Dict:
29
29
  ...
30
30
 
31
- @track(StepType.TOOL)
31
+ @track(step_type="tool")
32
32
  def execute_bash(*args, **kwargs)
33
33
  ...
34
34
 
@@ -0,0 +1,26 @@
1
+ from .llm_as_judge import JudgeResult, llm_as_judge
2
+ from .metric import (
3
+ accuracy,
4
+ evaluate,
5
+ f1,
6
+ get_metric,
7
+ list_metrics,
8
+ recall,
9
+ register_metric,
10
+ unregister_metric,
11
+ )
12
+ from . import metric
13
+
14
+ __all__ = [
15
+ "llm_as_judge",
16
+ "JudgeResult",
17
+ "accuracy",
18
+ "f1",
19
+ "recall",
20
+ "evaluate",
21
+ "register_metric",
22
+ "unregister_metric",
23
+ "get_metric",
24
+ "list_metrics",
25
+ "metric",
26
+ ]
@@ -0,0 +1,148 @@
1
+ """Simple LLM-as-judge utility."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import re
7
+ from dataclasses import dataclass
8
+ from typing import Any
9
+
10
+ from openai import OpenAI
11
+
12
+ from ..helper import serialize_helper
13
+ from ..models.common import LLMProvider
14
+ from ..track import track
15
+
16
+
17
+ @dataclass(frozen=True)
18
+ class JudgeResult:
19
+ score: float
20
+ reasoning: str
21
+ raw_response: str
22
+
23
+
24
+ def llm_as_judge(
25
+ *,
26
+ input_data: Any,
27
+ output_data: Any,
28
+ metric_name: str,
29
+ judge_prompt: str,
30
+ model: str = "gpt-4o-mini",
31
+ client: Any | None = None,
32
+ api_key: str | None = None,
33
+ base_url: str | None = None,
34
+ score_range_min: float = 0.0,
35
+ score_range_max: float = 1.0,
36
+ track_judge_call: bool = True,
37
+ judge_project_name: str = "mwin_judge",
38
+ llm_provider: LLMProvider = LLMProvider.OPENAI,
39
+ ) -> JudgeResult:
40
+ """Evaluate an output with an LLM judge and return score/reasoning."""
41
+ if len(metric_name.strip()) == 0:
42
+ raise ValueError("metric_name must be non-empty.")
43
+ if len(judge_prompt.strip()) == 0:
44
+ raise ValueError("judge_prompt must be non-empty.")
45
+ if score_range_min >= score_range_max:
46
+ raise ValueError(
47
+ "score_range_min must be smaller than score_range_max. "
48
+ f"Got {score_range_min} and {score_range_max}."
49
+ )
50
+
51
+ judge_client = client or OpenAI(api_key=api_key, base_url=base_url)
52
+
53
+ prompt = (
54
+ "You are evaluating an LLM output.\n\n"
55
+ f"Evaluation criterion: {metric_name}\n"
56
+ "Scoring guide:\n"
57
+ f"{judge_prompt}\n\n"
58
+ f"Score range: {score_range_min} (worst) to {score_range_max} (best)\n\n"
59
+ "=== Input ===\n"
60
+ f"{_dump_text(input_data)}\n\n"
61
+ "=== Output ===\n"
62
+ f"{_dump_text(output_data)}\n\n"
63
+ "Reply with ONLY valid JSON, no markdown fences:\n"
64
+ '{"score": <number within range>, "reasoning": "<one concise sentence>"}'
65
+ )
66
+
67
+ def _judge_call() -> Any:
68
+ return judge_client.chat.completions.create(
69
+ model=model,
70
+ messages=[{"role": "user", "content": prompt}],
71
+ temperature=0,
72
+ max_tokens=256,
73
+ )
74
+
75
+ if track_judge_call:
76
+ response = track(
77
+ project_name=judge_project_name,
78
+ tags=["eval", "llm_as_judge"],
79
+ step_type="llm",
80
+ model=model,
81
+ llm_provider=llm_provider,
82
+ description="mwin.eval.llm_as_judge",
83
+ )(_judge_call)()
84
+ else:
85
+ response = _judge_call()
86
+
87
+ raw_text = _extract_response_text(response)
88
+ payload = _parse_payload(raw_text)
89
+
90
+ if "score" not in payload:
91
+ raise ValueError("Judge response JSON must include 'score'.")
92
+ if "reasoning" not in payload:
93
+ raise ValueError("Judge response JSON must include 'reasoning'.")
94
+
95
+ score = float(payload["score"])
96
+ if score < score_range_min or score > score_range_max:
97
+ raise ValueError(
98
+ f"Judge score out of range: {score}. "
99
+ f"Expected within [{score_range_min}, {score_range_max}]."
100
+ )
101
+
102
+ return JudgeResult(
103
+ score=score,
104
+ reasoning=str(payload["reasoning"]).strip(),
105
+ raw_response=raw_text,
106
+ )
107
+
108
+
109
+ def _dump_text(value: Any) -> str:
110
+ serialized = serialize_helper.safe_serialize(value)
111
+ if isinstance(serialized, str):
112
+ return serialized
113
+ return json.dumps(serialized, ensure_ascii=False, indent=2)
114
+
115
+
116
+ def _extract_response_text(response: Any) -> str:
117
+ try:
118
+ content = response.choices[0].message.content
119
+ except Exception as exc:
120
+ raise ValueError("Unable to read judge response content.") from exc
121
+
122
+ if content is None:
123
+ raise ValueError("Judge response content is empty.")
124
+ return str(content).strip()
125
+
126
+
127
+ def _parse_payload(raw_text: str) -> dict[str, Any]:
128
+ text = raw_text.strip()
129
+
130
+ fenced_match = re.search(r"```(?:json)?\\s*(.*?)\\s*```", text, flags=re.DOTALL | re.IGNORECASE)
131
+ if fenced_match:
132
+ text = fenced_match.group(1).strip()
133
+
134
+ try:
135
+ parsed = json.loads(text)
136
+ if isinstance(parsed, dict):
137
+ return parsed
138
+ except json.JSONDecodeError:
139
+ pass
140
+
141
+ json_object_match = re.search(r"\{.*\}", text, flags=re.DOTALL)
142
+ if not json_object_match:
143
+ raise ValueError(f"Judge response is not valid JSON: {raw_text}")
144
+
145
+ parsed = json.loads(json_object_match.group(0))
146
+ if not isinstance(parsed, dict):
147
+ raise ValueError("Judge response JSON must be an object.")
148
+ return parsed
@@ -0,0 +1,187 @@
1
+ """Simple evaluation metrics for mwin."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any, Callable, Iterable, Mapping, Sequence
6
+
7
+ MetricFunc = Callable[[Sequence[Any], Sequence[Any]], float]
8
+
9
+
10
+ def _to_list(values: Sequence[Any] | Iterable[Any], *, field_name: str) -> list[Any]:
11
+ if isinstance(values, (str, bytes, bytearray)):
12
+ raise TypeError(f"{field_name} must be a sequence of labels, not string-like data.")
13
+
14
+ try:
15
+ return list(values)
16
+ except TypeError as exc:
17
+ raise TypeError(f"{field_name} must be iterable.") from exc
18
+
19
+
20
+ def _validate_pair_inputs(
21
+ y_true: Sequence[Any] | Iterable[Any],
22
+ y_pred: Sequence[Any] | Iterable[Any],
23
+ ) -> tuple[list[Any], list[Any]]:
24
+ true_list = _to_list(y_true, field_name="y_true")
25
+ pred_list = _to_list(y_pred, field_name="y_pred")
26
+
27
+ if len(true_list) != len(pred_list):
28
+ raise ValueError(
29
+ f"y_true and y_pred must have the same length. "
30
+ f"Got {len(true_list)} and {len(pred_list)}."
31
+ )
32
+ if len(true_list) == 0:
33
+ raise ValueError("y_true and y_pred must not be empty.")
34
+
35
+ return true_list, pred_list
36
+
37
+
38
+ def accuracy(
39
+ y_true: Sequence[Any] | Iterable[Any],
40
+ y_pred: Sequence[Any] | Iterable[Any],
41
+ ) -> float:
42
+ """Accuracy = correct / total."""
43
+ true_list, pred_list = _validate_pair_inputs(y_true, y_pred)
44
+ correct_count = sum(1 for target, prediction in zip(true_list, pred_list) if target == prediction)
45
+ return correct_count / len(true_list)
46
+
47
+
48
+ def recall(
49
+ y_true: Sequence[Any] | Iterable[Any],
50
+ y_pred: Sequence[Any] | Iterable[Any],
51
+ *,
52
+ positive_label: Any = 1,
53
+ ) -> float:
54
+ """Binary recall for a positive label."""
55
+ true_list, pred_list = _validate_pair_inputs(y_true, y_pred)
56
+
57
+ true_positive = 0
58
+ false_negative = 0
59
+ for target, prediction in zip(true_list, pred_list):
60
+ if target == positive_label and prediction == positive_label:
61
+ true_positive += 1
62
+ elif target == positive_label and prediction != positive_label:
63
+ false_negative += 1
64
+
65
+ denominator = true_positive + false_negative
66
+ return (true_positive / denominator) if denominator > 0 else 0.0
67
+
68
+
69
+ def f1(
70
+ y_true: Sequence[Any] | Iterable[Any],
71
+ y_pred: Sequence[Any] | Iterable[Any],
72
+ *,
73
+ positive_label: Any = 1,
74
+ ) -> float:
75
+ """Binary F1 for a positive label."""
76
+ true_list, pred_list = _validate_pair_inputs(y_true, y_pred)
77
+
78
+ true_positive = 0
79
+ false_positive = 0
80
+ false_negative = 0
81
+
82
+ for target, prediction in zip(true_list, pred_list):
83
+ if target == positive_label and prediction == positive_label:
84
+ true_positive += 1
85
+ elif target != positive_label and prediction == positive_label:
86
+ false_positive += 1
87
+ elif target == positive_label and prediction != positive_label:
88
+ false_negative += 1
89
+
90
+ precision_denominator = true_positive + false_positive
91
+ recall_denominator = true_positive + false_negative
92
+ precision_value = (true_positive / precision_denominator) if precision_denominator > 0 else 0.0
93
+ recall_value = (true_positive / recall_denominator) if recall_denominator > 0 else 0.0
94
+
95
+ denominator = precision_value + recall_value
96
+ return (2 * precision_value * recall_value / denominator) if denominator > 0 else 0.0
97
+
98
+
99
+ _BUILTIN_METRICS: dict[str, MetricFunc] = {
100
+ "accuracy": accuracy,
101
+ "f1": f1,
102
+ "recall": recall,
103
+ }
104
+
105
+ _CUSTOM_METRICS: dict[str, MetricFunc] = {}
106
+
107
+
108
+ def register_metric(name: str, metric_fn: MetricFunc, *, overwrite: bool = False) -> None:
109
+ """Register a custom metric globally."""
110
+ if not isinstance(name, str) or len(name.strip()) == 0:
111
+ raise ValueError("Metric name must be a non-empty string.")
112
+ if not callable(metric_fn):
113
+ raise TypeError("metric_fn must be callable.")
114
+
115
+ normalized_name = name.strip().lower()
116
+ if not overwrite and normalized_name in _CUSTOM_METRICS:
117
+ raise ValueError(
118
+ f"Metric '{normalized_name}' already exists. "
119
+ "Set overwrite=True to replace it."
120
+ )
121
+ _CUSTOM_METRICS[normalized_name] = metric_fn
122
+
123
+
124
+ def unregister_metric(name: str) -> None:
125
+ """Remove a custom metric from global registry."""
126
+ _CUSTOM_METRICS.pop(name.strip().lower(), None)
127
+
128
+
129
+ def get_metric(name: str) -> MetricFunc:
130
+ """Get metric by name (custom has higher priority than built-in)."""
131
+ normalized_name = name.strip().lower()
132
+
133
+ if normalized_name in _CUSTOM_METRICS:
134
+ return _CUSTOM_METRICS[normalized_name]
135
+ if normalized_name in _BUILTIN_METRICS:
136
+ return _BUILTIN_METRICS[normalized_name]
137
+
138
+ raise KeyError(f"Metric '{normalized_name}' does not exist.")
139
+
140
+
141
+ def list_metrics() -> list[str]:
142
+ """List all metric names."""
143
+ names = set(_BUILTIN_METRICS.keys())
144
+ names.update(_CUSTOM_METRICS.keys())
145
+ return sorted(names)
146
+
147
+
148
+ def evaluate(
149
+ y_true: Sequence[Any] | Iterable[Any],
150
+ y_pred: Sequence[Any] | Iterable[Any],
151
+ metrics: Sequence[str | MetricFunc] | None = None,
152
+ *,
153
+ custom_metrics: Mapping[str, MetricFunc] | None = None,
154
+ ) -> dict[str, float]:
155
+ """Evaluate one or more metrics.
156
+
157
+ Defaults to ["accuracy", "f1", "recall"].
158
+ """
159
+ true_list, pred_list = _validate_pair_inputs(y_true, y_pred)
160
+ targets = list(metrics) if metrics is not None else ["accuracy", "f1", "recall"]
161
+
162
+ merged_custom: dict[str, MetricFunc] = dict(_CUSTOM_METRICS)
163
+ if custom_metrics is not None:
164
+ for name, metric_fn in custom_metrics.items():
165
+ if not callable(metric_fn):
166
+ raise TypeError(f"custom_metrics['{name}'] must be callable.")
167
+ merged_custom[name.strip().lower()] = metric_fn
168
+
169
+ results: dict[str, float] = {}
170
+ for index, target in enumerate(targets):
171
+ if isinstance(target, str):
172
+ normalized_name = target.strip().lower()
173
+ metric_fn = merged_custom.get(normalized_name) or _BUILTIN_METRICS.get(normalized_name)
174
+ if metric_fn is None:
175
+ raise KeyError(f"Metric '{normalized_name}' does not exist.")
176
+ metric_name = normalized_name
177
+ elif callable(target):
178
+ metric_fn = target
179
+ metric_name = getattr(target, "__name__", f"custom_metric_{index}")
180
+ else:
181
+ raise TypeError(
182
+ "Each metric must be either a metric name string or a callable metric function."
183
+ )
184
+
185
+ results[metric_name] = float(metric_fn(true_list, pred_list))
186
+
187
+ return results
@@ -1,10 +1,10 @@
1
1
  from uuid import UUID
2
- from typing import Any, List, Dict
2
+ from typing import Any, List, Dict, Literal
3
3
 
4
4
  from pydantic import BaseModel
5
5
  from openai.types.completion_usage import CompletionUsage
6
6
  from . import id_helper
7
- from ..models import Step, Trace, StepType
7
+ from ..models import Step, Trace
8
8
  from .. import context
9
9
 
10
10
  class StartArguments(BaseModel):
@@ -27,7 +27,7 @@ def create_new_step(
27
27
  input: Any | None = None,
28
28
  output: Dict[str, Any] | None = None,
29
29
  name: str | None = None,
30
- type: StepType = StepType.GENERAL,
30
+ type: Literal["general", "llm", "retrieve", "tool"] = "general",
31
31
  tags: List[str] | None = None,
32
32
  model: str | None = None,
33
33
  usage: int | None = None,
@@ -45,7 +45,7 @@ def create_new_step(
45
45
  output(Dict[str, Any] | None): output of module. Default to `None`. None means it's logging input.
46
46
  name(str | None): the step name. Caller can set the name to define what the step role is. Default to ``None`. If it's None,
47
47
  AITrace will set step name based on step type.
48
- type(StepType): step type. Default to `StepType.GENERAL`.
48
+ type(Literal["general", "llm", "retrieve", "tool"]): step type. Default to `general`.
49
49
  tags(List[str] | None): step tags. Default to `None`. If it's None, it will be set an empty list.
50
50
  model(str | None): model name. Probably using a llm model in the step. Default to `None`.
51
51
  usage(int | None): llm token usage. Default to `None`.
@@ -0,0 +1,26 @@
1
+ from typing import Tuple
2
+
3
+
4
+ def parse_system_prompt_identifier(system_prompt: str | None) -> Tuple[str | None, str | None, str | None]:
5
+ """Parse system prompt identifier in `pipeline/name@version` format.
6
+
7
+ Returns `(pipeline, prompt_name, prompt_version)`.
8
+ Returns `(None, None, None)` when `system_prompt` is None.
9
+ """
10
+
11
+ if system_prompt is None:
12
+ return None, None, None
13
+
14
+ identifier = system_prompt.strip()
15
+ if not identifier or identifier.count("@") != 1:
16
+ raise ValueError("Invalid system_prompt. Expected format: pipeline/name@version")
17
+
18
+ path, prompt_version = identifier.rsplit("@", 1)
19
+ if not prompt_version or "/" not in path:
20
+ raise ValueError("Invalid system_prompt. Expected format: pipeline/name@version")
21
+
22
+ pipeline, prompt_name = path.split("/", 1)
23
+ if not pipeline or not prompt_name:
24
+ raise ValueError("Invalid system_prompt. Expected format: pipeline/name@version")
25
+
26
+ return pipeline, prompt_name, prompt_version