feat: backwards invoke app

This commit is contained in:
Yeuoly 2024-08-29 20:17:17 +08:00
parent 113ff27d07
commit 41ed2e0cc2
No known key found for this signature in database
GPG Key ID: A66E7E320FB19F61
7 changed files with 187 additions and 13 deletions

View File

@ -1,5 +1,4 @@
import time
from collections.abc import Generator
from flask_restful import Resource, reqparse
@ -30,15 +29,10 @@ class PluginInvokeLLMApi(Resource):
def post(self, user_id: str, tenant_model: Tenant, payload: RequestInvokeLLM):
def generator():
response = PluginModelBackwardsInvocation.invoke_llm(user_id, tenant_model, payload)
if isinstance(response, Generator):
for chunk in response:
yield chunk.model_dump_json().encode() + b'\n\n'
else:
yield response.model_dump_json().encode() + b'\n\n'
return PluginModelBackwardsInvocation.convert_to_event_stream(response)
return compact_generate_response(generator())
class PluginInvokeTextEmbeddingApi(Resource):
@setup_required
@plugin_inner_api_only

View File

@ -0,0 +1,141 @@
from collections.abc import Generator, Mapping
from typing import Literal, Union
from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator
from core.app.apps.workflow.app_generator import WorkflowAppGenerator
from core.app.entities.app_invoke_entities import InvokeFrom
from core.plugin.backwards_invocation.base import BaseBackwardsInvocation
from extensions.ext_database import db
from models.account import Account
from models.model import App, AppMode, EndUser
class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
@classmethod
def invoke_app(
cls, app_id: str,
user_id: str,
tenant_id: str,
query: str,
inputs: Mapping,
files: list[dict],
) -> Generator[dict, None, None] | dict:
"""
invoke app
"""
app = cls._get_app(app_id, tenant_id)
if app.mode in [AppMode.ADVANCED_CHAT.value, AppMode.AGENT_CHAT.value, AppMode.CHAT.value]:
return cls.invoke_chat_app(app, user_id, tenant_id, query, inputs, files)
elif app.mode in [AppMode.WORKFLOW.value]:
return cls.invoke_workflow_app(app, user_id, tenant_id, inputs, files)
elif app.mode in [AppMode.COMPLETION]:
return cls.invoke_completion_app(app, user_id, tenant_id, inputs, files)
raise ValueError("unexpected app type")
@classmethod
def invoke_chat_app(
cls,
app: App,
user: Account | EndUser,
tenant_id: str,
conversation_id: str,
query: str,
stream: bool,
inputs: Mapping,
files: list[dict],
) -> Generator[dict, None, None] | dict:
"""
invoke chat app
"""
if app.mode == AppMode.ADVANCED_CHAT.value:
workflow = app.workflow
if not workflow:
raise ValueError("unexpected app type")
generator = AdvancedChatAppGenerator()
response = generator.generate(
app_model=app,
workflow=workflow,
user=user,
args={
},
invoke_from=InvokeFrom.SERVICE_API,
stream=stream
)
@classmethod
def invoke_workflow_app(
cls,
app: App,
user_id: str,
tenant_id: str,
inputs: Mapping,
files: list[dict],
):
"""
invoke workflow app
"""
workflow = app.workflow
if not workflow:
raise ValueError("")
generator = WorkflowAppGenerator()
result = generator.generate(
app_model=app,
workflow=workflow,
user=cls._get_user(user_id),
args={
'inputs': tool_parameters,
'files': files
},
invoke_from=self.runtime.invoke_from,
stream=False,
call_depth=self.workflow_call_depth + 1,
)
@classmethod
def invoke_completion_app(
cls,
app: App,
user_id: str,
tenant_id: str,
inputs: Mapping,
files: list[dict],
):
"""
invoke completion app
"""
@classmethod
def _get_user(cls, user_id: str) -> Union[EndUser, Account]:
"""
get the user by user id
"""
user = db.session.query(EndUser).filter(EndUser.id == user_id).first()
if not user:
user = db.session.query(Account).filter(Account.id == user_id).first()
if not user:
raise ValueError('user not found')
return user
@classmethod
def _get_app(cls, app_id: str, tenant_id: str) -> App:
"""
get app
"""
app = db.session.query(App). \
filter(App.id == app_id). \
filter(App.tenant_id == tenant_id). \
first()
if not app:
raise ValueError("app not found")
return app

