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