feat: global filter

This commit is contained in:
Timothy J. Baek
2024-06-27 13:04:12 -07:00
parent c8c85ba7fc
commit edbd07f893
8 changed files with 212 additions and 29 deletions

View File

@@ -0,0 +1,49 @@
"""Peewee migrations -- 017_add_user_oauth_sub.py.
Some examples (model - class or model name)::
> Model = migrator.orm['table_name'] # Return model in current state by name
> Model = migrator.ModelClass # Return model in current state by name
> migrator.sql(sql) # Run custom SQL
> migrator.run(func, *args, **kwargs) # Run python function with the given args
> migrator.create_model(Model) # Create a model (could be used as decorator)
> migrator.remove_model(model, cascade=True) # Remove a model
> migrator.add_fields(model, **fields) # Add fields to a model
> migrator.change_fields(model, **fields) # Change fields
> migrator.remove_fields(model, *field_names, cascade=True)
> migrator.rename_field(model, old_field_name, new_field_name)
> migrator.rename_table(model, new_table_name)
> migrator.add_index(model, *col_names, unique=False)
> migrator.add_not_null(model, *field_names)
> migrator.add_default(model, field_name, default)
> migrator.add_constraint(model, name, sql)
> migrator.drop_index(model, *col_names)
> migrator.drop_not_null(model, *field_names)
> migrator.drop_constraints(model, *constraints)
"""
from contextlib import suppress
import peewee as pw
from peewee_migrate import Migrator
with suppress(ImportError):
import playhouse.postgres_ext as pw_pext
def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
"""Write your migrations here."""
migrator.add_fields(
"function",
is_global=pw.BooleanField(default=False),
)
def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
"""Write your rollback migrations here."""
migrator.remove_fields("function", "is_global")

View File

@@ -30,6 +30,7 @@ class Function(Model):
meta = JSONField()
valves = JSONField()
is_active = BooleanField(default=False)
is_global = BooleanField(default=False)
updated_at = BigIntegerField()
created_at = BigIntegerField()
@@ -50,6 +51,7 @@ class FunctionModel(BaseModel):
content: str
meta: FunctionMeta
is_active: bool = False
is_global: bool = False
updated_at: int # timestamp in epoch
created_at: int # timestamp in epoch
@@ -66,6 +68,7 @@ class FunctionResponse(BaseModel):
name: str
meta: FunctionMeta
is_active: bool
is_global: bool
updated_at: int # timestamp in epoch
created_at: int # timestamp in epoch
@@ -144,6 +147,16 @@ class FunctionsTable:
for function in Function.select().where(Function.type == type)
]
def get_global_filter_functions(self) -> List[FunctionModel]:
return [
FunctionModel(**model_to_dict(function))
for function in Function.select().where(
Function.type == "filter",
Function.is_active == True,
Function.is_global == True,
)
]
def get_function_valves_by_id(self, id: str) -> Optional[dict]:
try:
function = Function.get(Function.id == id)

View File

@@ -147,6 +147,33 @@ async def toggle_function_by_id(id: str, user=Depends(get_admin_user)):
)
############################
# ToggleGlobalById
############################
@router.post("/id/{id}/toggle/global", response_model=Optional[FunctionModel])
async def toggle_global_by_id(id: str, user=Depends(get_admin_user)):
function = Functions.get_function_by_id(id)
if function:
function = Functions.update_function_by_id(
id, {"is_global": not function.is_global}
)
if function:
return function
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT("Error updating function"),
)
else:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.NOT_FOUND,
)
############################
# UpdateFunctionById
############################

View File

@@ -416,21 +416,23 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
)
return 0
filter_ids = []
filter_ids = [
function.id for function in Functions.get_global_filter_functions()
]
if "info" in model and "meta" in model["info"]:
enabled_filter_ids = [
function.id
for function in Functions.get_functions_by_type(
"filter", active_only=True
)
]
filter_ids = [
filter_id
for filter_id in enabled_filter_ids
if filter_id in model["info"]["meta"].get("filterIds", [])
]
filter_ids.extend(model["info"]["meta"].get("filterIds", []))
filter_ids = list(set(filter_ids))
enabled_filter_ids = [
function.id
for function in Functions.get_functions_by_type(
"filter", active_only=True
)
]
filter_ids = [
filter_id for filter_id in filter_ids if filter_id in enabled_filter_ids
]
filter_ids.sort(key=get_priority)
for filter_id in filter_ids:
filter = Functions.get_function_by_id(filter_id)
@@ -919,7 +921,6 @@ async def generate_chat_completions(form_data: dict, user=Depends(get_verified_u
)
model = app.state.MODELS[model_id]
print(model)
pipe = model.get("pipe")
if pipe:
@@ -1010,21 +1011,19 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
return (function.valves if function.valves else {}).get("priority", 0)
return 0
filter_ids = []
filter_ids = [function.id for function in Functions.get_global_filter_functions()]
if "info" in model and "meta" in model["info"]:
enabled_filter_ids = [
function.id
for function in Functions.get_functions_by_type(
"filter", active_only=True
)
]
filter_ids = [
filter_id
for filter_id in enabled_filter_ids
if filter_id in model["info"]["meta"].get("filterIds", [])
]
filter_ids.extend(model["info"]["meta"].get("filterIds", []))
filter_ids = list(set(filter_ids))
enabled_filter_ids = [
function.id
for function in Functions.get_functions_by_type("filter", active_only=True)
]
filter_ids = [
filter_id for filter_id in filter_ids if filter_id in enabled_filter_ids
]
# Sort filter_ids by priority, using the get_priority function
filter_ids.sort(key=get_priority)