Skip to content

Tools

lazyllm.tools.IntentClassifier

Bases: ModuleBase

Intent classification module that classifies input text into a given intent list.
Supports automatic selection of Chinese or English prompt templates, and allows enhancement through examples, prompt text, constraints, and attention notes.

Parameters:

  • llm

    The large language model instance used for intent classification.

  • intent_list (list, default: None ) –

    Optional, list of intent categories, e.g., ["chat", "weather", "QA"].

  • prompt (str, default: '' ) –

    Optional, custom prompt inserted into the system prompt template.

  • constrain (str, default: '' ) –

    Optional, classification constraint description.

  • attention (str, default: '' ) –

    Optional, attention notes for classification.

  • examples (list[list[str, str]], default: None ) –

    Optional, classification examples, each element is [input text, label].

  • return_trace (bool, default: False ) –

    Whether to return execution trace. Default is False.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import IntentClassifier
>>> classifier_llm = lazyllm.OnlineChatModule(source="openai")
>>> chatflow_intent_list = ["Chat", "Financial Knowledge Q&A", "Employee Information Query", "Weather Query"]
>>> classifier = IntentClassifier(classifier_llm, intent_list=chatflow_intent_list)
>>> classifier.start()
>>> print(classifier('What is the weather today'))
Weather Query
>>>
>>> with IntentClassifier(classifier_llm) as ic:
>>>     ic.case['Weather Query', lambda x: '38.5°C']
>>>     ic.case['Chat', lambda x: 'permission denied']
>>>     ic.case['Financial Knowledge Q&A', lambda x: 'Calling Financial RAG']
>>>     ic.case['Employee Information Query', lambda x: 'Beijing']
...
>>> ic.start()
>>> print(ic('What is the weather today'))
38.5°C
Source code in lazyllm/tools/classifier/intent_classifier.py
class IntentClassifier(ModuleBase):
    """Intent classification module that classifies input text into a given intent list.  
Supports automatic selection of Chinese or English prompt templates, and allows enhancement through examples, prompt text, constraints, and attention notes.

Args:
    llm: The large language model instance used for intent classification.
    intent_list (list): Optional, list of intent categories, e.g., ["chat", "weather", "QA"].
    prompt (str): Optional, custom prompt inserted into the system prompt template.
    constrain (str): Optional, classification constraint description.
    attention (str): Optional, attention notes for classification.
    examples (list[list[str, str]]): Optional, classification examples, each element is [input text, label].
    return_trace (bool): Whether to return execution trace. Default is False.


Examples:
        >>> import lazyllm
        >>> from lazyllm.tools import IntentClassifier
        >>> classifier_llm = lazyllm.OnlineChatModule(source="openai")
        >>> chatflow_intent_list = ["Chat", "Financial Knowledge Q&A", "Employee Information Query", "Weather Query"]
        >>> classifier = IntentClassifier(classifier_llm, intent_list=chatflow_intent_list)
        >>> classifier.start()
        >>> print(classifier('What is the weather today'))
        Weather Query
        >>>
        >>> with IntentClassifier(classifier_llm) as ic:
        >>>     ic.case['Weather Query', lambda x: '38.5°C']
        >>>     ic.case['Chat', lambda x: 'permission denied']
        >>>     ic.case['Financial Knowledge Q&A', lambda x: 'Calling Financial RAG']
        >>>     ic.case['Employee Information Query', lambda x: 'Beijing']
        ...
        >>> ic.start()
        >>> print(ic('What is the weather today'))
        38.5°C
    """
    def __init__(self, llm, intent_list: list = None,
                 *, prompt: str = '', constrain: str = '', attention: str = '',
                 examples: Optional[list[list[str, str]]] = None, return_trace: bool = False) -> None:
        super().__init__(return_trace=return_trace)
        self._intent_list = intent_list or []
        self._llm = llm
        self._prompt, self._constrain, self._attention, self._examples = prompt, constrain, attention, examples or []
        if self._intent_list:
            self._init()

    def _init(self):
        def choose_prompt():
            # Use chinese prompt if intent elements have chinese character, otherwise use english version
            for ele in self._intent_list:
                for ch in ele:
                    # chinese unicode range
                    if "\u4e00" <= ch <= "\u9fff":
                        return ch_prompt_classifier_template
            return en_prompt_classifier_template

        example_template = '\nUser: {{{{"human_input": "{inp}", "intent_list": {intent}}}}}\nAssistant: {label}\n'
        examples = ''.join([example_template.format(
            inp=input, intent=self._intent_list, label=label) for input, label in self._examples])
        prompt = choose_prompt().replace(
            '{user_prompt}', f' {self._prompt}').replace('{attention}', self._attention).replace(
            '{user_constrains}', f' {self._constrain}').replace('{user_examples}', f' {examples}')
        self._llm = self._llm.share(prompt=AlpacaPrompter(dict(system=prompt, user='${input}')
                                                          ).pre_hook(self.intent_promt_hook)).used_by(self._module_id)
        self._impl = pipeline(self._llm, self.post_process_result)

    def intent_promt_hook(
        self,
        input: Union[str, List, Dict[str, str], None] = None,
        history: List[Union[List[str], Dict[str, Any]]] = [],  # noqa B006
        tools: Union[List[Dict[str, Any]], None] = None,
        label: Union[str, None] = None,
    ):
        """Pre-processing hook for intent classification.  
Packages the input text and intent list into JSON and generates a string of conversation history.

Args:
    input (str | List | Dict | None): The input text, only string type is supported.
    history (List): Conversation history, default empty list.
    tools (List[Dict] | None): Optional tool information.
    label (str | None): Optional label.

**Returns**

- tuple: (input data dict, history list, tools, label)
"""
        input_json = {}
        if isinstance(input, str):
            input_json = {"human_input": input, "intent_list": self._intent_list}
        else:
            raise ValueError(f"Unexpected type for input: {type(input)}")

        history_info = chat_history_to_str(history)
        history = []
        input_text = json.dumps(input_json, ensure_ascii=False)
        return dict(history_info=history_info, input=input_text), history, tools, label

    def post_process_result(self, input):
        """Post-processing of intent classification result.  
Returns the result directly if it is in the intent list, otherwise returns the first element of the intent list.

Args:
    input (str): Output result from the classification model.

**Returns**

- str: The final classification label.
"""
        input = input.strip()
        return input if input in self._intent_list else self._intent_list[0]

    def forward(self, input: str, llm_chat_history: List[Dict[str, Any]] = None):
        if llm_chat_history is not None and self._llm._module_id not in globals["chat_history"]:
            globals["chat_history"][self._llm._module_id] = llm_chat_history
        return self._impl(input)

    def __enter__(self):
        assert not self._intent_list, 'Intent list is already set'
        self._sw = switch()
        self._sw.__enter__()
        return self

    @property
    def case(self):
        return switch.Case(self)

    @property
    def submodules(self):
        submodule = []
        if isinstance(self._impl, switch):
            self._impl.for_each(lambda x: isinstance(x, ModuleBase), lambda x: submodule.append(x))
        return super().submodules + submodule

    # used by switch.Case
    def _add_case(self, cond, func):
        assert isinstance(cond, str), 'intent must be string'
        self._intent_list.append(cond)
        self._sw.case[cond, func]

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._sw.__exit__(exc_type, exc_val, exc_tb)
        self._init()
        self._sw._set_conversion(self._impl)
        self._impl = self._sw

intent_promt_hook(input=None, history=[], tools=None, label=None)

Pre-processing hook for intent classification.
Packages the input text and intent list into JSON and generates a string of conversation history.

Parameters:

  • input (str | List | Dict | None, default: None ) –

    The input text, only string type is supported.

  • history (List, default: [] ) –

    Conversation history, default empty list.

  • tools (List[Dict] | None, default: None ) –

    Optional tool information.

  • label (str | None, default: None ) –

    Optional label.

Returns

  • tuple: (input data dict, history list, tools, label)
Source code in lazyllm/tools/classifier/intent_classifier.py
    def intent_promt_hook(
        self,
        input: Union[str, List, Dict[str, str], None] = None,
        history: List[Union[List[str], Dict[str, Any]]] = [],  # noqa B006
        tools: Union[List[Dict[str, Any]], None] = None,
        label: Union[str, None] = None,
    ):
        """Pre-processing hook for intent classification.  
Packages the input text and intent list into JSON and generates a string of conversation history.

Args:
    input (str | List | Dict | None): The input text, only string type is supported.
    history (List): Conversation history, default empty list.
    tools (List[Dict] | None): Optional tool information.
    label (str | None): Optional label.

**Returns**

- tuple: (input data dict, history list, tools, label)
"""
        input_json = {}
        if isinstance(input, str):
            input_json = {"human_input": input, "intent_list": self._intent_list}
        else:
            raise ValueError(f"Unexpected type for input: {type(input)}")

        history_info = chat_history_to_str(history)
        history = []
        input_text = json.dumps(input_json, ensure_ascii=False)
        return dict(history_info=history_info, input=input_text), history, tools, label

post_process_result(input)

Post-processing of intent classification result.
Returns the result directly if it is in the intent list, otherwise returns the first element of the intent list.

Parameters:

  • input (str) –

    Output result from the classification model.

Returns

  • str: The final classification label.
Source code in lazyllm/tools/classifier/intent_classifier.py
    def post_process_result(self, input):
        """Post-processing of intent classification result.  
Returns the result directly if it is in the intent list, otherwise returns the first element of the intent list.

Args:
    input (str): Output result from the classification model.

**Returns**

- str: The final classification label.
"""
        input = input.strip()
        return input if input in self._intent_list else self._intent_list[0]

lazyllm.tools.Document

Bases: ModuleBase, BuiltinGroups

Initialize a document module with an optional user interface.

This constructor initializes a document module that can have an optional user interface. If the user interface is enabled, it also provides a UI to manage the document operation interface and offers a web page for user interface interaction.

Parameters:

  • dataset_path (str, default: None ) –

    The path to the dataset directory. This directory should contain the documents to be managed by the document module.

  • embed (Optional[Union[Callable, Dict[str, Callable]]], default: None ) –

    The object used to generate document embeddings. If you need to generate multiple embeddings for the text, you need to specify multiple embedding models in a dictionary format. The key identifies the name corresponding to the embedding, and the value is the corresponding embedding model.

  • create_ui (bool, default: False ) –

    [Deprecated] Whether to create a user interface. Use 'manager' parameter instead.

  • manager (bool, default: False ) –

    A flag indicating whether to create a user interface for the document module. Defaults to False.

  • server (Union[bool, int], default: False ) –

    Server configuration. True for default server, False for no server, or an integer port number for custom server.

  • name (Optional[str], default: None ) –

    Name identifier for this document collection. Required for cloud services.

  • launcher (optional, default: None ) –

    An object or function responsible for launching the server module. If not provided, the default asynchronous launcher from lazyllm.launchers is used (sync=False).

  • doc_fields (optional, default: None ) –

    Configure the fields that need to be stored and retrieved along with their corresponding types (currently only used by the Milvus backend).

  • doc_files (Optional[List[str]], default: None ) –

    List of temporary document files (alternative to dataset_path).When used, dataset_path must be None and only map store is supported.

  • store_conf (optional, default: None ) –

    Configure which storage backend, MapStore is the default choice.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Document
>>> m = lazyllm.OnlineEmbeddingModule(source="glm")
>>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)  # or documents = Document(dataset_path='your_doc_path', embed={"key": m}, manager=False)
>>> m1 = lazyllm.TrainableModule("bge-large-zh-v1.5").start()
>>> document1 = Document(dataset_path='your_doc_path', embed={"online": m, "local": m1}, manager=False)
>>> store_conf = {
>>>     "segment_store": {
>>>         "type": "map",
>>>         "kwargs": {
>>>             "uri": "/tmp/tmp_segments.db",
>>>         },
>>>     },
>>>     "vector_store": {
>>>         "type": "milvus",
>>>         "kwargs": {
>>>             "uri": "/tmp/tmp_milvus.db",
>>>             "index_kwargs": {
>>>                 "index_type": "FLAT",
>>>                 "metric_type": "COSINE",
>>>             },
>>>         },
>>>     },
>>> }
>>> doc_fields = {
>>>     'author': DocField(data_type=DataType.VARCHAR, max_size=128, default_value=' '),
>>>     'public_year': DocField(data_type=DataType.INT32),
>>> }
>>> document2 = Document(dataset_path='your_doc_path', embed={"online": m, "local": m1}, store_conf=store_conf, doc_fields=doc_fields)
Source code in lazyllm/tools/rag/document.py
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
class Document(ModuleBase, BuiltinGroups, metaclass=_MetaDocument):
    """Initialize a document module with an optional user interface.

This constructor initializes a document module that can have an optional user interface. If the user interface is enabled, it also provides a UI to manage the document operation interface and offers a web page for user interface interaction.

Args:
    dataset_path (str): The path to the dataset directory. This directory should contain the documents to be managed by the document module.
    embed (Optional[Union[Callable, Dict[str, Callable]]]): The object used to generate document embeddings. If you need to generate multiple embeddings for the text, you need to specify multiple embedding models in a dictionary format. The key identifies the name corresponding to the embedding, and the value is the corresponding embedding model.
    create_ui (bool): [Deprecated] Whether to create a user interface. Use 'manager' parameter instead.
    manager (bool, optional): A flag indicating whether to create a user interface for the document module. Defaults to False.
    server (Union[bool, int]): Server configuration. True for default server, False for no server, or an integer port number for custom server.
    name (Optional[str]): Name identifier for this document collection. Required for cloud services.
    launcher (optional): An object or function responsible for launching the server module. If not provided, the default asynchronous launcher from `lazyllm.launchers` is used (`sync=False`).
    doc_fields (optional): Configure the fields that need to be stored and retrieved along with their corresponding types (currently only used by the Milvus backend).
    doc_files (Optional[List[str]]): List of temporary document files (alternative to dataset_path).When used, dataset_path must be None and only map store is supported.
    store_conf (optional): Configure which storage backend, MapStore is the default choice.      


Examples:
    >>> import lazyllm
    >>> from lazyllm.tools import Document
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)  # or documents = Document(dataset_path='your_doc_path', embed={"key": m}, manager=False)
    >>> m1 = lazyllm.TrainableModule("bge-large-zh-v1.5").start()
    >>> document1 = Document(dataset_path='your_doc_path', embed={"online": m, "local": m1}, manager=False)

    >>> store_conf = {
    >>>     "segment_store": {
    >>>         "type": "map",
    >>>         "kwargs": {
    >>>             "uri": "/tmp/tmp_segments.db",
    >>>         },
    >>>     },
    >>>     "vector_store": {
    >>>         "type": "milvus",
    >>>         "kwargs": {
    >>>             "uri": "/tmp/tmp_milvus.db",
    >>>             "index_kwargs": {
    >>>                 "index_type": "FLAT",
    >>>                 "metric_type": "COSINE",
    >>>             },
    >>>         },
    >>>     },
    >>> }
    >>> doc_fields = {
    >>>     'author': DocField(data_type=DataType.VARCHAR, max_size=128, default_value=' '),
    >>>     'public_year': DocField(data_type=DataType.INT32),
    >>> }
    >>> document2 = Document(dataset_path='your_doc_path', embed={"online": m, "local": m1}, store_conf=store_conf, doc_fields=doc_fields)
    """
    class _Manager(ModuleBase):
        def __init__(self, dataset_path: Optional[str], embed: Optional[Union[Callable, Dict[str, Callable]]] = None,
                     manager: Union[bool, str] = False, server: Union[bool, int] = False, name: Optional[str] = None,
                     launcher: Optional[Launcher] = None, store_conf: Optional[Dict] = None,
                     doc_fields: Optional[Dict[str, DocField]] = None, cloud: bool = False,
                     doc_files: Optional[List[str]] = None, processor: Optional[DocumentProcessor] = None,
                     display_name: Optional[str] = "", description: Optional[str] = "algorithm description"):
            super().__init__()
            self._origin_path, self._doc_files, self._cloud = dataset_path, doc_files, cloud

            if dataset_path and not os.path.exists(dataset_path):
                defatult_path = os.path.join(lazyllm.config["data_path"], dataset_path)
                if os.path.exists(defatult_path):
                    dataset_path = defatult_path
            elif dataset_path:
                dataset_path = os.path.join(os.getcwd(), dataset_path)

            self._launcher: Launcher = launcher if launcher else lazyllm.launchers.remote(sync=False)
            self._dataset_path = dataset_path
            self._embed = self._get_embeds(embed)
            self._processor = processor
            name = name or DocListManager.DEFAULT_GROUP_NAME
            if not display_name: display_name = name

            self._dlm = None if (self._cloud or self._doc_files is not None) else DocListManager(
                dataset_path, name, enable_path_monitoring=False if manager else True)
            self._kbs = CallableDict({name: DocImpl(
                embed=self._embed, dlm=self._dlm, doc_files=doc_files, global_metadata_desc=doc_fields,
                store=store_conf, processor=processor, algo_name=name, display_name=display_name,
                description=description)})

            if manager: self._manager = ServerModule(DocManager(self._dlm), launcher=self._launcher)
            if manager == 'ui': self._docweb = DocWebModule(doc_server=self._manager)
            if server: self._kbs = ServerModule(self._kbs, port=(None if isinstance(server, bool) else int(server)))
            self._global_metadata_desc = doc_fields

        @property
        def url(self):
            if hasattr(self, '_manager'): return self._manager._url
            return None

        @property
        @deprecated('Document.manager.url')
        def _url(self):
            return self.url

        @property
        def web_url(self):
            if hasattr(self, '_docweb'): return self._docweb.url
            return None

        def _get_embeds(self, embed):
            embeds = embed if isinstance(embed, dict) else {EMBED_DEFAULT_KEY: embed} if embed else {}
            for embed in embeds.values():
                if isinstance(embed, ModuleBase):
                    self._submodules.append(embed)
            return embeds

        def add_kb_group(self, name, doc_fields: Optional[Dict[str, DocField]] = None,
                         store_conf: Optional[Dict] = None,
                         embed: Optional[Union[Callable, Dict[str, Callable]]] = None):
            embed = self._get_embeds(embed) if embed else self._embed
            if isinstance(self._kbs, ServerModule):
                self._kbs._impl._m[name] = DocImpl(dlm=self._dlm, embed=embed, kb_group_name=name,
                                                   global_metadata_desc=doc_fields, store=store_conf)
            else:
                self._kbs[name] = DocImpl(dlm=self._dlm, embed=self._embed, kb_group_name=name,
                                          global_metadata_desc=doc_fields, store=store_conf)
            self._dlm.add_kb_group(name=name)

        def get_doc_by_kb_group(self, name):
            return self._kbs._impl._m[name] if isinstance(self._kbs, ServerModule) else self._kbs[name]

        def stop(self):
            if hasattr(self, '_docweb'):
                self._docweb.stop()
            self._launcher.cleanup()

        def __call__(self, *args, **kw):
            return self._kbs(*args, **kw)

    def __new__(cls, *args, **kw):
        if url := kw.pop('url', None):
            name = kw.pop('name', None)
            assert not args and not kw, 'Only `name` is supported with `url`'
            return UrlDocument(url, name)
        else:
            return super().__new__(cls)

    def __init__(self, dataset_path: Optional[str] = None, embed: Optional[Union[Callable, Dict[str, Callable]]] = None,
                 create_ui: bool = False, manager: Union[bool, str, "Document._Manager", DocumentProcessor] = False,
                 server: Union[bool, int] = False, name: Optional[str] = None, launcher: Optional[Launcher] = None,
                 doc_files: Optional[List[str]] = None, doc_fields: Dict[str, DocField] = None,
                 store_conf: Optional[Dict] = None, display_name: Optional[str] = "",
                 description: Optional[str] = "algorithm description"):
        super().__init__()
        if create_ui:
            lazyllm.LOG.warning('`create_ui` for Document is deprecated, use `manager` instead')
            manager = create_ui
        if isinstance(dataset_path, (tuple, list)):
            doc_fields = dataset_path
            dataset_path = None
        if doc_files is not None:
            assert dataset_path is None and not manager, (
                'Manager and dataset_path are not supported for Document with temp-files')
            assert store_conf is None or store_conf['type'] == 'map', (
                'Only map store is supported for Document with temp-files')

        name = name or DocListManager.DEFAULT_GROUP_NAME

        if isinstance(manager, Document._Manager):
            assert not server, 'Server infomation is already set to by manager'
            assert not launcher, 'Launcher infomation is already set to by manager'
            assert not manager._cloud, 'manager is not allowed to share in cloud mode'
            assert manager._doc_files is None, 'manager is not allowed to share with temp files'
            if dataset_path != manager._dataset_path and dataset_path != manager._origin_path:
                raise RuntimeError(f'Document path mismatch, expected `{manager._dataset_path}`'
                                   f'while received `{dataset_path}`')
            manager.add_kb_group(name=name, doc_fields=doc_fields, store_conf=store_conf, embed=embed)
            self._manager = manager
            self._curr_group = name
        else:
            if isinstance(manager, DocumentProcessor):
                processor, cloud = manager, True
                processor._impl.start()
                manager = False
                assert name, '`Name` of Document is necessary when using cloud service'
                assert store_conf.get('type') != 'map', 'Cloud manager is not supported when using map store'
                assert not dataset_path, 'Cloud manager is not supported with local dataset path'
            else:
                cloud, processor = False, None
            self._manager = Document._Manager(dataset_path, embed, manager, server, name, launcher, store_conf,
                                              doc_fields, cloud=cloud, doc_files=doc_files, processor=processor,
                                              display_name=display_name, description=description)
            self._curr_group = name
        self._doc_to_db_processor: DocToDbProcessor = None

    def _list_all_files_in_dataset(self) -> List[str]:
        files_list = []
        for root, _, files in os.walk(self._manager._dataset_path):
            files = [os.path.join(root, file_path) for file_path in files]
            files_list.extend(files)
        return files_list

    def connect_sql_manager(
        self,
        sql_manager: SqlManager,
        schma: Optional[DocInfoSchema] = None,
        force_refresh: bool = True,
    ):
        def format_schema_to_dict(schema: DocInfoSchema):
            if schema is None:
                return None, None
            desc_dict = {ele["key"]: ele["desc"] for ele in schema}
            type_dict = {ele["key"]: ele["type"] for ele in schema}
            return desc_dict, type_dict

        def compare_schema(old_schema: DocInfoSchema, new_schema: DocInfoSchema):
            old_desc_dict, old_type_dict = format_schema_to_dict(old_schema)
            new_desc_dict, new_type_dict = format_schema_to_dict(new_schema)
            return old_desc_dict == new_desc_dict and old_type_dict == new_type_dict

        # 1. Check valid arguments
        if sql_manager.check_connection().status != DBStatus.SUCCESS:
            raise RuntimeError(f'Failed to connect to sql manager: {sql_manager._gen_conn_url()}')
        pre_doc_table_schema = None
        if self._doc_to_db_processor:
            pre_doc_table_schema = self._doc_to_db_processor.doc_info_schema
        assert pre_doc_table_schema or schma, "doc_table_schma must be given"

        schema_equal = compare_schema(pre_doc_table_schema, schma)
        assert (
            schema_equal or force_refresh is True
        ), "When changing doc_table_schema, force_refresh should be set to True"

        # 2. Init handler if needed
        need_init_processor = False
        if self._doc_to_db_processor is None:
            need_init_processor = True
        else:
            # avoid reinit for the same db
            if sql_manager != self._doc_to_db_processor.sql_manager:
                need_init_processor = True
        if need_init_processor:
            self._doc_to_db_processor = DocToDbProcessor(sql_manager)

        # 3. Reset doc_table_schema if needed
        if schma and not schema_equal:
            # This api call will clear existing db table "lazyllm_doc_elements"
            self._doc_to_db_processor._reset_doc_info_schema(schma)

    def get_sql_manager(self):
        if self._doc_to_db_processor is None:
            raise ValueError("Please call connect_sql_manager to init handler first")
        return self._doc_to_db_processor.sql_manager

    def extract_db_schema(
        self, llm: Union[OnlineChatModule, TrainableModule], print_schema: bool = False
    ) -> DocInfoSchema:
        file_paths = self._list_all_files_in_dataset()
        schema = extract_db_schema_from_files(file_paths, llm)
        if print_schema:
            lazyllm.LOG.info(f"Extracted Schema:\n\t{schema}\n")
        return schema

    def update_database(self, llm: Union[OnlineChatModule, TrainableModule]):
        assert self._doc_to_db_processor, "Please call connect_db to init handler first"
        file_paths = self._list_all_files_in_dataset()
        info_dicts = self._doc_to_db_processor.extract_info_from_docs(llm, file_paths)
        self._doc_to_db_processor.export_info_to_db(info_dicts)

    @deprecated('Document(dataset_path, manager=doc.manager, name=xx, doc_fields=xx, store_conf=xx)')
    def create_kb_group(self, name: str, doc_fields: Optional[Dict[str, DocField]] = None,
                        store_conf: Optional[Dict] = None) -> "Document":
        self._manager.add_kb_group(name=name, doc_fields=doc_fields, store_conf=store_conf)
        doc = copy.copy(self)
        doc._curr_group = name
        return doc

    @property
    @deprecated('Document._manager')
    def _impls(self): return self._manager

    @property
    def _impl(self) -> DocImpl: return self._manager.get_doc_by_kb_group(self._curr_group)

    @property
    def manager(self): return self._manager._processor or self._manager

    def activate_group(self, group_name: str, embed_keys: Optional[Union[str, List[str]]] = None):
        if isinstance(embed_keys, str): embed_keys = [embed_keys]
        elif embed_keys is None: embed_keys = []
        self._impl.activate_group(group_name, embed_keys)

    def activate_groups(self, groups: Union[str, List[str]]):
        if isinstance(groups, str): groups = [groups]
        for group in groups:
            self.activate_group(group)

    @DynamicDescriptor
    def create_node_group(self, name: str = None, *, transform: Callable, parent: str = LAZY_ROOT_NAME,
                          trans_node: bool = None, num_workers: int = 0, display_name: str = None,
                          group_type: NodeGroupType = NodeGroupType.CHUNK, **kwargs) -> None:
        """
Generate a node group produced by the specified rule.

Args:
    name (str): The name of the node group.
    transform (Callable): The transformation rule that converts a node into a node group. The function prototype is `(DocNode, group_name, **kwargs) -> List[DocNode]`. Currently built-in options include [SentenceSplitter][lazyllm.tools.SentenceSplitter], and users can define their own transformation rules.
    trans_node (bool): Determines whether the input and output of transform are `DocNode` or `str`, default is None. Can only be set to true when `transform` is `Callable`.
    num_workers (int): number of new threads used for transform. default: 0
    parent (str): The node that needs further transformation. The series of new nodes obtained after transformation will be child nodes of this parent node. If not specified, the transformation starts from the root node.
    kwargs: Parameters related to the specific implementation.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
    >>> documents.create_node_group(name="sentences", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    """
        if isinstance(self, type):
            DocImpl.create_global_node_group(name, transform=transform, parent=parent, trans_node=trans_node,
                                             num_workers=num_workers, display_name=display_name,
                                             group_type=group_type, **kwargs)
        else:
            self._impl.create_node_group(name, transform=transform, parent=parent, trans_node=trans_node,
                                         num_workers=num_workers, display_name=display_name, group_type=group_type,
                                         **kwargs)

    @DynamicDescriptor
    def add_reader(self, pattern: str, func: Optional[Callable] = None):
        """
Used to specify the file reader for an instance. The scope of action is visible only to the registered Document object. The registered file reader must be a Callable object. It can only be registered by calling a function. The priority of the file reader registered by the instance is higher than that of the file reader registered by the class, and the priority of the file reader registered by the instance and class is higher than the system default file reader. That is, the order of priority is: instance file reader > class file reader > system default file reader.

Args:
    pattern (str): Matching rules applied by the file reader.
    func (Callable): File reader, must be a Callable object.


Examples:

    >>> from lazyllm.tools.rag import Document, DocNode
    >>> from lazyllm.tools.rag.readers import ReaderBase
    >>> class YmlReader(ReaderBase):
    ...     def _load_data(self, file, fs=None):
    ...         try:
    ...             import yaml
    ...         except ImportError:
    ...             raise ImportError("yaml is required to read YAML file: `pip install pyyaml`")
    ...         with open(file, 'r') as f:
    ...             data = yaml.safe_load(f)
    ...         print("Call the class YmlReader.")
    ...         return [DocNode(text=data)]
    ...
    >>> def processYml(file):
    ...     with open(file, 'r') as f:
    ...         data = f.read()
    ...     print("Call the function processYml.")
    ...     return [DocNode(text=data)]
    ...
    >>> doc1 = Document(dataset_path="your_files_path", create_ui=False)
    >>> doc2 = Document(dataset_path="your_files_path", create_ui=False)
    >>> doc1.add_reader("**/*.yml", YmlReader)
    >>> print(doc1._impl._local_file_reader)
    {'**/*.yml': <class '__main__.YmlReader'>}
    >>> print(doc2._impl._local_file_reader)
    {}
    >>> files = ["your_yml_files"]
    >>> Document.register_global_reader("**/*.yml", processYml)
    >>> doc1._impl._reader.load_data(input_files=files)
    Call the class YmlReader.
    >>> doc2._impl._reader.load_data(input_files=files)
    Call the function processYml.
    """
        if isinstance(self, type):
            return DocImpl.register_global_reader(pattern=pattern, func=func)
        else:
            self._impl.add_reader(pattern, func)

    @classmethod
    def register_global_reader(cls, pattern: str, func: Optional[Callable] = None):
        """
Used to specify a file reader, which is visible to all Document objects. The registered file reader must be a Callable object. It can be registered using a decorator or by a function call.

Args:
    pattern (str): Matching rules applied by the file reader.
    func (Callable): File reader, must be a Callable object.


Examples:

    >>> from lazyllm.tools.rag import Document, DocNode
    >>> @Document.register_global_reader("**/*.yml")
    >>> def processYml(file):
    ...     with open(file, 'r') as f:
    ...         data = f.read()
    ...     return [DocNode(text=data)]
    ...
    >>> doc1 = Document(dataset_path="your_files_path", create_ui=False)
    >>> doc2 = Document(dataset_path="your_files_path", create_ui=False)
    >>> files = ["your_yml_files"]
    >>> docs1 = doc1._impl._reader.load_data(input_files=files)
    >>> docs2 = doc2._impl._reader.load_data(input_files=files)
    >>> print(docs1[0].text == docs2[0].text)
    # True
    """
        return cls.add_reader(pattern, func)

    def get_store(self):
        return StorePlaceholder()

    def get_embed(self):
        return EmbedPlaceholder()

    def register_index(self, index_type: str, index_cls: IndexBase, *args, **kwargs) -> None:
        self._impl.register_index(index_type, index_cls, *args, **kwargs)

    def _forward(self, func_name: str, *args, **kw):
        return self._manager(self._curr_group, func_name, *args, **kw)

    def find_parent(self, target) -> Callable:
        """
Find the parent node of the specified node.

Args:
    group (str): The name of the node group for which to find the parent.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
    >>> documents.create_node_group(name="parent", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    >>> documents.create_node_group(name="children", transform=SentenceSplitter, parent="parent", chunk_size=1024, chunk_overlap=100)
    >>> documents.find_parent('children')
    """
        return functools.partial(self._forward, 'find_parent', group=target)

    def find_children(self, target) -> Callable:
        """
Find the child nodes of the specified node.

Args:
    group (str): The name of the node group for which to find the children.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
    >>> documents.create_node_group(name="parent", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    >>> documents.create_node_group(name="children", transform=SentenceSplitter, parent="parent", chunk_size=1024, chunk_overlap=100)
    >>> documents.find_children('parent')
    """
        return functools.partial(self._forward, 'find_children', group=target)

    def find(self, target) -> Callable:
        return functools.partial(self._forward, 'find', group=target)

    def forward(self, *args, **kw) -> List[DocNode]:
        return self._forward('retrieve', *args, **kw)

    def clear_cache(self, group_names: Optional[List[str]] = None) -> None:
        return self._forward('clear_cache', group_names)

    def _get_post_process_tasks(self):
        return lazyllm.pipeline(lambda *a: self._forward('_lazy_init'))

    def __repr__(self):
        return lazyllm.make_repr("Module", "Document", manager=hasattr(self._manager, '_manager'),
                                 server=isinstance(self._manager._kbs, ServerModule))

add_reader(pattern, func=None)

Used to specify the file reader for an instance. The scope of action is visible only to the registered Document object. The registered file reader must be a Callable object. It can only be registered by calling a function. The priority of the file reader registered by the instance is higher than that of the file reader registered by the class, and the priority of the file reader registered by the instance and class is higher than the system default file reader. That is, the order of priority is: instance file reader > class file reader > system default file reader.

Parameters:

  • pattern (str) –

    Matching rules applied by the file reader.

  • func (Callable, default: None ) –

    File reader, must be a Callable object.

Examples:

>>> from lazyllm.tools.rag import Document, DocNode
>>> from lazyllm.tools.rag.readers import ReaderBase
>>> class YmlReader(ReaderBase):
...     def _load_data(self, file, fs=None):
...         try:
...             import yaml
...         except ImportError:
...             raise ImportError("yaml is required to read YAML file: `pip install pyyaml`")
...         with open(file, 'r') as f:
...             data = yaml.safe_load(f)
...         print("Call the class YmlReader.")
...         return [DocNode(text=data)]
...
>>> def processYml(file):
...     with open(file, 'r') as f:
...         data = f.read()
...     print("Call the function processYml.")
...     return [DocNode(text=data)]
...
>>> doc1 = Document(dataset_path="your_files_path", create_ui=False)
>>> doc2 = Document(dataset_path="your_files_path", create_ui=False)
>>> doc1.add_reader("**/*.yml", YmlReader)
>>> print(doc1._impl._local_file_reader)
{'**/*.yml': <class '__main__.YmlReader'>}
>>> print(doc2._impl._local_file_reader)
{}
>>> files = ["your_yml_files"]
>>> Document.register_global_reader("**/*.yml", processYml)
>>> doc1._impl._reader.load_data(input_files=files)
Call the class YmlReader.
>>> doc2._impl._reader.load_data(input_files=files)
Call the function processYml.
Source code in lazyllm/tools/rag/document.py
    @DynamicDescriptor
    def add_reader(self, pattern: str, func: Optional[Callable] = None):
        """
Used to specify the file reader for an instance. The scope of action is visible only to the registered Document object. The registered file reader must be a Callable object. It can only be registered by calling a function. The priority of the file reader registered by the instance is higher than that of the file reader registered by the class, and the priority of the file reader registered by the instance and class is higher than the system default file reader. That is, the order of priority is: instance file reader > class file reader > system default file reader.

Args:
    pattern (str): Matching rules applied by the file reader.
    func (Callable): File reader, must be a Callable object.


Examples:

    >>> from lazyllm.tools.rag import Document, DocNode
    >>> from lazyllm.tools.rag.readers import ReaderBase
    >>> class YmlReader(ReaderBase):
    ...     def _load_data(self, file, fs=None):
    ...         try:
    ...             import yaml
    ...         except ImportError:
    ...             raise ImportError("yaml is required to read YAML file: `pip install pyyaml`")
    ...         with open(file, 'r') as f:
    ...             data = yaml.safe_load(f)
    ...         print("Call the class YmlReader.")
    ...         return [DocNode(text=data)]
    ...
    >>> def processYml(file):
    ...     with open(file, 'r') as f:
    ...         data = f.read()
    ...     print("Call the function processYml.")
    ...     return [DocNode(text=data)]
    ...
    >>> doc1 = Document(dataset_path="your_files_path", create_ui=False)
    >>> doc2 = Document(dataset_path="your_files_path", create_ui=False)
    >>> doc1.add_reader("**/*.yml", YmlReader)
    >>> print(doc1._impl._local_file_reader)
    {'**/*.yml': <class '__main__.YmlReader'>}
    >>> print(doc2._impl._local_file_reader)
    {}
    >>> files = ["your_yml_files"]
    >>> Document.register_global_reader("**/*.yml", processYml)
    >>> doc1._impl._reader.load_data(input_files=files)
    Call the class YmlReader.
    >>> doc2._impl._reader.load_data(input_files=files)
    Call the function processYml.
    """
        if isinstance(self, type):
            return DocImpl.register_global_reader(pattern=pattern, func=func)
        else:
            self._impl.add_reader(pattern, func)

create_node_group(name=None, *, transform, parent=LAZY_ROOT_NAME, trans_node=None, num_workers=0, display_name=None, group_type=NodeGroupType.CHUNK, **kwargs)

Generate a node group produced by the specified rule.

Parameters:

  • name (str, default: None ) –

    The name of the node group.

  • transform (Callable) –

    The transformation rule that converts a node into a node group. The function prototype is (DocNode, group_name, **kwargs) -> List[DocNode]. Currently built-in options include SentenceSplitter, and users can define their own transformation rules.

  • trans_node (bool, default: None ) –

    Determines whether the input and output of transform are DocNode or str, default is None. Can only be set to true when transform is Callable.

  • num_workers (int, default: 0 ) –

    number of new threads used for transform. default: 0

  • parent (str, default: LAZY_ROOT_NAME ) –

    The node that needs further transformation. The series of new nodes obtained after transformation will be child nodes of this parent node. If not specified, the transformation starts from the root node.

  • kwargs

    Parameters related to the specific implementation.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Document, SentenceSplitter
>>> m = lazyllm.OnlineEmbeddingModule(source="glm")
>>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
>>> documents.create_node_group(name="sentences", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
Source code in lazyllm/tools/rag/document.py
    @DynamicDescriptor
    def create_node_group(self, name: str = None, *, transform: Callable, parent: str = LAZY_ROOT_NAME,
                          trans_node: bool = None, num_workers: int = 0, display_name: str = None,
                          group_type: NodeGroupType = NodeGroupType.CHUNK, **kwargs) -> None:
        """
Generate a node group produced by the specified rule.

Args:
    name (str): The name of the node group.
    transform (Callable): The transformation rule that converts a node into a node group. The function prototype is `(DocNode, group_name, **kwargs) -> List[DocNode]`. Currently built-in options include [SentenceSplitter][lazyllm.tools.SentenceSplitter], and users can define their own transformation rules.
    trans_node (bool): Determines whether the input and output of transform are `DocNode` or `str`, default is None. Can only be set to true when `transform` is `Callable`.
    num_workers (int): number of new threads used for transform. default: 0
    parent (str): The node that needs further transformation. The series of new nodes obtained after transformation will be child nodes of this parent node. If not specified, the transformation starts from the root node.
    kwargs: Parameters related to the specific implementation.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
    >>> documents.create_node_group(name="sentences", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    """
        if isinstance(self, type):
            DocImpl.create_global_node_group(name, transform=transform, parent=parent, trans_node=trans_node,
                                             num_workers=num_workers, display_name=display_name,
                                             group_type=group_type, **kwargs)
        else:
            self._impl.create_node_group(name, transform=transform, parent=parent, trans_node=trans_node,
                                         num_workers=num_workers, display_name=display_name, group_type=group_type,
                                         **kwargs)

find_children(target)

Find the child nodes of the specified node.

Parameters:

  • group (str) –

    The name of the node group for which to find the children.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Document, SentenceSplitter
>>> m = lazyllm.OnlineEmbeddingModule(source="glm")
>>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
>>> documents.create_node_group(name="parent", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
>>> documents.create_node_group(name="children", transform=SentenceSplitter, parent="parent", chunk_size=1024, chunk_overlap=100)
>>> documents.find_children('parent')
Source code in lazyllm/tools/rag/document.py
    def find_children(self, target) -> Callable:
        """
Find the child nodes of the specified node.

Args:
    group (str): The name of the node group for which to find the children.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
    >>> documents.create_node_group(name="parent", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    >>> documents.create_node_group(name="children", transform=SentenceSplitter, parent="parent", chunk_size=1024, chunk_overlap=100)
    >>> documents.find_children('parent')
    """
        return functools.partial(self._forward, 'find_children', group=target)

find_parent(target)

Find the parent node of the specified node.

Parameters:

  • group (str) –

    The name of the node group for which to find the parent.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Document, SentenceSplitter
>>> m = lazyllm.OnlineEmbeddingModule(source="glm")
>>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
>>> documents.create_node_group(name="parent", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
>>> documents.create_node_group(name="children", transform=SentenceSplitter, parent="parent", chunk_size=1024, chunk_overlap=100)
>>> documents.find_parent('children')
Source code in lazyllm/tools/rag/document.py
    def find_parent(self, target) -> Callable:
        """
Find the parent node of the specified node.

Args:
    group (str): The name of the node group for which to find the parent.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
    >>> documents.create_node_group(name="parent", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    >>> documents.create_node_group(name="children", transform=SentenceSplitter, parent="parent", chunk_size=1024, chunk_overlap=100)
    >>> documents.find_parent('children')
    """
        return functools.partial(self._forward, 'find_parent', group=target)

register_global_reader(pattern, func=None) classmethod

Used to specify a file reader, which is visible to all Document objects. The registered file reader must be a Callable object. It can be registered using a decorator or by a function call.

Parameters:

  • pattern (str) –

    Matching rules applied by the file reader.

  • func (Callable, default: None ) –

    File reader, must be a Callable object.

Examples:

>>> from lazyllm.tools.rag import Document, DocNode
>>> @Document.register_global_reader("**/*.yml")
>>> def processYml(file):
...     with open(file, 'r') as f:
...         data = f.read()
...     return [DocNode(text=data)]
...
>>> doc1 = Document(dataset_path="your_files_path", create_ui=False)
>>> doc2 = Document(dataset_path="your_files_path", create_ui=False)
>>> files = ["your_yml_files"]
>>> docs1 = doc1._impl._reader.load_data(input_files=files)
>>> docs2 = doc2._impl._reader.load_data(input_files=files)
>>> print(docs1[0].text == docs2[0].text)
# True
Source code in lazyllm/tools/rag/document.py
    @classmethod
    def register_global_reader(cls, pattern: str, func: Optional[Callable] = None):
        """
Used to specify a file reader, which is visible to all Document objects. The registered file reader must be a Callable object. It can be registered using a decorator or by a function call.

Args:
    pattern (str): Matching rules applied by the file reader.
    func (Callable): File reader, must be a Callable object.


Examples:

    >>> from lazyllm.tools.rag import Document, DocNode
    >>> @Document.register_global_reader("**/*.yml")
    >>> def processYml(file):
    ...     with open(file, 'r') as f:
    ...         data = f.read()
    ...     return [DocNode(text=data)]
    ...
    >>> doc1 = Document(dataset_path="your_files_path", create_ui=False)
    >>> doc2 = Document(dataset_path="your_files_path", create_ui=False)
    >>> files = ["your_yml_files"]
    >>> docs1 = doc1._impl._reader.load_data(input_files=files)
    >>> docs2 = doc2._impl._reader.load_data(input_files=files)
    >>> print(docs1[0].text == docs2[0].text)
    # True
    """
        return cls.add_reader(pattern, func)

lazyllm.tools.rag.store.ChromadbStore

Bases: LazyLLMStoreBase

Source code in lazyllm/tools/rag/store/vector/chroma_store.py
class ChromadbStore(LazyLLMStoreBase):
    capability = StoreCapability.VECTOR
    need_embedding = True
    supports_index_registration = False

    def __init__(self, uri: Optional[str] = None, dir: Optional[str] = None,
                 index_kwargs: Optional[Union[Dict, List]] = None, client_kwargs: Optional[Dict] = None,
                 **kwargs) -> None:
        assert uri or (dir), "uri or dir must be provided"
        self._index_kwargs = index_kwargs or DEFAULT_INDEX_CONFIG
        self._client_kwargs = client_kwargs or {}
        if dir:
            self._dir = dir
        else:
            self._dir, self._host, self._port = self._parse_uri(uri)
        self._primary_key = 'uid'

    @property
    def dir(self):
        if not self._dir: return None
        p = Path(self._dir)
        p = p if p.suffix else (p / "chroma.sqlite3")
        return str(p.resolve(strict=False))

    def _parse_uri(self, uri: str):
        windows_drive = re.match(r"^[a-zA-Z]:[\\/]", uri or "")
        if ("://" not in uri) and (windows_drive or os.path.isabs(uri)):
            return os.path.abspath(uri), None, None

        p = urlparse(uri)

        if p.scheme == "":
            return os.path.abspath(uri), None, None

        if p.scheme == "file":
            path = p.path
            if os.name == "nt" and path.startswith("/") and re.match(r"^/[a-zA-Z]:", path):
                path = path.lstrip("/")  # file:///C:/... -> C:/...
            return os.path.abspath(path), None, None

        scheme = p.scheme
        if scheme.startswith("chroma+"):
            scheme = scheme.split("+", 1)[1]  # http or https

        if scheme in ("http", "https"):
            host = p.hostname or "127.0.0.1"
            port = p.port or (443 if scheme == "https" else 80)
            return None, host, port

        raise ValueError(f"Unsupported URI scheme in '{uri}'. "
                         "Use file:///path or plain path for local; http(s)://host:port for remote.")

    @override
    def connect(self, embed_dims: Optional[Dict[str, int]] = None,
                embed_datatypes: Optional[Dict[str, DataType]] = None,
                global_metadata_desc: Optional[Dict[str, GlobalMetadataDesc]] = None, **kwargs):
        self._global_metadata_desc = global_metadata_desc or {}
        self._embed_dims = embed_dims or {}
        self._embed_datatypes = embed_datatypes or {}
        for k, v in self._global_metadata_desc.items():
            if v.data_type not in [DataType.VARCHAR, DataType.INT32, DataType.FLOAT, DataType.BOOLEAN]:
                raise ValueError(f"[Chromadb Store] Unsupported data type {v.data_type} for global metadata {k}"
                                 " (only string, int, float, bool are supported)")
        for k, v in self._embed_datatypes.items():
            if v not in [DataType.FLOAT_VECTOR, DataType.SPARSE_FLOAT_VECTOR]:
                raise ValueError(f"[Chromadb Store] Unsupported data type {v} for embed key {k}"
                                 " (only float vector and sparse float vector are supported)")
        if self._dir:
            self._client = chromadb.PersistentClient(path=self._dir, **self._client_kwargs)
            LOG.success(f"Initialzed chromadb in path: {self._dir}")
        else:
            self._client = chromadb.HttpClient(host=self._host, port=self._port, **self._client_kwargs)
            LOG.success(f"Initialzed chromadb in host: {self._host}, port: {self._port}")

    @override
    def upsert(self, collection_name: str, data: List[dict]) -> bool:
        try:
            # NOTE chromadb only support single embedding for each collection
            if not data: return
            data_embeddings = data[0].get('embedding', {})
            if not data_embeddings: return
            embed_keys = list(data_embeddings.keys())
            for embed_key in embed_keys:
                if embed_key not in self._embed_datatypes:
                    raise ValueError(f"Embed key {embed_key} not found in embed_datatypes")
                collection = self._client.get_or_create_collection(
                    name=self._gen_collection_name(collection_name, embed_key), configuration=self._index_kwargs)
                for i in range(0, len(data), INSERT_BATCH_SIZE):
                    collection.upsert(**self._serialize_data(data[i: i + INSERT_BATCH_SIZE], embed_key))
            return True
        except Exception as e:
            LOG.error(f"[Chromadb Store - upsert] Failed to create collection {collection_name}: {e}")
            LOG.error(traceback.format_exc())
            return False

    def _serialize_data(self, data: List[dict], embed_key: str) -> List[dict]:
        res = {'ids': [], 'embeddings': [], 'metadatas': []}
        for d in data:
            res['ids'].append(d.get('uid'))
            res['embeddings'].append(d.get('embedding', {}).get(embed_key))
            res['metadatas'].append({self._gen_global_meta_key(k): v for k, v in d.get('global_meta', {}).items()
                                     if k in self._global_metadata_desc})
        return res

    @override
    def delete(self, collection_name: str, criteria: Optional[dict] = None, **kwargs) -> bool:
        try:
            if not criteria:
                for embed_key in self._embed_datatypes.keys():
                    try:
                        self._client.delete_collection(name=self._gen_collection_name(collection_name, embed_key))
                    except Exception:
                        continue
                return True
            else:
                filters = self._construct_criteria(criteria)
                for embed_key in self._embed_datatypes.keys():
                    collection = self._client.get_collection(name=self._gen_collection_name(collection_name, embed_key))
                    collection.delete(**filters)
                return True
        except Exception as e:
            LOG.error(f"[Chromadb Store - delete] Failed to delete collection {collection_name}: {e}")
            LOG.error(traceback.format_exc())
            return False

    @override
    def get(self, collection_name: str, criteria: Optional[dict] = None, **kwargs) -> List[dict]:
        try:
            filters = self._construct_criteria(criteria) if criteria else {}
            all_data = []
            for key in self._embed_datatypes:
                try:
                    coll = self._client.get_collection(
                        name=self._gen_collection_name(collection_name, key)
                    )
                    data = coll.get(include=['metadatas', 'embeddings'], **filters)
                    all_data.append((key, data))
                except Exception:
                    continue

            res: Dict[str, Dict[str, Any]] = defaultdict(lambda: {
                'uid': None, 'global_meta': {}, 'embedding': {}})
            for embed_key, data in all_data:
                ids = data['ids']
                metas = data['metadatas']
                embs = data['embeddings']

                for uid, meta, emb in zip(ids, metas, embs):
                    entry = res[uid]
                    entry['uid'] = uid
                    if not entry['global_meta']:
                        entry['global_meta'] = {
                            k[len(GLOBAL_META_KEY_PREFIX):]: v
                            for k, v in meta.items()
                        }
                    entry['embedding'][embed_key] = list(emb)
            return list(res.values())
        except Exception as e:
            LOG.error(f"[ChromadbStore - get] task fail: {e}")
            LOG.error(traceback.format_exc())

    @override
    def search(self, collection_name: str, query_embedding: List[float], embed_key: str, topk: Optional[int] = 10,
               filters: Optional[Dict[str, Union[str, int, List, Set]]] = None,
               **kwargs) -> List[dict]:
        try:
            collection = self._client.get_collection(name=self._gen_collection_name(collection_name, embed_key))

            filters = self._construct_filter_expr(filters) if filters else {}
            query_results = collection.query(query_embeddings=[query_embedding], n_results=topk, **filters)
            res = []
            for i, r_list in enumerate(query_results['ids']):
                for j, uid in enumerate(r_list):
                    dis = query_results['distances'][i][j]
                    res.append({'uid': uid, 'score': 1 - dis})
            return res
        except Exception as e:
            LOG.error(f"[ChromadbStore - search] task fail: {e}")
            LOG.error(traceback.format_exc())

    def _construct_criteria(self, criteria: dict) -> dict:
        res = {}
        if self._primary_key in criteria:
            res['ids'] = criteria[self._primary_key]
        else:
            res['where'] = {}
            for key, vaule in criteria.items():
                if key not in self._global_metadata_desc:
                    continue
                field_key = self._gen_global_meta_key(key)
                if isinstance(vaule, list):
                    res['where'][field_key] = {'$in': vaule}
                elif isinstance(vaule, str):
                    res['where'][field_key] = {'$eq': vaule}
                else:
                    raise ValueError(f'invalid criteria type: {type(vaule)}')
        return res

    def _construct_filter_expr(self, filters: Dict[str, Union[str, int, List, Set]]) -> str:
        ret = {}
        for name, candidates in filters.items():
            desc = self._global_metadata_desc.get(name)
            if not desc:
                raise ValueError(f'cannot find desc of field [{name}]')
            key = self._gen_global_meta_key(name)
            if isinstance(candidates, str):
                candidates = [candidates]
            elif (not isinstance(candidates, List)) and (not isinstance(candidates, Set)):
                candidates = list(candidates)
            ret[key] = {'$in': candidates}
        return {'where': ret}

    def _gen_global_meta_key(self, k: str) -> str:
        return GLOBAL_META_KEY_PREFIX + k

    def _gen_collection_name(self, collection_name: str, embed_key: str) -> str:
        return collection_name + '_' + embed_key + "_embed"

lazyllm.tools.rag.store.MilvusStore

Bases: LazyLLMStoreBase

Source code in lazyllm/tools/rag/store/vector/milvus_store.py
class MilvusStore(LazyLLMStoreBase):
    capability = StoreCapability.VECTOR
    need_embedding = True
    supports_index_registration = False

    def __init__(self, uri: str = '', db_name: str = 'lazyllm', index_kwargs: Optional[Union[Dict, List]] = None,
                 client_kwargs: Optional[Dict] = None):
        # one database, different collection for each group (for standalone, add prefix to collection name)
        # when there's data need upsert, collection creation happen.
        self._uri = uri
        self._db_name = db_name
        self._index_kwargs = index_kwargs
        self._client_kwargs = client_kwargs or {}
        self._primary_key = 'uid'
        self._client = None
        if self._uri and parse.urlparse(self._uri).scheme.lower() in ['unix', 'http', 'https', 'tcp', 'grpc']:
            self._is_remote = True
        else:
            self._is_remote = False

    @property
    def dir(self):
        if self._is_remote: return None
        p = Path(self._uri)
        p = p if p.suffix else (p / "milvus.db")
        return str(p.resolve(strict=False))

    @override
    def connect(self, embed_dims: Optional[Dict[str, int]] = None,
                embed_datatypes: Optional[Dict[str, DataType]] = None,
                global_metadata_desc: Optional[Dict[str, GlobalMetadataDesc]] = None, **kwargs):
        self._embed_dims = embed_dims or {}
        self._embed_datatypes = embed_datatypes or {}
        self._global_metadata_desc = global_metadata_desc or {}
        self._set_constants()
        self._connect()
        LOG.info("[Milvus Vector Store] init success!")
        self._disconnect()

    def _connect(self):
        try:
            self._client = pymilvus.MilvusClient(uri=self._uri, **self._client_kwargs)
            if self._is_remote and self._db_name:
                existing_dbs = self._client.list_databases()
                if self._db_name not in existing_dbs:
                    self._client.create_database(self._db_name)
                self._client.using_database(self._db_name)
        except Exception as e:
            LOG.error(f'[Milvus Store - connect] error: {e}')

    def _disconnect(self):
        try:
            if self._client:
                self._client.close()
                self._client = None
        except Exception as e:
            LOG.error(f'[Milvus Store - disconnect] error: {e}')

    @override
    def upsert(self, collection_name: str, data: List[dict]) -> bool:
        try:
            if not data: return
            data_embeddings = data[0].get('embedding', {})
            if not data_embeddings: return
            self._connect()
            if not self._client.has_collection(collection_name):
                embed_kwargs = {}
                for embed_key in data_embeddings.keys():
                    assert self._embed_datatypes.get(embed_key), \
                        f'cannot find embedding params for embed [{embed_key}]'
                    if embed_key not in embed_kwargs:
                        embed_kwargs[embed_key] = {'dtype': self._type2milvus[self._embed_datatypes[embed_key]]}
                    if self._embed_dims.get(embed_key): embed_kwargs[embed_key]['dim'] = self._embed_dims[embed_key]
                self._create_collection(collection_name, embed_kwargs)

            for i in range(0, len(data), MILVUS_UPSERT_BATCH_SIZE):
                self._client.upsert(collection_name=collection_name,
                                    data=[self._serialize_data(d) for d in data[i:i + MILVUS_UPSERT_BATCH_SIZE]])
            self._disconnect()
            return True
        except Exception as e:
            LOG.error(f'[Milvus Store - upsert] error: {e}')
            LOG.error(traceback.format_exc())
            self._disconnect()
            return False

    @override
    def delete(self, collection_name: str, criteria: Optional[dict] = None, **kwargs) -> bool:
        try:
            self._connect()
            if not self._client.has_collection(collection_name):
                return True
            self._client.load_collection(collection_name)
            if not criteria:
                self._client.drop_collection(collection_name=collection_name)
            else:
                self._client.delete(collection_name=collection_name, **self._construct_criteria(criteria))
            self._disconnect()
            return True
        except Exception as e:
            LOG.error(f'[Milvus Store - delete] error: {e}')
            self._disconnect()
            return False

    @override
    def get(self, collection_name: str, criteria: Optional[dict] = None, **kwargs) -> List[dict]:
        try:
            self._connect()
            if not self._client.has_collection(collection_name):
                return []
            self._client.load_collection(collection_name)
            col_desc = self._client.describe_collection(collection_name=collection_name)
            field_names = [field.get('name') for field in col_desc.get('fields', [])
                           if field.get('name').startswith(EMBED_PREFIX)]
            if criteria and self._primary_key in criteria:
                res = self._client.get(collection_name=collection_name, ids=criteria[self._primary_key])
            else:
                filters = self._construct_criteria(criteria) if criteria else {}
                if version.parse(pymilvus.__version__) >= version.parse('2.4.11'):
                    iterator = self._client.query_iterator(collection_name=collection_name,
                                                           batch_size=MILVUS_PAGINATION_OFFSET,
                                                           output_fields=field_names, **filters)
                    res = []
                    while True:
                        result = iterator.next()
                        if not result:
                            iterator.close()
                            break
                        res += result
                else:
                    res = self._client.query(collection_name=collection_name, output_fields=field_names, **filters)
            self._disconnect()
            return [self._deserialize_data(r) for r in res]
        except Exception as e:
            LOG.error(f'[Milvus Store - get] error: {e}')
            self._disconnect()
            return []

    def _set_constants(self):
        self._type2milvus = {
            DataType.VARCHAR: pymilvus.DataType.VARCHAR,
            DataType.ARRAY: pymilvus.DataType.ARRAY,
            DataType.FLOAT_VECTOR: pymilvus.DataType.FLOAT_VECTOR,
            DataType.INT32: pymilvus.DataType.INT32,
            DataType.INT64: pymilvus.DataType.INT64,
            DataType.SPARSE_FLOAT_VECTOR: pymilvus.DataType.SPARSE_FLOAT_VECTOR,
            DataType.STRING: pymilvus.DataType.STRING,
        }
        self._builtin_keys = {
            'uid': {'dtype': pymilvus.DataType.VARCHAR, 'max_length': 256, 'is_primary': True}
        }
        self._constant_fields = self._get_constant_fields()

    def _get_constant_fields(self) -> list:
        field_list = []
        for k, kws in self._builtin_keys.items():
            field_list.append(pymilvus.FieldSchema(name=k, **kws))
        for k, desc in self._global_metadata_desc.items():
            field_name = self._gen_global_meta_key(k)
            if desc.data_type == DataType.ARRAY:
                if desc.element_type is None:
                    raise ValueError(f'Milvus field [{field_name}]: '
                                     '`element_type` is required when `data_type` is ARRAY.')
                field_args = {'element_type': self._type2milvus[desc.element_type], 'max_capacity': desc.max_size}
                if desc.element_type == DataType.VARCHAR: field_args['max_length'] = 65535
            elif desc.data_type == DataType.VARCHAR:
                field_args = {'max_length': desc.max_size}
            else:
                field_args = {}
            field_list.append(pymilvus.FieldSchema(name=field_name, dtype=self._type2milvus[desc.data_type],
                                                   default_value=desc.default_value, **field_args))
        return field_list

    def _create_collection(self, collection_name: str, embed_kwargs: Dict[str, Dict]):  # noqa: C901
        field_list = copy.deepcopy(self._constant_fields)
        index_params = self._client.prepare_index_params()
        for k, kws in embed_kwargs.items():
            embed_field_name = self._gen_embed_key(k)
            field_list.append(pymilvus.FieldSchema(name=embed_field_name, **kws))
            index_params.add_index(field_name=embed_field_name, **kws)
            if isinstance(self._index_kwargs, list):
                for item in self._index_kwargs:
                    embed_key = item.get('embed_key', None)
                    if not embed_key:
                        raise ValueError(f'cannot find `embed_key` in `index_kwargs` of `{item}`')
                    if embed_key == k:
                        index_kwarg = item.copy()
                        index_kwarg.pop('embed_key', None)
                        index_params.add_index(field_name=embed_field_name, **index_kwarg)
                        break
            elif isinstance(self._index_kwargs, dict):
                index_params.add_index(field_name=embed_field_name, **self._index_kwargs)
        schema = pymilvus.CollectionSchema(fields=field_list, auto_id=False, enable_dynamic_field=False)
        self._client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)

    def _serialize_data(self, d: dict) -> dict:
        # only keep primary_key, embedding and global_meta
        res = {
            self._primary_key: d.get(self._primary_key, '')
        }
        for embed_key, value in d.get('embedding', {}).items():
            res[self._gen_embed_key(embed_key)] = value
        global_meta = d.get('global_meta', {})
        for name, desc in self._global_metadata_desc.items():
            value = global_meta.get(name, desc.default_value)
            if value is not None:
                res[self._gen_global_meta_key(name)] = value
        return res

    def _deserialize_data(self, d: dict) -> dict:
        res = {
            self._primary_key: d.get(self._primary_key, ''),
            'embedding': {}
        }
        for k, v in d.items():
            if k.startswith(EMBED_PREFIX):
                res['embedding'][k[len(EMBED_PREFIX):]] = v
        return res

    def _gen_embed_key(self, k: str) -> str:
        return EMBED_PREFIX + k

    def _gen_global_meta_key(self, k: str) -> str:
        return GLOBAL_META_KEY_PREFIX + k

    def _construct_criteria(self, criteria: dict) -> dict:
        res = {}
        criteria = dict(criteria)
        if self._primary_key in criteria:
            res['ids'] = criteria[self._primary_key]
        else:
            filter_str = ''
            for key, vaule in criteria.items():
                if key not in self._global_metadata_desc:
                    continue
                field_name = self._gen_global_meta_key(key)
                if len(filter_str) > 0:
                    filter_str += ' and '
                if isinstance(vaule, list):
                    filter_str += f'{field_name} in {vaule}'
                elif isinstance(vaule, str):
                    filter_str += f'{field_name} == "{vaule}"'
                else:
                    raise ValueError(f'invalid criteria type: {type(vaule)}')
            res['filter'] = filter_str
        return res

    @override
    def search(self, collection_name: str, query_embedding: Union[dict, List[float]], topk: int,
               filters: Optional[Dict[str, Union[List, set]]] = None, embed_key: Optional[str] = None,
               filter_str: Optional[str] = '', **kwargs) -> List[dict]:
        self._connect()
        if not embed_key or embed_key not in self._embed_datatypes:
            raise ValueError(f'[Milvus Store - search] Not supported or None `embed_key`: {embed_key}')
        res = []
        filter_expr = self._construct_filter_expr(filters) if filters else filter_str
        results = self._client.search(collection_name=collection_name, data=[query_embedding], limit=topk,
                                      anns_field=self._gen_embed_key(embed_key),
                                      filter=filter_expr)
        if len(results) != 1:
            raise ValueError(f'number of results [{len(results)}] != expected [1]')
        for result in results[0]:
            score = result.get('distance', 0)
            uid = result.get('id', result.get(self._primary_key, ''))
            if not uid:
                continue
            res.append({'uid': uid, 'score': score})
        self._disconnect()
        return res

    def _construct_filter_expr(self, filters: Dict[str, Union[str, int, List, Set]]) -> str:
        ret_str = ''
        if not filters:
            return ret_str
        for name, candidates in filters.items():
            desc = self._global_metadata_desc.get(name)
            if not desc:
                raise ValueError(f'cannot find desc of field [{name}]')
            key = self._gen_global_meta_key(name)
            if isinstance(candidates, str):
                candidates = [candidates]
            elif (not isinstance(candidates, list)) and (not isinstance(candidates, set)):
                candidates = list(candidates)
            if desc.data_type == DataType.ARRAY:
                ret_str += f'array_contains_any({key}, {candidates}) and '
            else:
                ret_str += f'{key} in {candidates} and '
        if len(ret_str) > 0:
            return ret_str[:-5]  # truncate the last ' and '
        return ret_str

lazyllm.tools.rag.readers.ReaderBase

Bases: ModuleBase

Base document reader class that provides basic interfaces for document loading. Inherits from ModuleBase and uses LazyLLMRegisterMetaClass as metaclass.

Parameters:

  • return_trace (bool, default: True ) –

    Whether to return processing trace information. Defaults to True.

Notes: - Provides both lazy loading and regular loading methods - Subclasses need to implement _lazy_load_data method - Supports batch document processing - Automatically converts to standardized DocNode format

Examples:

```python
from lazyllm.tools.rag.readers.readerBase import LazyLLMReaderBase
from lazyllm.tools.rag.doc_node import DocNode
from typing import Iterable

class CustomReader(LazyLLMReaderBase):
    def _lazy_load_data(self, file_paths: list, **kwargs) -> Iterable[DocNode]:
        for file_path in file_paths:
            # Process each file and yield DocNode
            content = self._read_file(file_path)
            yield DocNode(
                text=content,
                metadata={"source": file_path}
            )

# Create reader instance
reader = CustomReader(return_trace=True)

# Load documents
documents = reader.forward(file_paths=["doc1.txt", "doc2.txt"])
```
Source code in lazyllm/tools/rag/readers/readerBase.py
class LazyLLMReaderBase(ModuleBase, metaclass=LazyLLMRegisterMetaClass):
    """
Base document reader class that provides basic interfaces for document loading. Inherits from ModuleBase and uses LazyLLMRegisterMetaClass as metaclass.

Args:
    return_trace (bool): Whether to return processing trace information. Defaults to True.

**Notes:**
- Provides both lazy loading and regular loading methods
- Subclasses need to implement _lazy_load_data method
- Supports batch document processing
- Automatically converts to standardized DocNode format


Examples:

    ```python
    from lazyllm.tools.rag.readers.readerBase import LazyLLMReaderBase
    from lazyllm.tools.rag.doc_node import DocNode
    from typing import Iterable

    class CustomReader(LazyLLMReaderBase):
        def _lazy_load_data(self, file_paths: list, **kwargs) -> Iterable[DocNode]:
            for file_path in file_paths:
                # Process each file and yield DocNode
                content = self._read_file(file_path)
                yield DocNode(
                    text=content,
                    metadata={"source": file_path}
                )

    # Create reader instance
    reader = CustomReader(return_trace=True)

    # Load documents
    documents = reader.forward(file_paths=["doc1.txt", "doc2.txt"])
    ```
    """
    def __init__(self, *args, return_trace: bool = True, **kwargs):
        super().__init__(return_trace=return_trace)

    def _lazy_load_data(self, *args, **load_kwargs) -> Iterable[DocNode]:
        raise NotImplementedError(f"{self.__class__.__name__} does not implement lazy_load_data method.")

    def _load_data(self, *args, **load_kwargs) -> List[DocNode]:
        return list(self._lazy_load_data(*args, **load_kwargs))

    def forward(self, *args, **kwargs) -> List[DocNode]:
        return self._load_data(*args, **kwargs)

lazyllm.tools.rag.readers.PandasCSVReader

Bases: LazyLLMReaderBase

Reader for parsing CSV files using pandas.

Parameters:

  • concat_rows (bool, default: True ) –

    Whether to concatenate all rows into a single text block. Default is True.

  • col_joiner (str, default: ', ' ) –

    String used to join column values.

  • row_joiner (str, default: '\n' ) –

    String used to join rows.

  • pandas_config (Optional[Dict], default: None ) –

    Optional config for pandas.read_csv.

  • return_trace (bool, default: True ) –

    Whether to return the processing trace.

Source code in lazyllm/tools/rag/readers/pandasReader.py
class PandasCSVReader(LazyLLMReaderBase):
    """Reader for parsing CSV files using pandas.

Args:
    concat_rows (bool): Whether to concatenate all rows into a single text block. Default is True.
    col_joiner (str): String used to join column values.
    row_joiner (str): String used to join rows.
    pandas_config (Optional[Dict]): Optional config for pandas.read_csv.
    return_trace (bool): Whether to return the processing trace.
"""
    def __init__(self, concat_rows: bool = True, col_joiner: str = ", ", row_joiner: str = "\n",
                 pandas_config: Optional[Dict] = None, return_trace: bool = True) -> None:
        super().__init__(return_trace=return_trace)
        self._concat_rows = concat_rows
        self._col_joiner = col_joiner
        self._row_joiner = row_joiner
        self._pandas_config = pandas_config or {}

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if not isinstance(file, Path): file = Path(file)

        if fs:
            with fs.open(file) as f:
                df = pd.read_csv(f, **self._pandas_config)
        else:
            df = pd.read_csv(file, **self._pandas_config)

        text_list = df.apply(lambda row: (self._col_joiner).join(row.astype(str).tolist()), axis=1).tolist()

        if self._concat_rows: return [DocNode(text=(self._row_joiner).join(text_list))]
        else: return [DocNode(text=text) for text in text_list]

lazyllm.tools.rag.readers.PandasExcelReader

Bases: LazyLLMReaderBase

Reader for extracting text content from Excel (.xlsx) files.

Parameters:

  • concat_rows (bool, default: True ) –

    Whether to concatenate all rows into a single block.

  • sheet_name (Optional[str], default: None ) –

    Name of the sheet to read. If None, all sheets will be read.

  • pandas_config (Optional[Dict], default: None ) –

    Optional config for pandas.read_excel.

  • return_trace (bool, default: True ) –

    Whether to return the processing trace.

Source code in lazyllm/tools/rag/readers/pandasReader.py
class PandasExcelReader(LazyLLMReaderBase):
    """Reader for extracting text content from Excel (.xlsx) files.

Args:
    concat_rows (bool): Whether to concatenate all rows into a single block.
    sheet_name (Optional[str]): Name of the sheet to read. If None, all sheets will be read.
    pandas_config (Optional[Dict]): Optional config for pandas.read_excel.
    return_trace (bool): Whether to return the processing trace.
"""
    def __init__(self, concat_rows: bool = True, sheet_name: Optional[str] = None,
                 pandas_config: Optional[Dict] = None, return_trace: bool = True) -> None:
        super().__init__(return_trace=return_trace)
        self._concat_rows = concat_rows
        self._sheet_name = sheet_name
        self._pandas_config = pandas_config or {}

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        openpyxl_spec = importlib.util.find_spec("openpyxl")
        if openpyxl_spec is not None: pass
        else: raise ImportError("Please install openpyxl to read Excel files. "
                                "You can install it with `pip install openpyxl`")

        if not isinstance(file, Path): file = Path(file)
        if fs:
            with fs.open(file) as f:
                dfs = pd.read_excel(f, self._sheet_name, **self._pandas_config)
        else:
            dfs = pd.read_excel(file, self._sheet_name, **self._pandas_config)

        documents = []
        if isinstance(dfs, pd.DataFrame):
            df = dfs.fillna("")
            text_list = (df.astype(str).apply(lambda row: " ".join(row.values), axis=1).tolist())

            if self._concat_rows: documents.append(DocNode(text="\n".join(text_list)))
            else: documents.extend([DocNode(text=text) for text in text_list])
        else:
            for df in dfs.values():
                df = df.fillna("")
                text_list = (df.astype(str).apply(lambda row: " ".join(row), axis=1).tolist())

                if self._concat_rows: documents.append(DocNode(text="\n".join(text_list)))
                else: documents.extend([DocNode(text=text) for text in text_list])

        return documents

lazyllm.tools.rag.readers.PDFReader

Bases: LazyLLMReaderBase

Reader for extracting text content from PDF files.

Parameters:

  • return_full_document (bool, default: False ) –

    Whether to merge the entire PDF into a single document node. If False, each page becomes a separate node.

  • return_trace (bool, default: True ) –

    Whether to return the processing trace. Default is True.

Source code in lazyllm/tools/rag/readers/pdfReader.py
class PDFReader(LazyLLMReaderBase):
    """Reader for extracting text content from PDF files.

Args:
    return_full_document (bool): Whether to merge the entire PDF into a single document node. If False, each page becomes a separate node.
    return_trace (bool): Whether to return the processing trace. Default is True.
"""
    def __init__(self, return_full_document: bool = False, return_trace: bool = True) -> None:
        super().__init__(return_trace=return_trace)
        self._return_full_document = return_full_document

    @retry(stop=stop_after_attempt(RETRY_TIMES))
    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if not isinstance(file, Path): file = Path(file)

        fs = fs or get_default_fs()
        with fs.open(file, 'rb') as fp:
            stream = fp if is_default_fs(fs) else io.BytesIO(fp.read())
            pdf = pypdf.PdfReader(stream)
            num_pages = len(pdf.pages)
            docs = []
            if self._return_full_document:
                text = "\n".join(pdf.pages[page].extract_text() for page in range(num_pages))
                docs.append(DocNode(text=text))
            else:
                for page in range(num_pages):
                    page_text = pdf.pages[page].extract_text()
                    page_label = pdf.page_labels[page]
                    metadata = {"page_label": page_label}
                    docs.append(DocNode(text=page_text, metadata=metadata))
            return docs

lazyllm.tools.rag.readers.PPTXReader

Bases: LazyLLMReaderBase

Reader for PPTX (PowerPoint) files. Extracts text from slides and generates captions for embedded images using a vision-language model.

Parameters:

  • return_trace (bool, default: True ) –

    Whether to record the processing trace. Default is True.

Source code in lazyllm/tools/rag/readers/pptxReader.py
class PPTXReader(LazyLLMReaderBase):
    """Reader for PPTX (PowerPoint) files. Extracts text from slides and generates captions for embedded images using a vision-language model.

Args:
    return_trace (bool): Whether to record the processing trace. Default is True.
"""
    def __init__(self, return_trace: bool = True) -> None:
        try:
            thirdparty.check_packages(['python-pptx', 'torch', 'Pillow', 'transformers'])
        except ImportError:
            raise ImportError("Please install extra dependencies that are required for the "
                              "PPTXReader: `pip install torch transformers python-pptx Pillow`")

        super().__init__(return_trace=return_trace)
        model = tf.VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
        feature_extractor = tf.ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
        tokenizer = tf.AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

        self._parser_config = {"feature_extractor": feature_extractor, "model": model, "tokenizer": tokenizer}

    def _caption_image(self, tmp_image_file: str) -> str:
        from PIL import Image

        model = self._parser_config['model']
        feature_extractor = self._parser_config['feature_extractor']
        tokenizer = self._parser_config['tokenizer']

        device = infer_torch_device()
        model.to(device)

        max_length = 16
        num_beams = 4
        gen_kwargs = {"max_length": max_length, "num_beams": num_beams}

        i_image = Image.open(tmp_image_file)
        if i_image.mode != "RGB": i_image = i_image.convert(mode="RGB")

        pixel_values = feature_extractor(images=[i_image], return_tensors="pt").pixel_values
        pixel_values = pixel_values.to(device)

        output_ids = model.generate(pixel_values, **gen_kwargs)

        preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
        return preds[0].strip()

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if not isinstance(file, Path): file = Path(file)

        if fs:
            with fs.open(file) as f:
                presentation = pptx.Presentation(f)
        else:
            presentation = pptx.Presentation(file)

        result = ""
        for i, slide in enumerate(presentation.slides):
            result += f"\n\nSlide #{i}: \n"
            for shape in slide.shapes:
                if hasattr(shape, "image"):
                    image = shape.image
                    image_bytes = image.blob
                    f = tempfile.NamedTemporaryFile("wb", delete=False)
                    try:
                        f.write(image_bytes)
                        f.close()
                        result += f"\n Image: {self._caption_image(f.name)}\n\n"
                    finally:
                        os.unlink(f.name)

                if hasattr(shape, "text"): result += f"{shape.text}\n"
        return [DocNode(text=result)]

lazyllm.tools.rag.readers.VideoAudioReader

Bases: LazyLLMReaderBase

Reader for extracting speech content from video or audio files using OpenAI's Whisper model for transcription.

Parameters:

  • model_version (str, default: 'base' ) –

    Whisper model version (e.g., "base", "small", "medium", "large"). Default is "base".

  • return_trace (bool, default: True ) –

    Whether to return the processing trace. Default is True.

Source code in lazyllm/tools/rag/readers/videoAudioReader.py
class VideoAudioReader(LazyLLMReaderBase):
    """Reader for extracting speech content from video or audio files using OpenAI's Whisper model for transcription.

Args:
    model_version (str): Whisper model version (e.g., "base", "small", "medium", "large"). Default is "base".
    return_trace (bool): Whether to return the processing trace. Default is True.
"""
    def __init__(self, model_version: str = "base", return_trace: bool = True) -> None:
        super().__init__(return_trace=return_trace)
        self._model_version = model_version

        try:
            import whisper
        except ImportError:
            raise ImportError("Please install OpenAI whisper model "
                              "`pip install openai-whisper` to use the model")

        model = whisper.load_model(self._model_version)
        self._parser_config = {"model": model}

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        import whisper

        if not isinstance(file, Path): file = Path(file)

        if file.name.endswith("mp4"):
            try:
                from pydub import AudioSegment
            except ImportError:
                raise ImportError("Please install pydub `pip install pydub`")

            if fs:
                with fs.open(file, 'rb') as f:
                    video = AudioSegment.from_file(f, format="mp4")
            else:
                video = AudioSegment.from_file(file, format="mp4")

            audio = video.split_to_mono()[0]
            file_str = str(file)[:-4] + ".mp3"
            audio.export(file_str, format="mp3")

        model = cast(whisper.Whisper, self._parser_config["model"])
        result = model.transcribe(str(file))

        transcript = result['text']
        return [DocNode(text=transcript)]

lazyllm.tools.SqlManager

Bases: DBManager

SqlManager is a specialized tool for interacting with databases. It provides methods for creating tables, executing queries, and performing updates on databases.

Parameters:

  • db_type (str) –

    "PostgreSQL","SQLite", "MySQL", "MSSQL". Note that when the type is "SQLite", db_name is a file path or ":memory:"

  • user (str) –

    Username for connection

  • password (str) –

    Password for connection

  • host (str) –

    Hostname or IP

  • port (int) –

    Port number

  • db_name (str) –

    Name of the database

  • **options_str (str, default: None ) –

    Options represented in the format k1=v1&k2=v2

Source code in lazyllm/tools/sql/sql_manager.py
class SqlManager(DBManager):
    """SqlManager is a specialized tool for interacting with databases.
It provides methods for creating tables, executing queries, and performing updates on databases.

Arguments:
    db_type (str): "PostgreSQL","SQLite", "MySQL", "MSSQL". Note that when the type is "SQLite", db_name is a file path or ":memory:"
    user (str): Username for connection
    password (str): Password for connection
    host (str): Hostname or IP
    port (int): Port number
    db_name (str): Name of the database
    **options_str (str): Options represented in the format k1=v1&k2=v2
"""
    DB_TYPE_SUPPORTED = set(["postgresql", "mysql", "mssql", "sqlite", "mysql+pymysql"])
    DB_DRIVER_MAP = {"mysql": "pymysql"}
    PYTYPE_TO_SQL_MAP = {
        "integer": sqlalchemy.Integer,
        "string": sqlalchemy.Text,
        "text": sqlalchemy.Text,
        "boolean": sqlalchemy.Boolean,
        "float": sqlalchemy.Float,
        "datetime": sqlalchemy.DateTime,
        "bytes": sqlalchemy.LargeBinary,
        "bool": sqlalchemy.Boolean,
        "date": sqlalchemy.Date,
        "time": sqlalchemy.Time,
        "list": sqlalchemy.ARRAY,
        "dict": sqlalchemy.JSON,
        "uuid": sqlalchemy.Uuid,
    }

    def __init__(self, db_type: str, user: str, password: str, host: str, port: int, db_name: str, *,
                 options_str: str = None, tables_info_dict: Dict = None):
        db_type = db_type.lower()
        if db_type not in self.DB_TYPE_SUPPORTED:
            raise ValueError(f"{db_type} not supported")
        super().__init__(db_type)
        self._user = user
        self._password = password
        self._host = host
        self._port = port
        self._db_name = db_name
        self._tables_desc_dict = {}
        self._engine = None
        self._visible_tables = None
        self._metadata = sqlalchemy.MetaData()
        self._options_str = options_str
        if tables_info_dict:
            self._init_tables_by_info(tables_info_dict)

    def _init_tables_by_info(self, tables_info_dict):
        try:
            tables_info = TablesInfo.model_validate(tables_info_dict)
            self._visible_tables = [table_info.name for table_info in tables_info.tables]
            # create table if not exist
            self._create_tables_by_info(tables_info)
            desc_dict = self._gen_desc_by_info(tables_info)
            self.set_desc(desc_dict)
        except pydantic.ValidationError as e:
            raise ValueError(f"Validate tables_info_dict failed: {str(e)}")

    def _create_tables_by_info(self, tables_info: TablesInfo):
        for table_info in tables_info.tables:
            attrs = {"__tablename__": table_info.name, "__table_args__": {"extend_existing": True},
                     "metadata": self._metadata}
            for column_info in table_info.columns:
                column_type = column_info.data_type.lower()
                is_nullable = column_info.nullable
                column_name = column_info.name
                is_primary = column_info.is_primary_key
                # Use text for unsupported column type
                real_type = self.PYTYPE_TO_SQL_MAP.get(column_type, sqlalchemy.Text)
                attrs[column_name] = sqlalchemy.Column(real_type, nullable=is_nullable, primary_key=is_primary)
            # When create dynamic class with same name, old version will be replaced
            TableClass = type(table_info.name.capitalize(), (TableBase,), attrs)
            self.create_table(TableClass)

    def _gen_desc_by_info(self, tables_info: TablesInfo) -> dict:
        desc_dict = {}
        for table_info in tables_info.tables:
            table_comment = ""
            if table_info.comment:
                table_comment += f"COMMENT ON TABLE '{table_info.name}': {table_info.comment}\n"
            for column_info in table_info.columns:
                table_comment += f"COMMENT ON COLUMN '{table_info.name}.{column_info.name}': {column_info.comment}\n"
            if table_comment:
                desc_dict[table_info.name] = table_comment
        return desc_dict

    def _gen_conn_url(self) -> str:
        if self._db_type == "sqlite":
            conn_url = f"sqlite:///{self._db_name}{('?' + self._options_str) if self._options_str else ''}"
        else:
            driver = self.DB_DRIVER_MAP.get(self._db_type, "")
            password = quote_plus(self._password)
            conn_url = (f"{self._db_type}{('+' + driver) if driver else ''}://{self._user}:{password}@{self._host}"
                        f":{self._port}/{self._db_name}{('?' + self._options_str) if self._options_str else ''}")
        return conn_url

    @property
    def engine(self):
        if self._engine is None:
            self._engine = sqlalchemy.create_engine(self._gen_conn_url())
        return self._engine

    @contextmanager
    def get_session(self):
        """This is a context manager that creates and returns a database session, yields it for use, and then automatically commits or rolls back changes and closes the session when done.

**Returns:**

- sqlalchemy.orm.Session: sqlalchemy database session
"""
        _Session = sessionmaker(bind=self.engine)
        session = _Session()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()

    def check_connection(self) -> DBResult:
        """Check the current connection status of the SqlManagerBase.

**Returns:**

- DBResult: DBResult.status True if the connection is successful, False if it fails. DBResult.detail contains failure information.
"""
        try:
            with self.engine.connect() as _:
                return DBResult()
        except SQLAlchemyError as e:
            return DBResult(status=DBStatus.FAIL, detail=str(e))

    @property
    def desc(self) -> str:
        if self._desc is None:
            self.set_desc(tables_desc_dict={})
        return self._desc

    def set_desc(self, tables_desc_dict: dict = {}):  # noqa B006
        """When using SqlManager with LLM to query table entries in natural language, set descriptions for better results, especially when table names, column names, and values are not self-explanatory.

Args:
    tables_desc_dict (dict): descriptive comment for tables
"""
        self._desc = ""
        if not isinstance(tables_desc_dict, dict):
            raise ValueError(f"desc type {type(tables_desc_dict)} not supported")
        self._tables_desc_dict = tables_desc_dict
        if len(self.visible_tables) == 0:
            return
        # Generate desc according to table schema and comment
        self._desc = "The tables description is as follows\n```\n"
        for table_name in self.visible_tables:
            self._desc += f"Table {table_name}\n(\n"
            TableCls = self.get_table_orm_class(table_name)
            if TableCls is None:
                # The table could be dropped in other session
                continue
            table_columns = TableCls.__table__.columns
            for i, column in enumerate(table_columns):
                self._desc += f" {column.name} {column.type}"
                if i != len(table_columns) - 1:
                    self._desc += ","
                self._desc += "\n"
            self._desc += ");\n"
            if table_name in tables_desc_dict:
                self._desc += tables_desc_dict[table_name] + "\n\n"
        self._desc += "```\n"

    @property
    def visible_tables(self):
        if self._visible_tables is None:
            self._visible_tables = self.get_all_tables()
        return self._visible_tables

    @visible_tables.setter
    def visible_tables(self, visible_tables: list):
        all_tables = set(self.get_all_tables())
        for ele in visible_tables:
            if ele not in all_tables:
                raise ValueError(f"Table {ele} not found in database")
        self._visible_tables = visible_tables
        self.set_desc(self._tables_desc_dict)

    def _refresh_metadata(self, only=None):
        # refresh metadata in case of deleting/creating table in other session
        self._metadata.clear()
        self._metadata.reflect(bind=self.engine, only=only)

    def get_all_tables(self) -> list:
        """Return all table names in the current database.
"""
        self._refresh_metadata()
        return list(self._metadata.tables.keys())

    def get_table_orm_class(self, table_name):
        """Return the sqlalchemy orm class corresponding to the given table name. Combine with get_session to perform orm operations.
"""
        self._refresh_metadata(only=[table_name])
        Base = automap_base(metadata=self._metadata)
        Base.prepare()
        return getattr(Base.classes, table_name, None)

    def execute_commit(self, statement: str):
        """Execute the SQL script without return and submit changes.
"""
        with self.get_session() as session:
            session.execute(sqlalchemy.text(statement))

    def execute_query(self, statement: str) -> str:
        """Execute the SQL query script and return the result as a JSON string.
"""
        statement = re.sub(r"/\*.*?\*/", "", statement, flags=re.DOTALL).strip()
        create_table_pattern = r".*\s*create\s+table\s+.*"
        drop_table_pattern = r".*\s*drop\s+table\s+.*"
        statement_lower = statement.lower()
        if re.match(create_table_pattern, statement_lower):
            return f"Create table not supported. Original statement: {statement}"
        elif re.match(drop_table_pattern, statement_lower):
            return f"Drop table not supported. Original statement: {statement}"
        try:
            result = []
            _Session = sessionmaker(bind=self.engine)
            # Use original session without post commit
            with _Session() as session:
                cursor_result = session.execute(sqlalchemy.text(statement))
                columns = list(cursor_result.keys())
                result = [dict(zip(columns, row)) for row in cursor_result]
            str_result = json.dumps(result, ensure_ascii=False, default=self._serialize_uncommon_type)
        except Exception as e:
            str_result = f"Execute SQL ERROR: {str(e)}"
        return str_result

    def _create_by_script(self, table: str) -> DBResult:
        status = DBStatus.SUCCESS
        detail = "Success"
        try:
            with self.engine.connect() as conn:
                conn.execute(sqlalchemy.text(table))
                conn.commit()
        except OperationalError as e:
            status = DBStatus.FAIL
            detail = f"ERROR: {str(e)}"
        return DBResult(status=status, detail=detail)

    def _create_by_api(self, table: Union[DeclarativeBase, DeclarativeMeta]) -> DBResult:
        table.metadata.create_all(bind=self.engine, checkfirst=True)
        return DBResult()

    def create_table(self, table: Union[str, Type[DeclarativeBase], DeclarativeMeta]) -> DBResult:
        """Create a table

Args:
    table (str/Type[DeclarativeBase]/DeclarativeMeta): table schema。Supports three types of parameters: SQL statements with type str, ORM classes that inherit from DeclarativeBase or declarative_base().
"""
        status = DBStatus.SUCCESS
        detail = "Success"
        if isinstance(table, str):
            return self._create_by_script(table)
        # Support DeclarativeMeta created by declarative_base() which is deprecated since: 2.0
        elif issubclass(table, DeclarativeBase) or isinstance(table, DeclarativeMeta):
            return self._create_by_api(table)
        else:
            status = DBStatus.FAIL
            detail += f"Failed: Unsupported Type: {table}"
        return DBResult(status=status, detail=detail)

    def drop_table(self, table: Union[str, Type[DeclarativeBase], DeclarativeMeta]) -> DBResult:
        """Delete a table

Args:
    table (str/Type[DeclarativeBase]/DeclarativeMeta): table schema。Supports three types of parameters: Table name with type str, ORM classes that inherit from DeclarativeBase or declarative_base().
"""
        metadata = self._metadata
        if isinstance(table, str):
            tablename = table
        elif issubclass(table, DeclarativeBase) or isinstance(table, DeclarativeMeta):
            tablename = table.__tablename__
        else:
            return DBResult(status=DBStatus.FAIL, detail=f"{table} type unsupported")
        Table = sqlalchemy.Table(tablename, metadata, autoload_with=self.engine)
        Table.drop(self.engine, checkfirst=True)
        return DBResult()

    def insert_values(self, table_name: str, vals: List[dict]) -> DBResult:
        """Bulk insert data

Args:
    table_name (str): Table name
    vals (List[dict]): data to be inserted, format as [{"col_name1": v01, "col_name2": v02, ...}, {"col_name1": v11, "col_name2": v12, ...}, ...]
"""
        # Refresh metadata in case of tables created by other api
        TableCls = self.get_table_orm_class(table_name)
        if TableCls is None:
            return DBResult(status=DBStatus.FAIL, detail=f"{table_name} not found in database")
        with self.get_session() as session:
            session.bulk_insert_mappings(TableCls, vals)
        return DBResult()

check_connection()

Check the current connection status of the SqlManagerBase.

Returns:

  • DBResult: DBResult.status True if the connection is successful, False if it fails. DBResult.detail contains failure information.
Source code in lazyllm/tools/sql/sql_manager.py
    def check_connection(self) -> DBResult:
        """Check the current connection status of the SqlManagerBase.

**Returns:**

- DBResult: DBResult.status True if the connection is successful, False if it fails. DBResult.detail contains failure information.
"""
        try:
            with self.engine.connect() as _:
                return DBResult()
        except SQLAlchemyError as e:
            return DBResult(status=DBStatus.FAIL, detail=str(e))

create_table(table)

Create a table

Parameters:

  • table (str / Type[DeclarativeBase] / DeclarativeMeta) –

    table schema。Supports three types of parameters: SQL statements with type str, ORM classes that inherit from DeclarativeBase or declarative_base().

Source code in lazyllm/tools/sql/sql_manager.py
    def create_table(self, table: Union[str, Type[DeclarativeBase], DeclarativeMeta]) -> DBResult:
        """Create a table

Args:
    table (str/Type[DeclarativeBase]/DeclarativeMeta): table schema。Supports three types of parameters: SQL statements with type str, ORM classes that inherit from DeclarativeBase or declarative_base().
"""
        status = DBStatus.SUCCESS
        detail = "Success"
        if isinstance(table, str):
            return self._create_by_script(table)
        # Support DeclarativeMeta created by declarative_base() which is deprecated since: 2.0
        elif issubclass(table, DeclarativeBase) or isinstance(table, DeclarativeMeta):
            return self._create_by_api(table)
        else:
            status = DBStatus.FAIL
            detail += f"Failed: Unsupported Type: {table}"
        return DBResult(status=status, detail=detail)

drop_table(table)

Delete a table

Parameters:

  • table (str / Type[DeclarativeBase] / DeclarativeMeta) –

    table schema。Supports three types of parameters: Table name with type str, ORM classes that inherit from DeclarativeBase or declarative_base().

Source code in lazyllm/tools/sql/sql_manager.py
    def drop_table(self, table: Union[str, Type[DeclarativeBase], DeclarativeMeta]) -> DBResult:
        """Delete a table

Args:
    table (str/Type[DeclarativeBase]/DeclarativeMeta): table schema。Supports three types of parameters: Table name with type str, ORM classes that inherit from DeclarativeBase or declarative_base().
"""
        metadata = self._metadata
        if isinstance(table, str):
            tablename = table
        elif issubclass(table, DeclarativeBase) or isinstance(table, DeclarativeMeta):
            tablename = table.__tablename__
        else:
            return DBResult(status=DBStatus.FAIL, detail=f"{table} type unsupported")
        Table = sqlalchemy.Table(tablename, metadata, autoload_with=self.engine)
        Table.drop(self.engine, checkfirst=True)
        return DBResult()

execute_commit(statement)

Execute the SQL script without return and submit changes.

Source code in lazyllm/tools/sql/sql_manager.py
    def execute_commit(self, statement: str):
        """Execute the SQL script without return and submit changes.
"""
        with self.get_session() as session:
            session.execute(sqlalchemy.text(statement))

execute_query(statement)

Execute the SQL query script and return the result as a JSON string.

Source code in lazyllm/tools/sql/sql_manager.py
    def execute_query(self, statement: str) -> str:
        """Execute the SQL query script and return the result as a JSON string.
"""
        statement = re.sub(r"/\*.*?\*/", "", statement, flags=re.DOTALL).strip()
        create_table_pattern = r".*\s*create\s+table\s+.*"
        drop_table_pattern = r".*\s*drop\s+table\s+.*"
        statement_lower = statement.lower()
        if re.match(create_table_pattern, statement_lower):
            return f"Create table not supported. Original statement: {statement}"
        elif re.match(drop_table_pattern, statement_lower):
            return f"Drop table not supported. Original statement: {statement}"
        try:
            result = []
            _Session = sessionmaker(bind=self.engine)
            # Use original session without post commit
            with _Session() as session:
                cursor_result = session.execute(sqlalchemy.text(statement))
                columns = list(cursor_result.keys())
                result = [dict(zip(columns, row)) for row in cursor_result]
            str_result = json.dumps(result, ensure_ascii=False, default=self._serialize_uncommon_type)
        except Exception as e:
            str_result = f"Execute SQL ERROR: {str(e)}"
        return str_result

get_all_tables()

Return all table names in the current database.

Source code in lazyllm/tools/sql/sql_manager.py
    def get_all_tables(self) -> list:
        """Return all table names in the current database.
"""
        self._refresh_metadata()
        return list(self._metadata.tables.keys())

get_session()

This is a context manager that creates and returns a database session, yields it for use, and then automatically commits or rolls back changes and closes the session when done.

Returns:

  • sqlalchemy.orm.Session: sqlalchemy database session
Source code in lazyllm/tools/sql/sql_manager.py
    @contextmanager
    def get_session(self):
        """This is a context manager that creates and returns a database session, yields it for use, and then automatically commits or rolls back changes and closes the session when done.

**Returns:**

- sqlalchemy.orm.Session: sqlalchemy database session
"""
        _Session = sessionmaker(bind=self.engine)
        session = _Session()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()

get_table_orm_class(table_name)

Return the sqlalchemy orm class corresponding to the given table name. Combine with get_session to perform orm operations.

Source code in lazyllm/tools/sql/sql_manager.py
    def get_table_orm_class(self, table_name):
        """Return the sqlalchemy orm class corresponding to the given table name. Combine with get_session to perform orm operations.
"""
        self._refresh_metadata(only=[table_name])
        Base = automap_base(metadata=self._metadata)
        Base.prepare()
        return getattr(Base.classes, table_name, None)

insert_values(table_name, vals)

Bulk insert data

Parameters:

  • table_name (str) –

    Table name

  • vals (List[dict]) –

    data to be inserted, format as [{"col_name1": v01, "col_name2": v02, ...}, {"col_name1": v11, "col_name2": v12, ...}, ...]

Source code in lazyllm/tools/sql/sql_manager.py
    def insert_values(self, table_name: str, vals: List[dict]) -> DBResult:
        """Bulk insert data

Args:
    table_name (str): Table name
    vals (List[dict]): data to be inserted, format as [{"col_name1": v01, "col_name2": v02, ...}, {"col_name1": v11, "col_name2": v12, ...}, ...]
"""
        # Refresh metadata in case of tables created by other api
        TableCls = self.get_table_orm_class(table_name)
        if TableCls is None:
            return DBResult(status=DBStatus.FAIL, detail=f"{table_name} not found in database")
        with self.get_session() as session:
            session.bulk_insert_mappings(TableCls, vals)
        return DBResult()

set_desc(tables_desc_dict={})

When using SqlManager with LLM to query table entries in natural language, set descriptions for better results, especially when table names, column names, and values are not self-explanatory.

Parameters:

  • tables_desc_dict (dict, default: {} ) –

    descriptive comment for tables

Source code in lazyllm/tools/sql/sql_manager.py
    def set_desc(self, tables_desc_dict: dict = {}):  # noqa B006
        """When using SqlManager with LLM to query table entries in natural language, set descriptions for better results, especially when table names, column names, and values are not self-explanatory.

Args:
    tables_desc_dict (dict): descriptive comment for tables
"""
        self._desc = ""
        if not isinstance(tables_desc_dict, dict):
            raise ValueError(f"desc type {type(tables_desc_dict)} not supported")
        self._tables_desc_dict = tables_desc_dict
        if len(self.visible_tables) == 0:
            return
        # Generate desc according to table schema and comment
        self._desc = "The tables description is as follows\n```\n"
        for table_name in self.visible_tables:
            self._desc += f"Table {table_name}\n(\n"
            TableCls = self.get_table_orm_class(table_name)
            if TableCls is None:
                # The table could be dropped in other session
                continue
            table_columns = TableCls.__table__.columns
            for i, column in enumerate(table_columns):
                self._desc += f" {column.name} {column.type}"
                if i != len(table_columns) - 1:
                    self._desc += ","
                self._desc += "\n"
            self._desc += ");\n"
            if table_name in tables_desc_dict:
                self._desc += tables_desc_dict[table_name] + "\n\n"
        self._desc += "```\n"

lazyllm.tools.Reranker

Bases: ModuleBase, _PostProcess

Initializes a Rerank module for postprocessing and reranking of nodes (documents). This constructor initializes a Reranker module that configures a reranking process based on a specified reranking type. It allows for the dynamic selection and instantiation of reranking kernels (algorithms) based on the type and provided keyword arguments.

Parameters:

  • name (str, default: 'ModuleReranker' ) –

    The type of reranker used for the postprocessing and reranking process. Defaults to 'ModuleReranker'.

  • target (str, default: None ) –

    Deprecated parameter, only used to notify users.

  • output_format (Optional[str], default: None ) –

    Specifies the output format. Defaults to None. Optional values include 'content' and 'dict'. - 'content' means the output is in string format. - 'dict' means the output is a dictionary.

  • join (Union[bool, str], default: False ) –

    Determines whether to join the top-k output nodes. - When output_format is 'content': - If set to True, returns a single long string. - If set to False, returns a list of strings, each representing one node’s content. - When output_format is 'dict': - Joining is not supported; join defaults to False. - Returns a dictionary with three keys: 'content', 'embedding', and 'metadata'.

  • kwargs

    Additional keyword arguments passed to the reranker upon instantiation.

Detailed explanation of reranker types

  • Reranker: Instantiates a SentenceTransformerRerank reranker with a list of document nodes and a query.

  • KeywordFilter: This registered reranking function instantiates a KeywordNodePostprocessor with specified required and excluded keywords. It filters nodes based on the presence or absence of these keywords.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Document, Reranker, Retriever, DocNode
>>> m = lazyllm.OnlineEmbeddingModule()
>>> documents = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
>>> retriever = Retriever(documents, group_name='CoarseChunk', similarity='bm25', similarity_cut_off=0.01, topk=6)
>>> reranker = Reranker(DocNode(text=user_data),query="user query")
>>> ppl = lazyllm.ActionModule(retriever, reranker)
>>> ppl.start()
>>> print(ppl("user query"))
Source code in lazyllm/tools/rag/rerank.py
class Reranker(ModuleBase, _PostProcess):
    """Initializes a Rerank module for postprocessing and reranking of nodes (documents).
This constructor initializes a Reranker module that configures a reranking process based on a specified reranking type. It allows for the dynamic selection and instantiation of reranking kernels (algorithms) based on the type and provided keyword arguments.

Args:
    name: The type of reranker used for the postprocessing and reranking process. Defaults to 'ModuleReranker'.
    target (str): **Deprecated** parameter, only used to notify users.
    output_format: Specifies the output format. Defaults to None. Optional values include 'content' and 'dict'. 
        - 'content' means the output is in string format.
        - 'dict' means the output is a dictionary.
    join: Determines whether to join the top-k output nodes.
        - When `output_format` is 'content':
            - If set to True, returns a single long string.
            - If set to False, returns a list of strings, each representing one node’s content.
        - When `output_format` is 'dict':
            - Joining is not supported; `join` defaults to False.
            - Returns a dictionary with three keys: 'content', 'embedding', and 'metadata'.
    kwargs: Additional keyword arguments passed to the reranker upon instantiation.
**Detailed explanation of reranker types**

- Reranker: Instantiates a `SentenceTransformerRerank` reranker with a list of document nodes and a query.

- KeywordFilter: This registered reranking function instantiates a KeywordNodePostprocessor with specified required and excluded keywords. It filters nodes based on the presence or absence of these keywords.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, Reranker, Retriever, DocNode
    >>> m = lazyllm.OnlineEmbeddingModule()
    >>> documents = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
    >>> retriever = Retriever(documents, group_name='CoarseChunk', similarity='bm25', similarity_cut_off=0.01, topk=6)
    >>> reranker = Reranker(DocNode(text=user_data),query="user query")
    >>> ppl = lazyllm.ActionModule(retriever, reranker)
    >>> ppl.start()
    >>> print(ppl("user query"))
    """
    registered_reranker = dict()

    def __new__(cls, name: str = "ModuleReranker", *args, **kwargs):
        assert name in cls.registered_reranker, f"Reranker: {name} is not registered, please register first."
        item = cls.registered_reranker[name]
        if isinstance(item, type) and issubclass(item, Reranker):
            return super(Reranker, cls).__new__(item)
        else:
            return super(Reranker, cls).__new__(cls)

    def __init__(self, name: str = "ModuleReranker", target: Optional[str] = None,
                 output_format: Optional[str] = None, join: Union[bool, str] = False, **kwargs) -> None:
        super().__init__()
        self._name = name
        self._kwargs = kwargs
        lazyllm.deprecated(bool(target), '`target` parameter of reranker')
        _PostProcess.__init__(self, output_format, join)

    def forward(self, nodes: List[DocNode], query: str = "") -> List[DocNode]:
        results = self.registered_reranker[self._name](nodes, query=query, **self._kwargs)
        LOG.debug(f"Rerank use `{self._name}` and get nodes: {results}")
        return self._post_process(results)

    @classmethod
    def register_reranker(
        cls: "Reranker", func: Optional[Callable] = None, batch: bool = False
    ):
        """A class decorator factory method that provides a flexible mechanism for registering custom reranking algorithms to the `Reranker` class.
Args:
    func (Optional[Callable]): The reranking function or class to register. This can be omitted when using decorator syntax (@).
    batch (bool): Whether to process nodes in batches. Defaults to False, meaning nodes are processed individually.


Examples:

    @Reranker.register_reranker
    def my_reranker(node: DocNode, **kwargs):
        return node.score * 0.8  # 自定义分数计算
    """
        def decorator(f):
            if isinstance(f, type):
                cls.registered_reranker[f.__name__] = f
                return f
            else:
                def wrapper(nodes, **kwargs):
                    if batch:
                        return f(nodes, **kwargs)
                    else:
                        results = [f(node, **kwargs) for node in nodes]
                        return [result for result in results if result]

                cls.registered_reranker[f.__name__] = wrapper
                return wrapper

        return decorator(func) if func else decorator

register_reranker(func=None, batch=False) classmethod

A class decorator factory method that provides a flexible mechanism for registering custom reranking algorithms to the Reranker class. Args: func (Optional[Callable]): The reranking function or class to register. This can be omitted when using decorator syntax (@). batch (bool): Whether to process nodes in batches. Defaults to False, meaning nodes are processed individually.

Examples:

@Reranker.register_reranker
def my_reranker(node: DocNode, **kwargs):
    return node.score * 0.8  # 自定义分数计算
Source code in lazyllm/tools/rag/rerank.py
    @classmethod
    def register_reranker(
        cls: "Reranker", func: Optional[Callable] = None, batch: bool = False
    ):
        """A class decorator factory method that provides a flexible mechanism for registering custom reranking algorithms to the `Reranker` class.
Args:
    func (Optional[Callable]): The reranking function or class to register. This can be omitted when using decorator syntax (@).
    batch (bool): Whether to process nodes in batches. Defaults to False, meaning nodes are processed individually.


Examples:

    @Reranker.register_reranker
    def my_reranker(node: DocNode, **kwargs):
        return node.score * 0.8  # 自定义分数计算
    """
        def decorator(f):
            if isinstance(f, type):
                cls.registered_reranker[f.__name__] = f
                return f
            else:
                def wrapper(nodes, **kwargs):
                    if batch:
                        return f(nodes, **kwargs)
                    else:
                        results = [f(node, **kwargs) for node in nodes]
                        return [result for result in results if result]

                cls.registered_reranker[f.__name__] = wrapper
                return wrapper

        return decorator(func) if func else decorator

lazyllm.tools.rag.readers.readerBase.LazyLLMReaderBase

Bases: ModuleBase

Base document reader class that provides basic interfaces for document loading. Inherits from ModuleBase and uses LazyLLMRegisterMetaClass as metaclass.

Parameters:

  • return_trace (bool, default: True ) –

    Whether to return processing trace information. Defaults to True.

Notes: - Provides both lazy loading and regular loading methods - Subclasses need to implement _lazy_load_data method - Supports batch document processing - Automatically converts to standardized DocNode format

Examples:

```python
from lazyllm.tools.rag.readers.readerBase import LazyLLMReaderBase
from lazyllm.tools.rag.doc_node import DocNode
from typing import Iterable

class CustomReader(LazyLLMReaderBase):
    def _lazy_load_data(self, file_paths: list, **kwargs) -> Iterable[DocNode]:
        for file_path in file_paths:
            # Process each file and yield DocNode
            content = self._read_file(file_path)
            yield DocNode(
                text=content,
                metadata={"source": file_path}
            )

# Create reader instance
reader = CustomReader(return_trace=True)

# Load documents
documents = reader.forward(file_paths=["doc1.txt", "doc2.txt"])
```
Source code in lazyllm/tools/rag/readers/readerBase.py
class LazyLLMReaderBase(ModuleBase, metaclass=LazyLLMRegisterMetaClass):
    """
Base document reader class that provides basic interfaces for document loading. Inherits from ModuleBase and uses LazyLLMRegisterMetaClass as metaclass.

Args:
    return_trace (bool): Whether to return processing trace information. Defaults to True.

**Notes:**
- Provides both lazy loading and regular loading methods
- Subclasses need to implement _lazy_load_data method
- Supports batch document processing
- Automatically converts to standardized DocNode format


Examples:

    ```python
    from lazyllm.tools.rag.readers.readerBase import LazyLLMReaderBase
    from lazyllm.tools.rag.doc_node import DocNode
    from typing import Iterable

    class CustomReader(LazyLLMReaderBase):
        def _lazy_load_data(self, file_paths: list, **kwargs) -> Iterable[DocNode]:
            for file_path in file_paths:
                # Process each file and yield DocNode
                content = self._read_file(file_path)
                yield DocNode(
                    text=content,
                    metadata={"source": file_path}
                )

    # Create reader instance
    reader = CustomReader(return_trace=True)

    # Load documents
    documents = reader.forward(file_paths=["doc1.txt", "doc2.txt"])
    ```
    """
    def __init__(self, *args, return_trace: bool = True, **kwargs):
        super().__init__(return_trace=return_trace)

    def _lazy_load_data(self, *args, **load_kwargs) -> Iterable[DocNode]:
        raise NotImplementedError(f"{self.__class__.__name__} does not implement lazy_load_data method.")

    def _load_data(self, *args, **load_kwargs) -> List[DocNode]:
        return list(self._lazy_load_data(*args, **load_kwargs))

    def forward(self, *args, **kwargs) -> List[DocNode]:
        return self._load_data(*args, **kwargs)

lazyllm.tools.rag.component.bm25

BM25

A BM25 retriever that uses the BM25 algorithm to retrieve nodes.

Source code in lazyllm/tools/rag/component/bm25.py
class BM25:
    """A BM25 retriever that uses the BM25 algorithm to retrieve nodes."""

    def __init__(
        self,
        nodes: List[DocNode],
        language: str = "en",
        topk: int = 2,
        **kwargs,
    ) -> None:
        if language == "en":
            self._stemmer = Stemmer.Stemmer("english")
            self._stopwords = language
            self._tokenizer = lambda t: t
        elif language == "zh":
            self._stemmer = None
            # TODO(ywt): after bm25s supports cn stopwards, update this
            self._stopwords = STOPWORDS_CHINESE
            self._tokenizer = lambda t: " ".join(jieba.lcut(t))
        self.topk = min(topk, len(nodes))
        self.nodes = nodes

        corpus_tokens = bm25s.tokenize(
            [self._tokenizer(node.get_text()) for node in nodes],
            stopwords=self._stopwords,
            stemmer=self._stemmer,
        )
        self.bm25 = bm25s.BM25()
        self.bm25.index(corpus_tokens)

    def retrieve(self, query: str) -> List[Tuple[DocNode, float]]:
        tokenized_query = bm25s.tokenize(
            self._tokenizer(query), stopwords=self._stopwords, stemmer=self._stemmer
        )
        indexs, scores = self.bm25.retrieve(tokenized_query, k=self.topk)
        results = []
        for idx, score in zip(indexs[0], scores[0]):
            results.append((self.nodes[idx], score))
        return results

lazyllm.tools.rag.doc_to_db.DocInfoSchemaItem

Bases: TypedDict

Definition of a single field in the document information schema.

Parameters:

  • key (str) –

    The name of the field.

  • desc (str) –

    The description of the field's meaning.

  • type (str) –

    The data type of the field.

Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
class DocInfoSchemaItem(TypedDict):
    """Definition of a single field in the document information schema.

Args:
    key (str): The name of the field.
    desc (str): The description of the field's meaning.
    type (str): The data type of the field.
"""
    key: str
    desc: str
    type: str

lazyllm.tools.rag.doc_to_db.DocGenreAnalyser

Used to analyze the genre/type of documents, such as contracts, resumes, invoices, etc. It reads the document content and uses a language model to classify its type.

Parameters:

  • maximum_doc_num (int, default: 3 ) –

    Maximum number of documents to analyze, default is 3.

Examples:

>>> import lazyllm
>>> from lazyllm.components.doc_info_extractor import DocGenreAnalyser
>>> from lazyllm import OnlineChatModule
>>> m = OnlineChatModule(source="openai")
>>> analyser = DocGenreAnalyser()
>>> genre = analyser.analyse_doc_genre(m, "path/to/document.txt")
>>> print(genre)
contract
Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
class DocGenreAnalyser:
    """Used to analyze the genre/type of documents, such as contracts, resumes, invoices, etc. It reads the document content and uses a language model to classify its type.

Args:
    maximum_doc_num (int): Maximum number of documents to analyze, default is 3.


Examples:
    >>> import lazyllm
    >>> from lazyllm.components.doc_info_extractor import DocGenreAnalyser
    >>> from lazyllm import OnlineChatModule
    >>> m = OnlineChatModule(source="openai")
    >>> analyser = DocGenreAnalyser()
    >>> genre = analyser.analyse_doc_genre(m, "path/to/document.txt")
    >>> print(genre)
    contract
    """
    ONE_DOC_TOKEN_LIMIT = 10000

    def __init__(self, maximum_doc_num=3):
        self._reader = DirectoryReader(None, {}, {})
        self._pattern = re.compile(r"```json(.+?)```", re.DOTALL)
        self._maximum_doc_num = maximum_doc_num
        self._tiktoken_tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
        assert self._maximum_doc_num > 0

    def gen_detection_query(self, doc_path: str):
        root_nodes = self._reader.load_data([doc_path], None)
        doc_content = ""
        for root_node in root_nodes:
            doc_content += root_node.text + "\n"
        doc_content = trim_content_by_token_num(self._tiktoken_tokenizer, doc_content, self.ONE_DOC_TOKEN_LIMIT)
        query = DOC_KWS_PROMPTS["doc_type_detection"].format(doc_content=doc_content)
        query += "\nBelow is the content of each document sample.\n\n"
        return query

    def _extract_doc_type_from_response(self, str_response: str) -> str:
        # Remove the triple backticks if present
        matches = self._pattern.findall(str_response)
        if matches:
            # Return the first match
            extracted_content = matches[0].strip()
            try:
                res_dict = json.loads(extracted_content)
                if not isinstance(res_dict, dict) or "doc_type" not in res_dict:
                    return ""
                return res_dict["doc_type"]
            except Exception as e:
                lazyllm.LOG.warning(f"Exception: {str(e)}, response_str: {str_response}")
                return ""
        else:
            return ""

    def analyse_doc_genre(self, llm: Union[OnlineChatModule, TrainableModule], doc_path: str) -> str:
        query = self.gen_detection_query(doc_path)
        response = llm(query)
        doc_genre = self._extract_doc_type_from_response(response)
        return doc_genre

lazyllm.tools.rag.doc_to_db.DocInfoSchemaAnalyser

Used to extract key-value schema from documents, such as field names, descriptions, and data types. Useful for building structured information extraction templates.

Parameters:

  • maximum_doc_num (int, default: 3 ) –

    Maximum number of documents to be used for generating schema, default is 3.

Examples:

>>> from lazyllm.components.doc_info_extractor import DocInfoSchemaAnalyser
>>> from lazyllm import OnlineChatModule
>>> analyser = DocInfoSchemaAnalyser()
>>> m = OnlineChatModule(source="openai")
>>> schema = analyser.analyse_info_schema(m, "contract", ["doc1.txt", "doc2.txt"])
>>> print(schema)
[{'key': 'party_a', 'desc': 'The first party', 'type': 'str'}, ...]
Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
class DocInfoSchemaAnalyser:
    """Used to extract key-value schema from documents, such as field names, descriptions, and data types. Useful for building structured information extraction templates.

Args:
    maximum_doc_num (int): Maximum number of documents to be used for generating schema, default is 3.


Examples:
    >>> from lazyllm.components.doc_info_extractor import DocInfoSchemaAnalyser
    >>> from lazyllm import OnlineChatModule
    >>> analyser = DocInfoSchemaAnalyser()
    >>> m = OnlineChatModule(source="openai")
    >>> schema = analyser.analyse_info_schema(m, "contract", ["doc1.txt", "doc2.txt"])
    >>> print(schema)
    [{'key': 'party_a', 'desc': 'The first party', 'type': 'str'}, ...]
    """
    ONE_DOC_TOKEN_LIMIT = 30000

    def __init__(self, maximum_doc_num=3):
        self._reader = DirectoryReader(None, {}, {})
        self._pattern = re.compile(r"```json(.+?)```", re.DOTALL)
        self._maximum_doc_num = maximum_doc_num
        self._tiktoken_tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
        assert self._maximum_doc_num > 0

    def _gen_first_round_query(self, doc_type: str, doc_paths: list[str]):
        doc_contents = []
        for doc_path in doc_paths:
            root_nodes = self._reader.load_data([doc_path], None)
            doc_content = ""
            for root_node in root_nodes:
                doc_content += root_node.text + "\n"
            doc_content = trim_content_by_token_num(self._tiktoken_tokenizer, doc_content, self.ONE_DOC_TOKEN_LIMIT)
            doc_contents.append(doc_content)
        query = DOC_KWS_PROMPTS["kws_generation"].format(number=len(doc_contents), doc_type=doc_type)
        query += "\nBelow is the content of each document sample.\n\n"
        for i, doc_content in enumerate(doc_contents):
            query += f"Document {i+1}:\n```\n{doc_content}\n```\n\n"
        return query

    def _extract_schema_from_response(self, str_response: str) -> List[dict]:
        # Remove the triple backticks if present
        matches = self._pattern.findall(str_response)
        empty_list = []
        if matches:
            # Return the first match
            extracted_content = matches[0].strip()
            try:
                kws_list = json.loads(extracted_content)
                # in case of the list is in a dict, unpack it
                if isinstance(kws_list, dict):
                    values = list(kws_list.values())
                    if len(values) == 1 and isinstance(values[0], list):
                        return values[0]
                if not isinstance(kws_list, list):
                    lazyllm.LOG.warning(f"Excepted original type list but got {type(kws_list)} value: {kws_list}")
                    return empty_list
                return kws_list
            except Exception as e:
                lazyllm.LOG.warning(f"Exception: {str(e)}, response_str: {str_response}")
                return empty_list
        else:
            return empty_list

    def analyse_info_schema(
        self, llm: Union[OnlineChatModule, TrainableModule], doc_type: str, doc_paths: list[str]
    ) -> DocInfoSchema:
        """Method for analyzing document information schema, used to extract structural definitions of key information fields from documents of a specified type.

Args:
    llm (Union[OnlineChatModule, TrainableModule]): LLM model used to generate information schema
    doc_type (str): Document type, used to guide the LLM in generating corresponding information schema
    doc_paths (list[str]): List of document paths, used as information sources for analysis

**Returns:**

- DocInfoSchema: List of schema containing key information field definitions, each field includes key, desc, and type attributes
"""
        RANDOM_SEED = 1331
        if len(doc_paths) > self._maximum_doc_num:
            doc_paths.sort()
            random.seed(RANDOM_SEED)
            doc_paths = random.sample(doc_paths, self._maximum_doc_num)
        first_round_query = self._gen_first_round_query(doc_type, doc_paths)
        first_response = llm(first_round_query)
        info_schema = self._extract_schema_from_response(first_response)
        for info_schema_item in info_schema:
            is_success, msg = validate_schema_item(info_schema_item, DocInfoSchemaItem)
            if not is_success:
                lazyllm.LOG.warning(f"Please Try Again! Invalid kws dict: {info_schema_item}, error_msg: {msg}")
                return []
        return info_schema

analyse_info_schema(llm, doc_type, doc_paths)

Method for analyzing document information schema, used to extract structural definitions of key information fields from documents of a specified type.

Parameters:

  • llm (Union[OnlineChatModule, TrainableModule]) –

    LLM model used to generate information schema

  • doc_type (str) –

    Document type, used to guide the LLM in generating corresponding information schema

  • doc_paths (list[str]) –

    List of document paths, used as information sources for analysis

Returns:

  • DocInfoSchema: List of schema containing key information field definitions, each field includes key, desc, and type attributes
Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
    def analyse_info_schema(
        self, llm: Union[OnlineChatModule, TrainableModule], doc_type: str, doc_paths: list[str]
    ) -> DocInfoSchema:
        """Method for analyzing document information schema, used to extract structural definitions of key information fields from documents of a specified type.

Args:
    llm (Union[OnlineChatModule, TrainableModule]): LLM model used to generate information schema
    doc_type (str): Document type, used to guide the LLM in generating corresponding information schema
    doc_paths (list[str]): List of document paths, used as information sources for analysis

**Returns:**

- DocInfoSchema: List of schema containing key information field definitions, each field includes key, desc, and type attributes
"""
        RANDOM_SEED = 1331
        if len(doc_paths) > self._maximum_doc_num:
            doc_paths.sort()
            random.seed(RANDOM_SEED)
            doc_paths = random.sample(doc_paths, self._maximum_doc_num)
        first_round_query = self._gen_first_round_query(doc_type, doc_paths)
        first_response = llm(first_round_query)
        info_schema = self._extract_schema_from_response(first_response)
        for info_schema_item in info_schema:
            is_success, msg = validate_schema_item(info_schema_item, DocInfoSchemaItem)
            if not is_success:
                lazyllm.LOG.warning(f"Please Try Again! Invalid kws dict: {info_schema_item}, error_msg: {msg}")
                return []
        return info_schema

lazyllm.tools.rag.doc_to_db.DocInfoExtractor

Extracts specific values for key fields from a document according to a provided schema. Returns a dictionary of key-value pairs.

Examples:

>>> from lazyllm.components.doc_info_extractor import DocInfoExtractor
>>> from lazyllm import OnlineChatModule
>>> extractor = DocInfoExtractor()
>>> m = OnlineChatModule(source="openai")
>>> schema = [{"key": "party_a", "desc": "Party A name", "type": "str"}]
>>> info = extractor.extract_doc_info(m, "contract.txt", schema)
>>> print(info)
{'party_a': 'ABC Corp'}
Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
class DocInfoExtractor:
    """Extracts specific values for key fields from a document according to a provided schema. Returns a dictionary of key-value pairs.

Args:
    None


Examples:
    >>> from lazyllm.components.doc_info_extractor import DocInfoExtractor
    >>> from lazyllm import OnlineChatModule
    >>> extractor = DocInfoExtractor()
    >>> m = OnlineChatModule(source="openai")
    >>> schema = [{"key": "party_a", "desc": "Party A name", "type": "str"}]
    >>> info = extractor.extract_doc_info(m, "contract.txt", schema)
    >>> print(info)
    {'party_a': 'ABC Corp'}
    """
    ONE_DOC_TOKEN_LIMIT = 50000

    def __init__(self):
        self._reader = DirectoryReader(None, {}, {})
        self._pattern = re.compile(r"```json(.+?)```", re.DOTALL)
        self._tiktoken_tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")

    def _gen_extraction_query(self, doc_path: str, info_schema: DocInfoSchema, extra_desc: str) -> str:
        root_nodes = self._reader.load_data([doc_path], None)
        doc_content = ""
        for root_node in root_nodes:
            doc_content += root_node.text + "\n"
        doc_content = trim_content_by_token_num(self._tiktoken_tokenizer, doc_content, self.ONE_DOC_TOKEN_LIMIT)
        if not extra_desc:
            extra_desc = f"Extra description: \n{extra_desc}"
        query = DOC_KWS_PROMPTS["kws_extraction"].format(
            kws_desc=json.dumps(info_schema), extra_desc=extra_desc, doc_content=doc_content
        )
        return query

    def _extract_kws_value_from_response(self, str_response: str) -> dict:
        # Remove the triple backticks if present
        matches = self._pattern.findall(str_response)
        empty_dict = {}
        if matches:
            # Return the first match
            extracted_content = matches[0].strip()
            try:
                kws_value = json.loads(extracted_content)
                if not isinstance(kws_value, dict):
                    lazyllm.LOG.warning(f"Excepted original type list but got {type(kws_value)}")
                    return empty_dict
                new_dict = {k: v for k, v in kws_value.items() if (isinstance(v, str) and v and v != "None")}
                return new_dict
            except Exception as e:
                lazyllm.LOG.warning(f"Exception: {str(e)}, response_str: {str_response}")
                return empty_dict
        else:
            return empty_dict

    def _format_info_by_schema(self, info: dict, info_schema: DocInfoSchema):
        valid_keys = set([info_schema_item["key"] for info_schema_item in info_schema])
        return {k: v for k, v in info.items() if k in valid_keys}

    def extract_doc_info(
        self,
        llm: Union[OnlineChatModule, TrainableModule],
        doc_path: str,
        info_schema: DocInfoSchema,
        extra_desc: str = "",
    ) -> dict:
        """Extracts specific key information values from a document according to a provided schema.

This method uses a large language model to analyze document content and extract corresponding information values based on predefined field structure, returning a key-value dictionary.

Args:
    llm (Union[OnlineChatModule, TrainableModule]): The large language model used for document information extraction.
    doc_path (str): Path to the document to be analyzed.
    info_schema (DocInfoSchema): Field structure definition containing the information to be extracted.
    extra_desc (str, optional): Additional description information to guide the extraction process. Defaults to empty string.

Returns:
    dict: Extracted key information dictionary with field names as keys and corresponding information values as values.
"""
        extraction_query = self._gen_extraction_query(doc_path, info_schema, extra_desc)
        response = llm(extraction_query)
        info: dict = self._extract_kws_value_from_response(response)
        info: dict = self._format_info_by_schema(info, info_schema)
        return info

extract_doc_info(llm, doc_path, info_schema, extra_desc='')

Extracts specific key information values from a document according to a provided schema.

This method uses a large language model to analyze document content and extract corresponding information values based on predefined field structure, returning a key-value dictionary.

Parameters:

  • llm (Union[OnlineChatModule, TrainableModule]) –

    The large language model used for document information extraction.

  • doc_path (str) –

    Path to the document to be analyzed.

  • info_schema (DocInfoSchema) –

    Field structure definition containing the information to be extracted.

  • extra_desc (str, default: '' ) –

    Additional description information to guide the extraction process. Defaults to empty string.

Returns:

  • dict ( dict ) –

    Extracted key information dictionary with field names as keys and corresponding information values as values.

Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
    def extract_doc_info(
        self,
        llm: Union[OnlineChatModule, TrainableModule],
        doc_path: str,
        info_schema: DocInfoSchema,
        extra_desc: str = "",
    ) -> dict:
        """Extracts specific key information values from a document according to a provided schema.

This method uses a large language model to analyze document content and extract corresponding information values based on predefined field structure, returning a key-value dictionary.

Args:
    llm (Union[OnlineChatModule, TrainableModule]): The large language model used for document information extraction.
    doc_path (str): Path to the document to be analyzed.
    info_schema (DocInfoSchema): Field structure definition containing the information to be extracted.
    extra_desc (str, optional): Additional description information to guide the extraction process. Defaults to empty string.

Returns:
    dict: Extracted key information dictionary with field names as keys and corresponding information values as values.
"""
        extraction_query = self._gen_extraction_query(doc_path, info_schema, extra_desc)
        response = llm(extraction_query)
        info: dict = self._extract_kws_value_from_response(response)
        info: dict = self._format_info_by_schema(info, info_schema)
        return info

lazyllm.tools.rag.doc_to_db.DocInfoExtractor

Extracts specific values for key fields from a document according to a provided schema. Returns a dictionary of key-value pairs.

Examples:

>>> from lazyllm.components.doc_info_extractor import DocInfoExtractor
>>> from lazyllm import OnlineChatModule
>>> extractor = DocInfoExtractor()
>>> m = OnlineChatModule(source="openai")
>>> schema = [{"key": "party_a", "desc": "Party A name", "type": "str"}]
>>> info = extractor.extract_doc_info(m, "contract.txt", schema)
>>> print(info)
{'party_a': 'ABC Corp'}
Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
class DocInfoExtractor:
    """Extracts specific values for key fields from a document according to a provided schema. Returns a dictionary of key-value pairs.

Args:
    None


Examples:
    >>> from lazyllm.components.doc_info_extractor import DocInfoExtractor
    >>> from lazyllm import OnlineChatModule
    >>> extractor = DocInfoExtractor()
    >>> m = OnlineChatModule(source="openai")
    >>> schema = [{"key": "party_a", "desc": "Party A name", "type": "str"}]
    >>> info = extractor.extract_doc_info(m, "contract.txt", schema)
    >>> print(info)
    {'party_a': 'ABC Corp'}
    """
    ONE_DOC_TOKEN_LIMIT = 50000

    def __init__(self):
        self._reader = DirectoryReader(None, {}, {})
        self._pattern = re.compile(r"```json(.+?)```", re.DOTALL)
        self._tiktoken_tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")

    def _gen_extraction_query(self, doc_path: str, info_schema: DocInfoSchema, extra_desc: str) -> str:
        root_nodes = self._reader.load_data([doc_path], None)
        doc_content = ""
        for root_node in root_nodes:
            doc_content += root_node.text + "\n"
        doc_content = trim_content_by_token_num(self._tiktoken_tokenizer, doc_content, self.ONE_DOC_TOKEN_LIMIT)
        if not extra_desc:
            extra_desc = f"Extra description: \n{extra_desc}"
        query = DOC_KWS_PROMPTS["kws_extraction"].format(
            kws_desc=json.dumps(info_schema), extra_desc=extra_desc, doc_content=doc_content
        )
        return query

    def _extract_kws_value_from_response(self, str_response: str) -> dict:
        # Remove the triple backticks if present
        matches = self._pattern.findall(str_response)
        empty_dict = {}
        if matches:
            # Return the first match
            extracted_content = matches[0].strip()
            try:
                kws_value = json.loads(extracted_content)
                if not isinstance(kws_value, dict):
                    lazyllm.LOG.warning(f"Excepted original type list but got {type(kws_value)}")
                    return empty_dict
                new_dict = {k: v for k, v in kws_value.items() if (isinstance(v, str) and v and v != "None")}
                return new_dict
            except Exception as e:
                lazyllm.LOG.warning(f"Exception: {str(e)}, response_str: {str_response}")
                return empty_dict
        else:
            return empty_dict

    def _format_info_by_schema(self, info: dict, info_schema: DocInfoSchema):
        valid_keys = set([info_schema_item["key"] for info_schema_item in info_schema])
        return {k: v for k, v in info.items() if k in valid_keys}

    def extract_doc_info(
        self,
        llm: Union[OnlineChatModule, TrainableModule],
        doc_path: str,
        info_schema: DocInfoSchema,
        extra_desc: str = "",
    ) -> dict:
        """Extracts specific key information values from a document according to a provided schema.

This method uses a large language model to analyze document content and extract corresponding information values based on predefined field structure, returning a key-value dictionary.

Args:
    llm (Union[OnlineChatModule, TrainableModule]): The large language model used for document information extraction.
    doc_path (str): Path to the document to be analyzed.
    info_schema (DocInfoSchema): Field structure definition containing the information to be extracted.
    extra_desc (str, optional): Additional description information to guide the extraction process. Defaults to empty string.

Returns:
    dict: Extracted key information dictionary with field names as keys and corresponding information values as values.
"""
        extraction_query = self._gen_extraction_query(doc_path, info_schema, extra_desc)
        response = llm(extraction_query)
        info: dict = self._extract_kws_value_from_response(response)
        info: dict = self._format_info_by_schema(info, info_schema)
        return info

extract_doc_info(llm, doc_path, info_schema, extra_desc='')

Extracts specific key information values from a document according to a provided schema.

This method uses a large language model to analyze document content and extract corresponding information values based on predefined field structure, returning a key-value dictionary.

Parameters:

  • llm (Union[OnlineChatModule, TrainableModule]) –

    The large language model used for document information extraction.

  • doc_path (str) –

    Path to the document to be analyzed.

  • info_schema (DocInfoSchema) –

    Field structure definition containing the information to be extracted.

  • extra_desc (str, default: '' ) –

    Additional description information to guide the extraction process. Defaults to empty string.

Returns:

  • dict ( dict ) –

    Extracted key information dictionary with field names as keys and corresponding information values as values.

Source code in lazyllm/tools/rag/doc_to_db/doc_analysis.py
    def extract_doc_info(
        self,
        llm: Union[OnlineChatModule, TrainableModule],
        doc_path: str,
        info_schema: DocInfoSchema,
        extra_desc: str = "",
    ) -> dict:
        """Extracts specific key information values from a document according to a provided schema.

This method uses a large language model to analyze document content and extract corresponding information values based on predefined field structure, returning a key-value dictionary.

Args:
    llm (Union[OnlineChatModule, TrainableModule]): The large language model used for document information extraction.
    doc_path (str): Path to the document to be analyzed.
    info_schema (DocInfoSchema): Field structure definition containing the information to be extracted.
    extra_desc (str, optional): Additional description information to guide the extraction process. Defaults to empty string.

Returns:
    dict: Extracted key information dictionary with field names as keys and corresponding information values as values.
"""
        extraction_query = self._gen_extraction_query(doc_path, info_schema, extra_desc)
        response = llm(extraction_query)
        info: dict = self._extract_kws_value_from_response(response)
        info: dict = self._format_info_by_schema(info, info_schema)
        return info

lazyllm.tools.rag.doc_to_db.DocToDbProcessor

Used to extract information from documents and export it to a database.

This class analyzes document topics, extracts schema structure, pulls out key information, and saves it into a database table.

Parameters:

  • sql_manager (SqlManager) –

    The SQL management module.

  • doc_table_name (str, default: 'lazyllm_doc_elements' ) –

    The table name to store document fields. Default is lazyllm_doc_elements.

Note
  • If the table already exists, it checks and avoids redundant creation.
  • Use reset_doc_info_schema to reset the schema if necessary.
Source code in lazyllm/tools/rag/doc_to_db/doc_processor.py
class DocToDbProcessor:
    """Used to extract information from documents and export it to a database.

This class analyzes document topics, extracts schema structure, pulls out key information, and saves it into a database table.

Args:
    sql_manager (SqlManager): The SQL management module.
    doc_table_name (str): The table name to store document fields. Default is ``lazyllm_doc_elements``.

Note:
    - If the table already exists, it checks and avoids redundant creation.
    - Use `reset_doc_info_schema` to reset the schema if necessary.
"""

    DB_TYPE_MAP = {
        "int": sqlalchemy.Integer,
        "text": sqlalchemy.Text,
        "float": sqlalchemy.Float,
    }
    UUID_COL_NAME = "lazyllm_uuid"
    CREATED_AT_COL_NAME = "lazyllm_created_at"
    DOC_PATH_COL_NAME = "lazyllm_doc_path"

    def __init__(self, sql_manager: SqlManager, doc_table_name="lazyllm_doc_elements"):
        self._doc_genre_analyser = DocGenreAnalyser()
        self._doc_info_schema_analyser = DocInfoSchemaAnalyser(maximum_doc_num=2)
        self._doc_info_extractor = DocInfoExtractor()
        self._sql_manager = sql_manager
        self._doc_info_schema: DocInfoSchema = None
        self._doc_table_name = doc_table_name
        self._table_class = None
        all_table_names = set(self._sql_manager.get_all_tables()) if sql_manager else {}
        # If doc_table exists, then desc_table must exist as well
        if self._doc_table_name in all_table_names:
            assert (
                LazyllmDocTableDesc.__tablename__ in all_table_names
            ), "LazyllmDocTableDesc table not found in database"
        # Create desc table for totally new database
        if sql_manager and LazyllmDocTableDesc.__tablename__ not in all_table_names:
            self._sql_manager.create_table(LazyllmDocTableDesc)

    @property
    def doc_info_schema(self):
        return self._doc_info_schema

    @property
    def doc_table_name(self):
        return self._doc_table_name

    @doc_table_name.setter
    def doc_table_name(self, doc_table_name: str):
        raise NotImplementedError("Invalid to change table name")

    @property
    def sql_manager(self):
        return self._sql_manager

    @sql_manager.setter
    def sql_manager(self, sql_manager: SqlManager):
        self._sql_manager = sql_manager

    @doc_info_schema.setter
    def doc_info_schema(self, doc_info_schema: DocInfoSchema):
        raise NotImplementedError("As it'a dangerous operation, please use reset_doc_info_schema instead")

    def _save_description_to_db(self, doc_info_schema: DocInfoSchema):
        assert self._sql_manager is not None, "sqlManager is not initialized"
        json_data = json.dumps(doc_info_schema)  # 直接存储为 JSON(如果数据库支持)
        with self._sql_manager.get_session() as session:
            existing = session.query(LazyllmDocTableDesc).filter_by(id=1).first()
            if existing:
                # 更新现有记录
                existing.desc = json_data
            else:
                # 插入新记录
                new_desc = LazyllmDocTableDesc(id=1, desc=json_data)
                session.add(new_desc)
            session.commit()

    def _clear_table_orm(self, drop_doc_table=True):
        if self._table_class is not None:
            if drop_doc_table:
                self._sql_manager.drop_table(self._table_class)
            TableBase.metadata.remove(self._table_class.__table__)
            TableBase.registry._dispose_cls(self._table_class)
            del self._table_class
            self._table_class = None

    def clear(self):
        self._clear_table_orm()
        self._table_class = None
        self._doc_info_schema = None

    # Alert, reset_doc_info_schema will drop old result in db
    def _reset_doc_info_schema(self, doc_info_schema: DocInfoSchema, recreate_doc_table=True):
        assert isinstance(doc_info_schema, list)
        self._save_description_to_db(doc_info_schema)
        self._clear_table_orm(drop_doc_table=recreate_doc_table)
        for schema_item in doc_info_schema:
            is_success, err_msg = validate_schema_item(schema_item, DocInfoSchemaItem)
            assert is_success, err_msg
        self._doc_info_schema = doc_info_schema
        attrs = {"__tablename__": self._doc_table_name, "__table_args__": {"extend_existing": True}}
        # use uuid as primary key
        attrs[self.UUID_COL_NAME] = sqlalchemy.Column(sqlalchemy.String(36), primary_key=True)
        attrs[self.CREATED_AT_COL_NAME] = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
        attrs[self.DOC_PATH_COL_NAME] = sqlalchemy.Column(
            sqlalchemy.Text, nullable=False, primary_key=False, index=True
        )
        for schema_item in doc_info_schema:
            real_type = self.DB_TYPE_MAP.get(schema_item["type"].lower(), sqlalchemy.Text)
            attrs[schema_item["key"]] = sqlalchemy.Column(real_type, nullable=True, primary_key=False)
        self._table_class = type(self._doc_table_name.capitalize(), (TableBase,), attrs)
        if recreate_doc_table:
            # After drop_table, create table in db
            db_result = self._sql_manager.create_table(self._table_class)
            if db_result.status != DBStatus.SUCCESS:
                lazyllm.LOG.warning(f"Create table failed: {db_result.detail}")
                self.clear()

    def analyze_info_schema_by_llm(
        self, llm: Union[OnlineChatModule, TrainableModule], doc_paths: List[str], doc_topic: str = ""
    ) -> DocInfoSchema:
        """Infer structured database information using a large language model from document nodes.

Args:
    nodes (list[DocNode]): List of document nodes.

Returns:
    dict: The inferred database schema, including table names, fields, and relationships.
"""
        assert len(doc_paths) > 0, "doc_paths should not be empty"
        if not doc_topic:
            doc_topic = self._doc_genre_analyser.analyse_doc_genre(llm, doc_paths[0])
            if doc_topic == "":
                raise ValueError("Failed to detect doc type")
        return self._doc_info_schema_analyser.analyse_info_schema(llm, doc_topic, doc_paths)

    def extract_info_from_docs(
        self, llm: Union[OnlineChatModule, TrainableModule], doc_paths: List[str], extra_desc: str = ""
    ) -> List[dict]:
        """Extract structured database-related information from documents.

This function uses embedding and retrieval techniques to identify relevant text fragments in the provided documents for schema generation.

Args:
    docs (list[DocNode]): List of input documents.
    num_nodes (int): Number of text fragments to retrieve. Default is 10.

Returns:
    list[DocNode]: The relevant extracted document nodes.
"""
        existent_doc_paths = self._list_existent_doc_paths_in_db(doc_paths)
        # skip docs already in db
        doc_paths = list(set(doc_paths) - set(existent_doc_paths))
        info_dicts = []
        for doc_path in doc_paths:
            kws_value = self._doc_info_extractor.extract_doc_info(llm, doc_path, self._doc_info_schema, extra_desc)
            if kws_value:
                kws_value[self.DOC_PATH_COL_NAME] = str(doc_path)
                info_dicts.append(kws_value)
            else:
                lazyllm.LOG.warning(f"Extract kws value failed for {doc_path}")
        return info_dicts

    def export_info_to_db(self, info_dicts: List[dict]):
        # Generate uuid explicitly because SQLite doesn't support auto gen uuid
        new_values = []
        for kws_value in info_dicts:
            if kws_value:
                kws_value[self.UUID_COL_NAME] = str(uuid.uuid4())
                kws_value[self.CREATED_AT_COL_NAME] = datetime.now()
                new_values.append(kws_value)
        db_result = self._sql_manager.insert_values(self._doc_table_name, new_values)
        if db_result.status != DBStatus.SUCCESS:
            raise ValueError(f"Insert values failed: {db_result.detail}")

    def _list_existent_doc_paths_in_db(self, doc_paths: list[str]) -> List[str]:
        doc_paths = [str(ele) for ele in doc_paths]
        with self._sql_manager.get_session() as session:
            stmt = sqlalchemy.select(getattr(self._table_class, self.DOC_PATH_COL_NAME)).where(
                getattr(self._table_class, self.DOC_PATH_COL_NAME).in_(doc_paths)
            )
            result = session.execute(stmt).fetchall()
            return [ele[0] for ele in result]

analyze_info_schema_by_llm(llm, doc_paths, doc_topic='')

Infer structured database information using a large language model from document nodes.

Parameters:

  • nodes (list[DocNode]) –

    List of document nodes.

Returns:

  • dict ( DocInfoSchema ) –

    The inferred database schema, including table names, fields, and relationships.

Source code in lazyllm/tools/rag/doc_to_db/doc_processor.py
    def analyze_info_schema_by_llm(
        self, llm: Union[OnlineChatModule, TrainableModule], doc_paths: List[str], doc_topic: str = ""
    ) -> DocInfoSchema:
        """Infer structured database information using a large language model from document nodes.

Args:
    nodes (list[DocNode]): List of document nodes.

Returns:
    dict: The inferred database schema, including table names, fields, and relationships.
"""
        assert len(doc_paths) > 0, "doc_paths should not be empty"
        if not doc_topic:
            doc_topic = self._doc_genre_analyser.analyse_doc_genre(llm, doc_paths[0])
            if doc_topic == "":
                raise ValueError("Failed to detect doc type")
        return self._doc_info_schema_analyser.analyse_info_schema(llm, doc_topic, doc_paths)

extract_info_from_docs(llm, doc_paths, extra_desc='')

Extract structured database-related information from documents.

This function uses embedding and retrieval techniques to identify relevant text fragments in the provided documents for schema generation.

Parameters:

  • docs (list[DocNode]) –

    List of input documents.

  • num_nodes (int) –

    Number of text fragments to retrieve. Default is 10.

Returns:

  • List[dict]

    list[DocNode]: The relevant extracted document nodes.

Source code in lazyllm/tools/rag/doc_to_db/doc_processor.py
    def extract_info_from_docs(
        self, llm: Union[OnlineChatModule, TrainableModule], doc_paths: List[str], extra_desc: str = ""
    ) -> List[dict]:
        """Extract structured database-related information from documents.

This function uses embedding and retrieval techniques to identify relevant text fragments in the provided documents for schema generation.

Args:
    docs (list[DocNode]): List of input documents.
    num_nodes (int): Number of text fragments to retrieve. Default is 10.

Returns:
    list[DocNode]: The relevant extracted document nodes.
"""
        existent_doc_paths = self._list_existent_doc_paths_in_db(doc_paths)
        # skip docs already in db
        doc_paths = list(set(doc_paths) - set(existent_doc_paths))
        info_dicts = []
        for doc_path in doc_paths:
            kws_value = self._doc_info_extractor.extract_doc_info(llm, doc_path, self._doc_info_schema, extra_desc)
            if kws_value:
                kws_value[self.DOC_PATH_COL_NAME] = str(doc_path)
                info_dicts.append(kws_value)
            else:
                lazyllm.LOG.warning(f"Extract kws value failed for {doc_path}")
        return info_dicts

lazyllm.tools.rag.doc_to_db.extract_db_schema_from_files(file_paths, llm)

Extract the schema information from documents using a given LLM.

Parameters:

Returns:

  • DocInfoSchema ( DocInfoSchema ) –

    The extracted field structure schema.

Examples:

>>> import lazyllm
>>> from lazyllm.components.document_to_db import extract_db_schema_from_files
>>> llm = lazyllm.OnlineChatModule()
>>> file_paths = ["doc1.pdf", "doc2.pdf"]
>>> schema = extract_db_schema_from_files(file_paths, llm)
>>> print(schema)
Source code in lazyllm/tools/rag/doc_to_db/doc_processor.py
def extract_db_schema_from_files(file_paths: List[str], llm: Union[OnlineChatModule, TrainableModule]) -> DocInfoSchema:
    """Extract the schema information from documents using a given LLM.

Args:
    file_paths (List[str]): Paths of the documents to analyze.
    llm (Union[OnlineChatModule, TrainableModule]): A chat-supported LLM module.

Returns:
    DocInfoSchema: The extracted field structure schema.


Examples:
    >>> import lazyllm
    >>> from lazyllm.components.document_to_db import extract_db_schema_from_files
    >>> llm = lazyllm.OnlineChatModule()
    >>> file_paths = ["doc1.pdf", "doc2.pdf"]
    >>> schema = extract_db_schema_from_files(file_paths, llm)
    >>> print(schema)
    """
    return DocToDbProcessor(sql_manager=None).analyze_info_schema_by_llm(llm, file_paths)

lazyllm.tools.rag.readers.DocxReader

Bases: LazyLLMReaderBase

A docx format file parser, reading text content from a .docx file and return a list of DocNode objects.

Parameters:

  • file (Path) –

    Path to the .docx file.

  • fs (Optional[AbstractFileSystem]) –

    Optional file system object for custom reading.

Returns:

  • List[DocNode]: A list containing the extracted text content as DocNode instances.

Source code in lazyllm/tools/rag/readers/docxReader.py
class DocxReader(LazyLLMReaderBase):
    """A docx format file parser, reading text content from a `.docx` file and return a list of `DocNode` objects.

Args:
    file (Path): Path to the `.docx` file.
    fs (Optional[AbstractFileSystem]): Optional file system object for custom reading.

Returns:
    List[DocNode]: A list containing the extracted text content as `DocNode` instances.
"""
    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if not isinstance(file, Path): file = Path(file)

        if fs:
            with fs.open(file) as f:
                text = docx2txt.process(f)
        else:
            text = docx2txt.process(file)

        return [DocNode(text=text)]

lazyllm.tools.rag.readers.EpubReader

Bases: LazyLLMReaderBase

A file reader for .epub format eBooks.

Inherits from LazyLLMReaderBase, and only needs to implement _load_data. The Document module can automatically use this class to load .epub files.

Note: Reading from fsspec file systems (e.g., remote paths) is not supported in this version. If fs is specified, it will fall back to reading from the local file system.

Returns:

  • List[DocNode]: A single node containing all merged chapter content from the EPUB file.

Source code in lazyllm/tools/rag/readers/epubReader.py
class EpubReader(LazyLLMReaderBase):
    """A file reader for `.epub` format eBooks.

Inherits from `LazyLLMReaderBase`, and only needs to implement `_load_data`. The `Document` module can automatically use this class to load `.epub` files.

Note: Reading from fsspec file systems (e.g., remote paths) is not supported in this version. If `fs` is specified, it will fall back to reading from the local file system.

Returns:
    List[DocNode]: A single node containing all merged chapter content from the EPUB file.
"""
    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if not isinstance(file, Path): file = Path(file)

        if fs:
            LOG.warning("fs was specified but EpubReader doesn't support loading from "
                        "fsspec filesystems. Will load from local filesystem instead.")

        text_list = []

        spec = importlib.util.find_spec("ebooklib.epub")
        if spec is None:
            raise ImportError(
                "Please install ebooklib to use ebooklib module. "
                "You can install it with `pip install ebooklib`"
            )
        epub_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(epub_module)

        book = epub_module.read_epub(file, options={"ignore_ncs": True})

        for item in book.get_items():
            if item.get_type() == ebooklib.ITEM_DOCUMENT:
                text_list.append(html2text.html2text(item.get_content().decode("utf-8")))
        text = "\n".join(text_list)
        return [DocNode(text=text)]

lazyllm.tools.rag.readers.HWPReader

Bases: LazyLLMReaderBase

A HWP format file parser. It supports loading from the local filesystem. It extracts body text from the .hwp file and returns it as a list of DocNode objects.

HWP is a proprietary binary document format used primarily in Korea. This reader focuses on extracting the plain text from the body sections of the document.

Parameters:

  • return_trace (bool, default: True ) –

    Whether to enable trace logging. Defaults to True.

Source code in lazyllm/tools/rag/readers/hwpReader.py
class HWPReader(LazyLLMReaderBase):
    """
A HWP format file parser. It supports loading from the local filesystem. It extracts body text from the `.hwp` file and returns it as a list of DocNode objects.

HWP is a proprietary binary document format used primarily in Korea. This reader focuses on extracting the plain text from the body sections of the document.

Args:
    return_trace (bool): Whether to enable trace logging. Defaults to ``True``.
"""
    def __init__(self, return_trace: bool = True) -> None:
        super().__init__(return_trace=return_trace)
        self._FILE_HEADER_SECTION = "FileHeader"
        self._HWP_SUMMARY_SECTION = "\x05HwpSummaryInformation"
        self._SECTION_NAME_LENGTH = len("Section")
        self._BODYTEXT_SECTION = "BodyText"
        self._HWP_TEXT_TAGS = [67]
        self._text = ""

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if fs:
            LOG.warning("fs was specified but HWPReader doesn't support loading from "
                        "fsspec filesystems. Will load from local filesystem instead.")

        if not isinstance(file, Path): file = Path(file)

        load_file = olefile.OleFileIO(file)
        file_dir = load_file.listdir()
        if self._is_valid(file_dir) is False: raise Exception("Not Valid HwpFile")

        result_text = self._get_text(load_file, file_dir)
        return [DocNode(text=result_text)]

    def _is_valid(self, dirs: List[str]) -> bool:
        if [self._FILE_HEADER_SECTION] not in dirs: return False
        return [self._HWP_SUMMARY_SECTION] in dirs

    def _get_text(self, load_file: Any, file_dirs: List[str]) -> str:
        sections = self._get_body_sections(file_dirs)
        text = ""
        for section in sections:
            text += self._get_text_from_section(load_file, section)
            text += "\n"

        self._text = text
        return self._text

    def _get_body_sections(self, dirs: List[str]) -> List[str]:
        m = []
        for d in dirs:
            if d[0] == self._BODYTEXT_SECTION:
                m.append(int(d[1][self._SECTION_NAME_LENGTH:]))

        return ["BodyText/Section" + str(x) for x in sorted(m)]

    def _is_compressed(self, load_file: Any) -> bool:
        header = load_file.openstream("FileHeader")
        header_data = header.read()
        return (header_data[36] & 1) == 1

    def _get_text_from_section(self, load_file: Any, section: str) -> str:
        bodytext = load_file.openstream(section)
        data = bodytext.read()

        unpacked_data = (zlib.decompress(data, -15) if self._is_compressed(load_file) else data)
        size = len(unpacked_data)

        i = 0
        text = ""
        while i < size:
            header = struct.unpack_from("<I", unpacked_data, i)[0]
            rec_type = header & 0x3FF
            (header >> 10) & 0x3FF
            rec_len = (header >> 20) & 0xFFF

            if rec_type in self._HWP_TEXT_TAGS:
                rec_data = unpacked_data[i + 4: i + 4 + rec_len]
                text += rec_data.decode("utf-16")
                text += "\n"

            i += 4 + rec_len
        return text

lazyllm.tools.rag.readers.ImageReader

Bases: LazyLLMReaderBase

Module for reading content from image files. Supports keeping the image as base64, parsing text from images using OCR or pretrained vision models, and returns a list of nodes with text and image path.

Parameters:

  • parser_config (Optional[Dict], default: None ) –

    Parser configuration containing the model and processor. Defaults to None. When parse_text=True and parser_config is None, relevant models will be auto-loaded based on text_type.

  • keep_image (bool, default: False ) –

    Whether to keep the image as base64 string. Default is False.

  • parse_text (bool, default: False ) –

    Whether to parse text from the image. Default is False.

  • text_type (str, default: 'text' ) –

    Type of text parsing. Supports text (default) and plain_text. If plain_text, pytesseract OCR is used; otherwise a pretrained vision encoder-decoder model is used.

  • pytesseract_model_kwargs (Optional[Dict], default: None ) –

    Optional arguments passed to pytesseract OCR. Defaults to empty dict.

  • return_trace (bool, default: True ) –

    Whether to record the processing trace. Default is True.

Source code in lazyllm/tools/rag/readers/imageReader.py
class ImageReader(LazyLLMReaderBase):
    """Module for reading content from image files. Supports keeping the image as base64, parsing text from images using OCR or pretrained vision models, and returns a list of nodes with text and image path.

Args:
    parser_config (Optional[Dict]): Parser configuration containing the model and processor. Defaults to None. When parse_text=True and parser_config is None, relevant models will be auto-loaded based on text_type.
    keep_image (bool): Whether to keep the image as base64 string. Default is False.
    parse_text (bool): Whether to parse text from the image. Default is False.
    text_type (str): Type of text parsing. Supports ``text`` (default) and ``plain_text``. If ``plain_text``, pytesseract OCR is used; otherwise a pretrained vision encoder-decoder model is used.
    pytesseract_model_kwargs (Optional[Dict]): Optional arguments passed to pytesseract OCR. Defaults to empty dict.
    return_trace (bool): Whether to record the processing trace. Default is True.
"""
    def __init__(self, parser_config: Optional[Dict] = None, keep_image: bool = False, parse_text: bool = False,
                 text_type: str = "text", pytesseract_model_kwargs: Optional[Dict] = None,
                 return_trace: bool = True) -> None:
        super().__init__(return_trace=return_trace)
        self._text_type = text_type
        if parser_config is None and parse_text:
            if text_type == "plain_text":
                try:
                    import pytesseract
                except ImportError:
                    raise ImportError("Please install extra dependencies that are required for the ImageReader "
                                      "when text_type is 'plain_text': `pip install pytesseract`")

                processor = None
                model = pytesseract
            else:
                thirdparty.check_packages(["sentencepiece", "torch", "transformers"])

                processor = tf.DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
                model = tf.VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
            parser_config = {'processor': processor, 'model': model}

        self._parser_config = parser_config
        self._keep_image = keep_image
        self._parse_text = parse_text
        self._pytesseract_model_kwargs = pytesseract_model_kwargs or {}

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[ImageDocNode]:
        if not isinstance(file, Path): file = Path(file)

        if fs:
            with fs.open(path=file) as f:
                image = PIL.Image.open(f.read())
        else:
            image = PIL.Image.open(file)

        if image.mode != "RGB": image = image.convert("RGB")

        image_str: Optional[str] = None  # noqa
        if self._keep_image: image_str = img_2_b64(image)  # noqa

        text_str: str = ""
        if self._parse_text:
            assert self._parser_config is not None
            model = self._parser_config["model"]
            processor = self._parser_config["processor"]

            if processor:
                device = infer_torch_device()
                model.to(device)

                task_prompt = "<s_cord-v2>"
                decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False,
                                                        return_tensors='pt').input_ids
                pixel_values = processor(image, return_tensors='pt').pixel_values

                output = model.generate(pixel_values.to(device), decoder_input_ids=decoder_input_ids.to(device),
                                        max_length=model.decoder.config.max_position_embeddings, early_stopping=True,
                                        pad_token_id=processor.tokenizer.pad_token_id,
                                        eos_token_id=processor.tokenizer.eos_token_id, use_cache=True, num_beams=3,
                                        bad_words_ids=[[processor.tokenizer.unk_token_id]],
                                        return_dict_in_generate=True)

                sequence = processor.batch_decode(output.sequences)[0]
                sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
                text_str = re.sub(r"<.*?>", "", sequence, count=1).strip()
            else:
                import pytesseract

                model = cast(pytesseract, self._parser_config['model'])
                text_str = model.image_to_string(image, **self._pytesseract_model_kwargs)

        return [ImageDocNode(text=text_str, image_path=str(file))]

lazyllm.tools.rag.readers.IPYNBReader

Bases: LazyLLMReaderBase

Module for reading and parsing Jupyter Notebook (.ipynb) files. Converts the notebook to script text, then splits it by code cells into multiple document nodes or concatenates into a single text node.

Parameters:

  • parser_config (Optional[Dict], default: None ) –

    Reserved parser configuration parameter, currently unused. Defaults to None.

  • concatenate (bool, default: False ) –

    Whether to concatenate all code cells into one text node. Defaults to False (split into multiple nodes).

  • return_trace (bool, default: True ) –

    Whether to record processing trace. Default is True.

Source code in lazyllm/tools/rag/readers/ipynbReader.py
class IPYNBReader(LazyLLMReaderBase):
    """Module for reading and parsing Jupyter Notebook (.ipynb) files. Converts the notebook to script text, then splits it by code cells into multiple document nodes or concatenates into a single text node.

Args:
    parser_config (Optional[Dict]): Reserved parser configuration parameter, currently unused. Defaults to None.
    concatenate (bool): Whether to concatenate all code cells into one text node. Defaults to False (split into multiple nodes).
    return_trace (bool): Whether to record processing trace. Default is True.
"""
    def __init__(self, parser_config: Optional[Dict] = None, concatenate: bool = False, return_trace: bool = True):
        super().__init__(return_trace=return_trace)
        self._parser_config = parser_config
        self._concatenate = concatenate

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if not isinstance(file, Path): file = Path(file)

        if file.name.endswith(".ipynb"):
            try:
                import nbconvert
            except ImportError:
                raise ImportError("Please install nbconvert `pip install nbconvert`")

        if fs:
            with fs.open(file, encoding='utf-8') as f:
                doc_str = nbconvert.exporters.ScriptExporter().from_file(f)[0]
        else:
            doc_str = nbconvert.exporters.ScriptExporter().from_file(file)[0]

        splits = re.split(r"In\[\d+\]:", doc_str)
        splits.pop(0)

        if self._concatenate: docs = [DocNode(text="\n\n".join(splits))]
        else: docs = [DocNode(text=s) for s in splits]

        return docs

lazyllm.tools.rag.readers.MagicPDFReader

Module to parse PDF content via the MagicPDF service. Supports file upload or URL-based parsing, with a callback to process the parsed elements into document nodes.

Parameters:

  • magic_url (str) –

    The MagicPDF service API URL.

  • callback (Optional[Callable[[List[dict], Path, dict], List[DocNode]]], default: None ) –

    A callback function that takes parsed element list, file path, and extra info, returns a list of DocNode. Defaults to merging all text into a single node.

  • upload_mode (bool, default: False ) –

    Whether to use file upload mode for the API call. Default is False, meaning JSON request with file path.

Source code in lazyllm/tools/rag/readers/magic_pdf_reader.py
class MagicPDFReader:
    """Module to parse PDF content via the MagicPDF service. Supports file upload or URL-based parsing, with a callback to process the parsed elements into document nodes.

Args:
    magic_url (str): The MagicPDF service API URL.
    callback (Optional[Callable[[List[dict], Path, dict], List[DocNode]]]): A callback function that takes parsed element list, file path, and extra info, returns a list of DocNode. Defaults to merging all text into a single node.
    upload_mode (bool): Whether to use file upload mode for the API call. Default is False, meaning JSON request with file path.
"""

    def __init__(self, magic_url, callback: Optional[Callable[[List[dict], Path, dict], List[DocNode]]] = None,
                 upload_mode: bool = False):
        self._magic_url = magic_url
        self._upload_mode = upload_mode
        if callback is not None:
            self._callback = callback
        else:
            def default_callback(elements: List[dict], file: Path, extra_info: Optional[Dict] = None) -> List[DocNode]:
                text_chunks = [el["text"] for el in elements if "text" in el]
                return [DocNode(text="\n".join(text_chunks), metadata={"file_name": file.name})]
            self._callback = default_callback

    def __call__(self, file: Path, **kwargs) -> List[DocNode]:
        try:
            return self._load_data(file, **kwargs)
        except Exception as e:
            LOG.error(f"[MagicPDFReader] Error loading data from {file}: {e}")
            return []

    def _load_data(self, file: Path, extra_info: Optional[Dict] = None, **kwargs) -> List[DocNode]:
        if isinstance(file, str):
            file = Path(file)
        if self._upload_mode:
            elements = self._upload_parse_pdf_elements(file)
        else:
            elements = self._parse_pdf_elements(file)
        docs: List[DocNode] = self._callback(elements, file, extra_info)
        return docs

    def _parse_pdf_elements(self, pdf_path: Path) -> List[dict]:
        payload = {"files": [str(pdf_path)], "reserve_image": True}
        try:
            response = requests.post(self._magic_url, json=payload)
            response.raise_for_status()
            res = response.json()
            if not isinstance(res, list) or not res:
                LOG.info(f"[MagicPDFReader] No elements found in PDF: {pdf_path}")
                return []
        except requests.exceptions.RequestException as e:
            LOG.error(f"[MagicPDFReader] POST failed: {e}")
            return []
        return self._extract_content_blocks(res[0])

    def _upload_parse_pdf_elements(self, pdf_path: Path) -> List[dict]:
        try:
            with open(pdf_path, "rb") as f:
                files = {'file': (os.path.basename(pdf_path), f)}
                response = requests.post(self._magic_url, files=files)
                response.raise_for_status()
                res = response.json()
                if not isinstance(res, list) or not res:
                    LOG.info(f"[MagicPDFReader] No elements found in PDF: {pdf_path}")
                    return []
        except requests.exceptions.RequestException as e:
            LOG.error(f"[MagicPDFReader] POST failed: {e}")
            return []
        return self._extract_content_blocks(res[0])

    def _extract_content_blocks(self, content_list) -> List[dict]:  # noqa: C901
        blocks = []
        cur_title = ""
        cur_level = -1
        for content in content_list:
            block = {}
            block["bbox"] = content["bbox"]
            block["lines"] = content["lines"] if 'lines' in content else []
            for line in block['lines']:
                line['content'] = self._clean_content(line['content'])
            if content["type"] == "text":
                content["text"] = self._clean_content(content["text"]).strip()
                if not content["text"]:
                    continue
                if "text_level" in content:
                    if cur_title and content["text_level"] > cur_level:
                        content["title"] = cur_title
                    cur_title = content["text"]
                    cur_level = content["text_level"]
                else:
                    if cur_title:
                        content["title"] = cur_title
                block = copy.deepcopy(content)
                block["page"] = content["page_idx"]
                del block["page_idx"]
                blocks.append(block)
            elif content["type"] == "image":
                if not content["img_path"]:
                    continue
                block["type"] = content["type"]
                block["page"] = content["page_idx"]
                block["image_path"] = os.path.basename(content["img_path"])
                block['img_caption'] = self._clean_content(content['img_caption'])
                block['img_footnote'] = self._clean_content(content['img_footnote'])
                if cur_title:
                    block["title"] = cur_title
                img_title = block["img_caption"][0] if len(block["img_caption"]) > 0 else ""
                block["text"] = f"![{img_title}]({block['image_path']})"
                blocks.append(block)
            elif content["type"] == "table":
                block["type"] = content["type"]
                block["page"] = content["page_idx"]
                if self.extract_table:
                    block["text"] = self._html_table_to_markdown(self._clean_content(content["table_body"])
                                                                 ) if "table_body" in content else ""
                else:
                    block['image_path'] = os.path.basename(content['img_path'])
                if cur_title:
                    block["title"] = cur_title
                block['table_caption'] = self._clean_content(content['table_caption'])
                block['table_footnote'] = self._clean_content(content['table_footnote'])
                blocks.append(block)
        return blocks

    def _clean_content(self, content) -> str:
        if isinstance(content, str):
            content = content.encode("utf-8", "replace").decode("utf-8")
            return unicodedata.normalize("NFKC", content)
        if isinstance(content, list):
            return [self._clean_content(t) for t in content]
        return content

    def _html_table_to_markdown(self, html_table) -> str:  # noqa: C901
        try:
            soup = BeautifulSoup(html_table.strip(), 'html.parser')
            table = soup.find('table')
            if not table:
                raise ValueError("No <table> found in the HTML.")

            rows = []
            max_cols = 0

            for row in table.find_all('tr'):
                cells = []
                for cell in row.find_all(['td', 'th']):
                    rowspan = int(cell.get('rowspan', 1))
                    colspan = int(cell.get('colspan', 1))
                    text = cell.get_text(strip=True)

                    for _ in range(colspan):
                        cells.append({'text': text, 'rowspan': rowspan})
                rows.append(cells)
                max_cols = max(max_cols, len(cells))

            expanded_rows = []
            rowspan_tracker = [0] * max_cols
            for row in rows:
                expanded_row = []
                col_idx = 0
                for cell in row:
                    while col_idx < max_cols and rowspan_tracker[col_idx] > 0:
                        expanded_row.append(None)
                        rowspan_tracker[col_idx] -= 1
                        col_idx += 1

                    expanded_row.append(cell['text'])
                    if cell['rowspan'] > 1:
                        rowspan_tracker[col_idx] = cell['rowspan'] - 1
                    col_idx += 1

                while col_idx < max_cols:
                    if rowspan_tracker[col_idx] > 0:
                        expanded_row.append(None)
                        rowspan_tracker[col_idx] -= 1
                    else:
                        expanded_row.append("")
                    col_idx += 1

                expanded_rows.append(expanded_row)

            markdown = ''
            if not expanded_rows:
                return ""

            headers = expanded_rows[0]
            body_rows = expanded_rows[1:]
            if headers:
                markdown += '| ' + ' | '.join(h if h else '' for h in headers) + ' |\n'
                markdown += '| ' + ' | '.join(['-' * (len(h) if h else 3) for h in headers]) + ' |\n'
            for row in body_rows:
                markdown += '| ' + ' | '.join(cell if cell else '' for cell in row) + ' |\n'

            return markdown

        except Exception as e:
            LOG.error(f"Error parsing table: {e}")
            return ''

lazyllm.tools.rag.readers.MarkdownReader

Bases: LazyLLMReaderBase

Module for reading and parsing Markdown files. Supports removing hyperlinks and images, and splits Markdown into text segments by headers, returning document nodes.

Parameters:

  • remove_hyperlinks (bool, default: True ) –

    Whether to remove hyperlinks, default is True.

  • remove_images (bool, default: True ) –

    Whether to remove image tags, default is True.

  • return_trace (bool, default: True ) –

    Whether to record processing trace, default is True.

Source code in lazyllm/tools/rag/readers/markdownReader.py
class MarkdownReader(LazyLLMReaderBase):
    """Module for reading and parsing Markdown files. Supports removing hyperlinks and images, and splits Markdown into text segments by headers, returning document nodes.

Args:
    remove_hyperlinks (bool): Whether to remove hyperlinks, default is True.
    remove_images (bool): Whether to remove image tags, default is True.
    return_trace (bool): Whether to record processing trace, default is True.
"""
    def __init__(self, remove_hyperlinks: bool = True, remove_images: bool = True, return_trace: bool = True) -> None:
        super().__init__(return_trace=return_trace)
        self._remove_hyperlinks = remove_hyperlinks
        self._remove_images = remove_images

    def _markdown_to_tups(self, markdown_text: str) -> List[Tuple[Optional[str], str]]:
        markdown_tups: List[Tuple[Optional[str], str]] = []
        lines = markdown_text.split("\n")

        current_header = None
        current_lines = []
        in_code_block = False

        for line in lines:
            if line.startswith("```"): in_code_block = not in_code_block

            header_match = re.match(r"^#+\s", line)
            if not in_code_block and header_match:
                if current_header is not None or len(current_lines) > 0:
                    markdown_tups.append((current_header, "\n".join(current_lines)))
                current_header = line
                current_lines.clear()
            else:
                current_lines.append(line)

        markdown_tups.append((current_header, "\n".join(current_lines)))
        return [(key if key is None else re.sub(r"#", "", key).strip(), re.sub(r"<.*?>", "", value),)
                for key, value in markdown_tups]

    def remove_images(self, content: str) -> str:
        """Remove custom image tags of the form ![[...]] from the content.

Args:
    content (str): Input markdown content.

Returns:
    str: Content with image tags removed.
"""
        pattern = r"!{1}\[\[(.*)\]\]"
        return re.sub(pattern, "", content)

    def remove_hyperlinks(self, content: str) -> str:
        """Remove markdown hyperlinks, converting [text](url) to just text.

Args:
    content (str): Input markdown content.

Returns:
    str: Content with hyperlinks removed, only link text retained.
"""
        pattern = r"\[(.*)\]\((.*)\)"
        return re.sub(pattern, r"\1", content)

    def _parse_tups(self, filepath: Path, errors: str = "ignore",
                    fs: Optional[AbstractFileSystem] = None) -> List[Tuple[Optional[str], str]]:
        fs = fs or LocalFileSystem()

        with fs.open(filepath, encoding="utf-8") as f:
            content = f.read().decode(encoding="utf-8")

        if self._remove_hyperlinks: content = self.remove_hyperlinks(content)
        if self._remove_images: content = self.remove_images(content)
        return self._markdown_to_tups(content)

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        if not isinstance(file, Path): file = Path(file)

        tups = self._parse_tups(file, fs=fs)
        results = [DocNode(
            content=[value if header is None else f"\n\n{header}\n{value}" for header, value in tups])]
        return results

Remove markdown hyperlinks, converting text to just text.

Parameters:

  • content (str) –

    Input markdown content.

Returns:

  • str ( str ) –

    Content with hyperlinks removed, only link text retained.

Source code in lazyllm/tools/rag/readers/markdownReader.py
    def remove_hyperlinks(self, content: str) -> str:
        """Remove markdown hyperlinks, converting [text](url) to just text.

Args:
    content (str): Input markdown content.

Returns:
    str: Content with hyperlinks removed, only link text retained.
"""
        pattern = r"\[(.*)\]\((.*)\)"
        return re.sub(pattern, r"\1", content)

remove_images(content)

Remove custom image tags of the form ![[...]] from the content.

Parameters:

  • content (str) –

    Input markdown content.

Returns:

  • str ( str ) –

    Content with image tags removed.

Source code in lazyllm/tools/rag/readers/markdownReader.py
    def remove_images(self, content: str) -> str:
        """Remove custom image tags of the form ![[...]] from the content.

Args:
    content (str): Input markdown content.

Returns:
    str: Content with image tags removed.
"""
        pattern = r"!{1}\[\[(.*)\]\]"
        return re.sub(pattern, "", content)

lazyllm.tools.rag.readers.MboxReader

Bases: LazyLLMReaderBase

Module to parse Mbox email archive files. Reads email messages and formats them into text. Supports limiting the maximum number of messages and custom message formatting.

Parameters:

  • max_count (int, default: 0 ) –

    Maximum number of emails to read. Default 0 means read all.

  • message_format (str, default: DEFAULT_MESSAGE_FORMAT ) –

    Template string for formatting each message, supports placeholders {_date}, {_from}, {_to}, {_subject}, and {_content}.

  • return_trace (bool, default: True ) –

    Whether to record processing trace. Default is True.

Source code in lazyllm/tools/rag/readers/mboxreader.py
class MboxReader(LazyLLMReaderBase):
    """Module to parse Mbox email archive files. Reads email messages and formats them into text. Supports limiting the maximum number of messages and custom message formatting.

Args:
    max_count (int): Maximum number of emails to read. Default 0 means read all.
    message_format (str): Template string for formatting each message, supports placeholders ``{_date}``, ``{_from}``, ``{_to}``, ``{_subject}``, and ``{_content}``.
    return_trace (bool): Whether to record processing trace. Default is True.
"""
    DEFAULT_MESSAGE_FORMAT: str = (
        "Date: {_date}\n"
        "From: {_from}\n"
        "To: {_to}\n"
        "Subject: {_subject}\n"
        "Content: {_content}"
    )

    def __init__(self, max_count: int = 0, message_format: str = DEFAULT_MESSAGE_FORMAT,
                 return_trace: bool = True) -> None:
        try:
            from bs4 import BeautifulSoup  # noqa
        except ImportError:
            raise ImportError("`BeautifulSoup` package not found: `pip install beautifulsoup4`")

        super().__init__(return_trace=return_trace)
        self._max_count = max_count
        self._message_format = message_format

    def _load_data(self, file: Path, fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        import mailbox
        from email.parser import BytesParser
        from email.policy import default
        from bs4 import BeautifulSoup

        if fs:
            LOG.warning("fs was specified but MboxReader doesn't support loading from "
                        "fsspec filesystems. Will load from local filesystem instead.")

        i = 0
        results: List[str] = []
        bytes_parser = BytesParser(policy=default).parse
        mbox = mailbox.mbox(file, factory=bytes_parser)

        for _, _msg in enumerate(mbox):
            try:
                msg: mailbox.mboxMessage = _msg
                if msg.is_multipart():
                    for part in msg.walk():
                        ctype = part.get_content_type()
                        cdispo = str(part.get("Content-Disposition"))
                        if ctype == "text/plain" and "attachment" not in cdispo:
                            content = part.get_payload(decode=True)
                            break
                else:
                    content = msg.get_payload(decode=True)

                soup = BeautifulSoup(content)
                stripped_content = " ".join(soup.get_text().split())
                msg_string = self._message_format.format(_date=msg["date"], _from=msg["from"], _to=msg["to"],
                                                         _subject=msg["subject"], _content=stripped_content)
                results.append(msg_string)
            except Exception as e:
                LOG.warning(f"Failed to parse message:\n{_msg}\n with exception {e}")

            i += 1
            if self._max_count > 0 and i >= self._max_count: break
        return [DocNode(text=result) for result in results]

lazyllm.tools.SqlCall

Bases: ModuleBase

SqlCall is a class that extends ModuleBase and provides an interface for generating and executing SQL queries using a language model (LLM). It is designed to interact with a SQL database, extract SQL queries from LLM responses, execute those queries, and return results or explanations.

Parameters:

  • llm

    A language model to be used for generating and interpreting SQL queries and explanations.

  • sql_manager (SqlManager) –

    An instance of SqlManager that handles interaction with the SQL database.

  • sql_examples (str, default: '' ) –

    An example of converting natural language represented by a JSON string into an SQL statement, formatted as: [{"Question": "Find the names of people in the same department as Smith", "Answer": "SELECT...;"}]

  • use_llm_for_sql_result (bool, default: True ) –

    Default is True. If set to False, the module will only output raw SQL results in JSON without further processing.

  • return_trace (bool, default: False ) –

    If set to True, the results will be recorded in the trace. Defaults to False.

Examples:

>>> # First, run SqlManager example
>>> import lazyllm
>>> from lazyllm.tools import SQLManger, SqlCall
>>> sql_tool = SQLManger("personal.db")
>>> sql_llm = lazyllm.OnlineChatModule(model="gpt-4o", source="openai", base_url="***")
>>> sql_call = SqlCall(sql_llm, sql_tool, use_llm_for_sql_result=True)
>>> print(sql_call("去年一整年销售额最多的员工是谁?"))
Source code in lazyllm/tools/sql_call/sql_call.py
class SqlCall(ModuleBase):
    """SqlCall is a class that extends ModuleBase and provides an interface for generating and executing SQL queries using a language model (LLM).
It is designed to interact with a SQL database, extract SQL queries from LLM responses, execute those queries, and return results or explanations.

Arguments:
    llm: A language model to be used for generating and interpreting SQL queries and explanations.
    sql_manager (SqlManager): An instance of SqlManager that handles interaction with the SQL database.
    sql_examples (str, optional): An example of converting natural language represented by a JSON string into an SQL statement, formatted as: [{"Question": "Find the names of people in the same department as Smith", "Answer": "SELECT...;"}]
    use_llm_for_sql_result (bool, optional): Default is True. If set to False, the module will only output raw SQL results in JSON without further processing.
    return_trace (bool, optional): If set to True, the results will be recorded in the trace. Defaults to False.


Examples:
        >>> # First, run SqlManager example
        >>> import lazyllm
        >>> from lazyllm.tools import SQLManger, SqlCall
        >>> sql_tool = SQLManger("personal.db")
        >>> sql_llm = lazyllm.OnlineChatModule(model="gpt-4o", source="openai", base_url="***")
        >>> sql_call = SqlCall(sql_llm, sql_tool, use_llm_for_sql_result=True)
        >>> print(sql_call("去年一整年销售额最多的员工是谁?"))
    """
    EXAMPLE_TITLE = "Here are some example: "

    def __init__(
        self,
        llm,
        sql_manager: DBManager,
        sql_examples: str = "",
        sql_post_func: Callable = None,
        use_llm_for_sql_result=True,
        return_trace: bool = False,
    ) -> None:
        super().__init__(return_trace=return_trace)
        if not sql_manager.desc:
            raise ValueError("Error: sql_manager found empty description.")
        self._sql_tool = sql_manager
        self.sql_post_func = sql_post_func

        if sql_manager.db_type == "mongodb":
            self._query_prompter = ChatPrompter(instruction=mongodb_query_instruct_template).pre_hook(
                self.sql_query_promt_hook
            )
            statement_type = "mongodb json pipeline"
            self._pattern = re.compile(r"```json(.+?)```", re.DOTALL)
        else:
            self._query_prompter = ChatPrompter(instruction=sql_query_instruct_template).pre_hook(
                self.sql_query_promt_hook
            )
            statement_type = "sql query"
            self._pattern = re.compile(r"```sql(.+?)```", re.DOTALL)

        self._llm_query = llm.share(prompt=self._query_prompter).used_by(self._module_id)
        self._answer_prompter = ChatPrompter(
            instruction=db_explain_instruct_template.format(statement_type=statement_type, db_type=sql_manager.db_type)
        ).pre_hook(self.sql_explain_prompt_hook)
        self._llm_answer = llm.share(prompt=self._answer_prompter).used_by(self._module_id)
        self.example = sql_examples
        with pipeline() as sql_execute_ppl:
            sql_execute_ppl.exec = self._sql_tool.execute_query
            if use_llm_for_sql_result:
                sql_execute_ppl.concate = (lambda q, r: [q, r]) | bind(sql_execute_ppl.input, _0)
                sql_execute_ppl.llm_answer = self._llm_answer
        with pipeline() as ppl:
            ppl.llm_query = self._llm_query
            ppl.sql_extractor = self.extract_sql_from_response
            with switch(judge_on_full_input=False) as ppl.sw:
                ppl.sw.case[False, lambda x: x]
                ppl.sw.case[True, sql_execute_ppl]
        self._impl = ppl

    def sql_query_promt_hook(
        self,
        input: Union[str, List, Dict[str, str], None] = None,
        history: Optional[List[Union[List[str], Dict[str, Any]]]] = None,
        tools: Union[List[Dict[str, Any]], None] = None,
        label: Union[str, None] = None,
    ):
        """Hook to prepare the prompt inputs for generating a database query from user input.

Args:
    input (Union[str, List, Dict[str, str], None]): The user's natural language query.
    history (List[Union[List[str], Dict[str, Any]]]): Conversation history.
    tools (Union[List[Dict[str, Any]], None]): Available tool descriptions.
    label (Union[str, None]): Optional label for the prompt.

Returns:
    Tuple: A tuple containing the formatted prompt dict (with current_date, db_type, desc, user_query), history, tools, and label.
"""
        current_date = datetime.datetime.now().strftime("%Y-%m-%d")
        schema_desc = self._sql_tool.desc
        if self.example:
            schema_desc += f"\n{self.EXAMPLE_TITLE}\n{self.example}\n"
        if not isinstance(input, str):
            raise ValueError(f"Unexpected type for input: {type(input)}")
        return (
            dict(current_date=current_date, db_type=self._sql_tool.db_type, desc=schema_desc, user_query=input),
            history or [],
            tools,
            label,
        )

    def sql_explain_prompt_hook(
        self,
        input: Union[str, List, Dict[str, str], None] = None,
        history: List[Union[List[str], Dict[str, Any]]] = [],  # noqa B006
        tools: Union[List[Dict[str, Any]], None] = None,
        label: Union[str, None] = None,
    ):
        """Hook to prepare the prompt for explaining the execution result of a database query.

Args:
    input (Union[str, List, Dict[str, str], None]): A list containing the query and its result.
    history (List[Union[List[str], Dict[str, Any]]]): Conversation history.
    tools (Union[List[Dict[str, Any]], None]): Available tool descriptions.
    label (Union[str, None]): Optional label for the prompt.

Returns:
    Tuple: A tuple containing the formatted prompt dict (history_info, desc, query, result, explain_query), history, tools, and label.
"""
        explain_query = "Tell the user based on the execution results, making sure to keep the language consistent \
            with the user's input and don't translate original result."
        if not isinstance(input, list) and len(input) != 2:
            raise ValueError(f"Unexpected type for input: {type(input)}")
        assert "root_input" in globals and self._llm_answer._module_id in globals["root_input"]
        user_query = globals["root_input"][self._llm_answer._module_id]
        globals.pop("root_input")
        history_info = chat_history_to_str(history, user_query)
        return (
            dict(
                history_info=history_info,
                desc=self._sql_tool.desc,
                query=input[0],
                result=input[1],
                explain_query=explain_query,
            ),
            history,
            tools,
            label,
        )

    def extract_sql_from_response(self, str_response: str) -> tuple[bool, str]:
        """Extract SQL (or MongoDB pipeline) statement from the raw LLM response.

Args:
    str_response (str): Raw text returned by the LLM which may contain code fences.

Returns:
    tuple[bool, str]: A tuple where the first element indicates whether extraction succeeded, and the second is the cleaned or original content. If sql_post_func is provided, it is applied to the extracted content.
"""
        # Remove the triple backticks if present
        matches = self._pattern.findall(str_response)
        if matches:
            # Return the first match
            extracted_content = matches[0].strip()
            return True, extracted_content if not self.sql_post_func else self.sql_post_func(extracted_content)
        else:
            return False, str_response

    def forward(self, input: str, llm_chat_history: List[Dict[str, Any]] = None):
        globals["root_input"] = {self._llm_answer._module_id: input}
        if self._module_id in globals["chat_history"]:
            globals["chat_history"][self._llm_query._module_id] = globals["chat_history"][self._module_id]
        return self._impl(input)

extract_sql_from_response(str_response)

Extract SQL (or MongoDB pipeline) statement from the raw LLM response.

Parameters:

  • str_response (str) –

    Raw text returned by the LLM which may contain code fences.

Returns:

  • tuple[bool, str]

    tuple[bool, str]: A tuple where the first element indicates whether extraction succeeded, and the second is the cleaned or original content. If sql_post_func is provided, it is applied to the extracted content.

Source code in lazyllm/tools/sql_call/sql_call.py
    def extract_sql_from_response(self, str_response: str) -> tuple[bool, str]:
        """Extract SQL (or MongoDB pipeline) statement from the raw LLM response.

Args:
    str_response (str): Raw text returned by the LLM which may contain code fences.

Returns:
    tuple[bool, str]: A tuple where the first element indicates whether extraction succeeded, and the second is the cleaned or original content. If sql_post_func is provided, it is applied to the extracted content.
"""
        # Remove the triple backticks if present
        matches = self._pattern.findall(str_response)
        if matches:
            # Return the first match
            extracted_content = matches[0].strip()
            return True, extracted_content if not self.sql_post_func else self.sql_post_func(extracted_content)
        else:
            return False, str_response

sql_explain_prompt_hook(input=None, history=[], tools=None, label=None)

Hook to prepare the prompt for explaining the execution result of a database query.

Parameters:

  • input (Union[str, List, Dict[str, str], None], default: None ) –

    A list containing the query and its result.

  • history (List[Union[List[str], Dict[str, Any]]], default: [] ) –

    Conversation history.

  • tools (Union[List[Dict[str, Any]], None], default: None ) –

    Available tool descriptions.

  • label (Union[str, None], default: None ) –

    Optional label for the prompt.

Returns:

  • Tuple

    A tuple containing the formatted prompt dict (history_info, desc, query, result, explain_query), history, tools, and label.

Source code in lazyllm/tools/sql_call/sql_call.py
    def sql_explain_prompt_hook(
        self,
        input: Union[str, List, Dict[str, str], None] = None,
        history: List[Union[List[str], Dict[str, Any]]] = [],  # noqa B006
        tools: Union[List[Dict[str, Any]], None] = None,
        label: Union[str, None] = None,
    ):
        """Hook to prepare the prompt for explaining the execution result of a database query.

Args:
    input (Union[str, List, Dict[str, str], None]): A list containing the query and its result.
    history (List[Union[List[str], Dict[str, Any]]]): Conversation history.
    tools (Union[List[Dict[str, Any]], None]): Available tool descriptions.
    label (Union[str, None]): Optional label for the prompt.

Returns:
    Tuple: A tuple containing the formatted prompt dict (history_info, desc, query, result, explain_query), history, tools, and label.
"""
        explain_query = "Tell the user based on the execution results, making sure to keep the language consistent \
            with the user's input and don't translate original result."
        if not isinstance(input, list) and len(input) != 2:
            raise ValueError(f"Unexpected type for input: {type(input)}")
        assert "root_input" in globals and self._llm_answer._module_id in globals["root_input"]
        user_query = globals["root_input"][self._llm_answer._module_id]
        globals.pop("root_input")
        history_info = chat_history_to_str(history, user_query)
        return (
            dict(
                history_info=history_info,
                desc=self._sql_tool.desc,
                query=input[0],
                result=input[1],
                explain_query=explain_query,
            ),
            history,
            tools,
            label,
        )

sql_query_promt_hook(input=None, history=None, tools=None, label=None)

Hook to prepare the prompt inputs for generating a database query from user input.

Parameters:

  • input (Union[str, List, Dict[str, str], None], default: None ) –

    The user's natural language query.

  • history (List[Union[List[str], Dict[str, Any]]], default: None ) –

    Conversation history.

  • tools (Union[List[Dict[str, Any]], None], default: None ) –

    Available tool descriptions.

  • label (Union[str, None], default: None ) –

    Optional label for the prompt.

Returns:

  • Tuple

    A tuple containing the formatted prompt dict (with current_date, db_type, desc, user_query), history, tools, and label.

Source code in lazyllm/tools/sql_call/sql_call.py
    def sql_query_promt_hook(
        self,
        input: Union[str, List, Dict[str, str], None] = None,
        history: Optional[List[Union[List[str], Dict[str, Any]]]] = None,
        tools: Union[List[Dict[str, Any]], None] = None,
        label: Union[str, None] = None,
    ):
        """Hook to prepare the prompt inputs for generating a database query from user input.

Args:
    input (Union[str, List, Dict[str, str], None]): The user's natural language query.
    history (List[Union[List[str], Dict[str, Any]]]): Conversation history.
    tools (Union[List[Dict[str, Any]], None]): Available tool descriptions.
    label (Union[str, None]): Optional label for the prompt.

Returns:
    Tuple: A tuple containing the formatted prompt dict (with current_date, db_type, desc, user_query), history, tools, and label.
"""
        current_date = datetime.datetime.now().strftime("%Y-%m-%d")
        schema_desc = self._sql_tool.desc
        if self.example:
            schema_desc += f"\n{self.EXAMPLE_TITLE}\n{self.example}\n"
        if not isinstance(input, str):
            raise ValueError(f"Unexpected type for input: {type(input)}")
        return (
            dict(current_date=current_date, db_type=self._sql_tool.db_type, desc=schema_desc, user_query=input),
            history or [],
            tools,
            label,
        )

lazyllm.tools.rag.default_index.DefaultIndex

Bases: IndexBase

Default index implementation responsible for querying, updating, and removing document nodes in the underlying store using embedding or text similarity. Supports multiple similarity metrics and performs embedding computation and node updates when needed.

Parameters:

  • embed (Dict[str, Callable]) –

    Mapping of embedding names to functions that generate vector representations from strings.

  • store (StoreBase) –

    Underlying storage to persist and retrieve DocNode objects.

  • **kwargs

    Reserved for future extension.

Source code in lazyllm/tools/rag/default_index.py
class DefaultIndex(IndexBase):
    """Default index implementation responsible for querying, updating, and removing document nodes in the underlying store using embedding or text similarity. Supports multiple similarity metrics and performs embedding computation and node updates when needed.

Args:
    embed (Dict[str, Callable]): Mapping of embedding names to functions that generate vector representations from strings.
    store (StoreBase): Underlying storage to persist and retrieve DocNode objects.
    **kwargs: Reserved for future extension.
"""
    def __init__(self, embed: Dict[str, Callable], store, **kwargs):
        self.embed = embed
        self.store = store

    @override
    def update(self, nodes: List[DocNode]) -> None:
        """Update the index with the given list of document nodes. This is a placeholder implementation and should be provided/extended in concrete usage.

Args:
    nodes (List[DocNode]): Document nodes to add or update in the index.
"""
        pass

    @override
    def remove(self, uids: List[str], group_name: Optional[str] = None) -> None:
        """Remove nodes with specified UIDs from the index. Optionally scoped to a group. This is a no-op placeholder and should be implemented in concrete usage.

Args:
    uids (List[str]): List of unique IDs of nodes to remove.
    group_name (Optional[str]): Optional group name to scope the removal.
"""
        pass

    @override
    def query(
        self,
        query: str,
        group_name: str,
        similarity_name: str,
        similarity_cut_off: Union[float, Dict[str, float]],
        topk: int,
        embed_keys: Optional[List[str]] = None,
        filters: Optional[Dict[str, List]] = None,
        **kwargs,
    ) -> List[DocNode]:
        """Perform a query against the index, supporting both embedding-based and text-based similarity modes. Filters and ranks nodes according to similarity functions and cutoffs.

Args:
    query (str): The raw query string.
    group_name (str): The group name from which to retrieve nodes.
    similarity_name (str): Name of the similarity metric to use; must be registered in registered_similarities.
    similarity_cut_off (Union[float, Dict[str, float]]): Similarity threshold(s) used to filter results; can be a single float or a mapping per embedding.
    topk (int): Maximum number of candidates to keep per similarity channel before final filtering.
    embed_keys (Optional[List[str]]): Specific embedding keys to use; defaults to all available if not provided.
    filters (Optional[Dict[str, List]]): Additional pre-filters applied to nodes before similarity computation.
    **kwargs: Extra keyword arguments forwarded to the similarity function.

**Returns**

    - list: List[DocNode]: Deduplicated list of document nodes passing similarity and cutoff criteria.
"""
        if similarity_name not in registered_similarities:
            raise ValueError(
                f"{similarity_name} not registered, please check your input. "
                f"Available options now: {registered_similarities.keys()}"
            )
        similarity_func, mode, descend = registered_similarities[similarity_name]

        nodes = self.store.get_nodes(group=group_name)
        if filters:
            nodes = generic_process_filters(nodes, filters)

        if mode == "embedding":
            assert self.embed, "Chosen similarity needs embed model."
            assert len(query) > 0, "Query should not be empty."
            if not embed_keys:
                embed_keys = list(self.embed.keys())
            query_embedding = {k: self.embed[k](query) for k in embed_keys}
            self._check_supported(similarity_name, query_embedding)
            modified_nodes = parallel_do_embedding(self.embed, embed_keys, nodes)
            self.store.update_nodes(modified_nodes)
            similarities = similarity_func(query_embedding, nodes, topk=topk, **kwargs)
        elif mode == "text":
            similarities = similarity_func(query, nodes, topk=topk, **kwargs)
        else:
            raise NotImplementedError(f"Mode {mode} is not supported.")

        if not isinstance(similarities, dict):
            results = self._filter_nodes_by_score(similarities, topk, similarity_cut_off, descend)
        else:
            results = []
            for key in (embed_keys or similarities.keys()):
                sims = similarities[key]
                sim_cut_off = similarity_cut_off if isinstance(similarity_cut_off, float) else similarity_cut_off[key]
                results.extend(self._filter_nodes_by_score(sims, topk, sim_cut_off, descend))
        results = list(set(results))
        LOG.debug(f"Retrieving query `{query}` and get results: {results}")
        return results

    def _filter_nodes_by_score(self, similarities: List[Tuple[DocNode, float]], topk: int,
                               similarity_cut_off: float, descend) -> List[DocNode]:
        similarities.sort(key=lambda x: x[1], reverse=descend)
        if topk is not None:
            similarities = similarities[:topk]

        return [node.with_sim_score(score) for node, score in similarities if score > similarity_cut_off]

    def _check_supported(self, similarity_name: str, query_embedding: Dict[str, Any]) -> None:
        if similarity_name.lower() == 'cosine':
            for k, e in query_embedding.items():
                if is_sparse(e):
                    raise NotImplementedError(f'embed `{k}`, which is sparse, is not supported.')

query(query, group_name, similarity_name, similarity_cut_off, topk, embed_keys=None, filters=None, **kwargs)

Perform a query against the index, supporting both embedding-based and text-based similarity modes. Filters and ranks nodes according to similarity functions and cutoffs.

Parameters:

  • query (str) –

    The raw query string.

  • group_name (str) –

    The group name from which to retrieve nodes.

  • similarity_name (str) –

    Name of the similarity metric to use; must be registered in registered_similarities.

  • similarity_cut_off (Union[float, Dict[str, float]]) –

    Similarity threshold(s) used to filter results; can be a single float or a mapping per embedding.

  • topk (int) –

    Maximum number of candidates to keep per similarity channel before final filtering.

  • embed_keys (Optional[List[str]], default: None ) –

    Specific embedding keys to use; defaults to all available if not provided.

  • filters (Optional[Dict[str, List]], default: None ) –

    Additional pre-filters applied to nodes before similarity computation.

  • **kwargs

    Extra keyword arguments forwarded to the similarity function.

Returns

- list: List[DocNode]: Deduplicated list of document nodes passing similarity and cutoff criteria.
Source code in lazyllm/tools/rag/default_index.py
    @override
    def query(
        self,
        query: str,
        group_name: str,
        similarity_name: str,
        similarity_cut_off: Union[float, Dict[str, float]],
        topk: int,
        embed_keys: Optional[List[str]] = None,
        filters: Optional[Dict[str, List]] = None,
        **kwargs,
    ) -> List[DocNode]:
        """Perform a query against the index, supporting both embedding-based and text-based similarity modes. Filters and ranks nodes according to similarity functions and cutoffs.

Args:
    query (str): The raw query string.
    group_name (str): The group name from which to retrieve nodes.
    similarity_name (str): Name of the similarity metric to use; must be registered in registered_similarities.
    similarity_cut_off (Union[float, Dict[str, float]]): Similarity threshold(s) used to filter results; can be a single float or a mapping per embedding.
    topk (int): Maximum number of candidates to keep per similarity channel before final filtering.
    embed_keys (Optional[List[str]]): Specific embedding keys to use; defaults to all available if not provided.
    filters (Optional[Dict[str, List]]): Additional pre-filters applied to nodes before similarity computation.
    **kwargs: Extra keyword arguments forwarded to the similarity function.

**Returns**

    - list: List[DocNode]: Deduplicated list of document nodes passing similarity and cutoff criteria.
"""
        if similarity_name not in registered_similarities:
            raise ValueError(
                f"{similarity_name} not registered, please check your input. "
                f"Available options now: {registered_similarities.keys()}"
            )
        similarity_func, mode, descend = registered_similarities[similarity_name]

        nodes = self.store.get_nodes(group=group_name)
        if filters:
            nodes = generic_process_filters(nodes, filters)

        if mode == "embedding":
            assert self.embed, "Chosen similarity needs embed model."
            assert len(query) > 0, "Query should not be empty."
            if not embed_keys:
                embed_keys = list(self.embed.keys())
            query_embedding = {k: self.embed[k](query) for k in embed_keys}
            self._check_supported(similarity_name, query_embedding)
            modified_nodes = parallel_do_embedding(self.embed, embed_keys, nodes)
            self.store.update_nodes(modified_nodes)
            similarities = similarity_func(query_embedding, nodes, topk=topk, **kwargs)
        elif mode == "text":
            similarities = similarity_func(query, nodes, topk=topk, **kwargs)
        else:
            raise NotImplementedError(f"Mode {mode} is not supported.")

        if not isinstance(similarities, dict):
            results = self._filter_nodes_by_score(similarities, topk, similarity_cut_off, descend)
        else:
            results = []
            for key in (embed_keys or similarities.keys()):
                sims = similarities[key]
                sim_cut_off = similarity_cut_off if isinstance(similarity_cut_off, float) else similarity_cut_off[key]
                results.extend(self._filter_nodes_by_score(sims, topk, sim_cut_off, descend))
        results = list(set(results))
        LOG.debug(f"Retrieving query `{query}` and get results: {results}")
        return results

remove(uids, group_name=None)

Remove nodes with specified UIDs from the index. Optionally scoped to a group. This is a no-op placeholder and should be implemented in concrete usage.

Parameters:

  • uids (List[str]) –

    List of unique IDs of nodes to remove.

  • group_name (Optional[str], default: None ) –

    Optional group name to scope the removal.

Source code in lazyllm/tools/rag/default_index.py
    @override
    def remove(self, uids: List[str], group_name: Optional[str] = None) -> None:
        """Remove nodes with specified UIDs from the index. Optionally scoped to a group. This is a no-op placeholder and should be implemented in concrete usage.

Args:
    uids (List[str]): List of unique IDs of nodes to remove.
    group_name (Optional[str]): Optional group name to scope the removal.
"""
        pass

update(nodes)

Update the index with the given list of document nodes. This is a placeholder implementation and should be provided/extended in concrete usage.

Parameters:

  • nodes (List[DocNode]) –

    Document nodes to add or update in the index.

Source code in lazyllm/tools/rag/default_index.py
    @override
    def update(self, nodes: List[DocNode]) -> None:
        """Update the index with the given list of document nodes. This is a placeholder implementation and should be provided/extended in concrete usage.

Args:
    nodes (List[DocNode]): Document nodes to add or update in the index.
"""
        pass

lazyllm.tools.Reranker

Bases: ModuleBase, _PostProcess

Initializes a Rerank module for postprocessing and reranking of nodes (documents). This constructor initializes a Reranker module that configures a reranking process based on a specified reranking type. It allows for the dynamic selection and instantiation of reranking kernels (algorithms) based on the type and provided keyword arguments.

Parameters:

  • name (str, default: 'ModuleReranker' ) –

    The type of reranker used for the postprocessing and reranking process. Defaults to 'ModuleReranker'.

  • target (str, default: None ) –

    Deprecated parameter, only used to notify users.

  • output_format (Optional[str], default: None ) –

    Specifies the output format. Defaults to None. Optional values include 'content' and 'dict'. - 'content' means the output is in string format. - 'dict' means the output is a dictionary.

  • join (Union[bool, str], default: False ) –

    Determines whether to join the top-k output nodes. - When output_format is 'content': - If set to True, returns a single long string. - If set to False, returns a list of strings, each representing one node’s content. - When output_format is 'dict': - Joining is not supported; join defaults to False. - Returns a dictionary with three keys: 'content', 'embedding', and 'metadata'.

  • kwargs

    Additional keyword arguments passed to the reranker upon instantiation.

Detailed explanation of reranker types

  • Reranker: Instantiates a SentenceTransformerRerank reranker with a list of document nodes and a query.

  • KeywordFilter: This registered reranking function instantiates a KeywordNodePostprocessor with specified required and excluded keywords. It filters nodes based on the presence or absence of these keywords.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Document, Reranker, Retriever, DocNode
>>> m = lazyllm.OnlineEmbeddingModule()
>>> documents = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
>>> retriever = Retriever(documents, group_name='CoarseChunk', similarity='bm25', similarity_cut_off=0.01, topk=6)
>>> reranker = Reranker(DocNode(text=user_data),query="user query")
>>> ppl = lazyllm.ActionModule(retriever, reranker)
>>> ppl.start()
>>> print(ppl("user query"))
Source code in lazyllm/tools/rag/rerank.py
class Reranker(ModuleBase, _PostProcess):
    """Initializes a Rerank module for postprocessing and reranking of nodes (documents).
This constructor initializes a Reranker module that configures a reranking process based on a specified reranking type. It allows for the dynamic selection and instantiation of reranking kernels (algorithms) based on the type and provided keyword arguments.

Args:
    name: The type of reranker used for the postprocessing and reranking process. Defaults to 'ModuleReranker'.
    target (str): **Deprecated** parameter, only used to notify users.
    output_format: Specifies the output format. Defaults to None. Optional values include 'content' and 'dict'. 
        - 'content' means the output is in string format.
        - 'dict' means the output is a dictionary.
    join: Determines whether to join the top-k output nodes.
        - When `output_format` is 'content':
            - If set to True, returns a single long string.
            - If set to False, returns a list of strings, each representing one node’s content.
        - When `output_format` is 'dict':
            - Joining is not supported; `join` defaults to False.
            - Returns a dictionary with three keys: 'content', 'embedding', and 'metadata'.
    kwargs: Additional keyword arguments passed to the reranker upon instantiation.
**Detailed explanation of reranker types**

- Reranker: Instantiates a `SentenceTransformerRerank` reranker with a list of document nodes and a query.

- KeywordFilter: This registered reranking function instantiates a KeywordNodePostprocessor with specified required and excluded keywords. It filters nodes based on the presence or absence of these keywords.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, Reranker, Retriever, DocNode
    >>> m = lazyllm.OnlineEmbeddingModule()
    >>> documents = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
    >>> retriever = Retriever(documents, group_name='CoarseChunk', similarity='bm25', similarity_cut_off=0.01, topk=6)
    >>> reranker = Reranker(DocNode(text=user_data),query="user query")
    >>> ppl = lazyllm.ActionModule(retriever, reranker)
    >>> ppl.start()
    >>> print(ppl("user query"))
    """
    registered_reranker = dict()

    def __new__(cls, name: str = "ModuleReranker", *args, **kwargs):
        assert name in cls.registered_reranker, f"Reranker: {name} is not registered, please register first."
        item = cls.registered_reranker[name]
        if isinstance(item, type) and issubclass(item, Reranker):
            return super(Reranker, cls).__new__(item)
        else:
            return super(Reranker, cls).__new__(cls)

    def __init__(self, name: str = "ModuleReranker", target: Optional[str] = None,
                 output_format: Optional[str] = None, join: Union[bool, str] = False, **kwargs) -> None:
        super().__init__()
        self._name = name
        self._kwargs = kwargs
        lazyllm.deprecated(bool(target), '`target` parameter of reranker')
        _PostProcess.__init__(self, output_format, join)

    def forward(self, nodes: List[DocNode], query: str = "") -> List[DocNode]:
        results = self.registered_reranker[self._name](nodes, query=query, **self._kwargs)
        LOG.debug(f"Rerank use `{self._name}` and get nodes: {results}")
        return self._post_process(results)

    @classmethod
    def register_reranker(
        cls: "Reranker", func: Optional[Callable] = None, batch: bool = False
    ):
        """A class decorator factory method that provides a flexible mechanism for registering custom reranking algorithms to the `Reranker` class.
Args:
    func (Optional[Callable]): The reranking function or class to register. This can be omitted when using decorator syntax (@).
    batch (bool): Whether to process nodes in batches. Defaults to False, meaning nodes are processed individually.


Examples:

    @Reranker.register_reranker
    def my_reranker(node: DocNode, **kwargs):
        return node.score * 0.8  # 自定义分数计算
    """
        def decorator(f):
            if isinstance(f, type):
                cls.registered_reranker[f.__name__] = f
                return f
            else:
                def wrapper(nodes, **kwargs):
                    if batch:
                        return f(nodes, **kwargs)
                    else:
                        results = [f(node, **kwargs) for node in nodes]
                        return [result for result in results if result]

                cls.registered_reranker[f.__name__] = wrapper
                return wrapper

        return decorator(func) if func else decorator

register_reranker(func=None, batch=False) classmethod

A class decorator factory method that provides a flexible mechanism for registering custom reranking algorithms to the Reranker class. Args: func (Optional[Callable]): The reranking function or class to register. This can be omitted when using decorator syntax (@). batch (bool): Whether to process nodes in batches. Defaults to False, meaning nodes are processed individually.

Examples:

@Reranker.register_reranker
def my_reranker(node: DocNode, **kwargs):
    return node.score * 0.8  # 自定义分数计算
Source code in lazyllm/tools/rag/rerank.py
    @classmethod
    def register_reranker(
        cls: "Reranker", func: Optional[Callable] = None, batch: bool = False
    ):
        """A class decorator factory method that provides a flexible mechanism for registering custom reranking algorithms to the `Reranker` class.
Args:
    func (Optional[Callable]): The reranking function or class to register. This can be omitted when using decorator syntax (@).
    batch (bool): Whether to process nodes in batches. Defaults to False, meaning nodes are processed individually.


Examples:

    @Reranker.register_reranker
    def my_reranker(node: DocNode, **kwargs):
        return node.score * 0.8  # 自定义分数计算
    """
        def decorator(f):
            if isinstance(f, type):
                cls.registered_reranker[f.__name__] = f
                return f
            else:
                def wrapper(nodes, **kwargs):
                    if batch:
                        return f(nodes, **kwargs)
                    else:
                        results = [f(node, **kwargs) for node in nodes]
                        return [result for result in results if result]

                cls.registered_reranker[f.__name__] = wrapper
                return wrapper

        return decorator(func) if func else decorator

lazyllm.tools.Retriever

Bases: ModuleBase, _PostProcess

Create a retrieval module for document querying and retrieval. This constructor initializes a retrieval module that configures the document retrieval process based on the specified similarity metric.

Parameters:

  • doc (object) –

    An instance of the document module. The document module can be a single instance or a list of instances. If it is a single instance, it means searching for a single Document, and if it is a list of instances, it means searching for multiple Documents.

  • group_name (str) –

    The name of the node group on which to perform the retrieval.

  • similarity (Optional[str], default: None ) –

    The similarity function to use for setting up document retrieval. Defaults to 'dummy'. Candidates include ["bm25", "bm25_chinese", "cosine"].

  • similarity_cut_off (Union[float, Dict[str, float]], default: float('-inf') ) –

    Discard the document when the similarity is below the specified value. In a multi-embedding scenario, if you need to specify different values for different embeddings, you need to specify them in a dictionary, where the key indicates which embedding is specified and the value indicates the corresponding threshold. If all embeddings use the same threshold, you only need to specify one value.

  • index (str, default: 'default' ) –

    The type of index to use for document retrieval. Currently, only 'default' is supported.

  • topk (int, default: 6 ) –

    The number of documents to retrieve with the highest similarity.

  • embed_keys (Optional[List[str]], default: None ) –

    Indicates which embeddings are used for retrieval. If not specified, all embeddings are used for retrieval.

  • target (Optional[str], default: None ) –

    The name of the target document group for result conversion

  • output_format (Optional[str], default: None ) –

    Represents the output format, with a default value of None. Optional values include 'content' and 'dict', where 'content' corresponds to a string output format and 'dict' corresponds to a dictionary.

  • join (Union[bool, str], default: False ) –

    Determines whether to concatenate the output of k nodes - when output format is 'content', setting True returns a single concatenated string while False returns a list of strings (each corresponding to a node's text content); when output format is 'dict', joining is unsupported (join defaults to False) and the output will be a dictionary containing 'content', 'embedding' and 'metadata' keys.

The group_name has three built-in splitting strategies, all of which use SentenceSplitter for splitting, with the difference being in the chunk size:

  • CoarseChunk: Chunk size is 1024, with an overlap length of 100
  • MediumChunk: Chunk size is 256, with an overlap length of 25
  • FineChunk: Chunk size is 128, with an overlap length of 12

Also, Image is available for group_name since LazyLLM supports image embedding and retrieval.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Retriever, Document, SentenceSplitter
>>> m = lazyllm.OnlineEmbeddingModule()
>>> documents = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
>>> rm = Retriever(documents, group_name='CoarseChunk', similarity='bm25', similarity_cut_off=0.01, topk=6)
>>> rm.start()
>>> print(rm("user query"))
>>> m1 = lazyllm.TrainableModule('bge-large-zh-v1.5').start()
>>> document1 = Document(dataset_path='/path/to/user/data', embed={'online':m , 'local': m1}, manager=False)
>>> document1.create_node_group(name='sentences', transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
>>> retriever = Retriever(document1, group_name='sentences', similarity='cosine', similarity_cut_off=0.4, embed_keys=['local'], topk=3)
>>> print(retriever("user query"))
>>> document2 = Document(dataset_path='/path/to/user/data', embed={'online':m , 'local': m1}, manager=False)
>>> document2.create_node_group(name='sentences', transform=SentenceSplitter, chunk_size=512, chunk_overlap=50)
>>> retriever2 = Retriever([document1, document2], group_name='sentences', similarity='cosine', similarity_cut_off=0.4, embed_keys=['local'], topk=3)
>>> print(retriever2("user query"))
>>>
>>> filters = {
>>>     "author": ["A", "B", "C"],
>>>     "public_year": [2002, 2003, 2004],
>>> }
>>> document3 = Document(dataset_path='/path/to/user/data', embed={'online':m , 'local': m1}, manager=False)
>>> document3.create_node_group(name='sentences', transform=SentenceSplitter, chunk_size=512, chunk_overlap=50)
>>> retriever3 = Retriever([document1, document3], group_name='sentences', similarity='cosine', similarity_cut_off=0.4, embed_keys=['local'], topk=3)
>>> print(retriever3(query="user query", filters=filters))
>>> document4 = Document(dataset_path='/path/to/user/data', embed=lazyllm.TrainableModule('siglip'))
>>> retriever4 = Retriever(document4, group_name='Image', similarity='cosine')
>>> nodes = retriever4("user query")
>>> print([node.get_content() for node in nodes])
>>> document5 = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
>>> rm = Retriever(document5, group_name='CoarseChunk', similarity='bm25_chinese', similarity_cut_off=0.01, topk=3, output_format='content')
>>> rm.start()
>>> print(rm("user query"))
>>> document6 = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
>>> rm = Retriever(document6, group_name='CoarseChunk', similarity='bm25_chinese', similarity_cut_off=0.01, topk=3, output_format='content', join=True)
>>> rm.start()
>>> print(rm("user query"))
>>> document7 = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
>>> rm = Retriever(document7, group_name='CoarseChunk', similarity='bm25_chinese', similarity_cut_off=0.01, topk=3, output_format='dict')
>>> rm.start()
>>> print(rm("user query"))
Source code in lazyllm/tools/rag/retriever.py
class Retriever(ModuleBase, _PostProcess):
    """
Create a retrieval module for document querying and retrieval. This constructor initializes a retrieval module that configures the document retrieval process based on the specified similarity metric.

Args:
    doc: An instance of the document module. The document module can be a single instance or a list of instances. If it is a single instance, it means searching for a single Document, and if it is a list of instances, it means searching for multiple Documents.
    group_name: The name of the node group on which to perform the retrieval.
    similarity: The similarity function to use for setting up document retrieval. Defaults to 'dummy'. Candidates include ["bm25", "bm25_chinese", "cosine"].
    similarity_cut_off: Discard the document when the similarity is below the specified value. In a multi-embedding scenario, if you need to specify different values for different embeddings, you need to specify them in a dictionary, where the key indicates which embedding is specified and the value indicates the corresponding threshold. If all embeddings use the same threshold, you only need to specify one value.
    index: The type of index to use for document retrieval. Currently, only 'default' is supported.
    topk: The number of documents to retrieve with the highest similarity.
    embed_keys: Indicates which embeddings are used for retrieval. If not specified, all embeddings are used for retrieval.
    target:The name of the target document group for result conversion
    output_format: Represents the output format, with a default value of None. Optional values include 'content' and 'dict', where 'content' corresponds to a string output format and 'dict' corresponds to a dictionary.
    join:  Determines whether to concatenate the output of k nodes - when output format is 'content', setting True returns a single concatenated string while False returns a list of strings (each corresponding to a node's text content); when output format is 'dict', joining is unsupported (join defaults to False) and the output will be a dictionary containing 'content', 'embedding' and 'metadata' keys.

The `group_name` has three built-in splitting strategies, all of which use `SentenceSplitter` for splitting, with the difference being in the chunk size:

- CoarseChunk: Chunk size is 1024, with an overlap length of 100
- MediumChunk: Chunk size is 256, with an overlap length of 25
- FineChunk: Chunk size is 128, with an overlap length of 12

Also, `Image` is available for `group_name` since LazyLLM supports image embedding and retrieval.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Retriever, Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule()
    >>> documents = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
    >>> rm = Retriever(documents, group_name='CoarseChunk', similarity='bm25', similarity_cut_off=0.01, topk=6)
    >>> rm.start()
    >>> print(rm("user query"))
    >>> m1 = lazyllm.TrainableModule('bge-large-zh-v1.5').start()
    >>> document1 = Document(dataset_path='/path/to/user/data', embed={'online':m , 'local': m1}, manager=False)
    >>> document1.create_node_group(name='sentences', transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    >>> retriever = Retriever(document1, group_name='sentences', similarity='cosine', similarity_cut_off=0.4, embed_keys=['local'], topk=3)
    >>> print(retriever("user query"))
    >>> document2 = Document(dataset_path='/path/to/user/data', embed={'online':m , 'local': m1}, manager=False)
    >>> document2.create_node_group(name='sentences', transform=SentenceSplitter, chunk_size=512, chunk_overlap=50)
    >>> retriever2 = Retriever([document1, document2], group_name='sentences', similarity='cosine', similarity_cut_off=0.4, embed_keys=['local'], topk=3)
    >>> print(retriever2("user query"))
    >>>
    >>> filters = {
    >>>     "author": ["A", "B", "C"],
    >>>     "public_year": [2002, 2003, 2004],
    >>> }
    >>> document3 = Document(dataset_path='/path/to/user/data', embed={'online':m , 'local': m1}, manager=False)
    >>> document3.create_node_group(name='sentences', transform=SentenceSplitter, chunk_size=512, chunk_overlap=50)
    >>> retriever3 = Retriever([document1, document3], group_name='sentences', similarity='cosine', similarity_cut_off=0.4, embed_keys=['local'], topk=3)
    >>> print(retriever3(query="user query", filters=filters))
    >>> document4 = Document(dataset_path='/path/to/user/data', embed=lazyllm.TrainableModule('siglip'))
    >>> retriever4 = Retriever(document4, group_name='Image', similarity='cosine')
    >>> nodes = retriever4("user query")
    >>> print([node.get_content() for node in nodes])
    >>> document5 = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
    >>> rm = Retriever(document5, group_name='CoarseChunk', similarity='bm25_chinese', similarity_cut_off=0.01, topk=3, output_format='content')
    >>> rm.start()
    >>> print(rm("user query"))
    >>> document6 = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
    >>> rm = Retriever(document6, group_name='CoarseChunk', similarity='bm25_chinese', similarity_cut_off=0.01, topk=3, output_format='content', join=True)
    >>> rm.start()
    >>> print(rm("user query"))
    >>> document7 = Document(dataset_path='/path/to/user/data', embed=m, manager=False)
    >>> rm = Retriever(document7, group_name='CoarseChunk', similarity='bm25_chinese', similarity_cut_off=0.01, topk=3, output_format='dict')
    >>> rm.start()
    >>> print(rm("user query"))
    """
    def __init__(self, doc: object, group_name: str, similarity: Optional[str] = None,
                 similarity_cut_off: Union[float, Dict[str, float]] = float("-inf"), index: str = "default",
                 topk: int = 6, embed_keys: Optional[List[str]] = None, target: Optional[str] = None,
                 output_format: Optional[str] = None, join: Union[bool, str] = False, **kwargs):
        super().__init__()

        if similarity:
            _, mode, _ = registered_similarities[similarity]
        else:
            mode = 'embedding'  # TODO FIXME XXX should be removed after similarity args refactor
        group_name, target = str(group_name), (str(target) if target else None)

        self._docs: List[Document] = [doc] if isinstance(doc, Document) else doc
        for doc in self._docs:
            assert isinstance(doc, (Document, UrlDocument)), 'Only Document or List[Document] are supported'
            if isinstance(doc, UrlDocument): continue
            self._submodules.append(doc)
            if mode == 'embedding' and embed_keys is None:
                embed_keys = list(doc._impl.embed.keys())
            doc.activate_group(group_name, embed_keys)
            if target: doc.activate_group(target)

        self._group_name = group_name
        self._similarity = similarity  # similarity function str
        self._similarity_cut_off = similarity_cut_off
        self._index = index
        self._topk = topk
        self._similarity_kw = kwargs  # kw parameters
        self._embed_keys = embed_keys
        self._target = target
        _PostProcess.__init__(self, output_format, join)

    @once_wrapper
    def _lazy_init(self):
        docs = [doc for doc in self._docs if isinstance(doc, UrlDocument) or self._group_name in doc._impl.node_groups
                or self._group_name in DocImpl._builtin_node_groups or self._group_name in DocImpl._global_node_groups]
        if not docs: raise RuntimeError(f'Group {self._group_name} not found in document {self._docs}')
        self._docs = docs

    def forward(
            self, query: str, filters: Optional[Dict[str, Union[str, int, List, Set]]] = None,
            **kwargs
    ) -> Union[List[DocNode], str]:
        self._lazy_init()
        all_nodes: List[DocNode] = []
        for doc in self._docs:
            nodes = doc.forward(query=query, group_name=self._group_name, similarity=self._similarity,
                                similarity_cut_off=self._similarity_cut_off, index=self._index,
                                topk=self._topk, similarity_kws=self._similarity_kw, embed_keys=self._embed_keys,
                                filters=filters, **kwargs)
            if nodes and self._target and self._target != nodes[0]._group:
                nodes = doc.find(self._target)(nodes)
            all_nodes.extend(nodes)
        return self._post_process(all_nodes)

lazyllm.tools.rag.retriever.TempDocRetriever

Bases: ModuleBase, _PostProcess

A temporary document retriever that inherits from ModuleBase and _PostProcess, used for quickly processing temporary files and performing retrieval tasks.

Parameters:

  • embed (Callable, default: None ) –

    The embedding function.

  • output_format (Optional[str], default: None ) –

    The format of the output result (e.g., JSON). Optional, defaults to None.

  • join (Union[bool, str], default: False ) –

    Whether to merge multiple result segments (set to True or specify a separator like "

").

Examples:

>>> import lazyllm
>>> from lazyllm.tools import TempDocRetriever, Document, SentenceSplitter
>>> retriever = TempDocRetriever(output_format="text", join="
---------------
")
    retriever.create_node_group(transform=lambda text: [s.strip() for s in text.split("。") if s] )
    retriever.add_subretriever(group=Document.MediumChunk, topk=3)
    files = ["机器学习是AI的核心领域。深度学习是其重要分支。"]
    results = retriever.forward(files, "什么是机器学习?")
    print(results)
Source code in lazyllm/tools/rag/retriever.py
class TempDocRetriever(ModuleBase, _PostProcess):
    """
A temporary document retriever that inherits from ModuleBase and _PostProcess, used for quickly processing temporary files and performing retrieval tasks.

Args:
    embed: The embedding function.
    output_format: The format of the output result (e.g., JSON). Optional, defaults to None.
    join: Whether to merge multiple result segments (set to True or specify a separator like "
").


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import TempDocRetriever, Document, SentenceSplitter
    >>> retriever = TempDocRetriever(output_format="text", join="
    ---------------
    ")
        retriever.create_node_group(transform=lambda text: [s.strip() for s in text.split("。") if s] )
        retriever.add_subretriever(group=Document.MediumChunk, topk=3)
        files = ["机器学习是AI的核心领域。深度学习是其重要分支。"]
        results = retriever.forward(files, "什么是机器学习?")
        print(results)
    """
    def __init__(self, embed: Callable = None, output_format: Optional[str] = None, join: Union[bool, str] = False):
        super().__init__()
        self._doc = Document(doc_files=[])
        self._embed = embed
        self._node_groups = []
        _PostProcess.__init__(self, output_format, join)

    def create_node_group(self, name: str = None, *, transform: Callable, parent: str = LAZY_ROOT_NAME,
                          trans_node: bool = None, num_workers: int = 0, **kwargs):
        """
Create a node group with specific processing pipeline.

Args:
    name (str): Name of the node group. Auto-generated if None.
    transform (Callable): Function to process documents in this group.
    parent (str): Parent group name. Defaults to root group.
    trans_node (bool): Whether to transform nodes. Inherits from parent if None.
    num_workers (int): Parallel workers for processing. Default 0 (sequential).
    **kwargs: Additional group parameters.
"""
        self._doc.create_node_group(name, transform=transform, parent=parent,
                                    trans_node=trans_node, num_workers=num_workers, **kwargs)
        return self

    def add_subretriever(self, group: str, **kwargs):
        """
Add a sub-retriever with search configuration.

Args:
    group (str): Target node group name.
    **kwargs: Retriever parameters (e.g., similarity='cosine').

**Returns:**

- self: For method chaining.
"""
        if 'similarity' not in kwargs: kwargs['similarity'] = ('cosine' if self._embed else 'bm25')
        self._node_groups.append((group, kwargs))
        return self

    @functools.lru_cache    # noqa: B019
    def _get_retrievers(self, doc_files: List[str]):
        active_node_groups = self._node_groups or [[Document.MediumChunk,
                                                    dict(similarity=('cosine' if self._embed else 'bm25'))]]
        doc = Document(embed=self._embed, doc_files=doc_files)
        doc._impl.node_groups = self._doc._impl.node_groups
        retrievers = [Retriever(doc, name, **kw) for (name, kw) in active_node_groups]
        return retrievers

    def forward(self, files: Union[str, List[str]], query: str):
        if isinstance(files, str): files = [files]
        retrievers = self._get_retrievers(doc_files=tuple(set(files)))
        r = lazyllm.parallel(*retrievers).sum
        return self._post_process(r(query))

add_subretriever(group, **kwargs)

Add a sub-retriever with search configuration.

Parameters:

  • group (str) –

    Target node group name.

  • **kwargs

    Retriever parameters (e.g., similarity='cosine').

Returns:

  • self: For method chaining.
Source code in lazyllm/tools/rag/retriever.py
    def add_subretriever(self, group: str, **kwargs):
        """
Add a sub-retriever with search configuration.

Args:
    group (str): Target node group name.
    **kwargs: Retriever parameters (e.g., similarity='cosine').

**Returns:**

- self: For method chaining.
"""
        if 'similarity' not in kwargs: kwargs['similarity'] = ('cosine' if self._embed else 'bm25')
        self._node_groups.append((group, kwargs))
        return self

create_node_group(name=None, *, transform, parent=LAZY_ROOT_NAME, trans_node=None, num_workers=0, **kwargs)

Create a node group with specific processing pipeline.

Parameters:

  • name (str, default: None ) –

    Name of the node group. Auto-generated if None.

  • transform (Callable) –

    Function to process documents in this group.

  • parent (str, default: LAZY_ROOT_NAME ) –

    Parent group name. Defaults to root group.

  • trans_node (bool, default: None ) –

    Whether to transform nodes. Inherits from parent if None.

  • num_workers (int, default: 0 ) –

    Parallel workers for processing. Default 0 (sequential).

  • **kwargs

    Additional group parameters.

Source code in lazyllm/tools/rag/retriever.py
    def create_node_group(self, name: str = None, *, transform: Callable, parent: str = LAZY_ROOT_NAME,
                          trans_node: bool = None, num_workers: int = 0, **kwargs):
        """
Create a node group with specific processing pipeline.

Args:
    name (str): Name of the node group. Auto-generated if None.
    transform (Callable): Function to process documents in this group.
    parent (str): Parent group name. Defaults to root group.
    trans_node (bool): Whether to transform nodes. Inherits from parent if None.
    num_workers (int): Parallel workers for processing. Default 0 (sequential).
    **kwargs: Additional group parameters.
"""
        self._doc.create_node_group(name, transform=transform, parent=parent,
                                    trans_node=trans_node, num_workers=num_workers, **kwargs)
        return self

lazyllm.tools.rag.retriever.UrlDocument

Bases: ModuleBase

UrlDocument class inherits from ModuleBase, used to manage remote document resources by specifying a URL and a name.
Internally delegates calls to lazyllm's UrlModule, supporting document find, retrieve, and querying active node groups.

Parameters:

  • url (str) –

    Access URL for the remote document resource.

  • name (str, default: None ) –

    Current document group name used to identify the document group.

Source code in lazyllm/tools/rag/document.py
class UrlDocument(ModuleBase):
    """UrlDocument class inherits from ModuleBase, used to manage remote document resources by specifying a URL and a name.  
Internally delegates calls to lazyllm's UrlModule, supporting document find, retrieve, and querying active node groups.

Args:
    url (str): Access URL for the remote document resource.
    name (str): Current document group name used to identify the document group.
"""
    def __init__(self, url: str, name: str = None):
        super().__init__()
        self._missing_keys = set(dir(Document)) - set(dir(UrlDocument))
        self._manager = lazyllm.UrlModule(url=ensure_call_endpoint(url))
        self._curr_group = name or DocListManager.DEFAULT_GROUP_NAME

    def _forward(self, func_name: str, *args, **kwargs):
        args = (self._curr_group, func_name, *args)
        return self._manager._call("__call__", *args, **kwargs)

    def find(self, target) -> Callable:
        """Creates a partially applied function to find a specified target within the current document group.

Args:
    target (str): The target identifier to find.

**Returns:**

- Callable: A partially applied function that executes the find operation when called.
"""
        return functools.partial(self._forward, 'find', group=target)

    def forward(self, *args, **kw):
        return self._forward('retrieve', *args, **kw)

    @cached_property
    def active_node_groups(self):
        return self._forward('active_node_groups')

    def __getattr__(self, name):
        if name in self._missing_keys:
            raise RuntimeError(f'Document generated with url and name has no attribute `{name}`')

find(target)

Creates a partially applied function to find a specified target within the current document group.

Parameters:

  • target (str) –

    The target identifier to find.

Returns:

  • Callable: A partially applied function that executes the find operation when called.
Source code in lazyllm/tools/rag/document.py
    def find(self, target) -> Callable:
        """Creates a partially applied function to find a specified target within the current document group.

Args:
    target (str): The target identifier to find.

**Returns:**

- Callable: A partially applied function that executes the find operation when called.
"""
        return functools.partial(self._forward, 'find', group=target)

lazyllm.tools.rag.DocManager

Bases: ModuleBase

The DocManager class manages document lists and related operations, providing APIs for uploading, deleting, and grouping documents.

Parameters:

  • dlm (DocListManager) –

    Document list manager responsible for handling document-related operations.

Source code in lazyllm/tools/rag/doc_manager.py
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
class DocManager(lazyllm.ModuleBase):
    """
The `DocManager` class manages document lists and related operations, providing APIs for uploading, deleting, and grouping documents.

Args:
    dlm (DocListManager): Document list manager responsible for handling document-related operations.

"""
    def __init__(self, dlm: DocListManager) -> None:
        super().__init__()
        # disable path monitoring in case of competition adding/deleting files
        self._manager = dlm
        self._manager.enable_path_monitoring = False

    def __reduce__(self):
        self._manager.enable_path_monitoring = False
        return (__class__, (self._manager,))

    @app.get("/", response_model=BaseResponse, summary="docs")
    def document(self):
        """
An endpoint to redirect to the default documentation page.

**Returns:**

- RedirectResponse: Redirects to the `/docs` page.
"""
        return RedirectResponse(url="/docs")

    @app.get("/list_kb_groups")
    def list_kb_groups(self):
        """
An endpoint to list all document groups.

**Returns:**

- BaseResponse: Contains the data of all document groups.
"""
        try:
            return BaseResponse(data=self._manager.list_all_kb_group())
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    # returns an error message if invalid
    @staticmethod
    def _validate_metadata(metadata: Dict) -> Optional[str]:
        if metadata.get(RAG_DOC_ID):
            return f"metadata MUST not contain key `{RAG_DOC_ID}`"
        if metadata.get(RAG_DOC_PATH):
            return f"metadata MUST not contain key `{RAG_DOC_PATH}`"
        return None

    def _gen_unique_filepath(self, file_path: str) -> str:
        suffix = os.path.splitext(file_path)[1]
        prefix = file_path[0: len(file_path) - len(suffix)]
        pattern = f"{prefix}%{suffix}"
        MAX_TRIES = 10000
        exist_paths = set(self._manager.get_existing_paths_by_pattern(pattern))
        if file_path not in exist_paths:
            return file_path
        for i in range(1, MAX_TRIES):
            new_path = f"{prefix}-{i}{suffix}"
            if new_path not in exist_paths:
                return new_path
        return f"{str(uuid.uuid4())}{suffix}"

    @app.post("/upload_files")
    def upload_files(self, files: List[UploadFile], override: bool = False,  # noqa C901
                     metadatas: Optional[str] = None, user_path: Optional[str] = None):
        """
An endpoint to upload files and update their status. Multiple files can be uploaded at once.

Args:
    files (List[UploadFile]): List of files to upload.
    override (bool): Whether to overwrite existing files. Default is False.
    metadatas (Optional[str]): Metadata for the files in JSON format.
    user_path (Optional[str]): User-defined path for file uploads.

**Returns:**

- BaseResponse: Upload results and file IDs.
"""
        try:
            if user_path: user_path = user_path.lstrip('/')
            if metadatas:
                metadatas: Optional[List[Dict[str, str]]] = json.loads(metadatas)
                if len(files) != len(metadatas):
                    return BaseResponse(code=400, msg='Length of files and metadatas should be the same',
                                        data=None)
                for idx, mt in enumerate(metadatas):
                    err_msg = self._validate_metadata(mt)
                    if err_msg:
                        return BaseResponse(code=400, msg=f'file [{files[idx].filename}]: {err_msg}', data=None)
            file_paths = [os.path.join(self._manager._path, user_path or '', file.filename) for file in files]
            paths_is_new = [True] * len(file_paths)
            if override is True:
                is_success, msg, paths_is_new = self._manager.validate_paths(file_paths)
                if not is_success:
                    return BaseResponse(code=500, msg=msg, data=None)
            directorys = set(os.path.dirname(path) for path in file_paths)
            [os.makedirs(directory, exist_ok=True) for directory in directorys if directory]
            ids, results = [], []
            for i in range(len(files)):
                file_path = file_paths[i]
                content = files[i].file.read()
                metadata = metadatas[i] if metadatas else None
                if override is False:
                    file_path = self._gen_unique_filepath(file_path)
                with open(file_path, 'wb') as f: f.write(content)
                msg = "success"
                doc_id = gen_docid(file_path)
                if paths_is_new[i]:
                    docs = self._manager.add_files(
                        [file_path], metadatas=[metadata], status=DocListManager.Status.success)
                    if not docs:
                        msg = f"Failed: path {file_path} already exists in Database."
                else:
                    self._manager.update_kb_group(cond_file_ids=[doc_id], new_need_reparse=True)
                    msg = f"Success: path {file_path} will be reparsed."
                ids.append(doc_id)
                results.append(msg)
            return BaseResponse(data=[ids, results])
        except Exception as e:
            lazyllm.LOG.error(f'upload_files exception: {e}')
            return BaseResponse(code=500, msg=str(e), data=None)

    class AddFilesRequest(BaseModel):
        files: List[str]
        group_name: Optional[str] = None
        metadatas: Optional[str] = None

    @app.post("/add_files")
    def add_files(self, request: AddFilesRequest):
        """
Batch add files.

Args:
    files (List[UploadFile]): List of uploaded files.
    group_name (str): Target knowledge base group name; if empty, files are not added to any group.
    metadatas (Optional[str]): Metadata of the files in JSON format.

**Returns:**

- BaseResponse: Returns a list of unique file IDs corresponding to all input files, including newly added and existing ones. In case of exceptions, returns error codes and exception information.
"""
        files = request.files
        group_name = request.group_name
        metadatas = request.metadatas
        try:
            if metadatas:
                metadatas: Optional[List[Dict[str, str]]] = json.loads(metadatas)
                assert len(files) == len(metadatas), 'Length of files and metadatas should be the same'

            exists_files_info = self._manager.list_files(limit=None, details=True, status=DocListManager.Status.all)
            exists_files_info = {row[2]: row[0] for row in exists_files_info}

            exist_ids = []
            new_files = []
            new_metadatas = []
            id_mapping = {}

            for idx, file in enumerate(files):
                if os.path.exists(file):
                    exist_id = exists_files_info.get(file, None)
                    if exist_id:
                        update_kws = dict(fileid=exist_id, status=DocListManager.Status.success)
                        if metadatas: update_kws["meta"] = json.dumps(metadatas[idx])
                        self._manager.update_file_message(**update_kws)
                        exist_ids.append(exist_id)
                        id_mapping[file] = exist_id
                    else:
                        new_files.append(file)
                        if metadatas:
                            new_metadatas.append(metadatas[idx])
                else:
                    id_mapping[file] = None

            new_ids = self._manager.add_files(new_files, metadatas=new_metadatas, status=DocListManager.Status.success)
            if group_name:
                self._manager.add_files_to_kb_group(new_ids + exist_ids, group=group_name)

            for file, new_id in zip(new_files, new_ids):
                id_mapping[file] = new_id
            return_ids = [id_mapping[file] for file in files]

            return BaseResponse(data=return_ids)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    @app.get("/list_files")
    def list_files(self, limit: Optional[int] = None, details: bool = True, alive: Optional[bool] = None):
        """
An endpoint to list uploaded files.

Args:
    limit (Optional[int]): Limit on the number of files returned. Default is None.
    details (bool): Whether to return detailed information. Default is True.
    alive (Optional[bool]): If True, only returns non-deleted files. Default is None.

**Returns:**

- BaseResponse: File list data.
"""
        try:
            status = [DocListManager.Status.success, DocListManager.Status.waiting, DocListManager.Status.working,
                      DocListManager.Status.failed] if alive else DocListManager.Status.all
            return BaseResponse(data=self._manager.list_files(limit=limit, details=details, status=status))
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    @app.get("/reparse_files")
    def reparse_files(self, file_ids: List[str], group_name: Optional[str] = None):
        try:
            self._manager.update_need_reparsing(file_ids, group_name)
            return BaseResponse()
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    @app.get("/list_files_in_group")
    def list_files_in_group(self, group_name: Optional[str] = None,
                            limit: Optional[int] = None, alive: Optional[bool] = None):
        """
An endpoint to list files in a specific group.

Args:
    group_name (Optional[str]): The name of the file group.
    limit (Optional[int]): Limit on the number of files returned. Default is None.
    alive (Optional[bool]): Whether to return only non-deleted files.

**Returns:**

- BaseResponse: List of files in the group.
"""
        try:
            status = [DocListManager.Status.success, DocListManager.Status.waiting, DocListManager.Status.working,
                      DocListManager.Status.failed] if alive else DocListManager.Status.all
            return BaseResponse(data=self._manager.list_kb_group_files(group_name, limit, details=True, status=status))
        except Exception as e:
            return BaseResponse(code=500, msg=str(e) + '\ntraceback:\n' + str(traceback.format_exc()), data=None)

    class FileGroupRequest(BaseModel):
        file_ids: List[str]
        group_name: Optional[str] = Field(None)

    @app.post("/add_files_to_group_by_id")
    def add_files_to_group_by_id(self, request: FileGroupRequest):
        """
An endpoint to add files to a specific group by file IDs.

Args:
    request (FileGroupRequest): Request containing file IDs and group name.

**Returns:**

- BaseResponse: Operation result.
"""
        try:
            self._manager.add_files_to_kb_group(request.file_ids, request.group_name)
            return BaseResponse()
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    @app.post("/add_files_to_group")
    def add_files_to_group(self, files: List[UploadFile], group_name: str, override: bool = False,
                           metadatas: Optional[str] = None, user_path: Optional[str] = None):
        """
An endpoint to upload files and directly add them to a specified group.

Args:
    files (List[UploadFile]): List of files to upload.
    group_name (str): Name of the group to add the files to.
    override (bool): Whether to overwrite existing files. Default is False.
    metadatas (Optional[str]): Metadata for the files in JSON format.
    user_path (Optional[str]): User-defined path for file uploads.

**Returns:**

- BaseResponse: Operation result and file IDs.
"""
        try:
            response = self.upload_files(files, override=override, metadatas=metadatas, user_path=user_path)
            if response.code != 200: return response
            ids = response.data[0]
            self._manager.add_files_to_kb_group(ids, group_name)
            return BaseResponse(data=ids)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    @app.post("/delete_files")
    def delete_files(self, request: FileGroupRequest):
        """
An endpoint to delete specified files.

Args:
    request (FileGroupRequest): Request containing file IDs and group name.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        try:
            if request.group_name:
                return self.delete_files_from_group(request)
            else:
                documents = self._manager.delete_files(request.file_ids)
                deleted_ids = set([ele.doc_id for ele in documents])
                for doc in documents:
                    if os.path.exists(path := doc.path):
                        os.remove(path)
                results = ["Success" if ele.doc_id in deleted_ids else "Failed" for ele in documents]
                return BaseResponse(data=[request.file_ids, results])
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    @app.post("/delete_files_from_group")
    def delete_files_from_group(self, request: FileGroupRequest):
        """
An endpoint to delete specified files in a group.

Args:
    request (FileGroupRequest): Request containing a list of file IDs and the group name.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        try:
            self._manager.update_kb_group(cond_file_ids=request.file_ids, cond_group=request.group_name,
                                          new_status=DocListManager.Status.deleting)
            return BaseResponse()
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    class AddMetadataRequest(BaseModel):
        doc_ids: List[str]
        kv_pair: Dict[str, Union[bool, int, float, str, list]]

    @app.post("/add_metadata")
    def add_metadata(self, add_metadata_request: AddMetadataRequest):
        """
An endpoint to add or update metadata for specified documents.

Args:
    add_metadata_request (AddMetadataRequest): Request containing list of document IDs and key-value metadata.

**Returns:**

- BaseResponse: Operation result information.
"""
        doc_ids = add_metadata_request.doc_ids
        kv_pair = add_metadata_request.kv_pair
        try:
            docs = self._manager.get_docs(doc_ids)
            if not docs:
                return BaseResponse(code=400, msg="Failed, no doc found")
            doc_meta = {}
            for doc in docs:
                meta_dict = json.loads(doc.meta) if doc.meta else {}
                for k, v in kv_pair.items():
                    if k not in meta_dict or not meta_dict[k]:
                        meta_dict[k] = v
                    elif isinstance(meta_dict[k], list):
                        meta_dict[k].extend(v) if isinstance(v, list) else meta_dict[k].append(v)
                    else:
                        meta_dict[k] = ([meta_dict[k]] + v) if isinstance(v, list) else [meta_dict[k], v]
                doc_meta[doc.doc_id] = meta_dict
            self._manager.set_docs_new_meta(doc_meta)
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    class DeleteMetadataRequest(BaseModel):
        doc_ids: List[str]
        keys: Optional[List[str]] = Field(None)
        kv_pair: Optional[Dict[str, Union[bool, int, float, str, list]]] = Field(None)

    def _inplace_del_meta(self, meta_dict, kv_pair: Dict[str, Union[None, bool, int, float, str, list]]):
        # alert: meta_dict is not a deepcopy
        for k, v in kv_pair.items():
            if k not in meta_dict:
                continue
            if v is None:
                meta_dict.pop(k, None)
            elif isinstance(meta_dict[k], list):
                if isinstance(v, (bool, int, float, str)):
                    v = [v]
                # delete v exists in meta_dict[k]
                meta_dict[k] = list(set(meta_dict[k]) - set(v))
            else:
                # old meta[k] not a list, use v as condition to delete the key
                if meta_dict[k] == v:
                    meta_dict.pop(k, None)

    @app.post("/delete_metadata_item")
    def delete_metadata_item(self, del_metadata_request: DeleteMetadataRequest):
        """
An endpoint to delete metadata fields or field values from specified documents.

Args:
    del_metadata_request (DeleteMetadataRequest): Request containing list of document IDs, field names, and/or deletion rules.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        doc_ids = del_metadata_request.doc_ids
        kv_pair = del_metadata_request.kv_pair
        keys = del_metadata_request.keys
        try:
            if keys is not None:
                # convert keys to kv_pair
                if kv_pair:
                    kv_pair.update({k: None for k in keys})
                else:
                    kv_pair = {k: None for k in keys}
            if not kv_pair:
                # clear metadata
                self._manager.set_docs_new_meta({doc_id: {} for doc_id in doc_ids})
            else:
                docs = self._manager.get_docs(doc_ids)
                if not docs:
                    return BaseResponse(code=400, msg="Failed, no doc found")
                doc_meta = {}
                for doc in docs:
                    meta_dict = json.loads(doc.meta) if doc.meta else {}
                    self._inplace_del_meta(meta_dict, kv_pair)
                    doc_meta[doc.doc_id] = meta_dict
                self._manager.set_docs_new_meta(doc_meta)
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    class UpdateMetadataRequest(BaseModel):
        doc_ids: List[str]
        kv_pair: Dict[str, Union[bool, int, float, str, list]]

    @app.post("/update_or_create_metadata_keys")
    def update_or_create_metadata_keys(self, update_metadata_request: UpdateMetadataRequest):
        """
An endpoint to update or create metadata fields for specified documents.

Args:
    update_metadata_request (UpdateMetadataRequest): Request containing a list of document IDs and metadata key-value pairs to update or create.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        doc_ids = update_metadata_request.doc_ids
        kv_pair = update_metadata_request.kv_pair
        try:
            docs = self._manager.get_docs(doc_ids)
            if not docs:
                return BaseResponse(code=400, msg="Failed, no doc found")
            for doc in docs:
                doc_meta = {}
                meta_dict = json.loads(doc.meta) if doc.meta else {}
                for k, v in kv_pair.items():
                    meta_dict[k] = v
                doc_meta[doc.doc_id] = meta_dict
            self._manager.set_docs_new_meta(doc_meta)
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    class ResetMetadataRequest(BaseModel):
        doc_ids: List[str]
        new_meta: Dict[str, Union[bool, int, float, str, list]]

    @app.post("/reset_metadata")
    def reset_metadata(self, reset_metadata_request: ResetMetadataRequest):
        """
An endpoint to reset all metadata fields of specified documents.

Args:
    reset_metadata_request (ResetMetadataRequest): Request containing a list of document IDs and the new metadata dictionary to apply.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        doc_ids = reset_metadata_request.doc_ids
        new_meta = reset_metadata_request.new_meta
        try:
            docs = self._manager.get_docs(doc_ids)
            if not docs:
                return BaseResponse(code=400, msg="Failed, no doc found")
            self._manager.set_docs_new_meta({doc.doc_id: new_meta for doc in docs})
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    class QueryMetadataRequest(BaseModel):
        doc_id: str
        key: Optional[str] = None

    @app.post("/query_metadata")
    def query_metadata(self, query_metadata_request: QueryMetadataRequest):
        """
An endpoint to query metadata of a specific document.

Args:
    query_metadata_request (QueryMetadataRequest): Request containing the document ID and an optional metadata field name.

**Returns:**

- BaseResponse: Returns the field value if key is specified and exists; otherwise returns full metadata. If the key does not exist, returns an error.
"""
        doc_id = query_metadata_request.doc_id
        key = query_metadata_request.key
        try:
            docs = self._manager.get_docs(doc_id)
            if not docs:
                return BaseResponse(data=None)
            doc = docs[0]
            meta_dict = json.loads(doc.meta) if doc.meta else {}
            if not key:
                return BaseResponse(data=meta_dict)
            if key not in meta_dict:
                return BaseResponse(code=400, msg=f"Failed, key {key} does not exist")
            return BaseResponse(data=meta_dict[key])
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

    def __repr__(self):
        return lazyllm.make_repr("Module", "DocManager")

add_files(request)

Batch add files.

Parameters:

  • files (List[UploadFile]) –

    List of uploaded files.

  • group_name (str) –

    Target knowledge base group name; if empty, files are not added to any group.

  • metadatas (Optional[str]) –

    Metadata of the files in JSON format.

Returns:

  • BaseResponse: Returns a list of unique file IDs corresponding to all input files, including newly added and existing ones. In case of exceptions, returns error codes and exception information.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/add_files")
    def add_files(self, request: AddFilesRequest):
        """
Batch add files.

Args:
    files (List[UploadFile]): List of uploaded files.
    group_name (str): Target knowledge base group name; if empty, files are not added to any group.
    metadatas (Optional[str]): Metadata of the files in JSON format.

**Returns:**

- BaseResponse: Returns a list of unique file IDs corresponding to all input files, including newly added and existing ones. In case of exceptions, returns error codes and exception information.
"""
        files = request.files
        group_name = request.group_name
        metadatas = request.metadatas
        try:
            if metadatas:
                metadatas: Optional[List[Dict[str, str]]] = json.loads(metadatas)
                assert len(files) == len(metadatas), 'Length of files and metadatas should be the same'

            exists_files_info = self._manager.list_files(limit=None, details=True, status=DocListManager.Status.all)
            exists_files_info = {row[2]: row[0] for row in exists_files_info}

            exist_ids = []
            new_files = []
            new_metadatas = []
            id_mapping = {}

            for idx, file in enumerate(files):
                if os.path.exists(file):
                    exist_id = exists_files_info.get(file, None)
                    if exist_id:
                        update_kws = dict(fileid=exist_id, status=DocListManager.Status.success)
                        if metadatas: update_kws["meta"] = json.dumps(metadatas[idx])
                        self._manager.update_file_message(**update_kws)
                        exist_ids.append(exist_id)
                        id_mapping[file] = exist_id
                    else:
                        new_files.append(file)
                        if metadatas:
                            new_metadatas.append(metadatas[idx])
                else:
                    id_mapping[file] = None

            new_ids = self._manager.add_files(new_files, metadatas=new_metadatas, status=DocListManager.Status.success)
            if group_name:
                self._manager.add_files_to_kb_group(new_ids + exist_ids, group=group_name)

            for file, new_id in zip(new_files, new_ids):
                id_mapping[file] = new_id
            return_ids = [id_mapping[file] for file in files]

            return BaseResponse(data=return_ids)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

add_files_to_group(files, group_name, override=False, metadatas=None, user_path=None)

An endpoint to upload files and directly add them to a specified group.

Parameters:

  • files (List[UploadFile]) –

    List of files to upload.

  • group_name (str) –

    Name of the group to add the files to.

  • override (bool, default: False ) –

    Whether to overwrite existing files. Default is False.

  • metadatas (Optional[str], default: None ) –

    Metadata for the files in JSON format.

  • user_path (Optional[str], default: None ) –

    User-defined path for file uploads.

Returns:

  • BaseResponse: Operation result and file IDs.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/add_files_to_group")
    def add_files_to_group(self, files: List[UploadFile], group_name: str, override: bool = False,
                           metadatas: Optional[str] = None, user_path: Optional[str] = None):
        """
An endpoint to upload files and directly add them to a specified group.

Args:
    files (List[UploadFile]): List of files to upload.
    group_name (str): Name of the group to add the files to.
    override (bool): Whether to overwrite existing files. Default is False.
    metadatas (Optional[str]): Metadata for the files in JSON format.
    user_path (Optional[str]): User-defined path for file uploads.

**Returns:**

- BaseResponse: Operation result and file IDs.
"""
        try:
            response = self.upload_files(files, override=override, metadatas=metadatas, user_path=user_path)
            if response.code != 200: return response
            ids = response.data[0]
            self._manager.add_files_to_kb_group(ids, group_name)
            return BaseResponse(data=ids)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

add_files_to_group_by_id(request)

An endpoint to add files to a specific group by file IDs.

Parameters:

  • request (FileGroupRequest) –

    Request containing file IDs and group name.

Returns:

  • BaseResponse: Operation result.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/add_files_to_group_by_id")
    def add_files_to_group_by_id(self, request: FileGroupRequest):
        """
An endpoint to add files to a specific group by file IDs.

Args:
    request (FileGroupRequest): Request containing file IDs and group name.

**Returns:**

- BaseResponse: Operation result.
"""
        try:
            self._manager.add_files_to_kb_group(request.file_ids, request.group_name)
            return BaseResponse()
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

add_metadata(add_metadata_request)

An endpoint to add or update metadata for specified documents.

Parameters:

  • add_metadata_request (AddMetadataRequest) –

    Request containing list of document IDs and key-value metadata.

Returns:

  • BaseResponse: Operation result information.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/add_metadata")
    def add_metadata(self, add_metadata_request: AddMetadataRequest):
        """
An endpoint to add or update metadata for specified documents.

Args:
    add_metadata_request (AddMetadataRequest): Request containing list of document IDs and key-value metadata.

**Returns:**

- BaseResponse: Operation result information.
"""
        doc_ids = add_metadata_request.doc_ids
        kv_pair = add_metadata_request.kv_pair
        try:
            docs = self._manager.get_docs(doc_ids)
            if not docs:
                return BaseResponse(code=400, msg="Failed, no doc found")
            doc_meta = {}
            for doc in docs:
                meta_dict = json.loads(doc.meta) if doc.meta else {}
                for k, v in kv_pair.items():
                    if k not in meta_dict or not meta_dict[k]:
                        meta_dict[k] = v
                    elif isinstance(meta_dict[k], list):
                        meta_dict[k].extend(v) if isinstance(v, list) else meta_dict[k].append(v)
                    else:
                        meta_dict[k] = ([meta_dict[k]] + v) if isinstance(v, list) else [meta_dict[k], v]
                doc_meta[doc.doc_id] = meta_dict
            self._manager.set_docs_new_meta(doc_meta)
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

delete_files(request)

An endpoint to delete specified files.

Parameters:

  • request (FileGroupRequest) –

    Request containing file IDs and group name.

Returns:

  • BaseResponse: Deletion operation result.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/delete_files")
    def delete_files(self, request: FileGroupRequest):
        """
An endpoint to delete specified files.

Args:
    request (FileGroupRequest): Request containing file IDs and group name.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        try:
            if request.group_name:
                return self.delete_files_from_group(request)
            else:
                documents = self._manager.delete_files(request.file_ids)
                deleted_ids = set([ele.doc_id for ele in documents])
                for doc in documents:
                    if os.path.exists(path := doc.path):
                        os.remove(path)
                results = ["Success" if ele.doc_id in deleted_ids else "Failed" for ele in documents]
                return BaseResponse(data=[request.file_ids, results])
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

delete_files_from_group(request)

An endpoint to delete specified files in a group.

Parameters:

  • request (FileGroupRequest) –

    Request containing a list of file IDs and the group name.

Returns:

  • BaseResponse: Deletion operation result.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/delete_files_from_group")
    def delete_files_from_group(self, request: FileGroupRequest):
        """
An endpoint to delete specified files in a group.

Args:
    request (FileGroupRequest): Request containing a list of file IDs and the group name.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        try:
            self._manager.update_kb_group(cond_file_ids=request.file_ids, cond_group=request.group_name,
                                          new_status=DocListManager.Status.deleting)
            return BaseResponse()
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

delete_metadata_item(del_metadata_request)

An endpoint to delete metadata fields or field values from specified documents.

Parameters:

  • del_metadata_request (DeleteMetadataRequest) –

    Request containing list of document IDs, field names, and/or deletion rules.

Returns:

  • BaseResponse: Deletion operation result.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/delete_metadata_item")
    def delete_metadata_item(self, del_metadata_request: DeleteMetadataRequest):
        """
An endpoint to delete metadata fields or field values from specified documents.

Args:
    del_metadata_request (DeleteMetadataRequest): Request containing list of document IDs, field names, and/or deletion rules.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        doc_ids = del_metadata_request.doc_ids
        kv_pair = del_metadata_request.kv_pair
        keys = del_metadata_request.keys
        try:
            if keys is not None:
                # convert keys to kv_pair
                if kv_pair:
                    kv_pair.update({k: None for k in keys})
                else:
                    kv_pair = {k: None for k in keys}
            if not kv_pair:
                # clear metadata
                self._manager.set_docs_new_meta({doc_id: {} for doc_id in doc_ids})
            else:
                docs = self._manager.get_docs(doc_ids)
                if not docs:
                    return BaseResponse(code=400, msg="Failed, no doc found")
                doc_meta = {}
                for doc in docs:
                    meta_dict = json.loads(doc.meta) if doc.meta else {}
                    self._inplace_del_meta(meta_dict, kv_pair)
                    doc_meta[doc.doc_id] = meta_dict
                self._manager.set_docs_new_meta(doc_meta)
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

document()

An endpoint to redirect to the default documentation page.

Returns:

  • RedirectResponse: Redirects to the /docs page.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.get("/", response_model=BaseResponse, summary="docs")
    def document(self):
        """
An endpoint to redirect to the default documentation page.

**Returns:**

- RedirectResponse: Redirects to the `/docs` page.
"""
        return RedirectResponse(url="/docs")

list_files(limit=None, details=True, alive=None)

An endpoint to list uploaded files.

Parameters:

  • limit (Optional[int], default: None ) –

    Limit on the number of files returned. Default is None.

  • details (bool, default: True ) –

    Whether to return detailed information. Default is True.

  • alive (Optional[bool], default: None ) –

    If True, only returns non-deleted files. Default is None.

Returns:

  • BaseResponse: File list data.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.get("/list_files")
    def list_files(self, limit: Optional[int] = None, details: bool = True, alive: Optional[bool] = None):
        """
An endpoint to list uploaded files.

Args:
    limit (Optional[int]): Limit on the number of files returned. Default is None.
    details (bool): Whether to return detailed information. Default is True.
    alive (Optional[bool]): If True, only returns non-deleted files. Default is None.

**Returns:**

- BaseResponse: File list data.
"""
        try:
            status = [DocListManager.Status.success, DocListManager.Status.waiting, DocListManager.Status.working,
                      DocListManager.Status.failed] if alive else DocListManager.Status.all
            return BaseResponse(data=self._manager.list_files(limit=limit, details=details, status=status))
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

list_files_in_group(group_name=None, limit=None, alive=None)

An endpoint to list files in a specific group.

Parameters:

  • group_name (Optional[str], default: None ) –

    The name of the file group.

  • limit (Optional[int], default: None ) –

    Limit on the number of files returned. Default is None.

  • alive (Optional[bool], default: None ) –

    Whether to return only non-deleted files.

Returns:

  • BaseResponse: List of files in the group.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.get("/list_files_in_group")
    def list_files_in_group(self, group_name: Optional[str] = None,
                            limit: Optional[int] = None, alive: Optional[bool] = None):
        """
An endpoint to list files in a specific group.

Args:
    group_name (Optional[str]): The name of the file group.
    limit (Optional[int]): Limit on the number of files returned. Default is None.
    alive (Optional[bool]): Whether to return only non-deleted files.

**Returns:**

- BaseResponse: List of files in the group.
"""
        try:
            status = [DocListManager.Status.success, DocListManager.Status.waiting, DocListManager.Status.working,
                      DocListManager.Status.failed] if alive else DocListManager.Status.all
            return BaseResponse(data=self._manager.list_kb_group_files(group_name, limit, details=True, status=status))
        except Exception as e:
            return BaseResponse(code=500, msg=str(e) + '\ntraceback:\n' + str(traceback.format_exc()), data=None)

list_kb_groups()

An endpoint to list all document groups.

Returns:

  • BaseResponse: Contains the data of all document groups.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.get("/list_kb_groups")
    def list_kb_groups(self):
        """
An endpoint to list all document groups.

**Returns:**

- BaseResponse: Contains the data of all document groups.
"""
        try:
            return BaseResponse(data=self._manager.list_all_kb_group())
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

query_metadata(query_metadata_request)

An endpoint to query metadata of a specific document.

Parameters:

  • query_metadata_request (QueryMetadataRequest) –

    Request containing the document ID and an optional metadata field name.

Returns:

  • BaseResponse: Returns the field value if key is specified and exists; otherwise returns full metadata. If the key does not exist, returns an error.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/query_metadata")
    def query_metadata(self, query_metadata_request: QueryMetadataRequest):
        """
An endpoint to query metadata of a specific document.

Args:
    query_metadata_request (QueryMetadataRequest): Request containing the document ID and an optional metadata field name.

**Returns:**

- BaseResponse: Returns the field value if key is specified and exists; otherwise returns full metadata. If the key does not exist, returns an error.
"""
        doc_id = query_metadata_request.doc_id
        key = query_metadata_request.key
        try:
            docs = self._manager.get_docs(doc_id)
            if not docs:
                return BaseResponse(data=None)
            doc = docs[0]
            meta_dict = json.loads(doc.meta) if doc.meta else {}
            if not key:
                return BaseResponse(data=meta_dict)
            if key not in meta_dict:
                return BaseResponse(code=400, msg=f"Failed, key {key} does not exist")
            return BaseResponse(data=meta_dict[key])
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

reset_metadata(reset_metadata_request)

An endpoint to reset all metadata fields of specified documents.

Parameters:

  • reset_metadata_request (ResetMetadataRequest) –

    Request containing a list of document IDs and the new metadata dictionary to apply.

Returns:

  • BaseResponse: Deletion operation result.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/reset_metadata")
    def reset_metadata(self, reset_metadata_request: ResetMetadataRequest):
        """
An endpoint to reset all metadata fields of specified documents.

Args:
    reset_metadata_request (ResetMetadataRequest): Request containing a list of document IDs and the new metadata dictionary to apply.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        doc_ids = reset_metadata_request.doc_ids
        new_meta = reset_metadata_request.new_meta
        try:
            docs = self._manager.get_docs(doc_ids)
            if not docs:
                return BaseResponse(code=400, msg="Failed, no doc found")
            self._manager.set_docs_new_meta({doc.doc_id: new_meta for doc in docs})
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

update_or_create_metadata_keys(update_metadata_request)

An endpoint to update or create metadata fields for specified documents.

Parameters:

  • update_metadata_request (UpdateMetadataRequest) –

    Request containing a list of document IDs and metadata key-value pairs to update or create.

Returns:

  • BaseResponse: Deletion operation result.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/update_or_create_metadata_keys")
    def update_or_create_metadata_keys(self, update_metadata_request: UpdateMetadataRequest):
        """
An endpoint to update or create metadata fields for specified documents.

Args:
    update_metadata_request (UpdateMetadataRequest): Request containing a list of document IDs and metadata key-value pairs to update or create.

**Returns:**

- BaseResponse: Deletion operation result.
"""
        doc_ids = update_metadata_request.doc_ids
        kv_pair = update_metadata_request.kv_pair
        try:
            docs = self._manager.get_docs(doc_ids)
            if not docs:
                return BaseResponse(code=400, msg="Failed, no doc found")
            for doc in docs:
                doc_meta = {}
                meta_dict = json.loads(doc.meta) if doc.meta else {}
                for k, v in kv_pair.items():
                    meta_dict[k] = v
                doc_meta[doc.doc_id] = meta_dict
            self._manager.set_docs_new_meta(doc_meta)
            return BaseResponse(data=None)
        except Exception as e:
            return BaseResponse(code=500, msg=str(e), data=None)

upload_files(files, override=False, metadatas=None, user_path=None)

An endpoint to upload files and update their status. Multiple files can be uploaded at once.

Parameters:

  • files (List[UploadFile]) –

    List of files to upload.

  • override (bool, default: False ) –

    Whether to overwrite existing files. Default is False.

  • metadatas (Optional[str], default: None ) –

    Metadata for the files in JSON format.

  • user_path (Optional[str], default: None ) –

    User-defined path for file uploads.

Returns:

  • BaseResponse: Upload results and file IDs.
Source code in lazyllm/tools/rag/doc_manager.py
    @app.post("/upload_files")
    def upload_files(self, files: List[UploadFile], override: bool = False,  # noqa C901
                     metadatas: Optional[str] = None, user_path: Optional[str] = None):
        """
An endpoint to upload files and update their status. Multiple files can be uploaded at once.

Args:
    files (List[UploadFile]): List of files to upload.
    override (bool): Whether to overwrite existing files. Default is False.
    metadatas (Optional[str]): Metadata for the files in JSON format.
    user_path (Optional[str]): User-defined path for file uploads.

**Returns:**

- BaseResponse: Upload results and file IDs.
"""
        try:
            if user_path: user_path = user_path.lstrip('/')
            if metadatas:
                metadatas: Optional[List[Dict[str, str]]] = json.loads(metadatas)
                if len(files) != len(metadatas):
                    return BaseResponse(code=400, msg='Length of files and metadatas should be the same',
                                        data=None)
                for idx, mt in enumerate(metadatas):
                    err_msg = self._validate_metadata(mt)
                    if err_msg:
                        return BaseResponse(code=400, msg=f'file [{files[idx].filename}]: {err_msg}', data=None)
            file_paths = [os.path.join(self._manager._path, user_path or '', file.filename) for file in files]
            paths_is_new = [True] * len(file_paths)
            if override is True:
                is_success, msg, paths_is_new = self._manager.validate_paths(file_paths)
                if not is_success:
                    return BaseResponse(code=500, msg=msg, data=None)
            directorys = set(os.path.dirname(path) for path in file_paths)
            [os.makedirs(directory, exist_ok=True) for directory in directorys if directory]
            ids, results = [], []
            for i in range(len(files)):
                file_path = file_paths[i]
                content = files[i].file.read()
                metadata = metadatas[i] if metadatas else None
                if override is False:
                    file_path = self._gen_unique_filepath(file_path)
                with open(file_path, 'wb') as f: f.write(content)
                msg = "success"
                doc_id = gen_docid(file_path)
                if paths_is_new[i]:
                    docs = self._manager.add_files(
                        [file_path], metadatas=[metadata], status=DocListManager.Status.success)
                    if not docs:
                        msg = f"Failed: path {file_path} already exists in Database."
                else:
                    self._manager.update_kb_group(cond_file_ids=[doc_id], new_need_reparse=True)
                    msg = f"Success: path {file_path} will be reparsed."
                ids.append(doc_id)
                results.append(msg)
            return BaseResponse(data=[ids, results])
        except Exception as e:
            lazyllm.LOG.error(f'upload_files exception: {e}')
            return BaseResponse(code=500, msg=str(e), data=None)

lazyllm.tools.rag.utils.SqliteDocListManager

Bases: DocListManager

SQLite-based document manager for persistent local file storage, status tracking, and metadata management.

This class inherits from DocListManager and uses a SQLite backend to store document records. It is suitable for managing locally identified documents with support for inserting, querying, updating, and filtering based on status. Optional file path monitoring is also supported.

Parameters:

  • path (str) –

    Directory path to store the database.

  • name (str) –

    Name of the SQLite database file (without path).

  • enable_path_monitoring (bool, default: True ) –

    Whether to enable path monitoring. Defaults to True.

Examples:

>>> from lazyllm.tools.rag.utils import SqliteDocListManager
>>> manager = SqliteDocListManager(path="./data", name="docs.sqlite")
>>> manager.insert({"uid": "doc_001", "name": "example.txt", "status": "ready"})
>>> print(manager.get("doc_001"))
>>> files = manager.list_files(limit=5, details=True)
>>> print(files)
Source code in lazyllm/tools/rag/utils.py
 589
 590
 591
 592
 593
 594
 595
 596
 597
 598
 599
 600
 601
 602
 603
 604
 605
 606
 607
 608
 609
 610
 611
 612
 613
 614
 615
 616
 617
 618
 619
 620
 621
 622
 623
 624
 625
 626
 627
 628
 629
 630
 631
 632
 633
 634
 635
 636
 637
 638
 639
 640
 641
 642
 643
 644
 645
 646
 647
 648
 649
 650
 651
 652
 653
 654
 655
 656
 657
 658
 659
 660
 661
 662
 663
 664
 665
 666
 667
 668
 669
 670
 671
 672
 673
 674
 675
 676
 677
 678
 679
 680
 681
 682
 683
 684
 685
 686
 687
 688
 689
 690
 691
 692
 693
 694
 695
 696
 697
 698
 699
 700
 701
 702
 703
 704
 705
 706
 707
 708
 709
 710
 711
 712
 713
 714
 715
 716
 717
 718
 719
 720
 721
 722
 723
 724
 725
 726
 727
 728
 729
 730
 731
 732
 733
 734
 735
 736
 737
 738
 739
 740
 741
 742
 743
 744
 745
 746
 747
 748
 749
 750
 751
 752
 753
 754
 755
 756
 757
 758
 759
 760
 761
 762
 763
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
class SqliteDocListManager(DocListManager):
    """SQLite-based document manager for persistent local file storage, status tracking, and metadata management.

This class inherits from DocListManager and uses a SQLite backend to store document records. It is suitable for managing locally identified documents with support for inserting, querying, updating, and filtering based on status. Optional file path monitoring is also supported.

Args:
    path (str): Directory path to store the database.
    name (str): Name of the SQLite database file (without path).
    enable_path_monitoring (bool): Whether to enable path monitoring. Defaults to True.


Examples:
    >>> from lazyllm.tools.rag.utils import SqliteDocListManager
    >>> manager = SqliteDocListManager(path="./data", name="docs.sqlite")
    >>> manager.insert({"uid": "doc_001", "name": "example.txt", "status": "ready"})
    >>> print(manager.get("doc_001"))
    >>> files = manager.list_files(limit=5, details=True)
    >>> print(files)
    """
    def __init__(self, path, name, enable_path_monitoring=True):
        super().__init__(path, name, enable_path_monitoring)

    def _init_sql(self):
        root_dir = os.path.expanduser(os.path.join(config['home'], '.dbs'))
        os.makedirs(root_dir, exist_ok=True)
        self._db_path = os.path.join(root_dir, f'.lazyllm_dlmanager.{self._id}.db')
        self._db_lock = FileLock(self._db_path + '.lock')
        # ensure that this connection is not used in another thread when sqlite3 is not threadsafe
        self._check_same_thread = not sqlite3_check_threadsafety()
        self._engine = sqlalchemy.create_engine(
            f"sqlite:///{self._db_path}?check_same_thread={self._check_same_thread}"
        )
        self._Session = sessionmaker(bind=self._engine)
        self.init_tables()

    def _init_tables(self):
        KBDataBase.metadata.create_all(bind=self._engine)

    def table_inited(self):
        """Checks whether the "documents" table has been initialized in the database.

The method queries the sqlite_master metadata table to verify if the "documents" table exists.

**Returns:**

- bool: True if the "documents" table exists, False otherwise.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='documents'")
            return cursor.fetchone() is not None

    @staticmethod
    def get_status_cond_and_params(status: Union[str, List[str]],
                                   exclude_status: Optional[Union[str, List[str]]] = None,
                                   prefix: str = None):
        conds, params = [], []
        prefix = f'{prefix}.' if prefix else ''
        if isinstance(status, str):
            if status != DocListManager.Status.all:
                conds.append(f'{prefix}status = ?')
                params.append(status)
        elif isinstance(status, (tuple, list)):
            conds.append(f'{prefix}status IN ({",".join("?" * len(status))})')
            params.extend(status)

        if isinstance(exclude_status, str):
            assert exclude_status != DocListManager.Status.all, 'Invalid status provided'
            conds.append(f'{prefix}status != ?')
            params.append(exclude_status)
        elif isinstance(exclude_status, (tuple, list)):
            conds.append(f'{prefix}status NOT IN ({",".join("?" * len(exclude_status))})')
            params.extend(exclude_status)

        return ' AND '.join(conds), params

    def _get_all_docs(self):
        with self._db_lock, self._Session() as session:
            return session.query(KBDocument).all()

    def _get_docs(self, to_be_added_doc_ids: List, to_be_deleted_doc_ids: List, filter_status_list: List):
        with self._db_lock, self._Session() as session:
            docs_not_expected = session.query(KBDocument).filter(KBDocument.doc_id.in_(to_be_added_doc_ids)).all()
            docs_expected = session.query(KBDocument).filter(KBDocument.doc_id.in_(to_be_deleted_doc_ids),
                                                             KBDocument.status.in_(filter_status_list)).all()
        return docs_not_expected, docs_expected

    def validate_paths(self, paths: List[str]) -> Tuple[bool, str, List[bool]]:
        """Validates whether the documents corresponding to the given paths can be safely added to the database.

The method checks if the document already exists. If it exists, it verifies whether the document is currently
being parsed, waiting to be parsed, or was not successfully re-parsed last time.

Args:
    paths (List[str]): A list of file paths to validate.

**Returns:**

- Tuple[bool, str, List[bool]]: 
    - bool: Whether all paths passed validation.
    - str: Description message of the validation result.
    - List[bool]: A boolean list corresponding to input paths, indicating whether each path is new (True) or already exists (False).
      If validation fails, this value is None.
"""
        # check and return: success, msg, path_is_new for each path
        unsafe_staus_set = set([DocListManager.Status.working, DocListManager.Status.waiting])
        paths_is_new = [True] * len(paths)
        doc_ids = [gen_docid(path) for path in paths]
        doc_id_to_path = {doc_id: path for doc_id, path in zip(doc_ids, paths)}
        found_doc_ids = []
        found_doc_group_rows = []
        with self._db_lock, self._Session() as session:
            rows = session.execute(
                select(KBDocument.doc_id).where(KBDocument.doc_id.in_(doc_ids))
            ).fetchall()
            if len(rows) == 0:
                return True, "Success", paths_is_new
            found_doc_ids = [row.doc_id for row in rows]
            found_doc_group_rows = session.execute(
                select(KBGroupDocuments.doc_id, KBGroupDocuments.need_reparse, KBGroupDocuments.status)
                .where(KBGroupDocuments.doc_id.in_(found_doc_ids))).fetchall()

        for doc_group_record in found_doc_group_rows:
            if doc_group_record.need_reparse:
                msg = f"Failed: {doc_id_to_path[doc_group_record.doc_id]} lasttime reparsing has not been finished"
                return False, msg, None
            if doc_group_record.status in unsafe_staus_set:
                return False, f"Failed: {doc_id_to_path[doc_group_record.doc_id]} is being parsed by kbgroup", None
        found_doc_ids = set(found_doc_ids)
        for i in range(len(paths)):
            cur_doc_id = doc_ids[i]
            if cur_doc_id in found_doc_ids:
                paths_is_new[i] = False
        return True, "Success", paths_is_new

    def update_need_reparsing(self, doc_id: str, need_reparse: bool, group_name: Optional[str] = None):
        """Updates the re-parsing flag for a specific document.

This method sets whether a document should be re-parsed. If a group name is provided, the update is scoped to that group only.

Args:
    doc_id (str): The unique identifier of the document.
    need_reparse (bool): Whether the document needs to be re-parsed.
    group_name (Optional[str]): Optional. The knowledge base group name to filter by. If provided, only documents in the specified group will be updated.
"""
        with self._db_lock, self._Session() as session:
            stmt = update(KBGroupDocuments).where(KBGroupDocuments.doc_id == doc_id)
            if group_name is not None: stmt = stmt.where(KBGroupDocuments.group_name == group_name)
            session.execute(stmt.values(need_reparse=need_reparse))
            session.commit()

    def list_files(self, limit: Optional[int] = None, details: bool = False,
                   status: Union[str, List[str]] = DocListManager.Status.all,
                   exclude_status: Optional[Union[str, List[str]]] = None):
        """Lists files in the document database based on status filters and returns either full records or file paths.

Args:
    limit (Optional[int]): The maximum number of records to return. If None, all matching records are returned.
    details (bool): Whether to return full database rows or just file paths (document IDs).
    status (Union[str, List[str]]): Status values to include in the result. Defaults to including all.
    exclude_status (Optional[Union[str, List[str]]]): Status values to exclude from the result.

**Returns:**

- list: A list of file records or document paths depending on the `details` flag.
"""
        query = "SELECT * FROM documents"
        params = []
        status_cond, status_params = self.get_status_cond_and_params(status, exclude_status, prefix=None)
        if status_cond:
            query += f' WHERE {status_cond}'
            params.extend(status_params)
        if limit:
            query += " LIMIT ?"
            params.append(limit)
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute(query, params)
            return cursor.fetchall() if details else [row[0] for row in cursor]

    def get_docs(self, doc_ids: List[str]) -> List[KBDocument]:
        """Fetches document objects from the database corresponding to the given list of document IDs.

Args:
    doc_ids (List[str]): A list of document IDs to query.

**Returns:**

- List[KBDocument]: A list of matching document objects. Returns an empty list if no matches found.
"""
        with self._db_lock, self._Session() as session:
            docs = session.query(KBDocument).filter(KBDocument.doc_id.in_(doc_ids)).all()
            return docs
        return []

    def set_docs_new_meta(self, doc_meta: Dict[str, dict]):
        """Batch updates the metadata (meta) of documents, and simultaneously updates the new_meta field of documents in knowledge base groups for documents that are not in waiting status.

Args:
    doc_meta (Dict[str, dict]): A dictionary mapping document IDs to their new metadata dictionaries.
"""
        data_to_update = [{"_doc_id": k, "_meta": json.dumps(v)} for k, v in doc_meta.items()]
        with self._db_lock, self._Session() as session:
            # Use sqlalchemy core bulk update
            stmt = KBDocument.__table__.update().where(
                KBDocument.doc_id == bindparam("_doc_id")).values(meta=bindparam("_meta"))
            session.execute(stmt, data_to_update)
            session.commit()

            stmt = KBGroupDocuments.__table__.update().where(
                KBGroupDocuments.doc_id == bindparam("_doc_id"),
                KBGroupDocuments.status != DocListManager.Status.waiting).values(new_meta=bindparam("_meta"))
            session.execute(stmt, data_to_update)
            session.commit()

    def fetch_docs_changed_meta(self, group: str) -> List[DocMetaChangedRow]:
        """Fetches the list of documents within a specified knowledge base group that have updated metadata, and resets the new_meta field for those documents.

Args:
    group (str): Name of the knowledge base group.

**Returns:**

- List[DocMetaChangedRow]: A list containing document IDs and their updated metadata.
"""
        rows = []
        conds = [KBGroupDocuments.group_name == group, KBGroupDocuments.new_meta.isnot(None)]
        with self._db_lock, self._Session() as session:
            rows = (
                session.query(KBDocument.doc_id, KBGroupDocuments.new_meta)
                .join(KBGroupDocuments, KBDocument.doc_id == KBGroupDocuments.doc_id)
                .filter(*conds).all()
            )
            stmt = update(KBGroupDocuments).where(sqlalchemy.and_(*conds)).values(new_meta=None)
            session.execute(stmt)
            session.commit()
        return rows

    def list_all_kb_group(self):
        """Lists all knowledge base group names stored in the database.

**Returns:**

- List[str]: A list of knowledge base group names.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute("SELECT group_name FROM document_groups")
            return [row[0] for row in cursor]

    def add_kb_group(self, name):
        """Adds a new knowledge base group name to the database; ignores if the group already exists.

Args:
    name (str): The name of the knowledge base group to add.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            conn.execute('INSERT OR IGNORE INTO document_groups (group_name) VALUES (?)', (name,))
            conn.commit()

    def list_kb_group_files(self, group: str = None, limit: Optional[int] = None, details: bool = False,
                            status: Union[str, List[str]] = DocListManager.Status.all,
                            exclude_status: Optional[Union[str, List[str]]] = None,
                            upload_status: Union[str, List[str]] = DocListManager.Status.all,
                            exclude_upload_status: Optional[Union[str, List[str]]] = None,
                            need_reparse: Optional[bool] = None):
        """Lists files in a specified knowledge base group, with support for multiple filters.

Args:
    group (str, optional): Knowledge base group name to filter by. If None, no group filtering is applied.
    limit (int, optional): Limit on the number of files to return.
    details (bool): Whether to return detailed file information.
    status (str or List[str], optional): Filter files by group document status.
    exclude_status (str or List[str], optional): Exclude files with these group document statuses.
    upload_status (str or List[str], optional): Filter files by upload document status.
    exclude_upload_status (str or List[str], optional): Exclude files with these upload document statuses.
    need_reparse (bool, optional): If set, only returns files marked as needing reparse.

**Returns:**

- list: 
    - If details is False, returns a list of tuples (doc_id, path).
    - If details is True, returns a list of tuples containing detailed file information:
      document ID, path, status, metadata, group name, group status, and group log.
"""
        query = """
            SELECT documents.doc_id, documents.path, documents.status, documents.meta,
                   kb_group_documents.group_name, kb_group_documents.status, kb_group_documents.log
            FROM kb_group_documents
            JOIN documents ON kb_group_documents.doc_id = documents.doc_id
        """
        conds, params = [], []
        if group:
            conds.append('kb_group_documents.group_name = ?')
            params.append(group)

        if need_reparse is not None:
            conds.append('kb_group_documents.need_reparse = ?')
            params.append(int(need_reparse))

        status_cond, status_params = self.get_status_cond_and_params(status, exclude_status, prefix='kb_group_documents')
        if status_cond:
            conds.append(status_cond)
            params.extend(status_params)

        status_cond, status_params = self.get_status_cond_and_params(
            upload_status, exclude_upload_status, prefix='documents')
        if status_cond:
            conds.append(status_cond)
            params.extend(status_params)

        if conds: query += ' WHERE ' + ' AND '.join(conds)

        if limit:
            query += ' LIMIT ?'
            params.append(limit)

        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute(query, params)
            rows = cursor.fetchall()

        if not details: return [row[:2] for row in rows]
        return rows

    def delete_unreferenced_doc(self):
        """Deletes documents from the database that are marked for deletion and are no longer referenced by any knowledge base group.

This method queries documents with status "deleting" and a reference count of zero, deletes them from the database,
and adds operation logs for these deletions.

"""
        with self._db_lock, self._Session() as session:
            docs_to_delete = (
                session.query(KBDocument)
                .filter(KBDocument.status == DocListManager.Status.deleting, KBDocument.count == 0)
                .all()
            )
            for doc in docs_to_delete:
                session.delete(doc)
                log = KBOperationLogs(log=f"Delete obsolete file, doc_id:{doc.doc_id}, path:{doc.path}.")
                session.add(log)
            session.commit()

    def _add_doc_records(self, files: List[str], metadatas: Optional[List[Dict[str, Any]]] = None,
                         status: Optional[str] = DocListManager.Status.waiting, batch_size: int = 64):
        documents = []

        for i in range(0, len(files), batch_size):
            batch_files = files[i:i + batch_size]
            batch_metadatas = metadatas[i:i + batch_size] if metadatas else [None] * batch_size
            vals = []

            for i, file_path in enumerate(batch_files):
                doc_id = gen_docid(file_path)

                metadata = batch_metadatas[i].copy() if batch_metadatas[i] else {}
                metadata.setdefault(RAG_DOC_ID, doc_id)
                metadata.setdefault(RAG_DOC_PATH, file_path)

                vals.append(
                    {
                        KBDocument.doc_id.name: doc_id,
                        KBDocument.filename.name: os.path.basename(file_path),
                        KBDocument.path.name: file_path,
                        KBDocument.meta.name: json.dumps(metadata),
                        KBDocument.status.name: status,
                        KBDocument.count.name: 0,
                    }
                )
            with self._db_lock, self._Session() as session:
                rows = session.execute(
                    insert(KBDocument)
                    .values(vals)
                    .prefix_with('OR IGNORE')
                    .returning(KBDocument.doc_id, KBDocument.path)
                ).fetchall()
                session.commit()
                documents.extend(rows)
        return documents

    def get_docs_need_reparse(self, group: str) -> List[KBDocument]:
        """Retrieves the list of documents that require re-parsing within a specified knowledge base group.

Only documents with status "success" or "failed" and marked as needing reparse in the group are returned.

Args:
    group (str): Name of the knowledge base group.

**Returns:**

- List[KBDocument]: List of documents that need to be re-parsed.
"""
        with self._db_lock, self._Session() as session:
            filter_status_list = [DocListManager.Status.success, DocListManager.Status.failed]
            documents = (
                session.query(KBDocument).join(KBGroupDocuments, KBDocument.doc_id == KBGroupDocuments.doc_id)
                .filter(KBGroupDocuments.need_reparse.is_(True),
                        KBGroupDocuments.group_name == group,
                        KBGroupDocuments.status.in_(filter_status_list)).all())
            return documents
        return []

    def get_existing_paths_by_pattern(self, pattern: str) -> List[str]:
        """Retrieves a list of existing document paths that match a given pattern.

Args:
    pattern (str): Path matching pattern, supports SQL LIKE wildcards.

**Returns:**

- List[str]: List of existing document paths matching the pattern.
"""
        exist_paths = []
        with self._db_lock, self._Session() as session:
            docs = session.query(KBDocument).filter(KBDocument.path.like(pattern)).all()
            exist_paths = [doc.path for doc in docs]
        return exist_paths

    # TODO(wangzhihong): set to metadatas and enable this function
    def update_file_message(self, fileid: str, **kw):
        """Updates fields of the specified file record.

Args:
    fileid (str): Unique identifier of the file (doc_id).
    **kw: Key-value pairs of fields to update and their new values.
"""
        set_clause = ", ".join([f"{k} = ?" for k in kw.keys()])
        params = list(kw.values()) + [fileid]
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            conn.execute(f"UPDATE documents SET {set_clause} WHERE doc_id = ?", params)
            conn.commit()

    def update_file_status(self, file_ids: List[str], status: str,
                           cond_status_list: Union[None, List[str]] = None) -> List[DocPartRow]:
        """Updates the status of multiple files, optionally filtered by current status.

Args:
    file_ids (List[str]): List of file IDs to update.
    status (str): New status to set.
    cond_status_list (Union[None, List[str]], optional): List of statuses to filter files that can be updated. Defaults to None.

**Returns:**

- List[DocPartRow]: List of updated file IDs and their paths.
"""
        rows = []
        if cond_status_list is None:
            sql_cond = KBDocument.doc_id.in_(file_ids)
        else:
            sql_cond = sqlalchemy.and_(KBDocument.status.in_(cond_status_list), KBDocument.doc_id.in_(file_ids))
        with self._db_lock, self._Session() as session:
            stmt = (
                update(KBDocument)
                .where(sql_cond)
                .values(status=status)
                .returning(KBDocument.doc_id, KBDocument.path)
            )
            rows = session.execute(stmt).fetchall()
            session.commit()
        return rows

    def add_files_to_kb_group(self, file_ids: List[str], group: str):
        """Adds multiple files to the specified knowledge base group.

This method sets the file status to waiting.
If successfully added, increments the document's count.

Args:
    file_ids (List[str]): List of file IDs to add.
    group (str): Name of the knowledge base group.
"""
        with self._db_lock, self._Session() as session:
            vals = []
            for doc_id in file_ids:
                vals = {
                    KBGroupDocuments.doc_id.name: doc_id,
                    KBGroupDocuments.group_name.name: group,
                    KBGroupDocuments.status.name: DocListManager.Status.waiting,
                }
                rows = session.execute(
                    insert(KBGroupDocuments).values(vals).prefix_with('OR IGNORE').returning(KBGroupDocuments.doc_id)
                ).fetchall()
                session.commit()
                if not rows:
                    continue
                doc = session.query(KBDocument).filter_by(doc_id=rows[0].doc_id).one()
                doc.count += 1
                session.commit()

    def delete_files_from_kb_group(self, file_ids: List[str], group: str):
        """Deletes multiple files from the specified knowledge base group.

After deletion, decrements the document's count but not below zero.
If the document is not found, logs a warning.

Args:
    file_ids (List[str]): List of file IDs to delete.
    group (str): Name of the knowledge base group.
"""
        with self._db_lock, self._Session() as session:
            for doc_id in file_ids:
                records_to_delete = (
                    session.query(KBGroupDocuments)
                    .filter(KBGroupDocuments.doc_id == doc_id, KBGroupDocuments.group_name == group)
                    .all()
                )
                for record in records_to_delete:
                    session.delete(record)
                session.commit()
                if not records_to_delete:
                    continue
                try:
                    doc = session.query(KBDocument).filter_by(doc_id=records_to_delete[0].doc_id).one()
                    doc.count = max(0, doc.count - 1)
                    session.commit()
                except NoResultFound:
                    lazyllm.LOG.warning(f"No document found for {doc_id}")

    def get_file_status(self, fileid: str):
        """Gets the status of a specified file.

Args:
    fileid (str): Unique identifier of the file.

**Returns:**

- Optional[Tuple]: A tuple containing the status, or None if the file does not exist.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute("SELECT status FROM documents WHERE doc_id = ?", (fileid,))
        return cursor.fetchone()

    def update_kb_group(self, cond_file_ids: List[str], cond_group: Optional[str] = None,
                        cond_status_list: Optional[List[str]] = None, new_status: Optional[str] = None,
                        new_need_reparse: Optional[bool] = None) -> List[GroupDocPartRow]:
        """Updates the status and reparse need flag of specified files in a knowledge base group.

Batch updates files' status and need_reparse flag within a knowledge base group based on file IDs, group name, and optional status filter.

Args:
    cond_file_ids (List[str]): List of file IDs to update.
    cond_group (Optional[str]): Group name to filter files, if specified only updates files in this group.
    cond_status_list (Optional[List[str]]): Only update files whose status is in this list.
    new_status (Optional[str]): New status to set.
    new_need_reparse (Optional[bool]): New flag indicating if reparse is needed.

**Returns:**

- List[Tuple]: List of tuples of updated files containing doc_id, group_name, and status.
"""
        rows = []
        conds = []
        if not cond_file_ids:
            return rows
        conds.append(KBGroupDocuments.doc_id.in_(cond_file_ids))
        if cond_group is not None:
            conds.append(KBGroupDocuments.group_name == cond_group)
        if cond_status_list:
            conds.append(KBGroupDocuments.status.in_(cond_status_list))

        vals = {}
        if new_status is not None:
            vals[KBGroupDocuments.status.name] = new_status
        if new_need_reparse is not None:
            vals[KBGroupDocuments.need_reparse.name] = new_need_reparse

        if not vals:
            return rows
        with self._db_lock, self._Session() as session:
            stmt = (
                update(KBGroupDocuments)
                .where(sqlalchemy.and_(*conds))
                .values(vals)
                .returning(KBGroupDocuments.doc_id, KBGroupDocuments.group_name, KBGroupDocuments.status)
            )
            rows = session.execute(stmt).fetchall()
            session.commit()
        return rows

    def release(self):
        """Clears all documents, groups, and operation logs from the database.

This operation deletes all records from documents, document_groups, kb_group_documents, and operation_logs tables.

"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            conn.execute('delete from documents')
            conn.execute('delete from document_groups')
            conn.execute('delete from kb_group_documents')
            conn.execute('delete from operation_logs')
            conn.commit()

    def __reduce__(self):
        return (__class__, (self._path, self._name, self._enable_path_monitoring))

add_files_to_kb_group(file_ids, group)

Adds multiple files to the specified knowledge base group.

This method sets the file status to waiting. If successfully added, increments the document's count.

Parameters:

  • file_ids (List[str]) –

    List of file IDs to add.

  • group (str) –

    Name of the knowledge base group.

Source code in lazyllm/tools/rag/utils.py
    def add_files_to_kb_group(self, file_ids: List[str], group: str):
        """Adds multiple files to the specified knowledge base group.

This method sets the file status to waiting.
If successfully added, increments the document's count.

Args:
    file_ids (List[str]): List of file IDs to add.
    group (str): Name of the knowledge base group.
"""
        with self._db_lock, self._Session() as session:
            vals = []
            for doc_id in file_ids:
                vals = {
                    KBGroupDocuments.doc_id.name: doc_id,
                    KBGroupDocuments.group_name.name: group,
                    KBGroupDocuments.status.name: DocListManager.Status.waiting,
                }
                rows = session.execute(
                    insert(KBGroupDocuments).values(vals).prefix_with('OR IGNORE').returning(KBGroupDocuments.doc_id)
                ).fetchall()
                session.commit()
                if not rows:
                    continue
                doc = session.query(KBDocument).filter_by(doc_id=rows[0].doc_id).one()
                doc.count += 1
                session.commit()

add_kb_group(name)

Adds a new knowledge base group name to the database; ignores if the group already exists.

Parameters:

  • name (str) –

    The name of the knowledge base group to add.

Source code in lazyllm/tools/rag/utils.py
    def add_kb_group(self, name):
        """Adds a new knowledge base group name to the database; ignores if the group already exists.

Args:
    name (str): The name of the knowledge base group to add.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            conn.execute('INSERT OR IGNORE INTO document_groups (group_name) VALUES (?)', (name,))
            conn.commit()

delete_files_from_kb_group(file_ids, group)

Deletes multiple files from the specified knowledge base group.

After deletion, decrements the document's count but not below zero. If the document is not found, logs a warning.

Parameters:

  • file_ids (List[str]) –

    List of file IDs to delete.

  • group (str) –

    Name of the knowledge base group.

Source code in lazyllm/tools/rag/utils.py
    def delete_files_from_kb_group(self, file_ids: List[str], group: str):
        """Deletes multiple files from the specified knowledge base group.

After deletion, decrements the document's count but not below zero.
If the document is not found, logs a warning.

Args:
    file_ids (List[str]): List of file IDs to delete.
    group (str): Name of the knowledge base group.
"""
        with self._db_lock, self._Session() as session:
            for doc_id in file_ids:
                records_to_delete = (
                    session.query(KBGroupDocuments)
                    .filter(KBGroupDocuments.doc_id == doc_id, KBGroupDocuments.group_name == group)
                    .all()
                )
                for record in records_to_delete:
                    session.delete(record)
                session.commit()
                if not records_to_delete:
                    continue
                try:
                    doc = session.query(KBDocument).filter_by(doc_id=records_to_delete[0].doc_id).one()
                    doc.count = max(0, doc.count - 1)
                    session.commit()
                except NoResultFound:
                    lazyllm.LOG.warning(f"No document found for {doc_id}")

delete_unreferenced_doc()

Deletes documents from the database that are marked for deletion and are no longer referenced by any knowledge base group.

This method queries documents with status "deleting" and a reference count of zero, deletes them from the database, and adds operation logs for these deletions.

Source code in lazyllm/tools/rag/utils.py
    def delete_unreferenced_doc(self):
        """Deletes documents from the database that are marked for deletion and are no longer referenced by any knowledge base group.

This method queries documents with status "deleting" and a reference count of zero, deletes them from the database,
and adds operation logs for these deletions.

"""
        with self._db_lock, self._Session() as session:
            docs_to_delete = (
                session.query(KBDocument)
                .filter(KBDocument.status == DocListManager.Status.deleting, KBDocument.count == 0)
                .all()
            )
            for doc in docs_to_delete:
                session.delete(doc)
                log = KBOperationLogs(log=f"Delete obsolete file, doc_id:{doc.doc_id}, path:{doc.path}.")
                session.add(log)
            session.commit()

fetch_docs_changed_meta(group)

Fetches the list of documents within a specified knowledge base group that have updated metadata, and resets the new_meta field for those documents.

Parameters:

  • group (str) –

    Name of the knowledge base group.

Returns:

  • List[DocMetaChangedRow]: A list containing document IDs and their updated metadata.
Source code in lazyllm/tools/rag/utils.py
    def fetch_docs_changed_meta(self, group: str) -> List[DocMetaChangedRow]:
        """Fetches the list of documents within a specified knowledge base group that have updated metadata, and resets the new_meta field for those documents.

Args:
    group (str): Name of the knowledge base group.

**Returns:**

- List[DocMetaChangedRow]: A list containing document IDs and their updated metadata.
"""
        rows = []
        conds = [KBGroupDocuments.group_name == group, KBGroupDocuments.new_meta.isnot(None)]
        with self._db_lock, self._Session() as session:
            rows = (
                session.query(KBDocument.doc_id, KBGroupDocuments.new_meta)
                .join(KBGroupDocuments, KBDocument.doc_id == KBGroupDocuments.doc_id)
                .filter(*conds).all()
            )
            stmt = update(KBGroupDocuments).where(sqlalchemy.and_(*conds)).values(new_meta=None)
            session.execute(stmt)
            session.commit()
        return rows

get_docs(doc_ids)

Fetches document objects from the database corresponding to the given list of document IDs.

Parameters:

  • doc_ids (List[str]) –

    A list of document IDs to query.

Returns:

  • List[KBDocument]: A list of matching document objects. Returns an empty list if no matches found.
Source code in lazyllm/tools/rag/utils.py
    def get_docs(self, doc_ids: List[str]) -> List[KBDocument]:
        """Fetches document objects from the database corresponding to the given list of document IDs.

Args:
    doc_ids (List[str]): A list of document IDs to query.

**Returns:**

- List[KBDocument]: A list of matching document objects. Returns an empty list if no matches found.
"""
        with self._db_lock, self._Session() as session:
            docs = session.query(KBDocument).filter(KBDocument.doc_id.in_(doc_ids)).all()
            return docs
        return []

get_docs_need_reparse(group)

Retrieves the list of documents that require re-parsing within a specified knowledge base group.

Only documents with status "success" or "failed" and marked as needing reparse in the group are returned.

Parameters:

  • group (str) –

    Name of the knowledge base group.

Returns:

  • List[KBDocument]: List of documents that need to be re-parsed.
Source code in lazyllm/tools/rag/utils.py
    def get_docs_need_reparse(self, group: str) -> List[KBDocument]:
        """Retrieves the list of documents that require re-parsing within a specified knowledge base group.

Only documents with status "success" or "failed" and marked as needing reparse in the group are returned.

Args:
    group (str): Name of the knowledge base group.

**Returns:**

- List[KBDocument]: List of documents that need to be re-parsed.
"""
        with self._db_lock, self._Session() as session:
            filter_status_list = [DocListManager.Status.success, DocListManager.Status.failed]
            documents = (
                session.query(KBDocument).join(KBGroupDocuments, KBDocument.doc_id == KBGroupDocuments.doc_id)
                .filter(KBGroupDocuments.need_reparse.is_(True),
                        KBGroupDocuments.group_name == group,
                        KBGroupDocuments.status.in_(filter_status_list)).all())
            return documents
        return []

get_existing_paths_by_pattern(pattern)

Retrieves a list of existing document paths that match a given pattern.

Parameters:

  • pattern (str) –

    Path matching pattern, supports SQL LIKE wildcards.

Returns:

  • List[str]: List of existing document paths matching the pattern.
Source code in lazyllm/tools/rag/utils.py
    def get_existing_paths_by_pattern(self, pattern: str) -> List[str]:
        """Retrieves a list of existing document paths that match a given pattern.

Args:
    pattern (str): Path matching pattern, supports SQL LIKE wildcards.

**Returns:**

- List[str]: List of existing document paths matching the pattern.
"""
        exist_paths = []
        with self._db_lock, self._Session() as session:
            docs = session.query(KBDocument).filter(KBDocument.path.like(pattern)).all()
            exist_paths = [doc.path for doc in docs]
        return exist_paths

get_file_status(fileid)

Gets the status of a specified file.

Parameters:

  • fileid (str) –

    Unique identifier of the file.

Returns:

  • Optional[Tuple]: A tuple containing the status, or None if the file does not exist.
Source code in lazyllm/tools/rag/utils.py
    def get_file_status(self, fileid: str):
        """Gets the status of a specified file.

Args:
    fileid (str): Unique identifier of the file.

**Returns:**

- Optional[Tuple]: A tuple containing the status, or None if the file does not exist.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute("SELECT status FROM documents WHERE doc_id = ?", (fileid,))
        return cursor.fetchone()

list_all_kb_group()

Lists all knowledge base group names stored in the database.

Returns:

  • List[str]: A list of knowledge base group names.
Source code in lazyllm/tools/rag/utils.py
    def list_all_kb_group(self):
        """Lists all knowledge base group names stored in the database.

**Returns:**

- List[str]: A list of knowledge base group names.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute("SELECT group_name FROM document_groups")
            return [row[0] for row in cursor]

list_files(limit=None, details=False, status=DocListManager.Status.all, exclude_status=None)

Lists files in the document database based on status filters and returns either full records or file paths.

Parameters:

  • limit (Optional[int], default: None ) –

    The maximum number of records to return. If None, all matching records are returned.

  • details (bool, default: False ) –

    Whether to return full database rows or just file paths (document IDs).

  • status (Union[str, List[str]], default: all ) –

    Status values to include in the result. Defaults to including all.

  • exclude_status (Optional[Union[str, List[str]]], default: None ) –

    Status values to exclude from the result.

Returns:

  • list: A list of file records or document paths depending on the details flag.
Source code in lazyllm/tools/rag/utils.py
    def list_files(self, limit: Optional[int] = None, details: bool = False,
                   status: Union[str, List[str]] = DocListManager.Status.all,
                   exclude_status: Optional[Union[str, List[str]]] = None):
        """Lists files in the document database based on status filters and returns either full records or file paths.

Args:
    limit (Optional[int]): The maximum number of records to return. If None, all matching records are returned.
    details (bool): Whether to return full database rows or just file paths (document IDs).
    status (Union[str, List[str]]): Status values to include in the result. Defaults to including all.
    exclude_status (Optional[Union[str, List[str]]]): Status values to exclude from the result.

**Returns:**

- list: A list of file records or document paths depending on the `details` flag.
"""
        query = "SELECT * FROM documents"
        params = []
        status_cond, status_params = self.get_status_cond_and_params(status, exclude_status, prefix=None)
        if status_cond:
            query += f' WHERE {status_cond}'
            params.extend(status_params)
        if limit:
            query += " LIMIT ?"
            params.append(limit)
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute(query, params)
            return cursor.fetchall() if details else [row[0] for row in cursor]

list_kb_group_files(group=None, limit=None, details=False, status=DocListManager.Status.all, exclude_status=None, upload_status=DocListManager.Status.all, exclude_upload_status=None, need_reparse=None)

Lists files in a specified knowledge base group, with support for multiple filters.

Parameters:

  • group (str, default: None ) –

    Knowledge base group name to filter by. If None, no group filtering is applied.

  • limit (int, default: None ) –

    Limit on the number of files to return.

  • details (bool, default: False ) –

    Whether to return detailed file information.

  • status (str or List[str], default: all ) –

    Filter files by group document status.

  • exclude_status (str or List[str], default: None ) –

    Exclude files with these group document statuses.

  • upload_status (str or List[str], default: all ) –

    Filter files by upload document status.

  • exclude_upload_status (str or List[str], default: None ) –

    Exclude files with these upload document statuses.

  • need_reparse (bool, default: None ) –

    If set, only returns files marked as needing reparse.

Returns:

  • list:
    • If details is False, returns a list of tuples (doc_id, path).
    • If details is True, returns a list of tuples containing detailed file information: document ID, path, status, metadata, group name, group status, and group log.
Source code in lazyllm/tools/rag/utils.py
    def list_kb_group_files(self, group: str = None, limit: Optional[int] = None, details: bool = False,
                            status: Union[str, List[str]] = DocListManager.Status.all,
                            exclude_status: Optional[Union[str, List[str]]] = None,
                            upload_status: Union[str, List[str]] = DocListManager.Status.all,
                            exclude_upload_status: Optional[Union[str, List[str]]] = None,
                            need_reparse: Optional[bool] = None):
        """Lists files in a specified knowledge base group, with support for multiple filters.

Args:
    group (str, optional): Knowledge base group name to filter by. If None, no group filtering is applied.
    limit (int, optional): Limit on the number of files to return.
    details (bool): Whether to return detailed file information.
    status (str or List[str], optional): Filter files by group document status.
    exclude_status (str or List[str], optional): Exclude files with these group document statuses.
    upload_status (str or List[str], optional): Filter files by upload document status.
    exclude_upload_status (str or List[str], optional): Exclude files with these upload document statuses.
    need_reparse (bool, optional): If set, only returns files marked as needing reparse.

**Returns:**

- list: 
    - If details is False, returns a list of tuples (doc_id, path).
    - If details is True, returns a list of tuples containing detailed file information:
      document ID, path, status, metadata, group name, group status, and group log.
"""
        query = """
            SELECT documents.doc_id, documents.path, documents.status, documents.meta,
                   kb_group_documents.group_name, kb_group_documents.status, kb_group_documents.log
            FROM kb_group_documents
            JOIN documents ON kb_group_documents.doc_id = documents.doc_id
        """
        conds, params = [], []
        if group:
            conds.append('kb_group_documents.group_name = ?')
            params.append(group)

        if need_reparse is not None:
            conds.append('kb_group_documents.need_reparse = ?')
            params.append(int(need_reparse))

        status_cond, status_params = self.get_status_cond_and_params(status, exclude_status, prefix='kb_group_documents')
        if status_cond:
            conds.append(status_cond)
            params.extend(status_params)

        status_cond, status_params = self.get_status_cond_and_params(
            upload_status, exclude_upload_status, prefix='documents')
        if status_cond:
            conds.append(status_cond)
            params.extend(status_params)

        if conds: query += ' WHERE ' + ' AND '.join(conds)

        if limit:
            query += ' LIMIT ?'
            params.append(limit)

        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute(query, params)
            rows = cursor.fetchall()

        if not details: return [row[:2] for row in rows]
        return rows

release()

Clears all documents, groups, and operation logs from the database.

This operation deletes all records from documents, document_groups, kb_group_documents, and operation_logs tables.

Source code in lazyllm/tools/rag/utils.py
    def release(self):
        """Clears all documents, groups, and operation logs from the database.

This operation deletes all records from documents, document_groups, kb_group_documents, and operation_logs tables.

"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            conn.execute('delete from documents')
            conn.execute('delete from document_groups')
            conn.execute('delete from kb_group_documents')
            conn.execute('delete from operation_logs')
            conn.commit()

set_docs_new_meta(doc_meta)

Batch updates the metadata (meta) of documents, and simultaneously updates the new_meta field of documents in knowledge base groups for documents that are not in waiting status.

Parameters:

  • doc_meta (Dict[str, dict]) –

    A dictionary mapping document IDs to their new metadata dictionaries.

Source code in lazyllm/tools/rag/utils.py
    def set_docs_new_meta(self, doc_meta: Dict[str, dict]):
        """Batch updates the metadata (meta) of documents, and simultaneously updates the new_meta field of documents in knowledge base groups for documents that are not in waiting status.

Args:
    doc_meta (Dict[str, dict]): A dictionary mapping document IDs to their new metadata dictionaries.
"""
        data_to_update = [{"_doc_id": k, "_meta": json.dumps(v)} for k, v in doc_meta.items()]
        with self._db_lock, self._Session() as session:
            # Use sqlalchemy core bulk update
            stmt = KBDocument.__table__.update().where(
                KBDocument.doc_id == bindparam("_doc_id")).values(meta=bindparam("_meta"))
            session.execute(stmt, data_to_update)
            session.commit()

            stmt = KBGroupDocuments.__table__.update().where(
                KBGroupDocuments.doc_id == bindparam("_doc_id"),
                KBGroupDocuments.status != DocListManager.Status.waiting).values(new_meta=bindparam("_meta"))
            session.execute(stmt, data_to_update)
            session.commit()

table_inited()

Checks whether the "documents" table has been initialized in the database.

The method queries the sqlite_master metadata table to verify if the "documents" table exists.

Returns:

  • bool: True if the "documents" table exists, False otherwise.
Source code in lazyllm/tools/rag/utils.py
    def table_inited(self):
        """Checks whether the "documents" table has been initialized in the database.

The method queries the sqlite_master metadata table to verify if the "documents" table exists.

**Returns:**

- bool: True if the "documents" table exists, False otherwise.
"""
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='documents'")
            return cursor.fetchone() is not None

update_file_message(fileid, **kw)

Updates fields of the specified file record.

Parameters:

  • fileid (str) –

    Unique identifier of the file (doc_id).

  • **kw

    Key-value pairs of fields to update and their new values.

Source code in lazyllm/tools/rag/utils.py
    def update_file_message(self, fileid: str, **kw):
        """Updates fields of the specified file record.

Args:
    fileid (str): Unique identifier of the file (doc_id).
    **kw: Key-value pairs of fields to update and their new values.
"""
        set_clause = ", ".join([f"{k} = ?" for k in kw.keys()])
        params = list(kw.values()) + [fileid]
        with self._db_lock, sqlite3.connect(self._db_path, check_same_thread=self._check_same_thread) as conn:
            conn.execute(f"UPDATE documents SET {set_clause} WHERE doc_id = ?", params)
            conn.commit()

update_file_status(file_ids, status, cond_status_list=None)

Updates the status of multiple files, optionally filtered by current status.

Parameters:

  • file_ids (List[str]) –

    List of file IDs to update.

  • status (str) –

    New status to set.

  • cond_status_list (Union[None, List[str]], default: None ) –

    List of statuses to filter files that can be updated. Defaults to None.

Returns:

  • List[DocPartRow]: List of updated file IDs and their paths.
Source code in lazyllm/tools/rag/utils.py
    def update_file_status(self, file_ids: List[str], status: str,
                           cond_status_list: Union[None, List[str]] = None) -> List[DocPartRow]:
        """Updates the status of multiple files, optionally filtered by current status.

Args:
    file_ids (List[str]): List of file IDs to update.
    status (str): New status to set.
    cond_status_list (Union[None, List[str]], optional): List of statuses to filter files that can be updated. Defaults to None.

**Returns:**

- List[DocPartRow]: List of updated file IDs and their paths.
"""
        rows = []
        if cond_status_list is None:
            sql_cond = KBDocument.doc_id.in_(file_ids)
        else:
            sql_cond = sqlalchemy.and_(KBDocument.status.in_(cond_status_list), KBDocument.doc_id.in_(file_ids))
        with self._db_lock, self._Session() as session:
            stmt = (
                update(KBDocument)
                .where(sql_cond)
                .values(status=status)
                .returning(KBDocument.doc_id, KBDocument.path)
            )
            rows = session.execute(stmt).fetchall()
            session.commit()
        return rows

update_kb_group(cond_file_ids, cond_group=None, cond_status_list=None, new_status=None, new_need_reparse=None)

Updates the status and reparse need flag of specified files in a knowledge base group.

Batch updates files' status and need_reparse flag within a knowledge base group based on file IDs, group name, and optional status filter.

Parameters:

  • cond_file_ids (List[str]) –

    List of file IDs to update.

  • cond_group (Optional[str], default: None ) –

    Group name to filter files, if specified only updates files in this group.

  • cond_status_list (Optional[List[str]], default: None ) –

    Only update files whose status is in this list.

  • new_status (Optional[str], default: None ) –

    New status to set.

  • new_need_reparse (Optional[bool], default: None ) –

    New flag indicating if reparse is needed.

Returns:

  • List[Tuple]: List of tuples of updated files containing doc_id, group_name, and status.
Source code in lazyllm/tools/rag/utils.py
    def update_kb_group(self, cond_file_ids: List[str], cond_group: Optional[str] = None,
                        cond_status_list: Optional[List[str]] = None, new_status: Optional[str] = None,
                        new_need_reparse: Optional[bool] = None) -> List[GroupDocPartRow]:
        """Updates the status and reparse need flag of specified files in a knowledge base group.

Batch updates files' status and need_reparse flag within a knowledge base group based on file IDs, group name, and optional status filter.

Args:
    cond_file_ids (List[str]): List of file IDs to update.
    cond_group (Optional[str]): Group name to filter files, if specified only updates files in this group.
    cond_status_list (Optional[List[str]]): Only update files whose status is in this list.
    new_status (Optional[str]): New status to set.
    new_need_reparse (Optional[bool]): New flag indicating if reparse is needed.

**Returns:**

- List[Tuple]: List of tuples of updated files containing doc_id, group_name, and status.
"""
        rows = []
        conds = []
        if not cond_file_ids:
            return rows
        conds.append(KBGroupDocuments.doc_id.in_(cond_file_ids))
        if cond_group is not None:
            conds.append(KBGroupDocuments.group_name == cond_group)
        if cond_status_list:
            conds.append(KBGroupDocuments.status.in_(cond_status_list))

        vals = {}
        if new_status is not None:
            vals[KBGroupDocuments.status.name] = new_status
        if new_need_reparse is not None:
            vals[KBGroupDocuments.need_reparse.name] = new_need_reparse

        if not vals:
            return rows
        with self._db_lock, self._Session() as session:
            stmt = (
                update(KBGroupDocuments)
                .where(sqlalchemy.and_(*conds))
                .values(vals)
                .returning(KBGroupDocuments.doc_id, KBGroupDocuments.group_name, KBGroupDocuments.status)
            )
            rows = session.execute(stmt).fetchall()
            session.commit()
        return rows

update_need_reparsing(doc_id, need_reparse, group_name=None)

Updates the re-parsing flag for a specific document.

This method sets whether a document should be re-parsed. If a group name is provided, the update is scoped to that group only.

Parameters:

  • doc_id (str) –

    The unique identifier of the document.

  • need_reparse (bool) –

    Whether the document needs to be re-parsed.

  • group_name (Optional[str], default: None ) –

    Optional. The knowledge base group name to filter by. If provided, only documents in the specified group will be updated.

Source code in lazyllm/tools/rag/utils.py
    def update_need_reparsing(self, doc_id: str, need_reparse: bool, group_name: Optional[str] = None):
        """Updates the re-parsing flag for a specific document.

This method sets whether a document should be re-parsed. If a group name is provided, the update is scoped to that group only.

Args:
    doc_id (str): The unique identifier of the document.
    need_reparse (bool): Whether the document needs to be re-parsed.
    group_name (Optional[str]): Optional. The knowledge base group name to filter by. If provided, only documents in the specified group will be updated.
"""
        with self._db_lock, self._Session() as session:
            stmt = update(KBGroupDocuments).where(KBGroupDocuments.doc_id == doc_id)
            if group_name is not None: stmt = stmt.where(KBGroupDocuments.group_name == group_name)
            session.execute(stmt.values(need_reparse=need_reparse))
            session.commit()

validate_paths(paths)

Validates whether the documents corresponding to the given paths can be safely added to the database.

The method checks if the document already exists. If it exists, it verifies whether the document is currently being parsed, waiting to be parsed, or was not successfully re-parsed last time.

Parameters:

  • paths (List[str]) –

    A list of file paths to validate.

Returns:

  • Tuple[bool, str, List[bool]]:
    • bool: Whether all paths passed validation.
    • str: Description message of the validation result.
    • List[bool]: A boolean list corresponding to input paths, indicating whether each path is new (True) or already exists (False). If validation fails, this value is None.
Source code in lazyllm/tools/rag/utils.py
    def validate_paths(self, paths: List[str]) -> Tuple[bool, str, List[bool]]:
        """Validates whether the documents corresponding to the given paths can be safely added to the database.

The method checks if the document already exists. If it exists, it verifies whether the document is currently
being parsed, waiting to be parsed, or was not successfully re-parsed last time.

Args:
    paths (List[str]): A list of file paths to validate.

**Returns:**

- Tuple[bool, str, List[bool]]: 
    - bool: Whether all paths passed validation.
    - str: Description message of the validation result.
    - List[bool]: A boolean list corresponding to input paths, indicating whether each path is new (True) or already exists (False).
      If validation fails, this value is None.
"""
        # check and return: success, msg, path_is_new for each path
        unsafe_staus_set = set([DocListManager.Status.working, DocListManager.Status.waiting])
        paths_is_new = [True] * len(paths)
        doc_ids = [gen_docid(path) for path in paths]
        doc_id_to_path = {doc_id: path for doc_id, path in zip(doc_ids, paths)}
        found_doc_ids = []
        found_doc_group_rows = []
        with self._db_lock, self._Session() as session:
            rows = session.execute(
                select(KBDocument.doc_id).where(KBDocument.doc_id.in_(doc_ids))
            ).fetchall()
            if len(rows) == 0:
                return True, "Success", paths_is_new
            found_doc_ids = [row.doc_id for row in rows]
            found_doc_group_rows = session.execute(
                select(KBGroupDocuments.doc_id, KBGroupDocuments.need_reparse, KBGroupDocuments.status)
                .where(KBGroupDocuments.doc_id.in_(found_doc_ids))).fetchall()

        for doc_group_record in found_doc_group_rows:
            if doc_group_record.need_reparse:
                msg = f"Failed: {doc_id_to_path[doc_group_record.doc_id]} lasttime reparsing has not been finished"
                return False, msg, None
            if doc_group_record.status in unsafe_staus_set:
                return False, f"Failed: {doc_id_to_path[doc_group_record.doc_id]} is being parsed by kbgroup", None
        found_doc_ids = set(found_doc_ids)
        for i in range(len(paths)):
            cur_doc_id = doc_ids[i]
            if cur_doc_id in found_doc_ids:
                paths_is_new[i] = False
        return True, "Success", paths_is_new

lazyllm.tools.rag.data_loaders.DirectoryReader

A directory reader class for loading and processing documents from file directories.

This class provides functionality to read documents from specified directories and convert them into document nodes. It supports both local and global file readers, and can handle different types of documents including images.

Parameters:

  • input_files (Optional[List[str]]) –

    A list of file paths to read. If None, files will be loaded when calling load_data method.

  • local_readers (Optional[Dict], default: None ) –

    A dictionary of local file readers specific to this instance. Keys are file patterns, values are reader functions.

  • global_readers (Optional[Dict], default: None ) –

    A dictionary of global file readers shared across all instances. Keys are file patterns, values are reader functions.

Examples:

>>> from lazyllm.tools.rag.data_loaders import DirectoryReader
>>> from lazyllm.tools.rag.readers import DocxReader, PDFReader
>>> local_readers = {
...     "**/*.docx": DocxReader,
...     "**/*.pdf": PDFReader
>>> }
>>> reader = DirectoryReader(
...     input_files=["path/to/documents"],
...     local_readers=local_readers,
...     global_readers={}
>>> )
>>> documents = reader.load_data()
>>> print(f"加载了 {len(documents)} 个文档")
Source code in lazyllm/tools/rag/data_loaders.py
class DirectoryReader:
    """A directory reader class for loading and processing documents from file directories.

This class provides functionality to read documents from specified directories and convert them into document nodes. It supports both local and global file readers, and can handle different types of documents including images.

Args:
    input_files (Optional[List[str]]): A list of file paths to read. If None, files will be loaded when calling load_data method.
    local_readers (Optional[Dict]): A dictionary of local file readers specific to this instance. Keys are file patterns, values are reader functions.
    global_readers (Optional[Dict]): A dictionary of global file readers shared across all instances. Keys are file patterns, values are reader functions.


Examples:
    >>> from lazyllm.tools.rag.data_loaders import DirectoryReader
    >>> from lazyllm.tools.rag.readers import DocxReader, PDFReader
    >>> local_readers = {
    ...     "**/*.docx": DocxReader,
    ...     "**/*.pdf": PDFReader
    >>> }
    >>> reader = DirectoryReader(
    ...     input_files=["path/to/documents"],
    ...     local_readers=local_readers,
    ...     global_readers={}
    >>> )
    >>> documents = reader.load_data()
    >>> print(f"加载了 {len(documents)} 个文档")
    """
    def __init__(self, input_files: Optional[List[str]], local_readers: Optional[Dict] = None,
                 global_readers: Optional[Dict] = None) -> None:
        self._input_files = input_files
        self._local_readers = local_readers
        self._global_readers = global_readers

    def load_data(self, input_files: Optional[List[str]] = None, metadatas: Optional[Dict] = None,
                  *, split_image_nodes: bool = False) -> List[DocNode]:
        """Load and process documents from the specified input files.

This method reads documents from the input files using the configured file readers (both local and global), processes them into document nodes, and optionally separates image nodes from text nodes.

Args:
    input_files (Optional[List[str]]): A list of file paths to read. If None, uses the files specified during initialization.
    metadatas (Optional[Dict]): Additional metadata to associate with the loaded documents.
    split_image_nodes (bool): Whether to separate image nodes from text nodes. If True, returns a tuple of (text_nodes, image_nodes). If False, returns all nodes together.

**Returns:**
- Union[List[DocNode], Tuple[List[DocNode], List[ImageDocNode]]]: If split_image_nodes is False, returns a list of all document nodes. If True, returns a tuple containing text nodes and image nodes separately.

"""
        input_files = input_files or self._input_files
        file_readers = self._local_readers.copy()
        for key, func in self._global_readers.items():
            if key not in file_readers: file_readers[key] = func
        LOG.info(f"DirectoryReader loads data, input files: {input_files}")
        reader = SimpleDirectoryReader(input_files=input_files, file_extractor=file_readers, metadatas=metadatas)
        nodes: List[DocNode] = []
        image_nodes: List[ImageDocNode] = []
        for doc in reader():
            doc._group = LAZY_IMAGE_GROUP if isinstance(doc, ImageDocNode) else LAZY_ROOT_NAME
            if not split_image_nodes or not isinstance(doc, ImageDocNode):
                nodes.append(doc)
            else:
                image_nodes.append(doc)
        if not nodes and not image_nodes:
            LOG.warning(
                f"No nodes load from path {input_files}, please check your data path."
            )
        LOG.info("DirectoryReader loads data done!")
        return (nodes, image_nodes) if split_image_nodes else nodes

load_data(input_files=None, metadatas=None, *, split_image_nodes=False)

Load and process documents from the specified input files.

This method reads documents from the input files using the configured file readers (both local and global), processes them into document nodes, and optionally separates image nodes from text nodes.

Parameters:

  • input_files (Optional[List[str]], default: None ) –

    A list of file paths to read. If None, uses the files specified during initialization.

  • metadatas (Optional[Dict], default: None ) –

    Additional metadata to associate with the loaded documents.

  • split_image_nodes (bool, default: False ) –

    Whether to separate image nodes from text nodes. If True, returns a tuple of (text_nodes, image_nodes). If False, returns all nodes together.

Returns: - Union[List[DocNode], Tuple[List[DocNode], List[ImageDocNode]]]: If split_image_nodes is False, returns a list of all document nodes. If True, returns a tuple containing text nodes and image nodes separately.

Source code in lazyllm/tools/rag/data_loaders.py
    def load_data(self, input_files: Optional[List[str]] = None, metadatas: Optional[Dict] = None,
                  *, split_image_nodes: bool = False) -> List[DocNode]:
        """Load and process documents from the specified input files.

This method reads documents from the input files using the configured file readers (both local and global), processes them into document nodes, and optionally separates image nodes from text nodes.

Args:
    input_files (Optional[List[str]]): A list of file paths to read. If None, uses the files specified during initialization.
    metadatas (Optional[Dict]): Additional metadata to associate with the loaded documents.
    split_image_nodes (bool): Whether to separate image nodes from text nodes. If True, returns a tuple of (text_nodes, image_nodes). If False, returns all nodes together.

**Returns:**
- Union[List[DocNode], Tuple[List[DocNode], List[ImageDocNode]]]: If split_image_nodes is False, returns a list of all document nodes. If True, returns a tuple containing text nodes and image nodes separately.

"""
        input_files = input_files or self._input_files
        file_readers = self._local_readers.copy()
        for key, func in self._global_readers.items():
            if key not in file_readers: file_readers[key] = func
        LOG.info(f"DirectoryReader loads data, input files: {input_files}")
        reader = SimpleDirectoryReader(input_files=input_files, file_extractor=file_readers, metadatas=metadatas)
        nodes: List[DocNode] = []
        image_nodes: List[ImageDocNode] = []
        for doc in reader():
            doc._group = LAZY_IMAGE_GROUP if isinstance(doc, ImageDocNode) else LAZY_ROOT_NAME
            if not split_image_nodes or not isinstance(doc, ImageDocNode):
                nodes.append(doc)
            else:
                image_nodes.append(doc)
        if not nodes and not image_nodes:
            LOG.warning(
                f"No nodes load from path {input_files}, please check your data path."
            )
        LOG.info("DirectoryReader loads data done!")
        return (nodes, image_nodes) if split_image_nodes else nodes

lazyllm.tools.SentenceSplitter

Bases: NodeTransform

Split sentences into chunks of a specified size. You can specify the size of the overlap between adjacent chunks.

Parameters:

  • chunk_size (int, default: 1024 ) –

    The size of the chunk after splitting.

  • chunk_overlap (int, default: 200 ) –

    The length of the overlapping content between two adjacent chunks.

  • num_workers (int, default: 0 ) –

    Controls the number of threads or processes used for parallel processing.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import Document, SentenceSplitter
>>> m = lazyllm.OnlineEmbeddingModule(source="glm")
>>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
>>> documents.create_node_group(name="sentences", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
Source code in lazyllm/tools/rag/transform.py
class SentenceSplitter(NodeTransform):
    """
Split sentences into chunks of a specified size. You can specify the size of the overlap between adjacent chunks.

Args:
    chunk_size (int): The size of the chunk after splitting.
    chunk_overlap (int): The length of the overlapping content between two adjacent chunks.
    num_workers (int): Controls the number of threads or processes used for parallel processing.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import Document, SentenceSplitter
    >>> m = lazyllm.OnlineEmbeddingModule(source="glm")
    >>> documents = Document(dataset_path='your_doc_path', embed=m, manager=False)
    >>> documents.create_node_group(name="sentences", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
    """
    def __init__(self, chunk_size: int = 1024, chunk_overlap: int = 200, num_workers: int = 0):
        super(__class__, self).__init__(num_workers=num_workers)
        if chunk_overlap > chunk_size:
            raise ValueError(
                f'Got a larger chunk overlap ({chunk_overlap}) than chunk size '
                f'({chunk_size}), should be smaller.'
            )

        assert (
            chunk_size > 0 and chunk_overlap >= 0
        ), 'chunk size should > 0 and chunk_overlap should >= 0'

        try:
            if 'TIKTOKEN_CACHE_DIR' not in os.environ and 'DATA_GYM_CACHE_DIR' not in os.environ:
                path = os.path.join(config['model_path'], 'tiktoken')
                os.makedirs(path, exist_ok=True)
                os.environ['TIKTOKEN_CACHE_DIR'] = path
            self._tiktoken_tokenizer = tiktoken.encoding_for_model('gpt-3.5-turbo')
            os.environ.pop('TIKTOKEN_CACHE_DIR')
        except requests.exceptions.ConnectionError:
            LOG.error(
                'Unable to download the vocabulary file for tiktoken `gpt-3.5-turbo`. '
                'Please check your internet connection. '
                'Alternatively, you can manually download the file '
                'and set the `TIKTOKEN_CACHE_DIR` environment variable.'
            )
            raise
        except Exception as e:
            LOG.error(f'Unable to build tiktoken tokenizer with error `{e}`')
            raise
        self._punkt_st_tokenizer = nltk.tokenize.PunktSentenceTokenizer()

        self._sentence_split_fns = [
            partial(split_text_keep_separator, separator='\n\n\n'),  # paragraph
            self._punkt_st_tokenizer.tokenize,
        ]

        self._sub_sentence_split_fns = [
            lambda t: re.findall(r'[^,.;。?!]+[,.;。?!]?', t),
            partial(split_text_keep_separator, separator=' '),
            list,  # split by character
        ]

        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def transform(self, node: DocNode, **kwargs) -> List[str]:
        return self.split_text(
            node.get_text(),
            metadata_size=self._get_metadata_size(node),
        )

    def _get_metadata_size(self, node: DocNode) -> int:
        # Return the bigger size to ensure chunk_size < limit
        return max(
            self._token_size(node.get_metadata_str(mode=MetadataMode.EMBED)),
            self._token_size(node.get_metadata_str(mode=MetadataMode.LLM)),
        )

    def split_text(self, text: str, metadata_size: int) -> List[str]:
        if text == '':
            return ['']
        effective_chunk_size = self.chunk_size - metadata_size
        if effective_chunk_size <= 0:
            raise ValueError(
                f'Metadata length ({metadata_size}) is longer than chunk size '
                f'({self.chunk_size}). Consider increasing the chunk size or '
                'decreasing the size of your metadata to avoid this.'
            )
        elif effective_chunk_size < 50:
            LOG.warning(
                f'Metadata length ({metadata_size}) is close to chunk size '
                f'({self.chunk_size}). Resulting chunks are less than 50 tokens. '
                'Consider increasing the chunk size or decreasing the size of '
                'your metadata to avoid this.'
            )

        splits = self._split(text, effective_chunk_size)
        chunks = self._merge(splits, effective_chunk_size)
        return chunks

    def _split(self, text: str, chunk_size: int) -> List[_Split]:
        """Break text into splits that are smaller than chunk size.

        The order of splitting is:
        1. split by paragraph separator
        2. split by chunking tokenizer
        3. split by second chunking regex
        4. split by default separator (' ')
        5. split by character
        """
        token_size = self._token_size(text)
        if token_size <= chunk_size:
            return [_Split(text, is_sentence=True, token_size=token_size)]

        text_splits_by_fns, is_sentence = self._get_splits_by_fns(text)

        text_splits = []
        for text in text_splits_by_fns:
            token_size = self._token_size(text)
            if token_size <= chunk_size:
                text_splits.append(
                    _Split(
                        text,
                        is_sentence=is_sentence,
                        token_size=token_size,
                    )
                )
            else:
                recursive_text_splits = self._split(text, chunk_size=chunk_size)
                text_splits.extend(recursive_text_splits)
        return text_splits

    def _merge(self, splits: List[_Split], chunk_size: int) -> List[str]:
        chunks: List[str] = []
        cur_chunk: List[Tuple[str, int]] = []  # list of (text, length)
        cur_chunk_len = 0
        is_chunk_new = True

        def close_chunk() -> None:
            nonlocal cur_chunk, cur_chunk_len, is_chunk_new

            chunks.append(''.join([text for text, _ in cur_chunk]))
            last_chunk = cur_chunk
            cur_chunk = []
            cur_chunk_len = 0
            is_chunk_new = True

            # Add overlap to the next chunk using the last one first
            overlap_len = 0
            for text, length in reversed(last_chunk):
                if overlap_len + length > self.chunk_overlap:
                    break
                cur_chunk.append((text, length))
                overlap_len += length
                cur_chunk_len += length
            cur_chunk.reverse()

        i = 0
        while i < len(splits):
            cur_split = splits[i]
            if cur_split.token_size > chunk_size:
                raise ValueError('Single token exceeded chunk size')
            if cur_chunk_len + cur_split.token_size > chunk_size and not is_chunk_new:
                # if adding split to current chunk exceeds chunk size
                close_chunk()
            else:
                if (
                    cur_split.is_sentence
                    or cur_chunk_len + cur_split.token_size <= chunk_size
                    or is_chunk_new  # new chunk, always add at least one split
                ):
                    # add split to chunk
                    cur_chunk_len += cur_split.token_size
                    cur_chunk.append((cur_split.text, cur_split.token_size))
                    i += 1
                    is_chunk_new = False
                else:
                    close_chunk()

        # handle the last chunk
        if not is_chunk_new:
            chunks.append(''.join([text for text, _ in cur_chunk]))

        # Remove whitespace only chunks and remove leading and trailing whitespace.
        return [stripped_chunk for chunk in chunks if (stripped_chunk := chunk.strip())]

    def _token_size(self, text: str) -> int:
        return len(self._tiktoken_tokenizer.encode(text, allowed_special='all'))

    def _get_splits_by_fns(self, text: str) -> Tuple[List[str], bool]:
        for split_fn in self._sentence_split_fns:
            splits = split_fn(text)
            if len(splits) > 1:
                return splits, True

        for split_fn in self._sub_sentence_split_fns:
            splits = split_fn(text)
            if len(splits) > 1:
                break

        return splits, False

lazyllm.tools.LLMParser

Bases: NodeTransform

A text summarizer and keyword extractor that is responsible for analyzing the text input by the user and providing concise summaries or extracting relevant keywords based on the requested task.

Parameters:

  • llm (TrainableModule) –

    A trainable module.

  • language (str) –

    The language type, currently only supports Chinese (zh) and English (en).

  • task_type (str) –

    Currently supports two types of tasks: summary and keyword extraction.

  • num_workers (int, default: 30 ) –

    Controls the number of threads or processes used for parallel processing.

Examples:

>>> from lazyllm import TrainableModule
>>> from lazyllm.tools.rag import LLMParser
>>> llm = TrainableModule("internlm2-chat-7b")
>>> summary_parser = LLMParser(llm, language="en", task_type="summary")
Source code in lazyllm/tools/rag/transform.py
class LLMParser(NodeTransform):
    """
A text summarizer and keyword extractor that is responsible for analyzing the text input by the user and providing concise summaries or extracting relevant keywords based on the requested task.

Args:
    llm (TrainableModule): A trainable module.
    language (str): The language type, currently only supports Chinese (zh) and English (en).
    task_type (str): Currently supports two types of tasks: summary and keyword extraction.
    num_workers (int): Controls the number of threads or processes used for parallel processing.


Examples:

    >>> from lazyllm import TrainableModule
    >>> from lazyllm.tools.rag import LLMParser
    >>> llm = TrainableModule("internlm2-chat-7b")
    >>> summary_parser = LLMParser(llm, language="en", task_type="summary")
    """
    def __init__(self, llm: TrainableModule, language: str, task_type: str, num_workers: int = 30):
        super(__class__, self).__init__(num_workers=num_workers)
        assert language in ['en', 'zh'], f'Not supported language {language}'
        assert task_type in ['summary', 'keywords', 'qa', 'qa_img'], f'Not supported task_type {task_type}'
        self._task_type = task_type
        if self._task_type == 'qa_img':
            prompt = dict(system=templates[language][task_type], user='{input}')
        else:
            prompt = dict(system=templates[language][task_type], user='#input:\n{input}\n#output:\n')
        self._llm = llm.share(prompt=AlpacaPrompter(prompt), stream=False, format=self._format)
        self._task_type = task_type

    def transform(self, node: DocNode, **kwargs) -> List[str]:
        """
Perform the set task on the specified document.

Args:
    node (DocNode): The document on which the extraction task needs to be performed.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import LLMParser
    >>> llm = lazyllm.TrainableModule("internlm2-chat-7b").start()
    >>> m = lazyllm.TrainableModule("bge-large-zh-v1.5").start()
    >>> summary_parser = LLMParser(llm, language="en", task_type="summary")
    >>> keywords_parser = LLMParser(llm, language="en", task_type="keywords")
    >>> documents = lazyllm.Document(dataset_path="/path/to/your/data", embed=m, manager=False)
    >>> rm = lazyllm.Retriever(documents, group_name='CoarseChunk', similarity='bm25', topk=6)
    >>> doc_nodes = rm("test")
    >>> summary_result = summary_parser.transform(doc_nodes[0])
    >>> keywords_result = keywords_parser.transform(doc_nodes[0])
    """
        if self._task_type == 'qa_img':
            inputs = encode_query_with_filepaths('Extract QA pairs from images.', [node.image_path])
        else:
            inputs = node.get_text()
        result = self._llm(inputs)
        return [result] if isinstance(result, str) else result

    def _format(self, input):
        if self._task_type == 'keywords':
            return [s.strip() for s in input.split(',')]
        elif self._task_type in ('qa', 'qa_img'):
            return [QADocNode(query=q.strip()[3:].strip(), answer=a.strip()[3:].strip()) for q, a in zip(
                list(filter(None, map(str.strip, input.split("\n"))))[::2],
                list(filter(None, map(str.strip, input.split("\n"))))[1::2])]
        return input

transform(node, **kwargs)

Perform the set task on the specified document.

Parameters:

  • node (DocNode) –

    The document on which the extraction task needs to be performed.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import LLMParser
>>> llm = lazyllm.TrainableModule("internlm2-chat-7b").start()
>>> m = lazyllm.TrainableModule("bge-large-zh-v1.5").start()
>>> summary_parser = LLMParser(llm, language="en", task_type="summary")
>>> keywords_parser = LLMParser(llm, language="en", task_type="keywords")
>>> documents = lazyllm.Document(dataset_path="/path/to/your/data", embed=m, manager=False)
>>> rm = lazyllm.Retriever(documents, group_name='CoarseChunk', similarity='bm25', topk=6)
>>> doc_nodes = rm("test")
>>> summary_result = summary_parser.transform(doc_nodes[0])
>>> keywords_result = keywords_parser.transform(doc_nodes[0])
Source code in lazyllm/tools/rag/transform.py
    def transform(self, node: DocNode, **kwargs) -> List[str]:
        """
Perform the set task on the specified document.

Args:
    node (DocNode): The document on which the extraction task needs to be performed.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools import LLMParser
    >>> llm = lazyllm.TrainableModule("internlm2-chat-7b").start()
    >>> m = lazyllm.TrainableModule("bge-large-zh-v1.5").start()
    >>> summary_parser = LLMParser(llm, language="en", task_type="summary")
    >>> keywords_parser = LLMParser(llm, language="en", task_type="keywords")
    >>> documents = lazyllm.Document(dataset_path="/path/to/your/data", embed=m, manager=False)
    >>> rm = lazyllm.Retriever(documents, group_name='CoarseChunk', similarity='bm25', topk=6)
    >>> doc_nodes = rm("test")
    >>> summary_result = summary_parser.transform(doc_nodes[0])
    >>> keywords_result = keywords_parser.transform(doc_nodes[0])
    """
        if self._task_type == 'qa_img':
            inputs = encode_query_with_filepaths('Extract QA pairs from images.', [node.image_path])
        else:
            inputs = node.get_text()
        result = self._llm(inputs)
        return [result] if isinstance(result, str) else result

lazyllm.tools.rag.transform.NodeTransform members: exclude-members:

lazyllm.tools.rag.transform.TransformArgs dataclass

A document transformation parameter container for centralized management of processing configurations. Args: f (Union[str, Callable]): Transformation function or registered function name.Can be either a callable function or a string identifier for registered functions. trans_node (bool): Whether to transform node types.When True, modifies the document node structure during processing. num_workers (int):Controls parallel processing threads.Values >0. kwargs (Dict):Additional parameters passed to the transformation function. pattern (Union[str, Callable[[str], bool]]):File name/content matching pattern.

Examples:

>>> from lazyllm.tools import TransformArgs
>>> args = TransformArgs(f=lambda text: text.lower(),num_workers=4,pattern=r'.*\.md$')
>>>config = {'f': 'parse_pdf','kwargs': {'engine': 'pdfminer'},'trans_node': True}
>>>args = TransformArgs.from_dict(config)
print(args['f'])
print(args.get('unknown'))
Source code in lazyllm/tools/rag/transform.py
@dataclass
class TransformArgs():
    """
A document transformation parameter container for centralized management of processing configurations.
Args:
    f (Union[str, Callable]): Transformation function or registered function name.Can be either a callable function or a string identifier for registered functions.
    trans_node (bool): Whether to transform node types.When True, modifies the document node structure during processing.
    num_workers (int):Controls parallel processing threads.Values >0.
    kwargs (Dict):Additional parameters passed to the transformation function.
    pattern (Union[str, Callable[[str], bool]]):File name/content matching pattern.


Examples:

    >>> from lazyllm.tools import TransformArgs
    >>> args = TransformArgs(f=lambda text: text.lower(),num_workers=4,pattern=r'.*\.md$')
    >>>config = {'f': 'parse_pdf','kwargs': {'engine': 'pdfminer'},'trans_node': True}
    >>>args = TransformArgs.from_dict(config)
    print(args['f'])
    print(args.get('unknown'))
    """
    f: Union[str, Callable]
    trans_node: Optional[bool] = None
    num_workers: int = 0
    kwargs: Dict = field(default_factory=dict)
    pattern: Optional[Union[str, Callable[[str], bool]]] = None

    @staticmethod
    def from_dict(d):
        return TransformArgs(f=d['f'], trans_node=d.get('trans_node'), num_workers=d.get(
            'num_workers', 0), kwargs=d.get('kwargs', dict()), pattern=d.get('pattern'))

    def __getitem__(self, key):
        if key in self.__dict__: return getattr(self, key)
        raise KeyError(f'Key {key} is not found in transform args')

    def get(self, key):
        if key in self.__dict__: return getattr(self, key)
        return None

lazyllm.tools.rag.similarity.register_similarity(func=None, mode=None, descend=True, batch=False)

Similarity computation registration decorator, used for unified registration and management of different types of similarity computation methods. Args: func (Callable): The name of the similarity computation function. mode (Literal['text', 'embedding']): 'text' indicates direct text matching, while 'embedding' indicates vector-based similarity computation. descend (bool): Controls whether multithreading is enabled (enabled when > 0). kwargs (Dict): Whether the results are sorted in descending order of similarity. batch (bool): Whether to process nodes in batch.

Source code in lazyllm/tools/rag/similarity.py
def register_similarity(
    func: Optional[Callable] = None,
    mode: Optional[Literal['text', 'embedding']] = None,
    descend: bool = True,
    batch: bool = False,
) -> Callable:
    """
Similarity computation registration decorator, used for unified registration and management of different types of similarity computation methods.
Args:
    func (Callable): The name of the similarity computation function.
    mode (Literal['text', 'embedding']): 'text' indicates direct text matching, while 'embedding' indicates vector-based similarity computation.
    descend (bool): Controls whether multithreading is enabled (enabled when > 0).
    kwargs (Dict): Whether the results are sorted in descending order of similarity.
    batch (bool): Whether to process nodes in batch.
"""
    def decorator(f):
        @functools.wraps(f)
        def wrapper(query, nodes, **kwargs):
            if mode != "embedding":
                if batch:
                    return f(query, nodes, **kwargs)
                else:
                    return [(node, f(query, node, **kwargs)) for node in nodes]
            else:
                assert isinstance(query, dict), "query must be of dict type, used for similarity calculation."
                similarity = {}
                if batch:
                    for key, val in query.items():
                        nodes_embed = [node.embedding[key] for node in nodes]
                        similarity[key] = [(node, sim) for node, sim in zip(nodes, f(val, nodes_embed, **kwargs))]
                else:
                    for key, val in query.items():
                        similarity[key] = [(node, f(val, node.embedding[key], **kwargs)) for node in nodes]
                return similarity
        registered_similarities[f.__name__] = (wrapper, mode, descend)
        return wrapper

    return decorator(func) if func else decorator

lazyllm.tools.rag.doc_node.DocNode

Execute assigned tasks on the specified document. Args: uid (str): Unique identifier. content (Union[str, List[Any]]): Node content. group (str): Document group name. embedding (Dict[str, List[float]]): Dictionary of embedding vectors. parent (Union[str, "DocNode"]): Reference to the parent node. store: Storage representation. node_groups (Dict[str, Dict]): Node storage groups. metadata (Dict[str, Any]): Node-level metadata. global_metadata (Dict[str, Any]): Document-level metadata. text (str): Node content, mutually exclusive with content.

Source code in lazyllm/tools/rag/doc_node.py
@reset_on_pickle(('_lock', threading.Lock))
class DocNode:
    """
Execute assigned tasks on the specified document.
Args:
    uid (str): Unique identifier.
    content (Union[str, List[Any]]): Node content.
    group (str): Document group name.
    embedding (Dict[str, List[float]]): Dictionary of embedding vectors.
    parent (Union[str, "DocNode"]): Reference to the parent node.
    store: Storage representation.
    node_groups (Dict[str, Dict]): Node storage groups.
    metadata (Dict[str, Any]): Node-level metadata.
    global_metadata (Dict[str, Any]): Document-level metadata.
    text (str): Node content, mutually exclusive with content.
"""
    def __init__(self, uid: Optional[str] = None, content: Optional[Union[str, List[Any]]] = None,
                 group: Optional[str] = None, embedding: Optional[Dict[str, List[float]]] = None,
                 parent: Optional[Union[str, "DocNode"]] = None, store=None,
                 node_groups: Optional[Dict[str, Dict]] = None, metadata: Optional[Dict[str, Any]] = None,
                 global_metadata: Optional[Dict[str, Any]] = None, *, text: Optional[str] = None):
        if text and content:
            raise ValueError('`text` and `content` cannot be set at the same time.')
        if not content and not text: content = ''
        self._uid: str = uid if uid else str(uuid.uuid4())
        self._content: Optional[Union[str, List[Any]]] = content if content is not None else text
        self._group: Optional[str] = group
        self._embedding: Optional[Dict[str, List[float]]] = embedding or {}
        # metadata: the chunk's meta
        self._metadata: Dict[str, Any] = metadata or {}
        # Global metadata: the file's global metadata (higher level)
        self._global_metadata = global_metadata or {}
        # Metadata keys that are excluded from text for the embed model.
        self._excluded_embed_metadata_keys: List[str] = []
        # Metadata keys that are excluded from text for the LLM.
        self._excluded_llm_metadata_keys: List[str] = []
        # NOTE: node in parent should be id when stored in db (use store to recover): parent: 'uid'
        self._parent: Optional[Union[str, "DocNode"]] = parent
        self._children: Dict[str, List["DocNode"]] = defaultdict(list)
        self._children_loaded = False
        self._store = store
        self._node_groups: Dict[str, Dict] = node_groups or {}
        self._lock = threading.Lock()
        self._embedding_state = set()
        self.relevance_score = None
        self.similarity_score = None

    @property
    def uid(self) -> str:
        return self._uid

    @property
    def group(self) -> str:
        return self._group

    @property
    def text(self) -> str:
        if isinstance(self._content, str):
            return self._content
        elif isinstance(self._content, list):
            if unexcepted := set([type(ele) for ele in self._content if not isinstance(ele, str)]):
                raise TypeError(f"Found non-string element in content: {unexcepted}")
            return '\n'.join(self._content)
        else:
            raise TypeError(f"content type '{type(self._content)}' is neither a str nor a list")

    @property
    def embedding(self):
        return self._embedding

    @embedding.setter
    def embedding(self, v: Optional[Dict[str, List[float]]]):
        self._embedding = v

    def _load_from_store(self, group_name: str, uids: Union[str, List[str]]) -> List["DocNode"]:
        if not self._store or not uids:
            return []
        if isinstance(uids, str):
            uids = [uids]
        nodes = self._store.get_nodes(group_name=group_name, uids=uids,
                                      kb_id=self.global_metadata.get(RAG_KB_ID), display=True)
        for n in nodes:
            n._store = self._store
            n._node_groups = self._node_groups
        return nodes

    @property
    def parent(self) -> Optional["DocNode"]:
        if self._parent and isinstance(self._parent, str) and self._node_groups:
            parent_group = self._node_groups[self._group]["parent"]
            loaded = self._load_from_store(parent_group, self._parent)
            self._parent = loaded[0] if loaded else None
        return self._parent

    @parent.setter
    def parent(self, v: Optional["DocNode"]):
        self._parent = v

    @property
    def children(self) -> Dict[str, List["DocNode"]]:
        if not self._children_loaded and self._store and self._node_groups:
            self._children_loaded = True
            kb_id = self.global_metadata.get(RAG_KB_ID)
            doc_id = self.global_metadata.get(RAG_DOC_ID)
            c_groups = [grp for grp in self._node_groups.keys() if self._node_groups[grp]['parent'] == self._group]
            for grp in c_groups:
                if not self._store.is_group_active(grp):
                    continue
                nodes = self._store.get_nodes(group_name=grp, kb_id=kb_id, doc_ids=[doc_id])
                c_nodes = [n for n in nodes if n._parent in {self, self._uid}]
                self._children[grp] = c_nodes
                for n in self._children[grp]:
                    n._store = self._store
                    n._node_groups = self._node_groups
        return self._children

    @children.setter
    def children(self, v: Dict[str, List["DocNode"]]):
        self._children = v

    @property
    def root_node(self) -> "DocNode":
        node = self
        while isinstance(node._parent, DocNode):
            node = node._parent
        return node

    @property
    def is_root_node(self) -> bool:
        return (not self.parent)

    @property
    def global_metadata(self) -> Dict[str, Any]:
        return self.root_node._global_metadata

    @global_metadata.setter
    def global_metadata(self, global_metadata: Dict) -> None:
        self._global_metadata = global_metadata

    @property
    def metadata(self) -> Dict:
        return self._metadata

    @metadata.setter
    def metadata(self, metadata: Dict) -> None:
        self._metadata = metadata

    @property
    def excluded_embed_metadata_keys(self) -> List:
        return list(set(self.root_node._excluded_embed_metadata_keys + self._excluded_embed_metadata_keys))

    @excluded_embed_metadata_keys.setter
    def excluded_embed_metadata_keys(self, excluded_embed_metadata_keys: List) -> None:
        self._excluded_embed_metadata_keys = excluded_embed_metadata_keys

    @property
    def excluded_llm_metadata_keys(self) -> List:
        return list(set(self.root_node._excluded_llm_metadata_keys + self._excluded_llm_metadata_keys))

    @excluded_llm_metadata_keys.setter
    def excluded_llm_metadata_keys(self, excluded_llm_metadata_keys: List) -> None:
        self._excluded_llm_metadata_keys = excluded_llm_metadata_keys

    @property
    def docpath(self) -> str:
        return self.root_node.global_metadata.get(RAG_DOC_PATH, '')

    @docpath.setter
    def docpath(self, path):
        assert not self.parent, 'Only root node can set docpath'
        self.global_metadata[RAG_DOC_PATH] = str(path)

    def get_children_str(self) -> str:
        return str(
            {key: [node._uid for node in nodes] for key, nodes in self.children.items()}
        )

    def get_parent_id(self) -> str:
        return self.parent._uid if self.parent else ''

    def __str__(self) -> str:
        return (
            f"DocNode(id: {self._uid}, group: {self._group}, content: {self._content}) parent: {self.get_parent_id()}, "
            f"children: {self.get_children_str()}"
        )

    def __repr__(self) -> str:
        return str(self) if config["debug"] else f'<Node id={self._uid}>'

    def __eq__(self, other):
        if isinstance(other, DocNode):
            return self._uid == other._uid
        return False

    def __hash__(self):
        return hash(self._uid)

    def __getstate__(self):
        st = self.__dict__.copy()
        for attr in _pickle_blacklist:
            st[attr] = None
        return st

    def has_missing_embedding(self, embed_keys: Union[str, List[str]]) -> List[str]:
        """
Check for missing embedding vectors.
Args:
    embed_keys (Union[str, List[str]]): List of target keys.
"""
        if isinstance(embed_keys, str): embed_keys = [embed_keys]
        assert len(embed_keys) > 0, "The ebmed_keys to be checked must be passed in."
        if self.embedding is None: return embed_keys
        return [k for k in embed_keys if k not in self.embedding]

    def do_embedding(self, embed: Dict[str, Callable]) -> None:
        """
Execute embedding computation.
Args:
    embed (Dict[str, Callable]): Target embedding objects.
"""
        generate_embed = {k: e(self.get_text(MetadataMode.EMBED)) for k, e in embed.items()}
        with self._lock:
            self.embedding = self.embedding or {}
            self.embedding = {**self.embedding, **generate_embed}

    def check_embedding_state(self, embed_key: str) -> None:
        """
Block to check the embedding status and ensure that asynchronous embedding computation is completed.
Args:
    embed_key (str): List of target keys.
"""
        while True:
            with self._lock:
                if not self.has_missing_embedding(embed_key):
                    self._embedding_state.discard(embed_key)
                    break
            time.sleep(1)

    def get_content(self) -> str:
        return self.get_text(MetadataMode.LLM)

    def get_metadata_str(self, mode: MetadataMode = MetadataMode.ALL) -> str:
        """Metadata info string."""
        if mode == MetadataMode.NONE:
            return ''

        metadata_keys = set(self.metadata.keys())
        if mode == MetadataMode.LLM:
            for key in self.excluded_llm_metadata_keys:
                if key in metadata_keys:
                    metadata_keys.remove(key)
        elif mode == MetadataMode.EMBED:
            for key in self.excluded_embed_metadata_keys:
                if key in metadata_keys:
                    metadata_keys.remove(key)

        return "\n".join([f"{key}: {self.metadata[key]}" for key in metadata_keys])

    def get_text(self, metadata_mode: MetadataMode = MetadataMode.NONE) -> str:
        """
Combine metadata and content.
Args:
    metadata_mode: Same as the parameter in get_metadata_str.
"""
        metadata_str = self.get_metadata_str(metadata_mode).strip()
        if not metadata_str:
            return self.text if self.text else ''
        return f"{metadata_str}\n\n{self.text}".strip()

    def to_dict(self) -> Dict:
        """
Convert to dictionary format
"""
        return dict(content=self._content, embedding=self.embedding, metadata=self.metadata)

    def with_score(self, score):
        """
Shallow copy the original node and add a semantic relevance score.
Args:
    score: Relevance score.
"""
        node = copy.copy(self)
        node.relevance_score = score
        return node

    def with_sim_score(self, score):
        """
Shallow copy the original node and add a similarity score.
Args:
    score: Similarity score.
"""
        node = copy.copy(self)
        node.similarity_score = score
        return node

check_embedding_state(embed_key)

Block to check the embedding status and ensure that asynchronous embedding computation is completed. Args: embed_key (str): List of target keys.

Source code in lazyllm/tools/rag/doc_node.py
    def check_embedding_state(self, embed_key: str) -> None:
        """
Block to check the embedding status and ensure that asynchronous embedding computation is completed.
Args:
    embed_key (str): List of target keys.
"""
        while True:
            with self._lock:
                if not self.has_missing_embedding(embed_key):
                    self._embedding_state.discard(embed_key)
                    break
            time.sleep(1)

do_embedding(embed)

Execute embedding computation. Args: embed (Dict[str, Callable]): Target embedding objects.

Source code in lazyllm/tools/rag/doc_node.py
    def do_embedding(self, embed: Dict[str, Callable]) -> None:
        """
Execute embedding computation.
Args:
    embed (Dict[str, Callable]): Target embedding objects.
"""
        generate_embed = {k: e(self.get_text(MetadataMode.EMBED)) for k, e in embed.items()}
        with self._lock:
            self.embedding = self.embedding or {}
            self.embedding = {**self.embedding, **generate_embed}

get_metadata_str(mode=MetadataMode.ALL)

Metadata info string.

Source code in lazyllm/tools/rag/doc_node.py
def get_metadata_str(self, mode: MetadataMode = MetadataMode.ALL) -> str:
    """Metadata info string."""
    if mode == MetadataMode.NONE:
        return ''

    metadata_keys = set(self.metadata.keys())
    if mode == MetadataMode.LLM:
        for key in self.excluded_llm_metadata_keys:
            if key in metadata_keys:
                metadata_keys.remove(key)
    elif mode == MetadataMode.EMBED:
        for key in self.excluded_embed_metadata_keys:
            if key in metadata_keys:
                metadata_keys.remove(key)

    return "\n".join([f"{key}: {self.metadata[key]}" for key in metadata_keys])

get_text(metadata_mode=MetadataMode.NONE)

Combine metadata and content. Args: metadata_mode: Same as the parameter in get_metadata_str.

Source code in lazyllm/tools/rag/doc_node.py
    def get_text(self, metadata_mode: MetadataMode = MetadataMode.NONE) -> str:
        """
Combine metadata and content.
Args:
    metadata_mode: Same as the parameter in get_metadata_str.
"""
        metadata_str = self.get_metadata_str(metadata_mode).strip()
        if not metadata_str:
            return self.text if self.text else ''
        return f"{metadata_str}\n\n{self.text}".strip()

has_missing_embedding(embed_keys)

Check for missing embedding vectors. Args: embed_keys (Union[str, List[str]]): List of target keys.

Source code in lazyllm/tools/rag/doc_node.py
    def has_missing_embedding(self, embed_keys: Union[str, List[str]]) -> List[str]:
        """
Check for missing embedding vectors.
Args:
    embed_keys (Union[str, List[str]]): List of target keys.
"""
        if isinstance(embed_keys, str): embed_keys = [embed_keys]
        assert len(embed_keys) > 0, "The ebmed_keys to be checked must be passed in."
        if self.embedding is None: return embed_keys
        return [k for k in embed_keys if k not in self.embedding]

to_dict()

Convert to dictionary format

Source code in lazyllm/tools/rag/doc_node.py
    def to_dict(self) -> Dict:
        """
Convert to dictionary format
"""
        return dict(content=self._content, embedding=self.embedding, metadata=self.metadata)

with_score(score)

Shallow copy the original node and add a semantic relevance score. Args: score: Relevance score.

Source code in lazyllm/tools/rag/doc_node.py
    def with_score(self, score):
        """
Shallow copy the original node and add a semantic relevance score.
Args:
    score: Relevance score.
"""
        node = copy.copy(self)
        node.relevance_score = score
        return node

with_sim_score(score)

Shallow copy the original node and add a similarity score. Args: score: Similarity score.

Source code in lazyllm/tools/rag/doc_node.py
    def with_sim_score(self, score):
        """
Shallow copy the original node and add a similarity score.
Args:
    score: Similarity score.
"""
        node = copy.copy(self)
        node.similarity_score = score
        return node

lazyllm.tools.rag.doc_processor.DocumentProcessor

Bases: ModuleBase

Document processor class for managing document addition, deletion and update operations.

Parameters:

  • server (bool, default: True ) –

    Whether to run in server mode. Defaults to True.

  • port (Optional[int], default: None ) –

    Server port number. Defaults to None.

  • url (Optional[str], default: None ) –

    Remote service URL. Defaults to None.

Notes: - Supports asynchronous document task processing - Provides document metadata update functionality - Supports task status callback notifications - Configurable database storage

Examples:

```python
# Create local document processor
processor = DocumentProcessor(server=False)

# Create server mode document processor
processor = DocumentProcessor(server=True, port=8080)

# Create remote document processor
processor = DocumentProcessor(url="http://remote-server:8080")
```
Source code in lazyllm/tools/rag/doc_processor.py
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
class DocumentProcessor(ModuleBase):
    """
Document processor class for managing document addition, deletion and update operations.

Args:
    server (bool): Whether to run in server mode. Defaults to True.
    port (Optional[int]): Server port number. Defaults to None.
    url (Optional[str]): Remote service URL. Defaults to None.

**Notes:**
- Supports asynchronous document task processing
- Provides document metadata update functionality
- Supports task status callback notifications
- Configurable database storage


Examples:

    ```python
    # Create local document processor
    processor = DocumentProcessor(server=False)

    # Create server mode document processor
    processor = DocumentProcessor(server=True, port=8080)

    # Create remote document processor
    processor = DocumentProcessor(url="http://remote-server:8080")
    ```
    """

    class Impl():
        def __init__(self, server: bool):
            self._processors: Dict[str, _Processor] = dict()
            self._server = server
            self._inited = False
            try:
                self._feedback_url = config['process_feedback_service']
                self._path_prefix = config['process_path_prefix']
            except Exception as e:
                LOG.warning(f"Failed to get config: {e}, use env variables instead")
                self._feedback_url = os.getenv("PROCESS_FEEDBACK_SERVICE", None)
                self._path_prefix = os.getenv("PROCESS_PATH_PREFIX", None)

        def _init_components(self, server: bool):
            if server and not self._inited:
                self._task_queue = queue.Queue()
                self._tasks = {}    # running tasks
                self._pending_task_ids = set()  # pending tasks
                self._add_executor = ThreadPoolExecutor(max_workers=4)
                self._add_futures = {}
                self._delete_executor = ThreadPoolExecutor(max_workers=4)
                self._update_executor = ThreadPoolExecutor(max_workers=4)
                self._update_futures = {}

                self._engines: dict[str, Engine] = {}
                self._inspectors: dict[str, inspect] = {}

                self._worker_thread = threading.Thread(target=self._worker, daemon=True)
                self._worker_thread.start()
            self._inited = True
            LOG.info(f"[DocumentProcessor] init done. feedback {self._feedback_url}, prefix {self._path_prefix}")

        def register_algorithm(self, name: str, store: _DocumentStore, reader: ReaderBase,
                               node_groups: Dict[str, Dict], display_name: Optional[str] = None,
                               description: Optional[str] = None, force_refresh: bool = False):
            self._init_components(server=self._server)
            if name in self._processors and not force_refresh:
                LOG.warning(f'There is already a processor with the same name {name}!')
                return
            self._processors[name] = _Processor(store, reader, node_groups, display_name, description)
            LOG.info(f'Processor {name} registered!')

        def drop_algorithm(self, name: str, clean_db: bool = False) -> None:
            if name not in self._processors:
                LOG.warning(f'Processor {name} not found!')
                return
            self._processors.pop(name)

        def _get_engine(self, url) -> Engine:
            if url not in self._engines:
                engine = create_engine(url, echo=False, pool_pre_ping=True)
                self._engines[url] = engine
                self._inspectors[url] = inspect(engine)
            return self._engines[url]

        def _get_inspector(self, url):
            self._get_engine(url=url)
            return self._inspectors[url]

        def _get_url_from_db_info(self, db_info: DBInfo):
            return (f"mysql+pymysql://{db_info.user}:{db_info.password}"
                    f"@{db_info.host}:{db_info.port}/{db_info.db_name}"
                    "?charset=utf8mb4")

        def create_table(self, db_info: DBInfo):
            if db_info.db_type == "mysql":
                try:
                    url = self._get_url_from_db_info(db_info)
                    engine = self._get_engine(url=url)
                    inspector = self._get_inspector(url=url)
                    tbl = db_info.table_name
                    schema = db_info.db_name

                    if not inspector.has_table(tbl, schema=schema):
                        metadata = MetaData()
                        table = Table(tbl, metadata, Column('document_id', String(255), primary_key=True),
                                      Column('file_name', String(255), nullable=False),
                                      Column('file_path', String(255), nullable=False),
                                      Column('description', String(255), nullable=True),
                                      Column('creater', String(255), nullable=False),
                                      Column('dataset_id', String(255), nullable=False),
                                      Column('tags', JSON, nullable=True),
                                      Column('created_at', TIMESTAMP, server_default=text("CURRENT_TIMESTAMP")))
                        metadata.create_all(engine, tables=[table])
                        LOG.info(f"Created table `{tbl}` in `{schema}`")
                except Exception as e:
                    LOG.error(f"Failed to create table `{tbl}` in `{schema}`: {e}")
                    return
            else:
                raise ValueError(f"Unsupported database type: {db_info.db_type}")

        def operate_db(self, db_info: DBInfo, operation: str,
                       file_infos: List[FileInfo] = None, params: Dict = None) -> None:
            db_type = db_info.db_type
            if db_type not in DB_TYPES:
                raise ValueError(f"Unsupported db_type: {db_type}")
            url = self._get_url_from_db_info(db_info)
            engine = self._get_engine(url=url)
            if operation == 'upsert':
                self._upsert_records(engine, db_info, file_infos)
            elif operation == 'delete':
                self._delete_records(engine, db_info, params)
            else:
                raise ValueError(f"Unsupported operation: {operation}")

        def _upsert_records(self, engine, db_info, file_infos):
            table_name = db_info['table_name']
            metadata = MetaData()
            metadata.reflect(bind=engine, only=[table_name])
            table = metadata.tables[table_name]
            with engine.begin() as conn:
                for file_info in file_infos:
                    document_id = file_info.get("doc_id")
                    file_path = file_info.get("file_path")
                    if not document_id or not file_path:
                        raise ValueError(f"Invalid file_info: {file_info}")

                    raw_infos = {"document_id": document_id, "file_name": os.path.basename(file_path),
                                 "file_path": file_path, "description": file_info["metadata"].get("description", None),
                                 "creater": file_info["metadata"].get("creater", None),
                                 "dataset_id": file_info["metadata"].get(RAG_KB_ID, None),
                                 "tags": file_info["metadata"].get("tags", []) or []}
                    infos = {}
                    for k, v in raw_infos.items():
                        if v is None:
                            continue
                        if isinstance(v, str) and not v.strip():
                            continue
                        if isinstance(v, (list, dict)) and not v:
                            continue
                        infos[k] = v
                    if "document_id" not in infos:
                        infos["document_id"] = document_id

                    stmt = mysql_insert(table).values(**infos)
                    update_dict = {k: stmt.inserted[k] for k in infos if k != 'document_id'}
                    upsert_stmt = stmt.on_duplicate_key_update(**update_dict)
                    conn.execute(upsert_stmt)

        def _delete_records(self, engine, db_info, params):
            table_name = db_info['table_name']
            metadata = MetaData()
            metadata.reflect(bind=engine, only=[table_name])
            table = metadata.tables[table_name]

            with engine.begin() as conn:  # 自动提交或回滚事务
                doc_ids = params.get("doc_ids", [])
                for document_id in doc_ids:
                    stmt = delete(table).where(table.c.document_id == document_id)
                    conn.execute(stmt)

        @app.get('/algo/list')
        async def get_algo_list(self) -> None:
            res = []
            for algo_id, processor in self._processors.items():
                res.append({"algo_id": algo_id, "display_name": processor._display_name,
                            "description": processor._description})
            return BaseResponse(code=200, msg='success', data=res)

        @app.get('/group/info')
        async def get_group_info(self, algo_id: str) -> None:
            if algo_id not in self._processors:
                return BaseResponse(code=400, msg=f"Invalid algo_id {algo_id}")
            processor = self._processors[algo_id]
            infos = []
            for group_name in processor._store.activated_groups():
                if group_name in processor._node_groups:
                    group_info = {"name": group_name, "type": processor._node_groups[group_name].get('group_type'),
                                  "display_name": processor._node_groups[group_name].get('display_name')}
                    infos.append(group_info)
            LOG.info(f"Get group info for {algo_id} success with {infos}")
            return BaseResponse(code=200, msg='success', data=infos)

        @app.post('/doc/add')
        async def async_add_doc(self, request: AddDocRequest):
            LOG.info(f"Add doc for {request.model_dump_json()}")
            task_id = request.task_id
            algo_id = request.algo_id
            file_infos = request.file_infos
            db_info = request.db_info
            feedback_url = request.feedback_url
            if algo_id not in self._processors:
                return BaseResponse(code=400, msg=f"Invalid algo_id {algo_id}")
            if task_id in self._pending_task_ids or task_id in self._tasks:
                return BaseResponse(code=400, msg=f'The task {task_id} already exists in queue', data=None)
            if self._path_prefix:
                for file_info in file_infos:
                    file_info.file_path = create_file_path(path=file_info.file_path, prefix=self._path_prefix)

            params = {"file_infos": file_infos, "db_info": db_info, "feedback_url": feedback_url}
            if ENABLE_DB:
                self.create_table(db_info=db_info)

            self._task_queue.put(('add', algo_id, task_id, params))
            self._pending_task_ids.add(task_id)
            return BaseResponse(code=200, msg='task submit successfully', data={"task_id": task_id})

        @app.post('/doc/meta/update')
        async def async_update_meta(self, request: UpdateMetaRequest):
            LOG.info(f"update doc meta for {request.model_dump_json()}")
            algo_id = request.algo_id
            file_infos = request.file_infos
            db_info = request.db_info

            if algo_id not in self._processors:
                return BaseResponse(code=400, msg=f"Invalid algo_id {algo_id}")

            for file_info in file_infos:
                doc_id = file_info.doc_id
                metadata = file_info.metadata
                old_fut = self._update_futures.get(doc_id)
                if old_fut and not old_fut.done():
                    cancelled = old_fut.cancel()
                    LOG.info(f"Canceled previous update for {doc_id}: {cancelled}")

                new_fut = self._update_executor.submit(self._processors[algo_id].update_doc_meta, doc_id=doc_id,
                                                       metadata=metadata)

                self._update_futures[doc_id] = new_fut

                def _cleanup(fut, doc_id=doc_id):
                    if self._update_futures.get(doc_id) is fut:
                        del self._update_futures[doc_id]
                new_fut.add_done_callback(_cleanup)
                if ENABLE_DB:
                    new_fut.add_done_callback(
                        lambda fut, dbi=db_info, fi=file_info: self.operate_db(dbi, 'upsert', file_infos=[fi]))

            return BaseResponse(code=200, msg='success')

        @app.delete('/doc/delete')
        async def async_delete_doc(self, request: DeleteDocRequest) -> None:
            LOG.info(f"Del doc for {request.model_dump_json()}")
            algo_id = request.algo_id
            dataset_id = request.dataset_id
            doc_ids = request.doc_ids
            db_info = request.db_info

            if algo_id not in self._processors:
                return BaseResponse(code=400, msg=f"Invalid algo_id {algo_id}")

            task_id = str(uuid.uuid4())
            self._task_queue.put(('delete', algo_id, task_id,
                                  {"dataset_id": dataset_id, "doc_ids": doc_ids, "db_info": db_info}))
            self._pending_task_ids.add(task_id)
            return BaseResponse(code=200, msg='task submit successfully', data={"task_id": task_id})

        @app.post('/doc/cancel')
        async def cancel_task(self, request: CancelDocRequest):
            task_id = request.task_id
            if task_id in self._pending_task_ids:
                self._pending_task_ids.remove(task_id)
                status = 1
            elif task_id in self._tasks:
                future = self._tasks.get(task_id)
                if future and not future.done():
                    cancelled = future.cancel()
                    status = 1 if cancelled else 0
                    if cancelled:
                        self._tasks.pop(task_id, None)
                else:
                    status = 0
            return BaseResponse(code=200, msg="success" if status else "failed",
                                data={"task_id": task_id, "status": status})

        def _send_status_message(self, task_id: str, callback_path: str, success: bool,
                                 error_code: str = "", error_msg: str = ""):
            if self._feedback_url:
                try:
                    full_url = self._feedback_url + callback_path
                    payload = {"task_id": task_id, "status": 1 if success else 0, "error_code": error_code,
                               "error_msg": error_msg}
                    headers = {"Content-Type": "application/json"}
                    res = None
                    for wait_time in fibonacci_backoff(max_retries=3):
                        try:
                            res = requests.post(full_url, json=payload, headers=headers, timeout=5)
                            if res.status_code == 200:
                                break
                            LOG.warning(
                                f"Task-{task_id}: Unexpected status {res.status_code}, retrying in {wait_time}s…")
                        except Exception as e:
                            LOG.error(f"Task-{task_id}: Request failed: {e}, retrying in {wait_time}s…")
                        time.sleep(wait_time)

                    if res is None:
                        raise RuntimeError("Failed to send feedback—no response received after retries")
                    res.raise_for_status()
                except Exception as e:
                    LOG.error(f"Task-{task_id}: Failed to send feedback to {full_url}: {e}")
            else:
                LOG.error("process_feedback_service is not set")

        def _exec_add_task(self, algo_id, task_id, params):
            try:
                file_infos: List[FileInfo] = params.get('file_infos')
                callback_path = params.get('feedback_url')
                db_info: DBInfo = params.get('db_info')

                input_files = []
                ids = []
                metadatas = []

                reparse_group = None
                reparse_doc_ids = []
                reparse_files = []
                reparse_metadatas = []

                for file_info in file_infos:
                    if file_info.reparse_group:
                        reparse_group = file_info.reparse_group
                        reparse_doc_ids.append(file_info.doc_id)
                        reparse_files.append(file_info.file_path)
                        reparse_metadatas.append(file_info.metadata)
                    else:
                        input_files.append(file_info.file_path)
                        ids.append(file_info.doc_id)
                        metadatas.append(file_info.metadata)

                if input_files:
                    future = self._add_executor.submit(self._processors[algo_id].add_doc, input_files=input_files,
                                                       ids=ids, metadatas=metadatas)
                    if ENABLE_DB:
                        future.add_done_callback(lambda fut: self.operate_db(db_info, 'upsert', file_infos=file_infos))
                elif reparse_group:
                    future = self._add_executor.submit(self._processors[algo_id].reparse, group_name=reparse_group,
                                                       doc_ids=reparse_doc_ids, doc_paths=reparse_files,
                                                       metadatas=reparse_metadatas)
                else:
                    LOG.error(
                        f"Task-{task_id}: add task error, no input files {input_files} or reparse group {reparse_group}"
                    )
                self._tasks[task_id] = (future, callback_path)
                self._pending_task_ids.remove(task_id)
            except Exception as e:
                LOG.error(f"Task-{task_id}: add task error {e}")

        def _exec_delete_task(self, algo_id, task_id, params):
            dataset_id = params.get("dataset_id")
            doc_ids = params.get("doc_ids")
            future = self._delete_executor.submit(
                self._processors[algo_id].delete_doc, dataset_id=dataset_id, doc_ids=doc_ids
            )
            if ENABLE_DB and params.get("db_info") is not None:
                db_info = params.get("db_info")
                future.add_done_callback(lambda fut: self.operate_db(db_info, 'delete', params=params))
            self._tasks[task_id] = (future, None)
            self._pending_task_ids.remove(task_id)

        def _worker(self):  # noqa: C901
            while True:
                try:
                    task_type, algo_id, task_id, params = self._task_queue.get(timeout=1)
                    if task_id not in self._pending_task_ids:
                        continue
                    if task_type == 'add':
                        self._exec_add_task(algo_id=algo_id, task_id=task_id, params=params)
                    elif task_type == 'delete':
                        self._exec_delete_task(algo_id=algo_id, task_id=task_id, params=params)
                    time.sleep(0.2)
                except queue.Empty:
                    task_need_pop = []
                    for task_id, (future, callback_path) in self._tasks.items():
                        if future.done():
                            task_need_pop.append(task_id)
                            ex = future.exception()
                            if callback_path and not ex:
                                self._send_status_message(task_id=task_id, callback_path=callback_path, success=True,
                                                          error_code="", error_msg="")
                            elif callback_path and ex:
                                self._send_status_message(task_id=task_id, callback_path=callback_path, success=False,
                                                          error_code=type(ex).__name__, error_msg=str(ex))
                                LOG.error(f"task {task_id} failed: {str(ex)}")
                            elif ex:
                                LOG.error(f"task {task_id} failed: {str(ex)}")
                    for task_id in task_need_pop:
                        self._tasks.pop(task_id)
                        LOG.info(f"task {task_id} done")
                    time.sleep(5)

        def __call__(self, func_name: str, *args, **kwargs):
            return getattr(self, func_name)(*args, **kwargs)

    def __init__(self, server: bool = True, port: int = None, url: str = None):
        super().__init__()
        if not url:
            self._impl = DocumentProcessor.Impl(server=server)
            if server:
                self._impl = ServerModule(self._impl, port=port)
        else:
            self._impl = UrlModule(url=ensure_call_endpoint(url))

    def _dispatch(self, method: str, *args, **kwargs):
        impl = self._impl
        if isinstance(impl, ServerModule):
            impl._call(method, *args, **kwargs)
        else:
            getattr(impl, method)(*args, **kwargs)

    def register_algorithm(self, name: str, store: _DocumentStore, reader: ReaderBase, node_groups: Dict[str, Dict],
                           display_name: Optional[str] = None, description: Optional[str] = None,
                           force_refresh: bool = False, **kwargs):
        """
Register an algorithm to the document processor.

Args:
    name (str): Algorithm name as unique identifier.
    store (StoreBase): Storage instance for managing document data.
    reader (ReaderBase): Reader instance for parsing document content.
    node_groups (Dict[str, Dict]): Node group configuration information.
    force_refresh (bool): Whether to force refresh existing algorithm. Defaults to False.

**Notes:**
- If algorithm name exists and force_refresh is False, registration will be skipped
- After successful registration, the algorithm can be used to process documents


Examples:

    ```python
    from lazyllm.rag import DocumentProcessor, FileStore, PDFReader

    # Create storage and reader instances
    store = FileStore(path="./data")
    reader = PDFReader()

    # Define node group configuration
    node_groups = {
        "text": {"transform": "text", "parent": "root"},
        "summary": {"transform": "summary", "parent": "text"}
    }

    # Register algorithm
    processor = DocumentProcessor()
    processor.register_algorithm(
        name="pdf_processor",
        store=store,
        reader=reader,
        node_groups=node_groups
    )
    ```
    """
        self._dispatch("register_algorithm", name, store, reader, node_groups,
                       display_name, description, force_refresh, **kwargs)

    def drop_algorithm(self, name: str, clean_db: bool = False) -> None:
        """
Remove specified algorithm from document processor.

Args:
    name (str): Name of the algorithm to remove.
    clean_db (bool): Whether to clean related database data. Defaults to False.

**Notes:**
- If algorithm name does not exist, a warning message will be output
- After removal, the algorithm will no longer be available


Examples:

    ```python
    # Remove algorithm
    processor.drop_algorithm("pdf_processor")

    # Remove algorithm and clean database
    processor.drop_algorithm("pdf_processor", clean_db=True)
    ```
    """
        return self._dispatch("drop_algorithm", name, clean_db)

drop_algorithm(name, clean_db=False)

Remove specified algorithm from document processor.

Parameters:

  • name (str) –

    Name of the algorithm to remove.

  • clean_db (bool, default: False ) –

    Whether to clean related database data. Defaults to False.

Notes: - If algorithm name does not exist, a warning message will be output - After removal, the algorithm will no longer be available

Examples:

```python
# Remove algorithm
processor.drop_algorithm("pdf_processor")

# Remove algorithm and clean database
processor.drop_algorithm("pdf_processor", clean_db=True)
```
Source code in lazyllm/tools/rag/doc_processor.py
    def drop_algorithm(self, name: str, clean_db: bool = False) -> None:
        """
Remove specified algorithm from document processor.

Args:
    name (str): Name of the algorithm to remove.
    clean_db (bool): Whether to clean related database data. Defaults to False.

**Notes:**
- If algorithm name does not exist, a warning message will be output
- After removal, the algorithm will no longer be available


Examples:

    ```python
    # Remove algorithm
    processor.drop_algorithm("pdf_processor")

    # Remove algorithm and clean database
    processor.drop_algorithm("pdf_processor", clean_db=True)
    ```
    """
        return self._dispatch("drop_algorithm", name, clean_db)

register_algorithm(name, store, reader, node_groups, display_name=None, description=None, force_refresh=False, **kwargs)

Register an algorithm to the document processor.

Parameters:

  • name (str) –

    Algorithm name as unique identifier.

  • store (StoreBase) –

    Storage instance for managing document data.

  • reader (ReaderBase) –

    Reader instance for parsing document content.

  • node_groups (Dict[str, Dict]) –

    Node group configuration information.

  • force_refresh (bool, default: False ) –

    Whether to force refresh existing algorithm. Defaults to False.

Notes: - If algorithm name exists and force_refresh is False, registration will be skipped - After successful registration, the algorithm can be used to process documents

Examples:

```python
from lazyllm.rag import DocumentProcessor, FileStore, PDFReader

# Create storage and reader instances
store = FileStore(path="./data")
reader = PDFReader()

# Define node group configuration
node_groups = {
    "text": {"transform": "text", "parent": "root"},
    "summary": {"transform": "summary", "parent": "text"}
}

# Register algorithm
processor = DocumentProcessor()
processor.register_algorithm(
    name="pdf_processor",
    store=store,
    reader=reader,
    node_groups=node_groups
)
```
Source code in lazyllm/tools/rag/doc_processor.py
    def register_algorithm(self, name: str, store: _DocumentStore, reader: ReaderBase, node_groups: Dict[str, Dict],
                           display_name: Optional[str] = None, description: Optional[str] = None,
                           force_refresh: bool = False, **kwargs):
        """
Register an algorithm to the document processor.

Args:
    name (str): Algorithm name as unique identifier.
    store (StoreBase): Storage instance for managing document data.
    reader (ReaderBase): Reader instance for parsing document content.
    node_groups (Dict[str, Dict]): Node group configuration information.
    force_refresh (bool): Whether to force refresh existing algorithm. Defaults to False.

**Notes:**
- If algorithm name exists and force_refresh is False, registration will be skipped
- After successful registration, the algorithm can be used to process documents


Examples:

    ```python
    from lazyllm.rag import DocumentProcessor, FileStore, PDFReader

    # Create storage and reader instances
    store = FileStore(path="./data")
    reader = PDFReader()

    # Define node group configuration
    node_groups = {
        "text": {"transform": "text", "parent": "root"},
        "summary": {"transform": "summary", "parent": "text"}
    }

    # Register algorithm
    processor = DocumentProcessor()
    processor.register_algorithm(
        name="pdf_processor",
        store=store,
        reader=reader,
        node_groups=node_groups
    )
    ```
    """
        self._dispatch("register_algorithm", name, store, reader, node_groups,
                       display_name, description, force_refresh, **kwargs)

lazyllm.tools.rag.doc_node.QADocNode

Bases: DocNode

Question-Answer document node class for storing QA pair data.

Parameters:

  • query (str) –

    The question text.

  • answer (str) –

    The answer text.

  • uid (str, default: None ) –

    Unique identifier.

  • group (str, default: None ) –

    Document group name.

  • embedding (Dict[str, List[float]], default: None ) –

    Dictionary of embedding vectors.

  • parent (DocNode, default: None ) –

    Reference to the parent node.

  • metadata (Dict[str, Any], default: None ) –

    Node-level metadata.

  • global_metadata (Dict[str, Any], default: None ) –

    Document-level metadata.

  • text (str, default: None ) –

    Node content, mutually exclusive with query.

Source code in lazyllm/tools/rag/doc_node.py
class QADocNode(DocNode):
    """Question-Answer document node class for storing QA pair data.

Args:
    query (str): The question text.
    answer (str): The answer text.
    uid (str): Unique identifier.
    group (str): Document group name.
    embedding (Dict[str, List[float]]): Dictionary of embedding vectors.
    parent (DocNode): Reference to the parent node.
    metadata (Dict[str, Any]): Node-level metadata.
    global_metadata (Dict[str, Any]): Document-level metadata.
    text (str): Node content, mutually exclusive with query.
"""
    def __init__(self, query: str, answer: str, uid: Optional[str] = None, group: Optional[str] = None,
                 embedding: Optional[Dict[str, List[float]]] = None, parent: Optional["DocNode"] = None,
                 metadata: Optional[Dict[str, Any]] = None, global_metadata: Optional[Dict[str, Any]] = None,
                 *, text: Optional[str] = None):
        super().__init__(uid, query, group, embedding, parent, metadata, global_metadata=global_metadata, text=text)
        self._answer = answer.strip()

    @property
    def answer(self) -> str:
        return self._answer

    def get_text(self, metadata_mode: MetadataMode = MetadataMode.NONE) -> str:
        """Get the text content of the node.

Args:
    metadata_mode (MetadataMode): Metadata mode, defaults to MetadataMode.NONE.
        When set to MetadataMode.LLM, returns formatted QA pair.
        For other modes, returns base class text format.

Returns:
    str: The formatted text content.
"""
        if metadata_mode == MetadataMode.LLM:
            return f'query:\n{self.text}\nanswer\n{self._answer}'
        return super().get_text(metadata_mode)

get_text(metadata_mode=MetadataMode.NONE)

Get the text content of the node.

Parameters:

  • metadata_mode (MetadataMode, default: NONE ) –

    Metadata mode, defaults to MetadataMode.NONE. When set to MetadataMode.LLM, returns formatted QA pair. For other modes, returns base class text format.

Returns:

  • str ( str ) –

    The formatted text content.

Source code in lazyllm/tools/rag/doc_node.py
    def get_text(self, metadata_mode: MetadataMode = MetadataMode.NONE) -> str:
        """Get the text content of the node.

Args:
    metadata_mode (MetadataMode): Metadata mode, defaults to MetadataMode.NONE.
        When set to MetadataMode.LLM, returns formatted QA pair.
        For other modes, returns base class text format.

Returns:
    str: The formatted text content.
"""
        if metadata_mode == MetadataMode.LLM:
            return f'query:\n{self.text}\nanswer\n{self._answer}'
        return super().get_text(metadata_mode)

lazyllm.tools.rag.dataReader.SimpleDirectoryReader

Bases: ModuleBase

A modular document directory reader that inherits from ModuleBase, supporting reading various document formats from the file system and converting them into standardized DocNode objects. Args: input_dir (Optional[str]): Input directory path. Mutually exclusive with input_files. input_files (Optional[List]): Directly specified list of files. Mutually exclusive with input_dir. exclude (Optional[List]): List of file patterns to exclude. exclude_hidden (bool): Whether to exclude hidden files. recursive (bool): Whether to recursively read subdirectories. encoding (str): Encoding format of text files. required_exts (Optional[List[str]]): Whitelist of file extensions to process. file_extractor (Optional[Dict[str, Callable]]): Dictionary of custom file readers. fs (Optional[AbstractFileSystem]): Custom file system. metadata_genf (Optional[Callable[[str], Dict]]): Metadata generation function that takes a file path and returns a metadata dictionary. num_files_limit (Optional[int]): Maximum number of files to read. return_trace (bool): Whether to return processing trace information. metadatas (Optional[Dict]): Predefined global metadata dictionary.

Examples:

>>> import lazyllm
>>> from lazyllm.tools.dataReader import SimpleDirectoryReader
>>> reader = SimpleDirectoryReader(input_dir="yourpath/",recursive=True,exclude=["*.tmp"],required_exts=[".pdf", ".docx"])
>>> documents = reader.load_data()
Source code in lazyllm/tools/rag/dataReader.py
class SimpleDirectoryReader(ModuleBase):
    """
A modular document directory reader that inherits from ModuleBase, supporting reading various document formats from the file system and converting them into standardized DocNode objects.
Args:
    input_dir (Optional[str]): Input directory path. Mutually exclusive with input_files.
    input_files (Optional[List]): Directly specified list of files. Mutually exclusive with input_dir.
    exclude (Optional[List]): List of file patterns to exclude.
    exclude_hidden (bool): Whether to exclude hidden files.
    recursive (bool): Whether to recursively read subdirectories.
    encoding (str): Encoding format of text files.
    required_exts (Optional[List[str]]): Whitelist of file extensions to process.
    file_extractor (Optional[Dict[str, Callable]]): Dictionary of custom file readers.
    fs (Optional[AbstractFileSystem]): Custom file system.
    metadata_genf (Optional[Callable[[str], Dict]]): Metadata generation function that takes a file path and returns a metadata dictionary.
    num_files_limit (Optional[int]): Maximum number of files to read.
    return_trace (bool): Whether to return processing trace information.
    metadatas (Optional[Dict]): Predefined global metadata dictionary.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools.dataReader import SimpleDirectoryReader
    >>> reader = SimpleDirectoryReader(input_dir="yourpath/",recursive=True,exclude=["*.tmp"],required_exts=[".pdf", ".docx"])
    >>> documents = reader.load_data()
    """
    default_file_readers: Dict[str, Type[ReaderBase]] = {
        "*.pdf": PDFReader,
        "*.docx": DocxReader,
        "*.hwp": HWPReader,
        "*.pptx": PPTXReader,
        "*.ppt": PPTXReader,
        "*.pptm": PPTXReader,
        "*.gif": ImageReader,
        "*.jpeg": ImageReader,
        "*.jpg": ImageReader,
        "*.png": ImageReader,
        "*.webp": ImageReader,
        "*.ipynb": IPYNBReader,
        "*.epub": EpubReader,
        "*.md": MarkdownReader,
        "*.mbox": MboxReader,
        "*.csv": PandasCSVReader,
        "*.xls": PandasExcelReader,
        "*.xlsx": PandasExcelReader,
        "*.mp3": VideoAudioReader,
        "*.mp4": VideoAudioReader,
    }

    def __init__(self, input_dir: Optional[str] = None, input_files: Optional[List] = None,
                 exclude: Optional[List] = None, exclude_hidden: bool = True, recursive: bool = False,
                 encoding: str = "utf-8", filename_as_id: bool = False, required_exts: Optional[List[str]] = None,
                 file_extractor: Optional[Dict[str, Callable]] = None, fs: Optional[AbstractFileSystem] = None,
                 metadata_genf: Optional[Callable[[str], Dict]] = None, num_files_limit: Optional[int] = None,
                 return_trace: bool = False, metadatas: Optional[Dict] = None) -> None:
        super().__init__(return_trace=return_trace)

        if (not input_dir and not input_files) or (input_dir and input_files):
            raise ValueError("Must provide either `input_dir` or `input_files`.")

        self._fs = fs or get_default_fs()
        self._encoding = encoding

        self._exclude = exclude
        self._recursive = recursive
        self._exclude_hidden = exclude_hidden
        self._required_exts = required_exts
        self._num_files_limit = num_files_limit
        self._Path = Path if is_default_fs(self._fs) else PurePosixPath
        self._metadatas = metadatas

        if input_files:
            self._input_files = []
            for path in input_files:
                if not self._fs.isfile(path):
                    path = os.path.join(config['data_path'], path)
                    if not self._fs.isfile(path):
                        raise ValueError(f"File {path} does not exist.")
                input_file = self._Path(path)
                self._input_files.append(input_file)
        elif input_dir:
            if not self._fs.isdir(input_dir):
                raise ValueError(f"Directory {input_dir} does not exist.")
            self._input_dir = self._Path(input_dir)
            self._input_files = self._add_files(self._input_dir)

        self._file_extractor = file_extractor or {}

        self._metadata_genf = metadata_genf or _DefaultFileMetadataFunc(self._fs)
        if filename_as_id: LOG.warning('Argument `filename_as_id` for DataReader is no longer used')

    def _add_files(self, input_dir: Path) -> List[Path]:  # noqa: C901
        all_files = set()
        rejected_files = set()
        rejected_dirs = set()

        if self._exclude is not None:
            for excluded_pattern in self._exclude:
                if self._recursive:
                    excluded_glob = self._Path(input_dir) / self._Path("**") / excluded_pattern
                else:
                    excluded_glob = self._Path(input_dir) / excluded_pattern
                for file in self._fs.glob(str(excluded_glob)):
                    if self._fs.isdir(file):
                        rejected_dirs.add(self._Path(file))
                    else:
                        rejected_files.add(self._Path(file))

        file_refs: List[str] = []
        if self._recursive:
            file_refs = self._fs.glob(str(input_dir) + "/**/*")
        else:
            file_refs = self._fs.glob(str(input_dir) + "/*")

        for ref in file_refs:
            ref = self._Path(ref)
            is_dir = self._fs.isdir(ref)
            skip_hidden = self._exclude_hidden and self._is_hidden(ref)
            skip_bad_exts = (self._required_exts is not None and ref.suffix not in self._required_exts)
            skip_excluded = ref in rejected_files
            if not skip_excluded:
                if is_dir:
                    ref_parent_dir = ref
                else:
                    ref_parent_dir = self._fs._parent(ref)
                for rejected_dir in rejected_dirs:
                    if str(ref_parent_dir).startswith(str(rejected_dir)):
                        skip_excluded = True
                        LOG.warning(f"Skipping {ref} because it in parent dir "
                                    f"{ref_parent_dir} which is in {rejected_dir}.")
                        break

            if is_dir or skip_hidden or skip_bad_exts or skip_excluded:
                continue
            else:
                all_files.add(ref)

        new_input_files = sorted(all_files)

        if len(new_input_files) == 0:
            raise ValueError(f"No files found in {input_dir}.")
        if self._num_files_limit is not None and self._num_files_limit > 0:
            new_input_files = new_input_files[0: self._num_files_limit]

        LOG.debug(f"[SimpleDirectoryReader] Total files add: {len(new_input_files)}")

        LOG.info(f"input_files: {new_input_files}")
        return new_input_files

    def _is_hidden(self, path: Path) -> bool:
        return any(part.startswith(".") and part not in [".", ".."] for part in path.parts)

    def _exclude_metadata(self, documents: List[DocNode]) -> List[DocNode]:
        for doc in documents:
            doc._excluded_embed_metadata_keys.extend(
                ["file_name", "file_type", "file_size", "creation_date",
                 "last_modified_date", "last_accessed_date", "lazyllm_store_num"])
            doc._excluded_llm_metadata_keys.extend(
                ["file_name", "file_type", "file_size", "creation_date",
                 "last_modified_date", "last_accessed_date", "lazyllm_store_num"])
        return documents

    @staticmethod
    def load_file(input_file: Path, metadata_genf: Callable[[str], Dict], file_extractor: Dict[str, Callable],
                  encoding: str = "utf-8", pathm: PurePath = Path, fs: Optional[AbstractFileSystem] = None,
                  metadata: Optional[Dict] = None) -> List[DocNode]:
        # metadata priority: user > reader > metadata_genf
        user_metadata: Dict = metadata or {}
        metadata_generated: Dict = metadata_genf(str(input_file)) if metadata_genf else {}
        documents: List[DocNode] = []

        filename_lower = str(input_file).lower()

        for pattern, extractor in file_extractor.items():
            pt_lower = str(pathm(pattern)).lower()
            match_pattern = pt_lower if pt_lower.endswith("*") else os.path.join(str(pathm.cwd()).lower(), pt_lower)
            if pt_lower.startswith("*"):
                match_pattern = pt_lower
            else:
                base = str(pathm.cwd()).lower()
                match_pattern = os.path.join(base, pt_lower)

            if fnmatch.fnmatch(filename_lower, match_pattern):
                reader = extractor() if isinstance(extractor, type) else extractor
                kwargs = {'fs': fs} if fs and not is_default_fs(fs) else {}
                docs = reader(input_file, **kwargs)
                if isinstance(docs, DocNode): docs = [docs]
                for doc in docs:
                    metadata = metadata_generated.copy()
                    metadata.update(doc._global_metadata or {})
                    metadata.update(user_metadata)
                    doc._global_metadata = metadata

                if config['rag_filename_as_id']:
                    for i, doc in enumerate(docs):
                        doc._uid = f"{input_file!s}_index_{i}"
                documents.extend(docs)
                break
        else:
            if not config['use_fallback_reader']:
                LOG.warning(f'no pattern found for {input_file}! '
                            'If you want fallback to default Reader, set `LAZYLLM_USE_FALLBACK_READER=True`.')
                return documents
            fs = fs or get_default_fs()
            with fs.open(input_file, encoding=encoding) as f:
                try:
                    data = f.read().decode(encoding)
                    doc = DocNode(text=data, global_metadata=user_metadata)
                    documents.append(doc)
                except Exception:
                    LOG.error(f'no pattern found for {input_file} and it is not utf-8, skip it!')
        return documents

    def _load_data(self, show_progress: bool = False, num_workers: Optional[int] = None,
                   fs: Optional[AbstractFileSystem] = None) -> List[DocNode]:
        documents = []

        fs = fs or self._fs
        process_file = self._input_files
        file_readers = self._file_extractor.copy()
        for key, func in self.default_file_readers.items():
            if key not in file_readers: file_readers[key] = func

        if num_workers and num_workers >= 1:
            if num_workers > multiprocessing.cpu_count():
                LOG.warning("Specified num_workers exceed number of CPUs in the system. "
                            "Setting `num_workers` down to the maximum CPU count.")
            with multiprocessing.get_context("spawn").Pool(num_workers) as p:
                results = p.starmap(SimpleDirectoryReader.load_file,
                                    zip(process_file, repeat(self._metadata_genf), repeat(file_readers),
                                        repeat(self._encoding), repeat(self._Path),
                                        repeat(self._fs), self._metadatas or repeat(None)))
                documents = reduce(lambda x, y: x + y, results)
        else:
            if show_progress:
                process_file = tqdm(self._input_files, desc="Loading files", unit="file")
            for input_file, metadata in zip(process_file, self._metadatas or repeat(None)):
                documents.extend(
                    SimpleDirectoryReader.load_file(
                        input_file=input_file, metadata_genf=self._metadata_genf, file_extractor=file_readers,
                        encoding=self._encoding, pathm=self._Path, fs=self._fs, metadata=metadata))

        return self._exclude_metadata(documents)

    def forward(self, *args, **kwargs) -> List[DocNode]:
        return self._load_data(*args, **kwargs)

lazyllm.tools.rag.dataReader.FileReader

Bases: object

File content reader whose main function is to convert various input file formats into concatenated plain text content. Args: input_files (Optional[List]): Directly specified list of input files.

Examples:

>>> import lazyllm
>>> from lazyllm.tools.dataReader import FileReader
>>> reader = FileReader()
>>> content = reader("yourpath/")
Source code in lazyllm/tools/rag/dataReader.py
class FileReader(object):
    """
File content reader whose main function is to convert various input file formats into concatenated plain text content.
Args:
    input_files (Optional[List]): Directly specified list of input files.


Examples:

    >>> import lazyllm
    >>> from lazyllm.tools.dataReader import FileReader
    >>> reader = FileReader()
    >>> content = reader("yourpath/") 
    """

    def __call__(self, input_files):
        file_list = _lazyllm_get_file_list(input_files)
        if isinstance(file_list, str) and file_list is not None:
            file_list = [file_list]
        if len(file_list) == 0:
            return []
        nodes = SimpleDirectoryReader(input_files=file_list)._load_data()
        txt = [node.get_text() for node in nodes]
        return "\n".join(txt)

lazyllm.tools.rag.transform.FuncNodeTransform

Bases: NodeTransform

Used for user defined function.

Wrapped the transform to: List[Docnode] -> List[Docnode]

This wrapper supports when trans_node is False:
    1. str -> list: transform=lambda t: t.split('

') 2. str -> str: transform=lambda t: t[:3]

This wrapper supports when trans_node is True:
    1. DocNode -> list: pipeline(lambda x:x, SentenceSplitter)
    2. DocNode -> DocNode: pipeline(LLMParser)
Source code in lazyllm/tools/rag/transform.py
class FuncNodeTransform(NodeTransform):
    """Used for user defined function.

    Wrapped the transform to: List[Docnode] -> List[Docnode]

    This wrapper supports when trans_node is False:
        1. str -> list: transform=lambda t: t.split('\n')
        2. str -> str: transform=lambda t: t[:3]

    This wrapper supports when trans_node is True:
        1. DocNode -> list: pipeline(lambda x:x, SentenceSplitter)
        2. DocNode -> DocNode: pipeline(LLMParser)
    """

    def __init__(self, func: Union[Callable[[str], List[str]], Callable[[DocNode], List[DocNode]]],
                 trans_node: bool = None, num_workers: int = 0):
        super(__class__, self).__init__(num_workers=num_workers)
        self._func, self._trans_node = func, trans_node

    def transform(self, node: DocNode, **kwargs) -> List[Union[str, DocNode]]:
        """
Transform a document node using the wrapped user-defined function.

This method applies the user-defined function to either the text content of the node (when trans_node=False) or the node itself (when trans_node=True).

Args:
    node (DocNode): The document node to be transformed.
    **kwargs: Additional keyword arguments passed to the transformation function.

Returns:
    List[Union[str, DocNode]]: The transformed results, which can be either strings or DocNode objects depending on the function implementation.
"""
        return self._func(node if self._trans_node else node.get_text())

transform(node, **kwargs)

Transform a document node using the wrapped user-defined function.

This method applies the user-defined function to either the text content of the node (when trans_node=False) or the node itself (when trans_node=True).

Parameters:

  • node (DocNode) –

    The document node to be transformed.

  • **kwargs

    Additional keyword arguments passed to the transformation function.

Returns:

  • List[Union[str, DocNode]]

    List[Union[str, DocNode]]: The transformed results, which can be either strings or DocNode objects depending on the function implementation.

Source code in lazyllm/tools/rag/transform.py
    def transform(self, node: DocNode, **kwargs) -> List[Union[str, DocNode]]:
        """
Transform a document node using the wrapped user-defined function.

This method applies the user-defined function to either the text content of the node (when trans_node=False) or the node itself (when trans_node=True).

Args:
    node (DocNode): The document node to be transformed.
    **kwargs: Additional keyword arguments passed to the transformation function.

Returns:
    List[Union[str, DocNode]]: The transformed results, which can be either strings or DocNode objects depending on the function implementation.
"""
        return self._func(node if self._trans_node else node.get_text())

lazyllm.tools.rag.web.DocWebModule

Bases: ModuleBase

Document Web Interface Module, inherits from ModuleBase, provides web-based document management interface.

Parameters:

  • doc_server (ServerModule) –

    Document server module instance providing backend API support

  • title (str, default: '文档管理演示终端' ) –

    Interface title, defaults to "文档管理演示终端"

  • port (int / range / list, default: None ) –

    Service port number or range, defaults to 20800-20999

  • history (list, default: None ) –

    Initial chat history, defaults to empty list

  • text_mode (Mode, default: None ) –

    Text processing mode, defaults to Mode.Dynamic

  • trace_mode (Mode, default: None ) –

    Trace mode, defaults to Mode.Refresh

Class Attributes

Mode: Mode enumeration class containing: - Dynamic: Dynamic mode - Refresh: Refresh mode - Appendix: Appendix mode

Notes
  • Requires a valid doc_server instance to work with
  • Automatically tries other ports in range when port conflict occurs
  • Releases resources when service is stopped

Examples:

>>> import lazyllm
>>> from lazyllm.tools.rag.web import DocWebModule
>>> from lazyllm import
>>> doc_server = ServerModule(url="your_url")
>>> doc_web = DocWebModule(
>>>   doc_server=doc_server,
>>>   title="文档管理演示终端",
>>>   port=range(20800, 20805)  # 自动寻找可用端口)
>>> deploy_task = doc_web._get_deploy_tasks()
>>> deploy_task()  
>>> print(doc_web.url)
>>> doc_web.stop()
Source code in lazyllm/tools/rag/web.py
class DocWebModule(ModuleBase):
    """Document Web Interface Module, inherits from ModuleBase, provides web-based document management interface.

Args:
    doc_server (ServerModule): Document server module instance providing backend API support
    title (str): Interface title, defaults to "文档管理演示终端"
    port (int/range/list): Service port number or range, defaults to 20800-20999
    history (list): Initial chat history, defaults to empty list
    text_mode (Mode): Text processing mode, defaults to Mode.Dynamic
    trace_mode (Mode): Trace mode, defaults to Mode.Refresh

Class Attributes:
    Mode: Mode enumeration class containing:
        - Dynamic: Dynamic mode
        - Refresh: Refresh mode
        - Appendix: Appendix mode

Notes:
    - Requires a valid doc_server instance to work with
    - Automatically tries other ports in range when port conflict occurs
    - Releases resources when service is stopped


Examples:
    >>> import lazyllm
    >>> from lazyllm.tools.rag.web import DocWebModule
    >>> from lazyllm import
    >>> doc_server = ServerModule(url="your_url")
    >>> doc_web = DocWebModule(
    >>>   doc_server=doc_server,
    >>>   title="文档管理演示终端",
    >>>   port=range(20800, 20805)  # 自动寻找可用端口)
    >>> deploy_task = doc_web._get_deploy_tasks()
    >>> deploy_task()  
    >>> print(doc_web.url)
    >>> doc_web.stop()
    """
    class Mode:
        Dynamic = 0
        Refresh = 1
        Appendix = 2

    def __init__(self, doc_server: ServerModule, title="文档管理演示终端", port=None,
                 history=None, text_mode=None, trace_mode=None) -> None:
        super().__init__()
        self.title = title
        self.port = port or range(20800, 20999)
        self.history = history or []
        self.trace_mode = trace_mode if trace_mode else DocWebModule.Mode.Refresh
        self.text_mode = text_mode if text_mode else DocWebModule.Mode.Dynamic
        self.doc_server = doc_server
        self._deploy_flag = lazyllm.once_flag()
        self.api_url = ""
        self.url = ""

    def _prepare(self, query, chat_history):
        if chat_history is None:
            chat_history = []
        return "", chat_history + [[query, None]]

    def _clear_history(self):
        return [], "", ""

    def _work(self):
        if isinstance(self.port, (range, tuple, list)):
            port = self._find_can_use_network_port()
        else:
            port = self.port
            assert self._verify_port_access(port), f"port {port} is occupied"

        self.api_url = self.doc_server._url.rsplit("/", 1)[0]
        self.web_ui = WebUi(self.api_url)
        self.demo = self.web_ui.create_ui()
        self.url = f'http://127.0.0.1:{port}'
        self.broadcast_url = f'http://0.0.0.0:{port}'

        self.demo.queue().launch(server_name="0.0.0.0", server_port=port, prevent_thread_lock=True)
        LOG.success('LazyLLM docwebmodule launched successfully: Running on: '
                    f'{self.broadcast_url}, local URL: {self.url}')

    def _get_deploy_tasks(self):
        return Pipeline(self._work)

    def _get_post_process_tasks(self):
        return Pipeline(self._print_url)

    def wait(self):
        """Blocks the current thread to keep the web interface running until manually stopped.

"""
        self.demo.block_thread()

    def stop(self):
        """Stops the web interface service and releases related resources.

"""
        if self.demo:
            self.demo.close()
            del self.demo
            self.demo, self.url = None, ''

    def _find_can_use_network_port(self):
        for port in self.port:
            if self._verify_port_access(port):
                return port
        raise RuntimeError(
            f"The ports in the range {self.port} are all occupied. "
            "Please change the port range or release the relevant ports."
        )

    def _print_url(self):
        lazyllm.LOG.success(f"LazyLLM DocWebModule launched successfully: Running on local URL: {self.url}")

    def _verify_port_access(self, port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            result = s.connect_ex(("127.0.0.1", port))
            return result != 0

    def __repr__(self):
        return lazyllm.make_repr("Module", "DocWebModule")

stop()

Stops the web interface service and releases related resources.

Source code in lazyllm/tools/rag/web.py
    def stop(self):
        """Stops the web interface service and releases related resources.

"""
        if self.demo:
            self.demo.close()
            del self.demo
            self.demo, self.url = None, ''

wait()

Blocks the current thread to keep the web interface running until manually stopped.

Source code in lazyllm/tools/rag/web.py
    def wait(self):
        """Blocks the current thread to keep the web interface running until manually stopped.

"""
        self.demo.block_thread()

lazyllm.tools.WebModule

Bases: ModuleBase

WebModule is a web-based interactive interface provided by LazyLLM for developers. After initializing and starting a WebModule, developers can see structure of the module they provides behind the WebModule, and transmit the input of the Chatbot component to their modules. The results and logs returned by the module will be displayed on the “Processing Logs” and Chatbot component on the web page. In addition, Checkbox or Text components can be added programmatically to the web page for additional parameters to the background module. Meanwhile, The WebModule page provides Checkboxes of “Use Context,” “Stream Output,” and “Append Output,” which can be used to adjust the interaction between the page and the module behind.

WebModule.init_web(component_descs) -> gradio.Blocks

Generate a demonstration web page based on gradio. The function initializes session-related data to save chat history and logs for different pages, then dynamically add Checkbox and Text components to the page according to component_descs parameter, and set the corresponding functions for the buttons and text boxes on the page at last. WebModule’s init function calls this method to generate the page.

Parameters:

  • component_descs (list) –

    A list used to add components to the page. Each element in the list is also a list containing

Examples:

>>> import lazyllm
>>> def func2(in_str, do_sample=True, temperature=0.0, *args, **kwargs):
...     return f"func2:{in_str}|do_sample:{str(do_sample)}|temp:{temperature}"
...
>>> m1=lazyllm.ActionModule(func2)
>>> m1.name="Module1"
>>> w = lazyllm.WebModule(m1, port=[20570, 20571, 20572], components={
...         m1:[('do_sample', 'Checkbox', True), ('temperature', 'Text', 0.1)]},
...                       text_mode=lazyllm.tools.WebModule.Mode.Refresh)
>>> w.start()
193703: 2024-06-07 10:26:00 lazyllm SUCCESS: ...
Source code in lazyllm/tools/webpages/webmodule.py
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
class WebModule(ModuleBase):
    """WebModule is a web-based interactive interface provided by LazyLLM for developers. After initializing and starting
a WebModule, developers can see structure of the module they provides behind the WebModule, and transmit the input
of the Chatbot component to their modules. The results and logs returned by the module will be displayed on the
“Processing Logs” and Chatbot component on the web page. In addition, Checkbox or Text components can be added
programmatically to the web page for additional parameters to the background module. Meanwhile, The WebModule page
provides Checkboxes of “Use Context,” “Stream Output,” and “Append Output,” which can be used to adjust the
interaction between the page and the module behind.

<span style="font-size: 20px;">&ensp;**`WebModule.init_web(component_descs) -> gradio.Blocks`**</span>

Generate a demonstration web page based on gradio. The function initializes session-related data to save chat history
and logs for different pages, then dynamically add Checkbox and Text components to the page according to component_descs
parameter, and set the corresponding functions for the buttons and text boxes on the page at last.
WebModule’s __init__ function calls this method to generate the page.

Args:
    component_descs (list): A list used to add components to the page. Each element in the list is also a list containing
    5 elements, which are the module ID, the module name, the component name, the component type (currently only
    supports Checkbox and Text), and the default value of the component.



Examples:
    >>> import lazyllm
    >>> def func2(in_str, do_sample=True, temperature=0.0, *args, **kwargs):
    ...     return f"func2:{in_str}|do_sample:{str(do_sample)}|temp:{temperature}"
    ...
    >>> m1=lazyllm.ActionModule(func2)
    >>> m1.name="Module1"
    >>> w = lazyllm.WebModule(m1, port=[20570, 20571, 20572], components={
    ...         m1:[('do_sample', 'Checkbox', True), ('temperature', 'Text', 0.1)]},
    ...                       text_mode=lazyllm.tools.WebModule.Mode.Refresh)
    >>> w.start()
    193703: 2024-06-07 10:26:00 lazyllm SUCCESS: ...
    """
    class Mode:
        Dynamic = 0
        Refresh = 1
        Appendix = 2

    def __init__(self, m: Any, *, components: Dict[Any, Any] = dict(), title: str = '对话演示终端',  # noqa B008
                 port: Optional[Union[int, range, tuple, list]] = None, history: List[Any] = [],  # noqa B006
                 text_mode: Optional[Mode] = None, trace_mode: Optional[Mode] = None, audio: bool = False,
                 stream: bool = False, files_target: Optional[Union[Any, List[Any]]] = None,
                 static_paths: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
                 encode_files: bool = False, share: bool = False) -> None:
        super().__init__()
        # Set the static directory of gradio so that gradio can access local resources in the directory
        if isinstance(static_paths, (str, Path)):
            self._static_paths = [static_paths]
        elif isinstance(static_paths, list) and all(isinstance(p, (str, Path)) for p in static_paths):
            self._static_paths = static_paths
        elif static_paths is None:
            self._static_paths = []
        else:
            raise ValueError(f"static_paths only supported str, path or list types. Not supported {static_paths}")
        self.m = lazyllm.ActionModule(m) if isinstance(m, lazyllm.FlowBase) else m
        self.pool = lazyllm.ThreadPoolExecutor(max_workers=50)
        self.title = title
        self.port = port or range(20500, 20799)
        components = sum([[([k._module_id, k._module_name] + list(v)) for v in vs]
                         for k, vs in components.items()], [])
        self.ckeys = [[c[0], c[2]] for c in components]
        if isinstance(m, (OnlineChatModule, TrainableModule)) and not history:
            history = [m]
        self.history = [h._module_id for h in history]
        if trace_mode:
            LOG.warn('trace_mode is deprecated')
        self.text_mode = text_mode if text_mode else WebModule.Mode.Dynamic
        self.cach_path = self._set_up_caching()
        self.audio = audio
        self.stream = stream
        self.files_target = files_target if isinstance(files_target, list) or files_target is None else [files_target]
        self.encode_files = encode_files
        self.share = share
        self.demo = self.init_web(components)
        self.url = None
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)

    def _get_all_file_submodule(self):
        if self.files_target: return
        self.files_target = []
        self.for_each(
            lambda x: getattr(x, 'template_message', None),
            lambda x: self.files_target.append(x)
        )

    def _signal_handler(self, signum, frame):
        LOG.info(f"Signal {signum} received, terminating subprocess.")
        atexit._run_exitfuncs()
        sys.exit(0)

    def _set_up_caching(self):
        if 'GRADIO_TEMP_DIR' in os.environ:
            cach_path = os.environ['GRADIO_TEMP_DIR']
        else:
            cach_path = os.path.join(lazyllm.config['temp_dir'], 'gradio_cach')
            os.environ['GRADIO_TEMP_DIR'] = cach_path
        if not os.path.exists(cach_path):
            os.makedirs(cach_path)
        return cach_path

    def init_web(self, component_descs):
        """Initialize the Web UI page.
This method uses Gradio to build the interactive chat interface and binds all components to the appropriate logic. It supports session selection, streaming output, context toggling, multimodal input, and control tools. The method returns the constructed Gradio Blocks object.
Args:
    component_descs (List[Tuple]): A list of component descriptors. Each element is a 5-tuple 
        (module, group_name, name, component_type, value), e.g. ('MyModule', 'GroupA', 'use_cache', 'Checkbox', True).
Returns:
    gr.Blocks: The constructed Gradio UI object, which can be launched via `.launch()`.
"""
        gr.set_static_paths(self._static_paths)
        with gr.Blocks(css=css, title=self.title, analytics_enabled=False) as demo:
            sess_data = gr.State(value={
                'sess_titles': [''],
                'sess_logs': {},
                'sess_history': {},
                'sess_num': 1,
                'curr_sess': '',
                'frozen_query': '',
            })
            with gr.Row():
                with gr.Column(scale=3):
                    with gr.Row():
                        with lazyllm.config.temp('repr_show_child', True):
                            gr.Textbox(elem_id='module', interactive=False, show_label=True,
                                       label="模型结构", value=repr(self.m))
                    with gr.Row():
                        chat_use_context = gr.Checkbox(interactive=True, value=False, label="使用上下文")
                    with gr.Row():
                        stream_output = gr.Checkbox(interactive=self.stream, value=self.stream, label="流式输出")
                        text_mode = gr.Checkbox(interactive=(self.text_mode == WebModule.Mode.Dynamic),
                                                value=(self.text_mode != WebModule.Mode.Refresh), label="追加输出")
                    components = []
                    for _, gname, name, ctype, value in component_descs:
                        if ctype in ('Checkbox', 'Text'):
                            components.append(getattr(gr, ctype)(interactive=True, value=value, label=f'{gname}.{name}'))
                        elif ctype == 'Dropdown':
                            components.append(getattr(gr, ctype)(interactive=True, choices=value,
                                                                 label=f'{gname}.{name}'))
                        else:
                            raise KeyError(f'invalid component type: {ctype}')
                    with gr.Row():
                        dbg_msg = gr.Textbox(show_label=True, label='处理日志',
                                             elem_id='logging', interactive=False, max_lines=10)
                    clear_btn = gr.Button(value="🗑️  Clear history", interactive=True)
                with gr.Column(scale=6):
                    with gr.Row():
                        add_sess_btn = gr.Button("添加新会话")
                        sess_drpdn = gr.Dropdown(choices=sess_data.value['sess_titles'], label="选择会话:", value='')
                        del_sess_btn = gr.Button("删除当前会话")
                    chatbot = gr.Chatbot(height=700)
                    query_box = gr.MultimodalTextbox(show_label=False, placeholder='输入内容并回车!!!', interactive=True)
                    recordor = gr.Audio(sources=["microphone"], type="filepath", visible=self.audio)

            query_box.submit(self._init_session, [query_box, sess_data, recordor],
                                                 [sess_drpdn, chatbot, dbg_msg, sess_data, recordor], queue=True
                ).then(lambda: gr.update(interactive=False), None, query_box, queue=False
                ).then(lambda: gr.update(interactive=False), None, add_sess_btn, queue=False
                ).then(lambda: gr.update(interactive=False), None, sess_drpdn, queue=False
                ).then(lambda: gr.update(interactive=False), None, del_sess_btn, queue=False
                ).then(self._prepare, [query_box, chatbot, sess_data], [query_box, chatbot], queue=True
                ).then(self._respond_stream, [chat_use_context, chatbot, stream_output, text_mode] + components,
                                             [chatbot, dbg_msg], queue=chatbot
                ).then(lambda: gr.update(interactive=True), None, query_box, queue=False
                ).then(lambda: gr.update(interactive=True), None, add_sess_btn, queue=False
                ).then(lambda: gr.update(interactive=True), None, sess_drpdn, queue=False
                ).then(lambda: gr.update(interactive=True), None, del_sess_btn, queue=False)
            clear_btn.click(self._clear_history, [sess_data], outputs=[chatbot, query_box, dbg_msg, sess_data])

            sess_drpdn.change(self._change_session, [sess_drpdn, chatbot, dbg_msg, sess_data],
                                                    [sess_drpdn, chatbot, query_box, dbg_msg, sess_data])
            add_sess_btn.click(self._add_session, [chatbot, dbg_msg, sess_data],
                                                  [sess_drpdn, chatbot, query_box, dbg_msg, sess_data])
            del_sess_btn.click(self._delete_session, [sess_drpdn, sess_data],
                                                     [sess_drpdn, chatbot, query_box, dbg_msg, sess_data])
            recordor.change(self._sub_audio, recordor, query_box)
            return demo

    def _sub_audio(self, audio):
        if audio:
            return {'text': '', 'files': [audio]}
        else:
            return {}

    def _init_session(self, query, session, audio):
        audio = None
        session['frozen_query'] = query
        if session['curr_sess'] != '':  # remain unchanged.
            return gr.Dropdown(), gr.Chatbot(), gr.Textbox(), session, audio

        if "text" in query and query["text"] is not None:
            id_name = query['text']
        else:
            id_name = id(id_name)
        session['curr_sess'] = f"({session['sess_num']})  {id_name}"
        session['sess_num'] += 1
        session['sess_titles'][0] = session['curr_sess']

        session['sess_logs'][session['curr_sess']] = []
        session['sess_history'][session['curr_sess']] = []
        return gr.update(choices=session['sess_titles'], value=session['curr_sess']), [], '', session, audio

    def _add_session(self, chat_history, log_history, session):
        if session['curr_sess'] == '':
            LOG.warning('Cannot create new session while current session is empty.')
            return gr.Dropdown(), gr.Chatbot(), {}, gr.Textbox(), session

        self._save_history(chat_history, log_history, session)

        session['curr_sess'] = ''
        session['sess_titles'].insert(0, session['curr_sess'])
        return gr.update(choices=session['sess_titles'], value=session['curr_sess']), [], {}, '', session

    def _save_history(self, chat_history, log_history, session):
        if session['curr_sess'] in session['sess_titles']:
            session['sess_history'][session['curr_sess']] = chat_history
            session['sess_logs'][session['curr_sess']] = log_history

    def _change_session(self, session_title, chat_history, log_history, session):
        if session['curr_sess'] == '':  # new session
            return gr.Dropdown(), [], {}, '', session

        if session_title not in session['sess_titles']:
            LOG.warning(f'{session_title} is not an existing session title.')
            return gr.Dropdown(), gr.Chatbot(), {}, gr.Textbox(), session

        self._save_history(chat_history, log_history, session)

        session['curr_sess'] = session_title
        return (gr.update(choices=session['sess_titles'], value=session['curr_sess']),
                session['sess_history'][session['curr_sess']], {},
                session['sess_logs'][session['curr_sess']], session)

    def _delete_session(self, session_title, session):
        if session_title not in session['sess_titles']:
            LOG.warning(f'session {session_title} does not exist.')
            return gr.Dropdown(), session
        session['sess_titles'].remove(session_title)

        if session_title != '':
            del session['sess_history'][session_title]
            del session['sess_logs'][session_title]
            session['curr_sess'] = session_title
        else:
            session['curr_sess'] = 'dummy session'
            # add_session and change_session cannot accept an uninitialized session.
            # Here we need to imitate removal of a real session so that
            # add_session and change_session could skip saving chat history.

        if len(session['sess_titles']) == 0:
            return self._add_session(None, None, session)
        else:
            return self._change_session(session['sess_titles'][0], None, {}, session)

    def _prepare(self, query, chat_history, session):
        if not query.get('text', '') and not query.get('files', []):
            query = session['frozen_query']
        if chat_history is None:
            chat_history = []
        for x in query["files"]:
            chat_history.append([[x,], None])
        if "text" in query and query["text"]:
            chat_history.append([query['text'], None])
        return {}, chat_history

    def _respond_stream(self, use_context, chat_history, stream_output, append_text, *args):  # noqa C901
        try:
            # TODO: move context to trainable module
            files = []
            chat_history[-1][1], log_history = '', []
            for file in chat_history[::-1]:
                if file[-1]: break  # not current chat
                if isinstance(file[0], (tuple, list)):
                    files.append(file[0][0])
                elif isinstance(file[0], str) and file[0].startswith('lazyllm_img::'):  # Just for pytest
                    files.append(file[0][13:])
            if isinstance(chat_history[-1][0], str):
                string = chat_history[-1][0]
            else:
                string = ''
            if self.files_target is None and not self.encode_files:
                self._get_all_file_submodule()
            if self.encode_files and files:
                string = encode_query_with_filepaths(string, files)
            if files and self.files_target:
                for module in self.files_target:
                    assert isinstance(module, ModuleBase)
                    if module._module_id in globals['lazyllm_files']:
                        globals['lazyllm_files'][module._module_id].extend(files)
                    else:
                        globals['lazyllm_files'][module._module_id] = files
                string += f' ## Get attachments: {os.path.basename(files[-1])}'
            elif self.files_target:
                for module in self.files_target:
                    assert isinstance(module, ModuleBase)
                    globals['lazyllm_files'][module._module_id] = []
            input = string
            history = chat_history[:-1] if use_context and len(chat_history) > 1 else list()

            for k, v in zip(self.ckeys, args):
                if k[0] not in globals['global_parameters']: globals['global_parameters'][k[0]] = dict()
                globals['global_parameters'][k[0]][k[1]] = v

            if use_context:
                for h in self.history:
                    if h not in globals['chat_history']: globals['chat_history'][h] = list()
                    globals['chat_history'][h] = history

            if FileSystemQueue().size() > 0: FileSystemQueue().clear()
            kw = dict(stream_output=stream_output) if isinstance(self.m, (TrainableModule, OnlineChatModule)) else {}
            func_future = self.pool.submit(self.m, input, **kw)
            while True:
                if value := FileSystemQueue().dequeue():
                    chat_history[-1][1] += ''.join(value) if append_text else ''.join(value)
                    if stream_output: yield chat_history, ''
                elif value := FileSystemQueue.get_instance('lazy_error').dequeue():
                    log_history.append(''.join(value))
                elif value := FileSystemQueue.get_instance('lazy_trace').dequeue():
                    log_history.append(''.join(value))
                elif func_future.done(): break
                time.sleep(0.01)
            result = func_future.result()
            if FileSystemQueue().size() > 0: FileSystemQueue().clear()

            def get_log_and_message(s):
                if isinstance(s, dict):
                    s = s.get("message", {}).get("content", "")
                else:
                    try:
                        r = decode_query_with_filepaths(s)
                        if isinstance(r, str):
                            r = json.loads(r)
                        if 'choices' in r:
                            if "type" not in r["choices"][0] or (
                                    "type" in r["choices"][0] and r["choices"][0]["type"] != "tool_calls"):
                                delta = r["choices"][0]["delta"]
                                if "content" in delta:
                                    s = delta["content"]
                                else:
                                    s = ""
                        elif isinstance(r, dict) and 'files' in r and 'query' in r:
                            return r['query'], ''.join(log_history), r['files'] if len(r['files']) > 0 else None
                        else:
                            s = s
                    except (ValueError, KeyError, TypeError):
                        s = s
                    except Exception as e:
                        LOG.error(f"Uncaptured error `{e}` when parsing `{s}`, please contact us if you see this.")
                return s, "".join(log_history), None

            def contains_markdown_image(text: str):
                pattern = r"!\[.*?\]\((.*?)\)"
                return bool(re.search(pattern, text))

            def extract_img_path(text: str):
                pattern = r"!\[.*?\]\((.*?)\)"
                urls = re.findall(pattern, text)
                return urls

            file_paths = None
            if isinstance(result, (str, dict)):
                result, log, file_paths = get_log_and_message(result)
            if file_paths:
                for i, file_path in enumerate(file_paths):
                    suffix = os.path.splitext(file_path)[-1].lower()
                    file = None
                    if suffix in PIL.Image.registered_extensions().keys():
                        file = gr.Image(file_path)
                    elif suffix in ('.mp3', '.wav'):
                        file = gr.Audio(file_path)
                    elif suffix in ('.mp4'):
                        file = gr.Video(file_path)
                    else:
                        LOG.error(f'Not supported typr: {suffix}, for file: {file}')
                    if i == 0:
                        chat_history[-1][1] = file
                    else:
                        chat_history.append([None, file])
                if result:
                    chat_history.append([None, result])
            else:
                assert isinstance(result, str), f'Result should only be str, but got {type(result)}'
                show_result = result
                if contains_markdown_image(show_result):
                    urls = extract_img_path(show_result)
                    for url in urls:
                        suffix = os.path.splitext(url)[-1].lower()
                        if suffix in PIL.Image.registered_extensions().keys() and os.path.exists(url):
                            show_result = show_result.replace(url, "file=" + url)
                if result:
                    count = (len(match.group(1)) if (match := re.search(r'(\n+)$', result)) else 0) + len(result) + 1
                    if not (result in chat_history[-1][1][-count:]):
                        chat_history[-1][1] += "\n\n" + show_result
                    elif show_result != result:
                        chat_history[-1][1] = chat_history[-1][1].replace(result, show_result)
        except requests.RequestException as e:
            chat_history = None
            log = str(e)
        except Exception as e:
            chat_history = None
            log = f'{str(e)}\n--- traceback ---\n{traceback.format_exc()}'
            LOG.error(log)
        globals['chat_history'].clear()
        yield chat_history, log

    def _clear_history(self, session):
        session['sess_history'][session['curr_sess']] = []
        session['sess_logs'][session['curr_sess']] = []
        return [], {}, '', session

    def _work(self):
        if isinstance(self.port, (range, tuple, list)):
            port = self._find_can_use_network_port()
        else:
            port = self.port
            assert self._verify_port_access(port), f'port {port} is occupied'

        self.url = f'http://127.0.0.1:{port}'
        self.broadcast_url = f'http://0.0.0.0:{port}'

        self.demo.queue().launch(server_name="0.0.0.0", server_port=port, prevent_thread_lock=True, share=self.share)
        LOG.success('LazyLLM webmodule launched successfully: Running on: '
                    f'{self.broadcast_url}, local URL: {self.url}')

    def _update(self, *, mode=None, recursive=True):
        super(__class__, self)._update(mode=mode, recursive=recursive)
        self._work()
        return self

    def wait(self):
        """Block the main thread until the web interface is closed.
This method blocks the current thread until the Gradio demo is closed. Useful in deployment scenarios to prevent premature program exit.
"""
        self.demo.block_thread()

    def stop(self):
        """Stop the web interface and clean up resources.
If the web demo has been initialized, this method closes the Gradio demo, frees related resources, and resets `demo` and `url` attributes.
"""
        if self.demo:
            self.demo.close()
            del self.demo
            self.demo, self.url = None, ''

    @property
    def status(self):
        return 'running' if self.url else 'waiting' if self.url is None else 'Cancelled'

    def __repr__(self):
        return lazyllm.make_repr('Module', 'Web', name=self._module_name, subs=[repr(self.m)])

    def _find_can_use_network_port(self):
        for port in self.port:
            if self._verify_port_access(port):
                return port
        raise RuntimeError(
            f'The ports in the range {self.port} are all occupied. '
            'Please change the port range or release the relevant ports.'
        )

    def _verify_port_access(self, port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            result = s.connect_ex(('127.0.0.1', port))
            return result != 0

init_web(component_descs)

Initialize the Web UI page. This method uses Gradio to build the interactive chat interface and binds all components to the appropriate logic. It supports session selection, streaming output, context toggling, multimodal input, and control tools. The method returns the constructed Gradio Blocks object. Args: component_descs (List[Tuple]): A list of component descriptors. Each element is a 5-tuple (module, group_name, name, component_type, value), e.g. ('MyModule', 'GroupA', 'use_cache', 'Checkbox', True). Returns: gr.Blocks: The constructed Gradio UI object, which can be launched via .launch().

Source code in lazyllm/tools/webpages/webmodule.py
    def init_web(self, component_descs):
        """Initialize the Web UI page.
This method uses Gradio to build the interactive chat interface and binds all components to the appropriate logic. It supports session selection, streaming output, context toggling, multimodal input, and control tools. The method returns the constructed Gradio Blocks object.
Args:
    component_descs (List[Tuple]): A list of component descriptors. Each element is a 5-tuple 
        (module, group_name, name, component_type, value), e.g. ('MyModule', 'GroupA', 'use_cache', 'Checkbox', True).
Returns:
    gr.Blocks: The constructed Gradio UI object, which can be launched via `.launch()`.
"""
        gr.set_static_paths(self._static_paths)
        with gr.Blocks(css=css, title=self.title, analytics_enabled=False) as demo:
            sess_data = gr.State(value={
                'sess_titles': [''],
                'sess_logs': {},
                'sess_history': {},
                'sess_num': 1,
                'curr_sess': '',
                'frozen_query': '',
            })
            with gr.Row():
                with gr.Column(scale=3):
                    with gr.Row():
                        with lazyllm.config.temp('repr_show_child', True):
                            gr.Textbox(elem_id='module', interactive=False, show_label=True,
                                       label="模型结构", value=repr(self.m))
                    with gr.Row():
                        chat_use_context = gr.Checkbox(interactive=True, value=False, label="使用上下文")
                    with gr.Row():
                        stream_output = gr.Checkbox(interactive=self.stream, value=self.stream, label="流式输出")
                        text_mode = gr.Checkbox(interactive=(self.text_mode == WebModule.Mode.Dynamic),
                                                value=(self.text_mode != WebModule.Mode.Refresh), label="追加输出")
                    components = []
                    for _, gname, name, ctype, value in component_descs:
                        if ctype in ('Checkbox', 'Text'):
                            components.append(getattr(gr, ctype)(interactive=True, value=value, label=f'{gname}.{name}'))
                        elif ctype == 'Dropdown':
                            components.append(getattr(gr, ctype)(interactive=True, choices=value,
                                                                 label=f'{gname}.{name}'))
                        else:
                            raise KeyError(f'invalid component type: {ctype}')
                    with gr.Row():
                        dbg_msg = gr.Textbox(show_label=True, label='处理日志',
                                             elem_id='logging', interactive=False, max_lines=10)
                    clear_btn = gr.Button(value="🗑️  Clear history", interactive=True)
                with gr.Column(scale=6):
                    with gr.Row():
                        add_sess_btn = gr.Button("添加新会话")
                        sess_drpdn = gr.Dropdown(choices=sess_data.value['sess_titles'], label="选择会话:", value='')
                        del_sess_btn = gr.Button("删除当前会话")
                    chatbot = gr.Chatbot(height=700)
                    query_box = gr.MultimodalTextbox(show_label=False, placeholder='输入内容并回车!!!', interactive=True)
                    recordor = gr.Audio(sources=["microphone"], type="filepath", visible=self.audio)

            query_box.submit(self._init_session, [query_box, sess_data, recordor],
                                                 [sess_drpdn, chatbot, dbg_msg, sess_data, recordor], queue=True
                ).then(lambda: gr.update(interactive=False), None, query_box, queue=False
                ).then(lambda: gr.update(interactive=False), None, add_sess_btn, queue=False
                ).then(lambda: gr.update(interactive=False), None, sess_drpdn, queue=False
                ).then(lambda: gr.update(interactive=False), None, del_sess_btn, queue=False
                ).then(self._prepare, [query_box, chatbot, sess_data], [query_box, chatbot], queue=True
                ).then(self._respond_stream, [chat_use_context, chatbot, stream_output, text_mode] + components,
                                             [chatbot, dbg_msg], queue=chatbot
                ).then(lambda: gr.update(interactive=True), None, query_box, queue=False
                ).then(lambda: gr.update(interactive=True), None, add_sess_btn, queue=False
                ).then(lambda: gr.update(interactive=True), None, sess_drpdn, queue=False
                ).then(lambda: gr.update(interactive=True), None, del_sess_btn, queue=False)
            clear_btn.click(self._clear_history, [sess_data], outputs=[chatbot, query_box, dbg_msg, sess_data])

            sess_drpdn.change(self._change_session, [sess_drpdn, chatbot, dbg_msg, sess_data],
                                                    [sess_drpdn, chatbot, query_box, dbg_msg, sess_data])
            add_sess_btn.click(self._add_session, [chatbot, dbg_msg, sess_data],
                                                  [sess_drpdn, chatbot, query_box, dbg_msg, sess_data])
            del_sess_btn.click(self._delete_session, [sess_drpdn, sess_data],
                                                     [sess_drpdn, chatbot, query_box, dbg_msg, sess_data])
            recordor.change(self._sub_audio, recordor, query_box)
            return demo

stop()

Stop the web interface and clean up resources. If the web demo has been initialized, this method closes the Gradio demo, frees related resources, and resets demo and url attributes.

Source code in lazyllm/tools/webpages/webmodule.py
    def stop(self):
        """Stop the web interface and clean up resources.
If the web demo has been initialized, this method closes the Gradio demo, frees related resources, and resets `demo` and `url` attributes.
"""
        if self.demo:
            self.demo.close()
            del self.demo
            self.demo, self.url = None, ''

wait()

Block the main thread until the web interface is closed. This method blocks the current thread until the Gradio demo is closed. Useful in deployment scenarios to prevent premature program exit.

Source code in lazyllm/tools/webpages/webmodule.py
    def wait(self):
        """Block the main thread until the web interface is closed.
This method blocks the current thread until the Gradio demo is closed. Useful in deployment scenarios to prevent premature program exit.
"""
        self.demo.block_thread()

lazyllm.tools.CodeGenerator

Bases: ModuleBase

Code Generation Module.

This module generates code based on a user-defined prompt. It automatically selects a Chinese or English system prompt based on the input, and extracts Python code snippets from the output.

__init__(self, base_model, prompt="") Initializes the code generator with a base model and prompt.

Parameters:

  • base_model (Union[str, TrainableModule, OnlineChatModuleBase]) –

    A path string to load the model, or an initialized model instance.

  • prompt (str, default: '' ) –

    A user-defined prompt to guide the code generation. May contain Chinese or English.

Examples:

>>> from lazyllm.components import CodeGenerator
>>> generator = CodeGenerator(base_model="deepseek-coder", prompt="写一个Python函数,计算斐波那契数列。")
>>> result = generator("请给出实现代码")
>>> print(result)
... def fibonacci(n):
...     if n <= 1:
...         return n
...     return fibonacci(n-1) + fibonacci(n-2)
Source code in lazyllm/tools/actors/code_generator.py
class CodeGenerator(ModuleBase):
    """Code Generation Module.

This module generates code based on a user-defined prompt. It automatically selects a Chinese or English system prompt based on the input, and extracts Python code snippets from the output.

`__init__(self, base_model, prompt="")`
Initializes the code generator with a base model and prompt.

Args:
    base_model (Union[str, TrainableModule, OnlineChatModuleBase]): A path string to load the model, or an initialized model instance.
    prompt (str): A user-defined prompt to guide the code generation. May contain Chinese or English.


Examples:
    >>> from lazyllm.components import CodeGenerator
    >>> generator = CodeGenerator(base_model="deepseek-coder", prompt="写一个Python函数,计算斐波那契数列。")
    >>> result = generator("请给出实现代码")
    >>> print(result)
    ... def fibonacci(n):
    ...     if n <= 1:
    ...         return n
    ...     return fibonacci(n-1) + fibonacci(n-2)
    """
    def __init__(
        self,
        base_model: Union[str, TrainableModule, OnlineChatModuleBase],
        prompt: str = "",
    ):
        super().__init__()
        self._prompt = self.choose_prompt(prompt).format(prompt=prompt)
        if isinstance(base_model, str):
            self._m = TrainableModule(base_model).start().prompt(self._prompt)
        else:
            self._m = base_model.share(self._prompt)

    def choose_prompt(self, prompt: str):
        """Selects an appropriate code generation prompt template based on the content of the input prompt.  
Returns the Chinese prompt template if Chinese characters are detected; otherwise returns the English prompt template.

Args:
    prompt (str): Input prompt text.

**Returns:**

- str: The selected code generation prompt template string.
"""
        # Use chinese prompt if intent elements have chinese character, otherwise use english version
        for ele in prompt:
            # chinese unicode range
            if "\u4e00" <= ele <= "\u9fff":
                return ch_code_generate_prompt
        return en_code_generate_prompt

    def forward(self, *args, **kw):
        res = self._m(*args, **kw)
        pattern = r"```python(.*?)\n```"
        matches = re.findall(pattern, res, re.DOTALL)
        if len(matches) > 0:
            return matches[0]
        return res

choose_prompt(prompt)

Selects an appropriate code generation prompt template based on the content of the input prompt.
Returns the Chinese prompt template if Chinese characters are detected; otherwise returns the English prompt template.

Parameters:

  • prompt (str) –

    Input prompt text.

Returns:

  • str: The selected code generation prompt template string.
Source code in lazyllm/tools/actors/code_generator.py
    def choose_prompt(self, prompt: str):
        """Selects an appropriate code generation prompt template based on the content of the input prompt.  
Returns the Chinese prompt template if Chinese characters are detected; otherwise returns the English prompt template.

Args:
    prompt (str): Input prompt text.

**Returns:**

- str: The selected code generation prompt template string.
"""
        # Use chinese prompt if intent elements have chinese character, otherwise use english version
        for ele in prompt:
            # chinese unicode range
            if "\u4e00" <= ele <= "\u9fff":
                return ch_code_generate_prompt
        return en_code_generate_prompt

lazyllm.tools.ParameterExtractor

Bases: ModuleBase

Parameter Extraction Module.

This module extracts structured parameters from a given text using a language model, based on the parameter names, types, descriptions, and whether they are required.

__init__(self, base_model, param, type, description, require) Initializes the parameter extractor with the parameter specification and base model.

Parameters:

  • base_model (Union[str, TrainableModule, OnlineChatModuleBase]) –

    A model path or model instance used for extraction.

  • param (list[str]) –

    List of parameter names to extract.

  • type (list[str]) –

    List of parameter types (e.g., "int", "str", "bool").

  • description (list[str]) –

    List of descriptions for each parameter.

  • require (list[bool]) –

    List indicating whether each parameter is required.

Examples:

>>> from lazyllm.components import ParameterExtractor
>>> extractor = ParameterExtractor(
...     base_model="deepseek-chat",
...     param=["name", "age"],
...     type=["str", "int"],
...     description=["The user's name", "The user's age"],
...     require=[True, True]
... )
>>> result = extractor("My name is Alice and I am 25 years old.")
>>> print(result)
... ['Alice', 25]
Source code in lazyllm/tools/actors/parameter_extractor.py
class ParameterExtractor(ModuleBase):
    """Parameter Extraction Module.

This module extracts structured parameters from a given text using a language model, based on the parameter names, types, descriptions, and whether they are required.

`__init__(self, base_model, param, type, description, require)`
Initializes the parameter extractor with the parameter specification and base model.

Args:
    base_model (Union[str, TrainableModule, OnlineChatModuleBase]): A model path or model instance used for extraction.
    param (list[str]): List of parameter names to extract.
    type (list[str]): List of parameter types (e.g., "int", "str", "bool").
    description (list[str]): List of descriptions for each parameter.
    require (list[bool]): List indicating whether each parameter is required.


Examples:
    >>> from lazyllm.components import ParameterExtractor
    >>> extractor = ParameterExtractor(
    ...     base_model="deepseek-chat",
    ...     param=["name", "age"],
    ...     type=["str", "int"],
    ...     description=["The user's name", "The user's age"],
    ...     require=[True, True]
    ... )
    >>> result = extractor("My name is Alice and I am 25 years old.")
    >>> print(result)
    ... ['Alice', 25]
    """
    type_map = {
        int.__name__: int,
        str.__name__: str,
        float.__name__: float,
        bool.__name__: bool,
        list.__name__: list,
        dict.__name__: dict,
    }

    def __init__(
        self,
        base_model: Union[str, TrainableModule, OnlineChatModuleBase],
        param: list[str],
        type: list[str],
        description: list[str],
        require: list[bool],
    ):
        super().__init__()
        assert len(param) == len(type) == len(description) == len(require) > 0
        self._param_dict = {p: ParameterExtractor.type_map[t] for p, t in zip(param, type)}
        param_prompt = repr([dict(name=p, type=t, description=d, require=r)
                             for p, t, d, r in zip(param, type, description, require)])
        self._prompt = self.choose_prompt(param_prompt).format(prompt=param_prompt)
        if isinstance(base_model, str):
            self._m = TrainableModule(base_model).start().prompt(self._prompt)
        else:
            self._m = base_model.share(self._prompt)

    def choose_prompt(self, prompt: str):
        # Use chinese prompt if intent elements have chinese character, otherwise use english version
        for ele in prompt:
            # chinese unicode range
            if "\u4e00" <= ele <= "\u9fff":
                return ch_parameter_extractor_prompt
        return en_parameter_extractor_prompt

    def forward(self, *args, **kw):
        res = self._m(*args, **kw)
        pattern = r"```json(.*?)\n```"
        matches = re.findall(pattern, res, re.DOTALL)
        if len(matches) > 0:
            res = matches[0]
            res.strip()
            try:
                res = json.loads(res)
            except Exception:
                pass
        else:
            res = res.split("\n")
            for param in res:
                try:
                    res = json.loads(param)
                except Exception:
                    continue
                if isinstance(res, dict): break
        if isinstance(res, dict):
            ret = [res.get(param_name, None) for param_name in self._param_dict]
        else:
            ret = [None] * len(self._param_dict)
        ret = package(ret)
        return ret

lazyllm.tools.QustionRewrite

Bases: ModuleBase

Question Rewrite Module.

This module rewrites or reformulates a user query using a language model. It supports both string and list output formats based on the formatter.

__init__(self, base_model, rewrite_prompt="", formatter="str") Initializes the question rewrite module with a prompt and model.

Parameters:

  • base_model (Union[str, TrainableModule, OnlineChatModuleBase]) –

    A path string or initialized model for question rewriting.

  • rewrite_prompt (str, default: '' ) –

    Custom prompt to guide the rewrite behavior.

  • formatter (str, default: 'str' ) –

    Output format type; either "str" or "list".

Examples:

>>> from lazyllm.components import QustionRewrite
>>> rewriter = QustionRewrite(base_model="chatglm", rewrite_prompt="请将问题改写为更适合检索的形式", formatter="list")
>>> result = rewriter("中国的最高山峰是什么?")
>>> print(result)
... ['中国的最高山峰是哪一座?', '中国海拔最高的山是什么?']
Source code in lazyllm/tools/actors/qustion_rewrite.py
class QustionRewrite(ModuleBase):
    """Question Rewrite Module.

This module rewrites or reformulates a user query using a language model. It supports both string and list output formats based on the formatter.

`__init__(self, base_model, rewrite_prompt="", formatter="str")`
Initializes the question rewrite module with a prompt and model.

Args:
    base_model (Union[str, TrainableModule, OnlineChatModuleBase]): A path string or initialized model for question rewriting.
    rewrite_prompt (str): Custom prompt to guide the rewrite behavior.
    formatter (str): Output format type; either "str" or "list".


Examples:
    >>> from lazyllm.components import QustionRewrite
    >>> rewriter = QustionRewrite(base_model="chatglm", rewrite_prompt="请将问题改写为更适合检索的形式", formatter="list")
    >>> result = rewriter("中国的最高山峰是什么?")
    >>> print(result)
    ... ['中国的最高山峰是哪一座?', '中国海拔最高的山是什么?']
    """
    def __init__(
        self,
        base_model: Union[str, TrainableModule, OnlineChatModuleBase],
        rewrite_prompt: str = "",
        formatter: str = "str",
    ):
        super().__init__()
        self._prompt = self.choose_prompt(rewrite_prompt).format(prompt=rewrite_prompt)
        if isinstance(base_model, str):
            self._m = TrainableModule(base_model).start().prompt(self._prompt)
        else:
            self._m = base_model.share(self._prompt)
        self.formatter = formatter

    def choose_prompt(self, prompt: str):
        """
Choose the appropriate prompt template based on the language of the input prompt.

This method analyzes the input prompt string and determines whether to use the Chinese or English prompt template. It checks each character in the prompt string and if any character falls within the Chinese Unicode range (\u4e00-\u9fff), it returns the Chinese prompt template; otherwise, it returns the English prompt template.

Args:
    prompt (str): The input prompt string to be analyzed for language detection.

Returns:
    str: The selected prompt template string (either Chinese or English version).


Examples:

    >>> from lazyllm.tools.actors.qustion_rewrite import QustionRewrite

    # Example 1: English prompt (no Chinese characters)
    >>> rewriter = QustionRewrite("gpt-3.5-turbo")
    >>> prompt_template = rewriter.choose_prompt("How to implement machine learning?")
    >>> print("Template contains Chinese:", "中文" in prompt_template)
    Template contains Chinese: False

    # Example 2: Chinese prompt (contains Chinese characters)
    >>> prompt_template = rewriter.choose_prompt("如何实现机器学习?")
    >>> print("Template contains Chinese:", "中文" in prompt_template)
    Template contains Chinese: True

    # Example 3: Mixed language prompt (contains Chinese characters)
    >>> prompt_template = rewriter.choose_prompt("What is 机器学习?")
    >>> print("Template contains Chinese:", "中文" in prompt_template)
    Template contains Chinese: True
    """
        # Use chinese prompt if intent elements have chinese character, otherwise use english version
        for ele in prompt:
            # chinese unicode range
            if "\u4e00" <= ele <= "\u9fff":
                return ch_qustion_rewrite_prompt
        return en_qustion_rewrite_prompt

    def forward(self, *args, **kw):
        res = self._m(*args, **kw)
        if self.formatter == "list":
            return list(filter(None, res.split('\n')))
        else:
            return res

choose_prompt(prompt)

Choose the appropriate prompt template based on the language of the input prompt.

This method analyzes the input prompt string and determines whether to use the Chinese or English prompt template. It checks each character in the prompt string and if any character falls within the Chinese Unicode range (一-鿿), it returns the Chinese prompt template; otherwise, it returns the English prompt template.

Parameters:

  • prompt (str) –

    The input prompt string to be analyzed for language detection.

Returns:

  • str

    The selected prompt template string (either Chinese or English version).

Examples:

>>> from lazyllm.tools.actors.qustion_rewrite import QustionRewrite

# Example 1: English prompt (no Chinese characters)
>>> rewriter = QustionRewrite("gpt-3.5-turbo")
>>> prompt_template = rewriter.choose_prompt("How to implement machine learning?")
>>> print("Template contains Chinese:", "中文" in prompt_template)
Template contains Chinese: False

# Example 2: Chinese prompt (contains Chinese characters)
>>> prompt_template = rewriter.choose_prompt("如何实现机器学习?")
>>> print("Template contains Chinese:", "中文" in prompt_template)
Template contains Chinese: True

# Example 3: Mixed language prompt (contains Chinese characters)
>>> prompt_template = rewriter.choose_prompt("What is 机器学习?")
>>> print("Template contains Chinese:", "中文" in prompt_template)
Template contains Chinese: True
Source code in lazyllm/tools/actors/qustion_rewrite.py
    def choose_prompt(self, prompt: str):
        """
Choose the appropriate prompt template based on the language of the input prompt.

This method analyzes the input prompt string and determines whether to use the Chinese or English prompt template. It checks each character in the prompt string and if any character falls within the Chinese Unicode range (\u4e00-\u9fff), it returns the Chinese prompt template; otherwise, it returns the English prompt template.

Args:
    prompt (str): The input prompt string to be analyzed for language detection.

Returns:
    str: The selected prompt template string (either Chinese or English version).


Examples:

    >>> from lazyllm.tools.actors.qustion_rewrite import QustionRewrite

    # Example 1: English prompt (no Chinese characters)
    >>> rewriter = QustionRewrite("gpt-3.5-turbo")
    >>> prompt_template = rewriter.choose_prompt("How to implement machine learning?")
    >>> print("Template contains Chinese:", "中文" in prompt_template)
    Template contains Chinese: False

    # Example 2: Chinese prompt (contains Chinese characters)
    >>> prompt_template = rewriter.choose_prompt("如何实现机器学习?")
    >>> print("Template contains Chinese:", "中文" in prompt_template)
    Template contains Chinese: True

    # Example 3: Mixed language prompt (contains Chinese characters)
    >>> prompt_template = rewriter.choose_prompt("What is 机器学习?")
    >>> print("Template contains Chinese:", "中文" in prompt_template)
    Template contains Chinese: True
    """
        # Use chinese prompt if intent elements have chinese character, otherwise use english version
        for ele in prompt:
            # chinese unicode range
            if "\u4e00" <= ele <= "\u9fff":
                return ch_qustion_rewrite_prompt
        return en_qustion_rewrite_prompt

lazyllm.tools.agent.toolsManager.ToolManager

Bases: ModuleBase

ToolManager is a tool management class used to provide tool information and tool calls to function call.

When constructing this management class, you need to pass in a list of tool name strings. The tool name here can be provided by LazyLLM or user-defined. If it is user-defined, it must first be registered in LazyLLM before it can be used. When registering, directly use the fc_register registrar, which has established the tool group, so when using the tool management class, all functions can be uniformly registered in the tool group. The function to be registered needs to annotate the function parameters, and add a functional description to the function, as well as the parameter type and function description. This is to facilitate the tool management class to parse the function and pass it to LLM for use.

Parameters:

  • tools (List[str]) –

    A list of tool name strings.

  • return_trace (bool, default: False ) –

    If True, return intermediate steps and tool calls.

  • stream (bool) –

    Whether to stream the planning and solving process.

Examples:

>>> from lazyllm.tools import ToolManager, fc_register
>>> import json
>>> from typing import Literal
>>> @fc_register("tool")
>>> def get_current_weather(location: str, unit: Literal["fahrenheit", "celsius"]="fahrenheit"):
...     '''
...     Get the current weather in a given location
...
...     Args:
...         location (str): The city and state, e.g. San Francisco, CA.
...         unit (str): The temperature unit to use. Infer this from the users location.
...     '''
...     if 'tokyo' in location.lower():
...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius'})
...     elif 'san francisco' in location.lower():
...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit'})
...     elif 'paris' in location.lower():
...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius'})
...     elif 'beijing' in location.lower():
...         return json.dumps({'location': 'Beijing', 'temperature': '90', 'unit': 'fahrenheit'})
...     else:
...         return json.dumps({'location': location, 'temperature': 'unknown'})
...
>>> @fc_register("tool")
>>> def get_n_day_weather_forecast(location: str, num_days: int, unit: Literal["celsius", "fahrenheit"]='fahrenheit'):
...     '''
...     Get an N-day weather forecast
...
...     Args:
...         location (str): The city and state, e.g. San Francisco, CA.
...         num_days (int): The number of days to forecast.
...         unit (Literal['celsius', 'fahrenheit']): The temperature unit to use. Infer this from the users location.
...     '''
...     if 'tokyo' in location.lower():
...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius', "num_days": num_days})
...     elif 'san francisco' in location.lower():
...         return json.dumps({'location': 'San Francisco', 'temperature': '75', 'unit': 'fahrenheit', "num_days": num_days})
...     elif 'paris' in location.lower():
...         return json.dumps({'location': 'Paris', 'temperature': '25', 'unit': 'celsius', "num_days": num_days})
...     elif 'beijing' in location.lower():
...         return json.dumps({'location': 'Beijing', 'temperature': '85', 'unit': 'fahrenheit', "num_days": num_days})
...     else:
...         return json.dumps({'location': location, 'temperature': 'unknown'})
...
>>> tools = ["get_current_weather", "get_n_day_weather_forecast"]
>>> tm = ToolManager(tools)
>>> print(tm([{'name': 'get_n_day_weather_forecast', 'arguments': {'location': 'Beijing', 'num_days': 3}}])[0])
'{"location": "Beijing", "temperature": "85", "unit": "fahrenheit", "num_days": 3}'
Source code in lazyllm/tools/agent/toolsManager.py
class ToolManager(ModuleBase):
    """ToolManager is a tool management class used to provide tool information and tool calls to function call.

When constructing this management class, you need to pass in a list of tool name strings. The tool name here can be provided by LazyLLM or user-defined. If it is user-defined, it must first be registered in LazyLLM before it can be used. When registering, directly use the `fc_register` registrar, which has established the `tool` group, so when using the tool management class, all functions can be uniformly registered in the `tool` group. The function to be registered needs to annotate the function parameters, and add a functional description to the function, as well as the parameter type and function description. This is to facilitate the tool management class to parse the function and pass it to LLM for use.

Args:
    tools (List[str]): A list of tool name strings.
    return_trace (bool): If True, return intermediate steps and tool calls.
    stream (bool): Whether to stream the planning and solving process.


Examples:
    >>> from lazyllm.tools import ToolManager, fc_register
    >>> import json
    >>> from typing import Literal
    >>> @fc_register("tool")
    >>> def get_current_weather(location: str, unit: Literal["fahrenheit", "celsius"]="fahrenheit"):
    ...     '''
    ...     Get the current weather in a given location
    ...
    ...     Args:
    ...         location (str): The city and state, e.g. San Francisco, CA.
    ...         unit (str): The temperature unit to use. Infer this from the users location.
    ...     '''
    ...     if 'tokyo' in location.lower():
    ...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius'})
    ...     elif 'san francisco' in location.lower():
    ...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit'})
    ...     elif 'paris' in location.lower():
    ...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius'})
    ...     elif 'beijing' in location.lower():
    ...         return json.dumps({'location': 'Beijing', 'temperature': '90', 'unit': 'fahrenheit'})
    ...     else:
    ...         return json.dumps({'location': location, 'temperature': 'unknown'})
    ...
    >>> @fc_register("tool")
    >>> def get_n_day_weather_forecast(location: str, num_days: int, unit: Literal["celsius", "fahrenheit"]='fahrenheit'):
    ...     '''
    ...     Get an N-day weather forecast
    ...
    ...     Args:
    ...         location (str): The city and state, e.g. San Francisco, CA.
    ...         num_days (int): The number of days to forecast.
    ...         unit (Literal['celsius', 'fahrenheit']): The temperature unit to use. Infer this from the users location.
    ...     '''
    ...     if 'tokyo' in location.lower():
    ...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius', "num_days": num_days})
    ...     elif 'san francisco' in location.lower():
    ...         return json.dumps({'location': 'San Francisco', 'temperature': '75', 'unit': 'fahrenheit', "num_days": num_days})
    ...     elif 'paris' in location.lower():
    ...         return json.dumps({'location': 'Paris', 'temperature': '25', 'unit': 'celsius', "num_days": num_days})
    ...     elif 'beijing' in location.lower():
    ...         return json.dumps({'location': 'Beijing', 'temperature': '85', 'unit': 'fahrenheit', "num_days": num_days})
    ...     else:
    ...         return json.dumps({'location': location, 'temperature': 'unknown'})
    ...
    >>> tools = ["get_current_weather", "get_n_day_weather_forecast"]
    >>> tm = ToolManager(tools)
    >>> print(tm([{'name': 'get_n_day_weather_forecast', 'arguments': {'location': 'Beijing', 'num_days': 3}}])[0])
    '{"location": "Beijing", "temperature": "85", "unit": "fahrenheit", "num_days": 3}'
    """
    def __init__(self, tools: List[Union[str, Callable]], return_trace: bool = False):
        super().__init__(return_trace=return_trace)
        self._tools = self._load_tools(tools)
        self._format_tools()
        self._tools_desc = self._transform_to_openai_function()

    def _load_tools(self, tools: List[Union[str, Callable]]):
        if "tmp_tool" not in LazyLLMRegisterMetaClass.all_clses:
            register.new_group('tmp_tool')

        _tools = []
        for element in tools:
            if isinstance(element, str):
                _tools.append(getattr(lazyllm.tool, element)())
            elif isinstance(element, Callable):
                # just to convert `element` to the internal type in `Register`
                register('tmp_tool')(element)
                _tools.append(getattr(lazyllm.tmp_tool, element.__name__)())
                lazyllm.tmp_tool.remove(element.__name__)

        return _tools

    @property
    def all_tools(self):
        return self._tools

    @property
    def tools_description(self):
        return self._tools_desc

    @property
    def tools_info(self):
        return self._tool_call

    def _validate_tool(self, tool_name: str, tool_arguments: Dict[str, Any]):
        tool = self._tool_call.get(tool_name)
        if not tool:
            LOG.error(f'cannot find tool named [{tool_name}]')
            return False

        return tool.validate_parameters(tool_arguments)

    def _format_tools(self):
        if isinstance(self._tools, List):
            self._tool_call = {tool.name: tool for tool in self._tools}

    @staticmethod
    def _gen_args_info_from_moduletool_and_docstring(tool, parsed_docstring):
        """
        returns a dict of param names containing at least
          1. `type`
          2. `description` of params

        for example:
            args = {
                "foo": {
                    "enum": ["baz", "bar"],
                    "type": "string",
                    "description": "a string",
                },
                "bar": {
                    "type": "integer",
                    "description": "an integer",
                }
            }
        """
        tool_args = tool.args
        assert len(tool_args) == len(parsed_docstring.params), ("The parameter description and the actual "
                                                                "number of input parameters are inconsistent.")

        args_description = {}
        for param in parsed_docstring.params:
            args_description[param.arg_name] = param.description

        args = {}
        for k, v in tool_args.items():
            val = copy.deepcopy(v)
            val.pop("title", None)
            val.pop("default", None)
            args[k] = val if val else {"type": "string"}
            desc = args_description.get(k, None)
            if desc:
                args[k].update({"description": desc})
            else:
                raise ValueError(f"The actual input parameter '{k}' is not found "
                                 f"in the parameter description of tool '{tool.name}'.")
        return args

    def _transform_to_openai_function(self):
        if not isinstance(self._tools, List):
            raise TypeError(f"The tools type should be List instead of {type(self._tools)}")

        format_tools = []
        for tool in self._tools:
            try:
                parsed_docstring = docstring_parser.parse(tool.description)
                args = self._gen_args_info_from_moduletool_and_docstring(tool, parsed_docstring)
                required_arg_list = tool.params_schema.model_json_schema().get("required", [])
                func = {
                    "type": "function",
                    "function": {
                        "name": tool.name,
                        "description": parsed_docstring.short_description,
                        "parameters": {
                            "type": "object",
                            "properties": args,
                            "required": required_arg_list,
                        }
                    }
                }
                format_tools.append(func)
            except Exception:
                typehints_template = """
                def myfunc(arg1: str, arg2: Dict[str, Any], arg3: Literal["aaa", "bbb", "ccc"]="aaa"):
                    '''
                    Function description ...

                    Args:
                        arg1 (str): arg1 description.
                        arg2 (Dict[str, Any]): arg2 description
                        arg3 (Literal["aaa", "bbb", "ccc"]): arg3 description
                    '''
                """
                raise TypeError("Function description must include function description and "
                                f"parameter description, the format is as follows: {typehints_template}")
        return format_tools

    def forward(self, tools: Union[Dict[str, Any], List[Dict[str, Any]]], verbose: bool = False):
        tool_calls = [tools,] if isinstance(tools, dict) else tools
        tool_calls = [{"name": tool['name'], "arguments": json.loads(tool['arguments'])
                      if isinstance(tool['arguments'], str) else tool['arguments']} for tool in tool_calls]
        output = []
        flag_val = [True if self._validate_tool(tool['name'], tool['arguments']) else False for tool in tool_calls]
        tool_inputs = [tool_calls[idx]['arguments'] for idx, val in enumerate(flag_val) if val]
        tools = [self._tool_call[tool_calls[idx]['name']] for idx, val in enumerate(flag_val) if val]
        tool_diverter = lazyllm.diverter(tuple(tools))
        rets = tool_diverter(tuple(tool_inputs))
        res = iter(rets)
        rets = [next(res) if ele else None for ele in flag_val]
        for idx, tool in enumerate(tool_calls):
            if flag_val[idx]:
                ret = rets[idx]
                output.append(json.dumps(ret, ensure_ascii=False) if not isinstance(ret, str) else ret)
            else:
                output.append(f"{tool} parameters error.")

        return output

lazyllm.tools.ModuleTool

Bases: ModuleBase

Base class for defining tools using callable Python functions.

This class automatically parses function signatures and docstrings to build a parameter schema using pydantic. It also performs input validation and handles standardized tool execution.

__init__(self, verbose=False, return_trace=True) Initializes a tool wrapper module.

Parameters:

  • verbose (bool, default: False ) –

    Whether to print verbose logs during execution.

  • return_trace (bool, default: True ) –

    Whether to keep intermediate execution trace in the result.

Examples:

>>> from lazyllm.components import ModuleTool
>>> class AddTool(ModuleTool):
...     def apply(self, a: int, b: int) -> int:
...         '''Add two integers.
...         
...         Args:
...             a (int): First number.
...             b (int): Second number.
...         
...         Returns:
...             int: The sum of a and b.
...         '''
...         return a + b
>>> tool = AddTool()
>>> result = tool({'a': 3, 'b': 5})
>>> print(result)
8
Source code in lazyllm/tools/agent/toolsManager.py
class ModuleTool(ModuleBase, metaclass=LazyLLMRegisterMetaClass):
    """Base class for defining tools using callable Python functions.

This class automatically parses function signatures and docstrings to build a parameter schema using `pydantic`. It also performs input validation and handles standardized tool execution.

`__init__(self, verbose=False, return_trace=True)`
Initializes a tool wrapper module.

Args:
    verbose (bool): Whether to print verbose logs during execution.
    return_trace (bool): Whether to keep intermediate execution trace in the result.


Examples:

    >>> from lazyllm.components import ModuleTool
    >>> class AddTool(ModuleTool):
    ...     def apply(self, a: int, b: int) -> int:
    ...         '''Add two integers.
    ...         
    ...         Args:
    ...             a (int): First number.
    ...             b (int): Second number.
    ...         
    ...         Returns:
    ...             int: The sum of a and b.
    ...         '''
    ...         return a + b
    >>> tool = AddTool()
    >>> result = tool({'a': 3, 'b': 5})
    >>> print(result)
    8
    """
    def __init__(self, verbose: bool = False, return_trace: bool = True):
        super().__init__(return_trace=return_trace)
        self._verbose = verbose
        self._name = self.apply.__name__\
            if hasattr(self.apply, "__name__") and self.apply.__name__ is not None\
            else (_ for _ in ()).throw(ValueError("Function must have a name."))
        self._description = self.apply.__doc__\
            if hasattr(self.apply, "__doc__") and self.apply.__doc__ is not None\
            else (_ for _ in ()).throw(ValueError("Function must have a docstring"))
        # strip space(s) and newlines before and after docstring, as RewooAgent requires
        self._description = self._description.strip(' \n')

        self._params_schema = self._load_function_schema(self.__class__.apply)

    def _load_function_schema(self, func: Callable) -> Type[BaseModel]:
        parsed_docstring = docstring_parser.parse(self._description)
        func_str_from_doc = _gen_empty_func_str_from_parsed_docstring(parsed_docstring)
        func_from_doc = _gen_func_from_str(func_str_from_doc, self._description)
        func_from_doc.__name__ = func.__name__
        doc_type_hints = get_type_hints(func_from_doc, globals(), locals())

        func_type_hints = get_type_hints(func, globals(), locals())

        _check_return_type_is_the_same(doc_type_hints, func_type_hints)

        signature = inspect.signature(func)
        has_var_args = False
        for _, param in signature.parameters.items():
            if param.kind == inspect.Parameter.VAR_POSITIONAL or\
               param.kind == inspect.Parameter.VAR_KEYWORD:
                has_var_args = True
                break

        if has_var_args:
            # we cannot get type hints from var args, so we get them from docstring
            self._type_hints = doc_type_hints
            signature = inspect.signature(func_from_doc)
        else:
            self._type_hints = func_type_hints
            # accomplish type_hints from docstring
            for name, type in doc_type_hints.items():
                self._type_hints.setdefault(name, type)

        self._return_type = self._type_hints.get('return') if self._type_hints else None

        fields = {
            name: (self._type_hints.get(name, Any), param.default
                   if param.default is not inspect.Parameter.empty
                   else ...)
            for name, param in signature.parameters.items()
        }

        return create_model(self._name, **fields)

    @property
    def name(self):
        return self._name

    @property
    def description(self):
        return self._description

    @property
    def params_schema(self) -> Type[BaseModel]:
        return self._params_schema

    @property
    def args(self) -> Dict[str, Any]:
        return self._params_schema.model_json_schema()["properties"]

    @property
    def required_args(self) -> Set[str]:
        return set(self._params_schema.model_json_schema().get("required", []))

    def apply(self, *args: Any, **kwargs: Any) -> Any:
        """
Abstract method to be implemented in subclasses.

This method should perform a specific task based on the provided arguments.

Raises:
    NotImplementedError: If the method is not overridden in a subclass.
"""
        raise NotImplementedError("Implement apply function in subclass")

    def _validate_input(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        input_params = self._params_schema
        if isinstance(tool_input, dict):
            if input_params is not None:
                ret = input_params.model_validate(tool_input)
                return {key: getattr(ret, key) for key in ret.model_dump().keys() if key in tool_input}
            return tool_input
        elif isinstance(tool_input, str):
            if input_params is not None:
                key = next(iter(input_params.model_fields.keys()))
                input_params.model_validate({key: tool_input})
                arg_type = self._type_hints.get(key)
                if arg_type:
                    return {key: arg_type(tool_input)}
                return {key: tool_input}

            if len(self._type_hints) != 1:
                return tool_input
            arg_type = self._type_hints.values()[0]
            return arg_type(tool_input)
        else:
            raise TypeError(f"tool_input {tool_input} only supports dict and str.")

    def validate_parameters(self, arguments: Dict[str, Any]) -> bool:
        """
Validate whether the provided arguments meet the required criteria.

This method checks if all required keys are present in the input dictionary and attempts format validation.

Args:
    arguments (Dict[str, Any]): Dictionary of input arguments.

Returns:
    bool: True if valid and complete; False otherwise.
"""
        if len(self.required_args.difference(set(arguments.keys()))) == 0:
            # contains all required parameters
            try:
                self._validate_input(arguments)
                return True
            except ValidationError:
                return False
        return False

    def forward(self, tool_input: Union[str, Dict[str, Any]], verbose: bool = False) -> Any:
        val_input = self._validate_input(tool_input)
        if isinstance(val_input, dict):
            ret = self.apply(**val_input)
        else:
            ret = self.apply(val_input)
        if verbose or self._verbose:
            lazyllm.LOG.debug(f"The output of tool {self.name} is {ret}")

        return ret

apply(*args, **kwargs)

Abstract method to be implemented in subclasses.

This method should perform a specific task based on the provided arguments.

Raises:

  • NotImplementedError

    If the method is not overridden in a subclass.

Source code in lazyllm/tools/agent/toolsManager.py
    def apply(self, *args: Any, **kwargs: Any) -> Any:
        """
Abstract method to be implemented in subclasses.

This method should perform a specific task based on the provided arguments.

Raises:
    NotImplementedError: If the method is not overridden in a subclass.
"""
        raise NotImplementedError("Implement apply function in subclass")

validate_parameters(arguments)

Validate whether the provided arguments meet the required criteria.

This method checks if all required keys are present in the input dictionary and attempts format validation.

Parameters:

  • arguments (Dict[str, Any]) –

    Dictionary of input arguments.

Returns:

  • bool ( bool ) –

    True if valid and complete; False otherwise.

Source code in lazyllm/tools/agent/toolsManager.py
    def validate_parameters(self, arguments: Dict[str, Any]) -> bool:
        """
Validate whether the provided arguments meet the required criteria.

This method checks if all required keys are present in the input dictionary and attempts format validation.

Args:
    arguments (Dict[str, Any]): Dictionary of input arguments.

Returns:
    bool: True if valid and complete; False otherwise.
"""
        if len(self.required_args.difference(set(arguments.keys()))) == 0:
            # contains all required parameters
            try:
                self._validate_input(arguments)
                return True
            except ValidationError:
                return False
        return False

lazyllm.tools.FunctionCall

Bases: ModuleBase

FunctionCall is a single-turn tool invocation class. It is used when the LLM alone cannot answer user queries and requires external knowledge through tool calls.
If the LLM output requires tool calls, the tools are invoked and the combined results (input, model output, tool output) are returned as a list.
If no tool calls are needed, the LLM output is returned directly as a string.

Parameters:

  • llm (ModuleBase) –

    The LLM instance to use, which can be either a TrainableModule or OnlineChatModule.

  • tools (List[Union[str, Callable]]) –

    A list of tool names or callable objects that the LLM can use.

  • return_trace (Optional[bool], default: False ) –

    Whether to return the invocation trace, defaults to False.

  • stream (Optional[bool], default: False ) –

    Whether to enable streaming output, defaults to False.

  • _prompt (Optional[str], default: None ) –

    Custom prompt for function call, defaults to automatic selection based on llm type.

Note: Tools in tools must include a __doc__ attribute and describe their purpose and parameters according to the Google Python Style.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import fc_register, FunctionCall
>>> import json
>>> from typing import Literal
>>> @fc_register("tool")
>>> def get_current_weather(location: str, unit: Literal["fahrenheit", "celsius"] = 'fahrenheit'):
...     '''
...     Get the current weather in a given location
...
...     Args:
...         location (str): The city and state, e.g. San Francisco, CA.
...         unit (str): The temperature unit to use. Infer this from the users location.
...     '''
...     if 'tokyo' in location.lower():
...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius'})
...     elif 'san francisco' in location.lower():
...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit'})
...     elif 'paris' in location.lower():
...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius'})
...     else:
...         return json.dumps({'location': location, 'temperature': 'unknown'})
...
>>> @fc_register("tool")
>>> def get_n_day_weather_forecast(location: str, num_days: int, unit: Literal["celsius", "fahrenheit"] = 'fahrenheit'):
...     '''
...     Get an N-day weather forecast
...
...     Args:
...         location (str): The city and state, e.g. San Francisco, CA.
...         num_days (int): The number of days to forecast.
...         unit (Literal['celsius', 'fahrenheit']): The temperature unit to use. Infer this from the users location.
...     '''
...     if 'tokyo' in location.lower():
...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius', "num_days": num_days})
...     elif 'san francisco' in location.lower():
...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit', "num_days": num_days})
...     elif 'paris' in location.lower():
...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius', "num_days": num_days})
...     else:
...         return json.dumps({'location': location, 'temperature': 'unknown'})
...
>>> tools=["get_current_weather", "get_n_day_weather_forecast"]
>>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()  # or llm = lazyllm.OnlineChatModule("openai", stream=False)
>>> query = "What's the weather like today in celsius in Tokyo."
>>> fc = FunctionCall(llm, tools)
>>> ret = fc(query)
>>> print(ret)
["What's the weather like today in celsius in Tokyo.", {'role': 'assistant', 'content': '
', 'tool_calls': [{'id': 'da19cddac0584869879deb1315356d2a', 'type': 'function', 'function': {'name': 'get_current_weather', 'arguments': {'location': 'Tokyo', 'unit': 'celsius'}}}]}, [{'role': 'tool', 'content': '{"location": "Tokyo", "temperature": "10", "unit": "celsius"}', 'tool_call_id': 'da19cddac0584869879deb1315356d2a', 'name': 'get_current_weather'}]]
>>> query = "Hello"
>>> ret = fc(query)
>>> print(ret)
'Hello! How can I assist you today?'
Source code in lazyllm/tools/agent/functionCall.py
class FunctionCall(ModuleBase):
    """FunctionCall is a single-turn tool invocation class. It is used when the LLM alone cannot answer user queries and requires external knowledge through tool calls.  
If the LLM output requires tool calls, the tools are invoked and the combined results (input, model output, tool output) are returned as a list.  
If no tool calls are needed, the LLM output is returned directly as a string.

Args:
    llm (ModuleBase): The LLM instance to use, which can be either a TrainableModule or OnlineChatModule.
    tools (List[Union[str, Callable]]): A list of tool names or callable objects that the LLM can use.
    return_trace (Optional[bool]): Whether to return the invocation trace, defaults to False.
    stream (Optional[bool]): Whether to enable streaming output, defaults to False.
    _prompt (Optional[str]): Custom prompt for function call, defaults to automatic selection based on llm type.

Note: Tools in `tools` must include a `__doc__` attribute and describe their purpose and parameters according to the [Google Python Style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings).


Examples:
    >>> import lazyllm
    >>> from lazyllm.tools import fc_register, FunctionCall
    >>> import json
    >>> from typing import Literal
    >>> @fc_register("tool")
    >>> def get_current_weather(location: str, unit: Literal["fahrenheit", "celsius"] = 'fahrenheit'):
    ...     '''
    ...     Get the current weather in a given location
    ...
    ...     Args:
    ...         location (str): The city and state, e.g. San Francisco, CA.
    ...         unit (str): The temperature unit to use. Infer this from the users location.
    ...     '''
    ...     if 'tokyo' in location.lower():
    ...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius'})
    ...     elif 'san francisco' in location.lower():
    ...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit'})
    ...     elif 'paris' in location.lower():
    ...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius'})
    ...     else:
    ...         return json.dumps({'location': location, 'temperature': 'unknown'})
    ...
    >>> @fc_register("tool")
    >>> def get_n_day_weather_forecast(location: str, num_days: int, unit: Literal["celsius", "fahrenheit"] = 'fahrenheit'):
    ...     '''
    ...     Get an N-day weather forecast
    ...
    ...     Args:
    ...         location (str): The city and state, e.g. San Francisco, CA.
    ...         num_days (int): The number of days to forecast.
    ...         unit (Literal['celsius', 'fahrenheit']): The temperature unit to use. Infer this from the users location.
    ...     '''
    ...     if 'tokyo' in location.lower():
    ...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius', "num_days": num_days})
    ...     elif 'san francisco' in location.lower():
    ...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit', "num_days": num_days})
    ...     elif 'paris' in location.lower():
    ...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius', "num_days": num_days})
    ...     else:
    ...         return json.dumps({'location': location, 'temperature': 'unknown'})
    ...
    >>> tools=["get_current_weather", "get_n_day_weather_forecast"]
    >>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()  # or llm = lazyllm.OnlineChatModule("openai", stream=False)
    >>> query = "What's the weather like today in celsius in Tokyo."
    >>> fc = FunctionCall(llm, tools)
    >>> ret = fc(query)
    >>> print(ret)
    ["What's the weather like today in celsius in Tokyo.", {'role': 'assistant', 'content': '
    ', 'tool_calls': [{'id': 'da19cddac0584869879deb1315356d2a', 'type': 'function', 'function': {'name': 'get_current_weather', 'arguments': {'location': 'Tokyo', 'unit': 'celsius'}}}]}, [{'role': 'tool', 'content': '{"location": "Tokyo", "temperature": "10", "unit": "celsius"}', 'tool_call_id': 'da19cddac0584869879deb1315356d2a', 'name': 'get_current_weather'}]]
    >>> query = "Hello"
    >>> ret = fc(query)
    >>> print(ret)
    'Hello! How can I assist you today?'
    """

    def __init__(self, llm, tools: List[Union[str, Callable]], *, return_trace: bool = False,
                 stream: bool = False, _prompt: str = None):
        super().__init__(return_trace=return_trace)
        if isinstance(llm, OnlineChatModule) and llm.series == "QWEN" and llm._stream is True:
            raise ValueError("The qwen platform does not currently support stream function calls.")
        if _prompt is None:
            _prompt = FC_PROMPT_ONLINE if isinstance(llm, OnlineChatModule) else FC_PROMPT_LOCAL

        self._tools_manager = ToolManager(tools, return_trace=return_trace)
        self._prompter = ChatPrompter(instruction=_prompt, tools=self._tools_manager.tools_description)\
            .pre_hook(function_call_hook)
        self._llm = llm.share(prompt=self._prompter, format=FunctionCallFormatter()).used_by(self._module_id)
        with pipeline() as self._impl:
            self._impl.ins = StreamResponse('Received instruction:', prefix_color=Color.yellow,
                                            color=Color.green, stream=stream)
            self._impl.m1 = self._llm
            self._impl.m2 = self._parser
            self._impl.dis = StreamResponse('Decision-making or result in this round:',
                                            prefix_color=Color.yellow, color=Color.green, stream=stream)
            self._impl.m3 = ifs(lambda x: isinstance(x, list),
                                pipeline(self._tools_manager, StreamResponse('Tool-Call result:',
                                         prefix_color=Color.yellow, color=Color.green, stream=stream)),
                                lambda out: out)
            self._impl.m4 = self._tool_post_action | bind(input=self._impl.input, llm_output=self._impl.m1)

    def _parser(self, llm_output: Union[str, List[Dict[str, Any]]]):
        LOG.debug(f"llm_output: {llm_output}")
        if isinstance(llm_output, list):
            res = []
            for item in llm_output:
                if isinstance(item, str):
                    continue
                arguments = item.get('function', {}).get('arguments', '')
                arguments = json.loads(arguments) if isinstance(arguments, str) else arguments
                res.append({"name": item.get('function', {}).get('name', ''), 'arguments': arguments})
            return res
        elif isinstance(llm_output, str):
            return llm_output
        else:
            raise TypeError(f"The {llm_output} type currently is only supports `list` and `str`,"
                            f" and does not support {type(llm_output)}.")

    def _tool_post_action(self, output: Union[str, List[str]], input: Union[str, List],
                          llm_output: List[Dict[str, Any]]):
        if isinstance(output, list):
            ret = []
            if isinstance(input, str):
                ret.append(input)
            elif isinstance(input, list):
                ret.append(input[-1])
            else:
                raise TypeError(f"The input type currently only supports `str` and `list`, "
                                f"and does not support {type(input)}.")

            content = "".join([item for item in llm_output if isinstance(item, str)])
            llm_output = [item for item in llm_output if not isinstance(item, str)]
            ret.append({"role": "assistant", "content": content, "tool_calls": llm_output})
            ret.append([{"role": "tool", "content": out, "tool_call_id": llm_output[idx]["id"],
                         "name": llm_output[idx]["function"]["name"]}
                        for idx, out in enumerate(output)])
            LOG.debug(f"functionCall result: {ret}")
            return ret
        elif isinstance(output, str):
            return output
        else:
            raise TypeError(f"The {output} type currently is only supports `list` and `str`,"
                            f" and does not support {type(output)}.")

    def forward(self, input: str, llm_chat_history: List[Dict[str, Any]] = None):
        globals['chat_history'].setdefault(self._llm._module_id, [])
        if llm_chat_history is not None:
            globals['chat_history'][self._llm._module_id] = llm_chat_history
        return self._impl(input)

lazyllm.tools.FunctionCallFormatter

Bases: JsonFormatter

Formatter for parsing structured function call messages.

This class extends JsonFormatter and is responsible for extracting JSON-based tool call structures from a mixed message string, optionally separating them using a global delimiter.

Private Method

_load(msg) Parses the input message string and extracts JSON-formatted tool calls, if present.

Examples:

>>> from lazyllm.components import FunctionCallFormatter
>>> formatter = FunctionCallFormatter()
>>> msg = "Please call this tool. <TOOL> [{"name": "search", "args": {"query": "weather"}}]"
>>> result = formatter._load(msg)
>>> print(result)
... [{'name': 'search', 'args': {'query': 'weather'}}, 'Please call this tool. ']
Source code in lazyllm/tools/agent/functionCallFormatter.py
class FunctionCallFormatter(JsonFormatter):
    """Formatter for parsing structured function call messages.

This class extends `JsonFormatter` and is responsible for extracting JSON-based tool call structures from a mixed message string, optionally separating them using a global delimiter.

Private Method:
    _load(msg)
        Parses the input message string and extracts JSON-formatted tool calls, if present.


Examples:
    >>> from lazyllm.components import FunctionCallFormatter
    >>> formatter = FunctionCallFormatter()
    >>> msg = "Please call this tool. <TOOL> [{\"name\": \"search\", \"args\": {\"query\": \"weather\"}}]"
    >>> result = formatter._load(msg)
    >>> print(result)
    ... [{'name': 'search', 'args': {'query': 'weather'}}, 'Please call this tool. ']
    """
    def _load(self, msg: str):
        if "{" not in msg:
            return msg
        if globals['tool_delimiter'] in msg:
            content, msg = msg.split(globals['tool_delimiter'])
            assert msg.count("{") == msg.count("}"), f"{msg} is not a valid json string."
            try:
                json_strs = json.loads(msg)
                res = []
                for json_str in json_strs:
                    res.append(json_str)
                if content:
                    res.append(content)
                return res
            except Exception:
                return msg

        return msg

lazyllm.tools.FunctionCallAgent

Bases: ModuleBase

FunctionCallAgent is an agent that uses the tool calling method to perform complete tool calls. That is, when answering uesr questions, if LLM needs to obtain external knowledge through the tool, it will call the tool and feed back the return results of the tool to LLM, which will finally summarize and output them.

Parameters:

  • llm (ModuleBase) –

    The LLM to be used can be either TrainableModule or OnlineChatModule.

  • tools (List[str]) –

    A list of tool names for LLM to use.

  • max_retries (int, default: 5 ) –

    The maximum number of tool call iterations. The default value is 5.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import fc_register, FunctionCallAgent
>>> import json
>>> from typing import Literal
>>> @fc_register("tool")
>>> def get_current_weather(location: str, unit: Literal["fahrenheit", "celsius"]='fahrenheit'):
...     '''
...     Get the current weather in a given location
...
...     Args:
...         location (str): The city and state, e.g. San Francisco, CA.
...         unit (str): The temperature unit to use. Infer this from the users location.
...     '''
...     if 'tokyo' in location.lower():
...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius'})
...     elif 'san francisco' in location.lower():
...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit'})
...     elif 'paris' in location.lower():
...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius'})
...     elif 'beijing' in location.lower():
...         return json.dumps({'location': 'Beijing', 'temperature': '90', 'unit': 'Fahrenheit'})
...     else:
...         return json.dumps({'location': location, 'temperature': 'unknown'})
...
>>> @fc_register("tool")
>>> def get_n_day_weather_forecast(location: str, num_days: int, unit: Literal["celsius", "fahrenheit"]='fahrenheit'):
...     '''
...     Get an N-day weather forecast
...
...     Args:
...         location (str): The city and state, e.g. San Francisco, CA.
...         num_days (int): The number of days to forecast.
...         unit (Literal['celsius', 'fahrenheit']): The temperature unit to use. Infer this from the users location.
...     '''
...     if 'tokyo' in location.lower():
...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius', "num_days": num_days})
...     elif 'san francisco' in location.lower():
...         return json.dumps({'location': 'San Francisco', 'temperature': '75', 'unit': 'fahrenheit', "num_days": num_days})
...     elif 'paris' in location.lower():
...         return json.dumps({'location': 'Paris', 'temperature': '25', 'unit': 'celsius', "num_days": num_days})
...     elif 'beijing' in location.lower():
...         return json.dumps({'location': 'Beijing', 'temperature': '85', 'unit': 'fahrenheit', "num_days": num_days})
...     else:
...         return json.dumps({'location': location, 'temperature': 'unknown'})
...
>>> tools = ['get_current_weather', 'get_n_day_weather_forecast']
>>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()  # or llm = lazyllm.OnlineChatModule(source="sensenova")
>>> agent = FunctionCallAgent(llm, tools)
>>> query = "What's the weather like today in celsius in Tokyo and Paris."
>>> res = agent(query)
>>> print(res)
'The current weather in Tokyo is 10 degrees Celsius, and in Paris, it is 22 degrees Celsius.'
>>> query = "Hello"
>>> res = agent(query)
>>> print(res)
'Hello! How can I assist you today?'
Source code in lazyllm/tools/agent/functionCall.py
class FunctionCallAgent(ModuleBase):
    """FunctionCallAgent is an agent that uses the tool calling method to perform complete tool calls. That is, when answering uesr questions, if LLM needs to obtain external knowledge through the tool, it will call the tool and feed back the return results of the tool to LLM, which will finally summarize and output them.

Args:
    llm (ModuleBase): The LLM to be used can be either TrainableModule or OnlineChatModule.
    tools (List[str]): A list of tool names for LLM to use.
    max_retries (int): The maximum number of tool call iterations. The default value is 5.


Examples:
    >>> import lazyllm
    >>> from lazyllm.tools import fc_register, FunctionCallAgent
    >>> import json
    >>> from typing import Literal
    >>> @fc_register("tool")
    >>> def get_current_weather(location: str, unit: Literal["fahrenheit", "celsius"]='fahrenheit'):
    ...     '''
    ...     Get the current weather in a given location
    ...
    ...     Args:
    ...         location (str): The city and state, e.g. San Francisco, CA.
    ...         unit (str): The temperature unit to use. Infer this from the users location.
    ...     '''
    ...     if 'tokyo' in location.lower():
    ...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius'})
    ...     elif 'san francisco' in location.lower():
    ...         return json.dumps({'location': 'San Francisco', 'temperature': '72', 'unit': 'fahrenheit'})
    ...     elif 'paris' in location.lower():
    ...         return json.dumps({'location': 'Paris', 'temperature': '22', 'unit': 'celsius'})
    ...     elif 'beijing' in location.lower():
    ...         return json.dumps({'location': 'Beijing', 'temperature': '90', 'unit': 'Fahrenheit'})
    ...     else:
    ...         return json.dumps({'location': location, 'temperature': 'unknown'})
    ...
    >>> @fc_register("tool")
    >>> def get_n_day_weather_forecast(location: str, num_days: int, unit: Literal["celsius", "fahrenheit"]='fahrenheit'):
    ...     '''
    ...     Get an N-day weather forecast
    ...
    ...     Args:
    ...         location (str): The city and state, e.g. San Francisco, CA.
    ...         num_days (int): The number of days to forecast.
    ...         unit (Literal['celsius', 'fahrenheit']): The temperature unit to use. Infer this from the users location.
    ...     '''
    ...     if 'tokyo' in location.lower():
    ...         return json.dumps({'location': 'Tokyo', 'temperature': '10', 'unit': 'celsius', "num_days": num_days})
    ...     elif 'san francisco' in location.lower():
    ...         return json.dumps({'location': 'San Francisco', 'temperature': '75', 'unit': 'fahrenheit', "num_days": num_days})
    ...     elif 'paris' in location.lower():
    ...         return json.dumps({'location': 'Paris', 'temperature': '25', 'unit': 'celsius', "num_days": num_days})
    ...     elif 'beijing' in location.lower():
    ...         return json.dumps({'location': 'Beijing', 'temperature': '85', 'unit': 'fahrenheit', "num_days": num_days})
    ...     else:
    ...         return json.dumps({'location': location, 'temperature': 'unknown'})
    ...
    >>> tools = ['get_current_weather', 'get_n_day_weather_forecast']
    >>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()  # or llm = lazyllm.OnlineChatModule(source="sensenova")
    >>> agent = FunctionCallAgent(llm, tools)
    >>> query = "What's the weather like today in celsius in Tokyo and Paris."
    >>> res = agent(query)
    >>> print(res)
    'The current weather in Tokyo is 10 degrees Celsius, and in Paris, it is 22 degrees Celsius.'
    >>> query = "Hello"
    >>> res = agent(query)
    >>> print(res)
    'Hello! How can I assist you today?'
    """
    def __init__(self, llm, tools: List[str], max_retries: int = 5, return_trace: bool = False, stream: bool = False):
        super().__init__(return_trace=return_trace)
        self._max_retries = max_retries
        self._fc = FunctionCall(llm, tools, return_trace=return_trace, stream=stream)
        self._agent = loop(self._fc, stop_condition=lambda x: isinstance(x, str), count=self._max_retries)
        self._fc._llm.used_by(self._module_id)

    def forward(self, query: str, llm_chat_history: List[Dict[str, Any]] = None):
        ret = self._agent(query, llm_chat_history) if llm_chat_history is not None else self._agent(query)
        return ret if isinstance(ret, str) else (_ for _ in ()).throw(
            ValueError(f"After retrying {self._max_retries} times, the function call agent still "
                       "failed to call successfully."))

lazyllm.tools.ReactAgent

Bases: ModuleBase

ReactAgent follows the process of Thought->Action->Observation->Thought...->Finish step by step through LLM and tool calls to display the steps to solve user questions and the final answer to the user.

Parameters:

  • llm (ModuleBase) –

    The LLM to be used can be either TrainableModule or OnlineChatModule.

  • tools (List[str]) –

    A list of tool names for LLM to use.

  • max_retries (int, default: 5 ) –

    The maximum number of tool call iterations. The default value is 5.

  • return_trace (bool, default: False ) –

    If True, return intermediate steps and tool calls.

  • stream (bool, default: False ) –

    Whether to stream the planning and solving process.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import fc_register, ReactAgent
>>> @fc_register("tool")
>>> def multiply_tool(a: int, b: int) -> int:
...     '''
...     Multiply two integers and return the result integer
...
...     Args:
...         a (int): multiplier
...         b (int): multiplier
...     '''
...     return a * b
...
>>> @fc_register("tool")
>>> def add_tool(a: int, b: int):
...     '''
...     Add two integers and returns the result integer
...
...     Args:
...         a (int): addend
...         b (int): addend
...     '''
...     return a + b
...
>>> tools = ["multiply_tool", "add_tool"]
>>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()   # or llm = lazyllm.OnlineChatModule(source="sensenova")
>>> agent = ReactAgent(llm, tools)
>>> query = "What is 20+(2*4)? Calculate step by step."
>>> res = agent(query)
>>> print(res)
'Answer: The result of 20+(2*4) is 28.'
Source code in lazyllm/tools/agent/reactAgent.py
class ReactAgent(ModuleBase):
    """ReactAgent follows the process of `Thought->Action->Observation->Thought...->Finish` step by step through LLM and tool calls to display the steps to solve user questions and the final answer to the user.

Args:
    llm (ModuleBase): The LLM to be used can be either TrainableModule or OnlineChatModule.
    tools (List[str]): A list of tool names for LLM to use.
    max_retries (int): The maximum number of tool call iterations. The default value is 5.
    return_trace (bool): If True, return intermediate steps and tool calls.
    stream (bool): Whether to stream the planning and solving process.


Examples:
    >>> import lazyllm
    >>> from lazyllm.tools import fc_register, ReactAgent
    >>> @fc_register("tool")
    >>> def multiply_tool(a: int, b: int) -> int:
    ...     '''
    ...     Multiply two integers and return the result integer
    ...
    ...     Args:
    ...         a (int): multiplier
    ...         b (int): multiplier
    ...     '''
    ...     return a * b
    ...
    >>> @fc_register("tool")
    >>> def add_tool(a: int, b: int):
    ...     '''
    ...     Add two integers and returns the result integer
    ...
    ...     Args:
    ...         a (int): addend
    ...         b (int): addend
    ...     '''
    ...     return a + b
    ...
    >>> tools = ["multiply_tool", "add_tool"]
    >>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()   # or llm = lazyllm.OnlineChatModule(source="sensenova")
    >>> agent = ReactAgent(llm, tools)
    >>> query = "What is 20+(2*4)? Calculate step by step."
    >>> res = agent(query)
    >>> print(res)
    'Answer: The result of 20+(2*4) is 28.'
    """
    def __init__(self, llm, tools: List[str], max_retries: int = 5, return_trace: bool = False,
                 prompt: str = None, stream: bool = False):
        super().__init__(return_trace=return_trace)
        self._max_retries = max_retries
        assert llm and tools, "llm and tools cannot be empty."

        if not prompt:
            prompt = INSTRUCTION.replace("{TOKENIZED_PROMPT}", WITHOUT_TOKEN_PROMPT if isinstance(llm, OnlineChatModule)
                                         else WITH_TOKEN_PROMPT)
            prompt = prompt.replace("{tool_names}", json.dumps([t.__name__ if callable(t) else t for t in tools],
                                                               ensure_ascii=False))
        self._agent = loop(FunctionCall(llm, tools, _prompt=prompt, return_trace=return_trace, stream=stream),
                           stop_condition=lambda x: isinstance(x, str), count=self._max_retries)

    def forward(self, query: str, llm_chat_history: List[Dict[str, Any]] = None):
        ret = self._agent(query, llm_chat_history) if llm_chat_history is not None else self._agent(query)
        return ret if isinstance(ret, str) else (_ for _ in ()).throw(ValueError(f"After retrying \
            {self._max_retries} times, the function call agent still failes to call successfully."))

lazyllm.tools.PlanAndSolveAgent

Bases: ModuleBase

PlanAndSolveAgent consists of two components. First, the planner breaks down the entire task into smaller subtasks, then the solver executes these subtasks according to the plan, which may involve tool calls, and finally returns the answer to the user.

Parameters:

  • llm (ModuleBase, default: None ) –

    The LLM to be used can be TrainableModule or OnlineChatModule. It is mutually exclusive with plan_llm and solve_llm. Either set llm(the planner and sovler share the same LLM), or set plan_llm and solve_llm,or only specify llm(to set the planner) and solve_llm. Other cases are considered invalid.

  • tools (List[str], default: [] ) –

    A list of tool names for LLM to use.

  • plan_llm (ModuleBase, default: None ) –

    The LLM to be used by the planner, which can be either TrainableModule or OnlineChatModule.

  • solve_llm (ModuleBase, default: None ) –

    The LLM to be used by the solver, which can be either TrainableModule or OnlineChatModule.

  • max_retries (int, default: 5 ) –

    The maximum number of tool call iterations. The default value is 5.

  • return_trace (bool, default: False ) –

    If True, return intermediate steps and tool calls.

  • stream (bool, default: False ) –

    Whether to stream the planning and solving process.

Examples:

>>> import lazyllm
>>> from lazyllm.tools import fc_register, PlanAndSolveAgent
>>> @fc_register("tool")
>>> def multiply(a: int, b: int) -> int:
...     '''
...     Multiply two integers and return the result integer
...
...     Args:
...         a (int): multiplier
...         b (int): multiplier
...     '''
...     return a * b
...
>>> @fc_register("tool")
>>> def add(a: int, b: int):
...     '''
...     Add two integers and returns the result integer
...
...     Args:
...         a (int): addend
...         b (int): addend
...     '''
...     return a + b
...
>>> tools = ["multiply", "add"]
>>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()  # or llm = lazyllm.OnlineChatModule(source="sensenova")
>>> agent = PlanAndSolveAgent(llm, tools)
>>> query = "What is 20+(2*4)? Calculate step by step."
>>> res = agent(query)
>>> print(res)
'The final answer is 28.'
Source code in lazyllm/tools/agent/planAndSolveAgent.py
class PlanAndSolveAgent(ModuleBase):
    """PlanAndSolveAgent consists of two components. First, the planner breaks down the entire task into smaller subtasks, then the solver executes these subtasks according to the plan, which may involve tool calls, and finally returns the answer to the user.

Args:
    llm (ModuleBase): The LLM to be used can be TrainableModule or OnlineChatModule. It is mutually exclusive with plan_llm and solve_llm. Either set llm(the planner and sovler share the same LLM), or set plan_llm and solve_llm,or only specify llm(to set the planner) and solve_llm. Other cases are considered invalid.
    tools (List[str]): A list of tool names for LLM to use.
    plan_llm (ModuleBase): The LLM to be used by the planner, which can be either TrainableModule or OnlineChatModule.
    solve_llm (ModuleBase): The LLM to be used by the solver, which can be either TrainableModule or OnlineChatModule.
    max_retries (int): The maximum number of tool call iterations. The default value is 5.
    return_trace (bool): If True, return intermediate steps and tool calls.
    stream (bool): Whether to stream the planning and solving process.


Examples:
    >>> import lazyllm
    >>> from lazyllm.tools import fc_register, PlanAndSolveAgent
    >>> @fc_register("tool")
    >>> def multiply(a: int, b: int) -> int:
    ...     '''
    ...     Multiply two integers and return the result integer
    ...
    ...     Args:
    ...         a (int): multiplier
    ...         b (int): multiplier
    ...     '''
    ...     return a * b
    ...
    >>> @fc_register("tool")
    >>> def add(a: int, b: int):
    ...     '''
    ...     Add two integers and returns the result integer
    ...
    ...     Args:
    ...         a (int): addend
    ...         b (int): addend
    ...     '''
    ...     return a + b
    ...
    >>> tools = ["multiply", "add"]
    >>> llm = lazyllm.TrainableModule("internlm2-chat-20b").start()  # or llm = lazyllm.OnlineChatModule(source="sensenova")
    >>> agent = PlanAndSolveAgent(llm, tools)
    >>> query = "What is 20+(2*4)? Calculate step by step."
    >>> res = agent(query)
    >>> print(res)
    'The final answer is 28.'
    """
    def __init__(self, llm: Union[ModuleBase, None] = None, tools: List[str] = [], *,  # noqa B006
                 plan_llm: Union[ModuleBase, None] = None, solve_llm: Union[ModuleBase, None] = None,
                 max_retries: int = 5, return_trace: bool = False, stream: bool = False):
        super().__init__(return_trace=return_trace)
        self._max_retries = max_retries
        assert (llm is None and plan_llm and solve_llm) or (llm and plan_llm is None), 'Either specify only llm \
               without specify plan and solve, or specify only plan and solve without specifying llm, or specify \
               both llm and solve. Other situations are not allowed.'
        assert tools, "tools cannot be empty."
        s = dict(prefix='I will give a plan first:\n', prefix_color=Color.blue, color=Color.green) if stream else False
        self._plan_llm = ((plan_llm or llm).share(prompt=ChatPrompter(instruction=PLANNER_PROMPT),
                                                  stream=s).used_by(self._module_id))
        self._solve_llm = (solve_llm or llm).share().used_by(self._module_id)
        self._tools = tools
        with pipeline() as self._agent:
            self._agent.plan = self._plan_llm
            self._agent.parse = (lambda text, query: package([], '', [v for v in re.split("\n\\s*\\d+\\. ", text)[1:]],
                                 query)) | bind(query=self._agent.input)
            with loop(stop_condition=lambda pre, res, steps, query: len(steps) == 0) as self._agent.lp:
                self._agent.lp.pre_action = self._pre_action
                self._agent.lp.solve = FunctionCallAgent(self._solve_llm, tools=self._tools,
                                                         return_trace=return_trace, stream=stream)
                self._agent.lp.post_action = self._post_action | bind(self._agent.lp.input[0][0], _0,
                                                                      self._agent.lp.input[0][2],
                                                                      self._agent.lp.input[0][3])

            self._agent.post_action = lambda pre, res, steps, query: res

    def _pre_action(self, pre_steps, response, steps, query):
        result = package(SOLVER_PROMPT.format(previous_steps="\n".join(pre_steps), current_step=steps[0],
                                              objective=query) + "input: " + response + "\n" + steps[0], [])
        return result

    def _post_action(self, pre_steps: List[str], response: str, steps: List[str], query: str):
        LOG.debug(f"current step: {steps[0]}, response: {response}")
        pre_steps.append(steps.pop(0))
        return package(pre_steps, response, steps, query)

    def forward(self, query: str):
        return self._agent(query)

lazyllm.tools.ReWOOAgent

Bases: ModuleBase

ReWOOAgent consists of three parts: Planer, Worker and Solver. The Planner uses predictive reasoning capabilities to create a solution blueprint for a complex task; the Worker interacts with the environment through tool calls and fills in actual evidence or observations into instructions; the Solver processes all plans and evidence to develop a solution to the original task or problem.

Parameters:

  • llm (ModuleBase, default: None ) –

    The LLM to be used can be TrainableModule or OnlineChatModule. It is mutually exclusive with plan_llm and solve_llm. Either set llm(the planner and sovler share the same LLM), or set plan_llm and solve_llm,or only specify llm(to set the planner) and solve_llm. Other cases are considered invalid.

  • tools (List[str], default: [] ) –

    A list of tool names for LLM to use.

  • plan_llm (ModuleBase, default: None ) –

    The LLM to be used by the planner, which can be either TrainableModule or OnlineChatModule.

  • solve_llm (ModuleBase, default: None ) –

    The LLM to be used by the solver, which can be either TrainableModule or OnlineChatModule.

  • max_retries (int) –

    The maximum number of tool call iterations. The default value is 5.

  • return_trace (bool, default: False ) –

    If True, return intermediate steps and tool calls.

  • stream (bool, default: False ) –

    Whether to stream the planning and solving process.

Examples:

>>> import lazyllm
>>> import wikipedia
>>> from lazyllm.tools import fc_register, ReWOOAgent
>>> @fc_register("tool")
>>> def WikipediaWorker(input: str):
...     '''
...     Worker that search for similar page contents from Wikipedia. Useful when you need to get holistic knowledge about people, places, companies, historical events, or other subjects. The response are long and might contain some irrelevant information. Input should be a search query.
...
...     Args:
...         input (str): search query.
...     '''
...     try:
...         evidence = wikipedia.page(input).content
...         evidence = evidence.split("\n\n")[0]
...     except wikipedia.PageError:
...         evidence = f"Could not find [{input}]. Similar: {wikipedia.search(input)}"
...     except wikipedia.DisambiguationError:
...         evidence = f"Could not find [{input}]. Similar: {wikipedia.search(input)}"
...     return evidence
...
>>> @fc_register("tool")
>>> def LLMWorker(input: str):
...     '''
...     A pretrained LLM like yourself. Useful when you need to act with general world knowledge and common sense. Prioritize it when you are confident in solving the problem yourself. Input can be any instruction.
...
...     Args:
...         input (str): instruction
...     '''
...     llm = lazyllm.OnlineChatModule(source="glm")
...     query = f"Respond in short directly with no extra words.\n\n{input}"
...     response = llm(query, llm_chat_history=[])
...     return response
...
>>> tools = ["WikipediaWorker", "LLMWorker"]
>>> llm = lazyllm.TrainableModule("GLM-4-9B-Chat").deploy_method(lazyllm.deploy.vllm).start()  # or llm = lazyllm.OnlineChatModule(source="sensenova")
>>> agent = ReWOOAgent(llm, tools)
>>> query = "What is the name of the cognac house that makes the main ingredient in The Hennchata?"
>>> res = agent(query)
>>> print(res)
'
Hennessy '
Source code in lazyllm/tools/agent/rewooAgent.py
class ReWOOAgent(ModuleBase):
    """ReWOOAgent consists of three parts: Planer, Worker and Solver. The Planner uses predictive reasoning capabilities to create a solution blueprint for a complex task; the Worker interacts with the environment through tool calls and fills in actual evidence or observations into instructions; the Solver processes all plans and evidence to develop a solution to the original task or problem.

Args:
    llm (ModuleBase): The LLM to be used can be TrainableModule or OnlineChatModule. It is mutually exclusive with plan_llm and solve_llm. Either set llm(the planner and sovler share the same LLM), or set plan_llm and solve_llm,or only specify llm(to set the planner) and solve_llm. Other cases are considered invalid.
    tools (List[str]): A list of tool names for LLM to use.
    plan_llm (ModuleBase): The LLM to be used by the planner, which can be either TrainableModule or OnlineChatModule.
    solve_llm (ModuleBase): The LLM to be used by the solver, which can be either TrainableModule or OnlineChatModule.
    max_retries (int): The maximum number of tool call iterations. The default value is 5.
    return_trace (bool): If True, return intermediate steps and tool calls.
    stream (bool): Whether to stream the planning and solving process.


Examples:
    >>> import lazyllm
    >>> import wikipedia
    >>> from lazyllm.tools import fc_register, ReWOOAgent
    >>> @fc_register("tool")
    >>> def WikipediaWorker(input: str):
    ...     '''
    ...     Worker that search for similar page contents from Wikipedia. Useful when you need to get holistic knowledge about people, places, companies, historical events, or other subjects. The response are long and might contain some irrelevant information. Input should be a search query.
    ...
    ...     Args:
    ...         input (str): search query.
    ...     '''
    ...     try:
    ...         evidence = wikipedia.page(input).content
    ...         evidence = evidence.split("\\n\\n")[0]
    ...     except wikipedia.PageError:
    ...         evidence = f"Could not find [{input}]. Similar: {wikipedia.search(input)}"
    ...     except wikipedia.DisambiguationError:
    ...         evidence = f"Could not find [{input}]. Similar: {wikipedia.search(input)}"
    ...     return evidence
    ...
    >>> @fc_register("tool")
    >>> def LLMWorker(input: str):
    ...     '''
    ...     A pretrained LLM like yourself. Useful when you need to act with general world knowledge and common sense. Prioritize it when you are confident in solving the problem yourself. Input can be any instruction.
    ...
    ...     Args:
    ...         input (str): instruction
    ...     '''
    ...     llm = lazyllm.OnlineChatModule(source="glm")
    ...     query = f"Respond in short directly with no extra words.\\n\\n{input}"
    ...     response = llm(query, llm_chat_history=[])
    ...     return response
    ...
    >>> tools = ["WikipediaWorker", "LLMWorker"]
    >>> llm = lazyllm.TrainableModule("GLM-4-9B-Chat").deploy_method(lazyllm.deploy.vllm).start()  # or llm = lazyllm.OnlineChatModule(source="sensenova")
    >>> agent = ReWOOAgent(llm, tools)
    >>> query = "What is the name of the cognac house that makes the main ingredient in The Hennchata?"
    >>> res = agent(query)
    >>> print(res)
    '
    Hennessy '
    """
    def __init__(self, llm: Union[ModuleBase, None] = None, tools: List[Union[str, Callable]] = [], *,  # noqa B006
                 plan_llm: Union[ModuleBase, None] = None, solve_llm: Union[ModuleBase, None] = None,
                 return_trace: bool = False, stream: bool = False):
        super().__init__(return_trace=return_trace)
        assert (llm is None and plan_llm and solve_llm) or (llm and plan_llm is None), 'Either specify only llm \
               without specify plan and solve, or specify only plan and solve without specifying llm, or specify \
               both llm and solve. Other situations are not allowed.'
        assert tools, "tools cannot be empty."
        self._planner = (plan_llm or llm).share(stream=dict(
            prefix='\nI will give a plan first:\n', prefix_color=Color.blue, color=Color.green) if stream else False)
        self._solver = (solve_llm or llm).share(stream=dict(
            prefix='\nI will solve the problem:\n', prefix_color=Color.blue, color=Color.green) if stream else False)
        self._name2tool = ToolManager(tools, return_trace=return_trace).tools_info
        with pipeline() as self._agent:
            self._agent.planner_pre_action = self._build_planner_prompt
            self._agent.planner = self._planner
            self._agent.parse_plan = self._parse_plan
            self._agent.woker = self._get_worker_evidences
            self._agent.solver_pre_action = self._build_solver_prompt | bind(input=self._agent.input)
            self._agent.solver = self._solver

    def _build_planner_prompt(self, input: str):
        prompt = P_PROMPT_PREFIX + "Tools can be one of the following:\n"
        for name, tool in self._name2tool.items():
            prompt += f"{name}[search query]: {tool.description}\n"
        prompt += P_FEWSHOT + "\n" + P_PROMPT_SUFFIX + input + "\n"
        globals['chat_history'][self._planner._module_id] = []
        return prompt

    def _parse_plan(self, response: str):
        LOG.debug(f"planner plans: {response}")
        plans = []
        evidence = {}
        for line in response.splitlines():
            if line.startswith("Plan"):
                plans.append(line)
            elif line.startswith("#") and line[1] == "E" and line[2].isdigit():
                e, tool_call = line.split("=", 1)
                e, tool_call = e.strip(), tool_call.strip()
                if len(e) == 3:
                    evidence[e] = tool_call
                else:
                    evidence[e] = "No evidence found"
        return package(plans, evidence)

    def _get_worker_evidences(self, plans: List[str], evidence: Dict[str, str]):
        worker_evidences = {}
        for e, tool_call in evidence.items():
            if "[" not in tool_call:
                worker_evidences[e] = tool_call
                continue
            tool, tool_input = tool_call.split("[", 1)
            tool_input = tool_input[:-1].strip("'").strip('"')
            # find variables in input and replace with previous evidences
            for var in re.findall(r"#E\d+", tool_input):
                if var in worker_evidences:
                    tool_input = tool_input.replace(var, "[" + worker_evidences[var] + "]")
            tool_instance = self._name2tool.get(tool)
            if tool_instance:
                worker_evidences[e] = tool_instance(tool_input)
            else:
                worker_evidences[e] = "No evidence found"

        worker_log = ""
        for idx, plan in enumerate(plans):
            e = f"#E{idx+1}"
            worker_log += f"{plan}\nEvidence:\n{worker_evidences[e]}\n"
        LOG.debug(f"worker_log: {worker_log}")
        return worker_log

    def _build_solver_prompt(self, worker_log, input):
        prompt = S_PROMPT_PREFIX + input + "\n" + worker_log + S_PROMPT_SUFFIX + input + "\n"
        globals['chat_history'][self._solver._module_id] = []
        return prompt

    def forward(self, query: str):
        return self._agent(query)

lazyllm.tools.rag.smart_embedding_index.SmartEmbeddingIndex

Bases: IndexBase

Source code in lazyllm/tools/rag/smart_embedding_index.py
class SmartEmbeddingIndex(IndexBase):
    def __init__(self, backend_type: str, **kwargs):
        if backend_type == 'milvus':
            self._store = MilvusStore(**kwargs)
        elif backend_type == 'map':
            self._store = MapStore(**kwargs)
        else:
            raise ValueError(f'unsupported backend [{backend_type}]')

    @override
    def update(self, nodes: List[DocNode]) -> None:
        self._store.update_nodes(nodes)

    @override
    def remove(self, uids: List[str], group_name: Optional[str] = None) -> None:
        self._store.remove_nodes(uids=uids)

    @override
    def query(self, *args, **kwargs) -> List[DocNode]:
        return self._store.query(*args, **kwargs)

lazyllm.tools.rag.doc_node.ImageDocNode

Bases: DocNode

A specialized document node for handling image content in RAG systems.

ImageDocNode extends DocNode to provide specialized functionality for image processing and embedding generation. It automatically handles image loading, base64 encoding for embedding, and PIL Image objects for LLM processing.

Parameters:

  • image_path (str) –

    The file path to the image file. This should be a valid path to an image file (e.g., .jpg, .png, .jpeg).

  • uid (Optional[str], default: None ) –

    Unique identifier for the document node. If not provided, a UUID will be automatically generated.

  • group (Optional[str], default: None ) –

    The group name this node belongs to. Used for organizing and filtering nodes.

  • embedding (Optional[Dict[str, List[float]]], default: None ) –

    Pre-computed embeddings for the image. Keys are embedding model names, values are embedding vectors.

  • parent (Optional[DocNode], default: None ) –

    Parent node in the document hierarchy. Used for building document trees.

  • metadata (Optional[Dict[str, Any]], default: None ) –

    Additional metadata associated with the image node.

  • global_metadata (Optional[Dict[str, Any]], default: None ) –

    Global metadata that applies to all nodes in the document.

  • text (Optional[str], default: None ) –

    Optional text description or caption for the image.

Examples:

>>> from lazyllm.tools.rag.doc_node import ImageDocNode, MetadataMode
>>> import numpy as np
>>> image_node = ImageDocNode(
...     image_path="/home/mnt/yehongfei/Code/Test/framework.jpg",
...     text="这是一张照片"
)
>>> def clip_emb(content, modality="image"):
...     if modality == "image":
...         return [np.random.rand(512).tolist()]
...     return [np.random.rand(256).tolist()]
>>> embed_functions = {"clip": clip_emb}
>>> image_node.do_embedding(embed_functions)
>>> print(f"嵌入维度: {len(image_node.embedding['clip'])}")
>>> text_representation = image_node.get_text()
>>> content_representation = image_node.get_content(MetadataMode.EMBED)
>>> print(f"text属性: {text_representation}")
>>> print(f"content属性: {content_representation}")
Source code in lazyllm/tools/rag/doc_node.py
class ImageDocNode(DocNode):
    """A specialized document node for handling image content in RAG systems.

ImageDocNode extends DocNode to provide specialized functionality for image processing and embedding generation. It automatically handles image loading, base64 encoding for embedding, and PIL Image objects for LLM processing.

Args:
    image_path (str): The file path to the image file. This should be a valid path to an image file (e.g., .jpg, .png, .jpeg).
    uid (Optional[str]): Unique identifier for the document node. If not provided, a UUID will be automatically generated.
    group (Optional[str]): The group name this node belongs to. Used for organizing and filtering nodes.
    embedding (Optional[Dict[str, List[float]]]): Pre-computed embeddings for the image. Keys are embedding model names, values are embedding vectors.
    parent (Optional[DocNode]): Parent node in the document hierarchy. Used for building document trees.
    metadata (Optional[Dict[str, Any]]): Additional metadata associated with the image node.
    global_metadata (Optional[Dict[str, Any]]): Global metadata that applies to all nodes in the document.
    text (Optional[str]): Optional text description or caption for the image.


Examples:
    >>> from lazyllm.tools.rag.doc_node import ImageDocNode, MetadataMode
    >>> import numpy as np
    >>> image_node = ImageDocNode(
    ...     image_path="/home/mnt/yehongfei/Code/Test/framework.jpg",
    ...     text="这是一张照片"
    )
    >>> def clip_emb(content, modality="image"):
    ...     if modality == "image":
    ...         return [np.random.rand(512).tolist()]
    ...     return [np.random.rand(256).tolist()]
    >>> embed_functions = {"clip": clip_emb}
    >>> image_node.do_embedding(embed_functions)
    >>> print(f"嵌入维度: {len(image_node.embedding['clip'])}")
    >>> text_representation = image_node.get_text()
    >>> content_representation = image_node.get_content(MetadataMode.EMBED)
    >>> print(f"text属性: {text_representation}")
    >>> print(f"content属性: {content_representation}")    
    """
    def __init__(self, image_path: str, uid: Optional[str] = None, group: Optional[str] = None,
                 embedding: Optional[Dict[str, List[float]]] = None, parent: Optional["DocNode"] = None,
                 metadata: Optional[Dict[str, Any]] = None, global_metadata: Optional[Dict[str, Any]] = None,
                 *, text: Optional[str] = None):
        super().__init__(uid, None, group, embedding, parent, metadata, global_metadata=global_metadata, text=text)
        self._image_path = image_path.strip()
        self._modality = 'image'

    def do_embedding(self, embed: Dict[str, Callable]) -> None:
        """Generate embeddings for the image using the provided embedding functions.

This method overrides the parent class method to handle image-specific embedding generation. It automatically converts the image to the appropriate format (base64 for embedding) and calls the embedding functions with the image modality.

Args:
    embed (Dict[str, Callable]): Dictionary of embedding functions. Keys are embedding model names, values are callable functions that accept (content, modality) and return embedding vectors.
"""
        for k, e in embed.items():
            emb = e(self.get_content(MetadataMode.EMBED), modality=self._modality)
            generate_embed = {k: emb[0]}

        with self._lock:
            self.embedding = self.embedding or {}
            self.embedding = {**self.embedding, **generate_embed}

    def get_content(self, metadata_mode=MetadataMode.LLM) -> str:
        """Get the image content in different formats based on the metadata mode.

This method returns the image content in different formats depending on the intended use case. For LLM processing, it returns a PIL Image object. For embedding generation, it returns a base64-encoded image string.

Args:
    metadata_mode (MetadataMode, optional): The mode for content retrieval. Defaults to MetadataMode.LLM.
        - MetadataMode.LLM: Returns PIL Image object for LLM processing
        - MetadataMode.EMBED: Returns base64-encoded image for embedding generation
        - Other modes: Returns the image path as text

**Returns:**

- Union[PIL.Image.Image, List[str], str]: The image content in the requested format.
"""
        if metadata_mode == MetadataMode.LLM:
            return Image.open(self._image_path)
        elif metadata_mode == MetadataMode.EMBED:
            image_base64, mime = _image_to_base64(self._image_path)
            return [f"data:{mime};base64,{image_base64}"]
        else:
            return self.get_text()

    @property
    def image_path(self):
        return self._image_path

    def get_text(self) -> str:  # Disable access to self._content
        """Get the image path as text representation.

This method overrides the parent class method to return the image path instead of the content field, since ImageDocNode doesn't use the content field for storing text.

**Returns:**

- str: The image file path.
"""
        return self._image_path

    @property
    def text(self) -> str:  # Disable access to self._content
        return self._image_path

do_embedding(embed)

Generate embeddings for the image using the provided embedding functions.

This method overrides the parent class method to handle image-specific embedding generation. It automatically converts the image to the appropriate format (base64 for embedding) and calls the embedding functions with the image modality.

Parameters:

  • embed (Dict[str, Callable]) –

    Dictionary of embedding functions. Keys are embedding model names, values are callable functions that accept (content, modality) and return embedding vectors.

Source code in lazyllm/tools/rag/doc_node.py
    def do_embedding(self, embed: Dict[str, Callable]) -> None:
        """Generate embeddings for the image using the provided embedding functions.

This method overrides the parent class method to handle image-specific embedding generation. It automatically converts the image to the appropriate format (base64 for embedding) and calls the embedding functions with the image modality.

Args:
    embed (Dict[str, Callable]): Dictionary of embedding functions. Keys are embedding model names, values are callable functions that accept (content, modality) and return embedding vectors.
"""
        for k, e in embed.items():
            emb = e(self.get_content(MetadataMode.EMBED), modality=self._modality)
            generate_embed = {k: emb[0]}

        with self._lock:
            self.embedding = self.embedding or {}
            self.embedding = {**self.embedding, **generate_embed}

get_content(metadata_mode=MetadataMode.LLM)

Get the image content in different formats based on the metadata mode.

This method returns the image content in different formats depending on the intended use case. For LLM processing, it returns a PIL Image object. For embedding generation, it returns a base64-encoded image string.

Parameters:

  • metadata_mode (MetadataMode, default: LLM ) –

    The mode for content retrieval. Defaults to MetadataMode.LLM. - MetadataMode.LLM: Returns PIL Image object for LLM processing - MetadataMode.EMBED: Returns base64-encoded image for embedding generation - Other modes: Returns the image path as text

Returns:

  • Union[PIL.Image.Image, List[str], str]: The image content in the requested format.
Source code in lazyllm/tools/rag/doc_node.py
    def get_content(self, metadata_mode=MetadataMode.LLM) -> str:
        """Get the image content in different formats based on the metadata mode.

This method returns the image content in different formats depending on the intended use case. For LLM processing, it returns a PIL Image object. For embedding generation, it returns a base64-encoded image string.

Args:
    metadata_mode (MetadataMode, optional): The mode for content retrieval. Defaults to MetadataMode.LLM.
        - MetadataMode.LLM: Returns PIL Image object for LLM processing
        - MetadataMode.EMBED: Returns base64-encoded image for embedding generation
        - Other modes: Returns the image path as text

**Returns:**

- Union[PIL.Image.Image, List[str], str]: The image content in the requested format.
"""
        if metadata_mode == MetadataMode.LLM:
            return Image.open(self._image_path)
        elif metadata_mode == MetadataMode.EMBED:
            image_base64, mime = _image_to_base64(self._image_path)
            return [f"data:{mime};base64,{image_base64}"]
        else:
            return self.get_text()

get_text()

Get the image path as text representation.

This method overrides the parent class method to return the image path instead of the content field, since ImageDocNode doesn't use the content field for storing text.

Returns:

  • str: The image file path.
Source code in lazyllm/tools/rag/doc_node.py
    def get_text(self) -> str:  # Disable access to self._content
        """Get the image path as text representation.

This method overrides the parent class method to return the image path instead of the content field, since ImageDocNode doesn't use the content field for storing text.

**Returns:**

- str: The image file path.
"""
        return self._image_path

lazyllm.tools.rag.transform.AdaptiveTransform

Bases: NodeTransform

A flexible document transformation system that applies different transforms based on document patterns.

AdaptiveTransform allows you to define multiple transformation strategies and automatically selects the appropriate one based on the document's file path or custom pattern matching. This is particularly useful when you have different types of documents that require different processing approaches.

Parameters:

  • transforms (Union[List[Union[TransformArgs, Dict]], Union[TransformArgs, Dict]]) –

    A list of transform configurations or a single transform configuration.

  • num_workers (int, default: 0 ) –

    Number of worker threads for parallel processing. Defaults to 0.

Examples:

>>> from lazyllm.tools.rag.transform import AdaptiveTransform, DocNode, SentenceSplitter
>>> doc1 = DocNode(text="这是第一个文档的内容。它包含多个句子。")
>>> doc2 = DocNode(text="这是第二个文档的内容。")
>>> transforms = [
...     {
...         'f': SentenceSplitter,
...         'pattern': '*.txt',
...         'kwargs': {'chunk_size': 50, 'chunk_overlap': 10}
...     },
...     {
...         'f': SentenceSplitter,
...         'pattern': '*.pdf',
...         'kwargs': {'chunk_size': 100, 'chunk_overlap': 20}
...     }
... ]
>>> adaptive = AdaptiveTransform(transforms)
>>> results1 = adaptive.transform(doc1)
>>> print(f"文档1转换结果: {len(results1)} 个块")
>>> for i, result in enumerate(results1):
...     print(f"  块 {i+1}: {result.text}")
>>> results2 = adaptive.transform(doc2)
>>> print(f"文档2转换结果: {len(results2)} 个块")
>>> for i, result in enumerate(results2):
...     print(f"  块 {i+1}: {result.text}")
Source code in lazyllm/tools/rag/transform.py
class AdaptiveTransform(NodeTransform):
    """A flexible document transformation system that applies different transforms based on document patterns.

AdaptiveTransform allows you to define multiple transformation strategies and automatically selects the appropriate one based on the document's file path or custom pattern matching. This is particularly useful when you have different types of documents that require different processing approaches.

Args:
    transforms (Union[List[Union[TransformArgs, Dict]], Union[TransformArgs, Dict]]): A list of transform configurations or a single transform configuration. 
    num_workers (int, optional): Number of worker threads for parallel processing. Defaults to 0.


Examples:
    >>> from lazyllm.tools.rag.transform import AdaptiveTransform, DocNode, SentenceSplitter
    >>> doc1 = DocNode(text="这是第一个文档的内容。它包含多个句子。")
    >>> doc2 = DocNode(text="这是第二个文档的内容。")
    >>> transforms = [
    ...     {
    ...         'f': SentenceSplitter,
    ...         'pattern': '*.txt',
    ...         'kwargs': {'chunk_size': 50, 'chunk_overlap': 10}
    ...     },
    ...     {
    ...         'f': SentenceSplitter,
    ...         'pattern': '*.pdf',
    ...         'kwargs': {'chunk_size': 100, 'chunk_overlap': 20}
    ...     }
    ... ]
    >>> adaptive = AdaptiveTransform(transforms)
    >>> results1 = adaptive.transform(doc1)
    >>> print(f"文档1转换结果: {len(results1)} 个块")
    >>> for i, result in enumerate(results1):
    ...     print(f"  块 {i+1}: {result.text}")
    >>> results2 = adaptive.transform(doc2)
    >>> print(f"文档2转换结果: {len(results2)} 个块")
    >>> for i, result in enumerate(results2):
    ...     print(f"  块 {i+1}: {result.text}")      
    """
    def __init__(self, transforms: Union[List[Union[TransformArgs, Dict]], Union[TransformArgs, Dict]],
                 num_workers: int = 0):
        super().__init__(num_workers=num_workers)
        if not isinstance(transforms, (tuple, list)): transforms = [transforms]
        self._transformers = [(t.get('pattern'), make_transform(t)) for t in transforms]

    def transform(self, document: DocNode, **kwargs) -> List[Union[str, DocNode]]:
        """Transform a document using the appropriate transformation strategy based on pattern matching.

This method evaluates each transform configuration in order and applies the first one that matches the document's path pattern. The matching logic supports both glob patterns and custom callable functions.

Args:
    document (DocNode): The document node to be transformed.
    **kwargs: Additional keyword arguments passed to the transform function.

**Returns:**

- List[Union[str, DocNode]]: A list of transformed results (strings or DocNode objects).
"""
        if not isinstance(document, DocNode): LOG.warning(f'Invalud document type {type(document)} got')
        for pt, transform in self._transformers:
            if pt and isinstance(pt, str) and not pt.startswith('*'): pt = os.path.join(str(os.cwd()), pt)
            if not pt or (callable(pt) and pt(document.docpath)) or (
                    isinstance(pt, str) and fnmatch.fnmatch(document.docpath, pt)):
                return transform(document, **kwargs)
        LOG.warning(f'No transform found for document {document.docpath} with group name `{self._name}`')
        return []

transform(document, **kwargs)

Transform a document using the appropriate transformation strategy based on pattern matching.

This method evaluates each transform configuration in order and applies the first one that matches the document's path pattern. The matching logic supports both glob patterns and custom callable functions.

Parameters:

  • document (DocNode) –

    The document node to be transformed.

  • **kwargs

    Additional keyword arguments passed to the transform function.

Returns:

  • List[Union[str, DocNode]]: A list of transformed results (strings or DocNode objects).
Source code in lazyllm/tools/rag/transform.py
    def transform(self, document: DocNode, **kwargs) -> List[Union[str, DocNode]]:
        """Transform a document using the appropriate transformation strategy based on pattern matching.

This method evaluates each transform configuration in order and applies the first one that matches the document's path pattern. The matching logic supports both glob patterns and custom callable functions.

Args:
    document (DocNode): The document node to be transformed.
    **kwargs: Additional keyword arguments passed to the transform function.

**Returns:**

- List[Union[str, DocNode]]: A list of transformed results (strings or DocNode objects).
"""
        if not isinstance(document, DocNode): LOG.warning(f'Invalud document type {type(document)} got')
        for pt, transform in self._transformers:
            if pt and isinstance(pt, str) and not pt.startswith('*'): pt = os.path.join(str(os.cwd()), pt)
            if not pt or (callable(pt) and pt(document.docpath)) or (
                    isinstance(pt, str) and fnmatch.fnmatch(document.docpath, pt)):
                return transform(document, **kwargs)
        LOG.warning(f'No transform found for document {document.docpath} with group name `{self._name}`')
        return []

lazyllm.tools.rag.rerank.ModuleReranker

Bases: Reranker

A reranker that uses trainable modules to reorder documents based on relevance to a query.

ModuleReranker is a specialized reranker that leverages trainable models (such as BGE-reranker, Cohere rerank, etc.) to improve the relevance of retrieved documents. It takes a list of documents and a query, then returns the documents reordered by their relevance scores.

Parameters:

  • name (str, default: 'ModuleReranker' ) –

    The name of the reranker. Defaults to "ModuleReranker".

  • model (Union[Callable, str], default: None ) –

    The reranking model. Can be either a model name (string) or a callable function.

  • target (Optional[str], default: None ) –

    Defaults to None.

  • output_format (Optional[str], default: None ) –

    The format for output processing. Defaults to None.

  • join (Union[bool, str], default: False ) –

    Whether to join the results. Defaults to False.

  • **kwargs

    Additional keyword arguments passed to the reranker model.

Examples:

>>> from lazyllm.tools.rag.rerank import ModuleReranker, DocNode
>>> def simple_reranker(query, documents, top_n):
...     query_lower = query.lower()
...     scores = []
...     for i, doc in enumerate(documents):
...         score = sum(1 for word in query_lower.split() if word in doc)
...         scores.append((i, score))
...     scores.sort(key=lambda x: x[1], reverse=True)
...     return scores[:top_n]
>>> reranker = ModuleReranker(
...     model=simple_reranker,
...     topk=2
... )
>>> docs = [
...     DocNode(text="机器学习算法在数据分析中应用广泛"),
...     DocNode(text="深度学习模型需要大量训练数据"),
...     DocNode(text="自然语言处理技术发展迅速"),
...     DocNode(text="计算机视觉在自动驾驶中的应用")
... ]
>>> query = "机器学习"
>>> results = reranker.forward(docs, query)
>>> for i, doc in enumerate(results):
...     print(f"  {i+1}. : {doc.text}")
...     print(f"     相关性分数: {doc.relevance_score:.4f}")
Source code in lazyllm/tools/rag/rerank.py
@Reranker.register_reranker()
class ModuleReranker(Reranker):
    """A reranker that uses trainable modules to reorder documents based on relevance to a query.

ModuleReranker is a specialized reranker that leverages trainable models (such as BGE-reranker, Cohere rerank, etc.) to improve the relevance of retrieved documents. It takes a list of documents and a query, then returns the documents reordered by their relevance scores.

Args:
    name (str): The name of the reranker. Defaults to "ModuleReranker".
    model (Union[Callable, str]): The reranking model. Can be either a model name (string) or a callable function.
    target (Optional[str]): Defaults to None.
    output_format (Optional[str]): The format for output processing. Defaults to None.
    join (Union[bool, str]): Whether to join the results. Defaults to False.
    **kwargs: Additional keyword arguments passed to the reranker model.


Examples:
    >>> from lazyllm.tools.rag.rerank import ModuleReranker, DocNode
    >>> def simple_reranker(query, documents, top_n):
    ...     query_lower = query.lower()
    ...     scores = []
    ...     for i, doc in enumerate(documents):
    ...         score = sum(1 for word in query_lower.split() if word in doc)
    ...         scores.append((i, score))
    ...     scores.sort(key=lambda x: x[1], reverse=True)
    ...     return scores[:top_n]
    >>> reranker = ModuleReranker(
    ...     model=simple_reranker,
    ...     topk=2
    ... )
    >>> docs = [
    ...     DocNode(text="机器学习算法在数据分析中应用广泛"),
    ...     DocNode(text="深度学习模型需要大量训练数据"),
    ...     DocNode(text="自然语言处理技术发展迅速"),
    ...     DocNode(text="计算机视觉在自动驾驶中的应用")
    ... ]
    >>> query = "机器学习"
    >>> results = reranker.forward(docs, query)
    >>> for i, doc in enumerate(results):
    ...     print(f"  {i+1}. : {doc.text}")
    ...     print(f"     相关性分数: {doc.relevance_score:.4f}")        
    """

    def __init__(self, name: str = "ModuleReranker", model: Union[Callable, str] = None, target: Optional[str] = None,
                 output_format: Optional[str] = None, join: Union[bool, str] = False, **kwargs) -> None:
        super().__init__(name, target, output_format, join, **kwargs)
        assert model is not None, "Reranker model must be specified as a model name or a callable."
        if isinstance(model, str):
            self._reranker = lazyllm.TrainableModule(model)
        else:
            self._reranker = model

    def forward(self, nodes: List[DocNode], query: str = "") -> List[DocNode]:
        """Forward pass of the reranker that reorders documents based on relevance to the query.

This method takes a list of documents and a query, then uses the underlying reranking model to score and reorder the documents by relevance. The documents are processed in MetadataMode.EMBED format to ensure compatibility with the reranking model.

Args:
    nodes (List[DocNode]): List of document nodes to be reranked.
    query (str): The query string to rank documents against. Defaults to "".

**Returns:**

- List[DocNode]: List of document nodes reordered by relevance score, with relevance_score attribute added.
"""
        if not nodes:
            return self._post_process([])

        docs = [node.get_text(metadata_mode=MetadataMode.EMBED) for node in nodes]
        top_n = self._kwargs['topk'] if 'topk' in self._kwargs else len(docs)
        sorted_indices = self._reranker(query, documents=docs, top_n=top_n)
        results = []
        for index, relevance_score in sorted_indices:
            results.append(nodes[index].with_score(relevance_score))
        LOG.debug(f"Rerank use `{self._name}` and get nodes: {results}")
        return self._post_process(results)

forward(nodes, query='')

Forward pass of the reranker that reorders documents based on relevance to the query.

This method takes a list of documents and a query, then uses the underlying reranking model to score and reorder the documents by relevance. The documents are processed in MetadataMode.EMBED format to ensure compatibility with the reranking model.

Parameters:

  • nodes (List[DocNode]) –

    List of document nodes to be reranked.

  • query (str, default: '' ) –

    The query string to rank documents against. Defaults to "".

Returns:

  • List[DocNode]: List of document nodes reordered by relevance score, with relevance_score attribute added.
Source code in lazyllm/tools/rag/rerank.py
    def forward(self, nodes: List[DocNode], query: str = "") -> List[DocNode]:
        """Forward pass of the reranker that reorders documents based on relevance to the query.

This method takes a list of documents and a query, then uses the underlying reranking model to score and reorder the documents by relevance. The documents are processed in MetadataMode.EMBED format to ensure compatibility with the reranking model.

Args:
    nodes (List[DocNode]): List of document nodes to be reranked.
    query (str): The query string to rank documents against. Defaults to "".

**Returns:**

- List[DocNode]: List of document nodes reordered by relevance score, with relevance_score attribute added.
"""
        if not nodes:
            return self._post_process([])

        docs = [node.get_text(metadata_mode=MetadataMode.EMBED) for node in nodes]
        top_n = self._kwargs['topk'] if 'topk' in self._kwargs else len(docs)
        sorted_indices = self._reranker(query, documents=docs, top_n=top_n)
        results = []
        for index, relevance_score in sorted_indices:
            results.append(nodes[index].with_score(relevance_score))
        LOG.debug(f"Rerank use `{self._name}` and get nodes: {results}")
        return self._post_process(results)

lazyllm.tools.rag.utils.DocListManager

Bases: ABC

Abstract base class for managing document lists and monitoring changes in a document directory.

Parameters:

  • path

    Path of the document directory to monitor.

  • name

    Name of the manager.

  • enable_path_monitoring

    Whether to enable path monitoring.

Examples:

>>> import lazyllm
>>> from lazyllm.rag.utils import DocListManager
>>> manager = DocListManager(path='your_file_path/', name="test_manager", enable_path_monitoring=False)
>>> added_docs = manager.add_files([test_file_list])
>>> manager.enable_path_monitoring(True)
>>> deleted = manager.delete_files([delete_file_list])
Source code in lazyllm/tools/rag/utils.py
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
class DocListManager(ABC):
    """Abstract base class for managing document lists and monitoring changes in a document directory.

Args:
    path: Path of the document directory to monitor.
    name: Name of the manager.
    enable_path_monitoring: Whether to enable path monitoring.


Examples:

    >>> import lazyllm
    >>> from lazyllm.rag.utils import DocListManager
    >>> manager = DocListManager(path='your_file_path/', name="test_manager", enable_path_monitoring=False)
    >>> added_docs = manager.add_files([test_file_list])
    >>> manager.enable_path_monitoring(True)
    >>> deleted = manager.delete_files([delete_file_list])
    """
    DEFAULT_GROUP_NAME = '__default__'
    __pool__ = dict()

    class Status:
        all = 'all'
        waiting = 'waiting'
        working = 'working'
        success = 'success'
        failed = 'failed'
        deleting = 'deleting'
        # deleted is no longer used
        deleted = 'deleted'

    def __init__(self, path, name, enable_path_monitoring=True):
        self._path = path
        self._name = name
        lazyllm.LOG.info(f'DocManager use file-system monitoring worker: {enable_path_monitoring}')
        self._id = hashlib.sha256(f'{name}@+@{path}'.encode()).hexdigest()
        if not os.path.isabs(path):
            raise ValueError(f"path [{path}] is not an absolute path")

        self._init_sql()
        self._delete_nonexistent_docs_on_startup()

        self._monitor_thread = threading.Thread(target=self._monitor_directory_worker)
        self._monitor_thread.daemon = True
        self._monitor_continue = True
        self._enable_path_monitoring = enable_path_monitoring
        self._init_monitor_event = threading.Event()
        if self._enable_path_monitoring:
            self._monitor_thread.start()
            self._init_monitor_event.wait()

    def _delete_nonexistent_docs_on_startup(self):
        ids = [row[0] for row in self.list_kb_group_files(details=True)
               if not Path(row[1]).exists()]
        if ids: self.delete_files(ids)

    def __new__(cls, *args, **kw):
        if cls is not DocListManager:
            return super().__new__(cls)
        return super().__new__(__class__.__pool__[config['default_dlmanager']])

    def init_tables(self) -> 'DocListManager':
        """Ensure that the default group exists in the database tables.
"""
        if not self.table_inited():
            self._init_tables()
        # in case of using after relase
        self.add_kb_group(DocListManager.DEFAULT_GROUP_NAME)
        return self

    def _monitor_directory(self) -> Set[str]:
        files_list = []
        for root, _, files in os.walk(self._path):
            files = [os.path.join(root, file_path) for file_path in files]
            files_list.extend(files)
        return set(files_list)

    # Actually it shoule be "set_docs_status_deleting"
    def delete_files(self, file_ids: List[str]) -> List[DocPartRow]:
        """Set the knowledge base entries associated with the document to "deleting," and have each knowledge base asynchronously delete parsed results and associated records.

Args:
    file_ids (list of str): List of file IDs to delete.
"""
        document_list = self.update_file_status(file_ids, DocListManager.Status.deleting)
        self.update_kb_group(cond_file_ids=file_ids, new_status=DocListManager.Status.deleting)
        return document_list

    @abstractmethod
    def table_inited(self):
        """Checks if the database table `documents` is initialized. This method ensures thread-safety when accessing the database.
Determines whether the `documents` table exists in the database.
Returns:
    bool: `True` if the `documents` table exists, `False` otherwise.
Notes:
    - Uses a thread-safe lock (`self._db_lock`) to ensure safe access to the database.
    - Establishes a connection to the SQLite database at `self._db_path` with the `check_same_thread` option.
    - Executes the SQL query: `SELECT name FROM sqlite_master WHERE type='table' AND name='documents'` to check for the table.
"""
        pass

    @abstractmethod
    def _init_tables(self): pass

    @abstractmethod
    def validate_paths(self, paths: List[str]) -> Tuple[bool, str, List[bool]]:
        """Validates a list of file paths to ensure they are ready for processing.
This method checks whether the provided paths are new, already processed, or currently being processed. It ensures there are no conflicts in processing the documents.
Args
    paths (List[str]): A list of file paths to validate.
Returns:
    Tuple[bool, str, List[bool]]: A tuple containing:
        - `bool`: `True` if all paths are valid, `False` otherwise.
        - `str`: A message indicating success or the reason for failure.
        - `List[bool]`: A list where each element corresponds to whether a path is new (`True`) or already exists (`False`).
Notes:
    - If any document is still being processed or needs reparsing, the method returns `False` with an appropriate error message.
    - The method uses a database session and thread-safe lock (`self._db_lock`) to retrieve document status information.
    - Unsafe statuses include `working` and `waiting`.

"""
        pass

    @abstractmethod
    def update_need_reparsing(self, doc_id: str, need_reparse: bool):
        """Updates the `need_reparse` status of a document in the `KBGroupDocuments` table.
This method sets the `need_reparse` flag for a specific document, optionally scoped to a given group.
Args:
    doc_id (str): The ID of the document to update.
    need_reparse (bool): The new value for the `need_reparse` flag.
    group_name (Optional[str]): If provided, the update will be applied only to the specified group.
Notes:
    - Uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - The `group_name` parameter allows scoping the update to a specific group; if not provided, the update applies to all groups containing the document.
    - The method commits the change to the database immediately.
"""
        pass

    @abstractmethod
    def list_files(self, limit: Optional[int] = None, details: bool = False,
                   status: Union[str, List[str]] = Status.all,
                   exclude_status: Optional[Union[str, List[str]]] = None):
        """Lists files from the `documents` table with optional filtering, limiting, and returning details.
This method retrieves file IDs or detailed file information from the database, based on the specified filtering conditions.
Args:
    limit (Optional[int]): Maximum number of files to return. If `None`, all matching files will be returned.
    details (bool): Whether to return detailed file information (`True`) or just file IDs (`False`).
    status (Union[str, List[str]]): The status or list of statuses to include in the results. Defaults to all statuses.
    exclude_status (Optional[Union[str, List[str]]]): The status or list of statuses to exclude from the results. Defaults to `None`.
Returns:
    List: A list of file IDs if `details=False`, or a list of detailed file rows if `details=True`.
Notes:
    - The method constructs a query dynamically based on the provided `status` and `exclude_status` conditions.
    - A thread-safe lock (`self._db_lock`) ensures safe database access.
    - The `LIMIT` clause is applied if `limit` is specified.
"""
        pass

    @abstractmethod
    def get_docs(self, doc_ids: List[str]) -> List[KBDocument]:
        """This method retrieves document objects of type `KBDocument` from the database for the provided list of document IDs.
Args:
    doc_ids (List[str]): A list of document IDs to fetch.
Returns:
    List[KBDocument]: A list of `KBDocument` objects corresponding to the provided document IDs. If no documents are found, an empty list is returned.
Notes:
    - The method uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - It performs a SQL join between `KBDocument` and `KBGroupDocuments` to retrieve the relevant rows.
    - After fetching, it updates the `new_meta` field of the affected rows to `None` and commits the changes to the database.
"""
        pass

    @abstractmethod
    def set_docs_new_meta(self, doc_meta: Dict[str, dict]):
        """Batch update metadata for documents.

Args:
    doc_meta (Dict[str, dict]): A dictionary mapping document IDs to their new metadata.
"""
        pass

    @abstractmethod
    def fetch_docs_changed_meta(self, group: str) -> List[DocMetaChangedRow]:
        """List files in a specific knowledge base (KB) group with optional filters, limiting, and details.
This method retrieves files from the `kb_group_documents` table, optionally filtering by group, document status, upload status, and whether reparsing is needed.
Args:
    group (str): The name of the group to filter documents by.
**Returns:**
    List[DocMetaChangedRow]: A list of rows, where each row contains the `doc_id` and the `new_meta` field of documents with changed metadata.
Notes:
    - This method constructs a SQL query dynamically based on the provided filters.
    - Uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - If `status` or `upload_status` are provided as lists, they are processed with SQL `IN` clauses.
"""
        pass

    @abstractmethod
    def list_all_kb_group(self):
        """Lists all the knowledge base group names.

**Returns:**
- list: List of knowledge base group names.
"""
        pass

    @abstractmethod
    def add_kb_group(self, name):
        """Adds a new knowledge base group.

Args:
    name (str): Name of the group to add.
"""
        pass

    @abstractmethod
    def list_kb_group_files(self, group: str = None, limit: Optional[int] = None, details: bool = False,
                            status: Union[str, List[str]] = Status.all,
                            exclude_status: Optional[Union[str, List[str]]] = None,
                            upload_status: Union[str, List[str]] = Status.all,
                            exclude_upload_status: Optional[Union[str, List[str]]] = None,
                            need_reparse: Optional[bool] = False):
        """List files in a specific knowledge base group .

Args:
    group (str): The name of the KB group to filter files by. Defaults to `None` .
    limit (Optional[int]): Maximum number of files to return. If `None`, returns all matching files.
    details (bool): Whether to return detailed file information or only file IDs and paths.
    status (Union[str, List[str]]): The KB group status or list of statuses to include in the results. Defaults to all statuses.
    exclude_status (Optional[Union[str, List[str]]): The KB group status or list of statuses to exclude from the results. Defaults to `None`.
    upload_status (Union[str, List[str]]): The document upload status or list of statuses to include in the results. Defaults to all statuses.
    exclude_upload_status (Optional[Union[str, List[str]]): The document upload status or list of statuses to exclude from the results. Defaults to `None`.
    need_reparse (Optional[bool]): Whether to filter files that need reparsing or not . Defaults to `None` .
**Returns:**:
    List: If `details=False`, returns a list of tuples containing `(doc_id, path)`. 
          If `details=True`, returns a list of detailed rows with additional metadata.
Notes:
    - The method first creates document records using the `_add_doc_records` helper function.
    - After the files are added, they are automatically linked to the default KB group (`DocListManager.DEFAULT_GROUP_NAME`).
    - Batch processing ensures scalability when adding a large number of files.
"""
        pass

    def add_files(
        self,
        files: List[str],
        metadatas: Optional[List[Dict[str, Any]]] = None,
        status: Optional[str] = Status.waiting,
        batch_size: int = 64,
    ) -> List[DocPartRow]:
        """Add multiple files to the document list with optional metadata, status, and batch processing.
This method adds a list of files to the database and sets optional metadata and initial status for each file. The files are processed in batches for efficiency. After the files are added, they are automatically associated with the default knowledge base (KB) group.
Args:
    files (List[str]): A list of file paths to add to the database.
    metadatas (Optional[List[Dict[str, Any]]]): A list of metadata dictionaries corresponding to the files. If `None`, no metadata will be associated. Defaults to `None`.
    status (Optional[str]): The initial status for the added files. Defaults to `Status.waiting`.
    batch_size (int): The number of files to process in each batch. Defaults to 64.
**Returns:**:
    List[DocPartRow]: A list of `DocPartRow` objects representing the added files and their associated information.
Notes:
- The method first creates document records using the helper function _add_doc_records.
- After the files are added, they are automatically linked to the default knowledge base group (DocListManager.DEFAULT_GROUP_NAME).
- Batch processing ensures good scalability when adding a large number of files.


"""
        documents = self._add_doc_records(files, metadatas, status, batch_size)
        if documents:
            self.add_files_to_kb_group([doc.doc_id for doc in documents], group=DocListManager.DEFAULT_GROUP_NAME)
        return documents

    @abstractmethod
    def _get_all_docs(self): pass

    @abstractmethod
    def _get_docs(self, to_be_added_doc_ids: List, to_be_deleted_doc_ids: List, filter_status_list: List): pass

    @abstractmethod
    def _add_doc_records(self, files: List[str], metadatas: Optional[List] = None,
                         status: Optional[str] = Status.waiting, batch_size: int = 64) -> List[DocPartRow]: pass

    @abstractmethod
    def delete_unreferenced_doc(self):
        """Delete documents marked as "deleting" and no longer referenced in the database.
This method removes documents from the database that meet the following conditions:
1. Their status is set to `DocListManager.Status.deleting`.
2. Their reference count (`count`) is 0.
"""
        pass

    @abstractmethod
    def get_docs_need_reparse(self, group: Optional[str] = None) -> List[KBDocument]:
        """Retrieve documents that require reparsing for a specific group.
This method fetches documents that are marked as needing reparsing (`need_reparse=True`) for the given group. Only documents with a status of `success` or `failed` are included in the results.
Args:
    group (str): The name of the group to filter documents by.
**Returns:**:
    List[KBDocument]: A list of `KBDocument` objects that need reparsing.
Notes:
    - The method uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - The query performs a SQL `JOIN` between `KBDocument` and `KBGroupDocuments` to filter by group and reparse status.
    - Documents with `need_reparse=True` and a status of `success` or `failed` are considered for reparsing.
"""
        pass

    @abstractmethod
    def get_existing_paths_by_pattern(self, file_path: str) -> List[str]:
        """Retrieve existing document paths that match a given pattern.
This method fetches all document paths from the database that match the provided SQL `LIKE` pattern.
Args:
    pattern (str): The SQL `LIKE` pattern to filter document paths. For example, `%example%` matches paths containing the word "example".
**Returns:**:
    List[str]: A list of document paths that match the given pattern. If no paths match, an empty list is returned.
Notes:
    - The method uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - The `LIKE` operator in the SQL query is used to perform pattern matching on document paths.
"""
        pass

    @abstractmethod
    def update_file_message(self, fileid: str, **kw):
        """Updates the message for a specified file.

Args:
    fileid (str): File ID.
    **kw: Additional key-value pairs to update.
"""
        pass

    @abstractmethod
    def update_file_status(self, file_ids: List[str], status: str,
                           cond_status_list: Union[None, List[str]] = None) -> List[DocPartRow]:
        """Update the status of specified files.

Args:
    file_ids (list of str): List of file IDs whose status needs to be updated.
    status (str): Target status to set.
    cond_status_list (Union[None, List[str]]): Optional. Only update files currently in these statuses.
"""
        pass

    @abstractmethod
    def add_files_to_kb_group(self, file_ids: List[str], group: str):
        """Adds files to the specified knowledge base group.

Args:
    file_ids (list of str): List of file IDs to add.
    group (str): Name of the group to add the files to.
"""
        pass

    @abstractmethod
    def delete_files_from_kb_group(self, file_ids: List[str], group: str):
        """Deletes files from the specified knowledge base group.

Args:
    file_ids (list of str): List of file IDs to delete.
    group (str): Name of the group.
"""
        pass

    @abstractmethod
    def get_file_status(self, fileid: str):
        """Retrieves the status of a specified file.

Args:
    fileid (str): File ID.

**Returns:**
- str: The current status of the file.
"""
        pass

    @abstractmethod
    def update_kb_group(self, cond_file_ids: List[str], cond_group: Optional[str] = None,
                        cond_status_list: Optional[List[str]] = None, new_status: Optional[str] = None,
                        new_need_reparse: Optional[bool] = None) -> List[GroupDocPartRow]:
        """Updates the record of kb_group_document.

Args:
    cond_file_ids (list of str, optional): a list of file IDs to filter by, default None.
    cond_group (str, optional): a kb_group name to filter by, default None.
    cond_status_list (list of str, optional): a list of statuses to filter by, default None.
    new_status (str, optional): the new status to update to, default None
    new_need_reparse (bool, optinoal): the new need_reparse flag to update to, default None

**Returns:**
- list: updated records, list of (doc_id, group_name)
"""
        pass

    @abstractmethod
    def release(self):
        """Releases the resources of the current manager.
"""
        pass

    @property
    def enable_path_monitoring(self):
        """Enable or disable path monitoring for the document manager.
This method enables or disables the path monitoring functionality in the document manager. When enabled, a monitoring thread starts to handle path-related operations. When disabled, the thread stops and joins (waits for it to terminate).
Args:
    val (bool): Whether to enable or disable path monitoring.
Notes:
    - If `val` is `True`, path monitoring is enabled by setting `_monitor_continue` to `True` and starting the `_monitor_thread`.
    - If `val` is `False`, path monitoring is disabled by setting `_monitor_continue` to `False` and joining the `_monitor_thread` if it is running.
    - This method ensures thread-safe operation when managing the monitoring thread.
"""
        return self._enable_path_monitoring

    @enable_path_monitoring.setter
    def enable_path_monitoring(self, val: bool):
        """Enable or disable path monitoring for the document manager.
This method enables or disables the path monitoring functionality in the document manager. When enabled, a monitoring thread starts to handle path-related operations. When disabled, the thread stops and joins (waits for it to terminate).
Args:
    val (bool): Whether to enable or disable path monitoring.
Notes:
    - If `val` is `True`, path monitoring is enabled by setting `_monitor_continue` to `True` and starting the `_monitor_thread`.
    - If `val` is `False`, path monitoring is disabled by setting `_monitor_continue` to `False` and joining the `_monitor_thread` if it is running.
    - This method ensures thread-safe operation when managing the monitoring thread.
"""
        self._enable_path_monitoring = (val is True)
        if val is True:
            self._monitor_continue = True
            self._monitor_thread.start()
        else:
            self._monitor_continue = False
            if self._monitor_thread.is_alive():
                self._monitor_thread.join()

    def _monitor_directory_worker(self):
        failed_files_count = defaultdict(int)
        docs_all = self._get_all_docs()

        previous_files = set([doc.path for doc in docs_all])
        skip_files = set()
        is_first_run = True
        while self._monitor_continue:
            # 1. Scan files in the directory, find added and deleted files
            current_files = set(self._monitor_directory())
            to_be_added_files = current_files - previous_files - skip_files
            to_be_deleted_files = previous_files - current_files - skip_files
            failed_files = set()

            to_be_added_doc_ids = set([gen_docid(ele) for ele in to_be_added_files])
            to_be_deleted_doc_ids = set([gen_docid(ele) for ele in to_be_deleted_files])
            failed_doc_ids = set()
            filter_status_list = [DocListManager.Status.success,
                                  DocListManager.Status.failed, DocListManager.Status.waiting]

            docs_not_expected, docs_expected = self._get_docs(to_be_added_doc_ids,
                                                              to_be_deleted_doc_ids, filter_status_list)

            # 2. Skip new files that are already in the database
            failed_files.update([doc.path for doc in docs_not_expected])
            failed_doc_ids.update([doc.doc_id for doc in docs_not_expected])
            to_be_added_files -= failed_files
            # Actually it is add to doc with success status, then add to kb_group with waiting status
            self.add_files(list(to_be_added_files), status=DocListManager.Status.success)

            # 3. Skip deleted files that are: 1. not in the database, 2. status not success/failed/waiting
            safe_to_delete_files = set([doc.path for doc in docs_expected])
            safe_to_delete_doc_ids = set([doc.doc_id for doc in docs_expected])
            failed_doc_ids.update(to_be_deleted_doc_ids - safe_to_delete_doc_ids)
            failed_files.update(to_be_deleted_files - safe_to_delete_files)
            to_be_deleted_files = safe_to_delete_files
            to_be_deleted_doc_ids = safe_to_delete_doc_ids
            self.delete_files(list(to_be_deleted_doc_ids))

            # 4. update skip_files
            for ele in failed_files:
                failed_files_count[ele] += 1
                if failed_files_count[ele] >= 3:
                    skip_files.add(ele)
            # update previous files, while failed files will be re-processed in the next loop
            previous_files = (current_files | to_be_added_files) - to_be_deleted_files
            if is_first_run:
                self._init_monitor_event.set()
            is_first_run = False
            time.sleep(10)
        lazyllm.LOG.warning("END MONITORING")

    def __del__(self):
        self.enable_path_monitoring = False

enable_path_monitoring property writable

Enable or disable path monitoring for the document manager. This method enables or disables the path monitoring functionality in the document manager. When enabled, a monitoring thread starts to handle path-related operations. When disabled, the thread stops and joins (waits for it to terminate). Args: val (bool): Whether to enable or disable path monitoring. Notes: - If val is True, path monitoring is enabled by setting _monitor_continue to True and starting the _monitor_thread. - If val is False, path monitoring is disabled by setting _monitor_continue to False and joining the _monitor_thread if it is running. - This method ensures thread-safe operation when managing the monitoring thread.

add_files(files, metadatas=None, status=Status.waiting, batch_size=64)

Add multiple files to the document list with optional metadata, status, and batch processing. This method adds a list of files to the database and sets optional metadata and initial status for each file. The files are processed in batches for efficiency. After the files are added, they are automatically associated with the default knowledge base (KB) group. Args: files (List[str]): A list of file paths to add to the database. metadatas (Optional[List[Dict[str, Any]]]): A list of metadata dictionaries corresponding to the files. If None, no metadata will be associated. Defaults to None. status (Optional[str]): The initial status for the added files. Defaults to Status.waiting. batch_size (int): The number of files to process in each batch. Defaults to 64. Returns:: List[DocPartRow]: A list of DocPartRow objects representing the added files and their associated information. Notes: - The method first creates document records using the helper function _add_doc_records. - After the files are added, they are automatically linked to the default knowledge base group (DocListManager.DEFAULT_GROUP_NAME). - Batch processing ensures good scalability when adding a large number of files.

Source code in lazyllm/tools/rag/utils.py
    def add_files(
        self,
        files: List[str],
        metadatas: Optional[List[Dict[str, Any]]] = None,
        status: Optional[str] = Status.waiting,
        batch_size: int = 64,
    ) -> List[DocPartRow]:
        """Add multiple files to the document list with optional metadata, status, and batch processing.
This method adds a list of files to the database and sets optional metadata and initial status for each file. The files are processed in batches for efficiency. After the files are added, they are automatically associated with the default knowledge base (KB) group.
Args:
    files (List[str]): A list of file paths to add to the database.
    metadatas (Optional[List[Dict[str, Any]]]): A list of metadata dictionaries corresponding to the files. If `None`, no metadata will be associated. Defaults to `None`.
    status (Optional[str]): The initial status for the added files. Defaults to `Status.waiting`.
    batch_size (int): The number of files to process in each batch. Defaults to 64.
**Returns:**:
    List[DocPartRow]: A list of `DocPartRow` objects representing the added files and their associated information.
Notes:
- The method first creates document records using the helper function _add_doc_records.
- After the files are added, they are automatically linked to the default knowledge base group (DocListManager.DEFAULT_GROUP_NAME).
- Batch processing ensures good scalability when adding a large number of files.


"""
        documents = self._add_doc_records(files, metadatas, status, batch_size)
        if documents:
            self.add_files_to_kb_group([doc.doc_id for doc in documents], group=DocListManager.DEFAULT_GROUP_NAME)
        return documents

add_files_to_kb_group(file_ids, group) abstractmethod

Adds files to the specified knowledge base group.

Parameters:

  • file_ids (list of str) –

    List of file IDs to add.

  • group (str) –

    Name of the group to add the files to.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def add_files_to_kb_group(self, file_ids: List[str], group: str):
        """Adds files to the specified knowledge base group.

Args:
    file_ids (list of str): List of file IDs to add.
    group (str): Name of the group to add the files to.
"""
        pass

add_kb_group(name) abstractmethod

Adds a new knowledge base group.

Parameters:

  • name (str) –

    Name of the group to add.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def add_kb_group(self, name):
        """Adds a new knowledge base group.

Args:
    name (str): Name of the group to add.
"""
        pass

delete_files(file_ids)

Set the knowledge base entries associated with the document to "deleting," and have each knowledge base asynchronously delete parsed results and associated records.

Parameters:

  • file_ids (list of str) –

    List of file IDs to delete.

Source code in lazyllm/tools/rag/utils.py
    def delete_files(self, file_ids: List[str]) -> List[DocPartRow]:
        """Set the knowledge base entries associated with the document to "deleting," and have each knowledge base asynchronously delete parsed results and associated records.

Args:
    file_ids (list of str): List of file IDs to delete.
"""
        document_list = self.update_file_status(file_ids, DocListManager.Status.deleting)
        self.update_kb_group(cond_file_ids=file_ids, new_status=DocListManager.Status.deleting)
        return document_list

delete_files_from_kb_group(file_ids, group) abstractmethod

Deletes files from the specified knowledge base group.

Parameters:

  • file_ids (list of str) –

    List of file IDs to delete.

  • group (str) –

    Name of the group.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def delete_files_from_kb_group(self, file_ids: List[str], group: str):
        """Deletes files from the specified knowledge base group.

Args:
    file_ids (list of str): List of file IDs to delete.
    group (str): Name of the group.
"""
        pass

delete_unreferenced_doc() abstractmethod

Delete documents marked as "deleting" and no longer referenced in the database. This method removes documents from the database that meet the following conditions: 1. Their status is set to DocListManager.Status.deleting. 2. Their reference count (count) is 0.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def delete_unreferenced_doc(self):
        """Delete documents marked as "deleting" and no longer referenced in the database.
This method removes documents from the database that meet the following conditions:
1. Their status is set to `DocListManager.Status.deleting`.
2. Their reference count (`count`) is 0.
"""
        pass

fetch_docs_changed_meta(group) abstractmethod

List files in a specific knowledge base (KB) group with optional filters, limiting, and details. This method retrieves files from the kb_group_documents table, optionally filtering by group, document status, upload status, and whether reparsing is needed. Args: group (str): The name of the group to filter documents by. Returns: List[DocMetaChangedRow]: A list of rows, where each row contains the doc_id and the new_meta field of documents with changed metadata. Notes: - This method constructs a SQL query dynamically based on the provided filters. - Uses a thread-safe lock (self._db_lock) to ensure safe database access. - If status or upload_status are provided as lists, they are processed with SQL IN clauses.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def fetch_docs_changed_meta(self, group: str) -> List[DocMetaChangedRow]:
        """List files in a specific knowledge base (KB) group with optional filters, limiting, and details.
This method retrieves files from the `kb_group_documents` table, optionally filtering by group, document status, upload status, and whether reparsing is needed.
Args:
    group (str): The name of the group to filter documents by.
**Returns:**
    List[DocMetaChangedRow]: A list of rows, where each row contains the `doc_id` and the `new_meta` field of documents with changed metadata.
Notes:
    - This method constructs a SQL query dynamically based on the provided filters.
    - Uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - If `status` or `upload_status` are provided as lists, they are processed with SQL `IN` clauses.
"""
        pass

get_docs(doc_ids) abstractmethod

This method retrieves document objects of type KBDocument from the database for the provided list of document IDs. Args: doc_ids (List[str]): A list of document IDs to fetch. Returns: List[KBDocument]: A list of KBDocument objects corresponding to the provided document IDs. If no documents are found, an empty list is returned. Notes: - The method uses a thread-safe lock (self._db_lock) to ensure safe database access. - It performs a SQL join between KBDocument and KBGroupDocuments to retrieve the relevant rows. - After fetching, it updates the new_meta field of the affected rows to None and commits the changes to the database.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def get_docs(self, doc_ids: List[str]) -> List[KBDocument]:
        """This method retrieves document objects of type `KBDocument` from the database for the provided list of document IDs.
Args:
    doc_ids (List[str]): A list of document IDs to fetch.
Returns:
    List[KBDocument]: A list of `KBDocument` objects corresponding to the provided document IDs. If no documents are found, an empty list is returned.
Notes:
    - The method uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - It performs a SQL join between `KBDocument` and `KBGroupDocuments` to retrieve the relevant rows.
    - After fetching, it updates the `new_meta` field of the affected rows to `None` and commits the changes to the database.
"""
        pass

get_docs_need_reparse(group=None) abstractmethod

Retrieve documents that require reparsing for a specific group. This method fetches documents that are marked as needing reparsing (need_reparse=True) for the given group. Only documents with a status of success or failed are included in the results. Args: group (str): The name of the group to filter documents by. Returns:: List[KBDocument]: A list of KBDocument objects that need reparsing. Notes: - The method uses a thread-safe lock (self._db_lock) to ensure safe database access. - The query performs a SQL JOIN between KBDocument and KBGroupDocuments to filter by group and reparse status. - Documents with need_reparse=True and a status of success or failed are considered for reparsing.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def get_docs_need_reparse(self, group: Optional[str] = None) -> List[KBDocument]:
        """Retrieve documents that require reparsing for a specific group.
This method fetches documents that are marked as needing reparsing (`need_reparse=True`) for the given group. Only documents with a status of `success` or `failed` are included in the results.
Args:
    group (str): The name of the group to filter documents by.
**Returns:**:
    List[KBDocument]: A list of `KBDocument` objects that need reparsing.
Notes:
    - The method uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - The query performs a SQL `JOIN` between `KBDocument` and `KBGroupDocuments` to filter by group and reparse status.
    - Documents with `need_reparse=True` and a status of `success` or `failed` are considered for reparsing.
"""
        pass

get_existing_paths_by_pattern(file_path) abstractmethod

Retrieve existing document paths that match a given pattern. This method fetches all document paths from the database that match the provided SQL LIKE pattern. Args: pattern (str): The SQL LIKE pattern to filter document paths. For example, %example% matches paths containing the word "example". Returns:: List[str]: A list of document paths that match the given pattern. If no paths match, an empty list is returned. Notes: - The method uses a thread-safe lock (self._db_lock) to ensure safe database access. - The LIKE operator in the SQL query is used to perform pattern matching on document paths.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def get_existing_paths_by_pattern(self, file_path: str) -> List[str]:
        """Retrieve existing document paths that match a given pattern.
This method fetches all document paths from the database that match the provided SQL `LIKE` pattern.
Args:
    pattern (str): The SQL `LIKE` pattern to filter document paths. For example, `%example%` matches paths containing the word "example".
**Returns:**:
    List[str]: A list of document paths that match the given pattern. If no paths match, an empty list is returned.
Notes:
    - The method uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - The `LIKE` operator in the SQL query is used to perform pattern matching on document paths.
"""
        pass

get_file_status(fileid) abstractmethod

Retrieves the status of a specified file.

Parameters:

  • fileid (str) –

    File ID.

Returns: - str: The current status of the file.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def get_file_status(self, fileid: str):
        """Retrieves the status of a specified file.

Args:
    fileid (str): File ID.

**Returns:**
- str: The current status of the file.
"""
        pass

init_tables()

Ensure that the default group exists in the database tables.

Source code in lazyllm/tools/rag/utils.py
    def init_tables(self) -> 'DocListManager':
        """Ensure that the default group exists in the database tables.
"""
        if not self.table_inited():
            self._init_tables()
        # in case of using after relase
        self.add_kb_group(DocListManager.DEFAULT_GROUP_NAME)
        return self

list_all_kb_group() abstractmethod

Lists all the knowledge base group names.

Returns: - list: List of knowledge base group names.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def list_all_kb_group(self):
        """Lists all the knowledge base group names.

**Returns:**
- list: List of knowledge base group names.
"""
        pass

list_files(limit=None, details=False, status=Status.all, exclude_status=None) abstractmethod

Lists files from the documents table with optional filtering, limiting, and returning details. This method retrieves file IDs or detailed file information from the database, based on the specified filtering conditions. Args: limit (Optional[int]): Maximum number of files to return. If None, all matching files will be returned. details (bool): Whether to return detailed file information (True) or just file IDs (False). status (Union[str, List[str]]): The status or list of statuses to include in the results. Defaults to all statuses. exclude_status (Optional[Union[str, List[str]]]): The status or list of statuses to exclude from the results. Defaults to None. Returns: List: A list of file IDs if details=False, or a list of detailed file rows if details=True. Notes: - The method constructs a query dynamically based on the provided status and exclude_status conditions. - A thread-safe lock (self._db_lock) ensures safe database access. - The LIMIT clause is applied if limit is specified.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def list_files(self, limit: Optional[int] = None, details: bool = False,
                   status: Union[str, List[str]] = Status.all,
                   exclude_status: Optional[Union[str, List[str]]] = None):
        """Lists files from the `documents` table with optional filtering, limiting, and returning details.
This method retrieves file IDs or detailed file information from the database, based on the specified filtering conditions.
Args:
    limit (Optional[int]): Maximum number of files to return. If `None`, all matching files will be returned.
    details (bool): Whether to return detailed file information (`True`) or just file IDs (`False`).
    status (Union[str, List[str]]): The status or list of statuses to include in the results. Defaults to all statuses.
    exclude_status (Optional[Union[str, List[str]]]): The status or list of statuses to exclude from the results. Defaults to `None`.
Returns:
    List: A list of file IDs if `details=False`, or a list of detailed file rows if `details=True`.
Notes:
    - The method constructs a query dynamically based on the provided `status` and `exclude_status` conditions.
    - A thread-safe lock (`self._db_lock`) ensures safe database access.
    - The `LIMIT` clause is applied if `limit` is specified.
"""
        pass

list_kb_group_files(group=None, limit=None, details=False, status=Status.all, exclude_status=None, upload_status=Status.all, exclude_upload_status=None, need_reparse=False) abstractmethod

List files in a specific knowledge base group .

Parameters:

  • group (str, default: None ) –

    The name of the KB group to filter files by. Defaults to None .

  • limit (Optional[int], default: None ) –

    Maximum number of files to return. If None, returns all matching files.

  • details (bool, default: False ) –

    Whether to return detailed file information or only file IDs and paths.

  • status (Union[str, List[str]], default: all ) –

    The KB group status or list of statuses to include in the results. Defaults to all statuses.

  • exclude_status (Optional[Union[str, List[str]], default: None ) –

    The KB group status or list of statuses to exclude from the results. Defaults to None.

  • upload_status (Union[str, List[str]], default: all ) –

    The document upload status or list of statuses to include in the results. Defaults to all statuses.

  • exclude_upload_status (Optional[Union[str, List[str]], default: None ) –

    The document upload status or list of statuses to exclude from the results. Defaults to None.

  • need_reparse (Optional[bool], default: False ) –

    Whether to filter files that need reparsing or not . Defaults to None .

Returns:: List: If details=False, returns a list of tuples containing (doc_id, path). If details=True, returns a list of detailed rows with additional metadata. Notes: - The method first creates document records using the _add_doc_records helper function. - After the files are added, they are automatically linked to the default KB group (DocListManager.DEFAULT_GROUP_NAME). - Batch processing ensures scalability when adding a large number of files.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def list_kb_group_files(self, group: str = None, limit: Optional[int] = None, details: bool = False,
                            status: Union[str, List[str]] = Status.all,
                            exclude_status: Optional[Union[str, List[str]]] = None,
                            upload_status: Union[str, List[str]] = Status.all,
                            exclude_upload_status: Optional[Union[str, List[str]]] = None,
                            need_reparse: Optional[bool] = False):
        """List files in a specific knowledge base group .

Args:
    group (str): The name of the KB group to filter files by. Defaults to `None` .
    limit (Optional[int]): Maximum number of files to return. If `None`, returns all matching files.
    details (bool): Whether to return detailed file information or only file IDs and paths.
    status (Union[str, List[str]]): The KB group status or list of statuses to include in the results. Defaults to all statuses.
    exclude_status (Optional[Union[str, List[str]]): The KB group status or list of statuses to exclude from the results. Defaults to `None`.
    upload_status (Union[str, List[str]]): The document upload status or list of statuses to include in the results. Defaults to all statuses.
    exclude_upload_status (Optional[Union[str, List[str]]): The document upload status or list of statuses to exclude from the results. Defaults to `None`.
    need_reparse (Optional[bool]): Whether to filter files that need reparsing or not . Defaults to `None` .
**Returns:**:
    List: If `details=False`, returns a list of tuples containing `(doc_id, path)`. 
          If `details=True`, returns a list of detailed rows with additional metadata.
Notes:
    - The method first creates document records using the `_add_doc_records` helper function.
    - After the files are added, they are automatically linked to the default KB group (`DocListManager.DEFAULT_GROUP_NAME`).
    - Batch processing ensures scalability when adding a large number of files.
"""
        pass

release() abstractmethod

Releases the resources of the current manager.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def release(self):
        """Releases the resources of the current manager.
"""
        pass

set_docs_new_meta(doc_meta) abstractmethod

Batch update metadata for documents.

Parameters:

  • doc_meta (Dict[str, dict]) –

    A dictionary mapping document IDs to their new metadata.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def set_docs_new_meta(self, doc_meta: Dict[str, dict]):
        """Batch update metadata for documents.

Args:
    doc_meta (Dict[str, dict]): A dictionary mapping document IDs to their new metadata.
"""
        pass

table_inited() abstractmethod

Checks if the database table documents is initialized. This method ensures thread-safety when accessing the database. Determines whether the documents table exists in the database. Returns: bool: True if the documents table exists, False otherwise. Notes: - Uses a thread-safe lock (self._db_lock) to ensure safe access to the database. - Establishes a connection to the SQLite database at self._db_path with the check_same_thread option. - Executes the SQL query: SELECT name FROM sqlite_master WHERE type='table' AND name='documents' to check for the table.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def table_inited(self):
        """Checks if the database table `documents` is initialized. This method ensures thread-safety when accessing the database.
Determines whether the `documents` table exists in the database.
Returns:
    bool: `True` if the `documents` table exists, `False` otherwise.
Notes:
    - Uses a thread-safe lock (`self._db_lock`) to ensure safe access to the database.
    - Establishes a connection to the SQLite database at `self._db_path` with the `check_same_thread` option.
    - Executes the SQL query: `SELECT name FROM sqlite_master WHERE type='table' AND name='documents'` to check for the table.
"""
        pass

update_file_message(fileid, **kw) abstractmethod

Updates the message for a specified file.

Parameters:

  • fileid (str) –

    File ID.

  • **kw

    Additional key-value pairs to update.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def update_file_message(self, fileid: str, **kw):
        """Updates the message for a specified file.

Args:
    fileid (str): File ID.
    **kw: Additional key-value pairs to update.
"""
        pass

update_file_status(file_ids, status, cond_status_list=None) abstractmethod

Update the status of specified files.

Parameters:

  • file_ids (list of str) –

    List of file IDs whose status needs to be updated.

  • status (str) –

    Target status to set.

  • cond_status_list (Union[None, List[str]], default: None ) –

    Optional. Only update files currently in these statuses.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def update_file_status(self, file_ids: List[str], status: str,
                           cond_status_list: Union[None, List[str]] = None) -> List[DocPartRow]:
        """Update the status of specified files.

Args:
    file_ids (list of str): List of file IDs whose status needs to be updated.
    status (str): Target status to set.
    cond_status_list (Union[None, List[str]]): Optional. Only update files currently in these statuses.
"""
        pass

update_kb_group(cond_file_ids, cond_group=None, cond_status_list=None, new_status=None, new_need_reparse=None) abstractmethod

Updates the record of kb_group_document.

Parameters:

  • cond_file_ids (list of str) –

    a list of file IDs to filter by, default None.

  • cond_group (str, default: None ) –

    a kb_group name to filter by, default None.

  • cond_status_list (list of str, default: None ) –

    a list of statuses to filter by, default None.

  • new_status (str, default: None ) –

    the new status to update to, default None

  • new_need_reparse ((bool, optinoal), default: None ) –

    the new need_reparse flag to update to, default None

Returns: - list: updated records, list of (doc_id, group_name)

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def update_kb_group(self, cond_file_ids: List[str], cond_group: Optional[str] = None,
                        cond_status_list: Optional[List[str]] = None, new_status: Optional[str] = None,
                        new_need_reparse: Optional[bool] = None) -> List[GroupDocPartRow]:
        """Updates the record of kb_group_document.

Args:
    cond_file_ids (list of str, optional): a list of file IDs to filter by, default None.
    cond_group (str, optional): a kb_group name to filter by, default None.
    cond_status_list (list of str, optional): a list of statuses to filter by, default None.
    new_status (str, optional): the new status to update to, default None
    new_need_reparse (bool, optinoal): the new need_reparse flag to update to, default None

**Returns:**
- list: updated records, list of (doc_id, group_name)
"""
        pass

update_need_reparsing(doc_id, need_reparse) abstractmethod

Updates the need_reparse status of a document in the KBGroupDocuments table. This method sets the need_reparse flag for a specific document, optionally scoped to a given group. Args: doc_id (str): The ID of the document to update. need_reparse (bool): The new value for the need_reparse flag. group_name (Optional[str]): If provided, the update will be applied only to the specified group. Notes: - Uses a thread-safe lock (self._db_lock) to ensure safe database access. - The group_name parameter allows scoping the update to a specific group; if not provided, the update applies to all groups containing the document. - The method commits the change to the database immediately.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def update_need_reparsing(self, doc_id: str, need_reparse: bool):
        """Updates the `need_reparse` status of a document in the `KBGroupDocuments` table.
This method sets the `need_reparse` flag for a specific document, optionally scoped to a given group.
Args:
    doc_id (str): The ID of the document to update.
    need_reparse (bool): The new value for the `need_reparse` flag.
    group_name (Optional[str]): If provided, the update will be applied only to the specified group.
Notes:
    - Uses a thread-safe lock (`self._db_lock`) to ensure safe database access.
    - The `group_name` parameter allows scoping the update to a specific group; if not provided, the update applies to all groups containing the document.
    - The method commits the change to the database immediately.
"""
        pass

validate_paths(paths) abstractmethod

Validates a list of file paths to ensure they are ready for processing. This method checks whether the provided paths are new, already processed, or currently being processed. It ensures there are no conflicts in processing the documents. Args paths (List[str]): A list of file paths to validate. Returns: Tuple[bool, str, List[bool]]: A tuple containing: - bool: True if all paths are valid, False otherwise. - str: A message indicating success or the reason for failure. - List[bool]: A list where each element corresponds to whether a path is new (True) or already exists (False). Notes: - If any document is still being processed or needs reparsing, the method returns False with an appropriate error message. - The method uses a database session and thread-safe lock (self._db_lock) to retrieve document status information. - Unsafe statuses include working and waiting.

Source code in lazyllm/tools/rag/utils.py
    @abstractmethod
    def validate_paths(self, paths: List[str]) -> Tuple[bool, str, List[bool]]:
        """Validates a list of file paths to ensure they are ready for processing.
This method checks whether the provided paths are new, already processed, or currently being processed. It ensures there are no conflicts in processing the documents.
Args
    paths (List[str]): A list of file paths to validate.
Returns:
    Tuple[bool, str, List[bool]]: A tuple containing:
        - `bool`: `True` if all paths are valid, `False` otherwise.
        - `str`: A message indicating success or the reason for failure.
        - `List[bool]`: A list where each element corresponds to whether a path is new (`True`) or already exists (`False`).
Notes:
    - If any document is still being processed or needs reparsing, the method returns `False` with an appropriate error message.
    - The method uses a database session and thread-safe lock (`self._db_lock`) to retrieve document status information.
    - Unsafe statuses include `working` and `waiting`.

"""
        pass

lazyllm.tools.rag.global_metadata.GlobalMetadataDesc

A descriptor for global metadata, defining its type, optional element type, default value, and size constraints. class GlobalMetadataDesc This class is used to describe metadata properties such as type, optional constraints, and default values. It supports scalar and array data types, with specific size limitations for certain types. Args: data_type (int): The type of the metadata as an integer, representing various data types (e.g., VARCHAR, ARRAY, etc.). element_type (Optional[int]): The type of individual elements if data_type is an array. Defaults to None. default_value (Optional[Any]): The default value for the metadata. If not provided, the default will be None. max_size (Optional[int]): The maximum size or length for the metadata. Required if data_type is VARCHAR or ARRAY.

Source code in lazyllm/tools/rag/global_metadata.py
class GlobalMetadataDesc:
    """A descriptor for global metadata, defining its type, optional element type, default value, and size constraints.
`class GlobalMetadataDesc`
This class is used to describe metadata properties such as type, optional constraints, and default values. It supports scalar and array data types, with specific size limitations for certain types.
Args:
    data_type (int): The type of the metadata as an integer, representing various data types (e.g., VARCHAR, ARRAY, etc.).
    element_type (Optional[int]): The type of individual elements if `data_type` is an array. Defaults to `None`.
    default_value (Optional[Any]): The default value for the metadata. If not provided, the default will be `None`.
    max_size (Optional[int]): The maximum size or length for the metadata. Required if `data_type` is `VARCHAR` or `ARRAY`.
"""
    # max_size MUST be set when data_type is DataType.VARCHAR or DataType.ARRAY
    def __init__(self, data_type: int, element_type: Optional[int] = None,
                 default_value: Optional[Any] = None, max_size: Optional[int] = None):
        self.data_type = data_type
        self.element_type = element_type
        self.default_value = default_value
        self.max_size = max_size

lazyllm.tools.rag.IndexBase.update(nodes) abstractmethod

Update index contents.

This method receives a list of document nodes and updates or inserts them into the index structure. Typically used for incremental indexing or refreshing data.

Parameters:

  • nodes (List[DocNode]) –

    A list of document nodes to update or insert.

Source code in lazyllm/tools/rag/index_base.py
    @abstractmethod
    def update(self, nodes: List[DocNode]) -> None:
        """Update index contents.

This method receives a list of document nodes and updates or inserts them into the index structure. Typically used for incremental indexing or refreshing data.

Args:
    nodes (List[DocNode]): A list of document nodes to update or insert.
"""
        pass

lazyllm.tools.rag.IndexBase.remove(uids, group_name=None) abstractmethod

Remove specific document nodes from the index.

Removes document nodes based on their unique identifiers, optionally scoped by group name.

Parameters:

  • uids (List[str]) –

    List of unique IDs corresponding to the document nodes to remove.

  • group_name (Optional[str], default: None ) –

    Optional group name to scope the removal operation.

Source code in lazyllm/tools/rag/index_base.py
    @abstractmethod
    def remove(self, uids: List[str], group_name: Optional[str] = None) -> None:
        """Remove specific document nodes from the index.

Removes document nodes based on their unique identifiers, optionally scoped by group name.

Args:
    uids (List[str]): List of unique IDs corresponding to the document nodes to remove.
    group_name (Optional[str]): Optional group name to scope the removal operation.
"""
        pass

lazyllm.tools.rag.IndexBase.query(*args, **kwargs) abstractmethod

Execute a query over the index.

Performs a query based on the given arguments and returns matching document nodes. The logic depends on the specific implementation.

Returns:

  • List[DocNode]

    List[DocNode]: A list of matched document nodes from the index.

Source code in lazyllm/tools/rag/index_base.py
    @abstractmethod
    def query(self, *args, **kwargs) -> List[DocNode]:
        """Execute a query over the index.

Performs a query based on the given arguments and returns matching document nodes. The logic depends on the specific implementation.

Returns:
    List[DocNode]: A list of matched document nodes from the index.
"""
        pass

lazyllm.tools.rag.index_base.IndexBase

Bases: ABC

An abstract base class for implementing indexing systems that support updating, removing, and querying document nodes. class IndexBase(ABC) This abstract base class defines the interface for an indexing system. It requires subclasses to implement methods for updating, removing, and querying document nodes.

Examples:

>>> from mymodule import IndexBase, DocNode
>>> class MyIndex(IndexBase):
...     def __init__(self):
...         self.nodes = []
...     def update(self, nodes):
...         self.nodes.extend(nodes)
...         print(f"Updated nodes: {nodes}")
...     def remove(self, uids, group_name=None):
...         self.nodes = [node for node in self.nodes if node.uid not in uids]
...         print(f"Removed nodes with uids: {uids}")
...     def query(self, *args, **kwargs):
...         print("Querying nodes...")
...         return self.nodes
>>> index = MyIndex()
>>> doc1 = DocNode(uid="1", content="Document 1")
>>> doc2 = DocNode(uid="2", content="Document 2")
>>> index.update([doc1, doc2])
Updated nodes: [DocNode(uid="1", content="Document 1"), DocNode(uid="2", content="Document 2")]
>>> index.query()
Querying nodes...
[DocNode(uid="1", content="Document 1"), DocNode(uid="2", content="Document 2")]
>>> index.remove(["1"])
Removed nodes with uids: ['1']
>>> index.query()
Querying nodes...
[DocNode(uid="2", content="Document 2")]
Source code in lazyllm/tools/rag/index_base.py
class IndexBase(ABC):
    """An abstract base class for implementing indexing systems that support updating, removing, and querying document nodes.
`class IndexBase(ABC)`
This abstract base class defines the interface for an indexing system. It requires subclasses to implement methods for updating, removing, and querying document nodes.


Examples:
    >>> from mymodule import IndexBase, DocNode
    >>> class MyIndex(IndexBase):
    ...     def __init__(self):
    ...         self.nodes = []
    ...     def update(self, nodes):
    ...         self.nodes.extend(nodes)
    ...         print(f"Updated nodes: {nodes}")
    ...     def remove(self, uids, group_name=None):
    ...         self.nodes = [node for node in self.nodes if node.uid not in uids]
    ...         print(f"Removed nodes with uids: {uids}")
    ...     def query(self, *args, **kwargs):
    ...         print("Querying nodes...")
    ...         return self.nodes
    >>> index = MyIndex()
    >>> doc1 = DocNode(uid="1", content="Document 1")
    >>> doc2 = DocNode(uid="2", content="Document 2")
    >>> index.update([doc1, doc2])
    Updated nodes: [DocNode(uid="1", content="Document 1"), DocNode(uid="2", content="Document 2")]
    >>> index.query()
    Querying nodes...
    [DocNode(uid="1", content="Document 1"), DocNode(uid="2", content="Document 2")]
    >>> index.remove(["1"])
    Removed nodes with uids: ['1']
    >>> index.query()
    Querying nodes...
    [DocNode(uid="2", content="Document 2")]
    """
    # TODO(chenjiahao): change params `nodes` to `segments`, index should be able to handle segments
    @abstractmethod
    def update(self, nodes: List[DocNode]) -> None:
        """Update index contents.

This method receives a list of document nodes and updates or inserts them into the index structure. Typically used for incremental indexing or refreshing data.

Args:
    nodes (List[DocNode]): A list of document nodes to update or insert.
"""
        pass

    @abstractmethod
    def remove(self, uids: List[str], group_name: Optional[str] = None) -> None:
        """Remove specific document nodes from the index.

Removes document nodes based on their unique identifiers, optionally scoped by group name.

Args:
    uids (List[str]): List of unique IDs corresponding to the document nodes to remove.
    group_name (Optional[str]): Optional group name to scope the removal operation.
"""
        pass

    @abstractmethod
    def query(self, *args, **kwargs) -> List[DocNode]:
        """Execute a query over the index.

Performs a query based on the given arguments and returns matching document nodes. The logic depends on the specific implementation.

Returns:
    List[DocNode]: A list of matched document nodes from the index.
"""
        pass

query(*args, **kwargs) abstractmethod

Execute a query over the index.

Performs a query based on the given arguments and returns matching document nodes. The logic depends on the specific implementation.

Returns:

  • List[DocNode]

    List[DocNode]: A list of matched document nodes from the index.

Source code in lazyllm/tools/rag/index_base.py
    @abstractmethod
    def query(self, *args, **kwargs) -> List[DocNode]:
        """Execute a query over the index.

Performs a query based on the given arguments and returns matching document nodes. The logic depends on the specific implementation.

Returns:
    List[DocNode]: A list of matched document nodes from the index.
"""
        pass

remove(uids, group_name=None) abstractmethod

Remove specific document nodes from the index.

Removes document nodes based on their unique identifiers, optionally scoped by group name.

Parameters:

  • uids (List[str]) –

    List of unique IDs corresponding to the document nodes to remove.

  • group_name (Optional[str], default: None ) –

    Optional group name to scope the removal operation.

Source code in lazyllm/tools/rag/index_base.py
    @abstractmethod
    def remove(self, uids: List[str], group_name: Optional[str] = None) -> None:
        """Remove specific document nodes from the index.

Removes document nodes based on their unique identifiers, optionally scoped by group name.

Args:
    uids (List[str]): List of unique IDs corresponding to the document nodes to remove.
    group_name (Optional[str]): Optional group name to scope the removal operation.
"""
        pass

update(nodes) abstractmethod

Update index contents.

This method receives a list of document nodes and updates or inserts them into the index structure. Typically used for incremental indexing or refreshing data.

Parameters:

  • nodes (List[DocNode]) –

    A list of document nodes to update or insert.

Source code in lazyllm/tools/rag/index_base.py
    @abstractmethod
    def update(self, nodes: List[DocNode]) -> None:
        """Update index contents.

This method receives a list of document nodes and updates or inserts them into the index structure. Typically used for incremental indexing or refreshing data.

Args:
    nodes (List[DocNode]): A list of document nodes to update or insert.
"""
        pass

lazyllm.tools.BaseEvaluator

Bases: ModuleBase

Abstract base class for evaluation modules.

This class defines the standard interface and retry logic for evaluating model outputs. It supports concurrent processing, input validation, and automatic result saving.

Parameters:

  • concurrency (int, default: 1 ) –

    Number of concurrent threads used during evaluation.

  • retry (int, default: 3 ) –

    Number of retry attempts for each evaluation item.

  • log_base_name (Optional[str], default: None ) –

    Optional log file name prefix for saving results.

Examples:

>>> from lazyllm.components import BaseEvaluator
>>> class SimpleAccuracyEvaluator(BaseEvaluator):
...     def _process_one_data_impl(self, data):
...         return {
...             "final_score": float(data["pred"] == data["label"])
...         }
>>> evaluator = SimpleAccuracyEvaluator()
>>> score = evaluator([
...     {"pred": "yes", "label": "yes"},
...     {"pred": "no", "label": "yes"}
... ])
>>> print(score)
... 0.5
Source code in lazyllm/tools/eval/eval_base.py
class BaseEvaluator(ModuleBase):
    """Abstract base class for evaluation modules.

This class defines the standard interface and retry logic for evaluating model outputs. It supports concurrent processing, input validation, and automatic result saving.

Args:
    concurrency (int): Number of concurrent threads used during evaluation.
    retry (int): Number of retry attempts for each evaluation item.
    log_base_name (Optional[str]): Optional log file name prefix for saving results.


Examples:
    >>> from lazyllm.components import BaseEvaluator
    >>> class SimpleAccuracyEvaluator(BaseEvaluator):
    ...     def _process_one_data_impl(self, data):
    ...         return {
    ...             "final_score": float(data["pred"] == data["label"])
    ...         }
    >>> evaluator = SimpleAccuracyEvaluator()
    >>> score = evaluator([
    ...     {"pred": "yes", "label": "yes"},
    ...     {"pred": "no", "label": "yes"}
    ... ])
    >>> print(score)
    ... 0.5
    """
    def __init__(self, concurrency=1, retry=3, log_base_name=None):
        super().__init__()
        self._concurrency = concurrency
        self._retry = retry
        self._lock = threading.Lock()
        self._warp = warp(self.process_one_data, _concurrent=self._concurrency)
        self._necessary_keys = []

    def _execute_with_retries(self, input_data, func, result_validator=None, post_processor=None):
        for attempt in range(1, self._retry + 1):
            try:
                result = func(input_data)
                if post_processor is not None:
                    result = post_processor(result)
                if result_validator is None or result_validator(result):
                    return result
                lazyllm.LOG.warning(f"Validation failed on attempt {attempt}/{self._retry}")
            except Exception as e:
                lazyllm.LOG.error(f"Attempt {attempt}/{self._retry} failed: {str(e)}")
        lazyllm.LOG.error(f"All {self._retry} attempts exhausted")
        return ''

    def forward(self, data):
        if not data:
            lazyllm.LOG.warning("Empty input data received")
            return 0.0

        with tqdm(total=len(data), desc=self.__class__.__name__.title()) as progress_bar:
            results = self.batch_process(data, progress_bar)

        if not results:
            return 0.0

        total_score = sum(item.get('final_score', 0) for item in results)
        return total_score / len(results)

    def process_one_data(self, data, progress_bar=None):
        res = self._process_one_data_impl(data)
        if progress_bar is not None:
            with self._lock:
                progress_bar.update(1)
        return res

    @abc.abstractmethod
    def _process_one_data_impl(self, data):
        pass

    def validate_inputs_key(self, data):
        if not isinstance(data, list):
            raise RuntimeError(f"The data should be a list, but got {type(data)}")
        for i, item in enumerate(data):
            if not isinstance(item, dict):
                raise RuntimeError(f"The item at index {i} should be a dict, but got {type(item)}")
            missing_keys = [key for key in self._necessary_keys if key not in item]
            if missing_keys:
                raise RuntimeError(
                    f"The dict at index {i} should contain "
                    f"keys: {self._necessary_keys}, but cannot find: {missing_keys}")

    def batch_process(self, data, progress_bar):
        self.validate_inputs_key(data)
        results = self._warp(data, progress_bar=progress_bar)
        self.save_res(results)
        return results

    def save_res(self, data, eval_res_save_name=None):
        save_dir = lazyllm.config['eval_result_dir']
        os.makedirs(save_dir, exist_ok=True)

        filename = eval_res_save_name or self.__class__.__name__
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        save_path = os.path.join(save_dir, f"{filename}_{timestamp}.json")
        try:
            with open(save_path, 'w') as file:
                json.dump(data, file, ensure_ascii=False, indent=4)
        except Exception as e:
            lazyllm.LOG.error(f"Dump Json error: {e}")

lazyllm.tools.ResponseRelevancy

Bases: BaseEvaluator

Evaluator for measuring the semantic relevancy between a user-generated question and a model-generated one.

This evaluator uses a language model to generate possible questions from an answer, and measures their semantic similarity to the original question using embeddings and cosine similarity.

Parameters:

  • llm (ModuleBase) –

    A language model used to generate inferred questions from the given answer.

  • embedding (ModuleBase) –

    An embedding module to encode questions for similarity comparison.

  • prompt (str, default: None ) –

    Custom prompt to guide the question generation. If not provided, a default will be used.

  • prompt_lang (str, default: 'en' ) –

    Language for the default prompt. Options: 'en' (default) or 'zh'.

  • num_infer_questions (int, default: 3 ) –

    Number of questions to generate and evaluate for each answer.

  • retry (int, default: 3 ) –

    Number of retry attempts if generation fails.

  • concurrency (int, default: 1 ) –

    Number of concurrent evaluations.

Examples:

>>> from lazyllm.components import ResponseRelevancy
>>> relevancy = ResponseRelevancy(
...     llm=YourLLM(),
...     embedding=YourEmbedding(),
...     prompt_lang="en",
...     num_infer_questions=3
... )
>>> result = relevancy([
...     {"question": "What is the capital of France?", "answer": "Paris is the capital city of France."}
... ])
>>> print(result)
... 0.95  # (a float score between 0 and 1)
Source code in lazyllm/tools/eval/rag_generator_metrics.py
class ResponseRelevancy(BaseEvaluator):
    """Evaluator for measuring the semantic relevancy between a user-generated question and a model-generated one.

This evaluator uses a language model to generate possible questions from an answer, and measures their semantic similarity to the original question using embeddings and cosine similarity.


Args:
    llm (ModuleBase): A language model used to generate inferred questions from the given answer.
    embedding (ModuleBase): An embedding module to encode questions for similarity comparison.
    prompt (str, optional): Custom prompt to guide the question generation. If not provided, a default will be used.
    prompt_lang (str): Language for the default prompt. Options: `'en'` (default) or `'zh'`.
    num_infer_questions (int): Number of questions to generate and evaluate for each answer.
    retry (int): Number of retry attempts if generation fails.
    concurrency (int): Number of concurrent evaluations.


Examples:
    >>> from lazyllm.components import ResponseRelevancy
    >>> relevancy = ResponseRelevancy(
    ...     llm=YourLLM(),
    ...     embedding=YourEmbedding(),
    ...     prompt_lang="en",
    ...     num_infer_questions=3
    ... )
    >>> result = relevancy([
    ...     {"question": "What is the capital of France?", "answer": "Paris is the capital city of France."}
    ... ])
    >>> print(result)
    ... 0.95  # (a float score between 0 and 1)
    """
    _default_generate_prompt_en = (
        'Please generate the most likely question based on '
        'the input, keeping it concise and to the point.')
    _default_generate_prompt_zh = ('请根据输入生成最可能的一个问题,保持简洁明了。')

    def __init__(self, llm, embedding, prompt=None, prompt_lang='en',
                 num_infer_questions=3, retry=3, concurrency=1):
        super().__init__(concurrency, retry)
        if prompt_lang.strip().lower() == 'zh':
            default_prompt = self._default_generate_prompt_zh
        else:
            default_prompt = self._default_generate_prompt_en
        self._llm = llm.prompt(prompt or default_prompt)
        self._embedding = embedding
        self._num_infer_questions = num_infer_questions
        self._necessary_keys = ['question', 'answer']

    def _cosine(self, x, y):
        product = np.dot(x, y)
        norm = np.linalg.norm(x) * np.linalg.norm(y)
        raw_cosine = product / norm if norm != 0 else 0.0
        return max(0.0, min(raw_cosine, 1.0))

    def _process_one_data_impl(self, data):
        one_total_score = 0
        res = copy.deepcopy(data)
        res['infer_questions'] = []
        for _ in range(self._num_infer_questions):
            # Generate Questions:
            guess_question = self._execute_with_retries(data['answer'], self._llm)

            # Calculate Similarity:
            try:
                if isinstance(self._embedding, lazyllm.module.OnlineEmbeddingModuleBase):
                    vector1 = self._embedding(guess_question)
                    vector2 = self._embedding(data['question'])
                else:
                    vector1, vector2 = json.loads(self._embedding([guess_question, data['question']]))
                score = self._cosine(vector1, vector2)
            except Exception as e:
                lazyllm.LOG.error(f'Eval-Infer Error: {e}')
                score = 0
            res['infer_questions'].append({
                'question': guess_question,
                'score': round(score, 4)
            })
            one_total_score += score
        res['final_score'] = round(one_total_score / self._num_infer_questions, 4)
        return res

lazyllm.tools.Faithfulness

Bases: BaseEvaluator

Evaluator that measures the factual consistency of an answer with the given context.

This evaluator splits the answer into atomic factual statements using a generation model, then verifies each against the context using binary (1/0) scoring. It computes a final score as the average of the individual statement scores.

Parameters:

  • llm (ModuleBase) –

    A language model capable of both generating statements and evaluating them.

  • generate_prompt (str, default: None ) –

    Custom prompt to generate factual statements from the answer.

  • eval_prompt (str, default: None ) –

    Custom prompt to evaluate statement support within the context.

  • prompt_lang (str, default: 'en' ) –

    Language of the default prompt, either 'en' or 'zh'.

  • retry (int, default: 3 ) –

    Number of retry attempts when generation or evaluation fails.

  • concurrency (int, default: 1 ) –

    Number of concurrent evaluations to run in parallel.

Examples:

>>> from lazyllm.components import Faithfulness
>>> evaluator = Faithfulness(llm=YourLLM(), prompt_lang="en")
>>> data = {
...     "question": "What is the role of ATP in cells?",
...     "answer": "ATP stores energy and transfers it within cells.",
...     "context": "ATP is the energy currency of the cell. It provides energy for many biochemical reactions."
... }
>>> result = evaluator([data])
>>> print(result)
... 1.0  # Average binary score of all factual statements
Source code in lazyllm/tools/eval/rag_generator_metrics.py
class Faithfulness(BaseEvaluator):
    """Evaluator that measures the factual consistency of an answer with the given context.

This evaluator splits the answer into atomic factual statements using a generation model, then verifies each against the context using binary (1/0) scoring. It computes a final score as the average of the individual statement scores.


Args:
    llm (ModuleBase): A language model capable of both generating statements and evaluating them.
    generate_prompt (str, optional): Custom prompt to generate factual statements from the answer.
    eval_prompt (str, optional): Custom prompt to evaluate statement support within the context.
    prompt_lang (str): Language of the default prompt, either 'en' or 'zh'.
    retry (int): Number of retry attempts when generation or evaluation fails.
    concurrency (int): Number of concurrent evaluations to run in parallel.


Examples:
    >>> from lazyllm.components import Faithfulness
    >>> evaluator = Faithfulness(llm=YourLLM(), prompt_lang="en")
    >>> data = {
    ...     "question": "What is the role of ATP in cells?",
    ...     "answer": "ATP stores energy and transfers it within cells.",
    ...     "context": "ATP is the energy currency of the cell. It provides energy for many biochemical reactions."
    ... }
    >>> result = evaluator([data])
    >>> print(result)
    ... 1.0  # Average binary score of all factual statements
    """
    _default_generate_prompt_en = (
        '[Task Description]\n'
        'Split the answer into independent factual statements using "|||" as '
        'the exclusive separator, following these rules:\n'
        '1. Each statement must be a complete sentence ending with proper punctuation\n'
        '2. Never use line breaks or other symbols as separators\n'
        '3. Statements containing "|||" must be rephrased\n'
        '4. Each statement must be clear, pronoun-free.\n'
        '[Output Format]\n'
        'statement_1|||statement_2|||statement_3\n'
        '[Example Input]\n'
        'Q: How does photosynthesis work?\n'
        'A: The process requires sunlight, then chlorophyll absorbs light energy. '
        'It converts water and CO2 into glucose.\n'
        '[Example Output]\n'
        'Photosynthesis requires sunlight.|||Chlorophyll absorbs light energy.'
        '|||Chlorophyll converts water and CO2 into glucose.\n'
    )
    _default_eval_prompt_en = (
        '[Task Description]\n'
        'Evaluate each "|||"-separated statement against provided context using binary scoring:\n'
        'Fully supported by context: 1\n'
        'Unsupported/contradictory: 0\n'
        '[Output Requirements]\n'
        '1. JSON format with array of objects\n'
        '2. Each object contains:\n'
        '    - "statement": Original text\n'
        '    - "score": 1 or 0\n'
        '3. Wrap output in ```json code block\n'
        '[Example Input]\n'
        'Context: Photosynthesis occurs in chloroplasts. Light reactions produce ATP using sunlight. '
        'Calvin cycle fixes CO2 into sugars.\n'
        'Statements: Photosynthesis requires sunlight.|||Chlorophyll absorbs light energy.'
        '|||Chlorophyll converts water and CO2 into glucose.\n'
        '[Example Output]\n'
        '[{"statement": "Photosynthesis requires sunlight.","score": 1},'
        '{"statement": "Chlorophyll absorbs light energy.", "score": 1},'
        '{"statement": "Chlorophyll converts water and CO2 into glucose.","score": 0}]\n'
    )
    _default_generate_prompt_zh = (
        '[任务描述]\n'
        '使用"|||"作为唯一分隔符,将答案分割成独立的事实陈述,遵循以下规则:\n'
        '1. 每个陈述必须是完整的句子,并以适当的标点结束\n'
        '2. 不要使用换行符或其他符号作为分隔符\n'
        '3. 包含"|||"的陈述必须重新措辞\n'
        '4. 每个陈述必须清晰,不包含代词。\n'
        '[输出格式]\n'
        'statement_1|||statement_2|||statement_3\n'
        '[示例输入]\n'
        'Q: 光合作用是如何工作的?\n'
        'A: 该过程需要阳光,然后叶绿素吸收光能。它将水和CO2转化为葡萄糖。\n'
        '[示例输出]\n'
        '光合作用需要阳光。|||叶绿素吸收光能。|||叶绿素将水和CO2转化为葡萄糖。\n'
    )
    _default_eval_prompt_zh = (
        '[任务描述]\n'
        '使用二进制评分对每个"|||"分隔的陈述与提供的内容进行评估:\n'
        '完全由内容支持:1\n'
        '不支持/矛盾:0\n'
        '[输出要求]\n'
        '1. JSON格式,包含对象数组\n'
        '2. 每个对象包含:\n'
        '    - "statement": 原始文本\n'
        '    - "score": 1或0\n'
        '3. 将输出包裹在```json代码块中\n'
        '[示例输入]\n'
        'Context: 光合作用发生在叶绿体中。光反应利用阳光产生ATP。卡尔文循环将CO2固定成糖。\n'
        'Statements: 光合作用需要阳光。|||叶绿素吸收光能。|||叶绿素将水和CO2转化为葡萄糖。\n'
        '[示例输出]\n'
        '[{"statement": "光合作用需要阳光。","score": 1},'
        '{"statement": "叶绿素吸收光能。", "score": 1},'
        '{"statement": "叶绿素将水和CO2转化为葡萄糖。","score": 0}]\n'
    )

    def __init__(self, llm, generate_prompt=None, eval_prompt=None, prompt_lang='en', retry=3, concurrency=1):
        super().__init__(concurrency, retry)
        self._base_llm = llm
        if prompt_lang == 'zh':
            default_generate_prompt = generate_prompt or self._default_generate_prompt_zh
            default_eval_prompt = eval_prompt or self._default_eval_prompt_zh
        else:
            default_generate_prompt = generate_prompt or self._default_generate_prompt_en
            default_eval_prompt = eval_prompt or self._default_eval_prompt_en
        self._build_llms(self._base_llm, default_generate_prompt, default_eval_prompt)
        self._necessary_keys = ['question', 'answer', 'context']

    def _build_llms(self, base_llm, generate_prompt, eval_prompt):
        self._gene_llm = base_llm.share(prompt=generate_prompt)
        self._eval_llm = base_llm.share(prompt=eval_prompt).formatter(JsonFormatter())

    def _validate_eval_result(self, result):
        return (
            isinstance(result, list)
            and len(result) > 0
            and all(isinstance(i, dict) and 'score' in i for i in result)
        )

    def _post_processor(self, eval_result):
        if isinstance(eval_result, dict):
            eval_result = [eval_result]
        return eval_result

    def _process_one_data_impl(self, data):
        res = copy.deepcopy(data)
        # Generate Statements:
        query1 = f'Q: {data["question"]}\nA: {data["answer"]}'
        statements = self._execute_with_retries(query1, self._gene_llm)
        res['statements'] = statements

        # Eval Statements in Context:
        query2 = f'Context: {data["context"]}\nStatements: {statements}'
        eval_result = self._execute_with_retries(
            query2, self._eval_llm, self._validate_eval_result, self._post_processor)
        if not self._validate_eval_result(eval_result):
            lazyllm.LOG.error("Invalid evaluation result format")
            res.update({'scores': [], 'final_score': 0.0})
            return res

        total_score = sum(
            int(entry.get('score', 0)) if entry.get('score') in (0, 1) else 0
            for entry in eval_result
        )
        res['scores'] = eval_result
        res['final_score'] = round(total_score / len(eval_result), 4) if eval_result else 0.0
        return res

lazyllm.tools.LLMContextRecall

Bases: BaseEvaluator

Source code in lazyllm/tools/eval/rag_retriever_metrics.py
class LLMContextRecall(BaseEvaluator):
    _default_eval_prompt_en = (
        '[Task Description]\n'
        'Given a context, and an answer, analyze each sentence in the answer and '
        'classify if the sentence can be attributed to the given context or not:\n'
        'Fully supported by context: 1\n'
        'Unsupported/contradictory: 0\n'
        '[Output Requirements]\n'
        '1. JSON format with array of objects\n'
        '2. Each object contains:\n'
        '    - "statement": Original text\n'
        '    - "reason": the reason why it is scored 1/0\n'
        '    - "score": 1 or 0\n'
        '3. Wrap output in ```json code block\n'
        '[Example Input]\n'
        'Question: What is Photosynthesis?'
        'Context: Photosynthesis occurs in chloroplasts. Light reactions produce ATP using sunlight.\n'
        'Statements: Photosynthesis was discovered in 1780s. It occurs in chloroplasts and produce ATP using sunlight.\n'
        '[Example Output]\n'
        '[{"statement": "Photosynthesis was discovered in 1780s", '
        '"reason": "The time when photosynthesis discovered was not mentioned in the given context","score": 0},'
        ' {"statement": "It occurs in chloroplasts and produce ATP using sunlight.", '
        '"reason": "The exact sentence is present in the given context", "score": 1}]\n'
    )
    _default_eval_prompt_zh = (
        '[任务描述]\n'
        '给定一个上下文和一个答案,分析答案中的每个句子并判断该句子是否可以归因于给定的上下文:\n'
        '完全受上下文支持:1\n'
        '不支持/矛盾:0\n'
        '[输出要求]\n'
        '1. 带有对象数组的 JSON 格式\n'
        '2. 每个对象包含:\n'
        ' - "statement":原始文本\n'
        ' - "reason":评分原因\n'
        ' - "score":1 或 0\n'
        '3. 将输出包裹在 ```json 代码块中\n'
        '[示例输入]\n'
        'question:什么是光合作用?'
        'context:光合作用发生在叶绿体中,利用阳光产生 ATP。\n'
        'statement:光合作用于 1780 年代被发现。光合作用发生在叶绿体中,并利用阳光产生 ATP。\n'
        '[示例输出]\n'
        '[{"statement": "光合作用于 1780 年代被发现", "reason": "给定上下文中未提及发现光合作用被发现的时间","score": 0},'
        ' {"statement": "光合作用发生在叶绿体中,并利用阳光产生 ATP。", "reason": "给定上下文中存在确切的句子", "score": 1}]\n'
    )

    def __init__(self, llm, eval_prompt=None, prompt_lang='en', retry=3, concurrency=1):
        super().__init__(concurrency, retry)
        if prompt_lang == 'zh':
            default_eval_prompt = eval_prompt or self._default_eval_prompt_zh
        else:
            default_eval_prompt = eval_prompt or self._default_eval_prompt_en
        self._llm = llm.prompt(default_eval_prompt).formatter(JsonFormatter()) if llm else None
        self._necessary_keys = ['question', 'answer', 'context_retrieved']

    def _validate_eval_result(self, result):
        return (
            isinstance(result, list)
            and len(result) > 0
            and all(isinstance(i, dict) and 'score' in i for i in result)
        )

    def _post_processor(self, eval_result):
        if isinstance(eval_result, dict):
            eval_result = [eval_result]
        return eval_result

    def _process_one_data_impl(self, data):
        res = copy.deepcopy(data)
        context = "\n".join(data['context_retrieved'])

        query = f'question: {data["question"]}\ncontext: {context}\nstatement: {data["answer"]}'
        eval_result = self._execute_with_retries(
            query, self._llm, self._validate_eval_result, self._post_processor)
        scores = [result["score"] for result in eval_result]

        res['final_score'] = round(sum(scores) / len(scores), 4) if scores else 0.0
        return res

lazyllm.tools.NonLLMContextRecall

Bases: BaseEvaluator

Source code in lazyllm/tools/eval/rag_retriever_metrics.py
class NonLLMContextRecall(BaseEvaluator):
    def __init__(self, th=0.5, binary=True, retry=3, concurrency=1):
        super().__init__(concurrency, retry)
        self._binary = binary
        self._threshold = th
        self._necessary_keys = ['context_retrieved', 'context_reference']

    def _calc_levenshtein_distance(self, reference, context):
        return 1 - rapidfuzz.distance.Levenshtein.normalized_distance(reference, context)

    def _calc_context_recall(self, data):
        contexts, reference = data["context"], data["reference"]
        scores = []
        for context in contexts:
            score = self._calc_levenshtein_distance(reference, context)
            scores.append(score)
        return scores

    def _compute_scores(self, scores):
        binary_scores = [1 if score > self._threshold else 0 for score in scores]

        if self._binary:
            return 1.0 if sum(binary_scores) > 0 else 0.0
        if len(binary_scores) > 0:
            return sum(binary_scores) / len(binary_scores)
        return 0

    def _process_one_data_impl(self, data):
        res = copy.deepcopy(data)
        scores = []
        for reference in data['context_reference']:
            input_data = {'context': data['context_retrieved'], 'reference': reference}
            eval_result = self._execute_with_retries(input_data, self._calc_context_recall)
            scores.append(self._compute_scores(eval_result))

        res['final_score'] = round(sum(scores) / len(scores), 4) if scores else 0.0
        return res

lazyllm.tools.ContextRelevance

Bases: BaseEvaluator

Source code in lazyllm/tools/eval/rag_retriever_metrics.py
class ContextRelevance(BaseEvaluator):
    def __init__(self, splitter="。", retry=3, concurrency=1):
        super().__init__(concurrency, retry)
        self._splitter = splitter
        self._necessary_keys = ['context_retrieved', 'context_reference']

    def _calc_context_relevance(self, data):
        sentences_retrieved, sentences_reference = data["context"], data["reference"]
        scores = [0] * len(sentences_retrieved)
        for i, sentence in enumerate(sentences_retrieved):
            if sentence in sentences_reference:
                scores[i] = 1
        return scores

    def _paragraphs_to_sentences(self, paragraphs):
        sentences = []
        pattern = rf'{re.escape(self._splitter)}+'
        for paragraph in paragraphs:
            sentences.extend([s.strip() for s in re.split(pattern, paragraph) if s.strip()])
        return sentences

    def _process_one_data_impl(self, data):
        res = copy.deepcopy(data)
        retrieved = self._paragraphs_to_sentences(data["context_retrieved"])
        reference = self._paragraphs_to_sentences(data["context_reference"])

        input_data = {'context': retrieved, 'reference': reference}
        eval_result = self._execute_with_retries(input_data, self._calc_context_relevance)
        total_score = sum(eval_result)

        res['final_score'] = round(total_score / len(eval_result), 4) if eval_result else 0.0
        return res

lazyllm.tools.HttpRequest

Bases: ModuleBase

General HTTP request executor.

This class builds and sends HTTP requests with support for dynamic variable substitution, API key injection, JSON or form data encoding, and file-aware response parsing.

Parameters:

  • method (str) –

    HTTP method, such as 'GET', 'POST', etc.

  • url (str) –

    The target URL for the HTTP request.

  • api_key (str) –

    Optional API key, inserted into query parameters.

  • headers (dict) –

    HTTP request headers.

  • params (dict) –

    URL query parameters.

  • body (Union[str, dict]) –

    HTTP request body (raw string or JSON-formatted dict).

  • timeout (int, default: 10 ) –

    Timeout duration for the request (in seconds).

  • proxies (dict, default: None ) –

    Proxy settings for the request, if needed.

Examples:

>>> from lazyllm.components import HttpRequest
>>> request = HttpRequest(
...     method="GET",
...     url="https://api.github.com/repos/openai/openai-python",
...     api_key="",
...     headers={"Accept": "application/json"},
...     params={},
...     body=None
... )
>>> result = request()
>>> print(result["status_code"])
... 200
>>> print(result["content"][:100])
... '{"id":123456,"name":"openai-python", ...}'
Source code in lazyllm/tools/http_request/http_request.py
class HttpRequest(ModuleBase):
    """General HTTP request executor.

This class builds and sends HTTP requests with support for dynamic variable substitution, API key injection, JSON or form data encoding, and file-aware response parsing.

Args:
    method (str): HTTP method, such as 'GET', 'POST', etc.
    url (str): The target URL for the HTTP request.
    api_key (str): Optional API key, inserted into query parameters.
    headers (dict): HTTP request headers.
    params (dict): URL query parameters.
    body (Union[str, dict]): HTTP request body (raw string or JSON-formatted dict).
    timeout (int): Timeout duration for the request (in seconds).
    proxies (dict, optional): Proxy settings for the request, if needed.


Examples:
    >>> from lazyllm.components import HttpRequest
    >>> request = HttpRequest(
    ...     method="GET",
    ...     url="https://api.github.com/repos/openai/openai-python",
    ...     api_key="",
    ...     headers={"Accept": "application/json"},
    ...     params={},
    ...     body=None
    ... )
    >>> result = request()
    >>> print(result["status_code"])
    ... 200
    >>> print(result["content"][:100])
    ... '{"id":123456,"name":"openai-python", ...}'
    """
    def __init__(self, method, url, api_key, headers, params, body, timeout=10, proxies=None):
        super().__init__()
        if not url:
            return

        self._method = method
        self._url = url
        self._api_key = api_key
        self._headers = headers
        self._params = params
        self._body = body
        self._timeout = timeout
        self._proxies = proxies

    def _process_api_key(self, headers, params):
        if self._api_key and self._api_key != '':
            params = params or {}
            params['api_key'] = self._api_key
        return headers, params

    def forward(self, *args, **kwargs):
        def _map_input(target_str):
            if not isinstance(target_str, str):
                return target_str

            # TODO: replacements could be more complex to create.
            replacements = {**kwargs, **(args[0] if args and isinstance(args[0], dict) else {})}
            if not replacements:
                return target_str

            pattern = r"\{\{([^}]+)\}\}"
            matches = re.findall(pattern, target_str)
            for match in matches:
                replacement = replacements.get(match)
                if replacement is not None:
                    if "{{" + match + "}}" == target_str:
                        return replacement
                    target_str = re.sub(r"\{\{" + re.escape(match) + r"\}\}", replacement, target_str)

            return target_str

        url = _map_input(self._url)
        params = {key: _map_input(value) for key, value in self._params.items()} if self._params else None
        headers = {key: _map_input(value) for key, value in self._headers.items()} if self._headers else None
        headers, params = self._process_api_key(headers, params)
        if isinstance(headers, dict) and headers.get("Content-Type") == "application/json":
            try:
                body = json.loads(self._body) if isinstance(self._body, str) else self._body
                body = {k: _map_input(v) for k, v in body.items()}

                http_response = httpx.request(method=self._method, url=url, headers=headers,
                                              params=params, json=body, timeout=self._timeout,
                                              proxies=self._proxies)
            except json.JSONDecodeError:
                raise ValueError(f"Invalid JSON format: {self._body}")
        else:
            body = (json.dumps({k: _map_input(v) for k, v in self._body.items()})
                    if isinstance(self._body, dict) else _map_input(self._body))

            http_response = httpx.request(method=self._method, url=url, headers=headers,
                                          params=params, data=body, timeout=self._timeout,
                                          proxies=self._proxies)

        response = HttpExecutorResponse(http_response)

        _, file_binary = response.extract_file()

        outputs = {
            'status_code': response.status_code,
            'content': response.content if len(file_binary) == 0 else None,
            'headers': response.headers,
            'file': file_binary
        }
        return outputs

lazyllm.tools.JobDescription

Bases: BaseModel

Model deployment job description schema.

Used to specify the configuration for creating a model inference job, including model name and GPU requirements.

Parameters:

  • deploy_model (str) –

    The model to be deployed. Default is "qwen1.5-0.5b-chat".

  • num_gpus (int) –

    Number of GPUs required for deployment. Default is 1.

Examples:

>>> from lazyllm.components import JobDescription
>>> job = JobDescription(deploy_model="deepseek-coder", num_gpus=2)
>>> print(job.dict())
... {'deploy_model': 'deepseek-coder', 'num_gpus': 2}
Source code in lazyllm/tools/infer_service/serve.py
class JobDescription(BaseModel):
    """Model deployment job description schema.

Used to specify the configuration for creating a model inference job, including model name and GPU requirements.

Args:
    deploy_model (str): The model to be deployed. Default is "qwen1.5-0.5b-chat".
    num_gpus (int): Number of GPUs required for deployment. Default is 1.


Examples:
    >>> from lazyllm.components import JobDescription
    >>> job = JobDescription(deploy_model="deepseek-coder", num_gpus=2)
    >>> print(job.dict())
    ... {'deploy_model': 'deepseek-coder', 'num_gpus': 2}
    """
    deploy_model: str = Field(default='qwen1.5-0.5b-chat')
    num_gpus: int = Field(default=1)

lazyllm.tools.DBManager

Bases: ABC, ModuleBase

Abstract base class for database managers.

This class defines the standard interface and helpers for building database connectors, including a required execute_query method and description property.

Parameters:

  • db_type (str) –

    Type identifier of the database (e.g., 'mysql', 'mongodb').

Examples:

>>> from lazyllm.components import DBManager
>>> class DummyDB(DBManager):
...     def __init__(self):
...         super().__init__(db_type="dummy")
...     def execute_query(self, statement):
...         return f"Executed: {statement}"
...     @property
...     def desc(self):
...         return "Dummy database for testing."
>>> db = DummyDB()
>>> print(db("SELECT * FROM test"))
... Executed: SELECT * FROM test
Source code in lazyllm/tools/sql/db_manager.py
class DBManager(ABC, ModuleBase, metaclass=CommonMeta):
    """Abstract base class for database managers.

This class defines the standard interface and helpers for building database connectors, including a required `execute_query` method and description property.

Args:
    db_type (str): Type identifier of the database (e.g., 'mysql', 'mongodb').


Examples:
    >>> from lazyllm.components import DBManager
    >>> class DummyDB(DBManager):
    ...     def __init__(self):
    ...         super().__init__(db_type="dummy")
    ...     def execute_query(self, statement):
    ...         return f"Executed: {statement}"
    ...     @property
    ...     def desc(self):
    ...         return "Dummy database for testing."
    >>> db = DummyDB()
    >>> print(db("SELECT * FROM test"))
    ... Executed: SELECT * FROM test
    """

    def __init__(self, db_type: str):
        ModuleBase.__init__(self)
        self._db_type = db_type
        self._desc = None

    @abstractmethod
    def execute_query(self, statement) -> str:
        """Abstract method for executing database query statements. This method needs to be implemented by specific database manager subclasses to execute various database operations.

Args:
    statement: The database query statement to execute, which can be SQL statements or other database-specific query languages

Features of this method:

- **Abstract Method**: Requires implementation of specific database operation logic in subclasses
- **Unified Interface**: Provides a unified query interface for different database types
- **Error Handling**: Subclass implementations should include appropriate error handling and status reporting
- **Result Formatting**: Returns formatted string results for subsequent processing

**Note**: This method is the core method of the database manager, and all specific database operations are executed through this method.

"""
        pass

    def forward(self, statement: str) -> str:
        return self.execute_query(statement)

    @property
    def db_type(self) -> str:
        return self._db_type

    @property
    @abstractmethod
    def desc(self) -> str: pass

    @staticmethod
    def _is_dict_all_str(d):
        if not isinstance(d, dict):
            return False
        return all(isinstance(key, str) and (isinstance(value, str) or DBManager._is_dict_all_str(value))
                   for key, value in d.items())

    @staticmethod
    def _serialize_uncommon_type(obj):
        if not isinstance(obj, (int, str, float, bool, tuple, list, dict)):
            return str(obj)

execute_query(statement) abstractmethod

Abstract method for executing database query statements. This method needs to be implemented by specific database manager subclasses to execute various database operations.

Parameters:

  • statement

    The database query statement to execute, which can be SQL statements or other database-specific query languages

Features of this method:

  • Abstract Method: Requires implementation of specific database operation logic in subclasses
  • Unified Interface: Provides a unified query interface for different database types
  • Error Handling: Subclass implementations should include appropriate error handling and status reporting
  • Result Formatting: Returns formatted string results for subsequent processing

Note: This method is the core method of the database manager, and all specific database operations are executed through this method.

Source code in lazyllm/tools/sql/db_manager.py
    @abstractmethod
    def execute_query(self, statement) -> str:
        """Abstract method for executing database query statements. This method needs to be implemented by specific database manager subclasses to execute various database operations.

Args:
    statement: The database query statement to execute, which can be SQL statements or other database-specific query languages

Features of this method:

- **Abstract Method**: Requires implementation of specific database operation logic in subclasses
- **Unified Interface**: Provides a unified query interface for different database types
- **Error Handling**: Subclass implementations should include appropriate error handling and status reporting
- **Result Formatting**: Returns formatted string results for subsequent processing

**Note**: This method is the core method of the database manager, and all specific database operations are executed through this method.

"""
        pass

lazyllm.tools.MongoDBManager

Bases: DBManager

Source code in lazyllm/tools/sql/mongodb_manager.py
class MongoDBManager(DBManager):
    MAX_TIMEOUT_MS = 5000

    def __init__(self, user: str, password: str, host: str, port: int, db_name: str, collection_name: str, **kwargs):
        super().__init__(db_type="mongodb")
        self._user = user
        self._password = password
        self._host = host
        self._port = port
        self._db_name = db_name
        self._collection_name = collection_name
        self._collection = None
        self._options_str = kwargs.get("options_str")
        self._conn_url = self._gen_conn_url()
        self._collection_desc_dict = kwargs.get("collection_desc_dict")

    @property
    def db_name(self):
        return self._db_name

    @property
    def collection_name(self):
        return self._collection_name

    def _gen_conn_url(self) -> str:
        password = quote_plus(self._password)
        conn_url = (f"{self._db_type}://{self._user}:{password}@{self._host}:{self._port}/"
                    f"{('?' + self._options_str) if self._options_str else ''}")
        return conn_url

    @contextmanager
    def get_client(self):
        client = pymongo.MongoClient(self._conn_url, serverSelectionTimeoutMS=self.MAX_TIMEOUT_MS)
        try:
            yield client
        finally:
            client.close()

    @property
    def desc(self):
        if self._desc is None:
            self.set_desc(schema_desc_dict=self._collection_desc_dict)
        return self._desc

    def set_desc(self, schema_desc_dict: dict):
        self._collection_desc_dict = schema_desc_dict
        if schema_desc_dict is None:
            with self.get_client() as client:
                egs_one = client[self._db_name][self._collection_name].find_one()
                if egs_one is not None:
                    self._desc = "Collection Example:\n"
                    self._desc += json.dumps(egs_one, ensure_ascii=False, indent=4)
        else:
            self._desc = ""
            try:
                collection_desc = CollectionDesc.model_validate(schema_desc_dict)
            except pydantic.ValidationError as e:
                raise ValueError(f"Validate input schema_desc_dict failed: {str(e)}")
            if not self._is_dict_all_str(collection_desc.schema_type):
                raise ValueError("schema_type shouble be str or nested str dict")
            if not self._is_dict_all_str(collection_desc.schema_desc):
                raise ValueError("schema_desc shouble be str or nested str dict")
            if collection_desc.summary:
                self._desc += f"Collection summary: {collection_desc.summary}\n"
            self._desc += "Collection schema:\n"
            self._desc += json.dumps(collection_desc.schema_type, ensure_ascii=False, indent=4)
            self._desc += "Collection schema description:\n"
            self._desc += json.dumps(collection_desc.schema_type, ensure_ascii=False, indent=4)

    def check_connection(self) -> DBResult:
        try:
            with pymongo.MongoClient(self._conn_url, serverSelectionTimeoutMS=self.MAX_TIMEOUT_MS) as client:
                _ = client.server_info()
            return DBResult()
        except Exception as e:
            return DBResult(status=DBStatus.FAIL, detail=str(e))

    def execute_query(self, statement) -> str:
        str_result = ""
        try:
            pipeline_list = json.loads(statement)
            with self.get_client() as client:
                collection = client[self._db_name][self._collection_name]
                result = list(collection.aggregate(pipeline_list))
                str_result = json.dumps(result, ensure_ascii=False, default=self._serialize_uncommon_type)
        except Exception as e:
            str_result = f"MongoDB ERROR: {str(e)}"
        return str_result

lazyllm.tools.HttpTool

Bases: HttpRequest

Module for accessing third-party services and executing custom code. The values in params and headers, as well as in body, can include template variables marked with double curly braces like {{variable}}, which are then replaced with actual values through parameters when called. Refer to the usage instructions in [[lazyllm.tools.HttpTool.forward]].

Parameters:

  • method (str, default: None ) –

    Specifies the HTTP request method, refer to https://developer.mozilla.org/en-US/docs/Web/HTTP/Methods.

  • url (str, default: None ) –

    The URL to access. If this field is empty, it indicates that the module does not need to access third-party services.

  • params (Dict[str, str], default: None ) –

    Params fields to be filled when requesting the URL. If the URL is empty, this field will be ignored.

  • headers (Dict[str, str], default: None ) –

    Header fields to be filled when accessing the URL. If the URL is empty, this field will be ignored.

  • body (Dict[str, str], default: None ) –

    Body fields to be filled when requesting the URL. If the URL is empty, this field will be ignored.

  • timeout (int, default: 10 ) –

    Request timeout in seconds, default value is 10.

  • proxies (Dict[str, str], default: None ) –

    Specifies the proxies to be used when requesting the URL. Proxy format refer to https://www.python-httpx.org/advanced/proxies.

  • code_str (str, default: None ) –

    A string containing a user-defined function. If the parameter url is empty, execute this function directly, forwarding all arguments to it; if url is not empty, the parameters of this function are the results returned from the URL request, and in this case, the function serves as a post-processing function for the URL response.

  • vars_for_code (Dict[str, Any], default: None ) –

    A dictionary that includes dependencies and variables required for running the code.

  • outputs (Optional[List[str]], default: None ) –

    Names of expected output fields.

  • extract_from_result (Optional[bool], default: None ) –

    Whether to extract fields directly from response dict using outputs.

Examples:

from lazyllm.tools import HttpTool

code_str = "def identity(content): return content"
tool = HttpTool(method='GET', url='http://www.sensetime.com/', code_str=code_str)
ret = tool()
Source code in lazyllm/tools/tools/http_tool.py
class HttpTool(HttpRequest):
    """
Module for accessing third-party services and executing custom code. The values in `params` and `headers`, as well as in body, can include template variables marked with double curly braces like `{{variable}}`, which are then replaced with actual values through parameters when called. Refer to the usage instructions in [[lazyllm.tools.HttpTool.forward]].

Args:
    method (str, optional): Specifies the HTTP request method, refer to `https://developer.mozilla.org/en-US/docs/Web/HTTP/Methods`.
    url (str, optional): The URL to access. If this field is empty, it indicates that the module does not need to access third-party services.
    params (Dict[str, str], optional): Params fields to be filled when requesting the URL. If the URL is empty, this field will be ignored.
    headers (Dict[str, str], optional): Header fields to be filled when accessing the URL. If the URL is empty, this field will be ignored.
    body (Dict[str, str], optional): Body fields to be filled when requesting the URL. If the URL is empty, this field will be ignored.
    timeout (int): Request timeout in seconds, default value is 10.
    proxies (Dict[str, str], optional): Specifies the proxies to be used when requesting the URL. Proxy format refer to `https://www.python-httpx.org/advanced/proxies`.
    code_str (str, optional): A string containing a user-defined function. If the parameter url is empty, execute this function directly, forwarding all arguments to it; if url is not empty, the parameters of this function are the results returned from the URL request, and in this case, the function serves as a post-processing function for the URL response.
    vars_for_code (Dict[str, Any]): A dictionary that includes dependencies and variables required for running the code.
    outputs (Optional[List[str]]): Names of expected output fields.
    extract_from_result (Optional[bool]): Whether to extract fields directly from response dict using `outputs`.


Examples:

    from lazyllm.tools import HttpTool

    code_str = "def identity(content): return content"
    tool = HttpTool(method='GET', url='http://www.sensetime.com/', code_str=code_str)
    ret = tool()
    """
    def __init__(self,
                 method: Optional[str] = None,
                 url: Optional[str] = None,
                 params: Optional[Dict[str, str]] = None,
                 headers: Optional[Dict[str, str]] = None,
                 body: Optional[str] = None,
                 timeout: int = 10,
                 proxies: Optional[Dict[str, str]] = None,
                 code_str: Optional[str] = None,
                 vars_for_code: Optional[Dict[str, Any]] = None,
                 outputs: Optional[List[str]] = None,
                 extract_from_result: Optional[bool] = None):
        super().__init__(method, url, '', headers, params, body, timeout, proxies)
        self._has_http = True if url else False
        self._compiled_func = (compile_func(code_str, vars_for_code) if code_str else
                               (lambda x: json.loads(x['content'])) if self._has_http else None)
        self._outputs, self._extract_from_result = outputs, extract_from_result
        if extract_from_result:
            assert outputs, 'Output information is necessary to extract output parameters'
            assert len(outputs) == 1, 'When the number of outputs is greater than 1, no manual setting is required'

    def _get_result(self, res):
        if self._extract_from_result or (isinstance(res, dict) and len(self._outputs) > 1):
            assert isinstance(res, dict), 'The result of the tool should be a dict type'
            r = package(res.get(key) for key in self._outputs)
            return r[0] if len(r) == 1 else r
        if len(self._outputs) > 1:
            assert isinstance(res, (tuple, list)), 'The result of the tool should be tuple or list'
            assert len(res) == len(self._outputs), 'The number of outputs is inconsistent with expectations'
            return package(res)
        return res

    def forward(self, *args, **kwargs):
        """
Used to perform operations specified during initialization: request the specified URL or execute the passed function. Generally not called directly, but through the base class's `__call__`. If the `url` parameter in the constructor is not empty, all passed parameters will be used as variables to replace template parameters marked with `{{}}` in the constructor; if the `url` parameter in the constructor is empty and `code_str` is not empty, all passed parameters will be used as arguments for the function defined in `code_str`.


Examples:

    from lazyllm.tools import HttpTool

    code_str = "def exp(v, n): return v ** n"
    tool = HttpTool(code_str=code_str)
    assert tool(v=10, n=2) == 100
    """
        if not self._compiled_func: return None
        if self._has_http:
            res = super().forward(*args, **kwargs)
            if int(res['status_code']) >= 400:
                raise RuntimeError(f'HttpRequest error, status code is {res["status_code"]}.')
            args, kwargs = (res,), {}
        res = self._compiled_func(*args, **kwargs)
        return self._get_result(res) if self._outputs else res

forward(*args, **kwargs)

Used to perform operations specified during initialization: request the specified URL or execute the passed function. Generally not called directly, but through the base class's __call__. If the url parameter in the constructor is not empty, all passed parameters will be used as variables to replace template parameters marked with {{}} in the constructor; if the url parameter in the constructor is empty and code_str is not empty, all passed parameters will be used as arguments for the function defined in code_str.

Examples:

from lazyllm.tools import HttpTool

code_str = "def exp(v, n): return v ** n"
tool = HttpTool(code_str=code_str)
assert tool(v=10, n=2) == 100
Source code in lazyllm/tools/tools/http_tool.py
    def forward(self, *args, **kwargs):
        """
Used to perform operations specified during initialization: request the specified URL or execute the passed function. Generally not called directly, but through the base class's `__call__`. If the `url` parameter in the constructor is not empty, all passed parameters will be used as variables to replace template parameters marked with `{{}}` in the constructor; if the `url` parameter in the constructor is empty and `code_str` is not empty, all passed parameters will be used as arguments for the function defined in `code_str`.


Examples:

    from lazyllm.tools import HttpTool

    code_str = "def exp(v, n): return v ** n"
    tool = HttpTool(code_str=code_str)
    assert tool(v=10, n=2) == 100
    """
        if not self._compiled_func: return None
        if self._has_http:
            res = super().forward(*args, **kwargs)
            if int(res['status_code']) >= 400:
                raise RuntimeError(f'HttpRequest error, status code is {res["status_code"]}.')
            args, kwargs = (res,), {}
        res = self._compiled_func(*args, **kwargs)
        return self._get_result(res) if self._outputs else res

lazyllm.tools.agent.functionCall.StreamResponse

StreamResponse class encapsulates streaming output behavior with configurable prefix and colors.
When streaming is enabled, calling the instance enqueues colored text to a filesystem queue for asynchronous processing or display.

Parameters:

  • prefix (str) –

    Prefix text before the output, typically used to indicate the source or category.

  • prefix_color (Optional[str], default: None ) –

    Color of the prefix text, supports terminal color codes, defaults to None.

  • color (Optional[str], default: None ) –

    Color of the main content text, supports terminal color codes, defaults to None.

  • stream (bool, default: False ) –

    Whether to enable streaming output mode, which enqueues text to the filesystem queue, defaults to False.

Examples:

>>> from lazyllm.tools.agent.functionCall import StreamResponse
>>> resp = StreamResponse(prefix="[INFO]", prefix_color="green", color="white", stream=True)
>>> resp("Hello, world!")
Hello, world!
Source code in lazyllm/tools/agent/functionCall.py
class StreamResponse():
    """StreamResponse class encapsulates streaming output behavior with configurable prefix and colors.  
When streaming is enabled, calling the instance enqueues colored text to a filesystem queue for asynchronous processing or display.

Args:
    prefix (str): Prefix text before the output, typically used to indicate the source or category.
    prefix_color (Optional[str]): Color of the prefix text, supports terminal color codes, defaults to None.
    color (Optional[str]): Color of the main content text, supports terminal color codes, defaults to None.
    stream (bool): Whether to enable streaming output mode, which enqueues text to the filesystem queue, defaults to False.


Examples:
    >>> from lazyllm.tools.agent.functionCall import StreamResponse
    >>> resp = StreamResponse(prefix="[INFO]", prefix_color="green", color="white", stream=True)
    >>> resp("Hello, world!")
    Hello, world!
    """
    def __init__(self, prefix: str, prefix_color: str = None, color: str = None, stream: bool = False):
        self.stream = stream
        self.prefix = prefix
        self.prefix_color = prefix_color
        self.color = color

    def __call__(self, *inputs):
        if self.stream: FileSystemQueue().enqueue(colored_text(f'\n{self.prefix}\n', self.prefix_color))
        if len(inputs) == 1:
            if self.stream: FileSystemQueue().enqueue(colored_text(f'{inputs[0]}', self.color))
            return inputs[0]
        if self.stream: FileSystemQueue().enqueue(colored_text(f'{inputs}', self.color))
        return package(*inputs)

lazyllm.tools.MCPClient

Bases: object

Source code in lazyllm/tools/mcp/client.py
class MCPClient(object):
    def __init__(
        self,
        command_or_url: str,
        args: Optional[list[str]] = None,
        env: dict[str, str] = None,
        headers: dict[str, Any] = None,
        timeout: float = 5,
    ):
        self._command_or_url = command_or_url
        self._args = args or []
        self._env = env
        self._headers = headers
        self._timeout = timeout

    @asynccontextmanager
    async def _run_session(self):
        if urlparse(self._command_or_url).scheme in ("http", "https"):
            spec = importlib.util.find_spec("mcp.client.sse")
            if spec is None:
                raise ImportError(
                    "Please install mcp to use mcp module. "
                    "You can install it with `pip install mcp`"
                )
            sse_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(sse_module)
            sse_client = sse_module.sse_client

            async with sse_client(
                url=self._command_or_url,
                headers=self._headers,
                timeout=self._timeout
            ) as streams:
                async with mcp.ClientSession(*streams) as session:
                    await session.initialize()
                    yield session
        else:
            server_parameters = mcp.StdioServerParameters(
                command=self._command_or_url, args=self._args, env=self._env
            )
            async with mcp.stdio_client(server_parameters) as streams:
                async with mcp.ClientSession(*streams) as session:
                    await session.initialize()
                    yield session

    async def call_tool(self, tool_name: str, arguments: dict):
        async with self._run_session() as session:
            return await session.call_tool(tool_name, arguments)

    async def list_tools(self):
        async with self._run_session() as session:
            return await session.list_tools()

    async def aget_tools(self, allowed_tools: list[str] = None):
        res = await self.list_tools()
        mcp_tools = getattr(res, "tools", [])
        if allowed_tools:
            mcp_tools = [tool for tool in mcp_tools if tool.name in allowed_tools]

        return [generate_lazyllm_tool(self, tool) for tool in mcp_tools]

    def get_tools(self, allowed_tools: list[str] = None):
        return patch_sync(self.aget_tools)(allowed_tools=allowed_tools)

    async def deploy(self, sse_settings: SseServerSettings):
        async with self._run_session() as session:
            await start_sse_server(session, sse_settings)

lazyllm.tools.tools.GoogleSearch

Bases: HttpTool

Search for specified keywords through Google.

Parameters:

  • custom_search_api_key (str) –

    The Google API key applied by the user.

  • search_engine_id (str) –

    The ID of the search engine created by the user for retrieval.

  • timeout (int, default: 10 ) –

    The timeout for the search request, in seconds, default is 10.

  • proxies (Dict[str, str], default: None ) –

    The proxy services used during the request. Format reference https://www.python-httpx.org/advanced/proxies.

Examples:

from lazyllm.tools.tools import GoogleSearch

key = '<your_google_search_api_key>'
cx = '<your_search_engine_id>'

google = GoogleSearch(custom_search_api_key=key, search_engine_id=cx)
Source code in lazyllm/tools/tools/google_search.py
class GoogleSearch(HttpTool):
    """
Search for specified keywords through Google.

Args:
    custom_search_api_key (str): The Google API key applied by the user.
    search_engine_id (str): The ID of the search engine created by the user for retrieval.
    timeout (int): The timeout for the search request, in seconds, default is 10.
    proxies (Dict[str, str], optional): The proxy services used during the request. Format reference `https://www.python-httpx.org/advanced/proxies`.


Examples:

    from lazyllm.tools.tools import GoogleSearch

    key = '<your_google_search_api_key>'
    cx = '<your_search_engine_id>'

    google = GoogleSearch(custom_search_api_key=key, search_engine_id=cx)
    """
    # @param proxies refer to https://www.python-httpx.org/advanced/proxies
    def __init__(self, custom_search_api_key: str, search_engine_id: str,
                 timeout=10, proxies: Optional[Dict] = None):
        # refer to https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list?hl=zh-cn
        params = {
            'key': custom_search_api_key,
            'cx': '{{search_engine_id}}',
            'q': '{{query}}',
            'dateRestrict': '{{date_restrict}}',
            'start': 0,
            'num': 10,
        }
        super().__init__(method='GET', url='https://customsearch.googleapis.com/customsearch/v1',
                         params=params, timeout=timeout, proxies=proxies)
        self._search_engine_id = search_engine_id

    def forward(self, query: str, date_restrict: str = 'm1',
                search_engine_id: Optional[str] = None) -> Optional[Dict]:
        """
Execute search request.

Args:
    query (str): Keywords to retrieve.
    date_restrict (str): Timeliness of the content to retrieve. Defaults to web pages within one month (m1). Refer to `https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list?hl=zh-cn` for parameter format.
    search_engine_id (str, optional): Search engine ID for retrieval. If this value is empty, the value passed in the constructor is used.


Examples:

    from lazyllm.tools.tools import GoogleSearch

    key = '<your_google_search_api_key>'
    cx = '<your_search_engine_id>'

    google = GoogleSearch(key, cx)
    res = google(query='商汤科技', date_restrict='m1')
    """
        if not search_engine_id:
            search_engine_id = self._search_engine_id

        return super().forward(query=query, search_engine_id=search_engine_id,
                               date_restrict=date_restrict)

forward(query, date_restrict='m1', search_engine_id=None)

Execute search request.

Parameters:

  • query (str) –

    Keywords to retrieve.

  • date_restrict (str, default: 'm1' ) –

    Timeliness of the content to retrieve. Defaults to web pages within one month (m1). Refer to https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list?hl=zh-cn for parameter format.

  • search_engine_id (str, default: None ) –

    Search engine ID for retrieval. If this value is empty, the value passed in the constructor is used.

Examples:

from lazyllm.tools.tools import GoogleSearch

key = '<your_google_search_api_key>'
cx = '<your_search_engine_id>'

google = GoogleSearch(key, cx)
res = google(query='商汤科技', date_restrict='m1')
Source code in lazyllm/tools/tools/google_search.py
    def forward(self, query: str, date_restrict: str = 'm1',
                search_engine_id: Optional[str] = None) -> Optional[Dict]:
        """
Execute search request.

Args:
    query (str): Keywords to retrieve.
    date_restrict (str): Timeliness of the content to retrieve. Defaults to web pages within one month (m1). Refer to `https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list?hl=zh-cn` for parameter format.
    search_engine_id (str, optional): Search engine ID for retrieval. If this value is empty, the value passed in the constructor is used.


Examples:

    from lazyllm.tools.tools import GoogleSearch

    key = '<your_google_search_api_key>'
    cx = '<your_search_engine_id>'

    google = GoogleSearch(key, cx)
    res = google(query='商汤科技', date_restrict='m1')
    """
        if not search_engine_id:
            search_engine_id = self._search_engine_id

        return super().forward(query=query, search_engine_id=search_engine_id,
                               date_restrict=date_restrict)

lazyllm.tools.tools.tencent_search.TencentSearch

Bases: ModuleBase

This is a search enhancement tool.

Examples:

from lazyllm.tools.tools import TencentSearch
secret_id = '<your_secret_id>'
secret_key = '<your_secret_key>'
searcher = TencentSearch(secret_id, secret_key)
Source code in lazyllm/tools/tools/tencent_search.py
class TencentSearch(ModuleBase):
    """
This is a search enhancement tool.


Examples:

    from lazyllm.tools.tools import TencentSearch
    secret_id = '<your_secret_id>'
    secret_key = '<your_secret_key>'
    searcher = TencentSearch(secret_id, secret_key)
    """
    def __init__(self, secret_id, secret_key):
        super().__init__()
        from tencentcloud.common.common_client import CommonClient
        from tencentcloud.common import credential
        from tencentcloud.common.profile.client_profile import ClientProfile
        from tencentcloud.common.profile.http_profile import HttpProfile

        self.cred = credential.Credential(secret_id, secret_key)
        httpProfile = HttpProfile()
        httpProfile.endpoint = "tms.tencentcloudapi.com"
        clientProfile = ClientProfile()
        clientProfile.httpProfile = httpProfile
        self.headers = {"X-TC-Action": "SearchPro"}
        self.common_client = CommonClient(
            "tms", '2020-12-29', self.cred, "", profile=clientProfile)

    def forward(self, query: str):
        """
Searches for the query entered by the user.

Args:
    query (str): The content that the user wants to query.


Examples:

    from lazyllm.tools.tools import TencentSearch
    secret_id = '<your_secret_id>'
    secret_key = '<your_secret_key>'
    searcher = TencentSearch(secret_id, secret_key)
    res = searcher('calculus')
    """
        try:
            res_dict = self.common_client.call_json("SearchPro", {'Query': query, 'Mode': 2}, headers=self.headers)
            res = package(res_dict["Response"]["Pages"])
        except Exception as err:
            lazyllm.LOG.error("Request Tencent Search meets error: ", err)
            res = package()
        return res

forward(query)

Searches for the query entered by the user.

Parameters:

  • query (str) –

    The content that the user wants to query.

Examples:

from lazyllm.tools.tools import TencentSearch
secret_id = '<your_secret_id>'
secret_key = '<your_secret_key>'
searcher = TencentSearch(secret_id, secret_key)
res = searcher('calculus')
Source code in lazyllm/tools/tools/tencent_search.py
    def forward(self, query: str):
        """
Searches for the query entered by the user.

Args:
    query (str): The content that the user wants to query.


Examples:

    from lazyllm.tools.tools import TencentSearch
    secret_id = '<your_secret_id>'
    secret_key = '<your_secret_key>'
    searcher = TencentSearch(secret_id, secret_key)
    res = searcher('calculus')
    """
        try:
            res_dict = self.common_client.call_json("SearchPro", {'Query': query, 'Mode': 2}, headers=self.headers)
            res = package(res_dict["Response"]["Pages"])
        except Exception as err:
            lazyllm.LOG.error("Request Tencent Search meets error: ", err)
            res = package()
        return res

lazyllm.tools.rag.web.WebUi

A Gradio-based web UI for managing knowledge base files.

This class provides an interactive UI to create/delete groups, upload files, list files, and perform deletion operations via RESTful APIs. It is designed for rapid integration of file and group management.

Parameters:

  • base_url (str) –

    Base URL of the backend API service.

Source code in lazyllm/tools/rag/web.py
class WebUi:
    """A Gradio-based web UI for managing knowledge base files.

This class provides an interactive UI to create/delete groups, upload files, list files, and perform deletion operations via RESTful APIs. It is designed for rapid integration of file and group management.

Args:
    base_url (str): Base URL of the backend API service.
"""
    def __init__(self, base_url) -> None:
        self.base_url = base_url

    def basic_headers(self, content_type=True):
        """
Generate standard HTTP headers.

Args:
    content_type (bool): Whether to include Content-Type in the headers (default: True).

Returns:
    dict: Dictionary of HTTP headers.
"""
        return {
            "accept": "application/json",
            "Content-Type": "application/json" if content_type else None,
        }

    def muti_headers(
        self,
    ):
        """
Generate HTTP headers for file upload.

Returns:
    dict: Dictionary of HTTP headers.
"""
        return {"accept": "application/json"}

    def post_request(self, url, data):
        """
Send a POST request.

Args:
    url (str): Target request URL.
    data (dict): Request data (will be serialized as JSON).

Returns:
    dict: JSON response from the server.
"""
        response = requests.post(
            url, headers=self.basic_headers(), data=json.dumps(data)
        )
        return response.json()

    def get_request(self, url):
        """
Send a GET request.

Args:
    url (str): Target request URL.

Returns:
    dict: JSON response from the server.
"""
        response = requests.get(url, headers=self.basic_headers(False))
        return response.json()

    def new_group(self, group_name: str):
        """
Create a new file group.

Args:
    group_name (str): Name of the new group.

Returns:
    str: Server message about the creation result.
"""
        response = requests.post(
            f"{self.base_url}/new_group?group_name={group_name}",
            headers=self.basic_headers(True),
        )
        return response.json()["msg"]

    def delete_group(self, group_name: str):
        """
Delete a specific file group.

Args:
    group_name (str): Name of the group to delete.

Returns:
    str: Server message about the deletion.
"""
        response = requests.post(
            f"{self.base_url}/delete_group?group_name={group_name}",
            headers=self.basic_headers(True),
        )
        return response.json()["msg"]

    def list_groups(self):
        """
List all available file groups.

Returns:
    List[str]: List of group names.
"""
        response = requests.get(
            f"{self.base_url}/list_kb_groups", headers=self.basic_headers(False)
        )
        return response.json()["data"]

    def upload_files(self, group_name: str, override: bool = True):
        """
Upload files to a specified group.

Args:
    group_name (str): Name of the group.
    override (bool): Whether to override existing files (default: True).

Returns:
    Any: Data returned by the backend.
"""
        response = requests.post(
            f"{self.base_url}/upload_files?group_name={group_name}&override={override}",
            headers=self.basic_headers(True),
        )
        return response.json()["data"]

    def list_files_in_group(self, group_name: str):
        """
List all files within a specific group.

Args:
    group_name (str): Name of the group.

Returns:
    List: List of file information.
"""
        response = requests.get(
            f"{self.base_url}/list_files_in_group?group_name={group_name}&alive=True",
            headers=self.basic_headers(False),
        )
        return response.json()["data"]

    def delete_file(self, group_name: str, file_ids: list[str]):
        """
Delete specific files from a group.

Args:
    group_name (str): Name of the group.
    file_ids (List[str]): IDs of files to delete.

Returns:
    str: Deletion result message.
"""
        response = requests.post(
            f"{self.base_url}/delete_files_from_group",
            headers=self.basic_headers(True),
            json={"group_name": group_name, "file_ids": file_ids}
        )
        return response.json()["msg"]

    def gr_show_list(self, str_list: list, list_name: Union[str, list]):
        """
Display a list of strings as a Gradio DataFrame.

Args:
    str_list (List): List of strings or rows.
    list_name (Union[str, List]): Column name(s) for the table.

Returns:
    gr.DataFrame: Gradio DataFrame component.
"""
        if isinstance(list_name, str):
            headers = ["index", list_name]
            value = [[index, str_list[index]] for index in range(len(str_list))]
        else:
            headers = ["index"] + list_name
            value = [[index] + str_list[index:index + len(list_name)] for index in range(len(str_list))]
        return gr.DataFrame(headers=headers, value=value)

    def create_ui(self):
        """
Build a Gradio-based file management UI, including tabs for group listing, file uploading, viewing, and deletion.

Returns:
    gr.Blocks: A complete Gradio application instance.
"""
        with gr.Blocks(analytics_enabled=False) as demo:
            with gr.Tabs():
                select_group_list = []

                with gr.TabItem("分组列表"):
                    select_group = self.gr_show_list(
                        self.list_groups(), list_name="group_name"
                    )
                    select_group_list.append(select_group)

                with gr.TabItem("上传文件"):

                    def _upload_files(group_name, files):

                        files_to_upload = [
                            ("files", (os.path.basename(file), open(file, "rb")))
                            for file in files
                        ]

                        url = f"{self.base_url}/add_files_to_group?group_name={group_name}&override=true"
                        response = requests.post(
                            url, files=files_to_upload, headers=self.muti_headers()
                        )
                        response.raise_for_status()
                        response_data = response.json()
                        gr.Info(str(response_data["msg"]))

                        for _, (_, file_obj) in files_to_upload:
                            file_obj.close()

                    select_group = gr.Dropdown(self.list_groups(), label="选择分组")
                    select_group.change(lambda x: x, inputs=select_group, outputs=None)

                    up_files = gr.Files(label="上传文件")
                    up_btn = gr.Button("上传")
                    up_btn.click(
                        _upload_files,
                        inputs=[select_group, up_files],
                        outputs=None,
                    )

                    select_group_list.append(select_group)

                with gr.TabItem("分组文件列表"):
                    def _list_group_files(group_name):
                        file_list = self.list_files_in_group(group_name)
                        values = [[i] + file_list[i][:2] for i in range(len(file_list))]
                        return gr.update(
                            value=values
                        )

                    select_group = gr.Dropdown(self.list_groups(), label="选择分组")
                    show_list = self.gr_show_list([], list_name=["file_id", "file_name"])
                    select_group.change(
                        fn=_list_group_files, inputs=select_group, outputs=show_list
                    )
                    select_group_list.append(select_group)

                with gr.TabItem("删除文件"):

                    def _list_group_files(group_name):
                        file_list = self.list_files_in_group(group_name)
                        file_list = [','.join(file[:2]) for file in file_list]
                        return gr.update(choices=file_list)

                    select_group = gr.Dropdown(self.list_groups(), label="选择分组")
                    select_file = gr.Dropdown([], label="选择文件")
                    select_group.change(
                        fn=_list_group_files, inputs=select_group, outputs=select_file
                    )
                    delete_btn = gr.Button("删除")

                    def _delete_file(group_name, select_file):
                        file_ids = [select_file.split(',')[0]]
                        gr.Info(self.delete_file(group_name, file_ids))
                        return _list_group_files(group_name)

                    delete_btn.click(
                        fn=_delete_file,
                        inputs=[select_group, select_file],
                        outputs=select_file,
                    )
                    select_group_list.append(select_group)

        return demo

basic_headers(content_type=True)

Generate standard HTTP headers.

Parameters:

  • content_type (bool, default: True ) –

    Whether to include Content-Type in the headers (default: True).

Returns:

  • dict

    Dictionary of HTTP headers.

Source code in lazyllm/tools/rag/web.py
    def basic_headers(self, content_type=True):
        """
Generate standard HTTP headers.

Args:
    content_type (bool): Whether to include Content-Type in the headers (default: True).

Returns:
    dict: Dictionary of HTTP headers.
"""
        return {
            "accept": "application/json",
            "Content-Type": "application/json" if content_type else None,
        }

create_ui()

Build a Gradio-based file management UI, including tabs for group listing, file uploading, viewing, and deletion.

Returns:

  • gr.Blocks: A complete Gradio application instance.

Source code in lazyllm/tools/rag/web.py
    def create_ui(self):
        """
Build a Gradio-based file management UI, including tabs for group listing, file uploading, viewing, and deletion.

Returns:
    gr.Blocks: A complete Gradio application instance.
"""
        with gr.Blocks(analytics_enabled=False) as demo:
            with gr.Tabs():
                select_group_list = []

                with gr.TabItem("分组列表"):
                    select_group = self.gr_show_list(
                        self.list_groups(), list_name="group_name"
                    )
                    select_group_list.append(select_group)

                with gr.TabItem("上传文件"):

                    def _upload_files(group_name, files):

                        files_to_upload = [
                            ("files", (os.path.basename(file), open(file, "rb")))
                            for file in files
                        ]

                        url = f"{self.base_url}/add_files_to_group?group_name={group_name}&override=true"
                        response = requests.post(
                            url, files=files_to_upload, headers=self.muti_headers()
                        )
                        response.raise_for_status()
                        response_data = response.json()
                        gr.Info(str(response_data["msg"]))

                        for _, (_, file_obj) in files_to_upload:
                            file_obj.close()

                    select_group = gr.Dropdown(self.list_groups(), label="选择分组")
                    select_group.change(lambda x: x, inputs=select_group, outputs=None)

                    up_files = gr.Files(label="上传文件")
                    up_btn = gr.Button("上传")
                    up_btn.click(
                        _upload_files,
                        inputs=[select_group, up_files],
                        outputs=None,
                    )

                    select_group_list.append(select_group)

                with gr.TabItem("分组文件列表"):
                    def _list_group_files(group_name):
                        file_list = self.list_files_in_group(group_name)
                        values = [[i] + file_list[i][:2] for i in range(len(file_list))]
                        return gr.update(
                            value=values
                        )

                    select_group = gr.Dropdown(self.list_groups(), label="选择分组")
                    show_list = self.gr_show_list([], list_name=["file_id", "file_name"])
                    select_group.change(
                        fn=_list_group_files, inputs=select_group, outputs=show_list
                    )
                    select_group_list.append(select_group)

                with gr.TabItem("删除文件"):

                    def _list_group_files(group_name):
                        file_list = self.list_files_in_group(group_name)
                        file_list = [','.join(file[:2]) for file in file_list]
                        return gr.update(choices=file_list)

                    select_group = gr.Dropdown(self.list_groups(), label="选择分组")
                    select_file = gr.Dropdown([], label="选择文件")
                    select_group.change(
                        fn=_list_group_files, inputs=select_group, outputs=select_file
                    )
                    delete_btn = gr.Button("删除")

                    def _delete_file(group_name, select_file):
                        file_ids = [select_file.split(',')[0]]
                        gr.Info(self.delete_file(group_name, file_ids))
                        return _list_group_files(group_name)

                    delete_btn.click(
                        fn=_delete_file,
                        inputs=[select_group, select_file],
                        outputs=select_file,
                    )
                    select_group_list.append(select_group)

        return demo

delete_file(group_name, file_ids)

Delete specific files from a group.

Parameters:

  • group_name (str) –

    Name of the group.

  • file_ids (List[str]) –

    IDs of files to delete.

Returns:

  • str

    Deletion result message.

Source code in lazyllm/tools/rag/web.py
    def delete_file(self, group_name: str, file_ids: list[str]):
        """
Delete specific files from a group.

Args:
    group_name (str): Name of the group.
    file_ids (List[str]): IDs of files to delete.

Returns:
    str: Deletion result message.
"""
        response = requests.post(
            f"{self.base_url}/delete_files_from_group",
            headers=self.basic_headers(True),
            json={"group_name": group_name, "file_ids": file_ids}
        )
        return response.json()["msg"]

delete_group(group_name)

Delete a specific file group.

Parameters:

  • group_name (str) –

    Name of the group to delete.

Returns:

  • str

    Server message about the deletion.

Source code in lazyllm/tools/rag/web.py
    def delete_group(self, group_name: str):
        """
Delete a specific file group.

Args:
    group_name (str): Name of the group to delete.

Returns:
    str: Server message about the deletion.
"""
        response = requests.post(
            f"{self.base_url}/delete_group?group_name={group_name}",
            headers=self.basic_headers(True),
        )
        return response.json()["msg"]

get_request(url)

Send a GET request.

Parameters:

  • url (str) –

    Target request URL.

Returns:

  • dict

    JSON response from the server.

Source code in lazyllm/tools/rag/web.py
    def get_request(self, url):
        """
Send a GET request.

Args:
    url (str): Target request URL.

Returns:
    dict: JSON response from the server.
"""
        response = requests.get(url, headers=self.basic_headers(False))
        return response.json()

gr_show_list(str_list, list_name)

Display a list of strings as a Gradio DataFrame.

Parameters:

  • str_list (List) –

    List of strings or rows.

  • list_name (Union[str, List]) –

    Column name(s) for the table.

Returns:

  • gr.DataFrame: Gradio DataFrame component.

Source code in lazyllm/tools/rag/web.py
    def gr_show_list(self, str_list: list, list_name: Union[str, list]):
        """
Display a list of strings as a Gradio DataFrame.

Args:
    str_list (List): List of strings or rows.
    list_name (Union[str, List]): Column name(s) for the table.

Returns:
    gr.DataFrame: Gradio DataFrame component.
"""
        if isinstance(list_name, str):
            headers = ["index", list_name]
            value = [[index, str_list[index]] for index in range(len(str_list))]
        else:
            headers = ["index"] + list_name
            value = [[index] + str_list[index:index + len(list_name)] for index in range(len(str_list))]
        return gr.DataFrame(headers=headers, value=value)

list_files_in_group(group_name)

List all files within a specific group.

Parameters:

  • group_name (str) –

    Name of the group.

Returns:

  • List

    List of file information.

Source code in lazyllm/tools/rag/web.py
    def list_files_in_group(self, group_name: str):
        """
List all files within a specific group.

Args:
    group_name (str): Name of the group.

Returns:
    List: List of file information.
"""
        response = requests.get(
            f"{self.base_url}/list_files_in_group?group_name={group_name}&alive=True",
            headers=self.basic_headers(False),
        )
        return response.json()["data"]

list_groups()

List all available file groups.

Returns:

  • List[str]: List of group names.

Source code in lazyllm/tools/rag/web.py
    def list_groups(self):
        """
List all available file groups.

Returns:
    List[str]: List of group names.
"""
        response = requests.get(
            f"{self.base_url}/list_kb_groups", headers=self.basic_headers(False)
        )
        return response.json()["data"]

muti_headers()

Generate HTTP headers for file upload.

Returns:

  • dict

    Dictionary of HTTP headers.

Source code in lazyllm/tools/rag/web.py
    def muti_headers(
        self,
    ):
        """
Generate HTTP headers for file upload.

Returns:
    dict: Dictionary of HTTP headers.
"""
        return {"accept": "application/json"}

new_group(group_name)

Create a new file group.

Parameters:

  • group_name (str) –

    Name of the new group.

Returns:

  • str

    Server message about the creation result.

Source code in lazyllm/tools/rag/web.py
    def new_group(self, group_name: str):
        """
Create a new file group.

Args:
    group_name (str): Name of the new group.

Returns:
    str: Server message about the creation result.
"""
        response = requests.post(
            f"{self.base_url}/new_group?group_name={group_name}",
            headers=self.basic_headers(True),
        )
        return response.json()["msg"]

post_request(url, data)

Send a POST request.

Parameters:

  • url (str) –

    Target request URL.

  • data (dict) –

    Request data (will be serialized as JSON).

Returns:

  • dict

    JSON response from the server.

Source code in lazyllm/tools/rag/web.py
    def post_request(self, url, data):
        """
Send a POST request.

Args:
    url (str): Target request URL.
    data (dict): Request data (will be serialized as JSON).

Returns:
    dict: JSON response from the server.
"""
        response = requests.post(
            url, headers=self.basic_headers(), data=json.dumps(data)
        )
        return response.json()

upload_files(group_name, override=True)

Upload files to a specified group.

Parameters:

  • group_name (str) –

    Name of the group.

  • override (bool, default: True ) –

    Whether to override existing files (default: True).

Returns:

  • Any

    Data returned by the backend.

Source code in lazyllm/tools/rag/web.py
    def upload_files(self, group_name: str, override: bool = True):
        """
Upload files to a specified group.

Args:
    group_name (str): Name of the group.
    override (bool): Whether to override existing files (default: True).

Returns:
    Any: Data returned by the backend.
"""
        response = requests.post(
            f"{self.base_url}/upload_files?group_name={group_name}&override={override}",
            headers=self.basic_headers(True),
        )
        return response.json()["data"]

lazyllm.tools.http_request.http_executor_response.HttpExecutorResponse

Source code in lazyllm/tools/http_request/http_executor_response.py
class HttpExecutorResponse:
    headers: dict[str, str]
    response: httpx.Response

    def __init__(self, response: httpx.Response = None):
        self.response = response
        self.headers = dict(response.headers) if isinstance(self.response, httpx.Response) else {}

    @property
    def is_file(self) -> bool:
        """
        check if response is file
        """
        content_type = self.get_content_type()
        file_content_types = ['image', 'audio', 'video']

        return any(v in content_type for v in file_content_types)

    def get_content_type(self) -> str:
        """Get the content type of the HTTP response.

Extracts the 'content-type' field value from the response headers to determine the type of response content.

Returns:
    str: The content type of the response, or empty string if not found.


Examples:
    >>> from lazyllm.tools.http_request.http_executor_response import HttpExecutorResponse
    >>> import httpx
    >>> response = httpx.Response(200, headers={'content-type': 'application/json'})
    >>> http_response = HttpExecutorResponse(response)
    >>> content_type = http_response.get_content_type()
    >>> print(content_type)
    ... 'application/json'
    """
        return self.headers.get('content-type', '')

    def extract_file(self) -> tuple[str, bytes]:
        """
        extract file from response if content type is file related
        """
        if self.is_file:
            return self.get_content_type(), self.body

        return '', b''

    @property
    def content(self) -> str:
        if isinstance(self.response, httpx.Response):
            return self.response.text
        else:
            raise ValueError(f'Invalid response type {type(self.response)}')

    @property
    def body(self) -> bytes:
        if isinstance(self.response, httpx.Response):
            return self.response.content
        else:
            raise ValueError(f'Invalid response type {type(self.response)}')

    @property
    def status_code(self) -> int:
        if isinstance(self.response, httpx.Response):
            return self.response.status_code
        else:
            raise ValueError(f'Invalid response type {type(self.response)}')

is_file property

check if response is file

extract_file()

extract file from response if content type is file related

Source code in lazyllm/tools/http_request/http_executor_response.py
def extract_file(self) -> tuple[str, bytes]:
    """
    extract file from response if content type is file related
    """
    if self.is_file:
        return self.get_content_type(), self.body

    return '', b''

get_content_type()

Get the content type of the HTTP response.

Extracts the 'content-type' field value from the response headers to determine the type of response content.

Returns:

  • str ( str ) –

    The content type of the response, or empty string if not found.

Examples:

>>> from lazyllm.tools.http_request.http_executor_response import HttpExecutorResponse
>>> import httpx
>>> response = httpx.Response(200, headers={'content-type': 'application/json'})
>>> http_response = HttpExecutorResponse(response)
>>> content_type = http_response.get_content_type()
>>> print(content_type)
... 'application/json'
Source code in lazyllm/tools/http_request/http_executor_response.py
    def get_content_type(self) -> str:
        """Get the content type of the HTTP response.

Extracts the 'content-type' field value from the response headers to determine the type of response content.

Returns:
    str: The content type of the response, or empty string if not found.


Examples:
    >>> from lazyllm.tools.http_request.http_executor_response import HttpExecutorResponse
    >>> import httpx
    >>> response = httpx.Response(200, headers={'content-type': 'application/json'})
    >>> http_response = HttpExecutorResponse(response)
    >>> content_type = http_response.get_content_type()
    >>> print(content_type)
    ... 'application/json'
    """
        return self.headers.get('content-type', '')