440 lines
15 KiB
Python
440 lines
15 KiB
Python
import threading
|
|
import asyncio
|
|
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
|
|
|
|
# TODO: Collect usage data
|
|
|
|
TEST_MOCK = False
|
|
VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning', 'usage']
|
|
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:
|
|
raise Exception("message role must be one of: " + ",".join(VALID_ROLES))
|
|
self.role = role
|
|
self.content = content
|
|
|
|
def to_dict(self):
|
|
return {"role": self.role, "content": self.content}
|
|
|
|
def __repr__(self):
|
|
return f"Message(role={self.role!r}, content={self.content!r})"
|
|
|
|
def format_tools(tools):
|
|
specs = []
|
|
for t in tools:
|
|
if isinstance(t, dict):
|
|
specs.append(t)
|
|
elif isinstance(t, ToolSpec):
|
|
specs.append(t.to_openai_dict())
|
|
elif isinstance(t, Tool):
|
|
specs.append(t.get_spec().to_openai_dict())
|
|
return specs
|
|
|
|
def format_openai(messages):
|
|
formatted = []
|
|
def last_assistant_message():
|
|
if formatted and formatted[-1]['role'] == 'assistant':
|
|
return formatted[-1]
|
|
else:
|
|
msg = {'role': 'assistant', 'content': ''}
|
|
formatted.append(msg)
|
|
return msg
|
|
|
|
# For debugging
|
|
_stream_count = 0
|
|
_badrole_count = 0
|
|
_badroles = set()
|
|
|
|
for msg in messages:
|
|
content = msg.content
|
|
if isinstance(content, Stream):
|
|
_stream_count += 1
|
|
content = content.content
|
|
|
|
if msg.role == 'tool':
|
|
last = last_assistant_message()
|
|
tool_calls = last.get('tool_calls')
|
|
if tool_calls == None:
|
|
tool_calls = []
|
|
last['tool_calls'] = tool_calls
|
|
tc = content
|
|
if isinstance(tc, dict):
|
|
tc = ToolCall(
|
|
tool_name=tc['tool_name'],
|
|
arguments=tc['arguments'],
|
|
result=tc['result'],
|
|
id=tc['id']
|
|
)
|
|
# Skip non-resolved calls
|
|
if not tc.result: continue
|
|
# TODO: Fix call id if missing
|
|
tool_calls.append({
|
|
'id': tc.id,
|
|
'type': 'function',
|
|
'function': {
|
|
'name': tc.tool_name,
|
|
'arguments': json.dumps(tc.arguments),
|
|
},
|
|
})
|
|
formatted.append({
|
|
'role': 'tool',
|
|
'tool_call_id': tc.id,
|
|
'content': json.dumps(tc.result),
|
|
})
|
|
elif msg.role == 'reasoning':
|
|
last = last_assistant_message()
|
|
reasoning = last.get('reasoning', '')
|
|
last['reasoning'] = reasoning + content
|
|
elif msg.role in ['user', 'assistant', 'system']:
|
|
fmsg = {'role': msg.role, 'content': content}
|
|
formatted.append(fmsg)
|
|
else:
|
|
_badrole_count += 1
|
|
_badroles.add(msg.role)
|
|
|
|
_warnings = []
|
|
if _stream_count: _warnings.append(f"{_stream_count} streams in message list")
|
|
if _badrole_count: _warnings.append(f"{_badrole_count} messages with bad roles {_badroles}")
|
|
if _warnings: print(f"[inference] WARNING: " + ", ".join(_warnings))
|
|
|
|
return formatted
|
|
|
|
|
|
# llama cpp models
|
|
model_cache = {}
|
|
last_model_file = None
|
|
last_model_instance = None
|
|
|
|
# Wrap an async function in a queue to send it safely from another thread
|
|
def safe_queue(callback):
|
|
queue = asyncio.Queue()
|
|
loop = asyncio.get_running_loop()
|
|
async def read_loop():
|
|
nonlocal callback
|
|
while True:
|
|
args = await queue.get()
|
|
task = callback(*args)
|
|
if inspect.isawaitable(task):
|
|
await task
|
|
loop.create_task(read_loop())
|
|
|
|
def send(*args):
|
|
def queue_put():
|
|
loop.create_task(queue.put(args))
|
|
loop.call_soon_threadsafe(queue_put)
|
|
return send
|
|
|
|
def get_llama_cpp_model(app, model_file):
|
|
raise Exception("llama-cpp disabled")
|
|
global last_model_file, last_model_instance
|
|
if last_model_file == model_file:
|
|
return last_model_instance
|
|
|
|
if last_model_instance is not None:
|
|
last_model_instance = None
|
|
#del last_model_instance
|
|
|
|
last_model_file = model_file
|
|
last_model_instance = Llama(
|
|
model_path=f"{app.data_directory}/models/{model_file}",
|
|
n_gpu_layers=-1, # Uncomment to use GPU acceleration
|
|
# seed=1337, # Uncomment to set a specific seed
|
|
n_ctx=CONTEXT_WINDOW,
|
|
)
|
|
return last_model_instance
|
|
|
|
|
|
|
|
# Generate a conversation based on agent inputs
|
|
def build_messages(system_prompt="", instruction=None, input=None, messages=None):
|
|
def normalize(message):
|
|
if isinstance(message, Message):
|
|
return message
|
|
elif isinstance(message, dict):
|
|
return Message(message['role'], message['content'])
|
|
else:
|
|
raise Exception("Unsupported message type: " + repr(type(message)))
|
|
|
|
# shouldn't be necessary. maybe an assert instead?
|
|
messages = [normalize(msg) for msg in messages or []]
|
|
# default user message
|
|
prompt = input
|
|
|
|
# Instruction extends system_prompt if there's input or messages
|
|
# Otherwise instruction becomes the user prompt
|
|
# Input is always added as a user message at the end, if present
|
|
|
|
if input is None and not messages:
|
|
prompt = instruction
|
|
elif instruction:
|
|
system_prompt = f"{system_prompt}\n{instruction}" if system_prompt else instruction
|
|
|
|
# Add default messages to list
|
|
if system_prompt:
|
|
messages.insert(0, Message("system", system_prompt))
|
|
if prompt:
|
|
messages.append(Message("user", prompt))
|
|
|
|
return messages
|
|
|
|
|
|
async def chat_completion(app,
|
|
model_id=None,
|
|
messages=None,
|
|
tools=[],
|
|
opts={},
|
|
):
|
|
db_model = (await app.get_model(model_id)).value
|
|
|
|
model_engine = db_model['engine']
|
|
model_name = db_model['model']
|
|
|
|
print(f"[inference] chat completion: {model_engine} '{model_name}', {len(messages)} messages, {len(tools)} tools")
|
|
|
|
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={}):
|
|
queue = asyncio.Queue()
|
|
loop = asyncio.get_running_loop()
|
|
emit = safe_queue(queue.put_nowait)
|
|
|
|
async def worker():
|
|
stream = None
|
|
try:
|
|
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)
|
|
|
|
spawn(worker())
|
|
|
|
while True:
|
|
message = await queue.get()
|
|
if message is None:
|
|
break
|
|
yield message
|
|
|
|
|
|
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
|
|
if event.usage:
|
|
usage = {}
|
|
for key in ['completion_tokens', 'prompt_tokens', 'total_tokens', 'cost']:
|
|
if hasattr(event.usage, key):
|
|
value = getattr(event.usage, key)
|
|
if value != None:
|
|
usage[key] = value
|
|
queue.put_nowait(Message('usage', usage))
|
|
|
|
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',
|
|
'buffers','threads','sockets','records','layers','streams',
|
|
'handles','slots','cursors','segments']
|
|
DELAY = 0.3
|
|
|
|
queue = asyncio.Queue()
|
|
emit = safe_queue(queue.put_nowait)
|
|
|
|
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)
|
|
|
|
spawn(worker())
|
|
|
|
while True:
|
|
message = await queue.get()
|
|
if message is None:
|
|
break
|
|
yield message
|