knowledge fix

This commit is contained in:
jyong 2024-03-18 21:22:16 +08:00
parent cc86850ad9
commit 3e810bc490
4 changed files with 11 additions and 13 deletions

View File

@ -3,7 +3,6 @@ from typing import Any, Literal, Optional
from pydantic import BaseModel
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.variable_entities import VariableSelector
class RerankingModelConfig(BaseModel):
@ -11,7 +10,7 @@ class RerankingModelConfig(BaseModel):
Reranking Model Config.
"""
provider: str
model: str
mode: str
class MultipleRetrievalConfig(BaseModel):
@ -45,8 +44,8 @@ class KnowledgeRetrievalNodeData(BaseNodeData):
Knowledge retrieval Node Data.
"""
type: str = 'knowledge-retrieval'
query_variable_selector: VariableSelector
query_variable_selector: list[str]
dataset_ids: list[str]
retrieval_mode: Literal['single', 'multiple']
multiple_retrieval_config: Optional[MultipleRetrievalConfig] = None
singleRetrievalConfig: Optional[SingleRetrievalConfig] = None
multiple_retrieval_config: Optional[MultipleRetrievalConfig]
singleRetrievalConfig: Optional[SingleRetrievalConfig]

View File

@ -43,9 +43,9 @@ class KnowledgeRetrievalNode(BaseNode):
node_data: KnowledgeRetrievalNodeData = cast(self._node_data_cls, self.node_data)
# extract variables
query = variable_pool.get_variable_value(variable_selector=node_data.query_variable_selector.value_selector)
query = variable_pool.get_variable_value(variable_selector=node_data.query_variable_selector)
variables = {
node_data.query_variable_selector.variable: query
'query': query
}
# retrieve knowledge
try:
@ -170,7 +170,7 @@ class KnowledgeRetrievalNode(BaseNode):
node_data = node_data
node_data = cast(cls._node_data_cls, node_data)
variable_mapping = {}
variable_mapping[node_data.query_variable_selector.variable] = node_data.query_variable_selector.value_selector
variable_mapping['query'] = node_data.query_variable_selector
return variable_mapping
def _single_retrieve(self, available_datasets, node_data, query):

View File

@ -3,7 +3,6 @@ from typing import Any
from pydantic import BaseModel
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.variable_entities import VariableSelector
class ModelConfig(BaseModel):
@ -43,7 +42,7 @@ class QuestionClassifierNodeData(BaseNodeData):
"""
Knowledge retrieval Node Data.
"""
query_variable_selector: VariableSelector
query_variable_selector: list[str]
type: str = 'question-classifier'
model: ModelConfig
classes: list[ClassConfig]

View File

@ -43,9 +43,9 @@ class QuestionClassifierNode(BaseNode):
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
node_data: QuestionClassifierNodeData = cast(self._node_data_cls, self.node_data)
# extract variables
query = variable_pool.get_variable_value(variable_selector=node_data.query_variable_selector.value_selector)
query = variable_pool.get_variable_value(variable_selector=node_data.query_variable_selector)
variables = {
node_data.query_variable_selector.variable: query
'query': query
}
# fetch model config
model_instance, model_config = self._fetch_model_config(node_data)
@ -104,7 +104,7 @@ class QuestionClassifierNode(BaseNode):
def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[str, list[str]]:
node_data = node_data
node_data = cast(cls._node_data_cls, node_data)
variable_mapping = {node_data.query_variable_selector.variable: node_data.query_variable_selector.value_selector}
variable_mapping = {'query': node_data.query_variable_selector}
return variable_mapping
@classmethod