adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(multimodel): support multimodal memory display and improve code style

Browse Source
This commit is contained in:
Eternity
2026-03-13 13:33:58 +08:00
parent cbc8714414
commit b71bc1f875
31 changed files with 877 additions and 543 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
api/app/core/workflow/nodes/base_node.py
View File

@@ -16,6 +16,7 @@ from app.core.workflow.variable.base_variable import VariableType, FileObject
from app.db import get_db_read
from app.models import ModelConfig, ModelApiKey, LoadBalanceStrategy
from app.schemas import FileInput
from app.schemas.model_schema import ModelInfo
from app.services.multimodal_service import MultimodalService
logger = logging.getLogger(__name__)
@@ -620,11 +621,12 @@ class BaseNode(ABC):
@staticmethod
async def process_message(
provider: str,
is_omni: bool,
api_config: ModelInfo,
content: str | dict | FileObject,
end_user_id: str,
enable_file=False
) -> list | str | None:
provider = api_config.provider
if isinstance(content, dict):
content = FileObject(
type=content.get("type"),
@@ -643,7 +645,7 @@ class BaseNode(ABC):
if content.content_cache.get(provider):
return content.content_cache[provider]
with get_db_read() as db:
multimodel_service = MultimodalService(db, provider, is_omni=is_omni)
multimodel_service = MultimodalService(db, api_config=api_config)
file_obj = FileInput(
type=content.type,
url=content.url,
@@ -653,7 +655,8 @@ class BaseNode(ABC):
)
file_obj.set_content(content.get_content())
message = await multimodel_service.process_files(
[file_obj]
end_user_id,
[file_obj],
)
content.set_content(file_obj.get_content())
if message:

22
api/app/core/workflow/nodes/if_else/node.py
View File

@@ -5,7 +5,7 @@ from typing import Any
from app.core.workflow.engine.state_manager import WorkflowState
from app.core.workflow.engine.variable_pool import VariablePool
from app.core.workflow.nodes.base_node import BaseNode
from app.core.workflow.nodes.enums import ComparisonOperator, LogicOperator
from app.core.workflow.nodes.enums import ComparisonOperator, LogicOperator, ValueInputType
from app.core.workflow.nodes.if_else import IfElseNodeConfig
from app.core.workflow.nodes.operators import ConditionExpressionResolver, CompareOperatorInstance
from app.core.workflow.variable.base_variable import VariableType
@@ -23,6 +23,26 @@ class IfElseNode(BaseNode):
"output": VariableType.STRING
}
def _extract_input(self, state: WorkflowState, variable_pool: VariablePool) -> dict[str, Any]:
result = []
for case in self.typed_config.cases:
expressions = []
for expression in case.expressions:
expressions.append({
"left": self.get_variable(expression.left, variable_pool, strict=False),
"right": expression.right
if expression.input_type == ValueInputType.CONSTANT
else self.get_variable(expression.right, variable_pool, strict=False),
"operator": expression.operator,
})
result.append({
"expressions": expressions,
"logical_operator": case.logical_operator,
})
return {
"cases": result
}
@staticmethod
def _evaluate(operator, instance: CompareOperatorInstance) -> Any:
match operator:

6
api/app/core/workflow/nodes/knowledge/node.py
View File

