Skip to main content

Instruction-Tool Retrieval for efficient agentic LLMs

Project description

ITR: Instruction-Tool Retrieval

Python 3.8+ License: MIT GitHub stars GitHub contributors Code style: black Ruff uv

ITR (Instruction-Tool Retrieval) is a sophisticated system for efficiently retrieving and assembling the most relevant instructions and tools for agentic Large Language Models (LLMs). It enables intelligent context management for AI agents by dynamically selecting the optimal subset of instructions and tools based on user queries and token budget constraints.

💡 Key Benefits: Reduce context bloat, improve tool selection and latency, and optimize token usage through intelligent hybrid retrieval and budget-aware selection algorithms.

📑 Table of Contents

🚀 Features

  • Hybrid Retrieval: Combines dense (embedding-based) and sparse (BM25) retrieval methods
  • Budget-Aware Selection: Intelligent token budget management with greedy optimization
  • Dynamic Chunking: Smart text chunking with configurable size ranges
  • Fallback Mechanisms: Automatic expansion of tool sets for better coverage
  • CLI Interface: Easy-to-use command-line interface for testing and integration
  • Flexible Configuration: Configurable parameters for different use cases
  • Type Safety: Full type hints for better development experience

📦 Installation

Using uv (Recommended)

# Install ITR
uv add itr

# Or install from source
git clone https://github.com/uriafranko/ITR.git
cd ITR
uv sync

Using pip

pip install itr

Development Installation

git clone https://github.com/uriafranko/ITR.git
cd ITR
uv sync --dev

⚡ Quick Start

Basic Usage

from itr import ITR, ITRConfig

# Initialize ITR with custom configuration
config = ITRConfig(
    k_a_instructions=3,  # Max instructions to select
    k_b_tools=2,         # Max tools to select
    token_budget=1500    # Total token budget
)
itr = ITR(config)

# Add instructions
itr.add_instruction(
    "You are a helpful AI assistant. Be concise and accurate.",
    metadata={"source": "base_personality", "priority": 1}
)
itr.add_instruction(
    "Always prioritize safety and ethical considerations in your responses.",
    metadata={"source": "safety_guidelines", "priority": 2}
)

# Add tools
calculator_tool = {
    "name": "calculator",
    "description": "Perform mathematical calculations and arithmetic operations",
    "schema": {
        "type": "object",
        "properties": {
            "expression": {"type": "string", "description": "Math expression to evaluate"}
        }
    }
}
itr.add_tool(calculator_tool)

# Perform retrieval
query = "What is 15% of 240?"
result = itr.step(query)

print(f"Selected {len(result.instructions)} instructions and {len(result.tools)} tools")
print(f"Total tokens: {result.total_tokens}")
print(f"Confidence: {result.confidence_score:.2f}")

# Get assembled prompt
prompt = itr.get_prompt(query)
print("\\nAssembled Prompt:")
print(prompt)

Loading from Files

from itr import ITR

itr = ITR()

# Load instructions from text file
itr.load_instructions("path/to/instructions.txt")

# Load tools from JSON file
itr.load_tools("path/to/tools.json")

# Perform retrieval
result = itr.step("How do I process a CSV file?")

Using Pre-created Fragments

from itr import ITR, InstructionFragment, FragmentType

# Create fragments manually
fragments = [
    InstructionFragment(
        id="custom_1",
        content="Use clear and concise language",
        token_count=6,
        fragment_type=FragmentType.STYLE_RULE,
        metadata={"custom": True}
    ),
    InstructionFragment(
        id="custom_2",
        content="Provide step-by-step explanations for complex topics",
        token_count=9,
        fragment_type=FragmentType.DOMAIN_SPECIFIC
    )
]

itr = ITR()
itr.add_instruction_fragments(fragments)

result = itr.step("Explain machine learning")

🖥️ CLI Usage

ITR provides a command-line interface for easy testing and integration:

Basic Retrieval

itr retrieve --query "How do I calculate compound interest?" \\
             --instructions instructions.txt \\
             --tools tools.json \\
             --show-prompt

Interactive Mode

itr interactive

Generate Configuration File

itr init-config --output my-config.json

Using Custom Configuration

itr retrieve --config my-config.json \\
             --query "Analyze this dataset" \\
             --instructions data_instructions.txt \\
             --tools analysis_tools.json

⚙️ Configuration

ITR uses a flexible configuration system. You can customize behavior through the ITRConfig class:

from itr import ITRConfig

