How to Control the Agent Loop Manually with `complete()` in Needle

You can control the Needle agent loop manually by calling agent.complete() directly instead of agent.run(), allowing you to parse intermediate JSON responses, execute tools from agent._functions, and feed results back into the model on your own schedule.

The Needle package from cactus-compute/needle ships with a lightweight C-based inference engine exposed through Python. While the automatic run() method manages the entire tool-calling lifecycle, using complete() directly gives you granular control over every inference step. This approach is essential when you need custom termination logic, intermediate logging, or non-standard execution flows that the built-in loop does not support.

Understanding complete() vs. run()

According to the source code in needle/__init__.py, the Needle class exposes two primary entry points for inference:

• Needle.complete(text, max_new_tokens): Sends a raw prompt to the native engine and returns a JSON response. This method performs no tool execution and runs no loop—it simply returns the model's output, including any "function_calls" requested by the model.

• Needle.run(query, max_steps, max_new_tokens): Implements the automatic agent loop by calling complete(), checking for function calls in the response, invoking the matching Python tools, and feeding results back into complete() up to max_steps times.

When you control the agent loop manually with complete(), you replace the automatic logic in run() with your own implementation, gaining the ability to pause between steps, inject side effects, or abort early based on custom criteria.

Step-by-Step Manual Loop Implementation

To manually drive the agent loop, instantiate a Needle object and manage the call-and-response cycle yourself. The core pattern involves parsing the JSON response for "function_calls", executing the corresponding functions stored in Needle._functions, and serializing the results back to the model.

import json
from needle import Needle, tool

@tool
def search(query: str):
    """Dummy search tool."""
    return {"result": f"found {query}"}

# 1. Initialize the agent with available tools

agent = Needle(tools=[search])

# 2. Initial prompt

prompt = "What is the capital of France?"
response = agent.complete(prompt)
print(response)  # JSON response from native engine

# 3. Check for function calls and execute manually

calls = response.get("function_calls", [])
if calls:
    results = []
    for call in calls:
        # Retrieve the Python callable from the internal registry

        fn = agent._functions[call["name"]]
        args = call.get("arguments", {})
        results.append(fn(**args))
    
    # 4. Feed results back into the model

    next_prompt = json.dumps(results, default=str)
    response = agent.complete(next_prompt)
    print(response)

You can repeat steps 3 and 4 indefinitely, implementing your own max_steps counter, token limits, or domain-specific termination conditions.

Practical Example: Custom Stop Condition

This example demonstrates a manual loop that stops when the model no longer requests tool calls, using the add tool defined in needle/agent/tools.py:

import json
from needle import Needle, tool

@tool
def add(a: int, b: int):
    """Returns the sum of two integers."""
    return {"sum": a + b}

agent = Needle(tools=[add])

prompt = "Compute 2 + 3 and then multiply the result by 4."
while True:
    resp = agent.complete(prompt)
    print("Model:", resp["output"])

    calls = resp.get("function_calls")
    if not calls:
        break  # No more tool calls → exit loop

    # Execute the first tool call manually

    call = calls[0]
    fn = agent._functions[call["name"]]
    args = call.get("arguments", {})
    tool_result = fn(**args)

    # Feed the result back as the next prompt

    prompt = json.dumps([tool_result], default=str)

Using the Playground Engine for Manual Control

The Engine class in needle/playground/server.py provides a thread-safe wrapper around Needle that is used by the HTTP playground server. Its complete() method forwards requests to a cached Needle instance while handling model weight caching and reset logic.

from needle.playground.server import Engine
import json

engine = Engine()
engine.load()  # Loads default base weights

tools = json.dumps([{"name": "echo", "description": "returns its input"}])

# Manual interaction bypassing the automatic loop

resp = engine.complete(tools, "Say hello")
print(resp)

# Handle function calls manually

calls = resp.get("function_calls")
if calls:
    result = {"echo": calls[0]["arguments"]["message"]}
    resp = engine.complete(tools, json.dumps([result]))
    print(resp)

The Engine.complete() implementation (lines 30-38 in server.py) ensures thread safety with a lock and manages the Needle instance lifecycle, making it suitable for server environments where you still want manual control over the agent loop.

Key Source Files

Understanding the following files is essential when implementing manual control over the agent loop:

• needle/__init__.py: Contains the core Needle class, the complete() method (lines 98-103), and the automatic run() implementation (lines 107-127) that you are replacing with manual logic.

• needle/playground/server.py: Houses the Engine class used by the Playground server, providing a thread-safe complete() method that wraps a cached Needle instance.

• needle/agent/tools.py: Provides the @tool decorator and Field helpers for converting Python callables into JSON schemas that the native engine can understand.

• tests/test_inference.py: Contains unit tests exercising Needle.complete() and demonstrating manual loop behaviors.

Summary

• complete() provides low-level access to the native inference engine without automatic tool execution.
• run() wraps complete() in an automatic loop; bypass it by calling complete() directly to control the agent loop manually.
• Access registered tools through agent._functions when handling function calls yourself.
• Serialize tool results with json.dumps() before feeding them back into subsequent complete() calls.
• The Engine class in the Playground server offers a thread-safe alternative for manual loops in server contexts.

Frequently Asked Questions

What is the difference between complete() and run() in Needle?

complete() sends a prompt to the native C engine and returns the raw JSON response without executing any tools, while run() implements a full agent loop that automatically detects function calls, executes tools, and continues the conversation up to max_steps. Use complete() when you need manual control over the loop flow.

How do I access tool functions when using complete() manually?

Registered tools are stored in the _functions dictionary of the Needle instance. When parsing a response containing "function_calls", retrieve the callable using agent._functions[call["name"]] and invoke it with the arguments provided in the JSON payload.

Can I use the Playground Engine for manual loops?

Yes. The Engine class in needle/playground/server.py exposes a complete() method that forwards requests to a cached Needle instance. This approach provides thread safety and model caching while still allowing you to bypass the automatic run() loop and manage tool execution manually.

How do I handle tool results when using manual control?

After executing a tool function, serialize the return value using json.dumps(result, default=str) to ensure proper handling of non-JSON types. Feed this JSON string back to complete() as the next prompt, which the model will interpret as the tool execution result.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}