monorepoly 0.2.0__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.
@@ -0,0 +1,466 @@
1
+ """Splits YAML, TOML and Dockerfiles into a file symbol plus objects, keys, tables or stages."""
2
+
3
+ import collections
4
+ import itertools
5
+ import logging
6
+ import re
7
+ import tomllib
8
+ from collections.abc import Iterable, Iterator, Sequence
9
+
10
+ import yaml
11
+
12
+ from monorepoly.models import Symbol, SymbolKind
13
+
14
+ _LOGGER = logging.getLogger(__name__)
15
+
16
+ _DOCKER_ASSIGNMENT = re.compile(
17
+ r'(?P<key>[A-Za-z_][\w.-]*)=(?P<value>"(?:[^"\\]|\\.)*"|\'[^\']*\'|\S+)'
18
+ )
19
+ _DOCKER_DIRECTIVE = re.compile(r'^#\s*(?P<name>\w+)\s*=\s*(?P<value>\S+)\s*$')
20
+ _DOCKER_HEREDOC = re.compile(r'<<-?\s*(?P<quote>["\']?)(?P<word>[A-Za-z_]\w*)(?P=quote)')
21
+ _DOCKER_INSTRUCTION = re.compile(r'^\s*(?P<keyword>[A-Za-z]+)(?:\s|$)')
22
+ # The pre-1.0 `ENV KEY value with spaces` form, which has no `=`.
23
+ _DOCKER_LEGACY_ENV = re.compile(r'^(?P<head>\s*ENV\s+(?P<key>[^\s=]+)\s+)(?P<value>.*?)\s*$', re.I)
24
+ _DOCKER_SECRET_KEYWORDS = frozenset({'ARG', 'ENV', 'LABEL'})
25
+ _DOCKER_STAGE_NAME = re.compile(r'\sAS\s+(?P<name>\S+)\s*$', re.IGNORECASE)
26
+ _GENERATED = re.compile(r'generated|do not edit', re.IGNORECASE)
27
+ _GENERATED_HEAD_LINES = 5
28
+ # The line breaks YAML 1.1 counts, so line numbers agree with the parser's marks.
29
+ _LINE_BREAK = re.compile('\r\n|[\r\n\x85\u2028\u2029]')
30
+ _LINE = re.compile(r'[^\r\n]*')
31
+ _MASK = 'x'
32
+ # A Go template action, as in Helm charts; `{{/* comments */}}` may span lines.
33
+ _TEMPLATE_ACTION = re.compile(r'\{\{.*?\}\}', re.DOTALL)
34
+ _LOCK_FILE = re.compile(r'[-.]lock\.ya?ml$')
35
+ _MAX_BYTES = 512 * 1024
36
+ _REDACTED = '<redacted>'
37
+ _SECRET_KEY = re.compile(r'password|secret|token|private_key|credentials', re.IGNORECASE)
38
+ _TOML_KEY = r'(?:[A-Za-z0-9_-]+|"(?:[^"\\]|\\.)*"|\'[^\']*\')'
39
+ _TOML_HEADER = re.compile(
40
+ rf'^\s*(?P<open>\[\[?)\s*(?P<name>{_TOML_KEY}(?:\s*\.\s*{_TOML_KEY})*)\s*\]\]?\s*(?:#.*)?$'
41
+ )
42
+
43
+ # (kind, 1-based first line, 1-based last line before trimming, name, signature)
44
+ type _Part = tuple[SymbolKind, int, int, str, str]
45
+
46
+
47
+ def chunk_yaml(*, path: str, project_id: int, ref: str, source: str) -> list[Symbol]:
48
+ """Returns the file symbol followed by its Kubernetes objects or top-level keys."""
49
+ if _skipped(path=path, source=source):
50
+ return []
51
+ try:
52
+ documents, template = _compose_yaml(source)
53
+ except yaml.YAMLError as err:
54
+ _LOGGER.warning("Skipping unparsable '%s' in project %s: %s", path, project_id, err)
55
+ return []
56
+ if not documents:
57
+ return []
58
+ return _symbols(
59
+ parts=_yaml_parts(documents=documents, template=template),
60
+ path=path,
61
+ project_id=project_id,
62
+ ref=ref,
63
+ text=_redact_yaml(documents=documents, source=source),
64
+ )
65
+
66
+
67
+ def chunk_toml(*, path: str, project_id: int, ref: str, source: str) -> list[Symbol]:
68
+ """Returns the file symbol followed by its tables."""
69
+ if _skipped(path=path, source=source):
70
+ return []
71
+ try:
72
+ data = tomllib.loads(source)
73
+ except tomllib.TOMLDecodeError as err:
74
+ _LOGGER.warning("Skipping unparsable '%s' in project %s: %s", path, project_id, err)
75
+ return []
76
+ if not data:
77
+ return []
78
+ text = _redact_toml(data=data, source=source)
79
+ if text is None:
80
+ _LOGGER.warning(
81
+ "Skipping '%s' in project %s: a secret is not written literally.", path, project_id
82
+ )
83
+ return []
84
+ # Headers come from the original: a redacted value may also occur in a header.
85
+ return _symbols(
86
+ parts=_toml_parts(_split_lines(source)),
87
+ path=path,
88
+ project_id=project_id,
89
+ ref=ref,
90
+ text=text,
91
+ )
92
+
93
+
94
+ def chunk_dockerfile(*, path: str, project_id: int, ref: str, source: str) -> list[Symbol]:
95
+ """Returns the file symbol followed by its build stages."""
96
+ if _skipped(path=path, source=source):
97
+ return []
98
+ lines = _split_lines(source)
99
+ instructions = list(_docker_instructions(lines))
100
+ if not instructions:
101
+ return []
102
+ stages = [start for start, _, keyword in instructions if keyword == 'FROM']
103
+ if not stages:
104
+ _LOGGER.warning("Skipping '%s' in project %s: no FROM.", path, project_id)
105
+ return []
106
+ return _symbols(
107
+ parts=_docker_parts(lines=lines, stages=stages),
108
+ path=path,
109
+ project_id=project_id,
110
+ ref=ref,
111
+ text=_redact_dockerfile(instructions=instructions, lines=lines, source=source),
112
+ )
113
+
114
+
115
+ def _skipped(*, path: str, source: str) -> bool:
116
+ head = '\n'.join(source.splitlines()[:_GENERATED_HEAD_LINES])
117
+ skipped = bool(
118
+ _LOCK_FILE.search(path) or len(source.encode()) > _MAX_BYTES or _GENERATED.search(head)
119
+ )
120
+ if skipped:
121
+ _LOGGER.debug("Skipping lock, generated or oversized '%s'.", path)
122
+ return skipped
123
+
124
+
125
+ def _symbols(
126
+ *,
127
+ parts: Iterable[_Part],
128
+ path: str,
129
+ project_id: int,
130
+ ref: str,
131
+ text: str,
132
+ ) -> list[Symbol]:
133
+ lines = _split_lines(text)
134
+ symbols = [
135
+ Symbol(
136
+ docstring=None,
137
+ end_line=max(len(lines), 1),
138
+ kind=SymbolKind.CONFIG_FILE,
139
+ path=path,
140
+ project_id=project_id,
141
+ qualname=path,
142
+ ref=ref,
143
+ signature=path,
144
+ source=text,
145
+ start_line=1,
146
+ )
147
+ ]
148
+ for kind, start, end, name, signature in parts:
149
+ end = _trim_end(end=end, lines=lines, start=start)
150
+ symbols.append(
151
+ Symbol(
152
+ docstring=_comment_above(line=start, lines=lines),
153
+ end_line=end,
154
+ kind=kind,
155
+ path=path,
156
+ project_id=project_id,
157
+ qualname=f'{path}::{name}',
158
+ ref=ref,
159
+ signature=signature,
160
+ source='\n'.join(lines[start - 1 : end]),
161
+ start_line=start,
162
+ )
163
+ )
164
+ return symbols
165
+
166
+
167
+ def _split_lines(text: str) -> list[str]:
168
+ lines = _LINE_BREAK.split(text)
169
+ # Drop trailing empty element to keep the file's end_line == len(lines).
170
+ if lines and lines[-1] == '':
171
+ lines.pop()
172
+ return lines
173
+
174
+
175
+ def _trim_end(*, end: int, lines: Sequence[str], start: int) -> int:
176
+ """Returns `end` moved up past trailing blank and comment lines, which belong to the next."""
177
+ end = min(end, len(lines))
178
+ while end > start and (not lines[end - 1].strip() or lines[end - 1].lstrip().startswith('#')):
179
+ end -= 1
180
+ return end
181
+
182
+
183
+ def _comment_above(*, line: int, lines: Sequence[str]) -> str | None:
184
+ """Returns the comment block directly above the 1-based `line`, past a `---` marker."""
185
+ index = line - 2
186
+ while index >= 0 and lines[index].startswith('---'):
187
+ index -= 1
188
+ comments = []
189
+ while index >= 0 and lines[index].lstrip().startswith('#'):
190
+ comments.append(lines[index].lstrip().removeprefix('#').strip())
191
+ index -= 1
192
+ return '\n'.join(reversed(comments)) or None
193
+
194
+
195
+ def _compose_yaml(source: str) -> tuple[list[yaml.Node], str | None]:
196
+ """Returns the documents, and the original source if they were composed from a masked one."""
197
+ try:
198
+ # Composing, unlike loading, accepts custom tags such as GitLab's `!reference`.
199
+ return list(yaml.compose_all(source, Loader=yaml.SafeLoader)), None
200
+ except yaml.YAMLError:
201
+ # Only what does not parse as is gets masked: valid YAML may hold `{{` in its strings.
202
+ if '{{' not in source:
203
+ raise
204
+ return list(yaml.compose_all(_mask_template_actions(source), Loader=yaml.SafeLoader)), source
205
+
206
+
207
+ def _mask_template_actions(source: str) -> str:
208
+ """Returns `source` with Go template actions masked, every character keeping its position.
209
+
210
+ A line holding only actions becomes blank; an inline action becomes a plain `x…` scalar.
211
+ """
212
+ masked = list(source)
213
+ in_action = [False] * len(source)
214
+ for action in _TEMPLATE_ACTION.finditer(source):
215
+ for index in range(action.start(), action.end()):
216
+ if masked[index] not in '\r\n':
217
+ masked[index], in_action[index] = _MASK, True
218
+ for line in _LINE.finditer(source):
219
+ span = range(line.start(), line.end())
220
+ if any(in_action[i] for i in span) and all(
221
+ in_action[i] or source[i].isspace() for i in span
222
+ ):
223
+ masked[line.start() : line.end()] = ' ' * len(span)
224
+ return ''.join(masked)
225
+
226
+
227
+ def _yaml_parts(*, documents: Sequence[yaml.Node], template: str | None) -> Iterator[_Part]:
228
+ for document in documents:
229
+ if not isinstance(document, yaml.MappingNode) or _is_secret(document):
230
+ continue
231
+ kind, metadata = _value(document, 'kind'), _value(document, 'metadata')
232
+ name = _value(metadata, 'name') if isinstance(metadata, yaml.MappingNode) else None
233
+ if isinstance(kind, yaml.ScalarNode) and isinstance(name, yaml.ScalarNode):
234
+ namespace = _value(metadata, 'namespace')
235
+ suffix = (
236
+ f' ({_text(namespace, template)})' if isinstance(namespace, yaml.ScalarNode) else ''
237
+ )
238
+ kind_text, name_text = _text(kind, template), _text(name, template)
239
+ yield (
240
+ SymbolKind.CONFIG_DOCUMENT,
241
+ document.start_mark.line + 1,
242
+ _last_line(document.end_mark),
243
+ f'{kind_text}/{name_text}',
244
+ f'{kind_text} {name_text}{suffix}',
245
+ )
246
+ continue
247
+ for key, value in document.value:
248
+ if isinstance(key, yaml.ScalarNode):
249
+ start, text = key.start_mark.line + 1, _text(key, template)
250
+ yield SymbolKind.CONFIG_KEY, start, _last_line(value.end_mark), text, text
251
+
252
+
253
+ def _text(node: yaml.ScalarNode, template: str | None) -> str:
254
+ """Returns the scalar's value, or its template text where an action was masked in it."""
255
+ if template is None:
256
+ return node.value
257
+ written = template[node.start_mark.index : node.end_mark.index]
258
+ return written.strip('\'"') if '{{' in written else node.value
259
+
260
+
261
+ def _last_line(mark: yaml.Mark) -> int:
262
+ # A block ends where whatever follows it starts, at column 0 of the next line.
263
+ return mark.line + 1 if mark.column else mark.line
264
+
265
+
266
+ def _redact_yaml(*, documents: Sequence[yaml.Node], source: str) -> str:
267
+ """Returns `source` with sensitive values replaced, keeping every line number."""
268
+ # A block ends where the next token starts, so trim the indentation of whatever follows.
269
+ spans = sorted(
270
+ (start, start + len(source[start:end].rstrip()))
271
+ for document in documents
272
+ for start, end in _sensitive_spans(node=document, seen=set())
273
+ )
274
+ merged: list[tuple[int, int]] = []
275
+ for start, end in spans:
276
+ # An alias can point into a value that is already redacted.
277
+ if merged and start <= merged[-1][1]:
278
+ merged[-1] = (merged[-1][0], max(end, merged[-1][1]))
279
+ else:
280
+ merged.append((start, end))
281
+ parts, position = [], 0
282
+ for start, end in merged:
283
+ parts += [source[position:start], _redacted(source[start:end])]
284
+ position = end
285
+ parts.append(source[position:])
286
+ return ''.join(parts)
287
+
288
+
289
+ def _sensitive_spans(*, node: yaml.Node, seen: set[int]) -> Iterator[tuple[int, int]]:
290
+ # Aliases share nodes, and may even refer to their own ancestor.
291
+ if id(node) in seen:
292
+ return
293
+ seen.add(id(node))
294
+ match node:
295
+ case yaml.MappingNode() if _is_secret(node):
296
+ # Fail closed: nothing of a Kubernetes Secret is indexed.
297
+ yield node.start_mark.index, node.end_mark.index
298
+ case yaml.MappingNode():
299
+ # A Kubernetes env entry names its secret instead: `- name: DB_PASSWORD`, `value: …`.
300
+ name = _value(node, 'name')
301
+ secret_name = isinstance(name, yaml.ScalarNode) and bool(_SECRET_KEY.search(name.value))
302
+ for key, value in node.value:
303
+ if isinstance(key, yaml.ScalarNode) and (
304
+ _SECRET_KEY.search(key.value) or (secret_name and key.value == 'value')
305
+ ):
306
+ if value.start_mark.index < value.end_mark.index:
307
+ yield value.start_mark.index, value.end_mark.index
308
+ else:
309
+ yield from _sensitive_spans(node=value, seen=seen)
310
+ case yaml.SequenceNode():
311
+ for item in node.value:
312
+ yield from _sensitive_spans(node=item, seen=seen)
313
+
314
+
315
+ def _is_secret(node: yaml.MappingNode) -> bool:
316
+ kind = _value(node, 'kind')
317
+ return isinstance(kind, yaml.ScalarNode) and kind.value == 'Secret'
318
+
319
+
320
+ def _value(node: yaml.MappingNode, key: str) -> yaml.Node | None:
321
+ return next(
322
+ (v for k, v in node.value if isinstance(k, yaml.ScalarNode) and k.value == key), None
323
+ )
324
+
325
+
326
+ def _toml_parts(lines: Sequence[str]) -> Iterator[_Part]:
327
+ """Yields one part per table, `[[array]]` tables numbered from 0."""
328
+ headers = [*_toml_headers(lines), (len(lines), '', False)]
329
+ seen: collections.Counter[str] = collections.Counter()
330
+ for (index, name, is_array), (next_index, _, _) in itertools.pairwise(headers):
331
+ signature = f'[[{name}]]' if is_array else f'[{name}]'
332
+ qualname = f'{name}[{seen[name]}]' if is_array else name
333
+ seen[name] += 1
334
+ yield SymbolKind.CONFIG_KEY, index + 1, next_index, qualname, signature
335
+
336
+
337
+ def _toml_headers(lines: Sequence[str]) -> Iterator[tuple[int, str, bool]]:
338
+ """Yields (0-based line, dotted name, is array) per table header outside values."""
339
+ depth, quote = 0, None
340
+ for index, line in enumerate(lines):
341
+ if depth == 0 and quote is None and (match := _TOML_HEADER.match(line)):
342
+ yield index, re.sub(r'\s*\.\s*', '.', match['name']), match['open'] == '[['
343
+ continue
344
+ depth, quote = _scan_toml_line(depth=depth, line=line, quote=quote)
345
+
346
+
347
+ def _scan_toml_line(*, depth: int, line: str, quote: str | None) -> tuple[int, str | None]:
348
+ """Returns the bracket depth and open multi-line string delimiter after `line`."""
349
+ index = 0
350
+ while index < len(line):
351
+ if quote:
352
+ end = line.find(quote, index)
353
+ if end < 0:
354
+ return depth, quote
355
+ index, quote = end + 3, None
356
+ elif line.startswith(('"""', "'''"), index):
357
+ quote = line[index : index + 3]
358
+ index += 3
359
+ elif line[index] == '#':
360
+ break
361
+ elif line[index] in '"\'':
362
+ delimiter, index = line[index], index + 1
363
+ while index < len(line) and line[index] != delimiter:
364
+ index += 2 if delimiter == '"' and line[index] == '\\' else 1
365
+ index += 1
366
+ else:
367
+ depth += {'[': 1, '{': 1, ']': -1, '}': -1}.get(line[index], 0)
368
+ index += 1
369
+ return depth, quote
370
+
371
+
372
+ def _redact_toml(*, data: dict, source: str) -> str | None:
373
+ """Returns `source` with sensitive strings replaced, or None if one cannot be found."""
374
+ text = source
375
+ # Longest first, so a secret that contains another is replaced whole.
376
+ for secret in sorted(set(_sensitive_strings(data)), key=len, reverse=True):
377
+ # An escaped or line-continued secret cannot be located, so the file is not indexed.
378
+ if secret not in source:
379
+ return None
380
+ text = text.replace(secret, _redacted(secret))
381
+ return text
382
+
383
+
384
+ def _sensitive_strings(value: object, *, sensitive: bool = False) -> Iterator[str]:
385
+ match value:
386
+ case dict():
387
+ for key, item in value.items():
388
+ yield from _sensitive_strings(
389
+ item, sensitive=sensitive or bool(_SECRET_KEY.search(key))
390
+ )
391
+ case list():
392
+ for item in value:
393
+ yield from _sensitive_strings(item, sensitive=sensitive)
394
+ case str() if sensitive and value:
395
+ yield value
396
+
397
+
398
+ def _redacted(value: str) -> str:
399
+ return _REDACTED + '\n' * len(_LINE_BREAK.findall(value))
400
+
401
+
402
+ def _docker_instructions(lines: Sequence[str]) -> Iterator[tuple[int, int, str]]:
403
+ """Yields (0-based first line, 0-based last line, keyword) per instruction."""
404
+ escape, index = '\\', 0
405
+ # Parser directives count only at the very top, before any other line.
406
+ while index < len(lines) and (match := _DOCKER_DIRECTIVE.match(lines[index])):
407
+ if match['name'].lower() == 'escape':
408
+ escape = match['value']
409
+ index += 1
410
+ while index < len(lines):
411
+ match = _DOCKER_INSTRUCTION.match(lines[index])
412
+ if not match:
413
+ index += 1
414
+ continue
415
+ start = index
416
+ while lines[index].rstrip().endswith(escape) and index + 1 < len(lines):
417
+ index += 1
418
+ words = [
419
+ heredoc['word']
420
+ for line in lines[start : index + 1]
421
+ for heredoc in _DOCKER_HEREDOC.finditer(line)
422
+ ]
423
+ # Heredoc bodies follow the instruction in order, each up to its own terminator.
424
+ for word in words:
425
+ index += 1
426
+ while index < len(lines) and lines[index].strip() != word:
427
+ index += 1
428
+ index = min(index, len(lines) - 1)
429
+ yield start, index, match['keyword'].upper()
430
+ index += 1
431
+
432
+
433
+ def _docker_parts(*, lines: Sequence[str], stages: Sequence[int]) -> Iterator[_Part]:
434
+ """Yields one part per stage, named like `--from=` refers to it: by name, else by index."""
435
+ for number, (start, next_start) in enumerate(itertools.pairwise([*stages, len(lines)])):
436
+ signature = lines[start].strip()
437
+ match = _DOCKER_STAGE_NAME.search(signature)
438
+ name = match['name'] if match else str(number)
439
+ yield SymbolKind.CONFIG_DOCUMENT, start + 1, next_start, name, signature
440
+
441
+
442
+ def _redact_dockerfile(
443
+ *,
444
+ instructions: Iterable[tuple[int, int, str]],
445
+ lines: Sequence[str],
446
+ source: str,
447
+ ) -> str:
448
+ """Returns `source` with sensitive ARG, ENV and LABEL values replaced, line by line."""
449
+ redacted = list(lines)
450
+ for start, end, keyword in instructions:
451
+ if keyword not in _DOCKER_SECRET_KEYWORDS:
452
+ continue
453
+ legacy = _DOCKER_LEGACY_ENV.match(lines[start])
454
+ if legacy and _SECRET_KEY.search(legacy['key']):
455
+ redacted[start] = f'{legacy["head"]}{_REDACTED}'
456
+ continue
457
+ for index in range(start, end + 1):
458
+ redacted[index] = _DOCKER_ASSIGNMENT.sub(_redact_assignment, redacted[index])
459
+ trailing = '\n' if _LINE_BREAK.search(source[-2:]) else ''
460
+ return '\n'.join(redacted) + trailing
461
+
462
+
463
+ def _redact_assignment(match: re.Match[str]) -> str:
464
+ if _SECRET_KEY.search(match['key']):
465
+ return f'{match["key"]}={_REDACTED}'
466
+ return match[0]
monorepoly/config.py ADDED
@@ -0,0 +1,54 @@
1
+ """Settings read from the environment."""
2
+
3
+ import dataclasses
4
+ import logging
5
+ import os
6
+ from collections.abc import Mapping
7
+
8
+ from monorepoly.exceptions import MonorepolyConfigError
9
+
10
+ _LOGGER = logging.getLogger(__name__)
11
+
12
+
13
+ @dataclasses.dataclass(frozen=True, kw_only=True)
14
+ class Settings:
15
+ gitlab_group_id: str
16
+ gitlab_token: str
17
+ gitlab_url: str
18
+ google_cloud_project: str
19
+ mongodb_db: str
20
+ mongodb_uri: str
21
+ typesafe_api_key: str
22
+
23
+
24
+ @dataclasses.dataclass(frozen=True, kw_only=True)
25
+ class ServerSettings:
26
+ """Settings of the hosted server; each user brings their own GitLab token."""
27
+
28
+ gitlab_group_id: str
29
+ gitlab_url: str
30
+ google_cloud_project: str
31
+ mongodb_db: str
32
+ mongodb_uri: str
33
+ public_url: str
34
+ typesafe_api_key: str
35
+
36
+
37
+ def load_settings(environ: Mapping[str, str] = os.environ) -> Settings:
38
+ """Returns the settings, raising if any is missing."""
39
+ return _load(cls=Settings, environ=environ)
40
+
41
+
42
+ def load_server_settings(environ: Mapping[str, str] = os.environ) -> ServerSettings:
43
+ """Returns the hosted server's settings, raising if any is missing."""
44
+ return _load(cls=ServerSettings, environ=environ)
45
+
46
+
47
+ def _load[T](*, cls: type[T], environ: Mapping[str, str]) -> T:
48
+ names = [field.name for field in dataclasses.fields(cls)]
49
+ missing = [name.upper() for name in names if not environ.get(name.upper())]
50
+ if missing:
51
+ raise MonorepolyConfigError(
52
+ f'Missing settings {", ".join(f"`{name}`" for name in missing)}!'
53
+ )
54
+ return cls(**{name: environ[name.upper()] for name in names})
monorepoly/embed.py ADDED
@@ -0,0 +1,66 @@
1
+ """Vertex AI embeddings; the one switch for model, dimensions and input format."""
2
+
3
+ import concurrent.futures
4
+ import enum
5
+ import logging
6
+ from collections.abc import Sequence
7
+
8
+ from google import genai
9
+ from google.genai import types
10
+
11
+ from monorepoly.exceptions import MonorepolyValueError
12
+
13
+ _LOGGER = logging.getLogger(__name__)
14
+
15
+ _LOCATION = 'eu' # gemini-embedding-2 is served only from multi-regions (global, us, eu).
16
+ _MAX_INPUT_CHARS = 20_000 # The 8k token limit, assuming the worst case of ~2.5 chars per token.
17
+ _MODEL = 'gemini-embedding-2'
18
+ # Waits 2, 4, ... 60 s: about 7 minutes in all, outlasting several per-minute quota windows.
19
+ _RETRY = types.HttpRetryOptions(attempts=12, exp_base=2, initial_delay=2, max_delay=60)
20
+ _TIMEOUT_MS = 60_000
21
+ _WORKERS = 8
22
+
23
+ DIMENSIONS = 768
24
+
25
+
26
+ class QueryTask(enum.StrEnum):
27
+ """Task instructions; `gemini-embedding-2` takes them in the input, not as a parameter."""
28
+
29
+ CODE = 'code retrieval'
30
+ SEARCH = 'search result'
31
+
32
+
33
+ class Embedder:
34
+ def __init__(self, *, client: genai.Client) -> None:
35
+ self._client = client
36
+ self._pool = concurrent.futures.ThreadPoolExecutor(max_workers=_WORKERS)
37
+
38
+ def documents(self, items: Sequence[tuple[str, str]]) -> list[list[float]]:
39
+ """Returns one vector per (title, text) pair."""
40
+ if not items:
41
+ return []
42
+ inputs = [f'title: {title} | text: {text}'[:_MAX_INPUT_CHARS] for title, text in items]
43
+ return list(self._pool.map(self._embed, inputs))
44
+
45
+ def query(self, *, task: QueryTask, text: str) -> list[float]:
46
+ """Returns the vector for a search query."""
47
+ return self._embed(f'task: {task} | query: {text}'[:_MAX_INPUT_CHARS])
48
+
49
+ def _embed(self, text: str) -> list[float]:
50
+ # The model only accepts one content per request.
51
+ response = self._client.models.embed_content(
52
+ config=types.EmbedContentConfig(output_dimensionality=DIMENSIONS),
53
+ contents=text,
54
+ model=_MODEL,
55
+ )
56
+ vectors = [embedding.values for embedding in (response.embeddings or [])]
57
+ if len(vectors) != 1:
58
+ raise MonorepolyValueError(f'Expected 1 embedding, got {len(vectors)}!')
59
+ return vectors[0]
60
+
61
+
62
+ def make_embedder(*, project: str) -> Embedder:
63
+ """Returns an embedder authenticated with Application Default Credentials."""
64
+ options = types.HttpOptions(retry_options=_RETRY, timeout=_TIMEOUT_MS)
65
+ client = genai.Client(http_options=options, location=_LOCATION, project=project, vertexai=True)
66
+ return Embedder(client=client)
@@ -0,0 +1,29 @@
1
+ """Exceptions raised by monorepoly."""
2
+
3
+
4
+ class MonorepolyError(Exception):
5
+ """Base class for monorepoly errors."""
6
+
7
+
8
+ class MonorepolyConfigError(MonorepolyError):
9
+ """Raised when required configuration is missing or malformed."""
10
+
11
+
12
+ class MonorepolyAccessError(MonorepolyError):
13
+ """Raised when GitLab rejects the token."""
14
+
15
+
16
+ class MonorepolyGitError(MonorepolyError):
17
+ """Raised when a git command fails."""
18
+
19
+
20
+ class MonorepolyGitTransientError(MonorepolyGitError):
21
+ """Raised when a git command fails on a server error or connection failure worth retrying."""
22
+
23
+
24
+ class MonorepolyValueError(MonorepolyError, ValueError):
25
+ """Raised when external data violates a structural invariant."""
26
+
27
+
28
+ class MonorepolyLicenseError(MonorepolyError):
29
+ """Raised when a license file is malformed, tampered with or signed by an untrusted key."""