config = ITRConfig(
    # Retrieval parameters
    top_m_instructions=20,      # Candidates to retrieve
    top_m_tools=15,            # Tool candidates to retrieve
    k_a_instructions=4,        # Max instructions to select
    k_b_tools=2,              # Max tools to select

    # Scoring weights for hybrid retrieval
    dense_weight=0.4,         # Embedding similarity weight
    sparse_weight=0.3,        # BM25 score weight
    rerank_weight=0.3,        # Reranking weight

    # Budget management
    token_budget=2000,        # Total token limit
    safety_overlay_tokens=200, # Reserved tokens

    # Fallback settings
    confidence_threshold=0.7,         # Trigger fallback below this
    discovery_expansion_factor=2.0,   # Tool expansion multiplier

    # Model settings
    embedding_model="all-MiniLM-L6-v2",
    reranker_model="cross-encoder/ms-marco-MiniLM-L-6-v2"
)

Configuration Files

ITR supports JSON, YAML and .env configuration files:

{
  "k_a_instructions": 3,
  "k_b_tools": 2,
  "token_budget": 1500,
  "confidence_threshold": 0.8,
  "embedding_model": "all-MiniLM-L6-v2"
}

🔍 How It Works

ITR uses a multi-stage pipeline for instruction and tool retrieval:

  1. Indexing: Instructions are chunked and indexed with both dense embeddings and sparse representations
  2. Retrieval: User queries are processed through hybrid retrieval (dense + sparse)
  3. Selection: Budget-aware selection algorithm picks optimal subset within token limits
  4. Assembly: Selected instructions and tools are assembled into final prompt
  5. Fallback: If confidence is low, expanded tool sets are used

Fragment Types

ITR automatically categorizes instruction fragments:

  • ROLE_GUIDANCE: Defines AI agent roles and personas
  • STYLE_RULE: Formatting and communication style guidelines
  • SAFETY_POLICY: Safety and ethical constraints
  • DOMAIN_SPECIFIC: Task-specific instructions
  • EXEMPLAR: Examples and demonstrations

🎯 Performance

ITR is designed for efficiency:

  • Fast Retrieval: Optimized hybrid search with caching
  • Memory Efficient: Lazy loading and smart chunking
  • Scalable: Handles large instruction/tool corpora
  • Token Aware: Precise token counting and budget management

Token Optimization Results

ITR Performance Results

ITR achieves significant token reduction while maintaining high confidence and functionality across diverse analysis types.

🤝 Contributing

We welcome contributions! Please see our Contributing Guidelines for details.

Development Setup

git clone https://github.com/uria-franko/ITR.git
cd ITR
uv sync --dev

# Run tests
uv run pytest

# Format code
uv run black .
uv run isort .

# Type checking
uv run mypy itr/

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

🔗 Links

📚 Examples

Check out the examples directory for more detailed usage examples:

  • Basic retrieval workflows
  • Custom instruction creation
  • Tool integration patterns
  • Configuration examples
  • Performance optimization tips

🔧 Troubleshooting

Common Issues

"Module not found" errors: Make sure you've installed all dependencies:

uv sync

Memory issues with large corpora: Use chunking and increase available memory:

config = ITRConfig(chunk_size_range=(100, 400))  # Smaller chunks

Poor retrieval quality: Try adjusting scoring weights:

config = ITRConfig(
    dense_weight=0.6,    # Increase embedding weight
    sparse_weight=0.4    # Decrease keyword weight
)

For more help, please open an issue on GitHub.


Built with ❤️ for the AI Agent community

⭐ Star this repo if you find ITR useful!

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

instruction_tool_retrieval-0.1.1.tar.gz (27.1 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

instruction_tool_retrieval-0.1.1-py3-none-any.whl (31.6 kB view details)

Uploaded Python 3

File details

Details for the file instruction_tool_retrieval-0.1.1.tar.gz.

File metadata

File hashes

Hashes for instruction_tool_retrieval-0.1.1.tar.gz
Algorithm Hash digest
SHA256 d6520b7e69f6e71196858e4eb0fbd70593c0c00da36ff5d9af77cef89639dd82
MD5 3a63befbb8ba92a74189264db89c492e
BLAKE2b-256 57a61f7aab980908831b11a829bac01c867975afbfc0f407a10efb8358133ff8

See more details on using hashes here.

File details

Details for the file instruction_tool_retrieval-0.1.1-py3-none-any.whl.

File metadata

File hashes

Hashes for instruction_tool_retrieval-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 6725fb2f0cc774bfd4e6da1f3e2c507638dcc9376f9823fd6e0a3d0842001d9e
MD5 c3c3b58fa2f84c57e514930a9637afa7
BLAKE2b-256 177f45e58bf25a099dd571337ee2719ea790c0e3f86528b125e5c9c074c21fda

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page