commons 0.1.0b1__py3-none-any.whl

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 (64) hide show
  1. commons/__init__.py +42 -0
  2. commons/_agent.py +494 -0
  3. commons/_backends.py +151 -0
  4. commons/_catalog/__init__.py +65 -0
  5. commons/_catalog/_core.py +594 -0
  6. commons/_catalog/_databricks.py +264 -0
  7. commons/_catalog/_import.py +157 -0
  8. commons/_catalog/_security.py +398 -0
  9. commons/_catalog/_snowflake.py +173 -0
  10. commons/_citation_scan.py +320 -0
  11. commons/_citations.py +375 -0
  12. commons/_context_layer.py +185 -0
  13. commons/_data_dictionary.py +457 -0
  14. commons/_data_source.py +628 -0
  15. commons/_definitions/__init__.py +40 -0
  16. commons/_definitions/_compile.py +498 -0
  17. commons/_definitions/_emit_duckdb.py +346 -0
  18. commons/_definitions/_emit_sql.py +534 -0
  19. commons/_definitions/_export.py +845 -0
  20. commons/_definitions/_expression.py +529 -0
  21. commons/_definitions/_registry.py +311 -0
  22. commons/_display.py +293 -0
  23. commons/_duckdb.py +62 -0
  24. commons/_execution/__init__.py +1 -0
  25. commons/_execution/_backend.py +272 -0
  26. commons/_execution/_env.py +129 -0
  27. commons/_frames.py +149 -0
  28. commons/_handles.py +65 -0
  29. commons/_icons.py +56 -0
  30. commons/_measures.py +844 -0
  31. commons/_pool.py +419 -0
  32. commons/_prompt.py +313 -0
  33. commons/_provenance.py +161 -0
  34. commons/_reminders.py +53 -0
  35. commons/_rows.py +94 -0
  36. commons/_sample_summary.py +267 -0
  37. commons/_sql_guard.py +141 -0
  38. commons/_tools.py +852 -0
  39. commons/_tracing.py +45 -0
  40. commons/_ui/__init__.py +46 -0
  41. commons/_ui/_app.py +58 -0
  42. commons/_ui/_assets.py +52 -0
  43. commons/_ui/_server.py +78 -0
  44. commons/_ui/_theme.py +61 -0
  45. commons/prompts/README.md +1 -0
  46. commons/prompts/citation-request.md +1 -0
  47. commons/prompts/system-prompt.md +158 -0
  48. commons/py.typed +0 -0
  49. commons/ui.py +5 -0
  50. commons/www/README.md +1 -0
  51. commons/www/commons-chat/commons-chat.css +681 -0
  52. commons/www/commons-chat/commons-chat.js +94 -0
  53. commons/www/commons-chat/figs/citation-definition.svg +4 -0
  54. commons/www/commons-chat/figs/citation-mark.svg +5 -0
  55. commons/www/commons-chat/figs/citation-prose.svg +4 -0
  56. commons/www/commons-chat/figs/citation-schema.svg +4 -0
  57. commons/www/commons-chat/figs/trusted-icon.svg +6 -0
  58. commons/www/commons-chat/figs/warning-icon.svg +7 -0
  59. commons/www/commons-viewer/commons-viewer.css +447 -0
  60. commons/www/commons-viewer/commons-viewer.js +222 -0
  61. commons-0.1.0b1.dist-info/METADATA +88 -0
  62. commons-0.1.0b1.dist-info/RECORD +64 -0
  63. commons-0.1.0b1.dist-info/WHEEL +4 -0
  64. commons-0.1.0b1.dist-info/licenses/LICENSE.md +21 -0
