functions.py

from pydantic import BaseModel, ConfigDict
from typing import List, Union, Optional
import time
import logging

from sqlalchemy import Column, String, Text, BigInteger, Boolean
from sqlalchemy.orm import Session

from apps.webui.internal.db import JSONField, Base
from apps.webui.models.users import Users

import json
import copy


from config import SRC_LOG_LEVELS

log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])

####################
# Functions DB Schema
####################


class Function(Base):
    __tablename__ = "function"

    id = Column(String, primary_key=True)
    user_id = Column(String)
    name = Column(Text)
    type = Column(Text)
    content = Column(Text)
    meta = Column(JSONField)
    valves = Column(JSONField)
    is_active = Column(Boolean)
    updated_at = Column(BigInteger)
    created_at = Column(BigInteger)


class FunctionMeta(BaseModel):
    description: Optional[str] = None
    manifest: Optional[dict] = {}


class FunctionModel(BaseModel):
    id: str
    user_id: str
    name: str
    type: str
    content: str
    meta: FunctionMeta
    is_active: bool = False
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch

    model_config = ConfigDict(from_attributes=True)


####################
# Forms
####################


class FunctionResponse(BaseModel):
    id: str
    user_id: str
    type: str
    name: str
    meta: FunctionMeta
    is_active: bool
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch


class FunctionForm(BaseModel):
    id: str
    name: str
    content: str
    meta: FunctionMeta


class FunctionValves(BaseModel):
    valves: Optional[dict] = None


class FunctionsTable:

    def insert_new_function(
        self, db: Session, user_id: str, type: str, form_data: FunctionForm
    ) -> Optional[FunctionModel]:
        function = FunctionModel(
            **{
                **form_data.model_dump(),
                "user_id": user_id,
                "type": type,
                "updated_at": int(time.time()),
                "created_at": int(time.time()),
            }
        )

        try:
            result = Function(**function.model_dump())
            db.add(result)
            db.commit()
            db.refresh(result)
            if result:
                return FunctionModel.model_validate(result)
            else:
                return None
        except Exception as e:
            print(f"Error creating tool: {e}")
            return None

    def get_function_by_id(self, db: Session, id: str) -> Optional[FunctionModel]:
        try:
            function = db.get(Function, id)
            return FunctionModel.model_validate(function)
        except:
            return None

    def get_functions(self, active_only=False) -> List[FunctionModel]:
        if active_only:
            return [
                FunctionModel(**model_to_dict(function))
                for function in Function.select().where(Function.is_active == True)
            ]
        else:
            return [
                FunctionModel(**model_to_dict(function))
                for function in Function.select()
            ]

    def get_functions_by_type(
        self, type: str, active_only=False
    ) -> List[FunctionModel]:
        if active_only:
            return [
                FunctionModel(**model_to_dict(function))
                for function in Function.select().where(
                    Function.type == type, Function.is_active == True
                )
            ]
        else:
            return [
                FunctionModel(**model_to_dict(function))
                for function in Function.select().where(Function.type == type)
            ]

    def get_function_valves_by_id(self, id: str) -> Optional[dict]:
        try:
            function = Function.get(Function.id == id)
            return function.valves if function.valves else {}
        except Exception as e:
            print(f"An error occurred: {e}")
            return None

    def update_function_valves_by_id(
        self, id: str, valves: dict
    ) -> Optional[FunctionValves]:
        try:
            query = Function.update(
                **{"valves": valves},
                updated_at=int(time.time()),
            ).where(Function.id == id)
            query.execute()

            function = Function.get(Function.id == id)
            return FunctionValves(**model_to_dict(function))
        except:
            return None

    def get_user_valves_by_id_and_user_id(
        self, id: str, user_id: str
    ) -> Optional[dict]:
        try:
            user = Users.get_user_by_id(user_id)
            user_settings = user.settings.model_dump()

            # Check if user has "functions" and "valves" settings
            if "functions" not in user_settings:
                user_settings["functions"] = {}
            if "valves" not in user_settings["functions"]:
                user_settings["functions"]["valves"] = {}

            return user_settings["functions"]["valves"].get(id, {})
        except Exception as e:
            print(f"An error occurred: {e}")
            return None

    def update_user_valves_by_id_and_user_id(
        self, id: str, user_id: str, valves: dict
    ) -> Optional[dict]:
        try:
            user = Users.get_user_by_id(user_id)
            user_settings = user.settings.model_dump()

            # Check if user has "functions" and "valves" settings
            if "functions" not in user_settings:
                user_settings["functions"] = {}
            if "valves" not in user_settings["functions"]:
                user_settings["functions"]["valves"] = {}

            user_settings["functions"]["valves"][id] = valves

            # Update the user settings in the database
            query = Users.update_user_by_id(user_id, {"settings": user_settings})
            query.execute()

            return user_settings["functions"]["valves"][id]
        except Exception as e:
            print(f"An error occurred: {e}")
            return None

    def update_function_by_id(self, id: str, updated: dict) -> Optional[FunctionModel]:
        try:
            db.query(Function).filter_by(id=id).update({
                **updated,
                "updated_at": int(time.time()),
            })
            return self.get_function_by_id(db, id)
        except:
            return None

    def deactivate_all_functions(self) -> Optional[bool]:
        try:
            query = Function.update(
                **{"is_active": False},
                updated_at=int(time.time()),
            )

            query.execute()

            return True
        except:
            return None

    def delete_function_by_id(self, db: Session, id: str) -> bool:
        try:
            db.query(Function).filter_by(id=id).delete()
            return True
        except:
            return False


Functions = FunctionsTable()