485 lines
17 KiB
Python
485 lines
17 KiB
Python
import threading
|
|
import asyncio
|
|
import json
|
|
import os
|
|
#import ollama
|
|
import traceback
|
|
from stream import Stream
|
|
from tool import format_tools_for_engine, build_tool_prompt_prefix, ToolSpec, ToolCall
|
|
|
|
#from llama_cpp import Llama
|
|
from openrouter import OpenRouter
|
|
|
|
VALID_ROLES = ['user', 'assistant', 'system', 'tool']
|
|
CONTEXT_WINDOW = 10_000
|
|
|
|
# TODO: Rename to ChatMessage. Add methods for constructing from instruct/openai/etc formatted objects.
|
|
class LlamaMessage:
|
|
def __init__(self, role, content, reasoning=None):
|
|
if role not in VALID_ROLES:
|
|
raise Exception("message role must be one of: " + ",".join(VALID_ROLES))
|
|
self.role = role
|
|
self.content = content
|
|
self.reasoning = reasoning
|
|
|
|
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
|
|
|
|
def __repr__(self):
|
|
return f"LlamaMessage(role={self.role!r}, content={self.content!r}, reasoning={self.reasoning!r})"
|
|
|
|
class Models:
|
|
llama_8b = "meta/meta-llama-3-8b-instruct"
|
|
llama_70b = "meta/meta-llama-3-70b-instruct"
|
|
qwen3_8b = "lucataco/qwen3-vl-8b-instruct:39e893666996acf464cff75688ad49ac95ef54e9f1c688fbc677330acc478e11"
|
|
qwen3_235b = "qwen/qwen3-235b-a22b-instruct-2507"
|
|
|
|
def format_openai(messages):
|
|
formatted = []
|
|
last_assistant_message = None
|
|
for msg in messages:
|
|
if msg.role == 'tool':
|
|
if not last_assistant_message:
|
|
raise Exception("tool message before any assistant message")
|
|
tool_calls = last_assistant_message.get('tool_calls')
|
|
if tool_calls == None:
|
|
tool_calls = []
|
|
last_assistant_message['tool_calls'] = tool_calls
|
|
tool_calls.append({
|
|
'id': msg.content.id,
|
|
'type': 'function',
|
|
'function': {
|
|
'name': msg.content.tool_name,
|
|
'arguments': json.dumps(msg.content.arguments),
|
|
},
|
|
})
|
|
formatted.append({
|
|
'role': 'tool',
|
|
'tool_call_id': msg.content.id,
|
|
'content': json.dumps(msg.content.result),
|
|
})
|
|
else:
|
|
fmsg = {'role': msg.role, 'content': msg.content}
|
|
formatted.append(fmsg)
|
|
if msg.role == 'assistant':
|
|
last_assistant_message = fmsg
|
|
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():
|
|
while True:
|
|
args = await queue.get()
|
|
await callback(*args)
|
|
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, LlamaMessage):
|
|
return message
|
|
elif isinstance(message, dict):
|
|
return LlamaMessage(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, LlamaMessage("system", system_prompt))
|
|
if prompt:
|
|
messages.append(LlamaMessage("user", prompt))
|
|
|
|
return messages
|
|
|
|
|
|
async def chat_completion(app,
|
|
model_id=None,
|
|
messages=None,
|
|
tools=[],
|
|
opts={},
|
|
agent_run=None
|
|
):
|
|
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 == "replicate":
|
|
# Clone opts to avoid mutating the caller's dict
|
|
opts = dict(opts)
|
|
if 'parameters' in db_model:
|
|
opts['parameters'] = db_model['parameters']
|
|
if 'format' in db_model:
|
|
opts['format'] = db_model['format']
|
|
result = await replicate_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
|
|
elif model_engine == "ollama":
|
|
result = await ollama_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))
|
|
|
|
if agent_run:
|
|
return None
|
|
return result
|
|
|
|
|
|
# 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_for_engine(tools, "llama-cpp") 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 ollama_list(app, messages, model, opts={}):
|
|
# TODO: test the output, and use this
|
|
return ollama.list()
|
|
|
|
async def ollama_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
|
|
raise Exception("ollama disabled")
|
|
if agent_run:
|
|
content_stream = agent_run.stream
|
|
else:
|
|
content_stream = Stream(app)
|
|
|
|
async def worker():
|
|
print("START OLLAMA INFERENCE")
|
|
kwargs = {
|
|
'model': model,
|
|
'messages': [msg.to_dict() for msg in messages],
|
|
'stream': True,
|
|
}
|
|
|
|
# Pass tools in Ollama format if provided
|
|
formatted_tools = format_tools_for_engine(tools, "ollama") if tools else None
|
|
if formatted_tools:
|
|
kwargs['tools'] = formatted_tools
|
|
|
|
try:
|
|
stream_response = await ollama.AsyncClient().chat(**kwargs)
|
|
|
|
async for part in stream_response:
|
|
if agent_run and agent_run.cancelled:
|
|
print(f"Cancelling ollama inference for run {agent_run.id}")
|
|
break
|
|
|
|
if part.message.content:
|
|
content_stream.feed(part.message.content)
|
|
except asyncio.CancelledError:
|
|
print(f"Ollama task cancelled for run {agent_run.id}")
|
|
raise
|
|
|
|
content_stream.close()
|
|
print("END OLLAMA INFERENCE")
|
|
|
|
task = asyncio.create_task(worker())
|
|
if agent_run:
|
|
agent_run.task = task
|
|
else:
|
|
task.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
|
|
async def message_loop(name, value=None):
|
|
nonlocal content_stream, reasoning_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)
|
|
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 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_for_engine(tools, "openrouter") if tools else None
|
|
if formatted_tools:
|
|
kwargs['tools'] = formatted_tools
|
|
|
|
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:
|
|
print(f"Cancelling openrouter inference for run {agent_run.id}")
|
|
# TODO: Need to cancell?
|
|
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):
|
|
try:
|
|
task.result()
|
|
except Exception as e:
|
|
print(f"[openrouter] Worker error: {e}")
|
|
print(traceback.format_exc())
|
|
|
|
task = asyncio.create_task(worker())
|
|
task.add_done_callback(handle_done)
|
|
agent_run.task = task
|
|
|
|
|
|
async def replicate_completion(app, messages, model=None, tools=[], opts={}, agent_run=None):
|
|
if model is None:
|
|
model = Models.llama_8b
|
|
|
|
parameter_names = opts.get('parameters', {})
|
|
|
|
arguments = {}
|
|
|
|
# Build the prompt, injecting tool descriptions for models that rely on
|
|
# prompt-based tool usage rather than native function calling.
|
|
prompt_text = format_instruct(messages)
|
|
if tools:
|
|
tool_prefix = build_tool_prompt_prefix(tools)
|
|
if tool_prefix:
|
|
prompt_text = tool_prefix + "\n\n" + prompt_text
|
|
|
|
arguments[parameter_names.get('prompt', 'prompt')] = prompt_text
|
|
arguments[parameter_names.get('temperature', 'temperature')] = opts.get('temperature') or 0.75
|
|
arguments[parameter_names.get('max_tokens', 'max_tokens')] = opts.get('max_tokens', 2048)
|
|
|
|
arguments[parameter_names.get('system_prompt', 'system_prompt')] = ""
|
|
|
|
result = await replicate_base_completion(app, model, arguments, agent_run=agent_run)
|
|
if agent_run:
|
|
return None
|
|
return result
|
|
|
|
async def replicate_base_completion(app, model, input, agent_run=None):
|
|
if agent_run:
|
|
content_stream = agent_run.stream
|
|
else:
|
|
content_stream = Stream(app)
|
|
thread = threading.Thread(
|
|
target=run_replicate_inference_loop,
|
|
kwargs={
|
|
'replicate_client': app.replicate_client,
|
|
'model': model,
|
|
'content_stream': content_stream,
|
|
'input': input,
|
|
'agent_run': agent_run
|
|
}
|
|
)
|
|
thread.start()
|
|
#response = await content_stream.content_wait()
|
|
#return response.lstrip()
|
|
if agent_run:
|
|
agent_run.thread = thread
|
|
return None
|
|
return content_stream
|
|
|
|
|
|
def run_replicate_inference_loop(replicate_client, model, input, content_stream, agent_run=None):
|
|
print("START INFERENCE")
|
|
completion = replicate_client.run(model, input=input, stream=True)
|
|
for event in completion:
|
|
if agent_run and agent_run.cancelled:
|
|
print(f"Cancelling replicate inference for run {agent_run.id}")
|
|
break
|
|
# TODO: does completion need explicit cancellation?
|
|
|
|
if event:
|
|
content_stream.feed(str(event))
|
|
print("END INFERENCE")
|
|
content_stream.close()
|
|
|
|
def format_instruct(messages):
|
|
prompt = "<|begin_of_text|>"
|
|
for message in messages:
|
|
prompt += f"<|start_header_id|>{message.role}<|end_header_id|>\n\n{message.content}<|eot_id|>"
|
|
prompt += f"<|start_header_id|>assistant<|end_header_id|>"
|
|
return prompt
|