From 3139b62d768d92c9fc095e31aab545e1128a4356 Mon Sep 17 00:00:00 2001
From: fit2cloud-chenyw <yawen.chen@fit2cloud.com>
Date: Wed, 3 Sep 2025 17:52:43 +0800
Subject: [PATCH] perf: Assistant's basic application records are stored in the
 workspace of their respective data source

---
 backend/apps/system/api/assistant.py   |  8 +++++++-
 backend/apps/system/middleware/auth.py | 11 +++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/backend/apps/system/api/assistant.py b/backend/apps/system/api/assistant.py
index f50cfd9b..ea41e6d7 100644
--- a/backend/apps/system/api/assistant.py
+++ b/backend/apps/system/api/assistant.py
@@ -64,9 +64,15 @@ async def validator(session: SessionDep, id: int, virtual: Optional[int] = Query
     if not db_model:
         return AssistantValidator()
     db_model = AssistantModel.model_validate(db_model)
+    assistant_oid = 1
+    if(db_model.type == 0):
+        configuration = db_model.configuration
+        config_obj = json.loads(configuration) if configuration else {}
+        assistant_oid = config_obj.get('oid', 1)
+    
     access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
     assistantDict = {
-        "id": virtual, "account": 'sqlbot-inner-assistant', "oid": 1, "assistant_id": id
+        "id": virtual, "account": 'sqlbot-inner-assistant', "oid": assistant_oid, "assistant_id": id
     }
     access_token = create_access_token(
         assistantDict, expires_delta=access_token_expires
diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py
index 6da8cfcc..9a2e3fe1 100644
--- a/backend/apps/system/middleware/auth.py
+++ b/backend/apps/system/middleware/auth.py
@@ -1,5 +1,6 @@
 
 import base64
+import json
 from typing import Optional
 from fastapi import Request
 from fastapi.responses import JSONResponse
@@ -112,6 +113,16 @@ class TokenMiddleware(BaseHTTPMiddleware):
                 assistant_info = await get_assistant_info(session=session, assistant_id=payload['assistant_id'])
                 assistant_info = AssistantModel.model_validate(assistant_info)
                 assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
+                if assistant_info and assistant_info.type == 0:
+                    if payload['oid']:
+                        session_user.oid = payload['oid']
+                    else:
+                        assistant_oid = 1
+                        configuration = assistant_info.configuration
+                        config_obj = json.loads(configuration) if configuration else {}
+                        assistant_oid = config_obj.get('oid', 1)
+                        session_user.oid = assistant_oid
+                        
                 return True, session_user, assistant_info
         except Exception as e:
             SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}")