View File

@ -0,0 +1,20 @@
import json
from collections.abc import Generator
from pydantic import BaseModel
class BaseBackwardsInvocation:
@classmethod
def convert_to_event_stream(cls, response: Generator[BaseModel | dict, None, None] | BaseModel | dict):
if isinstance(response, Generator):
for chunk in response:
if isinstance(chunk, BaseModel):
yield chunk.model_dump_json().encode() + b'\n\n'
else:
yield json.dumps(chunk).encode() + b'\n\n'
else:
if isinstance(response, BaseModel):
yield response.model_dump_json().encode() + b'\n\n'
else:
yield json.dumps(response).encode() + b'\n\n'

View File

@ -2,12 +2,13 @@ from collections.abc import Generator
from core.model_manager import ModelManager
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk
from core.plugin.backwards_invocation.base import BaseBackwardsInvocation
from core.plugin.entities.request import RequestInvokeLLM
from core.workflow.nodes.llm.llm_node import LLMNode
from models.account import Tenant
class PluginBackwardsInvocation:
class PluginModelBackwardsInvocation(BaseBackwardsInvocation):
@classmethod
def invoke_llm(
cls, user_id: str, tenant: Tenant, payload: RequestInvokeLLM
@ -47,3 +48,5 @@ class PluginBackwardsInvocation:
if response.usage:
LLMNode.deduct_llm_quota(tenant_id=tenant.id, model_instance=model_instance, usage=response.usage)
return response

View File

@ -1,4 +1,4 @@
from typing import Any, Optional
from typing import Any, Literal, Optional
from pydantic import BaseModel, Field, field_validator
@ -93,3 +93,15 @@ class RequestInvokeNode(BaseModel):
"""
Request to invoke node
"""
class RequestInvokeApp(BaseModel):
"""
Request to invoke app
"""
app_id: str
inputs: dict[str, Any]
query: Optional[str] = None
response_mode: Literal["blocking", "streaming"]
conversation_id: Optional[str] = None
user: Optional[str] = None
files: list[dict] = Field(default_factory=list)

View File

@ -9,6 +9,7 @@ from werkzeug.local import LocalProxy
from extensions.ext_database import db
from models.account import Account, Tenant, TenantAccountJoin
from models.model import EndUser
#: A proxy for the current user. If no user is logged in, this will be an
#: anonymous user
@ -96,7 +97,7 @@ def login_required(func):
return decorated_view
def _get_user():
def _get_user() -> EndUser | Account | None:
if has_request_context():
if "_login_user" not in g:
current_app.login_manager._load_user()

View File

@ -278,7 +278,10 @@ class ToolConversationVariables(db.Model):
def variables(self) -> dict:
return json.loads(self.variables_str)
class ToolFile(DeclarativeBase):
class Base(DeclarativeBase):
pass
class ToolFile(Base):
"""
store the file created by agent
"""
@ -293,9 +296,9 @@ class ToolFile(DeclarativeBase):
# conversation user id
user_id: Mapped[str] = mapped_column(StringUUID)
# tenant id
tenant_id: Mapped[StringUUID] = mapped_column(StringUUID)
tenant_id: Mapped[str] = mapped_column(StringUUID)
# conversation id
conversation_id: Mapped[StringUUID] = mapped_column(nullable=True)
conversation_id: Mapped[str] = mapped_column(StringUUID, nullable=True)
# file key
file_key: Mapped[str] = mapped_column(db.String(255), nullable=False)
# mime type