refac
This commit is contained in:
Timothy Jaeryang Baek
@@ -57,17 +57,15 @@ log = logging.getLogger(__name__)
|
||||
def get_function_module_by_id(request: Request, pipe_id: str):
|
||||
function_module, _, _ = get_function_module_from_cache(request, pipe_id)
|
||||
|
||||
if hasattr(function_module, "valves") and hasattr(function_module, "Valves"):
|
||||
if hasattr(function_module, 'valves') and hasattr(function_module, 'Valves'):
|
||||
Valves = function_module.Valves
|
||||
valves = Functions.get_function_valves_by_id(pipe_id)
|
||||
|
||||
if valves:
|
||||
try:
|
||||
function_module.valves = Valves(
|
||||
**{k: v for k, v in valves.items() if v is not None}
|
||||
)
|
||||
function_module.valves = Valves(**{k: v for k, v in valves.items() if v is not None})
|
||||
except Exception as e:
|
||||
log.exception(f"Error loading valves for function {pipe_id}: {e}")
|
||||
log.exception(f'Error loading valves for function {pipe_id}: {e}')
|
||||
raise e
|
||||
else:
|
||||
function_module.valves = Valves()
|
||||
@@ -76,7 +74,7 @@ def get_function_module_by_id(request: Request, pipe_id: str):
|
||||
|
||||
|
||||
async def get_function_models(request):
|
||||
pipes = Functions.get_functions_by_type("pipe", active_only=True)
|
||||
pipes = Functions.get_functions_by_type('pipe', active_only=True)
|
||||
pipe_models = []
|
||||
|
||||
for pipe in pipes:
|
||||
@@ -84,11 +82,11 @@ async def get_function_models(request):
|
||||
function_module = get_function_module_by_id(request, pipe.id)
|
||||
|
||||
has_user_valves = False
|
||||
if hasattr(function_module, "UserValves"):
|
||||
if hasattr(function_module, 'UserValves'):
|
||||
has_user_valves = True
|
||||
|
||||
# Check if function is a manifold
|
||||
if hasattr(function_module, "pipes"):
|
||||
if hasattr(function_module, 'pipes'):
|
||||
sub_pipes = []
|
||||
|
||||
# Handle pipes being a list, sync function, or async function
|
||||
@@ -104,32 +102,30 @@ async def get_function_models(request):
|
||||
log.exception(e)
|
||||
sub_pipes = []
|
||||
|
||||
log.debug(
|
||||
f"get_function_models: function '{pipe.id}' is a manifold of {sub_pipes}"
|
||||
)
|
||||
log.debug(f"get_function_models: function '{pipe.id}' is a manifold of {sub_pipes}")
|
||||
|
||||
for p in sub_pipes:
|
||||
sub_pipe_id = f'{pipe.id}.{p["id"]}'
|
||||
sub_pipe_name = p["name"]
|
||||
sub_pipe_name = p['name']
|
||||
|
||||
if hasattr(function_module, "name"):
|
||||
sub_pipe_name = f"{function_module.name}{sub_pipe_name}"
|
||||
if hasattr(function_module, 'name'):
|
||||
sub_pipe_name = f'{function_module.name}{sub_pipe_name}'
|
||||
|
||||
pipe_flag = {"type": pipe.type}
|
||||
pipe_flag = {'type': pipe.type}
|
||||
|
||||
pipe_models.append(
|
||||
{
|
||||
"id": sub_pipe_id,
|
||||
"name": sub_pipe_name,
|
||||
"object": "model",
|
||||
"created": pipe.created_at,
|
||||
"owned_by": "openai",
|
||||
"pipe": pipe_flag,
|
||||
"has_user_valves": has_user_valves,
|
||||
'id': sub_pipe_id,
|
||||
'name': sub_pipe_name,
|
||||
'object': 'model',
|
||||
'created': pipe.created_at,
|
||||
'owned_by': 'openai',
|
||||
'pipe': pipe_flag,
|
||||
'has_user_valves': has_user_valves,
|
||||
}
|
||||
)
|
||||
else:
|
||||
pipe_flag = {"type": "pipe"}
|
||||
pipe_flag = {'type': 'pipe'}
|
||||
|
||||
log.debug(
|
||||
f"get_function_models: function '{pipe.id}' is a single pipe {{ 'id': {pipe.id}, 'name': {pipe.name} }}"
|
||||
@@ -137,13 +133,13 @@ async def get_function_models(request):
|
||||
|
||||
pipe_models.append(
|
||||
{
|
||||
"id": pipe.id,
|
||||
"name": pipe.name,
|
||||
"object": "model",
|
||||
"created": pipe.created_at,
|
||||
"owned_by": "openai",
|
||||
"pipe": pipe_flag,
|
||||
"has_user_valves": has_user_valves,
|
||||
'id': pipe.id,
|
||||
'name': pipe.name,
|
||||
'object': 'model',
|
||||
'created': pipe.created_at,
|
||||
'owned_by': 'openai',
|
||||
'pipe': pipe_flag,
|
||||
'has_user_valves': has_user_valves,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -153,9 +149,7 @@ async def get_function_models(request):
|
||||
return pipe_models
|
||||
|
||||
|
||||
async def generate_function_chat_completion(
|
||||
request, form_data, user, models: dict = {}
|
||||
):
|
||||
async def generate_function_chat_completion(request, form_data, user, models: dict = {}):
|
||||
async def execute_pipe(pipe, params):
|
||||
if inspect.iscoroutinefunction(pipe):
|
||||
return await pipe(**params)
|
||||
@@ -166,32 +160,32 @@ async def generate_function_chat_completion(
|
||||
if isinstance(res, str):
|
||||
return res
|
||||
if isinstance(res, Generator):
|
||||
return "".join(map(str, res))
|
||||
return ''.join(map(str, res))
|
||||
if isinstance(res, AsyncGenerator):
|
||||
return "".join([str(stream) async for stream in res])
|
||||
return ''.join([str(stream) async for stream in res])
|
||||
|
||||
def process_line(form_data: dict, line):
|
||||
if isinstance(line, BaseModel):
|
||||
line = line.model_dump_json()
|
||||
line = f"data: {line}"
|
||||
line = f'data: {line}'
|
||||
if isinstance(line, dict):
|
||||
line = f"data: {json.dumps(line)}"
|
||||
line = f'data: {json.dumps(line)}'
|
||||
|
||||
try:
|
||||
line = line.decode("utf-8")
|
||||
line = line.decode('utf-8')
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if line.startswith("data:"):
|
||||
return f"{line}\n\n"
|
||||
if line.startswith('data:'):
|
||||
return f'{line}\n\n'
|
||||
else:
|
||||
line = openai_chat_chunk_message_template(form_data["model"], line)
|
||||
return f"data: {json.dumps(line)}\n\n"
|
||||
line = openai_chat_chunk_message_template(form_data['model'], line)
|
||||
return f'data: {json.dumps(line)}\n\n'
|
||||
|
||||
def get_pipe_id(form_data: dict) -> str:
|
||||
pipe_id = form_data["model"]
|
||||
if "." in pipe_id:
|
||||
pipe_id, _ = pipe_id.split(".", 1)
|
||||
pipe_id = form_data['model']
|
||||
if '.' in pipe_id:
|
||||
pipe_id, _ = pipe_id.split('.', 1)
|
||||
return pipe_id
|
||||
|
||||
def get_function_params(function_module, form_data, user, extra_params=None):
|
||||
@@ -202,27 +196,25 @@ async def generate_function_chat_completion(
|
||||
|
||||
# Get the signature of the function
|
||||
sig = inspect.signature(function_module.pipe)
|
||||
params = {"body": form_data} | {
|
||||
k: v for k, v in extra_params.items() if k in sig.parameters
|
||||
}
|
||||
params = {'body': form_data} | {k: v for k, v in extra_params.items() if k in sig.parameters}
|
||||
|
||||
if "__user__" in params and hasattr(function_module, "UserValves"):
|
||||
if '__user__' in params and hasattr(function_module, 'UserValves'):
|
||||
user_valves = Functions.get_user_valves_by_id_and_user_id(pipe_id, user.id)
|
||||
try:
|
||||
params["__user__"]["valves"] = function_module.UserValves(**user_valves)
|
||||
params['__user__']['valves'] = function_module.UserValves(**user_valves)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
params["__user__"]["valves"] = function_module.UserValves()
|
||||
params['__user__']['valves'] = function_module.UserValves()
|
||||
|
||||
return params
|
||||
|
||||
model_id = form_data.get("model")
|
||||
model_id = form_data.get('model')
|
||||
model_info = Models.get_model_by_id(model_id)
|
||||
|
||||
metadata = form_data.pop("metadata", {})
|
||||
metadata = form_data.pop('metadata', {})
|
||||
|
||||
files = metadata.get("files", [])
|
||||
tool_ids = metadata.get("tool_ids", [])
|
||||
files = metadata.get('files', [])
|
||||
tool_ids = metadata.get('tool_ids', [])
|
||||
# Check if tool_ids is None
|
||||
if tool_ids is None:
|
||||
tool_ids = []
|
||||
@@ -233,56 +225,56 @@ async def generate_function_chat_completion(
|
||||
__task_body__ = None
|
||||
|
||||
if metadata:
|
||||
if all(k in metadata for k in ("session_id", "chat_id", "message_id")):
|
||||
if all(k in metadata for k in ('session_id', 'chat_id', 'message_id')):
|
||||
__event_emitter__ = get_event_emitter(metadata)
|
||||
__event_call__ = get_event_call(metadata)
|
||||
__task__ = metadata.get("task", None)
|
||||
__task_body__ = metadata.get("task_body", None)
|
||||
__task__ = metadata.get('task', None)
|
||||
__task_body__ = metadata.get('task_body', None)
|
||||
|
||||
oauth_token = None
|
||||
try:
|
||||
if request.cookies.get("oauth_session_id", None):
|
||||
if request.cookies.get('oauth_session_id', None):
|
||||
oauth_token = await request.app.state.oauth_manager.get_oauth_token(
|
||||
user.id,
|
||||
request.cookies.get("oauth_session_id", None),
|
||||
request.cookies.get('oauth_session_id', None),
|
||||
)
|
||||
except Exception as e:
|
||||
log.error(f"Error getting OAuth token: {e}")
|
||||
log.error(f'Error getting OAuth token: {e}')
|
||||
|
||||
extra_params = {
|
||||
"__event_emitter__": __event_emitter__,
|
||||
"__event_call__": __event_call__,
|
||||
"__chat_id__": metadata.get("chat_id", None),
|
||||
"__session_id__": metadata.get("session_id", None),
|
||||
"__message_id__": metadata.get("message_id", None),
|
||||
"__task__": __task__,
|
||||
"__task_body__": __task_body__,
|
||||
"__files__": files,
|
||||
"__user__": user.model_dump() if isinstance(user, UserModel) else {},
|
||||
"__metadata__": metadata,
|
||||
"__oauth_token__": oauth_token,
|
||||
"__request__": request,
|
||||
'__event_emitter__': __event_emitter__,
|
||||
'__event_call__': __event_call__,
|
||||
'__chat_id__': metadata.get('chat_id', None),
|
||||
'__session_id__': metadata.get('session_id', None),
|
||||
'__message_id__': metadata.get('message_id', None),
|
||||
'__task__': __task__,
|
||||
'__task_body__': __task_body__,
|
||||
'__files__': files,
|
||||
'__user__': user.model_dump() if isinstance(user, UserModel) else {},
|
||||
'__metadata__': metadata,
|
||||
'__oauth_token__': oauth_token,
|
||||
'__request__': request,
|
||||
}
|
||||
extra_params["__tools__"] = await get_tools(
|
||||
extra_params['__tools__'] = await get_tools(
|
||||
request,
|
||||
tool_ids,
|
||||
user,
|
||||
{
|
||||
**extra_params,
|
||||
"__model__": models.get(form_data["model"], None),
|
||||
"__messages__": form_data["messages"],
|
||||
"__files__": files,
|
||||
'__model__': models.get(form_data['model'], None),
|
||||
'__messages__': form_data['messages'],
|
||||
'__files__': files,
|
||||
},
|
||||
)
|
||||
|
||||
if model_info:
|
||||
if model_info.base_model_id:
|
||||
form_data["model"] = model_info.base_model_id
|
||||
form_data['model'] = model_info.base_model_id
|
||||
|
||||
params = model_info.params.model_dump()
|
||||
|
||||
if params:
|
||||
system = params.pop("system", None)
|
||||
system = params.pop('system', None)
|
||||
form_data = apply_model_params_to_body_openai(params, form_data)
|
||||
form_data = apply_system_prompt_to_body(system, form_data, metadata, user)
|
||||
|
||||
@@ -292,7 +284,7 @@ async def generate_function_chat_completion(
|
||||
pipe = function_module.pipe
|
||||
params = get_function_params(function_module, form_data, user, extra_params)
|
||||
|
||||
if form_data.get("stream", False):
|
||||
if form_data.get('stream', False):
|
||||
|
||||
async def stream_content():
|
||||
try:
|
||||
@@ -304,17 +296,17 @@ async def generate_function_chat_completion(
|
||||
yield data
|
||||
return
|
||||
if isinstance(res, dict):
|
||||
yield f"data: {json.dumps(res)}\n\n"
|
||||
yield f'data: {json.dumps(res)}\n\n'
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
log.error(f"Error: {e}")
|
||||
yield f"data: {json.dumps({'error': {'detail':str(e)}})}\n\n"
|
||||
log.error(f'Error: {e}')
|
||||
yield f'data: {json.dumps({"error": {"detail": str(e)}})}\n\n'
|
||||
return
|
||||
|
||||
if isinstance(res, str):
|
||||
message = openai_chat_chunk_message_template(form_data["model"], res)
|
||||
yield f"data: {json.dumps(message)}\n\n"
|
||||
message = openai_chat_chunk_message_template(form_data['model'], res)
|
||||
yield f'data: {json.dumps(message)}\n\n'
|
||||
|
||||
if isinstance(res, Iterator):
|
||||
for line in res:
|
||||
@@ -325,21 +317,19 @@ async def generate_function_chat_completion(
|
||||
yield process_line(form_data, line)
|
||||
|
||||
if isinstance(res, str) or isinstance(res, Generator):
|
||||
finish_message = openai_chat_chunk_message_template(
|
||||
form_data["model"], ""
|
||||
)
|
||||
finish_message["choices"][0]["finish_reason"] = "stop"
|
||||
yield f"data: {json.dumps(finish_message)}\n\n"
|
||||
yield "data: [DONE]"
|
||||
finish_message = openai_chat_chunk_message_template(form_data['model'], '')
|
||||
finish_message['choices'][0]['finish_reason'] = 'stop'
|
||||
yield f'data: {json.dumps(finish_message)}\n\n'
|
||||
yield 'data: [DONE]'
|
||||
|
||||
return StreamingResponse(stream_content(), media_type="text/event-stream")
|
||||
return StreamingResponse(stream_content(), media_type='text/event-stream')
|
||||
else:
|
||||
try:
|
||||
res = await execute_pipe(pipe, params)
|
||||
|
||||
except Exception as e:
|
||||
log.error(f"Error: {e}")
|
||||
return {"error": {"detail": str(e)}}
|
||||
log.error(f'Error: {e}')
|
||||
return {'error': {'detail': str(e)}}
|
||||
|
||||
if isinstance(res, StreamingResponse) or isinstance(res, dict):
|
||||
return res
|
||||
@@ -347,4 +337,4 @@ async def generate_function_chat_completion(
|
||||
return res.model_dump()
|
||||
|
||||
message = await get_message_content(res)
|
||||
return openai_chat_completion_message_template(form_data["model"], message)
|
||||
return openai_chat_completion_message_template(form_data['model'], message)
|
||||
|
||||
Reference in New Issue
Block a user