fabricatio 0.3.14.dev1__cp312-cp312-win_amd64.whl → 0.3.14.dev4__cp312-cp312-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- fabricatio/__init__.py +5 -6
- fabricatio/actions/article.py +31 -31
- fabricatio/actions/article_rag.py +58 -58
- fabricatio/actions/output.py +58 -24
- fabricatio/actions/rag.py +2 -3
- fabricatio/capabilities/advanced_judge.py +4 -7
- fabricatio/capabilities/advanced_rag.py +2 -1
- fabricatio/capabilities/censor.py +5 -4
- fabricatio/capabilities/check.py +27 -27
- fabricatio/capabilities/correct.py +22 -22
- fabricatio/capabilities/extract.py +33 -33
- fabricatio/capabilities/persist.py +103 -0
- fabricatio/capabilities/propose.py +2 -2
- fabricatio/capabilities/rag.py +37 -37
- fabricatio/capabilities/rating.py +66 -70
- fabricatio/capabilities/review.py +12 -11
- fabricatio/capabilities/task.py +19 -18
- fabricatio/decorators.py +9 -9
- fabricatio/{core.py → emitter.py} +17 -19
- fabricatio/journal.py +2 -1
- fabricatio/models/action.py +9 -11
- fabricatio/models/extra/aricle_rag.py +15 -12
- fabricatio/models/extra/article_base.py +4 -5
- fabricatio/models/extra/article_essence.py +2 -1
- fabricatio/models/extra/article_main.py +5 -4
- fabricatio/models/extra/article_outline.py +2 -1
- fabricatio/models/extra/article_proposal.py +1 -1
- fabricatio/models/extra/rag.py +2 -2
- fabricatio/models/extra/rule.py +2 -1
- fabricatio/models/generic.py +48 -131
- fabricatio/models/kwargs_types.py +1 -9
- fabricatio/models/role.py +14 -13
- fabricatio/models/task.py +3 -4
- fabricatio/models/tool.py +5 -6
- fabricatio/models/usages.py +137 -147
- fabricatio/parser.py +59 -99
- fabricatio/rust.cp312-win_amd64.pyd +0 -0
- fabricatio/rust.pyi +39 -59
- fabricatio/utils.py +6 -170
- fabricatio-0.3.14.dev4.data/scripts/tdown.exe +0 -0
- {fabricatio-0.3.14.dev1.data → fabricatio-0.3.14.dev4.data}/scripts/ttm.exe +0 -0
- {fabricatio-0.3.14.dev1.dist-info → fabricatio-0.3.14.dev4.dist-info}/METADATA +3 -7
- fabricatio-0.3.14.dev4.dist-info/RECORD +64 -0
- fabricatio-0.3.14.dev1.data/scripts/tdown.exe +0 -0
- fabricatio-0.3.14.dev1.dist-info/RECORD +0 -63
- {fabricatio-0.3.14.dev1.dist-info → fabricatio-0.3.14.dev4.dist-info}/WHEEL +0 -0
- {fabricatio-0.3.14.dev1.dist-info → fabricatio-0.3.14.dev4.dist-info}/licenses/LICENSE +0 -0
fabricatio/capabilities/check.py
CHANGED
@@ -1,10 +1,9 @@
|
|
1
1
|
"""A class that provides the capability to check strings and objects against rules and guidelines."""
|
2
2
|
|
3
|
+
from abc import ABC
|
3
4
|
from asyncio import gather
|
4
5
|
from typing import List, Optional, Unpack
|
5
6
|
|
6
|
-
from fabricatio.rust import CONFIG, TEMPLATE_MANAGER, detect_language
|
7
|
-
|
8
7
|
from fabricatio.capabilities.advanced_judge import AdvancedJudge
|
9
8
|
from fabricatio.capabilities.propose import Propose
|
10
9
|
from fabricatio.journal import logger
|
@@ -13,10 +12,11 @@ from fabricatio.models.extra.problem import Improvement
|
|
13
12
|
from fabricatio.models.extra.rule import Rule, RuleSet
|
14
13
|
from fabricatio.models.generic import Display, WithBriefing
|
15
14
|
from fabricatio.models.kwargs_types import ValidateKwargs
|
15
|
+
from fabricatio.rust import CONFIG, TEMPLATE_MANAGER, detect_language
|
16
16
|
from fabricatio.utils import override_kwargs
|
17
17
|
|
18
18
|
|
19
|
-
class Check(AdvancedJudge, Propose):
|
19
|
+
class Check(AdvancedJudge, Propose, ABC):
|
20
20
|
"""Class for validating strings/objects against predefined rules and guidelines.
|
21
21
|
|
22
22
|
This capability combines rule-based judgment and proposal generation to provide
|
@@ -24,7 +24,7 @@ class Check(AdvancedJudge, Propose):
|
|
24
24
|
"""
|
25
25
|
|
26
26
|
async def draft_ruleset(
|
27
|
-
|
27
|
+
self, ruleset_requirement: str, rule_count: int = 0, **kwargs: Unpack[ValidateKwargs[Rule]]
|
28
28
|
) -> Optional[RuleSet]:
|
29
29
|
"""Generate rule set based on requirement description.
|
30
30
|
|
@@ -80,11 +80,11 @@ class Check(AdvancedJudge, Propose):
|
|
80
80
|
return RuleSet(rules=rules, **ruleset_patch.as_kwargs())
|
81
81
|
|
82
82
|
async def check_string_against_rule(
|
83
|
-
|
84
|
-
|
85
|
-
|
86
|
-
|
87
|
-
|
83
|
+
self,
|
84
|
+
input_text: str,
|
85
|
+
rule: Rule,
|
86
|
+
reference: str = "",
|
87
|
+
**kwargs: Unpack[ValidateKwargs[Improvement]],
|
88
88
|
) -> Optional[Improvement]:
|
89
89
|
"""Validate text against specific rule.
|
90
90
|
|
@@ -103,9 +103,9 @@ class Check(AdvancedJudge, Propose):
|
|
103
103
|
- Proposes Improvement only when violation is confirmed
|
104
104
|
"""
|
105
105
|
if judge := await self.evidently_judge(
|
106
|
-
|
107
|
-
|
108
|
-
|
106
|
+
f"# Content to exam\n{input_text}\n\n# Rule Must to follow\n{rule.display()}\nDoes `Content to exam` provided above violate the `{rule.name}` provided above?"
|
107
|
+
f"should I take some measure to fix that violation? true for I do need, false for I don't need.",
|
108
|
+
**override_kwargs(kwargs, default=None),
|
109
109
|
):
|
110
110
|
logger.info(f"Rule `{rule.name}` violated: \n{judge.display()}")
|
111
111
|
return await self.propose(
|
@@ -119,11 +119,11 @@ class Check(AdvancedJudge, Propose):
|
|
119
119
|
return None
|
120
120
|
|
121
121
|
async def check_obj_against_rule[M: (Display, WithBriefing)](
|
122
|
-
|
123
|
-
|
124
|
-
|
125
|
-
|
126
|
-
|
122
|
+
self,
|
123
|
+
obj: M,
|
124
|
+
rule: Rule,
|
125
|
+
reference: str = "",
|
126
|
+
**kwargs: Unpack[ValidateKwargs[Improvement]],
|
127
127
|
) -> Optional[Improvement]:
|
128
128
|
"""Validate object against rule using text representation.
|
129
129
|
|
@@ -151,11 +151,11 @@ class Check(AdvancedJudge, Propose):
|
|
151
151
|
return await self.check_string_against_rule(input_text, rule, reference, **kwargs)
|
152
152
|
|
153
153
|
async def check_string(
|
154
|
-
|
155
|
-
|
156
|
-
|
157
|
-
|
158
|
-
|
154
|
+
self,
|
155
|
+
input_text: str,
|
156
|
+
ruleset: RuleSet,
|
157
|
+
reference: str = "",
|
158
|
+
**kwargs: Unpack[ValidateKwargs[Improvement]],
|
159
159
|
) -> Optional[List[Improvement]]:
|
160
160
|
"""Validate text against full ruleset.
|
161
161
|
|
@@ -182,11 +182,11 @@ class Check(AdvancedJudge, Propose):
|
|
182
182
|
return [imp for imp in imp_seq if imp]
|
183
183
|
|
184
184
|
async def check_obj[M: (Display, WithBriefing)](
|
185
|
-
|
186
|
-
|
187
|
-
|
188
|
-
|
189
|
-
|
185
|
+
self,
|
186
|
+
obj: M,
|
187
|
+
ruleset: RuleSet,
|
188
|
+
reference: str = "",
|
189
|
+
**kwargs: Unpack[ValidateKwargs[Improvement]],
|
190
190
|
) -> Optional[List[Improvement]]:
|
191
191
|
"""Validate object against full ruleset.
|
192
192
|
|
@@ -1,9 +1,8 @@
|
|
1
1
|
"""A module containing the Correct capability for reviewing, validating, and improving objects."""
|
2
2
|
|
3
|
+
from abc import ABC
|
3
4
|
from asyncio import gather
|
4
|
-
from typing import Optional, Unpack, cast
|
5
|
-
|
6
|
-
from fabricatio.rust import CONFIG, TEMPLATE_MANAGER
|
5
|
+
from typing import Optional, Type, Unpack, cast
|
7
6
|
|
8
7
|
from fabricatio.capabilities.propose import Propose
|
9
8
|
from fabricatio.capabilities.rating import Rating
|
@@ -15,14 +14,15 @@ from fabricatio.models.kwargs_types import (
|
|
15
14
|
BestKwargs,
|
16
15
|
ValidateKwargs,
|
17
16
|
)
|
17
|
+
from fabricatio.rust import CONFIG, TEMPLATE_MANAGER
|
18
18
|
from fabricatio.utils import fallback_kwargs, ok, override_kwargs
|
19
19
|
|
20
20
|
|
21
|
-
class Correct(Rating, Propose):
|
21
|
+
class Correct(Rating, Propose, ABC):
|
22
22
|
"""A class that provides the capability to correct objects."""
|
23
23
|
|
24
24
|
async def decide_solution(
|
25
|
-
|
25
|
+
self, problem_solutions: ProblemSolutions, **kwargs: Unpack[BestKwargs]
|
26
26
|
) -> ProblemSolutions:
|
27
27
|
"""Decide the best solution from a list of problem solutions.
|
28
28
|
|
@@ -72,11 +72,11 @@ class Correct(Rating, Propose):
|
|
72
72
|
return improvement
|
73
73
|
|
74
74
|
async def fix_troubled_obj[M: SketchedAble](
|
75
|
-
|
76
|
-
|
77
|
-
|
78
|
-
|
79
|
-
|
75
|
+
self,
|
76
|
+
obj: M,
|
77
|
+
problem_solutions: ProblemSolutions,
|
78
|
+
reference: str = "",
|
79
|
+
**kwargs: Unpack[ValidateKwargs[M]],
|
80
80
|
) -> Optional[M]:
|
81
81
|
"""Fix a troubled object based on problem solutions.
|
82
82
|
|
@@ -106,11 +106,11 @@ class Correct(Rating, Propose):
|
|
106
106
|
)
|
107
107
|
|
108
108
|
async def fix_troubled_string(
|
109
|
-
|
110
|
-
|
111
|
-
|
112
|
-
|
113
|
-
|
109
|
+
self,
|
110
|
+
input_text: str,
|
111
|
+
problem_solutions: ProblemSolutions,
|
112
|
+
reference: str = "",
|
113
|
+
**kwargs: Unpack[ValidateKwargs[str]],
|
114
114
|
) -> Optional[str]:
|
115
115
|
"""Fix a troubled string based on problem solutions.
|
116
116
|
|
@@ -140,11 +140,11 @@ class Correct(Rating, Propose):
|
|
140
140
|
)
|
141
141
|
|
142
142
|
async def correct_obj[M: SketchedAble](
|
143
|
-
|
144
|
-
|
145
|
-
|
146
|
-
|
147
|
-
|
143
|
+
self,
|
144
|
+
obj: M,
|
145
|
+
improvement: Improvement,
|
146
|
+
reference: str = "",
|
147
|
+
**kwargs: Unpack[ValidateKwargs[M]],
|
148
148
|
) -> Optional[M]:
|
149
149
|
"""Review and correct an object based on defined criteria and templates.
|
150
150
|
|
@@ -178,7 +178,7 @@ class Correct(Rating, Propose):
|
|
178
178
|
return obj
|
179
179
|
|
180
180
|
async def correct_string(
|
181
|
-
|
181
|
+
self, input_text: str, improvement: Improvement, reference: str = "", **kwargs: Unpack[ValidateKwargs[str]]
|
182
182
|
) -> Optional[str]:
|
183
183
|
"""Review and correct a string based on defined criteria and templates.
|
184
184
|
|
@@ -210,7 +210,7 @@ class Correct(Rating, Propose):
|
|
210
210
|
return input_text
|
211
211
|
|
212
212
|
async def correct_obj_inplace[M: ProposedUpdateAble](
|
213
|
-
|
213
|
+
self, obj: M, **kwargs: Unpack[CorrectKwargs[M]]
|
214
214
|
) -> Optional[M]:
|
215
215
|
"""Correct an object in place based on defined criteria and templates.
|
216
216
|
|
@@ -1,65 +1,65 @@
|
|
1
1
|
"""A module that provide capabilities for extracting information from a given source to a model."""
|
2
2
|
|
3
|
+
from abc import ABC
|
3
4
|
from typing import List, Optional, Type, Unpack, overload
|
4
5
|
|
5
|
-
from fabricatio.rust import CONFIG
|
6
|
-
|
7
6
|
from fabricatio import TEMPLATE_MANAGER
|
8
7
|
from fabricatio.capabilities.propose import Propose
|
9
8
|
from fabricatio.models.generic import ProposedAble
|
10
9
|
from fabricatio.models.kwargs_types import ValidateKwargs
|
10
|
+
from fabricatio.rust import CONFIG
|
11
11
|
|
12
12
|
|
13
|
-
class Extract(Propose):
|
13
|
+
class Extract(Propose, ABC):
|
14
14
|
"""A class that extract information from a given source to a model."""
|
15
15
|
|
16
16
|
@overload
|
17
17
|
async def extract[M: ProposedAble](
|
18
|
-
|
19
|
-
|
20
|
-
|
21
|
-
|
22
|
-
|
23
|
-
|
18
|
+
self,
|
19
|
+
cls: Type[M],
|
20
|
+
source: str,
|
21
|
+
extract_requirement: Optional[str] = None,
|
22
|
+
align_language: bool = True,
|
23
|
+
**kwargs: Unpack[ValidateKwargs[M]],
|
24
24
|
) -> M: ...
|
25
25
|
|
26
26
|
@overload
|
27
27
|
async def extract[M: ProposedAble](
|
28
|
-
|
29
|
-
|
30
|
-
|
31
|
-
|
32
|
-
|
33
|
-
|
28
|
+
self,
|
29
|
+
cls: Type[M],
|
30
|
+
source: str,
|
31
|
+
extract_requirement: Optional[str] = None,
|
32
|
+
align_language: bool = True,
|
33
|
+
**kwargs: Unpack[ValidateKwargs[None]],
|
34
34
|
) -> Optional[M]: ...
|
35
35
|
|
36
36
|
@overload
|
37
37
|
async def extract[M: ProposedAble](
|
38
|
-
|
39
|
-
|
40
|
-
|
41
|
-
|
42
|
-
|
43
|
-
|
38
|
+
self,
|
39
|
+
cls: Type[M],
|
40
|
+
source: List[str],
|
41
|
+
extract_requirement: Optional[str] = None,
|
42
|
+
align_language: bool = True,
|
43
|
+
**kwargs: Unpack[ValidateKwargs[M]],
|
44
44
|
) -> List[M]: ...
|
45
45
|
|
46
46
|
@overload
|
47
47
|
async def extract[M: ProposedAble](
|
48
|
-
|
49
|
-
|
50
|
-
|
51
|
-
|
52
|
-
|
53
|
-
|
48
|
+
self,
|
49
|
+
cls: Type[M],
|
50
|
+
source: List[str],
|
51
|
+
extract_requirement: Optional[str] = None,
|
52
|
+
align_language: bool = True,
|
53
|
+
**kwargs: Unpack[ValidateKwargs[None]],
|
54
54
|
) -> List[Optional[M]]: ...
|
55
55
|
|
56
56
|
async def extract[M: ProposedAble](
|
57
|
-
|
58
|
-
|
59
|
-
|
60
|
-
|
61
|
-
|
62
|
-
|
57
|
+
self,
|
58
|
+
cls: Type[M],
|
59
|
+
source: List[str] | str,
|
60
|
+
extract_requirement: Optional[str] = None,
|
61
|
+
align_language: bool = True,
|
62
|
+
**kwargs: Unpack[ValidateKwargs[Optional[M]]],
|
63
63
|
) -> M | List[M] | Optional[M] | List[Optional[M]]:
|
64
64
|
"""Extract information from a given source to a model."""
|
65
65
|
return await self.propose(
|
@@ -0,0 +1,103 @@
|
|
1
|
+
"""Persistence capabilities for model instances."""
|
2
|
+
|
3
|
+
from abc import ABC
|
4
|
+
from datetime import datetime
|
5
|
+
from pathlib import Path
|
6
|
+
from typing import Optional, Self
|
7
|
+
|
8
|
+
from loguru import logger
|
9
|
+
|
10
|
+
from fabricatio.fs import safe_text_read
|
11
|
+
from fabricatio.models.generic import Base
|
12
|
+
from fabricatio.rust import blake3_hash
|
13
|
+
|
14
|
+
|
15
|
+
class PersistentAble(Base, ABC):
|
16
|
+
"""Class providing file persistence capabilities.
|
17
|
+
|
18
|
+
Enables saving model instances to disk with timestamped filenames and loading from persisted files.
|
19
|
+
Implements basic versioning through filename hashing and timestamping.
|
20
|
+
"""
|
21
|
+
|
22
|
+
def persist(self, path: str | Path) -> Self:
|
23
|
+
"""Save model instance to disk with versioned filename.
|
24
|
+
|
25
|
+
Args:
|
26
|
+
path (str | Path): Target directory or file path. If directory, filename is auto-generated.
|
27
|
+
|
28
|
+
Returns:
|
29
|
+
Self: Current instance for method chaining
|
30
|
+
|
31
|
+
Notes:
|
32
|
+
- Filename format: <ClassName>_<YYYYMMDD_HHMMSS>_<6-char_hash>.json
|
33
|
+
- Hash generated from JSON content ensures uniqueness
|
34
|
+
"""
|
35
|
+
p = Path(path)
|
36
|
+
out = self.model_dump_json(indent=1, by_alias=True)
|
37
|
+
|
38
|
+
# Generate a timestamp in the format YYYYMMDD_HHMMSS
|
39
|
+
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
40
|
+
|
41
|
+
# Generate the hash
|
42
|
+
file_hash = blake3_hash(out.encode())[:6]
|
43
|
+
|
44
|
+
# Construct the file name with timestamp and hash
|
45
|
+
file_name = f"{self.__class__.__name__}_{timestamp}_{file_hash}.json"
|
46
|
+
|
47
|
+
if p.is_dir():
|
48
|
+
p.joinpath(file_name).write_text(out, encoding="utf-8")
|
49
|
+
else:
|
50
|
+
p.mkdir(exist_ok=True, parents=True)
|
51
|
+
p.write_text(out, encoding="utf-8")
|
52
|
+
|
53
|
+
logger.info(f"Persisted `{self.__class__.__name__}` to {p.as_posix()}")
|
54
|
+
return self
|
55
|
+
|
56
|
+
@classmethod
|
57
|
+
def from_latest_persistent(cls, dir_path: str | Path) -> Optional[Self]:
|
58
|
+
"""Load most recent persisted instance from directory.
|
59
|
+
|
60
|
+
Args:
|
61
|
+
dir_path (str | Path): Directory containing persisted files
|
62
|
+
|
63
|
+
Returns:
|
64
|
+
Self: Most recently modified instance
|
65
|
+
|
66
|
+
Raises:
|
67
|
+
NotADirectoryError: If path is not a valid directory
|
68
|
+
FileNotFoundError: If no matching files found
|
69
|
+
"""
|
70
|
+
dir_path = Path(dir_path)
|
71
|
+
if not dir_path.is_dir():
|
72
|
+
return None
|
73
|
+
|
74
|
+
pattern = f"{cls.__name__}_*.json"
|
75
|
+
files = list(dir_path.glob(pattern))
|
76
|
+
|
77
|
+
if not files:
|
78
|
+
return None
|
79
|
+
|
80
|
+
def _get_timestamp(file_path: Path) -> datetime:
|
81
|
+
stem = file_path.stem
|
82
|
+
parts = stem.split("_")
|
83
|
+
return datetime.strptime(f"{parts[1]}_{parts[2]}", "%Y%m%d_%H%M%S")
|
84
|
+
|
85
|
+
files.sort(key=lambda f: _get_timestamp(f), reverse=True)
|
86
|
+
|
87
|
+
return cls.from_persistent(files.pop(0))
|
88
|
+
|
89
|
+
@classmethod
|
90
|
+
def from_persistent(cls, path: str | Path) -> Self:
|
91
|
+
"""Load an instance from a specific persisted file.
|
92
|
+
|
93
|
+
Args:
|
94
|
+
path (str | Path): Path to the JSON file.
|
95
|
+
|
96
|
+
Returns:
|
97
|
+
Self: The loaded instance from the file.
|
98
|
+
|
99
|
+
Raises:
|
100
|
+
FileNotFoundError: If the specified file does not exist.
|
101
|
+
ValueError: If the file content is invalid for the model.
|
102
|
+
"""
|
103
|
+
return cls.model_validate_json(safe_text_read(path))
|
@@ -1,5 +1,5 @@
|
|
1
1
|
"""A module for the task capabilities of the Fabricatio library."""
|
2
|
-
|
2
|
+
from abc import ABC
|
3
3
|
from typing import List, Optional, Type, Unpack, overload
|
4
4
|
|
5
5
|
from fabricatio.models.generic import ProposedAble
|
@@ -7,7 +7,7 @@ from fabricatio.models.kwargs_types import ValidateKwargs
|
|
7
7
|
from fabricatio.models.usages import LLMUsage
|
8
8
|
|
9
9
|
|
10
|
-
class Propose(LLMUsage):
|
10
|
+
class Propose(LLMUsage,ABC):
|
11
11
|
"""A class that proposes an Obj based on a prompt."""
|
12
12
|
|
13
13
|
@overload
|
fabricatio/capabilities/rag.py
CHANGED
@@ -1,5 +1,7 @@
|
|
1
1
|
"""A module for the RAG (Retrieval Augmented Generation) model."""
|
2
2
|
|
3
|
+
from abc import ABC
|
4
|
+
|
3
5
|
try:
|
4
6
|
from pymilvus import MilvusClient
|
5
7
|
except ImportError as e:
|
@@ -10,7 +12,6 @@ from functools import lru_cache
|
|
10
12
|
from operator import itemgetter
|
11
13
|
from typing import List, Optional, Self, Type, Unpack
|
12
14
|
|
13
|
-
from fabricatio.rust import CONFIG, TEMPLATE_MANAGER
|
14
15
|
from more_itertools.recipes import flatten, unique
|
15
16
|
from pydantic import Field, PrivateAttr
|
16
17
|
|
@@ -19,6 +20,7 @@ from fabricatio.models.adv_kwargs_types import CollectionConfigKwargs, FetchKwar
|
|
19
20
|
from fabricatio.models.extra.rag import MilvusDataBase
|
20
21
|
from fabricatio.models.kwargs_types import ChooseKwargs
|
21
22
|
from fabricatio.models.usages import EmbeddingUsage
|
23
|
+
from fabricatio.rust import CONFIG, TEMPLATE_MANAGER
|
22
24
|
from fabricatio.utils import ok
|
23
25
|
|
24
26
|
|
@@ -32,7 +34,7 @@ def create_client(uri: str, token: str = "", timeout: Optional[float] = None) ->
|
|
32
34
|
)
|
33
35
|
|
34
36
|
|
35
|
-
class RAG(EmbeddingUsage):
|
37
|
+
class RAG(EmbeddingUsage, ABC):
|
36
38
|
"""A class representing the RAG (Retrieval Augmented Generation) model."""
|
37
39
|
|
38
40
|
target_collection: Optional[str] = Field(default=None)
|
@@ -49,16 +51,16 @@ class RAG(EmbeddingUsage):
|
|
49
51
|
return self._client
|
50
52
|
|
51
53
|
def init_client(
|
52
|
-
|
53
|
-
|
54
|
-
|
55
|
-
|
54
|
+
self,
|
55
|
+
milvus_uri: Optional[str] = None,
|
56
|
+
milvus_token: Optional[str] = None,
|
57
|
+
milvus_timeout: Optional[float] = None,
|
56
58
|
) -> Self:
|
57
59
|
"""Initialize the Milvus client."""
|
58
60
|
self._client = create_client(
|
59
|
-
uri=milvus_uri or ok(self.milvus_uri or CONFIG.rag.milvus_uri)
|
61
|
+
uri=milvus_uri or ok(self.milvus_uri or CONFIG.rag.milvus_uri),
|
60
62
|
token=milvus_token
|
61
|
-
|
63
|
+
or (token.get_secret_value() if (token := (self.milvus_token or CONFIG.rag.milvus_token)) else ""),
|
62
64
|
timeout=milvus_timeout or self.milvus_timeout or CONFIG.rag.milvus_timeout,
|
63
65
|
)
|
64
66
|
return self
|
@@ -72,7 +74,7 @@ class RAG(EmbeddingUsage):
|
|
72
74
|
return self
|
73
75
|
|
74
76
|
def view(
|
75
|
-
|
77
|
+
self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionConfigKwargs]
|
76
78
|
) -> Self:
|
77
79
|
"""View the specified collection.
|
78
80
|
|
@@ -114,7 +116,7 @@ class RAG(EmbeddingUsage):
|
|
114
116
|
return ok(self.target_collection, "No collection is being viewed. Have you called `self.view()`?")
|
115
117
|
|
116
118
|
async def add_document[D: MilvusDataBase](
|
117
|
-
|
119
|
+
self, data: List[D] | D, collection_name: Optional[str] = None, flush: bool = False
|
118
120
|
) -> Self:
|
119
121
|
"""Adds a document to the specified collection.
|
120
122
|
|
@@ -141,15 +143,15 @@ class RAG(EmbeddingUsage):
|
|
141
143
|
return self
|
142
144
|
|
143
145
|
async def afetch_document[D: MilvusDataBase](
|
144
|
-
|
145
|
-
|
146
|
-
|
147
|
-
|
148
|
-
|
149
|
-
|
150
|
-
|
151
|
-
|
152
|
-
|
146
|
+
self,
|
147
|
+
query: List[str],
|
148
|
+
document_model: Type[D],
|
149
|
+
collection_name: Optional[str] = None,
|
150
|
+
similarity_threshold: float = 0.37,
|
151
|
+
result_per_query: int = 10,
|
152
|
+
tei_endpoint: Optional[str] = None,
|
153
|
+
reranker_threshold: float = 0.7,
|
154
|
+
filter_expr: str = "",
|
153
155
|
) -> List[D]:
|
154
156
|
"""Asynchronously fetches documents from a Milvus database based on input vectors.
|
155
157
|
|
@@ -162,7 +164,7 @@ class RAG(EmbeddingUsage):
|
|
162
164
|
result_per_query (int): The maximum number of results to return per query. Defaults to 10.
|
163
165
|
tei_endpoint (str): the endpoint of the TEI api.
|
164
166
|
reranker_threshold (float): The threshold used to filtered low relativity document.
|
165
|
-
|
167
|
+
filter_expr (str) : The filter expression used to filter out unwanted documents.
|
166
168
|
|
167
169
|
Returns:
|
168
170
|
List[D]: A list of document objects created from the fetched data.
|
@@ -177,9 +179,9 @@ class RAG(EmbeddingUsage):
|
|
177
179
|
limit=result_per_query,
|
178
180
|
)
|
179
181
|
if tei_endpoint is not None:
|
180
|
-
from fabricatio.
|
182
|
+
from fabricatio.rust import TEIClient
|
181
183
|
|
182
|
-
reranker =
|
184
|
+
reranker = TEIClient(base_url=tei_endpoint)
|
183
185
|
|
184
186
|
retrieved_id = set()
|
185
187
|
raw_result = []
|
@@ -191,9 +193,7 @@ class RAG(EmbeddingUsage):
|
|
191
193
|
if not models:
|
192
194
|
continue
|
193
195
|
rank_scores = await reranker.arerank(q, [m.prepare_vectorization() for m in models], truncate=True)
|
194
|
-
raw_result.extend(
|
195
|
-
(models[s["index"]], s["score"]) for s in rank_scores if s["score"] > reranker_threshold
|
196
|
-
)
|
196
|
+
raw_result.extend((models[idx], scr) for (idx, scr) in rank_scores if scr > reranker_threshold)
|
197
197
|
|
198
198
|
raw_result_sorted = sorted(raw_result, key=lambda x: x[1], reverse=True)
|
199
199
|
return [r[0] for r in raw_result_sorted]
|
@@ -214,11 +214,11 @@ class RAG(EmbeddingUsage):
|
|
214
214
|
return document_model.from_sequence(resp)
|
215
215
|
|
216
216
|
async def aretrieve[D: MilvusDataBase](
|
217
|
-
|
218
|
-
|
219
|
-
|
220
|
-
|
221
|
-
|
217
|
+
self,
|
218
|
+
query: List[str] | str,
|
219
|
+
document_model: Type[D],
|
220
|
+
max_accepted: int = 20,
|
221
|
+
**kwargs: Unpack[FetchKwargs],
|
222
222
|
) -> List[D]:
|
223
223
|
"""Retrieve data from the collection.
|
224
224
|
|
@@ -235,15 +235,15 @@ class RAG(EmbeddingUsage):
|
|
235
235
|
query = [query]
|
236
236
|
|
237
237
|
return (
|
238
|
-
|
239
|
-
|
240
|
-
|
241
|
-
|
242
|
-
|
243
|
-
|
238
|
+
await self.afetch_document(
|
239
|
+
query=query,
|
240
|
+
document_model=document_model,
|
241
|
+
**kwargs,
|
242
|
+
)
|
243
|
+
)[:max_accepted]
|
244
244
|
|
245
245
|
async def arefined_query(
|
246
|
-
|
246
|
+
self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs[Optional[List[str]]]]
|
247
247
|
) -> Optional[List[str]]:
|
248
248
|
"""Refines the given question using a template.
|
249
249
|
|