commons/__init__.py ADDED
@@ -0,0 +1,42 @@
1
+ """Build trustworthy data agents.
2
+
3
+ Give an LLM data, semantic, and context layers to work with, tools for
4
+ querying them, and A/B/C provenance semantics so every answer carries a
5
+ classification as to how much it can be trusted.
6
+ """
7
+
8
+ import importlib
9
+ from typing import Any
10
+
11
+ from ._agent import Commons
12
+ from ._context_layer import ContextLayer, context_layer
13
+ from ._data_source import DataSource, data_source, list_tables
14
+ from ._measures import Injected, Measure, SemanticLayer, measure, semantic_layer
15
+ from ._provenance import Tag
16
+
17
+ __all__: list[str] = [
18
+ "Commons",
19
+ "ContextLayer",
20
+ "DataSource",
21
+ "Injected",
22
+ "Measure",
23
+ "SemanticLayer",
24
+ "Tag",
25
+ "context_layer",
26
+ "data_source",
27
+ "list_tables",
28
+ "measure",
29
+ "semantic_layer",
30
+ ]
31
+
32
+
33
+ def __getattr__(name: str) -> Any:
34
+ # `commons.ui` is resolved on first use so that importing commons does
35
+ # not import shiny for users who may never build a UI.
36
+ if name == "ui":
37
+ return importlib.import_module(f"{__name__}.ui")
38
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
39
+
40
+
41
+ def __dir__() -> list[str]:
42
+ return sorted([*__all__, "ui"])
commons/_agent.py ADDED
@@ -0,0 +1,494 @@
1
+ """The agent: its layers, the tools they earn, and the rules a turn follows.
2
+
3
+ `pkg-r/R/commons.R` assembles the same agent for R, in the order this follows.
4
+ A `Commons` agent inherits directly from `chatlas.Chat`,
5
+ in the same way it inherits from `ellmer::Chat` in R. A `Commons` agent
6
+ will reject `Chat` methods it does not explicitly support to prevent
7
+ interaction without provenance and citation tracking.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import copy
13
+ import warnings
14
+ from collections.abc import AsyncGenerator, Mapping, Sequence
15
+ from typing import Any, Literal, NoReturn
16
+
17
+ from chatlas import Chat, StreamController, Tool, Turn, UserTurn
18
+ from chatlas.types import ChatResponse, Content, SubmitInputArgsT
19
+ from pydantic import BaseModel
20
+
21
+ from ._backends import DuckDBBackend, EngineBackend
22
+ from ._citation_scan import CitationScanner
23
+ from ._citations import (
24
+ CitationRequest,
25
+ CorpusEntry,
26
+ build_citation_corpus,
27
+ turn_has_user_message,
28
+ )
29
+ from ._context_layer import ContextLayer, augment_context_layer
30
+ from ._data_source import DataSource
31
+ from ._definitions import Registry, build_registry
32
+ from ._handles import HandleStore
33
+ from ._measures import SemanticLayer, resolve_injections, semantic_layer
34
+ from ._prompt import (
35
+ check_instructions,
36
+ read_instructions,
37
+ render_system_prompt,
38
+ system_prompt_data,
39
+ system_prompt_template,
40
+ )
41
+ from ._provenance import collect_appended_tags, derive_provenance_tag, provenance_aside
42
+ from ._reminders import append_restored_conversation_reminder, append_turn_reminder
43
+ from ._tools import FirstTouch, ToolContext, build_commons_tools
44
+
45
+ __all__ = ["Commons"]
46
+
47
+ EchoOptions = Literal["output", "all", "none", "text"]
48
+
49
+ # The label a lone source is filed under. An agent with one source never shows
50
+ # the model a `source` argument and never labels its prompt sections, so this
51
+ # is only what the agent's own bookkeeping keys on.
52
+ SOLE_SOURCE = "data"
53
+
54
+ # Frozen and empty, so every agent built without a semantic layer can share it.
55
+ _NO_MEASURES = semantic_layer()
56
+
57
+
58
+ class Commons(Chat[Any, Any]):
59
+ """A trustworthy agent that answers questions about its data.
60
+
61
+ Given a `chatlas.Chat` for the provider and model, the data sources it
62
+ can query, and optionally a semantic layer of trusted calculations and a
63
+ context layer of prose, a `Commons` agent will allow for agent interactions
64
+ with answers classified by how they were produced.
65
+
66
+ A `Commons` agent inherits directly from `chatlas.Chat` and relies on the
67
+ chatlas infrastructure to set up the LLM provider and model. `Commons`
68
+ initializes its own chat state and system prompt to ensure provenance
69
+ and citation tracking. Passing a custom system prompt in the `Commons`
70
+ constructor is ignored with a warning; use `instructions` to add to
71
+ commons' prompt instead. For best results, enable thinking where the
72
+ provider and model support it.
73
+
74
+ `chat()` and `stream_async()` are the currently supported ways to
75
+ interact with a `Commons` agent. The other entry points chatlas offers
76
+ (`chat_async()`, `stream()`, `chat_structured()`, etc.)
77
+ are disabled and raise `NotImplementedError`s because they are not (yet)
78
+ tied in to the commons framework. The rest of chatlas's surface works as
79
+ it does on any chat.
80
+
81
+ `data_sources` is a `DataSource`, or a mapping of name to `DataSource`;
82
+ a measure can take a named source's connection as an argument named
83
+ after it. `instructions` is extra text placed under an
84
+ `## Additional instructions` heading at the end of commons' built-in
85
+ system prompt, as a string or the path to a text or Markdown file.
86
+
87
+ Construction raises a TypeError if `client` is not a `chatlas.Chat`, if
88
+ an entry of `data_sources` is not a `DataSource`, or if a layer is not
89
+ the layer its argument claims; a ValueError if `data_sources` names no
90
+ source or a measure asks for an injection no named source can fill; and
91
+ a FileNotFoundError if `instructions` names a file that does not exist.
92
+ """
93
+
94
+ def __init__(
95
+ self,
96
+ client: Chat,
97
+ data_sources: DataSource | Mapping[str, DataSource],
98
+ semantic_layer: SemanticLayer | None = None,
99
+ context_layer: ContextLayer | None = None,
100
+ *,
101
+ instructions: str | None = None,
102
+ ) -> None:
103
+ if not isinstance(client, Chat):
104
+ raise TypeError(
105
+ "client must be a chatlas.Chat, e.g. from chatlas.ChatAnthropic(), "
106
+ f"not {type(client).__name__}."
107
+ )
108
+ # What the client carried is warned about before any other argument
109
+ # is checked, as pkg-r/R/commons.R does, so a bad later argument
110
+ # does not eat the warning.
111
+ _warn_ignored_client_state(client)
112
+ sources = _as_data_sources(data_sources)
113
+ if context_layer is not None and not isinstance(context_layer, ContextLayer):
114
+ raise TypeError(
115
+ "context_layer must be a ContextLayer from commons.context_layer(), "
116
+ f"or None, not {type(context_layer).__name__}."
117
+ )
118
+ if semantic_layer is None:
119
+ semantic_layer = _NO_MEASURES
120
+ if not isinstance(semantic_layer, SemanticLayer):
121
+ raise TypeError(
122
+ "semantic_layer must be a SemanticLayer from "
123
+ f"commons.semantic_layer(), or None, not "
124
+ f"{type(semantic_layer).__name__}."
125
+ )
126
+ check_instructions(instructions)
127
+
128
+ # Share the provider, which carries the chosen model; shallow-copy
129
+ # the chat kwargs so later changes don't cross between the two.
130
+ super().__init__(
131
+ provider=client.provider, kwargs_chat=copy.copy(client.kwargs_chat)
132
+ )
133
+ # chatlas never generates one, so an id the caller chose is explicitly kept
134
+ self.conversation_id = client.conversation_id
135
+ _carry_model_params(self, client)
136
+
137
+ self._sources = sources
138
+ self._context_layer = augment_context_layer(context_layer, sources.values())
139
+ self._definitions = build_registry(sources)
140
+ self._measures = semantic_layer.measures
141
+ # The injectables side stays here because it is the only part that
142
+ # knows what a DataSource is: a measure asks for a source by name and
143
+ # receives whatever that source is queried through.
144
+ self._injections = resolve_injections(
145
+ self._measures, _measure_injectables(data_sources, sources)
146
+ )
147
+ self._first_touch = FirstTouch()
148
+ self._handles = HandleStore()
149
+ self._citation_request = CitationRequest()
150
+ self._corpus = build_citation_corpus(
151
+ self._context_layer, self._measures.values(), sources
152
+ )
153
+ self._restore_reminder_pending = False
154
+
155
+ tools = build_commons_tools(
156
+ ToolContext(
157
+ sources=sources,
158
+ measures=self._measures,
159
+ definitions=self._definitions,
160
+ context_layer=self._context_layer,
161
+ handles=self._handles,
162
+ citation_request=self._citation_request,
163
+ injections=self._injections,
164
+ first_touch=self._first_touch,
165
+ )
166
+ )
167
+ self.set_tools(list(tools))
168
+ self.system_prompt = _system_prompt(
169
+ sources,
170
+ self._definitions,
171
+ instructions=instructions,
172
+ tools=tools,
173
+ model=self.model,
174
+ )
175
+
176
+ def __repr__(self) -> str:
177
+ count = len(self._sources)
178
+ plural = "" if count == 1 else "s"
179
+ return f"A commons agent over {count} data source{plural}."
180
+
181
+ def __deepcopy__(self, memo: dict[int, Any]) -> NoReturn:
182
+ # chatlas Chat objects can be deep-copied to fork a conversation,
183
+ # Because a Commons agent may have database connections (which can't be copied)
184
+ # we explicitly forbid deep copying with a clear error.
185
+ raise NotImplementedError(
186
+ "A commons agent cannot be copied: it holds database connections "
187
+ "that copying cannot reach. Build a second agent instead."
188
+ )
189
+
190
+ # ---- asking it something ---------------------------------------------
191
+
192
+ def chat(
193
+ self,
194
+ *args: Content | str,
195
+ echo: EchoOptions = "output",
196
+ stream: bool = True,
197
+ kwargs: SubmitInputArgsT | None = None,
198
+ ) -> ChatResponse:
199
+ """Ask a question and wait for the whole answer.
200
+
201
+ A reminder queued with `queue_restore_reminder()` rides this turn,
202
+ and a turn that fails leaves it queued for the next one.
203
+ """
204
+ was_pending = self._restore_reminder_pending
205
+ inputs = self._prepare_turn_inputs(args)
206
+ self._citation_request.reset()
207
+ response = super().chat(*inputs, echo=echo, stream=stream, kwargs=kwargs)
208
+ self._consume_restore_reminder(was_pending)
209
+ return response
210
+
211
+ async def stream_async(
212
+ self,
213
+ *args: Content | str,
214
+ content: Literal["text", "all"] = "text",
215
+ echo: EchoOptions = "none",
216
+ data_model: type[BaseModel] | None = None,
217
+ kwargs: SubmitInputArgsT | None = None,
218
+ controller: StreamController | None = None,
219
+ ) -> AsyncGenerator[Any, None]:
220
+ """Ask a question and stream the answer as it arrives.
221
+
222
+ The signature is identical to chatlas's, so a chat UI can drive this agent
223
+ directly and needs the attachment content, the mode, and the controller its
224
+ stop button cancels through.
225
+
226
+ The Commons agent does not accept `data_model`. If you pass it, this
227
+ method raises NotImplementedError. In chatlas, using `data_model` means
228
+ the chunks are JSON that the caller parses as one document. Commons adds
229
+ provenance markers and citations to the stream that are not compatible with
230
+ `data_model`, so it is explicitly forbidden.
231
+ """
232
+ if data_model is not None:
233
+ raise NotImplementedError(
234
+ "stream_async(data_model=...) is not available on a commons "
235
+ "agent: the provenance marker and citations it appends would leave the "
236
+ "streamed JSON unparseable."
237
+ )
238
+ from_index = len(self.get_turns())
239
+ was_pending = self._restore_reminder_pending
240
+ inputs = self._prepare_turn_inputs(args)
241
+ self._citation_request.reset()
242
+ raw = await super().stream_async(
243
+ *inputs,
244
+ content=content,
245
+ echo=echo,
246
+ kwargs=kwargs,
247
+ controller=controller,
248
+ )
249
+ return self._projected(raw, from_index, was_pending)
250
+
251
+ # Citations are always projected, so reserved model markup cannot reach
252
+ # the browser whatever the display layer does with the stream.
253
+ async def _projected(
254
+ self,
255
+ raw: AsyncGenerator[Any, None],
256
+ from_index: int,
257
+ was_pending: bool,
258
+ ) -> AsyncGenerator[str | Content, None]:
259
+ scanner = CitationScanner(self._corpus)
260
+ try:
261
+ async for chunk in raw:
262
+ # A provider's structured content passes through untouched;
263
+ # only the model's own text can carry the reserved dialect.
264
+ if not isinstance(chunk, str):
265
+ yield chunk
266
+ continue
267
+ projected = scanner.feed(chunk)
268
+ if projected:
269
+ yield projected
270
+ finally:
271
+ # A consumer that walks away early without cancelling still
272
+ # closes the provider's stream.
273
+ await raw.aclose()
274
+
275
+ self._consume_restore_reminder(was_pending)
276
+
277
+ tail = scanner.finish()
278
+ if tail:
279
+ yield tail
280
+
281
+ tag = derive_provenance_tag(
282
+ collect_appended_tags(self.get_turns(), from_index),
283
+ scanner.any_verified,
284
+ )
285
+ aside = provenance_aside(tag)
286
+ if aside:
287
+ yield aside
288
+
289
+ # ---- chatlas.Chat entry points a Commons agent does not support ----------------
290
+
291
+ def chat_async(self, *args: Any, **kwargs: Any) -> NoReturn:
292
+ raise NotImplementedError(_unrouted("chat_async"))
293
+
294
+ def stream(self, *args: Any, **kwargs: Any) -> NoReturn:
295
+ raise NotImplementedError(_unrouted("stream"))
296
+
297
+ def chat_structured(self, *args: Any, **kwargs: Any) -> NoReturn:
298
+ raise NotImplementedError(_unrouted("chat_structured"))
299
+
300
+ def chat_structured_async(self, *args: Any, **kwargs: Any) -> NoReturn:
301
+ raise NotImplementedError(_unrouted("chat_structured_async"))
302
+
303
+ def extract_data(self, *args: Any, **kwargs: Any) -> NoReturn:
304
+ raise NotImplementedError(_unrouted("extract_data"))
305
+
306
+ def extract_data_async(self, *args: Any, **kwargs: Any) -> NoReturn:
307
+ raise NotImplementedError(_unrouted("extract_data_async"))
308
+
309
+ def to_solver(self, *args: Any, **kwargs: Any) -> NoReturn:
310
+ # chatlas's solver answers each eval sample through `chat_async()` or
311
+ # `chat_structured_async()`, so it would raise mid-eval anyway.
312
+ raise NotImplementedError(
313
+ "A commons agent has no to_solver(): the solver it returns "
314
+ "answers through chat_async(), which an agent does not provide."
315
+ )
316
+
317
+ # ---- what the agent knows --------------------------------------------
318
+
319
+ def citation_corpus(self) -> list[CorpusEntry]:
320
+ """The trusted text this agent's citations are verified against."""
321
+ return list(self._corpus)
322
+
323
+ def prewarm(self) -> None:
324
+ """Build the caches the first question would otherwise pay for.
325
+
326
+ Failures propagate: a direct call is typically warming caches ahead
327
+ of a deployment, so a cold cache should fail the deploy.
328
+ """
329
+ if self._context_layer is not None:
330
+ self._context_layer.prewarm()
331
+ for source in self._sources.values():
332
+ source.ensure_loaded()
333
+
334
+ # ---- the turn rules ---------------------------------------------------
335
+
336
+ def add_turn(self, turn: Turn) -> None:
337
+ """Add a turn, restarting the citation request if a person spoke.
338
+
339
+ A user turn of nothing but tool results is the same question still
340
+ running, and an assistant turn is nobody asking anything.
341
+ """
342
+ if isinstance(turn, UserTurn) and turn_has_user_message(turn):
343
+ self._citation_request.reset()
344
+ super().add_turn(turn)
345
+
346
+ def set_turns(self, turns: Sequence[Turn]) -> None:
347
+ """Replace the conversation, dropping any reminder queued for it."""
348
+ self._restore_reminder_pending = False
349
+ super().set_turns(turns)
350
+
351
+ def queue_restore_reminder(self) -> None:
352
+ """Tell the next turn that the session behind its history is gone."""
353
+ self._restore_reminder_pending = True
354
+
355
+ def _prepare_turn_inputs(
356
+ self, inputs: Sequence[Content | str]
357
+ ) -> list[Content | str]:
358
+ prepared = append_turn_reminder(inputs, self.model)
359
+ if self._restore_reminder_pending:
360
+ prepared = append_restored_conversation_reminder(prepared)
361
+ return prepared
362
+
363
+ # Only a reminder that was pending when the turn started is spent, so a
364
+ # turn that failed, or one whose stream was never consumed, leaves the
365
+ # next turn to deliver it.
366
+ def _consume_restore_reminder(self, was_pending: bool) -> None:
367
+ if was_pending:
368
+ self._restore_reminder_pending = False
369
+
370
+
371
+ # Called straight from __init__, so one stacklevel reaches whoever built the
372
+ # agent from every warning below.
373
+ _CALLER = 3
374
+
375
+
376
+ def _warn_ignored_client_state(client: Chat) -> None:
377
+ """Warn about whatever the agent's own chat will not carry over."""
378
+ if client.system_prompt is not None:
379
+ warnings.warn(
380
+ "The system prompt set on client is ignored; commons builds its "
381
+ "own. Use `instructions` to add to commons' prompt.",
382
+ stacklevel=_CALLER,
383
+ )
384
+ if client.get_turns():
385
+ warnings.warn(
386
+ "An agent starts a new conversation, so the turns on client are "
387
+ "not carried over. Restore them with the agent's set_turns().",
388
+ stacklevel=_CALLER,
389
+ )
390
+
391
+
392
+ def _unrouted(name: str) -> str:
393
+ """Why an inherited entry point is closed, and what to ask instead."""
394
+ return (
395
+ f"A commons agent has no {name}(): it would submit a turn outside "
396
+ "commons' turn handling, and the answer would carry neither the "
397
+ "citation scanner's work nor a provenance marker. Ask the agent "
398
+ "with chat() or stream_async()."
399
+ )
400
+
401
+
402
+ def _carry_model_params(agent: Chat, client: Chat) -> None:
403
+ """Move whatever `set_model_params()` put on the caller's chat.
404
+
405
+ The agent brings its own system prompt and tools and starts an empty
406
+ conversation, as `pkg-r/R/commons.R` does when it initializes from the
407
+ client's provider, so the model parameters are all there is to carry.
408
+ """
409
+ # chatlas has the setter for these and no getter, so they are read off the
410
+ # attribute behind it and written back through the public setter, which
411
+ # checks them against the provider again. An attribute that is missing, or
412
+ # no longer a mapping, has to be told apart from an empty one, which means
413
+ # nothing was set: dropping a temperature in silence is worse than saying
414
+ # that this chatlas does not show what was set.
415
+ params = getattr(client, "_standard_model_params", None)
416
+ if not isinstance(params, Mapping):
417
+ warnings.warn(
418
+ "Any model parameters set on client with set_model_params() are "
419
+ "not carried onto the agent: this version of chatlas does not "
420
+ "expose them.",
421
+ stacklevel=_CALLER,
422
+ )
423
+ elif params:
424
+ agent.set_model_params(**dict(params))
425
+
426
+
427
+ def _as_data_sources(
428
+ data_sources: DataSource | Mapping[str, DataSource],
429
+ ) -> dict[str, DataSource]:
430
+ if isinstance(data_sources, DataSource):
431
+ return {SOLE_SOURCE: data_sources}
432
+ if not isinstance(data_sources, Mapping):
433
+ raise TypeError(
434
+ "data_sources must be a DataSource from commons.data_source(), or a "
435
+ f"mapping of name to DataSource, not {type(data_sources).__name__}."
436
+ )
437
+ if not data_sources:
438
+ raise ValueError("data_sources must name at least one DataSource.")
439
+ wrong = [
440
+ name
441
+ for name, source in data_sources.items()
442
+ if not isinstance(source, DataSource)
443
+ ]
444
+ if wrong:
445
+ raise TypeError(
446
+ f"Every entry in data_sources must be a DataSource from "
447
+ f"commons.data_source(); {', '.join(sorted(wrong))} "
448
+ f"{'is' if len(wrong) == 1 else 'are'} not."
449
+ )
450
+ return dict(data_sources)
451
+
452
+
453
+ # A measure can take a named source's connection as an argument named after
454
+ # the source. A lone source passed on its own has no name to be asked for, so
455
+ # it offers nothing to inject, as in pkg-r/R/commons.R.
456
+ def _measure_injectables(
457
+ data_sources: DataSource | Mapping[str, DataSource],
458
+ sources: Mapping[str, DataSource],
459
+ ) -> dict[str, Any]:
460
+ if isinstance(data_sources, DataSource):
461
+ return {}
462
+ return {name: _connection(source) for name, source in sources.items()}
463
+
464
+
465
+ def _connection(source: DataSource) -> Any:
466
+ """What a measure queries a source through.
467
+
468
+ The engine for a caller's database and the DuckDB connection for the
469
+ in-process one commons builds, which is what each was queried through
470
+ before it reached the measure.
471
+ """
472
+ backend = source.backend
473
+ if isinstance(backend, EngineBackend):
474
+ return backend.engine
475
+ if isinstance(backend, DuckDBBackend):
476
+ return backend.connection
477
+ raise TypeError(f"No connection to inject for a {type(backend).__name__}.")
478
+
479
+
480
+ def _system_prompt(
481
+ sources: Mapping[str, DataSource],
482
+ definitions: Registry,
483
+ instructions: str | None,
484
+ tools: Sequence[Tool],
485
+ model: str | None,
486
+ ) -> str:
487
+ data = system_prompt_data(
488
+ sources,
489
+ definitions,
490
+ instructions=read_instructions(instructions),
491
+ tools=[tool.name for tool in tools],
492
+ model=model,
493
+ )
494
+ return render_system_prompt(system_prompt_template(), data)
commons/_backends.py ADDED
@@ -0,0 +1,151 @@
1
+ """What a data source needs from whatever holds its tables.
2
+
3
+ Two implementations: a SQLAlchemy `Engine` for connections a caller supplies
4
+ (D2 makes the engine the connection currency), and a raw DuckDB connection for
5
+ the in-process database commons builds from frames and pins. The raw path is
6
+ not routed through SQLAlchemy because the lockdown and the deferred pin writes
7
+ are DuckDB-specific and gain nothing from it.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from collections.abc import Callable
13
+ from typing import TYPE_CHECKING, Any, Protocol
14
+
15
+ import duckdb
16
+ import sqlalchemy
17
+
18
+ from ._duckdb import quote_identifier
19
+
20
+ if TYPE_CHECKING:
21
+ from ._data_source import TableId
22
+
23
+ __all__ = ["Backend", "DuckDBBackend", "EngineBackend"]
24
+
25
+
26
+ class Backend(Protocol):
27
+ def query(self, sql: str) -> list[dict[str, Any]]: ...
28
+
29
+ def list_tables(self) -> list[str]: ...
30
+
31
+ def quote(self, table_id: TableId) -> str: ...
32
+
33
+ def columns(self, table_id: TableId) -> list[dict[str, Any]]:
34
+ """A relation's columns, in the shape a catalog listing reports them.
35
+
36
+ For the backends a warehouse catalog never describes, since there the
37
+ listing itself already carries the columns.
38
+ """
39
+ ...
40
+
41
+ def dialect(self) -> str: ...
42
+
43
+ def inspector(self) -> Callable[[TableId], bool] | None:
44
+ """A predicate answering whether a table exists, if the backend can.
45
+
46
+ Used to tell a table that is genuinely absent from one that a failing
47
+ query only made look absent. A backend that cannot answer returns
48
+ None, and the caller re-raises the original error rather than guessing.
49
+ """
50
+ ...
51
+
52
+
53
+ class DuckDBBackend:
54
+ def __init__(self, con: duckdb.DuckDBPyConnection) -> None:
55
+ self._con = con
56
+
57
+ @property
58
+ def connection(self) -> duckdb.DuckDBPyConnection:
59
+ return self._con
60
+
61
+ def query(self, sql: str) -> list[dict[str, Any]]:
62
+ cursor = self._con.execute(sql)
63
+ if cursor.description is None:
64
+ return []
65
+ columns = [column[0] for column in cursor.description]
66
+ return [dict(zip(columns, row)) for row in cursor.fetchall()]
67
+
68
+ def list_tables(self) -> list[str]:
69
+ rows = self._con.execute(
70
+ "SELECT table_name FROM information_schema.tables "
71
+ "WHERE table_schema = 'main' ORDER BY table_name"
72
+ ).fetchall()
73
+ return [row[0] for row in rows]
74
+
75
+ def quote(self, table_id: TableId) -> str:
76
+ return ".".join(quote_identifier(part) for part in table_id.parts)
77
+
78
+ def columns(self, table_id: TableId) -> list[dict[str, Any]]:
79
+ # A zero-row select, because the cursor description carries DuckDB's
80
+ # own type names and no round trip to a metadata table is needed.
81
+ cursor = self._con.execute(f"SELECT * FROM {self.quote(table_id)} LIMIT 0")
82
+ description = cursor.description or ()
83
+ return [{"column": name, "type": str(kind)} for name, kind, *_ in description]
84
+
85
+ def dialect(self) -> str:
86
+ return "duckdb"
87
+
88
+ def inspector(self) -> Callable[[TableId], bool] | None:
89
+ # Frame and pin sources build their own tables, so nothing reaches
90
+ # the existence check by this route.
91
+ return None
92
+
93
+
94
+ class EngineBackend:
95
+ def __init__(self, engine: sqlalchemy.Engine) -> None:
96
+ self._engine = engine
97
+
98
+ @property
99
+ def engine(self) -> sqlalchemy.Engine:
100
+ return self._engine
101
+
102
+ def query(self, sql: str) -> list[dict[str, Any]]:
103
+ with self._engine.connect() as connection:
104
+ result = connection.execute(sqlalchemy.text(sql))
105
+ return [dict(row) for row in result.mappings()]
106
+
107
+ def list_tables(self) -> list[str]:
108
+ return list(sqlalchemy.inspect(self._engine).get_table_names())
109
+
110
+ def quote(self, table_id: TableId) -> str:
111
+ preparer = self._engine.dialect.identifier_preparer
112
+ quoted = preparer.quote(table_id.table)
113
+ if table_id.schema is None:
114
+ return quoted
115
+ # quote_schema() is given one component at a time: handed
116
+ # "ANALYTICS.PUBLIC" it produces one identifier containing a dot.
117
+ outer = ".".join(
118
+ preparer.quote_schema(part)
119
+ for part in (table_id.catalog, table_id.schema)
120
+ if part is not None
121
+ )
122
+ return f"{outer}.{quoted}"
123
+
124
+ def columns(self, table_id: TableId) -> list[dict[str, Any]]:
125
+ inspector = sqlalchemy.inspect(self._engine)
126
+ # SQLAlchemy takes every level above the table as one dotted
127
+ # `schema`, so a catalog is joined onto it rather than dropped.
128
+ outer = ".".join(table_id.parts[:-1])
129
+ return [
130
+ {
131
+ "column": column["name"],
132
+ "type": str(column["type"]),
133
+ "nullable": column.get("nullable"),
134
+ "description": column.get("comment"),
135
+ }
136
+ for column in inspector.get_columns(table_id.table, schema=outer or None)
137
+ ]
138
+
139
+ def dialect(self) -> str:
140
+ return self._engine.dialect.name
141
+
142
+ def inspector(self) -> Callable[[TableId], bool] | None:
143
+ inspector = sqlalchemy.inspect(self._engine)
144
+
145
+ def exists(table_id: TableId) -> bool:
146
+ # SQLAlchemy takes every level above the table as one dotted
147
+ # `schema`, so a catalog is joined onto it rather than dropped.
148
+ outer = ".".join(table_id.parts[:-1])
149
+ return inspector.has_table(table_id.table, schema=outer or None)
150
+
151
+ return exists