Merge remote-tracking branch 'origin/main'

# Conflicts:
#	frontend/src/router/index.ts

Author: ziyujiahao

backend/alembic/versions/028_ds_oid.py

@@ -0,0 +1,28 @@
+"""028_ds_oid
+
+Revision ID: e96b16d3daab
+Revises: b049c9f8ca5b
+Create Date: 2025-07-17 14:40:48.522033
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'e96b16d3daab'
+down_revision = 'b049c9f8ca5b'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('core_datasource', sa.Column('oid', sa.BigInteger(), nullable=True))
+    op.execute('update core_datasource set oid = 1')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('core_datasource', 'oid')
+    # ### end Alembic commands ###

backend/alembic/versions/029_modify_chat.py

@@ -0,0 +1,54 @@
+"""029_modify_chat
+
+Revision ID: 77d4c39ec22f
+Revises: e96b16d3daab
+Create Date: 2025-07-17 17:05:13.392973
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '77d4c39ec22f'
+down_revision = 'e96b16d3daab'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('chat', sa.Column('oid', sa.BigInteger(), nullable=True))
+    op.execute('update chat set oid = 1')
+    op.alter_column('chat', 'create_time',
+               existing_type=postgresql.TIMESTAMP(timezone=True),
+               type_=sa.DateTime(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'create_time',
+               existing_type=postgresql.TIMESTAMP(timezone=True),
+               type_=sa.DateTime(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'finish_time',
+               existing_type=postgresql.TIMESTAMP(timezone=True),
+               type_=sa.DateTime(),
+               existing_nullable=True)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('chat_record', 'finish_time',
+               existing_type=sa.DateTime(),
+               type_=postgresql.TIMESTAMP(timezone=True),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'create_time',
+               existing_type=sa.DateTime(),
+               type_=postgresql.TIMESTAMP(timezone=True),
+               existing_nullable=True)
+    op.alter_column('chat', 'create_time',
+               existing_type=sa.DateTime(),
+               type_=postgresql.TIMESTAMP(timezone=True),
+               existing_nullable=True)
+    op.drop_column('chat', 'oid')
+    # ### end Alembic commands ###

backend/apps/chat/curd/chat.py

@@ -4,7 +4,6 @@
 import orjson
 from sqlalchemy import and_
 from sqlalchemy.orm import load_only
-from sqlmodel import select
 
 from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion
 from apps.datasource.models.datasource import CoreDatasource
@@ -13,8 +12,9 @@
 
 
 def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
-    chart_list = session.exec(
-        select(Chat).where(Chat.create_by == current_user.id).order_by(Chat.create_time.desc())).all()
+    oid = current_user.oid if current_user.oid is not None else 1
+    chart_list = session.query(Chat).filter(and_(Chat.create_by == current_user.id, Chat.oid == oid)).order_by(
+        Chat.create_time.desc()).all()
     return chart_list
 
 
@@ -149,6 +149,7 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
 
     chat = Chat(create_time=datetime.datetime.now(),
                 create_by=current_user.id,
+                oid=current_user.oid if current_user.oid is not None else 1,
                 brief=create_chat_obj.question.strip()[:20])
     ds: CoreDatasource | None = None
     if create_chat_obj.datasource:

backend/apps/chat/models/chat_model.py

@@ -17,7 +17,8 @@
 class Chat(SQLModel, table=True):
     __tablename__ = "chat"
     id: Optional[int] = Field(sa_column=Column(BigInteger, Identity(always=True), primary_key=True))
-    create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
+    oid: Optional[int] = Field(sa_column=Column(BigInteger, nullable=True, default=1))
+    create_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
     create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
     brief: str = Field(max_length=64, nullable=True)
     chat_type: str = Field(max_length=20, default="chat")  # chat, datasource
@@ -31,8 +32,8 @@ class ChatRecord(SQLModel, table=True):
     chat_id: int = Field(sa_column=Column(BigInteger, nullable=False))
     ai_modal_id: Optional[int] = Field(sa_column=Column(BigInteger))
     first_chat: bool = Field(sa_column=Column(Boolean, nullable=True, default=False))
-    create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
-    finish_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
+    create_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
+    finish_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
     create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
     datasource: int = Field(sa_column=Column(BigInteger, nullable=True))
     engine_type: str = Field(max_length=64)
@@ -66,6 +67,7 @@ class ChatRecord(SQLModel, table=True):
     analysis_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
     predict_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
 
+
 class CreateChat(BaseModel):
     id: int = None
     question: str = None

backend/apps/chat/task/llm.py

@@ -846,6 +846,8 @@ def run_analysis_or_predict_task(llm_service: LLMService, action_type: str, base
     try:
         llm_service.set_record(save_analysis_predict_record(llm_service.session, base_record, action_type))
 
+        yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n'
+
         if action_type == 'analysis':
             # generate analysis
             analysis_res = llm_service.generate_analysis()

backend/apps/datasource/api/datasource.py

@@ -19,8 +19,8 @@
 
 
 @router.get("/list")
-async def datasource_list(session: SessionDep):
-    return get_datasource_list(session=session)
+async def datasource_list(session: SessionDep, user: CurrentUser):
+    return get_datasource_list(session=session, user=user)
 
 
 @router.post("/get/{id}")
@@ -44,8 +44,8 @@ async def choose_tables(session: SessionDep, id: int, tables: List[CoreTable]):
 
 
 @router.post("/update", response_model=CoreDatasource)
-async def update(session: SessionDep, ds: CoreDatasource):
-    return update_ds(session, ds)
+async def update(session: SessionDep,user: CurrentUser, ds: CoreDatasource):
+    return update_ds(session,user, ds)
 
 
 @router.post("/delete/{id}", response_model=CoreDatasource)

backend/apps/datasource/crud/datasource.py

@@ -22,10 +22,10 @@
     DatasourceConf, TableAndFields
 
 
-def get_datasource_list(session: SessionDep):
-    statement = select(CoreDatasource).order_by(CoreDatasource.create_time.desc())
-    datasource_list = session.exec(statement).fetchall()
-    return datasource_list
+def get_datasource_list(session: SessionDep, user: CurrentUser):
+    oid = user.oid if user.oid is not None else 1
+    return session.query(CoreDatasource).filter(CoreDatasource.oid == oid).order_by(
+        CoreDatasource.create_time.desc()).all()
 
 
 def get_ds(session: SessionDep, id: int):
@@ -45,25 +45,27 @@ def check_status(session: SessionDep, ds: CoreDatasource):
         return False
 
 
-def check_name(session: SessionDep, ds: CoreDatasource):
+def check_name(session: SessionDep, user: CurrentUser, ds: CoreDatasource):
     if ds.id is not None:
         ds_list = session.query(CoreDatasource).filter(
-            and_(CoreDatasource.name == ds.name, CoreDatasource.id != ds.id)).all()
+            and_(CoreDatasource.name == ds.name, CoreDatasource.id != ds.id, CoreDatasource.oid == user.oid)).all()
         if ds_list is not None and len(ds_list) > 0:
             raise 'Name exist'
     else:
-        ds_list = session.query(CoreDatasource).filter(CoreDatasource.name == ds.name).all()
+        ds_list = session.query(CoreDatasource).filter(
+            and_(CoreDatasource.name == ds.name, CoreDatasource.oid == user.oid)).all()
         if ds_list is not None and len(ds_list) > 0:
             raise 'Name exist'
 
 
 def create_ds(session: SessionDep, user: CurrentUser, create_ds: CreateDatasource):
     ds = CoreDatasource()
     deepcopy_ignore_extra(create_ds, ds)
-    check_name(session, ds)
+    check_name(session, user, ds)
     ds.create_time = datetime.datetime.now()
     # status = check_status(session, ds)
     ds.create_by = user.id
+    ds.oid = user.oid if user.oid is not None else 1
     ds.status = "Success"
     ds.type_name = db_type_relation()[ds.type]
     record = CoreDatasource(**ds.model_dump())
@@ -85,9 +87,9 @@ def chooseTables(session: SessionDep, id: int, tables: List[CoreTable]):
     updateNum(session, ds)
 
 
-def update_ds(session: SessionDep, ds: CoreDatasource):
+def update_ds(session: SessionDep, user: CurrentUser, ds: CoreDatasource):
     ds.id = int(ds.id)
-    check_name(session, ds)
+    check_name(session, user, ds)
     status = check_status(session, ds)
     ds.status = "Success" if status is True else "Fail"
     record = session.exec(select(CoreDatasource).where(CoreDatasource.id == ds.id)).first()

backend/apps/datasource/models/datasource.py

@@ -18,6 +18,7 @@ class CoreDatasource(SQLModel, table=True):
     create_by: int = Field(sa_column=Column(BigInteger()))
     status: str = Field(max_length=64, nullable=True)
     num: str = Field(max_length=256, nullable=True)
+    oid: int = Field(sa_column=Column(BigInteger()))
 
 
 class CoreTable(SQLModel, table=True):
@@ -53,6 +54,7 @@ class CreateDatasource(BaseModel):
     create_by: int = 0
     status: str = ''
     num: str = ''
+    oid: int = 1
     tables: List[CoreTable] = []
 
 

backend/apps/mcp/mcp.py

@@ -3,20 +3,33 @@
 
 from datetime import timedelta
 
-from fastapi import APIRouter, HTTPException
+import jwt
+from fastapi import HTTPException, status, APIRouter
 from fastapi.responses import StreamingResponse
+# from fastapi.security import OAuth2PasswordBearer
+from jwt.exceptions import InvalidTokenError
+from pydantic import ValidationError
+from sqlmodel import select
 
 from apps.chat.api.chat import create_chat
 from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart
 from apps.chat.task.llm import LLMService, run_task
-from apps.datasource.crud.datasource import get_datasource_list
 from apps.system.crud.user import authenticate
-from apps.system.models.system_model import AiModelDetail
+from apps.system.crud.user import get_db_user
+from apps.system.models.system_model import UserWsModel
+from apps.system.models.user import UserModel
+from apps.system.schemas.system_schema import BaseUserDTO
+from apps.system.schemas.system_schema import UserInfoDTO
+from common.core import security
 from common.core.config import settings
-from common.core.deps import SessionDep, get_current_user
-from common.core.schemas import Token
+from common.core.deps import SessionDep
+from common.core.schemas import TokenPayload, XOAuth2PasswordBearer, Token
 from common.core.security import create_access_token
 
+reusable_oauth2 = XOAuth2PasswordBearer(
+    tokenUrl=f"{settings.API_V1_STR}/login/access-token"
+)
+
 router = APIRouter(tags=["mcp"], prefix="/mcp")
 
 
@@ -35,21 +48,24 @@
 #     ))
 
 
-@router.get("/ds_list", operation_id="get_datasource_list")
-async def datasource_list(session: SessionDep):
-    return get_datasource_list(session=session)
-
-
-@router.get("/model_list", operation_id="get_model_list")
-async def get_model_list(session: SessionDep):
-    return session.query(AiModelDetail).all()
+# @router.get("/ds_list", operation_id="get_datasource_list")
+# async def datasource_list(session: SessionDep):
+#     return get_datasource_list(session=session)
+#
+#
+# @router.get("/model_list", operation_id="get_model_list")
+# async def get_model_list(session: SessionDep):
+#     return session.query(AiModelDetail).all()
 
 
 @router.post("/mcp_start", operation_id="mcp_start")
 async def mcp_start(session: SessionDep, chat: ChatStart):
-    user = authenticate(session=session, account=chat.username, password=chat.password)
+    user: BaseUserDTO = authenticate(session=session, account=chat.username, password=chat.password)
     if not user:
         raise HTTPException(status_code=400, detail="Incorrect account or password")
+
+    if not user.oid or user.oid == 0:
+        raise HTTPException(status_code=400, detail="No associated workspace, Please contact the administrator")
     access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
     user_dict = user.to_dict()
     t = Token(access_token=create_access_token(
@@ -61,9 +77,36 @@ async def mcp_start(session: SessionDep, chat: ChatStart):
 
 @router.post("/mcp_question", operation_id="mcp_question")
 async def mcp_question(session: SessionDep, chat: ChatMcp):
-    user = await get_current_user(session, chat.token)
-
-    llm_service = LLMService(session, user, chat)
+    try:
+        payload = jwt.decode(
+            chat.token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
+        )
+        token_data = TokenPayload(**payload)
+    except (InvalidTokenError, ValidationError):
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Could not validate credentials",
+        )
+    # session_user = await get_user_info(session=session, user_id=token_data.id)
+
+    db_user: UserModel = get_db_user(session=session, user_id=token_data.id)
+    session_user = UserInfoDTO.model_validate(db_user.model_dump())
+    session_user.isAdmin = session_user.id == 1 and session_user.account == 'admin'
+    if session_user.isAdmin:
+        session_user = session_user
+    ws_model: UserWsModel = session.exec(
+        select(UserWsModel).where(UserWsModel.uid == session_user.id, UserWsModel.oid == session_user.oid)).first()
+    session_user.weight = ws_model.weight if ws_model else -1
+
+    session_user = UserInfoDTO.model_validate(session_user)
+    if not session_user:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    if session_user.status != 1:
+        raise HTTPException(status_code=400, detail="Inactive user")
+
+    # ask
+    llm_service = LLMService(session, session_user, chat)
     llm_service.init_record()
 
     return StreamingResponse(run_task(llm_service, False), media_type="text/event-stream")