memoturn 0.1.0__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.
@@ -0,0 +1,34 @@
1
+ node_modules/
2
+ dist/
3
+ .next/
4
+ .turbo/
5
+ coverage/
6
+ *.tsbuildinfo
7
+
8
+ # env
9
+ .env
10
+ .env.local
11
+ .env.*.local
12
+
13
+ # os / editor
14
+ .DS_Store
15
+ *.log
16
+
17
+ # claude code — personal/local settings (team config is .claude/settings.json)
18
+ .claude/settings.local.json
19
+
20
+ # prisma
21
+ packages/db/prisma/*.db
22
+ # generated by the TanStack Router plugin
23
+ **/routeTree.gen.ts
24
+
25
+ # python
26
+ __pycache__/
27
+ *.egg-info/
28
+ .venv/
29
+ sdks/python/dist/
30
+
31
+ # playwright (e2e)
32
+ **/playwright-report/
33
+ **/test-results/
34
+ **/.playwright/
@@ -0,0 +1,82 @@
1
+ Metadata-Version: 2.4
2
+ Name: memoturn
3
+ Version: 0.1.0
4
+ Summary: memoturn Python SDK — tracing, @observe, OpenAI wrapper, prompts.
5
+ Project-URL: Homepage, https://github.com/memoturn/memoturn
6
+ Project-URL: Repository, https://github.com/memoturn/memoturn
7
+ Project-URL: Issues, https://github.com/memoturn/memoturn/issues
8
+ License: Apache-2.0
9
+ Keywords: evals,langchain,llm,memoturn,observability,openai,opentelemetry,tracing
10
+ Classifier: Development Status :: 4 - Beta
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: License :: OSI Approved :: Apache Software License
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
15
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
16
+ Requires-Python: >=3.9
17
+ Provides-Extra: dev
18
+ Requires-Dist: pytest>=8; extra == 'dev'
19
+ Provides-Extra: openai
20
+ Requires-Dist: openai>=1.0; extra == 'openai'
21
+ Description-Content-Type: text/markdown
22
+
23
+ # memoturn Python SDK
24
+
25
+ Tracing, prompts, and the OpenAI wrapper for [memoturn](https://github.com/memoturn/memoturn).
26
+ Stdlib-only (no required dependencies).
27
+
28
+ ```bash
29
+ pip install memoturn # or: uv add memoturn
30
+ ```
31
+
32
+ ## Trace with the decorator
33
+
34
+ ```python
35
+ from memoturn import Memoturn, configure, observe
36
+
37
+ configure(Memoturn(base_url="http://localhost:3001", public_key="pk-...", secret_key="sk-..."))
38
+
39
+ @observe()
40
+ def retrieve(q): ...
41
+
42
+ @observe(as_type="generation")
43
+ def answer(q, docs): ...
44
+
45
+ @observe(name="rag-pipeline")
46
+ def rag(q):
47
+ return answer(q, retrieve(q)) # nested spans under one trace
48
+ ```
49
+
50
+ The outermost `@observe` opens a trace; nested calls become child spans (a waterfall in
51
+ the console). Env vars `MEMOTURN_BASE_URL` / `MEMOTURN_PUBLIC_KEY` / `MEMOTURN_SECRET_KEY`
52
+ are used when not passed explicitly.
53
+
54
+ ## Low-level API
55
+
56
+ ```python
57
+ mt = Memoturn()
58
+ trace = mt.trace(name="chat", user_id="u1")
59
+ gen = trace.generation(name="answer", model="claude-sonnet-4-6", input=messages)
60
+ gen.end(output=reply, usage={"promptTokens": 100, "completionTokens": 20})
61
+ trace.score("user-feedback", value=1, comment="helpful")
62
+ mt.shutdown() # flush
63
+ ```
64
+
65
+ ## OpenAI wrapper
66
+
67
+ ```python
68
+ from openai import OpenAI
69
+ from memoturn import wrap_openai
70
+
71
+ client = wrap_openai(OpenAI())
72
+ client.chat.completions.create(model="gpt-4o-mini", messages=[...]) # recorded automatically
73
+ ```
74
+
75
+ ## Prompts
76
+
77
+ ```python
78
+ from memoturn import get_prompt, compile_prompt
79
+
80
+ prompt = get_prompt("support-reply", channel="production")
81
+ messages = compile_prompt(prompt, product="memoturn", question="How do I trace a call?")
82
+ ```
@@ -0,0 +1,60 @@
1
+ # memoturn Python SDK
2
+
3
+ Tracing, prompts, and the OpenAI wrapper for [memoturn](https://github.com/memoturn/memoturn).
4
+ Stdlib-only (no required dependencies).
5
+
6
+ ```bash
7
+ pip install memoturn # or: uv add memoturn
8
+ ```
9
+
10
+ ## Trace with the decorator
11
+
12
+ ```python
13
+ from memoturn import Memoturn, configure, observe
14
+
15
+ configure(Memoturn(base_url="http://localhost:3001", public_key="pk-...", secret_key="sk-..."))
16
+
17
+ @observe()
18
+ def retrieve(q): ...
19
+
20
+ @observe(as_type="generation")
21
+ def answer(q, docs): ...
22
+
23
+ @observe(name="rag-pipeline")
24
+ def rag(q):
25
+ return answer(q, retrieve(q)) # nested spans under one trace
26
+ ```
27
+
28
+ The outermost `@observe` opens a trace; nested calls become child spans (a waterfall in
29
+ the console). Env vars `MEMOTURN_BASE_URL` / `MEMOTURN_PUBLIC_KEY` / `MEMOTURN_SECRET_KEY`
30
+ are used when not passed explicitly.
31
+
32
+ ## Low-level API
33
+
34
+ ```python
35
+ mt = Memoturn()
36
+ trace = mt.trace(name="chat", user_id="u1")
37
+ gen = trace.generation(name="answer", model="claude-sonnet-4-6", input=messages)
38
+ gen.end(output=reply, usage={"promptTokens": 100, "completionTokens": 20})
39
+ trace.score("user-feedback", value=1, comment="helpful")
40
+ mt.shutdown() # flush
41
+ ```
42
+
43
+ ## OpenAI wrapper
44
+
45
+ ```python
46
+ from openai import OpenAI
47
+ from memoturn import wrap_openai
48
+
49
+ client = wrap_openai(OpenAI())
50
+ client.chat.completions.create(model="gpt-4o-mini", messages=[...]) # recorded automatically
51
+ ```
52
+
53
+ ## Prompts
54
+
55
+ ```python
56
+ from memoturn import get_prompt, compile_prompt
57
+
58
+ prompt = get_prompt("support-reply", channel="production")
59
+ messages = compile_prompt(prompt, product="memoturn", question="How do I trace a call?")
60
+ ```
@@ -0,0 +1,44 @@
1
+ """Emit a trace from Python end-to-end. Run after `bun run dev` + `bun run seed`:
2
+
3
+ cd sdks/python && uv run examples/quickstart.py
4
+ """
5
+ import os
6
+ import time
7
+
8
+ from memoturn import Memoturn, configure, observe
9
+
10
+ configure(
11
+ Memoturn(
12
+ base_url=os.environ.get("MEMOTURN_BASE_URL", "http://localhost:3001"),
13
+ public_key=os.environ.get("MEMOTURN_PUBLIC_KEY", "pk-mt-dev"),
14
+ secret_key=os.environ.get("MEMOTURN_SECRET_KEY", "sk-mt-dev"),
15
+ )
16
+ )
17
+
18
+
19
+ @observe()
20
+ def retrieve(query: str) -> list[str]:
21
+ time.sleep(0.02)
22
+ return ["memoturn is an open-source AI engineering platform."]
23
+
24
+
25
+ @observe(as_type="generation")
26
+ def answer(question: str, docs: list[str]) -> str:
27
+ time.sleep(0.05)
28
+ return "memoturn is an open-source AI engineering platform."
29
+
30
+
31
+ @observe(name="rag-pipeline")
32
+ def rag(question: str) -> str:
33
+ docs = retrieve(question)
34
+ return answer(question, docs)
35
+
36
+
37
+ if __name__ == "__main__":
38
+ from memoturn import get_client
39
+
40
+ result = rag("What is memoturn?")
41
+ get_client().shutdown() # flush
42
+ print("answer:", result)
43
+ print("emitted a nested trace (rag-pipeline → retrieve, answer)")
44
+ print(" open the console Traces view to see the waterfall")
@@ -0,0 +1,37 @@
1
+ [project]
2
+ name = "memoturn"
3
+ version = "0.1.0"
4
+ description = "memoturn Python SDK — tracing, @observe, OpenAI wrapper, prompts."
5
+ readme = "README.md"
6
+ requires-python = ">=3.9"
7
+ license = { text = "Apache-2.0" }
8
+ keywords = ["memoturn", "llm", "observability", "tracing", "openai", "langchain", "evals", "opentelemetry"]
9
+ dependencies = []
10
+ classifiers = [
11
+ "Development Status :: 4 - Beta",
12
+ "Intended Audience :: Developers",
13
+ "License :: OSI Approved :: Apache Software License",
14
+ "Programming Language :: Python :: 3",
15
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
16
+ "Topic :: Software Development :: Libraries :: Python Modules",
17
+ ]
18
+
19
+ [project.urls]
20
+ Homepage = "https://github.com/memoturn/memoturn"
21
+ Repository = "https://github.com/memoturn/memoturn"
22
+ Issues = "https://github.com/memoturn/memoturn/issues"
23
+
24
+ [project.optional-dependencies]
25
+ openai = ["openai>=1.0"]
26
+ dev = ["pytest>=8"]
27
+
28
+ [build-system]
29
+ requires = ["hatchling"]
30
+ build-backend = "hatchling.build"
31
+
32
+ [tool.hatch.build.targets.wheel]
33
+ packages = ["src/memoturn"]
34
+
35
+ [tool.pytest.ini_options]
36
+ testpaths = ["tests"]
37
+ pythonpath = ["src"]
@@ -0,0 +1,19 @@
1
+ """memoturn Python SDK — LLM observability, prompts, and evals."""
2
+ from .client import Memoturn, Span, Trace
3
+ from .decorator import configure, get_client, observe
4
+ from .openai import wrap_openai
5
+ from .prompt import compile_prompt, get_prompt
6
+
7
+ __all__ = [
8
+ "Memoturn",
9
+ "Trace",
10
+ "Span",
11
+ "observe",
12
+ "configure",
13
+ "get_client",
14
+ "get_prompt",
15
+ "compile_prompt",
16
+ "wrap_openai",
17
+ ]
18
+
19
+ __version__ = "0.1.0"
@@ -0,0 +1,152 @@
1
+ """memoturn client — batches events and flushes to POST /v1/ingest.
2
+
3
+ Stdlib-only (urllib). Create trace/span/generation handles and call .end() as work
4
+ completes; the client handles ids, timestamps, batching, and Basic auth.
5
+ """
6
+ from __future__ import annotations
7
+
8
+ import atexit
9
+ import base64
10
+ import datetime as _dt
11
+ import json
12
+ import os
13
+ import threading
14
+ import urllib.error
15
+ import urllib.request
16
+ import uuid
17
+ from typing import Any, Optional
18
+
19
+
20
+ def _id() -> str:
21
+ return str(uuid.uuid4())
22
+
23
+
24
+ def _now() -> str:
25
+ # ISO-8601 with millisecond precision + 'Z', matching the JS SDK.
26
+ return _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")
27
+
28
+
29
+ class Memoturn:
30
+ def __init__(
31
+ self,
32
+ base_url: Optional[str] = None,
33
+ public_key: Optional[str] = None,
34
+ secret_key: Optional[str] = None,
35
+ environment: Optional[str] = None,
36
+ flush_at: int = 20,
37
+ ) -> None:
38
+ self.base_url = (base_url or os.environ.get("MEMOTURN_BASE_URL", "http://localhost:3001")).rstrip("/")
39
+ self.public_key = public_key or os.environ.get("MEMOTURN_PUBLIC_KEY", "")
40
+ self.secret_key = secret_key or os.environ.get("MEMOTURN_SECRET_KEY", "")
41
+ self.environment = environment or os.environ.get("MEMOTURN_ENVIRONMENT", "default")
42
+ self.flush_at = flush_at
43
+ self._buffer: list[dict[str, Any]] = []
44
+ self._lock = threading.Lock()
45
+ atexit.register(self.flush)
46
+
47
+ # ── public API ────────────────────────────────────────────────────────────
48
+ def trace(self, **body: Any) -> "Trace":
49
+ tid = body.pop("id", None) or _id()
50
+ body.setdefault("environment", self.environment)
51
+ self._enqueue("trace-create", {**body, "id": tid})
52
+ return Trace(self, tid)
53
+
54
+ def flush(self) -> None:
55
+ with self._lock:
56
+ batch, self._buffer = self._buffer, []
57
+ if not batch:
58
+ return
59
+ auth = base64.b64encode(f"{self.public_key}:{self.secret_key}".encode()).decode()
60
+ req = urllib.request.Request(
61
+ f"{self.base_url}/v1/ingest",
62
+ data=json.dumps({"batch": batch}, default=str).encode(),
63
+ headers={"content-type": "application/json", "authorization": f"Basic {auth}"},
64
+ method="POST",
65
+ )
66
+ try:
67
+ urllib.request.urlopen(req, timeout=10).read()
68
+ except urllib.error.HTTPError as e:
69
+ if e.code != 207:
70
+ with self._lock:
71
+ self._buffer[0:0] = batch # re-buffer for next flush
72
+ raise
73
+
74
+ def shutdown(self) -> None:
75
+ self.flush()
76
+
77
+ # ── internal ──────────────────────────────────────────────────────────────
78
+ def _enqueue(self, type_: str, body: dict[str, Any]) -> None:
79
+ with self._lock:
80
+ self._buffer.append({"id": _id(), "type": type_, "timestamp": _now(), "body": body})
81
+ should = len(self._buffer) >= self.flush_at
82
+ if should:
83
+ self.flush()
84
+
85
+
86
+ class Trace:
87
+ def __init__(self, client: Memoturn, trace_id: str) -> None:
88
+ self._c = client
89
+ self.id = trace_id
90
+
91
+ def update(self, **body: Any) -> "Trace":
92
+ self._c._enqueue("trace-create", {**body, "id": self.id, "environment": self._c.environment})
93
+ return self
94
+
95
+ def span(self, **body: Any) -> "Span":
96
+ return self._observe("span-create", "span", body)
97
+
98
+ def generation(self, **body: Any) -> "Span":
99
+ return self._observe("generation-create", "generation", body)
100
+
101
+ def event(self, **body: Any) -> None:
102
+ self._c._enqueue(
103
+ "event-create",
104
+ {**body, "id": body.pop("id", None) or _id(), "traceId": self.id,
105
+ "environment": self._c.environment, "startTime": _now()},
106
+ )
107
+
108
+ def score(self, name: str, value: Optional[float] = None, **body: Any) -> "Trace":
109
+ self._c._enqueue(
110
+ "score-create",
111
+ {"id": _id(), "traceId": self.id, "name": name, "value": value,
112
+ "environment": self._c.environment, **body},
113
+ )
114
+ return self
115
+
116
+ def _observe(self, type_: str, kind: str, body: dict[str, Any]) -> "Span":
117
+ oid = body.pop("id", None) or _id()
118
+ self._c._enqueue(
119
+ type_,
120
+ {**body, "id": oid, "traceId": self.id, "environment": self._c.environment, "startTime": _now()},
121
+ )
122
+ return Span(self._c, self.id, oid, kind)
123
+
124
+
125
+ class Span:
126
+ def __init__(self, client: Memoturn, trace_id: str, obs_id: str, kind: str) -> None:
127
+ self._c = client
128
+ self._trace_id = trace_id
129
+ self.id = obs_id
130
+ self._kind = kind
131
+
132
+ def span(self, **body: Any) -> "Span":
133
+ return self._child("span-create", "span", body)
134
+
135
+ def generation(self, **body: Any) -> "Span":
136
+ return self._child("generation-create", "generation", body)
137
+
138
+ def _child(self, type_: str, kind: str, body: dict[str, Any]) -> "Span":
139
+ oid = body.pop("id", None) or _id()
140
+ self._c._enqueue(
141
+ type_,
142
+ {**body, "id": oid, "traceId": self._trace_id, "parentObservationId": self.id,
143
+ "environment": self._c.environment, "startTime": _now()},
144
+ )
145
+ return Span(self._c, self._trace_id, oid, kind)
146
+
147
+ def end(self, **body: Any) -> None:
148
+ type_ = "generation-update" if self._kind == "generation" else "span-update"
149
+ self._c._enqueue(
150
+ type_,
151
+ {**body, "id": self.id, "traceId": self._trace_id, "environment": self._c.environment, "endTime": _now()},
152
+ )
@@ -0,0 +1,87 @@
1
+ """@observe decorator — trace any function with automatic nesting.
2
+
3
+ The outermost @observe creates a trace + root observation; nested @observe calls become
4
+ child spans. Uses a contextvar so nesting works across sync and async call stacks.
5
+ """
6
+ from __future__ import annotations
7
+
8
+ import contextvars
9
+ import functools
10
+ import inspect
11
+ from typing import Any, Callable, Optional
12
+
13
+ from .client import Memoturn, Span, Trace
14
+
15
+ _default: Optional[Memoturn] = None
16
+ _ctx: contextvars.ContextVar[Optional[tuple[Trace, Span]]] = contextvars.ContextVar("memoturn_ctx", default=None)
17
+
18
+
19
+ def configure(client: Memoturn) -> Memoturn:
20
+ """Set the default client used by @observe (and returned by get_client)."""
21
+ global _default
22
+ _default = client
23
+ return client
24
+
25
+
26
+ def get_client() -> Memoturn:
27
+ global _default
28
+ if _default is None:
29
+ _default = Memoturn()
30
+ return _default
31
+
32
+
33
+ def _begin(name: str, as_type: str, inp: Any) -> tuple[Trace, Span, Any]:
34
+ client = get_client()
35
+ cur = _ctx.get()
36
+ if cur is None:
37
+ trace = client.trace(name=name, input=inp)
38
+ span = trace.generation(name=name, input=inp) if as_type == "generation" else trace.span(name=name, input=inp)
39
+ else:
40
+ trace, parent = cur
41
+ span = parent.generation(name=name, input=inp) if as_type == "generation" else parent.span(name=name, input=inp)
42
+ token = _ctx.set((trace, span))
43
+ return trace, span, (token, cur is None)
44
+
45
+
46
+ def _finish(trace: Trace, span: Span, state: Any, output: Any) -> None:
47
+ token, is_root = state
48
+ span.end(output=output)
49
+ if is_root:
50
+ trace.update(output=output)
51
+ _ctx.reset(token)
52
+
53
+
54
+ def observe(_fn: Optional[Callable] = None, *, name: Optional[str] = None, as_type: str = "span") -> Callable:
55
+ def deco(fn: Callable) -> Callable:
56
+ obs_name = name or fn.__name__
57
+
58
+ if inspect.iscoroutinefunction(fn):
59
+ @functools.wraps(fn)
60
+ async def awrapper(*args: Any, **kwargs: Any) -> Any:
61
+ trace, span, state = _begin(obs_name, as_type, {"args": args, "kwargs": kwargs})
62
+ try:
63
+ out = await fn(*args, **kwargs)
64
+ _finish(trace, span, state, out)
65
+ return out
66
+ except Exception as e: # noqa: BLE001
67
+ span.end(level="ERROR", statusMessage=str(e))
68
+ _ctx.reset(state[0])
69
+ raise
70
+
71
+ return awrapper
72
+
73
+ @functools.wraps(fn)
74
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
75
+ trace, span, state = _begin(obs_name, as_type, {"args": args, "kwargs": kwargs})
76
+ try:
77
+ out = fn(*args, **kwargs)
78
+ _finish(trace, span, state, out)
79
+ return out
80
+ except Exception as e: # noqa: BLE001
81
+ span.end(level="ERROR", statusMessage=str(e))
82
+ _ctx.reset(state[0])
83
+ raise
84
+
85
+ return wrapper
86
+
87
+ return deco(_fn) if callable(_fn) else deco
@@ -0,0 +1,89 @@
1
+ """LangChain callback handler — records chains, LLM calls, and tools as a trace tree.
2
+
3
+ Pass an instance in ``callbacks=[...]``:
4
+
5
+ from memoturn.langchain import MemoturnCallbackHandler
6
+ chain.invoke(input, config={"callbacks": [MemoturnCallbackHandler()]})
7
+
8
+ Implemented without importing langchain (duck-typed) so the SDK has no hard dependency;
9
+ LangChain invokes these methods by name at runtime.
10
+ """
11
+ from __future__ import annotations
12
+
13
+ from typing import Any, Optional
14
+ from uuid import UUID
15
+
16
+ from .client import Memoturn, Span, Trace
17
+ from .decorator import get_client
18
+
19
+
20
+ class MemoturnCallbackHandler:
21
+ # LangChain reads these attributes off the handler.
22
+ raise_error = False
23
+ run_inline = True
24
+
25
+ def __init__(self, client: Optional[Memoturn] = None, trace_name: str = "langchain") -> None:
26
+ self._c = client or get_client()
27
+ self._trace_name = trace_name
28
+ self._trace: Optional[Trace] = None
29
+ self._spans: dict[str, Span] = {}
30
+
31
+ def _ensure_trace(self) -> Trace:
32
+ if self._trace is None:
33
+ self._trace = self._c.trace(name=self._trace_name)
34
+ return self._trace
35
+
36
+ # ── chains ──────────────────────────────────────────────────────────────
37
+ def on_chain_start(self, serialized: Any, inputs: Any, *, run_id: UUID, **kwargs: Any) -> None:
38
+ self._spans[str(run_id)] = self._ensure_trace().span(name="chain", input=inputs)
39
+
40
+ def on_chain_end(self, outputs: Any, *, run_id: UUID, **kwargs: Any) -> None:
41
+ self._end(run_id, output=outputs)
42
+
43
+ def on_chain_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
44
+ self._end(run_id, level="ERROR", statusMessage=str(error))
45
+
46
+ # ── LLMs ────────────────────────────────────────────────────────────────
47
+ def on_llm_start(self, serialized: Any, prompts: list[str], *, run_id: UUID, **kwargs: Any) -> None:
48
+ model = (kwargs.get("invocation_params") or {}).get("model") or (serialized or {}).get("name")
49
+ self._spans[str(run_id)] = self._ensure_trace().generation(name="llm", model=model, input=prompts)
50
+
51
+ def on_chat_model_start(self, serialized: Any, messages: Any, *, run_id: UUID, **kwargs: Any) -> None:
52
+ model = (kwargs.get("invocation_params") or {}).get("model") or (serialized or {}).get("name")
53
+ self._spans[str(run_id)] = self._ensure_trace().generation(name="chat", model=model, input=messages)
54
+
55
+ def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
56
+ usage = None
57
+ try:
58
+ tu = (response.llm_output or {}).get("token_usage") or {}
59
+ usage = {
60
+ "promptTokens": tu.get("prompt_tokens"),
61
+ "completionTokens": tu.get("completion_tokens"),
62
+ "totalTokens": tu.get("total_tokens"),
63
+ }
64
+ except Exception: # noqa: BLE001
65
+ pass
66
+ self._end(run_id, output=getattr(response, "generations", response), usage=usage)
67
+
68
+ def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
69
+ self._end(run_id, level="ERROR", statusMessage=str(error))
70
+
71
+ # ── tools ───────────────────────────────────────────────────────────────
72
+ def on_tool_start(self, serialized: Any, input_str: str, *, run_id: UUID, **kwargs: Any) -> None:
73
+ name = (serialized or {}).get("name", "tool")
74
+ self._spans[str(run_id)] = self._ensure_trace().span(name=name, input=input_str)
75
+
76
+ def on_tool_end(self, output: Any, *, run_id: UUID, **kwargs: Any) -> None:
77
+ self._end(run_id, output=output)
78
+
79
+ def on_tool_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
80
+ self._end(run_id, level="ERROR", statusMessage=str(error))
81
+
82
+ # ── helpers ─────────────────────────────────────────────────────────────
83
+ def _end(self, run_id: UUID, **body: Any) -> None:
84
+ span = self._spans.pop(str(run_id), None)
85
+ if span is not None:
86
+ span.end(**{k: v for k, v in body.items() if v is not None})
87
+
88
+ def flush(self) -> None:
89
+ self._c.flush()
@@ -0,0 +1,51 @@
1
+ """Drop-in OpenAI wrapper — records each chat completion as a memoturn generation."""
2
+ from __future__ import annotations
3
+
4
+ from typing import Any, Optional
5
+
6
+ from .client import Memoturn, Trace
7
+ from .decorator import get_client
8
+
9
+
10
+ def wrap_openai(client: Any, memoturn: Optional[Memoturn] = None, *, trace: Optional[Trace] = None) -> Any:
11
+ """Patch ``client.chat.completions.create`` to trace calls. Returns the same client."""
12
+ mt = memoturn or get_client()
13
+ completions = client.chat.completions
14
+ original = completions.create
15
+
16
+ def create(*args: Any, **kwargs: Any) -> Any:
17
+ t = trace or mt.trace(name="openai.chat")
18
+ gen = t.generation(
19
+ name="openai.chat.completions",
20
+ model=kwargs.get("model"),
21
+ provider="openai",
22
+ input=kwargs.get("messages"),
23
+ modelParameters={k: v for k, v in kwargs.items() if k not in ("model", "messages")},
24
+ )
25
+ try:
26
+ resp = original(*args, **kwargs)
27
+ usage = getattr(resp, "usage", None)
28
+ gen.end(
29
+ output=_message(resp),
30
+ usage={
31
+ "promptTokens": getattr(usage, "prompt_tokens", None),
32
+ "completionTokens": getattr(usage, "completion_tokens", None),
33
+ "totalTokens": getattr(usage, "total_tokens", None),
34
+ }
35
+ if usage
36
+ else None,
37
+ )
38
+ return resp
39
+ except Exception as e: # noqa: BLE001
40
+ gen.end(level="ERROR", statusMessage=str(e))
41
+ raise
42
+
43
+ completions.create = create # type: ignore[assignment]
44
+ return client
45
+
46
+
47
+ def _message(resp: Any) -> Any:
48
+ try:
49
+ return resp.choices[0].message.model_dump()
50
+ except Exception: # noqa: BLE001
51
+ return str(resp)