SQLBot/backend/apps/system/crud/assistant.py

import json
from typing import Optional
from fastapi import FastAPI
import requests
from sqlalchemy import Engine, create_engine
from sqlmodel import Session, select
from apps.datasource.models.datasource import CoreDatasource, DatasourceConf

from apps.system.models.system_model import AssistantModel
from apps.system.schemas.auth import CacheName, CacheNamespace
from apps.system.schemas.system_schema import AssistantHeader, AssistantOutDsSchema, UserInfoDTO
from common.core.sqlbot_cache import cache
from common.core.db import engine
from starlette.middleware.cors import CORSMiddleware
from common.core.config import settings
from common.utils.utils import string_to_numeric_hash

@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id")
async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None:
    db_model = session.get(AssistantModel, assistant_id)
    return db_model

def get_assistant_user(*, id: int):
    return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", email="[email protected]")

def get_assistant_ds(llm_service) -> list[dict]:
    assistant: AssistantHeader = llm_service.current_assistant
    session: Session = llm_service.session
    type = assistant.type
    if type == 0:
        configuration = assistant.configuration
        if configuration:
            config: dict[any] = json.loads(configuration)
            oid: str = config['oid']
            stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid)
            private_list:list[int] = config.get('private_list') or None
            if private_list:
                stmt.where(~CoreDatasource.id.in_(private_list))
        db_ds_list = session.exec(stmt)
        
        result_list = [
            {
                "id": ds.id,
                "name": ds.name,
                "description": ds.description
            }
            for ds in db_ds_list
        ]
    
        # filter private ds if offline
        return result_list
    out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant)
    llm_service.out_ds_instance = out_ds_instance
    dslist = out_ds_instance.get_simple_ds_list()
    # format?
    return dslist

def init_dynamic_cors(app: FastAPI):
    try: 
        with Session(engine) as session:
            list_result = session.exec(select(AssistantModel).order_by(AssistantModel.create_time)).all()
            seen = set()
            unique_domains = []
            for item in list_result:
                if item.domain:
                    for domain in item.domain.split(','):
                        domain = domain.strip()
                        if domain and domain not in seen:
                            seen.add(domain)
                            unique_domains.append(domain)
            cors_middleware = None
            for middleware in app.user_middleware:
                if middleware.cls == CORSMiddleware:
                    cors_middleware = middleware
                    break
            if cors_middleware:
                updated_origins = list(set(settings.all_cors_origins + unique_domains))
                cors_middleware.kwargs['allow_origins'] = updated_origins
    except Exception as e:
        return False, e
    
    

class AssistantOutDs:
    assistant: AssistantHeader
    ds_list: Optional[list[AssistantOutDsSchema]] = None
    certificate: Optional[str] = None
    def __init__(self, assistant: AssistantHeader):
        self.assistant = assistant
        self.ds_list = None
        self.certificate = assistant.certificate
        self.get_ds_from_api()
        
    #@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id")    
    def get_ds_from_api(self):
        config: dict[any] = json.loads(self.assistant.configuration)
        endpoint: str = config['endpoint']
        certificateList: list[any] = json.loads(self.certificate)
        header = {}
        cookies = {}
        for item in certificateList:
            if item['target'] == 'header':
                header[item['key']] = item['value']
            if item['target'] == 'cookie':
                cookies[item['key']] = item['value']
        
        res = requests.get(url=endpoint, headers=header, cookies=cookies, timeout=10)
        if res.status_code == 200:
            result_json: dict[any] = json.loads(res.text)
            if result_json.get('code') == 0:
                temp_list = result_json.get('data', [])
                self.ds_list = [
                    self.convert2schema(item)
                    for item in temp_list
                ]
                
                return self.ds_list
            else:
                raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}")
        else:
            raise Exception(f"Failed to get datasource list from {endpoint}, status code: {res.status_code}")
    
    def get_simple_ds_list(self):
        if self.ds_list:
            return [{'id': ds.id, 'name': ds.name, 'description': ds.comment} for ds in self.ds_list]
        else:
            raise Exception("Datasource list is not found.")
       
    def get_db_schema(self, ds_id: int) -> str:
        ds = self.get_ds(ds_id)
        schema_str = ""
        db_name = ds.schema
        schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
        for table in ds.tables:
            schema_str += f"# Table: {db_name}.{table.name}"
            schema_str += f", {table.comment}\n[\n"
            field_list = []
            for field in table.fields:
                field_list.append(f"({field.name}:{field.type}, {field.comment})")
            schema_str += ",\n".join(field_list)
            schema_str += '\n]\n'
        return schema_str
    
    def get_ds(self, ds_id: int):
        if self.ds_list:
            for ds in self.ds_list:
                if ds.id == ds_id:
                    return ds
        else:
            raise Exception("Datasource list is not found.")
        raise Exception(f"Datasource with id {ds_id} not found.")

    def convert2schema(self, ds_dict: dict) -> AssistantOutDsSchema:
        id_marker: str = ''
        attr_list = ['name', 'type', 'host', 'port', 'user', 'dataBase', 'schema']
        for attr in attr_list:
            if attr in ds_dict:
                id_marker += str(ds_dict.get(attr, '')) + '--sqlbot--'
        id = string_to_numeric_hash(id_marker)
        return AssistantOutDsSchema(**{**ds_dict, "id": id})
    
class AssistantOutDsFactory:
    @staticmethod
    def get_instance(assistant: AssistantHeader) -> AssistantOutDs:
        return AssistantOutDs(assistant)

def get_ds_engine(ds: AssistantOutDsSchema) -> Engine:
    timeout: int = 30
    connect_args = {"connect_timeout": timeout}
    conf = DatasourceConf(
        host=ds.host, 
        port=ds.port, 
        username=ds.user,
        password=ds.password,   
        database=ds.dataBase,
        driver='',
        extraJdbc=ds.extraParams,
        dbSchema=ds.schema or ''
    )
    from apps.db.db import get_uri_from_config
    uri = get_uri_from_config(ds.type, conf)
    if ds.type == "pg" and ds.schema:
        connect_args.update({"options": f"-c search_path={ds.schema}"})
    engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10)
    return engine