@@ -30,6 +30,12 @@ class KnowledgeRetrievalNode(BaseNode):
"output": VariableType.ARRAY_STRING
}
def _extract_input(self, state: WorkflowState, variable_pool: VariablePool) -> dict[str, Any]:
return {
"query": self._render_template(self.typed_config.query, variable_pool),
"knowledge_bases": [kb_config.model_dump(mode="json") for kb_config in self.typed_config.knowledge_bases],
}
@staticmethod
def _build_kb_filter(kb_ids: list[uuid.UUID], permission: knowledge_model.PermissionType):
"""

59
api/app/core/workflow/nodes/llm/node.py
View File

@@ -20,6 +20,7 @@ from app.core.workflow.nodes.llm.config import LLMNodeConfig
from app.core.workflow.variable.base_variable import VariableType
from app.db import get_db_context
from app.models import ModelType
from app.schemas.model_schema import ModelInfo
from app.services.model_service import ModelConfigService
logger = logging.getLogger(__name__)
@@ -113,12 +114,15 @@ class LLMNode(BaseNode):
# 在 Session 关闭前提取所有需要的数据
api_config = self.model_balance(config)
model_name = api_config.model_name
provider = api_config.provider
api_key = api_config.api_key
api_base = api_config.api_base
is_omni = api_config.is_omni
model_type = config.type
model_info = ModelInfo(
model_name=api_config.model_name,
model_type=ModelType(config.type),
api_key=api_config.api_key,
api_base=api_config.api_base,
provider=api_config.provider,
is_omni=api_config.is_omni,
capability=api_config.capability
)
# 4. 创建 LLM 实例（使用已提取的数据）
# 注意：对于流式输出，需要在模型初始化时设置 streaming=True
@@ -126,17 +130,18 @@ class LLMNode(BaseNode):
llm = RedBearLLM(
RedBearModelConfig(
model_name=model_name,
provider=provider,
api_key=api_key,
base_url=api_base,
model_name=model_info.model_name,
provider=model_info.provider,
api_key=model_info.api_key,
base_url=model_info.api_base,
extra_params=extra_params,
is_omni=is_omni
is_omni=model_info.is_omni
),
type=ModelType(model_type)
type=model_info.model_type
)
logger.debug(f"创建 LLM 实例: provider={provider}, model={model_name}, streaming={stream}")
logger.debug(
f"创建 LLM 实例: provider={model_info.provider}, model={model_info.model_name}, streaming={stream}")
messages_config = self.typed_config.messages
@@ -148,35 +153,40 @@ class LLMNode(BaseNode):
content_template = msg_config.content
content_template = self._render_context(content_template, variable_pool)
content = self._render_template(content_template, variable_pool)
user_id = self.get_variable("sys.user_id", variable_pool)
# 根据角色创建对应的消息对象
if role == "system":
messages.append({
"role": "system",
"content": await self.process_message(provider, is_omni, content, self.typed_config.vision)
"content": await self.process_message(
model_info,
content,
user_id,
self.typed_config.vision,
)
})
elif role in ["user", "human"]:
messages.append({
"role": "user",
"content": await self.process_message(provider, is_omni, content, self.typed_config.vision)
"content": await self.process_message(model_info, content, user_id, self.typed_config.vision)
})
elif role in ["ai", "assistant"]:
messages.append({
"role": "assistant",
"content": await self.process_message(provider, is_omni, content, self.typed_config.vision)
"content": await self.process_message(model_info, content, user_id, self.typed_config.vision)
})
else:
logger.warning(f"未知的消息角色: {role}，默认使用 user")
messages.append({
"role": "user",
"content": await self.process_message(provider, is_omni, content, self.typed_config.vision)
"content": await self.process_message(model_info, content, user_id, self.typed_config.vision)
})
if self.typed_config.vision_input and self.typed_config.vision:
file_content = []
files = variable_pool.get_instance(self.typed_config.vision_input)
for file in files.value:
content = await self.process_message(provider, is_omni, file.value, self.typed_config.vision)
content = await self.process_message(model_info, file.value, user_id, self.typed_config.vision)
if content:
file_content.extend(content)
if messages and messages[-1]["role"] == 'user':
@@ -190,14 +200,19 @@ class LLMNode(BaseNode):
if isinstance(message["content"], list):
file_content = []
for file in message["content"]:
content = await self.process_message(provider, is_omni, file, self.typed_config.vision)
content = await self.process_message(model_info, file, user_id, self.typed_config.vision)
if content:
file_content.extend(content)
history_message.append(
{"role": message["role"], "content": file_content}
)
else:
message["content"] = await self.process_message(provider, is_omni, message["content"], self.typed_config.vision)
message["content"] = await self.process_message(
model_info,
message["content"],
user_id,
self.typed_config.vision
)
history_message.append(message)
messages = messages[:-1] + history_message + messages[-1:]
self.messages = messages
@@ -293,7 +308,7 @@ class LLMNode(BaseNode):
# 调用 LLM（流式，支持字符串或消息列表）
last_meta_data = {}
async for chunk in llm.astream(self.messages, stream_usage=True):
async for chunk in llm.astream(self.messages):
# 提取内容
if hasattr(chunk, 'content'):
content = self.process_model_output(chunk.content)

8
api/app/core/workflow/nodes/parameter_extractor/node.py
View File

@@ -37,6 +37,14 @@ class ParameterExtractorNode(BaseNode):
}
return None
def _extract_input(self, state: WorkflowState, variable_pool: VariablePool) -> dict[str, Any]:
return {
"text": self._render_template(self.typed_config.text, variable_pool),
"prompt": self._render_template(self.typed_config.prompt, variable_pool),
"params": [param.model_dump(mode="json") for param in self.typed_config.params],
"model_id": str(self.typed_config.model_id),
}
def _output_types(self) -> dict[str, VariableType]:
outputs = {}
for param in self.typed_config.params: