refactor: refactor server as a package
This commit is contained in:
Shivering Ocelot committed
1 parent
75e49fcbcd
commit
e1063fc733
5 files changed
+374
-363
No files matched your search
@@ -0,0 +1 @@
|
||||
from server.start import start_server
|
||||
@@ -0,0 +1,260 @@
|
||||
import json, time, traceback, asyncio
|
||||
from quart import Response, request
|
||||
from quart_auth import login_required, current_user
|
||||
from tool import ToolSpec, ToolParameter
|
||||
from agent import AgentRun
|
||||
from utils import encode_json
|
||||
|
||||
def register(server, app):
|
||||
# https://github.com/hrgarber/openai-compatible-api/blob/main/app/main.py
|
||||
@server.post('/api/openai/v1/chat/completions')
|
||||
@login_required
|
||||
async def openai_chat_completion():
|
||||
data = await request.get_json()
|
||||
|
||||
print("Completion: " + json.dumps(data))
|
||||
|
||||
# Extract required fields
|
||||
agent_id = data.get("model")
|
||||
openai_messages = data.get("messages", [])
|
||||
openai_tools = data.get("tools", [])
|
||||
stream_mode = data.get("stream", False)
|
||||
|
||||
if not agent_id or not openai_messages:
|
||||
return {"error": "model and messages are required"}, 400
|
||||
|
||||
user_id = current_user.auth_id
|
||||
if not user_id: pass
|
||||
|
||||
# --- Parse messages ---
|
||||
message_errors = []
|
||||
messages = []
|
||||
for msg in openai_messages:
|
||||
role = msg['role']
|
||||
if role not in ['user', 'assistant', 'system']:
|
||||
message_errors.append("Bad role: " + role)
|
||||
continue
|
||||
content = msg.get('content', None)
|
||||
|
||||
if type(content) == list:
|
||||
parts = content
|
||||
content = ""
|
||||
for part in parts:
|
||||
if part['type'] == 'text':
|
||||
content += part['text']
|
||||
|
||||
if type(content) != str:
|
||||
message_errors.append("Bad content")
|
||||
continue
|
||||
|
||||
messages.append({'role': role, 'content': content})
|
||||
|
||||
if message_errors:
|
||||
print("[openai] ERROR: bad messages input\n" + "\n".join([f"\t- {e}" for e in message_errors]))
|
||||
|
||||
# --- Parse tools (OpenAI format -> ToolSpec) ---
|
||||
tools = []
|
||||
for tool in openai_tools:
|
||||
if tool.get('type') != 'function':
|
||||
continue
|
||||
func = tool['function']
|
||||
params_schema = func.get('parameters', {})
|
||||
required_params = params_schema.get('required', [])
|
||||
parameters = []
|
||||
for param_name, param_spec in params_schema.get('properties', {}).items():
|
||||
parameters.append(ToolParameter(
|
||||
name=param_name,
|
||||
type=param_spec.get('type', 'string'),
|
||||
description=param_spec.get('description', ''),
|
||||
required=param_name in required_params,
|
||||
enum=param_spec.get('enum'),
|
||||
))
|
||||
tools.append(ToolSpec(
|
||||
name=func['name'],
|
||||
description=func.get('description', ''),
|
||||
parameters=parameters,
|
||||
))
|
||||
|
||||
# --- Get agent control and create AgentRun directly ---
|
||||
control = await app.get_agent_control(agent_id)
|
||||
|
||||
run = AgentRun.create(
|
||||
app,
|
||||
agent_id=agent_id,
|
||||
user_id=user_id,
|
||||
input={
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
},
|
||||
)
|
||||
print("run input: " + repr(run.input))
|
||||
|
||||
# --- Streaming ---
|
||||
if stream_mode:
|
||||
queue = asyncio.Queue()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
# Listen to the stream for both content and tool_call chunks
|
||||
|
||||
def on_output(stream):
|
||||
async def on_chunk(chunk):
|
||||
await queue.put(('content', chunk))
|
||||
async def on_eof():
|
||||
# on_done doesn't work?
|
||||
await queue.put(('eof', None))
|
||||
stream.listen(on_chunk, on_eof)
|
||||
|
||||
def on_reasoning(stream):
|
||||
async def on_chunk(chunk):
|
||||
await queue.put(('reasoning', chunk))
|
||||
async def on_eof(): pass
|
||||
stream.listen(on_chunk, on_eof)
|
||||
|
||||
def on_tool_call(tool_call):
|
||||
print('agent tool call ' + repr(tool_call))
|
||||
loop.create_task(queue.put(('tool_call', tool_call)))
|
||||
|
||||
def on_done(output):
|
||||
print('agent done ' + repr(output))
|
||||
loop.create_task(queue.put(('eof', None)))
|
||||
|
||||
run.on('output', on_output)
|
||||
run.on('reasoning', on_reasoning)
|
||||
run.on('tool-call', on_tool_call)
|
||||
run.on('done', on_done)
|
||||
|
||||
# Start the agent run
|
||||
await control.run(run)
|
||||
|
||||
async def generate():
|
||||
try:
|
||||
tool_call_count = 0
|
||||
while True:
|
||||
msg_type, data = await queue.get()
|
||||
|
||||
delta = None
|
||||
if msg_type == 'content':
|
||||
delta = {'content': data}
|
||||
elif msg_type == 'reasoning':
|
||||
delta = {'reasoning': data}
|
||||
|
||||
elif msg_type == 'tool_call':
|
||||
tool_call_count = tool_call_count + 1
|
||||
arguments = encode_json(data.arguments) if data.arguments else '{}'
|
||||
tc_entry = {"index": tool_call_count, "function": {'name': data.tool_name, 'arguments': arguments}}
|
||||
if data.id:
|
||||
tc_entry['id'] = data.id
|
||||
delta = {"tool_calls": [tc_entry]}
|
||||
|
||||
elif msg_type == 'eof':
|
||||
# Determine finish_reason — was there a tool call?
|
||||
finish_reason = "tool_calls" if any(
|
||||
e['name'] == 'tool_call' for e in run.events
|
||||
) else "stop"
|
||||
|
||||
chunk_body = {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{"delta": {}, "index": 0, "finish_reason": finish_reason}]
|
||||
}
|
||||
yield f"data: {json.dumps(chunk_body)}\n\n"
|
||||
break
|
||||
|
||||
if delta:
|
||||
chunk_body = {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{"delta": delta, "index": 0, "finish_reason": None}]
|
||||
}
|
||||
yield f"data: {json.dumps(chunk_body)}\n\n"
|
||||
else:
|
||||
raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
|
||||
|
||||
except Exception as e:
|
||||
print("[openai] ERROR: " + str(e))
|
||||
print(traceback.format_exc())
|
||||
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
||||
|
||||
return Response(generate(), content_type="text/event-stream")
|
||||
|
||||
# --- Non-streaming ---
|
||||
else:
|
||||
await control.run(run)
|
||||
await run.get_output_stream().content_wait()
|
||||
|
||||
# Collect any tool calls that were emitted
|
||||
tool_events = [e['data'] for e in run.events if e['name'] == 'tool_call']
|
||||
|
||||
if tool_events:
|
||||
# Aggregate tool calls: combine partial argument chunks by id
|
||||
tc_map = {}
|
||||
for evt in tool_events:
|
||||
tid = evt.get('id', '')
|
||||
if tid not in tc_map:
|
||||
tc_map[tid] = {
|
||||
"id": tid,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": evt.get('function', {}).get('name', ''),
|
||||
"arguments": evt.get('function', {}).get('arguments', ''),
|
||||
}
|
||||
}
|
||||
else:
|
||||
tc_map[tid]['function']['arguments'] += \
|
||||
evt.get('function', {}).get('arguments', '')
|
||||
|
||||
return {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": list(tc_map.values()),
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
}]
|
||||
}
|
||||
else:
|
||||
# No tool calls: return text content
|
||||
content = run.stream.content
|
||||
return {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}]
|
||||
}
|
||||
|
||||
@server.get("/api/openai/v1/models")
|
||||
@login_required
|
||||
async def openai_models():
|
||||
"""
|
||||
Return a list of available models.
|
||||
Customize this list based on your available models.
|
||||
"""
|
||||
|
||||
user_id = current_user.auth_id
|
||||
if not user_id: pass
|
||||
|
||||
return {
|
||||
"object": "list",
|
||||
"data": [{
|
||||
"id": agent.id,
|
||||
"object": "model",
|
||||
"created": int(time.time()),
|
||||
"owned_by": user_id,
|
||||
} for agent in await app.get_agents(user_id)]
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
from quart import Quart, request, send_from_directory
|
||||
from quart_auth import QuartAuth, AuthUser, login_user
|
||||
from server.websocket import register as register_websocket
|
||||
from server.openai import register as register_openai
|
||||
|
||||
# Vue frontend
|
||||
# relative to the module file???? Quart being unserious
|
||||
frontend_path = '../../../frontend/dist'
|
||||
|
||||
async def start_server(app, port, shutdown_event):
|
||||
server = Quart(__name__, static_folder=frontend_path)
|
||||
server.secret_key = 'horde-secret'
|
||||
auth_manager = QuartAuth(server, mode="bearer")
|
||||
|
||||
stream_listeners = {}
|
||||
|
||||
def cleanup(self):
|
||||
stream_listeners = []
|
||||
|
||||
@server.after_request
|
||||
def after_request(response):
|
||||
response.headers.add('Access-Control-Allow-Origin', '*')
|
||||
response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
|
||||
response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
|
||||
return response
|
||||
|
||||
@server.route('/')
|
||||
async def index():
|
||||
return await send_from_directory(server.static_folder, 'index.html')
|
||||
|
||||
# These paths must match the routes in the frontend SPA
|
||||
# file: frontend/src/routes/index.js
|
||||
@server.route('/user')
|
||||
@server.route('/models')
|
||||
@server.route('/remotes')
|
||||
@server.route('/runs')
|
||||
@server.route('/chat')
|
||||
@server.route('/agent')
|
||||
@server.route('/chat/<path:subpath>')
|
||||
@server.route('/agent/<path:subpath>')
|
||||
async def serve_spa_routes(subpath=None):
|
||||
# Ignore subpath and render index, Vue router will take care of everything
|
||||
return await send_from_directory(server.static_folder, 'index.html')
|
||||
|
||||
@server.route('/<path:path>')
|
||||
async def serve_static(path):
|
||||
return await send_from_directory(server.static_folder, path)
|
||||
|
||||
@server.route('/user/login', methods=['POST'])
|
||||
async def login():
|
||||
data = await request.get_json()
|
||||
|
||||
# token has precedence if it's present
|
||||
if "token" in data:
|
||||
user_id = auth_manager.load_token(data['token'])
|
||||
elif "username" in data and "password" in data:
|
||||
user_id = await app.verify_user_credentials(data['username'], data['password'])
|
||||
else:
|
||||
return "invalid login arguments", 400
|
||||
|
||||
if user_id == None:
|
||||
return "bad login credentials", 401
|
||||
|
||||
token = auth_manager.dump_token(user_id)
|
||||
login_user(AuthUser(auth_id=user_id))
|
||||
return {"token": token}
|
||||
|
||||
@server.route('/user/signup', methods=['POST'])
|
||||
async def signup():
|
||||
data = await request.get_json()
|
||||
|
||||
user_id = None
|
||||
|
||||
username = data['username']
|
||||
email = data['email']
|
||||
password = data['password']
|
||||
|
||||
user_id = await app.create_new_user(username, email, password)
|
||||
if user_id is None:
|
||||
return "existing or incorrect credentials", 400
|
||||
|
||||
await app.generate_default_user_data(user_id)
|
||||
|
||||
token = auth_manager.dump_token(user_id)
|
||||
return {"token": token}
|
||||
|
||||
register_websocket(server, app, auth_manager)
|
||||
register_openai(server, app)
|
||||
|
||||
await server.run_task(debug=True, shutdown_trigger=shutdown_event.wait, port=port)
|
||||
@@ -1,114 +1,15 @@
|
||||
from quart import Quart, Response, request, send_from_directory, websocket
|
||||
from quart_auth import (
|
||||
QuartAuth, current_user,
|
||||
AuthUser, login_user, logout_user, login_required
|
||||
)
|
||||
import time
|
||||
import json
|
||||
import asyncio
|
||||
import traceback
|
||||
from stream import Stream
|
||||
from tool import ToolSpec, ToolParameter
|
||||
from agent import AgentRun
|
||||
|
||||
# Vue frontend
|
||||
# relative to the module file???? wtf
|
||||
frontend_path = '../../frontend/dist'
|
||||
|
||||
def encode_json(value):
|
||||
def encode_unknown(value):
|
||||
if isinstance(value, ToolSpec):
|
||||
return value.to_openai_dict()
|
||||
elif isinstance(value, Stream):
|
||||
return {'streamId': value.id}
|
||||
else:
|
||||
return repr(value)
|
||||
return json.dumps(value, default=encode_unknown)
|
||||
|
||||
async def start_server(app, port, shutdown_event):
|
||||
server = Quart(__name__, static_folder=frontend_path)
|
||||
server.secret_key = 'horde-secret'
|
||||
auth_manager = QuartAuth(server, mode="bearer")
|
||||
|
||||
stream_listeners = {}
|
||||
|
||||
def cleanup(self):
|
||||
stream_listeners = []
|
||||
|
||||
@server.after_request
|
||||
def after_request(response):
|
||||
response.headers.add('Access-Control-Allow-Origin', '*')
|
||||
response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
|
||||
response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
|
||||
return response
|
||||
|
||||
@server.route('/')
|
||||
async def index():
|
||||
return await send_from_directory(server.static_folder, 'index.html')
|
||||
|
||||
# These paths must match the routes in the frontend SPA
|
||||
# file: frontend/src/routes/index.js
|
||||
@server.route('/user')
|
||||
@server.route('/models')
|
||||
@server.route('/remotes')
|
||||
@server.route('/runs')
|
||||
@server.route('/chat')
|
||||
@server.route('/agent')
|
||||
@server.route('/chat/<path:subpath>')
|
||||
@server.route('/agent/<path:subpath>')
|
||||
async def serve_spa_routes(subpath=None):
|
||||
# Ignore subpath and render index, Vue router will take care of everything
|
||||
return await send_from_directory(server.static_folder, 'index.html')
|
||||
|
||||
@server.route('/<path:path>')
|
||||
async def serve_static(path):
|
||||
return await send_from_directory(server.static_folder, path)
|
||||
|
||||
@server.route('/user/login', methods=['POST'])
|
||||
async def login():
|
||||
data = await request.get_json()
|
||||
|
||||
# token has precedence if it's present
|
||||
if "token" in data:
|
||||
user_id = auth_manager.load_token(data['token'])
|
||||
elif "username" in data and "password" in data:
|
||||
user_id = await app.verify_user_credentials(data['username'], data['password'])
|
||||
else:
|
||||
return "invalid login arguments", 400
|
||||
|
||||
if user_id == None:
|
||||
return "bad login credentials", 401
|
||||
|
||||
token = auth_manager.dump_token(user_id)
|
||||
login_user(AuthUser(auth_id=user_id))
|
||||
return {"token": token}
|
||||
|
||||
@server.route('/user/signup', methods=['POST'])
|
||||
async def signup():
|
||||
data = await request.get_json()
|
||||
|
||||
user_id = None
|
||||
|
||||
username = data['username']
|
||||
email = data['email']
|
||||
password = data['password']
|
||||
|
||||
user_id = await app.create_new_user(username, email, password)
|
||||
if user_id is None:
|
||||
return "existing or incorrect credentials", 400
|
||||
|
||||
await app.generate_default_user_data(user_id)
|
||||
|
||||
token = auth_manager.dump_token(user_id)
|
||||
return {"token": token}
|
||||
import json, asyncio
|
||||
from quart import websocket
|
||||
from utils import encode_json
|
||||
from stream import Stream # isinstance check in send-agent-prompt
|
||||
|
||||
def register(server, app, auth_manager):
|
||||
@server.websocket('/ws')
|
||||
async def ws():
|
||||
print("incoming ws")
|
||||
|
||||
# Authenticate before dedicating any resources
|
||||
auth_req = json.loads(await websocket.receive())
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
if 'token' in auth_req:
|
||||
try:
|
||||
@@ -146,7 +47,7 @@ async def start_server(app, port, shutdown_event):
|
||||
async def send(msg):
|
||||
await send_queue.put(msg)
|
||||
|
||||
asyncio.create_task(sender_impl())
|
||||
sender_task = asyncio.create_task(sender_impl())
|
||||
ws_task = asyncio.current_task()
|
||||
app.ws_tasks.append(ws_task)
|
||||
|
||||
@@ -161,7 +62,7 @@ async def start_server(app, port, shutdown_event):
|
||||
# Listen for events and broadcast them
|
||||
async def on_event(name, data):
|
||||
await send({"action": name, "data": data})
|
||||
app.on_event(on_event)
|
||||
stop_listening = app.on_event(on_event)
|
||||
|
||||
while True:
|
||||
req = json.loads(await websocket.receive())
|
||||
@@ -226,7 +127,7 @@ async def start_server(app, port, shutdown_event):
|
||||
)
|
||||
if model:
|
||||
await send({'message': "model saved", 'modelId': model.id, 'nonce': nonce})
|
||||
elif modelId in body:
|
||||
elif "modelId" in body:
|
||||
await send({'message': "model deleted", 'modelId': body["modelId"], 'nonce': nonce})
|
||||
|
||||
elif action == "get-agent-config-fields":
|
||||
@@ -396,6 +297,7 @@ async def start_server(app, port, shutdown_event):
|
||||
await send({'error': 'unknown action: ' + action, 'nonce': nonce})
|
||||
finally:
|
||||
print("User logged out:", user_id)
|
||||
sender_task.cancel()
|
||||
if ws_task in app.ws_tasks:
|
||||
app.ws_tasks.remove(ws_task)
|
||||
# Unregister the event listener: the send closure would
|
||||
@@ -408,259 +310,3 @@ async def start_server(app, port, shutdown_event):
|
||||
for remotename in list(session_connections):
|
||||
app.disconnectRemote(user_id, remotename)
|
||||
app.remove_user_connection(user_id, send_queue)
|
||||
|
||||
# https://github.com/hrgarber/openai-compatible-api/blob/main/app/main.py
|
||||
@server.post('/api/openai/v1/chat/completions')
|
||||
@login_required
|
||||
async def openai_chat_completion():
|
||||
data = await request.get_json()
|
||||
|
||||
print("Completion: " + json.dumps(data))
|
||||
|
||||
# Extract required fields
|
||||
agent_id = data.get("model")
|
||||
openai_messages = data.get("messages", [])
|
||||
openai_tools = data.get("tools", [])
|
||||
stream_mode = data.get("stream", False)
|
||||
|
||||
if not agent_id or not openai_messages:
|
||||
return {"error": "model and messages are required"}, 400
|
||||
|
||||
user_id = current_user.auth_id
|
||||
if not user_id: pass
|
||||
|
||||
# --- Parse messages ---
|
||||
message_errors = []
|
||||
messages = []
|
||||
for msg in openai_messages:
|
||||
role = msg['role']
|
||||
if role not in ['user', 'assistant', 'system']:
|
||||
message_errors.append("Bad role: " + role)
|
||||
continue
|
||||
content = msg.get('content', None)
|
||||
|
||||
if type(content) == list:
|
||||
parts = content
|
||||
content = ""
|
||||
for part in parts:
|
||||
if part['type'] == 'text':
|
||||
content += part['text']
|
||||
|
||||
if type(content) != str:
|
||||
message_errors.append("Bad content")
|
||||
continue
|
||||
|
||||
messages.append({'role': role, 'content': content})
|
||||
|
||||
if message_errors:
|
||||
print("[openai] ERROR: bad messages input\n" + "\n".join([f"\t- {e}" for e in message_errors]))
|
||||
|
||||
# --- Parse tools (OpenAI format -> ToolSpec) ---
|
||||
tools = []
|
||||
for tool in openai_tools:
|
||||
if tool.get('type') != 'function':
|
||||
continue
|
||||
func = tool['function']
|
||||
params_schema = func.get('parameters', {})
|
||||
required_params = params_schema.get('required', [])
|
||||
parameters = []
|
||||
for param_name, param_spec in params_schema.get('properties', {}).items():
|
||||
parameters.append(ToolParameter(
|
||||
name=param_name,
|
||||
type=param_spec.get('type', 'string'),
|
||||
description=param_spec.get('description', ''),
|
||||
required=param_name in required_params,
|
||||
enum=param_spec.get('enum'),
|
||||
))
|
||||
tools.append(ToolSpec(
|
||||
name=func['name'],
|
||||
description=func.get('description', ''),
|
||||
parameters=parameters,
|
||||
))
|
||||
|
||||
# --- Get agent control and create AgentRun directly ---
|
||||
control = await app.get_agent_control(agent_id)
|
||||
|
||||
run = AgentRun.create(
|
||||
app,
|
||||
agent_id=agent_id,
|
||||
user_id=user_id,
|
||||
input={
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
},
|
||||
)
|
||||
print("run input: " + repr(run.input))
|
||||
|
||||
# --- Streaming ---
|
||||
if stream_mode:
|
||||
queue = asyncio.Queue()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
# Listen to the stream for both content and tool_call chunks
|
||||
|
||||
def on_output(stream):
|
||||
async def on_chunk(chunk):
|
||||
await queue.put(('content', chunk))
|
||||
async def on_eof():
|
||||
# on_done doesn't work?
|
||||
await queue.put(('eof', None))
|
||||
stream.listen(on_chunk, on_eof)
|
||||
|
||||
def on_reasoning(stream):
|
||||
async def on_chunk(chunk):
|
||||
await queue.put(('reasoning', chunk))
|
||||
async def on_eof(): pass
|
||||
stream.listen(on_chunk, on_eof)
|
||||
|
||||
def on_tool_call(tool_call):
|
||||
print('agent tool call ' + repr(tool_call))
|
||||
loop.create_task(queue.put(('tool_call', tool_call)))
|
||||
|
||||
def on_done(output):
|
||||
print('agent done ' + repr(output))
|
||||
loop.create_task(queue.put(('eof', None)))
|
||||
|
||||
run.on('output', on_output)
|
||||
run.on('reasoning', on_reasoning)
|
||||
run.on('tool-call', on_tool_call)
|
||||
run.on('done', on_done)
|
||||
|
||||
# Start the agent run
|
||||
await control.run(run)
|
||||
|
||||
async def generate():
|
||||
try:
|
||||
tool_call_count = 0
|
||||
while True:
|
||||
msg_type, data = await queue.get()
|
||||
|
||||
delta = None
|
||||
if msg_type == 'content':
|
||||
delta = {'content': data}
|
||||
elif msg_type == 'reasoning':
|
||||
delta = {'reasoning': data}
|
||||
|
||||
elif msg_type == 'tool_call':
|
||||
tool_call_count = tool_call_count + 1
|
||||
arguments = encode_json(data.arguments) if data.arguments else '{}'
|
||||
tc_entry = {"index": tool_call_count, "function": {'name': data.tool_name, 'arguments': arguments}}
|
||||
if data.id:
|
||||
tc_entry['id'] = data.id
|
||||
delta = {"tool_calls": [tc_entry]}
|
||||
|
||||
elif msg_type == 'eof':
|
||||
# Determine finish_reason — was there a tool call?
|
||||
finish_reason = "tool_calls" if any(
|
||||
e['name'] == 'tool_call' for e in run.events
|
||||
) else "stop"
|
||||
|
||||
chunk_body = {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{"delta": {}, "index": 0, "finish_reason": finish_reason}]
|
||||
}
|
||||
yield f"data: {json.dumps(chunk_body)}\n\n"
|
||||
break
|
||||
|
||||
if delta:
|
||||
chunk_body = {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{"delta": delta, "index": 0, "finish_reason": None}]
|
||||
}
|
||||
yield f"data: {json.dumps(chunk_body)}\n\n"
|
||||
else:
|
||||
raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
|
||||
|
||||
except Exception as e:
|
||||
print("[openai] ERROR: " + str(e))
|
||||
print(traceback.format_exc())
|
||||
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
||||
|
||||
return Response(generate(), content_type="text/event-stream")
|
||||
|
||||
# --- Non-streaming ---
|
||||
else:
|
||||
await control.run(run)
|
||||
await run.get_output_stream().content_wait()
|
||||
|
||||
# Collect any tool calls that were emitted
|
||||
tool_events = [e['data'] for e in run.events if e['name'] == 'tool_call']
|
||||
|
||||
if tool_events:
|
||||
# Aggregate tool calls: combine partial argument chunks by id
|
||||
tc_map = {}
|
||||
for evt in tool_events:
|
||||
tid = evt.get('id', '')
|
||||
if tid not in tc_map:
|
||||
tc_map[tid] = {
|
||||
"id": tid,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": evt.get('function', {}).get('name', ''),
|
||||
"arguments": evt.get('function', {}).get('arguments', ''),
|
||||
}
|
||||
}
|
||||
else:
|
||||
tc_map[tid]['function']['arguments'] += \
|
||||
evt.get('function', {}).get('arguments', '')
|
||||
|
||||
return {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": list(tc_map.values()),
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
}]
|
||||
}
|
||||
else:
|
||||
# No tool calls: return text content
|
||||
content = run.stream.content
|
||||
return {
|
||||
"id": "chatcmpl-" + str(time.time()),
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": agent_id,
|
||||
"choices": [{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}]
|
||||
}
|
||||
|
||||
@server.get("/api/openai/v1/models")
|
||||
@login_required
|
||||
async def openai_models():
|
||||
"""
|
||||
Return a list of available models.
|
||||
Customize this list based on your available models.
|
||||
"""
|
||||
|
||||
user_id = current_user.auth_id
|
||||
if not user_id: pass
|
||||
|
||||
return {
|
||||
"object": "list",
|
||||
"data": [{
|
||||
"id": agent.id,
|
||||
"object": "model",
|
||||
"created": int(time.time()),
|
||||
"owned_by": user_id,
|
||||
} for agent in await app.get_agents(user_id)]
|
||||
}
|
||||
|
||||
|
||||
await server.run_task(debug=True, shutdown_trigger=shutdown_event.wait, port=port)
|
||||
@@ -0,0 +1,14 @@
|
||||
import json
|
||||
|
||||
from stream import Stream
|
||||
from tool import ToolSpec
|
||||
|
||||
def encode_json(value):
|
||||
def encode_unknown(value):
|
||||
if isinstance(value, ToolSpec):
|
||||
return value.to_openai_dict()
|
||||
elif isinstance(value, Stream):
|
||||
return {'streamId': value.id}
|
||||
else:
|
||||
return repr(value)
|
||||
return json.dumps(value, default=encode_unknown)
|
||||
Reference in new issue
Block a user