fabricatio 0.2.6.dev6__tar.gz → 0.2.6.dev8__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (123) hide show
  1. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/PKG-INFO +1 -1
  2. fabricatio-0.2.6.dev8/examples/extract_and_inject/.gitignore +2 -0
  3. fabricatio-0.2.6.dev8/examples/extract_and_inject/extract_and_inject.py +59 -0
  4. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/pyproject.toml +1 -1
  5. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/actions/rag.py +4 -3
  6. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/capabilities/rag.py +23 -8
  7. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/config.py +5 -3
  8. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/action.py +12 -0
  9. fabricatio-0.2.6.dev8/python/fabricatio/models/extra.py +655 -0
  10. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/kwargs_types.py +16 -6
  11. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/usages.py +5 -1
  12. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/utils.py +4 -0
  13. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/create_json_obj.hbs +2 -2
  14. fabricatio-0.2.6.dev8/templates.tar.gz +0 -0
  15. fabricatio-0.2.6.dev8/uv.lock +1549 -0
  16. fabricatio-0.2.6.dev6/examples/extract_and_review/.gitignore +0 -1
  17. fabricatio-0.2.6.dev6/examples/extract_and_review/extract_and_review.py +0 -46
  18. fabricatio-0.2.6.dev6/python/fabricatio/models/extra.py +0 -339
  19. fabricatio-0.2.6.dev6/templates.tar.gz +0 -0
  20. fabricatio-0.2.6.dev6/uv.lock +0 -1547
  21. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/.github/workflows/build-package.yaml +0 -0
  22. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/.github/workflows/ruff.yaml +0 -0
  23. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/.github/workflows/tests.yaml +0 -0
  24. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/.gitignore +0 -0
  25. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/.python-version +0 -0
  26. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/Cargo.lock +0 -0
  27. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/Cargo.toml +0 -0
  28. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/LICENSE +0 -0
  29. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/Makefile +0 -0
  30. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/README.md +0 -0
  31. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/correct/correct.py +0 -0
  32. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/correct/correct_loop.py +0 -0
  33. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/extract_article/extract.py +0 -0
  34. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/llm_usages/llm_usage.py +0 -0
  35. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/make_a_rating/rating.py +0 -0
  36. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/make_diary/commits.json +0 -0
  37. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/make_diary/diary.py +0 -0
  38. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/minor/hello_fabricatio.py +0 -0
  39. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/minor/write_a_poem.py +0 -0
  40. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/propose_task/propose.py +0 -0
  41. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/reviewer/review.py +0 -0
  42. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/search_bibtex/.gitignore +0 -0
  43. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/search_bibtex/search.py +0 -0
  44. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/simple_chat/chat.py +0 -0
  45. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/simple_rag/simple_rag.py +0 -0
  46. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/task_handle/handle_task.py +0 -0
  47. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/write_outline/.gitignore +0 -0
  48. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/write_outline/write_outline.py +0 -0
  49. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/examples/write_outline/write_outline_corrected.py +0 -0
  50. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/__init__.py +0 -0
  51. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/_rust.pyi +0 -0
  52. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/_rust_instances.py +0 -0
  53. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/actions/article.py +0 -0
  54. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/actions/output.py +0 -0
  55. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/capabilities/correct.py +0 -0
  56. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/capabilities/propose.py +0 -0
  57. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/capabilities/rating.py +0 -0
  58. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/capabilities/review.py +0 -0
  59. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/capabilities/task.py +0 -0
  60. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/core.py +0 -0
  61. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/decorators.py +0 -0
  62. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/fs/__init__.py +0 -0
  63. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/fs/curd.py +0 -0
  64. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/fs/readers.py +0 -0
  65. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/journal.py +0 -0
  66. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/events.py +0 -0
  67. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/generic.py +0 -0
  68. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/role.py +0 -0
  69. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/task.py +0 -0
  70. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/models/tool.py +0 -0
  71. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/parser.py +0 -0
  72. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/py.typed +0 -0
  73. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/toolboxes/__init__.py +0 -0
  74. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/toolboxes/arithmetic.py +0 -0
  75. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/toolboxes/fs.py +0 -0
  76. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/workflows/articles.py +0 -0
  77. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/python/fabricatio/workflows/rag.py +0 -0
  78. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/src/bib_tools.rs +0 -0
  79. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/src/hash.rs +0 -0
  80. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/src/hbs_helpers.rs +0 -0
  81. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/src/lib.rs +0 -0
  82. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/src/templates.rs +0 -0
  83. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/binary-exploitation-ctf-solver.hbs +0 -0
  84. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/claude-xml.hbs +0 -0
  85. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/clean-up-code.hbs +0 -0
  86. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/co_validation.hbs +0 -0
  87. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/correct.hbs +0 -0
  88. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/cryptography-ctf-solver.hbs +0 -0
  89. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/dependencies.hbs +0 -0
  90. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/document-the-code.hbs +0 -0
  91. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/draft_rating_criteria.hbs +0 -0
  92. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/draft_rating_manual.hbs +0 -0
  93. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/draft_rating_weights_klee.hbs +0 -0
  94. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/draft_tool_usage_code.hbs +0 -0
  95. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/extract_criteria_from_reasons.hbs +0 -0
  96. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/extract_reasons_from_examples.hbs +0 -0
  97. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/find-security-vulnerabilities.hbs +0 -0
  98. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/fix-bugs.hbs +0 -0
  99. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/generic_string.hbs +0 -0
  100. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/improve-performance.hbs +0 -0
  101. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/liststr.hbs +0 -0
  102. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/make_choice.hbs +0 -0
  103. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/make_judgment.hbs +0 -0
  104. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/pathstr.hbs +0 -0
  105. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/rate_fine_grind.hbs +0 -0
  106. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/refactor.hbs +0 -0
  107. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/refined_query.hbs +0 -0
  108. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/retrieved_display.hbs +0 -0
  109. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/reverse-engineering-ctf-solver.hbs +0 -0
  110. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/review_string.hbs +0 -0
  111. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/task_briefing.hbs +0 -0
  112. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/web-ctf-solver.hbs +0 -0
  113. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/write-git-commit.hbs +0 -0
  114. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/write-github-pull-request.hbs +0 -0
  115. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/templates/built-in/write-github-readme.hbs +0 -0
  116. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_config.py +0 -0
  117. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_models/test_action.py +0 -0
  118. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_models/test_advanced.py +0 -0
  119. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_models/test_generic.py +0 -0
  120. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_models/test_role.py +0 -0
  121. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_models/test_task.py +0 -0
  122. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_models/test_tool.py +0 -0
  123. {fabricatio-0.2.6.dev6 → fabricatio-0.2.6.dev8}/tests/test_models/test_usages.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fabricatio
3
- Version: 0.2.6.dev6
3
+ Version: 0.2.6.dev8
4
4
  Classifier: License :: OSI Approved :: MIT License
5
5
  Classifier: Programming Language :: Rust
6
6
  Classifier: Programming Language :: Python :: 3.12
@@ -0,0 +1,2 @@
1
+ bpdf_out
2
+ output*
@@ -0,0 +1,59 @@
1
+ """Example of proposing a task to a role."""
2
+
3
+ import asyncio
4
+ from pathlib import Path
5
+ from typing import List, Optional
6
+
7
+ from fabricatio import Action, Event, Role, Task, WorkFlow, logger
8
+ from fabricatio.actions.article import ExtractArticleEssence
9
+ from fabricatio.actions.rag import InjectToDB
10
+ from fabricatio.fs.curd import dump_text, gather_files
11
+ from fabricatio.models.extra import ArticleEssence
12
+
13
+
14
+ class SaveToFS(Action):
15
+ """Save to file system."""
16
+
17
+ async def _execute(self, to_inject: List[Optional[ArticleEssence]], output_dir: Path, **_) -> None:
18
+ output_dir.mkdir(parents=True, exist_ok=True)
19
+ empty_count = 0
20
+ for i, e in enumerate(to_inject):
21
+ if e is None:
22
+ empty_count += 1
23
+ logger.error(f"Invalid essence at index {i}")
24
+ continue
25
+ dump_text(output_dir / f"{i}.json", e.display())
26
+ logger.info(f"Passed {empty_count} empty essence ")
27
+
28
+
29
+ async def main() -> None:
30
+ """Main function."""
31
+ role = Role(
32
+ name="Researcher",
33
+ description="Extract article essence",
34
+ llm_model="openai/deepseek-r1-distill-llama-70b",
35
+ llm_rpm=50,
36
+ llm_tpm=100000,
37
+ registry={
38
+ Event.quick_instantiate("article"): WorkFlow(
39
+ name="extract",
40
+ steps=(ExtractArticleEssence(output_key="to_inject"), SaveToFS, InjectToDB(output_key="task_output")),
41
+ extra_init_context={"collection_name": "article_essence", "output_dir": Path("output")},
42
+ )
43
+ },
44
+ )
45
+
46
+ task: Task[str] = await role.propose_task(
47
+ "Extract the essence of the article from the files in './bpdf_out'",
48
+ )
49
+
50
+ col_name = await task.override_dependencies(gather_files("bpdf_out", "md")).delegate("article")
51
+
52
+ if col_name is None:
53
+ logger.error("No essence found")
54
+ return
55
+ logger.success(f"Injected to collection: {col_name}")
56
+
57
+
58
+ if __name__ == "__main__":
59
+ asyncio.run(main())
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "fabricatio"
3
- version = "0.2.6-dev6"
3
+ version = "0.2.6-dev8"
4
4
  description = "A LLM multi-agent framework."
