UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-07 00:59:52 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-08-17 16:32:39 +02:00
parent fe747382c1
commit e71f55e58f
1 changed files with 28 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
backend/main.py
View File

@ -369,22 +369,22 @@ async def chat_completion_filter_functions_handler(body, model, extra_params):
return body, {} return body, {}
def get_tool_with_custom_params( def apply_extra_params_to_tool_function(
tool: Callable, custom_params: dict function: Callable, custom_params: dict
) -> Callable[..., Awaitable]: ) -> Callable[..., Awaitable]:
sig = inspect.signature(tool) sig = inspect.signature(function)
extra_params = { extra_params = {
key: value for key, value in custom_params.items() if key in sig.parameters key: value for key, value in custom_params.items() if key in sig.parameters
} }
is_coroutine = inspect.iscoroutinefunction(tool) is_coroutine = inspect.iscoroutinefunction(function)
async def new_tool(**kwargs): async def new_function(**kwargs):
extra_kwargs = kwargs | extra_params extra_kwargs = kwargs | extra_params
if is_coroutine: if is_coroutine:
return await tool(**extra_kwargs) return await function(**extra_kwargs)
return tool(**extra_kwargs) return function(**extra_kwargs)
return new_tool return new_function
# Mutation on extra_params # Mutation on extra_params
@ -403,6 +403,7 @@ def get_tools(
webui_app.state.TOOLS[tool_id] = module webui_app.state.TOOLS[tool_id] = module
extra_params["__id__"] = tool_id extra_params["__id__"] = tool_id
has_citation = hasattr(module, "citation") and module.citation has_citation = hasattr(module, "citation") and module.citation
handles_files = hasattr(module, "file_handler") and module.file_handler handles_files = hasattr(module, "file_handler") and module.file_handler
@ -420,11 +421,12 @@ def get_tools(
for val in spec.get("parameters", {}).get("properties", {}).values(): for val in spec.get("parameters", {}).get("properties", {}).values():
if val["type"] == "str": if val["type"] == "str":
val["type"] = "string" val["type"] = "string"
name = spec["name"] function_name = spec["name"]
callable = getattr(module, name)
# convert to function that takes only model params and inserts custom params # convert to function that takes only model params and inserts custom params
custom_callable = get_tool_with_custom_params(callable, extra_params) callable = apply_extra_params_to_tool_function(
getattr(module, function_name), extra_params
)
# TODO: This needs to be a pydantic model # TODO: This needs to be a pydantic model
tool_dict = { tool_dict = {
@ -432,16 +434,16 @@ def get_tools(
"citation": has_citation, "citation": has_citation,
"file_handler": handles_files, "file_handler": handles_files,
"toolkit_id": tool_id, "toolkit_id": tool_id,
"callable": custom_callable, "callable": callable,
} }
# TODO: if collision, prepend toolkit name
if name in tools:
log.warning(f"Tool {name} already exists in another toolkit!")
log.warning(f"Collision between {toolkit} and {tool_id}.")
log.warning(f"Discarding {toolkit}.{name}")
else:
tools[name] = tool_dict
# TODO: if collision, prepend toolkit name
if function_name in tools:
log.warning(f"Tool {function_name} already exists in another toolkit!")
log.warning(f"Collision between {toolkit} and {tool_id}.")
log.warning(f"Discarding {toolkit}.{function_name}")
else:
tools[function_name] = tool_dict
return tools return tools
@ -487,8 +489,12 @@ async def chat_completion_tools_handler(
specs = [tool["spec"] for tool in tools.values()] specs = [tool["spec"] for tool in tools.values()]
tools_specs = json.dumps(specs) tools_specs = json.dumps(specs)
content = tool_calling_generation_template(app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE, tools_specs) content = tool_calling_generation_template(
payload = get_tools_function_calling_payload(body["messages"], task_model_id, content) app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE, tools_specs
)
payload = get_tools_function_calling_payload(
body["messages"], task_model_id, content
)
try: try:
payload = filter_pipeline(payload, user) payload = filter_pipeline(payload, user)
@ -501,7 +507,6 @@ async def chat_completion_tools_handler(
content = await get_content_from_response(response) content = await get_content_from_response(response)
log.debug(f"{content=}") log.debug(f"{content=}")
if content is None: if content is None:
return body, {} return body, {}
@ -512,7 +517,7 @@ async def chat_completion_tools_handler(
tool_params = result.get("parameters", {}) tool_params = result.get("parameters", {})
toolkit_id = tools[tool_name]["toolkit_id"] toolkit_id = tools[tool_name]["toolkit_id"]
try: try:
tool_output = await tools[tool_name]["callable"](**tool_params) tool_output = await tools[tool_name]["callable"](**tool_params)
except Exception as e: except Exception as e: