Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

Dockerfile

@@ -54,6 +54,12 @@ ENV SQLBOT_HOME=/opt/sqlbot
 ENV PYTHONPATH=${SQLBOT_HOME}/app
 ENV PATH="${SQLBOT_HOME}/app/.venv/bin:$PATH"
 
+ENV SQLBOT_DB_HOST=localhost
+ENV SQLBOT_DB_PORT=5432
+ENV SQLBOT_DB_DB=sqlbot
+ENV SQLBOT_DB_USER=root
+ENV SQLBOT_DB_PASSWORD=Password123@pg
+
 # Copy necessary files from builder
 COPY start.sh /opt/sqlbot/app/start.sh
 COPY g2-ssr/*.ttf /usr/share/fonts/truetype/liberation/

backend/apps/chat/api/chat.py

@@ -3,6 +3,7 @@
 import traceback
 
 import numpy as np
+import orjson
 import pandas as pd
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
@@ -147,56 +148,52 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
         llm_service.run_task_async()
     except Exception as e:
         traceback.print_exc()
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
-        )
+
+        def _err(_e: Exception):
+            yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
+
+        return StreamingResponse(_err(e), media_type="text/event-stream")
 
     return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
 
 
 @router.post("/record/{chat_record_id}/{action_type}")
 async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str,
                               current_assistant: CurrentAssistant):
-    if action_type != 'analysis' and action_type != 'predict':
-        raise HTTPException(
-            status_code=404,
-            detail="Not Found"
-        )
-    record: ChatRecord | None = None
-
-    stmt = select(ChatRecord.id, ChatRecord.question, ChatRecord.chat_id, ChatRecord.datasource, ChatRecord.engine_type,
-                  ChatRecord.ai_modal_id, ChatRecord.create_by, ChatRecord.chart, ChatRecord.data).where(
-        and_(ChatRecord.id == chat_record_id))
-    result = session.execute(stmt)
-    for r in result:
-        record = ChatRecord(id=r.id, question=r.question, chat_id=r.chat_id, datasource=r.datasource,
-                            engine_type=r.engine_type, ai_modal_id=r.ai_modal_id, create_by=r.create_by, chart=r.chart,
-                            data=r.data)
-
-    if not record:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Chat record with id {chat_record_id} not found"
-        )
+    try:
+        if action_type != 'analysis' and action_type != 'predict':
+            raise Exception(f"Type {action_type} Not Found")
+        record: ChatRecord | None = None
+
+        stmt = select(ChatRecord.id, ChatRecord.question, ChatRecord.chat_id, ChatRecord.datasource,
+                      ChatRecord.engine_type,
+                      ChatRecord.ai_modal_id, ChatRecord.create_by, ChatRecord.chart, ChatRecord.data).where(
+            and_(ChatRecord.id == chat_record_id))
+        result = session.execute(stmt)
+        for r in result:
+            record = ChatRecord(id=r.id, question=r.question, chat_id=r.chat_id, datasource=r.datasource,
+                                engine_type=r.engine_type, ai_modal_id=r.ai_modal_id, create_by=r.create_by,
+                                chart=r.chart,
+                                data=r.data)
 
-    if not record.chart:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Chat record with id {chat_record_id} has not generated chart, do not support to analyze it"
-        )
+        if not record:
+            raise Exception(f"Chat record with id {chat_record_id} not found")
 
-    request_question = ChatQuestion(chat_id=record.chat_id, question=record.question)
+        if not record.chart:
+            raise Exception(
+                f"Chat record with id {chat_record_id} has not generated chart, do not support to analyze it")
+
+        request_question = ChatQuestion(chat_id=record.chat_id, question=record.question)
 
-    try:
         llm_service = await LLMService.create(current_user, request_question, current_assistant)
         llm_service.run_analysis_or_predict_task_async(action_type, record)
     except Exception as e:
         traceback.print_exc()
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
-        )
+
+        def _err(_e: Exception):
+            yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
+
+        return StreamingResponse(_err(e), media_type="text/event-stream")
 
     return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
 

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

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 @@ async def validateAssistant(self, assistantToken: Optional[str], trans: I18n) ->
                 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)}")