5
5
  readme = "README.md"
6
6
  license = { file = "LICENSE" }
@@ -3,6 +3,7 @@
3
3
  from typing import List, Optional
4
4
 
5
5
  from fabricatio.capabilities.rag import RAG
6
+ from fabricatio.journal import logger
6
7
  from fabricatio.models.action import Action
7
8
  from fabricatio.models.generic import PrepareVectorization
8
9
 
@@ -13,13 +14,13 @@ class InjectToDB(Action, RAG):
13
14
  output_key: str = "collection_name"
14
15
 
15
16
  async def _execute[T: PrepareVectorization](
16
- self, to_inject: T | List[T], collection_name: Optional[str] = "my_collection", **_
17
+ self, to_inject: Optional[T] | List[Optional[T]], collection_name: Optional[str] = "my_collection", **_
17
18
  ) -> Optional[str]:
18
19
  if not isinstance(to_inject, list):
19
20
  to_inject = [to_inject]
20
-
21
+ logger.info(f"Injecting {len(to_inject)} items into the collection '{collection_name}'")
21
22
  await self.view(collection_name, create=True).consume_string(
22
- [t.prepare_vectorization(self.embedding_max_sequence_length) for t in to_inject],
23
+ [t.prepare_vectorization(self.embedding_max_sequence_length) for t in to_inject if isinstance(t,PrepareVectorization)],
23
24
  )
24
25
 
25
26
  return collection_name
@@ -15,7 +15,7 @@ from fabricatio.config import configs
15
15
  from fabricatio.journal import logger
16
16
  from fabricatio.models.kwargs_types import (
17
17
  ChooseKwargs,
18
- CollectionSimpleConfigKwargs,
18
+ CollectionConfigKwargs,
19
19
  EmbeddingKwargs,
20
20
  FetchKwargs,
21
21
  LLMKwargs,
@@ -63,10 +63,18 @@ class RAG(EmbeddingUsage):
63
63
  uri=milvus_uri or ok(self.milvus_uri or configs.rag.milvus_uri).unicode_string(),
64
64
  token=milvus_token
65
65
  or (token.get_secret_value() if (token := (self.milvus_token or configs.rag.milvus_token)) else ""),
66
- timeout=milvus_timeout or self.milvus_timeout,
66
+ timeout=milvus_timeout or self.milvus_timeout or configs.rag.milvus_timeout,
67
67
  )
68
68
  return self
69
69
 
70
+ def check_client(self, init: bool = True) -> Self:
71
+ """Check if the client is initialized, and if not, initialize it."""
72
+ if self._client is None and init:
73
+ return self.init_client()
74
+ if self._client is None and not init:
75
+ raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
76
+ return self
77
+
70
78
  @overload
71
79
  async def pack(
72
80
  self, input_text: List[str], subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
@@ -102,17 +110,24 @@ class RAG(EmbeddingUsage):
102
110
  ]
103
111
 
