hordelab/hub
2
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: inference generates messages, simplify inference run events, fix encoding bug

Browse Source
This commit is contained in:
Shivering Ocelot committed 2026-10-03 11:46:16 -05:00
1 parent ee323fe660
commit 868e9b0545
5 files changed
+258 -309
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/src/inference.py
+218 -231
View File
@@ -4,15 +4,28 @@ import json
import os
import traceback
import random
import inspect
from stream import Stream
from tool import Tool, ToolSpec, ToolCall
#from llama_cpp import Llama
from openrouter import OpenRouter
TEST_MOCK = False
VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning']
CONTEXT_WINDOW = 10_000
# Background inference tasks are kept referenced here so they aren't GC'd mid-flight.
_background_tasks = set()
def spawn(coro):
task = asyncio.create_task(coro)
_background_tasks.add(task)
task.add_done_callback(_background_tasks.discard)
return task
# These reflect "emitted messages" from any actor, individual turn parts, rather than entire turns.
# role acts as a type and a sender simultaneously, and content is dynamic
class Message:
def __init__(self, role, content):
if role not in VALID_ROLES:
@@ -21,11 +34,7 @@ class Message:
self.content = content
def to_dict(self):
"""Serialize to the dict format expected by llama-cpp / ollama."""
d = {"role": self.role, "content": self.content}
if self.reasoning:
d["reasoning"] = self.reasoning
return d
return {"role": self.role, "content": self.content}
def __repr__(self):
return f"Message(role={self.role!r}, content={self.content!r})"
@@ -42,7 +51,6 @@ def format_tools(tools):
return specs
def format_openai(messages):
print('format openai ' + repr(messages))
formatted = []
def last_assistant_message():
if formatted and formatted[-1]['role'] == 'assistant':
@@ -103,9 +111,12 @@ def safe_queue(callback):
queue = asyncio.Queue()
loop = asyncio.get_running_loop()
async def read_loop():
nonlocal callback
while True:
args = await queue.get()
await callback(*args)
task = callback(*args)
if inspect.isawaitable(task):
await task
loop.create_task(read_loop())
def send(*args):
@@ -173,232 +184,200 @@ async def chat_completion(app,
messages=None,
tools=[],
opts={},
agent_run=None
):
# TESTING OVERRIDE
# return await mock_completion(app, agent_run)
db_model = (await app.get_model(model_id)).value
model_engine = db_model['engine']
model_name = db_model['model']
if model_engine == "llama-cpp":
result = await local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
elif model_engine == "openrouter":
result = await openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
else:
raise Exception("Unsupported model: " + repr(model))
print(f"[inference] chat completion: {model_engine} '{model_name}', {len(messages)} messages, {len(tools)} tools")
if agent_run:
return None
return result
if TEST_MOCK:
generator = mock_completion(app)
elif model_engine == "llama-cpp":
generator = local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts)
elif model_engine == "openrouter":
generator = openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts)
else:
raise Exception("Unsupported model engine: " + repr(model_engine))
async for message in generator:
yield message
print(f"[inference] chat completion end: {model_engine} '{model_name}'")
# llama.cpp
async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
if not agent_run:
raise Exception("agent_run required")
content_stream = None
async def message_loop(name, value=None):
nonlocal content_stream, agent_run
if name == 'output':
if content_stream == None:
content_stream = Stream(app)
agent_run.emit_event('output', content_stream)
content_stream.feed(value)
elif name == 'close':
if content_stream:
content_stream.close()
agent_run.emit_event('inference-end', {})
send_message = safe_queue(message_loop)
def worker():
llm = get_llama_cpp_model(app, model)
kwargs = {}
temp = opts.get("temperature")
if temp is not None:
kwargs['temperature'] = temp
repeat = opts.get("repetition_penalty")
if repeat is not None:
kwargs['repeat_penalty'] = repeat
# Format tools for llama-cpp (OpenAI-compatible function format)
formatted_tools = format_tools(tools) if tools else None
if formatted_tools:
kwargs['tools'] = formatted_tools
kwargs['tool_choice'] = 'auto'
# With stream=True, returns an iterator that infers tokens on demand (openai compatible)
llm_output = llm.create_chat_completion(
messages=[msg.to_dict() for msg in messages],
stream=True,
**kwargs
)
for chunk in llm_output:
# Check for cancellation
if agent_run and agent_run.cancelled:
print(f"Cancelling local inference for run {agent_run.id}")
break
delta = chunk['choices'][0]['delta']
if 'content' in delta:
send_message('output', delta['content'])
# Handle tool_call deltas if the model emits them
if 'tool_calls' in delta:
# Forward tool call chunks as JSON strings for the consumer
send_message("tool-call", {
"id": tc.get("id"),
"function": tc.get("function"),
})
send_message('close')
thread = threading.Thread(target=worker)
if agent_run:
agent_run.thread = thread
thread.start()
if agent_run:
return None
return content_stream
async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
if not agent_run:
raise Exception("agent_run required")
# TODO: Separate as InferenceRun class
content_stream = None
reasoning_stream = None
last_channel = None
def open_channel(channel):
nonlocal last_channel, content_stream, reasoning_stream
if channel != last_channel:
if content_stream: content_stream.close()
if reasoning_stream: reasoning_stream.close()
content_stream = None
reasoning_stream = None
last_channel = channel
async def message_loop(name, value=None):
nonlocal content_stream, reasoning_stream, agent_run
open_channel(name) # Hack: use the message as the channel name
if name == 'output':
if content_stream == None:
content_stream = Stream(app)
agent_run.emit_event('output', content_stream)
content_stream.feed(value)
if name == 'reasoning':
if reasoning_stream == None:
reasoning_stream = Stream(app)
agent_run.emit_event('reasoning', reasoning_stream)
reasoning_stream.feed(value)
if name == 'tool-call':
agent_run.emit_event('tool-call', value)
elif name == 'close':
if content_stream: content_stream.close()
if reasoning_stream: reasoning_stream.close()
agent_run.emit_event('inference-end', {})
emit = safe_queue(message_loop)
async def local_chat_completion(app, messages, model, tools=[], opts={}):
queue = asyncio.Queue()
loop = asyncio.get_running_loop()
emit = safe_queue(queue.put_nowait)
async def worker():
print("START OPENROUTER INFERENCE")
if not app.openrouter_client:
raise Exception("OPENROUTER_API_KEY environment variable is not set")
kwargs = {
'model': model,
'messages': format_openai(messages),
'stream': True,
# Options
'timeout_ms': 45_000,
'retries': 3,
}
# Pass tools in OpenAI-compatible format if provided
formatted_tools = format_tools(tools) if tools else None
if formatted_tools:
kwargs['tools'] = formatted_tools
print(f"Tools: {tools} and {formatted_tools}")
temp = agent_run.input.get("temperature")
if temp is not None:
kwargs['temperature'] = temp
reasoning = agent_run.input.get("reasoning_effort")
if reasoning is not None:
kwargs['reasoning'] = {"effort": reasoning}
stream_response = await app.openrouter_client.chat.send_async(**kwargs)
# OpenAI format streams tool calls in parts
# We could, but I won't assume any tool call is complete until the message is over
tool_calls = []
async for event in stream_response:
if agent_run and agent_run.cancelled:
stream_response.cancel()
break
if event.choices and len(event.choices) > 0:
delta = event.choices[0].delta
if delta.content:
emit('output', delta.content)
if delta.reasoning:
emit('reasoning', delta.reasoning)
if delta.tool_calls:
print("[openrouter] tool call:" + repr(delta.tool_calls))
for tc in delta.tool_calls:
if tc.index >= len(tool_calls):
tool_calls.append(ToolCall(tool_name="", arguments="", id=""))
tool_call = tool_calls[tc.index]
if tc.id:
tool_call.id += tc.id
if tc.function:
if tc.function.name:
tool_call.tool_name += tc.function.name
if tc.function.arguments:
tool_call.arguments += tc.function.arguments
else:
print("[openrouter] no event choices: " + repr(event))
for tc in tool_calls:
try:
arguments = json.loads(tc.arguments)
tc.arguments = arguments
emit('tool-call', tc)
except Exception:
print("[openrouter] Bad tool call arguments for " + repr(tc))
print(traceback.format_exc())
emit('close')
print("END OPENROUTER INFERENCE")
def handle_done(task):
stream = None
try:
task.result()
except asyncio.CancelledError:
print("[openrouter] inference cancelled")
except Exception as e:
print(f"[openrouter] Worker error: {e}")
llm = get_llama_cpp_model(app, model)
kwargs = {}
temp = opts.get("temperature")
if temp is not None:
kwargs['temperature'] = temp
repeat = opts.get("repetition_penalty")
if repeat is not None:
kwargs['repeat_penalty'] = repeat
# Format tools for llama-cpp (OpenAI-compatible function format)
formatted_tools = format_tools(tools) if tools else None
if formatted_tools:
kwargs['tools'] = formatted_tools
kwargs['tool_choice'] = 'auto'
# With stream=True, returns an iterator that infers tokens on demand (openai compatible)
iterator = await loop.run_in_executor(None, lambda: llm.create_chat_completion(
messages=[msg.to_dict() for msg in messages],
stream=True,
**kwargs
))
sentinel = object()
while True:
chunk = await loop.run_in_executor(None, next, iterator, sentinel)
if chunk is sentinel:
break
delta = chunk['choices'][0]['delta']
if 'content' in delta:
if stream is None:
stream = Stream(app)
emit(Message('assistant', stream))
stream.feed(delta['content'])
if 'tool_calls' in delta:
for tc in delta['tool_calls']:
emit(Message('tool', tc))
except Exception:
print(traceback.format_exc())
finally:
if stream: stream.close()
emit(None)
task = asyncio.create_task(worker())
task.add_done_callback(handle_done)
agent_run.task = task
spawn(worker())
while True:
message = await queue.get()
if message is None:
break
yield message
async def mock_completion(app, agent_run=None):
async def openrouter_chat_completion(app, messages, model, tools=[], opts={}):
if not app.openrouter_client:
raise Exception("OPENROUTER_API_KEY environment variable is not set")
queue = asyncio.Queue()
async def worker():
content_stream = None
reasoning_stream = None
last_channel = None
def open_channel(channel):
nonlocal last_channel, content_stream, reasoning_stream
if channel != last_channel:
if content_stream: content_stream.close()
if reasoning_stream: reasoning_stream.close()
content_stream = None
reasoning_stream = None
last_channel = channel
try:
kwargs = {
'model': model,
'messages': format_openai(messages),
'stream': True,
# Options
'timeout_ms': 45_000,
'retries': 3,
}
# Pass tools in OpenAI-compatible format if provided
formatted_tools = format_tools(tools) if tools else None
if formatted_tools:
kwargs['tools'] = formatted_tools
temp = opts.get("temperature")
if temp is not None:
kwargs['temperature'] = temp
reasoning = opts.get("reasoning_effort")
if reasoning is not None:
kwargs['reasoning'] = {"effort": reasoning}
stream_response = await app.openrouter_client.chat.send_async(**kwargs)
# OpenAI format streams tool calls in parts
# We could, but I won't assume any tool call is complete until the message is over
tool_calls = []
async for event in stream_response:
if event.choices and len(event.choices) > 0:
delta = event.choices[0].delta
if delta.reasoning:
open_channel('reasoning')
if reasoning_stream is None:
reasoning_stream = Stream(app)
queue.put_nowait(Message('reasoning', reasoning_stream))
reasoning_stream.feed(delta.reasoning)
if delta.content:
open_channel('output')
if content_stream is None:
content_stream = Stream(app)
queue.put_nowait(Message('assistant', content_stream))
content_stream.feed(delta.content)
if delta.tool_calls:
for tc in delta.tool_calls:
if tc.index >= len(tool_calls):
tool_calls.append(ToolCall(tool_name="", arguments="", id=""))
tool_call = tool_calls[tc.index]
if tc.id:
tool_call.id += tc.id
if tc.function:
if tc.function.name:
tool_call.tool_name += tc.function.name
if tc.function.arguments:
tool_call.arguments += tc.function.arguments
else:
print("[openrouter] no event choices: " + repr(event))
if content_stream: content_stream.close()
if reasoning_stream: reasoning_stream.close()
for tc in tool_calls:
try:
tc.arguments = json.loads(tc.arguments)
queue.put_nowait(Message('tool', tc))
except Exception:
print("[openrouter] Bad tool call arguments for " + repr(tc))
print(traceback.format_exc())
except Exception:
print(traceback.format_exc())
finally:
queue.put_nowait(None)
spawn(worker())
while True:
message = await queue.get()
if message is None:
break
yield message
async def mock_completion(app):
VERBS = ['fix','scan','sync','deploy','parse','render','cache','trace',
'merge','purge','index','probe','flush','bind','warm','drain']
OBJECTS = ['nodes','packets','tokens','frames','shards','clusters',
@@ -406,20 +385,28 @@ async def mock_completion(app, agent_run=None):
'handles','slots','cursors','segments']
DELAY = 0.3
content_stream = Stream(app)
if agent_run:
agent_run.emit_event('output', content_stream)
queue = asyncio.Queue()
emit = safe_queue(queue.put_nowait)
async def content_task ():
await asyncio.sleep(DELAY)
content_stream.feed(random.choice(VERBS))
await asyncio.sleep(DELAY)
content_stream.feed(f" {random.randint(1, 99)} ")
await asyncio.sleep(DELAY)
content_stream.feed(random.choice(VERBS))
await asyncio.sleep(DELAY)
content_stream.close()
agent_run.emit_event('inference-end', {})
async def worker():
stream = Stream(app)
emit(Message('assistant', stream))
try:
await asyncio.sleep(DELAY)
stream.feed(random.choice(VERBS))
await asyncio.sleep(DELAY)
stream.feed(f" {random.randint(1, 99)} ")
await asyncio.sleep(DELAY)
stream.feed(random.choice(VERBS))
await asyncio.sleep(DELAY)
finally:
stream.close()
emit(None)
asyncio.create_task(content_task())
if not agent_run: return content_stream
spawn(worker())
while True:
message = await queue.get()
if message is None:
break
yield message