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