104
112
  def view(
105
- self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionSimpleConfigKwargs]
113
+ self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionConfigKwargs]
106
114
  ) -> Self:
107
115
  """View the specified collection.
108
116
 
109
117
  Args:
110
118
  collection_name (str): The name of the collection.
111
119
  create (bool): Whether to create the collection if it does not exist.
112
- **kwargs (Unpack[CollectionSimpleConfigKwargs]): Additional keyword arguments for collection configuration.
120
+ **kwargs (Unpack[CollectionConfigKwargs]): Additional keyword arguments for collection configuration.
113
121
  """
114
- if create and collection_name and self.client.has_collection(collection_name):
115
- kwargs["dimension"] = kwargs.get("dimension") or self.milvus_dimensions or configs.rag.milvus_dimensions
122
+ if create and collection_name and not self.check_client().client.has_collection(collection_name):
123
+ kwargs["dimension"] = ok(
124
+ kwargs.get("dimension")
125
+ or self.milvus_dimensions
126
+ or configs.rag.milvus_dimensions
127
+ or self.embedding_dimensions
128
+ or configs.embedding.dimensions,
129
+ "`dimension` is not set at any level.",
130
+ )
116
131
  self.client.create_collection(collection_name, auto_id=True, **kwargs)
117
132
  logger.info(f"Creating collection {collection_name}")
118
133
 
@@ -158,7 +173,7 @@ class RAG(EmbeddingUsage):
158
173
  else:
159
174
  raise TypeError(f"Expected MilvusData or list of MilvusData, got {type(data)}")
160
175
  c_name = collection_name or self.safe_target_collection
161
- self.client.insert(c_name, prepared_data)
176
+ self.check_client().client.insert(c_name, prepared_data)
162
177
 
163
178
  if flush:
164
179
  logger.debug(f"Flushing collection {c_name}")
@@ -219,7 +234,7 @@ class RAG(EmbeddingUsage):
219
234
  List[Dict[str, Any]] | List[Any]: The retrieved data.
220
235
  """
221
236
  # Step 1: Search for vectors
222
- search_results = self.client.search(
237
+ search_results = self.check_client().client.search(
223
238
  collection_name or self.safe_target_collection,
224
239
  vecs,
225
240
  search_params={"radius": similarity_threshold},
@@ -277,7 +277,7 @@ class RagConfig(BaseModel):
277
277
 
278
278
  milvus_uri: Optional[HttpUrl] = Field(default=HttpUrl("http://localhost:19530"))
279
279
  """The URI of the Milvus server."""
280
- milvus_timeout: Optional[PositiveFloat] = Field(default=None)
280
+ milvus_timeout: Optional[PositiveFloat] = Field(default=30.0)
281
281
  """The timeout of the Milvus server."""
282
282
  milvus_token: Optional[SecretStr] = Field(default=None)
283
283
  """The token of the Milvus server."""
@@ -303,11 +303,13 @@ class RoutingConfig(BaseModel):
303
303
 
304
304
  model_config = ConfigDict(use_attribute_docstrings=True)
305
305
 
306
- allowed_fails: Optional[int] = 1
306
+ max_parallel_requests: Optional[int] = 60
307
+ """The maximum number of parallel requests. None means not checked."""
308
+ allowed_fails: Optional[int] = 3
307
309
  """The number of allowed fails before the routing is considered failed."""
308
310
  retry_after: int = 15
309
311
  """The time in seconds to wait before retrying the routing after a fail."""
310
- cooldown_time: Optional[int] = 120
312
+ cooldown_time: Optional[int] = 30
311
313
  """The time in seconds to wait before retrying the routing after a cooldown."""
312
314
 
313
315
 
@@ -216,3 +216,15 @@ class WorkFlow(WithBriefing, ToolBoxUsage):
216
216
  """
217
217
  self.provide_tools_to(self._instances)
218
218
  return self
219
+
220
+ def update_init_context(self, **kwargs) -> Self:
221
+ """Update the initial context with additional key-value pairs.
222
+
223
+ Args:
224
+ **kwargs: Key-value pairs to add to the initial context.
225
+
226
+ Returns:
227
+ Self: The workflow instance for method chaining.
228
+ """
229
+ self.extra_init_context.update(kwargs)
230
+ return self