# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Processors operating on text and regular expressions.""" from collections.abc import AsyncIterable, Callable import re from typing import Type from genai_processors import content_api from genai_processors import processor _MAX_LOOP_COUNT = 1000 class MatchProcessor(processor.Processor): r"""Processor finding text patterns and extracting them from the input stream. There are two modes controlled by the `remove_from_input_stream` parameter. 1. [True] When the pattern is found, the matched text is removed from the input stream and returned as a single part with the substream name `substream_output` and the mimetype text/plain. 2. [False] When the pattern is found, it is not removed from the input stream. It is still returned as a single part with the substream name `substream_output` and the mimetype text/plain. Example [mode 1]: ```python p = MatchProcessor( word_start='[', pattern=r'\[.*\]', substream_output='from_regex' ) output = processor.apply_sync(p, ['a', 'b', 'c[d', 'e]f','g']) ``` output will contain 6 parts: - content_api.ProcessorPart('a', mimetype='text/plain') - content_api.ProcessorPart('b', mimetype='text/plain') - content_api.ProcessorPart('c', mimetype='text/plain') - content_api.ProcessorPart('[de]', mimetype='text/plain', substream_name='from_regex', ) - content_api.ProcessorPart('f', mimetype='text/plain') - content_api.ProcessorPart('g', mimetype='text/plain') The text part of `output` will be `abcfg`. In [mode 2], the output will be: - content_api.ProcessorPart('a', mimetype='text/plain') - content_api.ProcessorPart('b', mimetype='text/plain') - content_api.ProcessorPart('c[d', mimetype='text/plain') - content_api.ProcessorPart('e]f', mimetype='text/plain') - content_api.ProcessorPart('[de]', mimetype='text/plain', substream_name='from_regex', ) - content_api.ProcessorPart('g', mimetype='text/plain') When using this processor in a real-time setting, it is recommended to set `remove_from_input_stream` to False. This will allow the processor to output the parts as soon as possible and not block the stream. If you still need to remove the matched text from the input stream, it is advised to define a short `word_start` and/or flush conditions (via `flush_fn`) that are met often to reset the buffer and to keep producing parts in the output stream. """ def __init__( self, *, pattern: str, word_start: str | None = None, substream_input: str = '', substream_output: str = '', flush_fn: Callable[[content_api.ProcessorPart], bool] | None = (None), remove_from_input_stream: bool = True, transform: ( Callable[[content_api.ProcessorPart], content_api.ProcessorPart] | None ) = None, ): """Extracts text parts from the input stream that match the pattern. Only considers text parts from the input stream `substream_input`. See class docstring for more details. This processor buffers input parts until one of the following happens: 1. the `pattern` is found, then it outputs all parts in the buffer up to the end of the match. 2. when the `flush_fn` returns True, then it outputs all parts in the buffer. When `word_start` is set, the following happens: 3. the word_start is found, then it outputs all the parts before the `word_start`. 4. `word_start` is not found in buffer_text[-len(word_start):], where buffer_text is the concatenation of all the text parts in the buffer. All parts whose text is in buffer_text[-len(word_start):] are returned. While in the buffer, the parts are not output which means they can lead to delays in the output stream. To avoid such delays, set the `remove_from_input_stream` to False and/or define a `flush_fn` that returns True often to discard the parts in the buffer frequently. Args: pattern: pattern to match a text to extract into a part. When `remove_from_input_stream` is True, the matched text will be removed from the stream and will be replaced by a single extracted part. The parts before and after this match will be returned as is. Note that re.DOTALL is used to match newlines. word_start: text to match the start of the text that needs to be captured. `word_start` is not a regular expression but a plain string that will be matched exactly. `word_start` should be a substring of the pattern and should indicate that the pattern is about to be matched. Whenever `word_start` is found, the parts after it will be buffered (not returned) until the pattern is found, the `flush_fn` returns True, or the input stream is exhausted. When set to None (default), this logic is not applied. substream_input: name of the substream to use for the input part. substream_output: name of the substream to use for the extracted part. flush_fn: function to check when to reset the buffer and yield all the parts in the buffer. The part where `flush_fn` returns True will be returned as is and will not be matched against the pattern. remove_from_input_stream: if True, the processor will remove the matched parts from the input stream. If False, the input stream will be preserved and the parts will be returned as is quickly. The processor will output into its `substream_output` substream once a match is found. transform: A transformation to be applied to the matched Parts. """ self._word_start = word_start self._pattern = re.compile(pattern, re.DOTALL) self._substream_input = substream_input self._substream_output = substream_output self._flush_fn = flush_fn or (lambda _: False) self._remove_from_input_stream = remove_from_input_stream if transform: self._transform = transform else: self._transform = lambda part: part def _extract_part( self, text_buffer: str, part_buffer: list[content_api.ProcessorPart] ) -> tuple[list[content_api.ProcessorPart], list[content_api.ProcessorPart]]: """Returns the list of parts to yield and the remaining parts to process.""" to_yield = [] to_process = part_buffer left_over = [] if (match := self._pattern.search(text_buffer)) is None: return to_yield, to_process # We have found the pattern, we can yield all the parts up to the # beginning of the match and then yield the parts after the match. offset = 0 part_idx = -1 for c in part_buffer: part_idx += 1 if ( not content_api.is_text(c.mimetype) or c.substream_name != self._substream_input ): if self._remove_from_input_stream: to_yield.append(c) continue # Find the start of the part to extract. Yields all parts until # this start. if (offset + len(c.text)) <= match.start(): part_end = len(c.text) else: part_end = match.start() - offset if part_end > 0 and self._remove_from_input_stream: to_yield.append( content_api.ProcessorPart( c.text[:part_end], metadata=c.metadata, substream_name=c.substream_name, mimetype=c.mimetype, ) ) if match.start() < offset + len(c.text) and match.start() >= offset: to_yield.append( self._transform( content_api.ProcessorPart( match.group(0), metadata=c.metadata, substream_name=self._substream_output, mimetype=c.mimetype, ) ) ) # Find the start of the parts after the match. part_start = match.end() - offset if part_start < len(c.text) and self._remove_from_input_stream: # We have reached the end of the match, there can be another one later # in the buffer. We keep the part after the match and stop. left_over = [ content_api.ProcessorPart( c.text[part_start:], metadata=c.metadata, substream_name=c.substream_name, mimetype=c.mimetype, ) ] break offset += len(c.text) to_process = left_over + part_buffer[part_idx + 1 :] return to_yield, to_process async def call( self, content: AsyncIterable[content_api.ProcessorPart] ) -> AsyncIterable[content_api.ProcessorPartTypes]: part_buffer = [] async for part in content: if not self._remove_from_input_stream: yield part part_buffer.append(part) if self._flush_fn(part): if self._remove_from_input_stream: for part_b in part_buffer: yield part_b part_buffer = [] text_buffer = content_api.as_text( part_buffer, substream_name=self._substream_input ) # If the word_start is not in the buffer, we can already yield all the # parts up to the last ones that could contain word_start. # This is a quick check to avoid buffering more than necessary. # Only applies when word_start is set. if self._word_start is not None and self._word_start not in text_buffer: offset = 0 idx = 0 # index of first part not yielded. for c in part_buffer: if ( not content_api.is_text(c.mimetype) and self._remove_from_input_stream ): yield c idx += 1 continue offset += len(content_api.as_text(c)) if ( offset < len(text_buffer) - len(self._word_start) and self._remove_from_input_stream ): yield c idx += 1 else: break part_buffer = part_buffer[idx:] else: to_yield, part_buffer = self._extract_part(text_buffer, part_buffer) for part in to_yield: yield part # Process the last part which can contain the pattern many times. loop_count = 0 while part_buffer and loop_count < _MAX_LOOP_COUNT: text_buffer = content_api.as_text( part_buffer, substream_name=self._substream_input ) to_yield, part_buffer = self._extract_part(text_buffer, part_buffer) if not to_yield: # No match found, we can yield all the parts. if self._remove_from_input_stream: for part in part_buffer: yield part break else: for part in to_yield: yield part if loop_count >= _MAX_LOOP_COUNT: raise RuntimeError( 'Max loop count reached, the pattern or the input stream is probably' ' malformed.' ) class UrlExtractor(MatchProcessor): """Replaces encountered text URLs with strongly typed Parts. In some scenarios it is useful to replace URLs mentioned in the prompt with the content they point to. In many cases it can be handled by tool calls, but if we want to avoid additional roundtrip or the underlying model does not support tools, hardwired logic might be preferrable. We recommend splitting detecting the URLs and fetching them into separate processors. The processor that does the fetching should act on Parts with a special MIME type to avoid fetching them unintentionally. And a separate processor should decide which URLs should be processed. This processor turns each URL in the prompt text in-to that Part with a special MIME type. Define a dataclass for each URL: @dataclasses_json.dataclass_json @dataclasses.dataclass(frozen=True) class YouTubeUrl: url: str And then tell UrlExtractor to extract them: UrlExtractor({ 'https://youtube.': YouTubeUrl, 'https://github.com': GithubUrl }) Note that all URLs must have the same scheme to allow efficient matching. """ def __init__( self, urls: dict[str, Type], # pylint: disable=g-bare-generic *, substream_input: str = '', substream_output: str = '', ): """Initiallizes the extractor. Args: urls: A map from URL prefix (e.g. 'https://github.com') to a Dataclass) substream_input: name of the substream to use for the input part. substream_output: name of the substream to use for the extracted part. """ scheme = None for prefix in urls.keys(): next_scheme = prefix.split(':')[0] if scheme and scheme != next_scheme: raise ValueError( 'All URL prefixes must have the same scheme e.g. https. Got' f' {scheme!r} and {next_scheme!r}' ) scheme = next_scheme def transform(part: content_api.ProcessorPart): for prefix, dataclass in urls.items(): if part.text.startswith(prefix): return content_api.ProcessorPart.from_dataclass( dataclass=dataclass(part.text), metadata=part.metadata, substream_name=part.substream_name, ) super().__init__( pattern='(' + '|'.join(urls.keys()) + ")[0-9a-zA-Z$\\-_\\.\\+!*'\\(\\);/\\?:@=&]*", word_start=scheme, substream_input=substream_input, substream_output=substream_output, remove_from_input_stream=True, transform=transform, )