How to Control the Agent Loop Manually Using `agent.complete()` in Needle

Call agent.complete() directly instead of agent.run() to control each inference step manually, parse function calls yourself, and feed tool results back into subsequent prompts.

The Needle package provides a lightweight C-based inference engine for building LLM agents with tool use. While agent.run() offers an automatic loop for common workflows, many applications require finer control over the agent execution flow. This guide explains how to control the agent loop manually using agent.complete() in the cactus-compute/needle repository.

Understanding the Two Entry Points

The Needle class in needle/__init__.py exposes two primary methods for inference:

Method	Behavior	Location
Needle.complete(text, max_new_tokens)	Single-shot inference to the native engine; returns raw JSON with no automatic tool execution	needle/__init__.py:98-103
Needle.run(query, max_steps, max_new_tokens)	Automatic agent loop: calls complete(), executes any function_calls, feeds results back, repeats	needle/__init__.py:107-127

The run() method implements this automatic loop internally:

def run(self, query, max_steps=10, max_new_tokens=256):
    current = query
    for step in range(max_steps):
        response = self.complete(current, max_new_tokens)
        calls = response.get("function_calls", [])
        if not calls:
            break
        results = []
        for call in calls:
            fn = self._functions[call["name"]]
            results.append(fn(**call.get("arguments", {})))
        current = json.dumps(results, default=str)
    return response

To control the agent loop manually using agent.complete(), you replicate this logic yourself with custom logic at each step.

Setting Up the Manual Control Loop

Step 1: Instantiate the Agent with Tools

Create a Needle instance with any tools you want available. Tools are automatically converted to JSON schemas using the @tool decorator from needle/agent/tools.py.

from needle import Needle, tool, Field
import json

@tool
def search(query: str):
    """Search for information."""
    return {"results": f"results for {query}"}

@tool
def calculate(expression: str):
    """Evaluate a mathematical expression."""
    return {"value": eval(expression)}

# Create agent with tool registry

agent = Needle(tools=[search, calculate])

Step 2: Call complete() for Initial Inference

Send your original prompt directly to the engine:

response = agent.complete("What is 15 * 23?")
print(response)  # Raw JSON from native C engine

Step 3: Parse and Execute Function Calls

Inspect the response for function_calls and invoke the corresponding Python functions stored in agent._functions:

calls = response.get("function_calls", [])
if calls:
    results = []
    for call in calls:
        fn = agent._functions[call["name"]]
        args = call.get("arguments", {})
        result = fn(**args)
        results.append(result)

Step 4: Feed Results Back to Continue the Loop

Serialize tool results and pass them to complete() for the next inference round:

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

Steps 3-4 can be repeated indefinitely with your own termination logic.

Complete Manual Loop Example

This example demonstrates how to control the agent loop manually using agent.complete() with custom stop conditions and logging:

import json
from needle import Needle, tool

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

@tool
def multiply(a: int, b: int):
    """Multiply two integers."""
    return {"product": a * b}

agent = Needle(tools=[add, multiply])

prompt = "Calculate (2 + 3) * 4 step by step."
max_iterations = 5

for step in range(max_iterations):
    response = agent.complete(prompt, max_new_tokens=128)
    
    print(f"\n--- Step {step} ---")
    print(f"Model output: {response.get('output', '')}")
    
    calls = response.get("function_calls", [])
    
    # Custom termination: stop if no tools requested or explicit "FINAL" in output

    if not calls or "FINAL" in response.get("output", ""):
        print("Terminating: no more tool calls or final answer reached")
        break
    
    # Execute all requested tools

    results = []
    for call in calls:
        print(f"Executing: {call['name']}({call.get('arguments', {})})")
        fn = agent._functions[call["name"]]
        args = call.get("arguments", {})
        result = fn(**args)
        results.append(result)
        print(f"Result: {result}")
    
    # Prepare next prompt with tool results

    prompt = json.dumps(results, default=str)
else:
    print("Reached max_iterations without completion")

Using the Playground Engine for Manual Control

The Engine class in needle/playground/server.py provides the same manual control pattern with thread-safe caching. Its complete() method forwards to a cached Needle instance:

from needle.playground.server import Engine
import json

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

tools_json = json.dumps([
    {"name": "weather", "description": "Get current weather"},
    {"name": "calendar", "description": "Check calendar availability"}
])

# Manual loop with Engine wrapper

query = "What's the weather and do I have meetings today?"
for i in range(3):
    resp = engine.complete(tools_json, query)
    print(f"Round {i}: {resp.get('output')}")
    
    calls = resp.get("function_calls", [])
    if not calls:
        break
    
    # Simulate tool execution (normally would call actual implementations)

    mock_results = [{"tool": c["name"], "status": "ok"} for c in calls]
    query = json.dumps(mock_results)

The Engine.complete() implementation shows the pattern (source):

def complete(self, tools_json, query):
    from .. import Needle, _lib
    with self.lock:
        if self.agent is None or tools_json != self.tools_json:
            self.agent = Needle(tools=tools_json, weights=self.weights)
            self.tools_json = tools_json
        else:
            _lib().needle_reset()
        return self.agent.complete(query)

Key Differences: Manual vs. Automatic Control

Aspect	agent.run() (Automatic)	agent.complete() (Manual)
Loop control	Fixed max_steps parameter	Fully customizable termination
Intermediate inspection	No access to intermediate states	Full visibility into each response
Side effects	Limited hooks	Arbitrary code between iterations
Error handling	Built-in or fails	Custom retry and fallback logic
Performance	Minimal Python overhead	Slightly more overhead, maximum flexibility

When to Use Manual Loop Control

Control the agent loop manually using agent.complete() when you need to:

• Pause execution between steps for user confirmation or external triggers
• Log or audit every intermediate model response and tool call
• Modify prompts dynamically based on context not available to the model
• Implement complex retry logic for failed tool executions
• Integrate with external systems that require synchronous handoffs
• Enforce custom safety policies that inspect responses before continuing

Summary

• agent.complete() provides single-shot inference without automatic tool execution
• agent.run() wraps complete() in a fixed loop; reimplement this yourself for full control
• agent._functions maps tool names to Python callables for manual invocation
• Tool results must be JSON-serialized and passed as the next prompt to continue the conversation
• Engine.complete() in needle/playground/server.py offers the same pattern with thread-safe caching

Frequently Asked Questions

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

complete() performs a single inference pass to the C engine and returns the raw JSON response with no automatic tool execution. run() implements a loop that calls complete(), checks for function_calls, executes the corresponding Python tools, and feeds results back—repeating up to max_steps. Use complete() when you need manual control over each iteration.

How do I access the tool functions registered with a Needle agent?

Registered tools are stored in the dictionary agent._functions, which maps tool names (strings) to the decorated Python callables. When parsing a function_calls array from the model response, look up the function with agent._functions[call["name"]] and invoke it with call.get("arguments", {}).

Can I mix manual and automatic control in the same application?

Yes. You can call run() for straightforward queries where the automatic loop suffices, and switch to manual complete() calls for complex multi-step workflows requiring custom logic. Both methods operate on the same Needle instance and share the same tool registry and model weights.

Does complete() handle token limits differently than run()?

Both methods accept a max_new_tokens parameter that is passed directly to the C engine. In manual mode, you control when to stop based on response content, token usage, or iteration count. The C engine enforces the token limit per call; your Python code manages the overall conversation length.

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