Agentic RAG Cost Optimization: 7 Ways to Slash Token Costs & Latency

Step 1: "Let me check tracking" → 2,000 tokens
Step 2: "Let me check weather" → 2,000 tokens + history = 4,500 tokens  
Step 3: "Let me analyze" → 2,000 tokens + history = 6,500 tokens
Step 4: "Let me respond" → 2,000 tokens + history = 8,500 tokens
Total: ~21,500 tokens for one query

def build_prompt_with_cache(query, conversation_history):
    # STATIC: Gets cached once
    system_prompt = """
    You are a helpful assistant. You have access to these tools:
    1. get_tracking(package_id) - Check package status
    2. get_weather(city) - Check current weather
    """
    
    # STATIC: Also cached
    tool_definitions = """
    <tools>
        <tool name="get_tracking">
            <parameter name="package_id" type="string"/>
        </tool>
        <tool name="get_weather">
            <parameter name="city" type="string"/>
        </tool>
    </tools>
    """
    
    # This entire block above is cached!
    # Only the query and new history are charged full price
    
    return system_prompt, tool_definitions, query, conversation_history

def compress_tracking_response(raw_data):
    """Extract ONLY what the LLM needs. Discard the rest."""
    return {
        "status": raw_data.get("status"),
        "hub": raw_data.get("hub"),
        "estimated_delivery": raw_data.get("estimated_delivery")
    }
    # 20+ fields → just 3 fields

def compress_weather_response(raw_data):
    """Extract ONLY what the LLM needs."""
    return {
        "condition": raw_data.get("condition"),
        "temp": raw_data.get("temp"),
        "alert": raw_data.get("alerts", [{}])[0].get("message") if raw_data.get("alerts") else None
    }
    # 15+ fields → just 3 fields

# In your agent loop:
raw_tracking = await get_tracking_api(package_id)
compressed_tracking = compress_tracking_response(raw_tracking)
# Now feed compressed_tracking to the LLM, not raw_tracking

Sequential: Tool_A (1.5s) + Tool_B (1.2s) + Tool_C (2.0s) = 4.7s
Parallel: max(1.5s, 1.2s, 2.0s) + overhead = ~2.1s

import asyncio

async def execute_independent_tools(package_id, city):
    # Fire both at the SAME TIME
    tracking_task = get_tracking_api(package_id)
    weather_task = get_weather_api(city)
    
    # Wait for both to complete (they were running in parallel)
    tracking_result, weather_result = await asyncio.gather(
        tracking_task, 
        weather_task,
        return_exceptions=True  # Handle failures gracefully
    )
    
    # Check for failures
    if isinstance(tracking_result, Exception):
        # Handle tracking failure
        tracking_result = {"status": "unknown", "error": str(tracking_result)}
    if isinstance(weather_result, Exception):
        weather_result = {"condition": "unknown", "error": str(weather_result)}
    
    return tracking_result, weather_result

from sentence_transformers import SentenceTransformer, util

class SemanticRouter:
    def __init__(self, threshold=0.7):
        # Option 1: Lightweight embedding model (fast, cheap)
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
        self.threshold = threshold
        self.simple_embeddings = self.model.encode([
            "What's the weather?",
            "Hello", 
            "What's 2+2?",
            "Tell me a joke",
            "What time is it?"
        ])
    
    def route(self, query):
        query_embedding = self.model.encode(query)
        similarities = util.cos_sim(query_embedding, self.simple_embeddings)
        max_similarity = max(similarities[0])
        
        if max_similarity > self.threshold:
            return "simple"
        return "complex"

# In production
router = SemanticRouter()
for query in incoming_queries:
    if router.route(query) == "simple":
        response = cheap_model.generate(query)  # GPT-4o-mini or similar
    else:
        response = expensive_agent_pipeline(query)

def budgeted_context(retrieved_docs, tokenizer, budget=500):
    """Select only the most relevant documents up to the budget."""
    selected = []
    current_tokens = 0
    omitted = 0
    
    for doc in sorted(retrieved_docs, key=lambda x: x.get('relevance', 0), reverse=True):
        doc_tokens = len(tokenizer.encode(doc.get('text', '')))
        if current_tokens + doc_tokens <= budget:
            selected.append(doc)
            current_tokens += doc_tokens
        else:
            omitted += 1
    
    # CRITICAL: Tell the model something was omitted
    if omitted > 0:
        note = f"[Note: {omitted} additional memories were omitted due to context limits]"
        selected.append({"text": note, "relevance": 0, "is_note": True})
    
    return selected

import uuid

class MemoryCompressor:
    def __init__(self):
        self.external_store = {}
    
    def compress_observation(self, observation, key_summary):
        """Replace verbose output with compact reference."""
        obs_id = str(uuid.uuid4())[:8]
        self.external_store[obs_id] = observation
        return f"[Obs:{obs_id}] {key_summary}"
    
    def get_full(self, obs_id):
        """Retrieve full content if needed."""
        return self.external_store.get(obs_id)

# Usage
compressor = MemoryCompressor()

# Before: Sending full verbose output
history.append(f"Assistant: Tracking status: {raw_tracking}")  # 500+ tokens

# After: Compressed reference
compressed = compressor.compress_observation(
    raw_tracking, 
    "Package in Chicago, delayed due to snow, ETA Aug 10"
)
history.append(f"Assistant: Tracking: {compressed}")  # 50 tokens

import numpy as np
from sentence_transformers import SentenceTransformer
from typing import Optional, Dict, Any
from dataclasses import dataclass
from datetime import datetime, timedelta
import hashlib

@dataclass
class CacheEntry:
    response: Any
    timestamp: datetime
    embedding: np.ndarray
    ttl_seconds: int = 3600  # 1 hour default

class SemanticCache:
    def __init__(self, threshold: float = 0.85, max_size: int = 1000):
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
        self.threshold = threshold
        self.max_size = max_size
        self.cache: Dict[str, CacheEntry] = {}
    
    def _get_embedding(self, text: str) -> np.ndarray:
        return self.model.encode(text, normalize_embeddings=True)
    
    def _hash_embedding(self, embedding: np.ndarray) -> str:
        # Stable hash for embedding (not tobytes() which can be inconsistent)
        return hashlib.sha256(embedding.tobytes()).hexdigest()
    
    def get(self, query: str) -> Optional[Any]:
        query_embedding = self._get_embedding(query)
        
        # Check for expired entries
        current_time = datetime.now()
        for key, entry in list(self.cache.items()):
            if (current_time - entry.timestamp).total_seconds() > entry.ttl_seconds:
                del self.cache[key]
        
        # Find closest match
        best_match = None
        best_score = -1
        
        for key, entry in self.cache.items():
            similarity = np.dot(query_embedding, entry.embedding)
            if similarity > best_score:
                best_score = similarity
                best_match = entry
        
        if best_score > self.threshold:
            # Update timestamp (refresh TTL)
            best_match.timestamp = current_time
            return best_match.response
        
        return None
    
    def set(self, query: str, response: Any, ttl_seconds: int = 3600):
        embedding = self._get_embedding(query)
        key = self._hash_embedding(embedding)
        
        # LRU eviction
        if len(self.cache) >= self.max_size:
            oldest = min(self.cache.keys(), key=lambda k: self.cache[k].timestamp)
            del self.cache[oldest]
        
        self.cache[key] = CacheEntry(
            response=response,
            timestamp=datetime.now(),
            embedding=embedding,
            ttl_seconds=ttl_seconds
        )

# Usage in production
cache = SemanticCache(threshold=0.85, max_size=5000)

def process_query(query):
    # Check cache first
    cached = cache.get(query)
    if cached:
        return f"[Cached] {cached}"
    
    # Run the expensive pipeline
    response = expensive_agent_pipeline(query)
    
    # Cache for 1 hour
    cache.set(query, response, ttl_seconds=3600)
    return response

import asyncio
import time
import logging
from typing import Optional, Dict, Any, List
from tenacity import retry, stop_after_attempt, wait_exponential
from sentence_transformers import SentenceTransformer

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AgenticRAGPipeline:
    """Production-ready agentic RAG pipeline with all optimizations."""
    
    def __init__(self):
        self.cache = SemanticCache(threshold=0.85, max_size=5000)
        self.router = SemanticRouter(threshold=0.7)
        self.compressor = MemoryCompressor()
        self.tokenizer = self._get_tokenizer()
    
    def _get_tokenizer(self):
        """Get tokenizer for token counting."""
        try:
            from transformers import AutoTokenizer
            return AutoTokenizer.from_pretrained("gpt2")
        except:
            # Fallback: approximate token counting
            return None
    
    def _count_tokens(self, text: str) -> int:
        """Count tokens with fallback."""
        if self.tokenizer:
            return len(self.tokenizer.encode(text))
        # Rough approximation: 4 chars per token
        return len(text) // 4
    
    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
    async def _call_api_with_retry(self, api_func, *args, **kwargs):
        """Call API with retry logic."""
        try:
            return await api_func(*args, **kwargs)
        except Exception as e:
            logger.warning(f"API call failed: {e}. Retrying...")
            raise
    
    async def _execute_tools_parallel(self, tools: List[Dict]) -> Dict[str, Any]:
        """Execute multiple independent tools in parallel."""
        async def execute_one(tool):
            name = tool.get('name')
            args = tool.get('args', {})
            
            # Dispatch to the right tool handler
            handlers = {
                'get_tracking': self._call_tracking_api,
                'get_weather': self._call_weather_api,
                # Add more handlers...
            }
            
            handler = handlers.get(name)
            if not handler:
                return {name: {"error": f"Unknown tool: {name}"}}
            
            try:
                result = await self._call_api_with_retry(handler, **args)
                return {name: result}
            except Exception as e:
                logger.error(f"Tool {name} failed: {e}")
                return {name: {"error": str(e)}}
        
        tasks = [execute_one(tool) for tool in tools]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Merge results
        merged = {}
        for result in results:
            if isinstance(result, Exception):
                logger.error(f"Tool execution error: {result}")
                continue
            merged.update(result)
        
        return merged
    
    def _compress_tool_output(self, tool_name: str, raw_output: Dict) -> Dict:
        """Compress tool output based on tool type."""
        compressors = {
            'get_tracking': self._compress_tracking,
            'get_weather': self._compress_weather,
        }
        
        compressor = compressors.get(tool_name)
        if compressor:
            return compressor(raw_output)
        
        # Default: keep only essential fields
        return {"result": raw_output.get("result")}
    
    def _compress_tracking(self, raw: Dict) -> Dict:
        return {
            "status": raw.get("status"),
            "hub": raw.get("hub"),
            "estimated_delivery": raw.get("estimated_delivery")
        }
    
    def _compress_weather(self, raw: Dict) -> Dict:
        return {
            "condition": raw.get("condition"),
            "temp": raw.get("temp")
        }
    
    async def process_query(self, query: str, context: Optional[Dict] = None) -> Dict:
        """Main entry point for query processing."""
        start_time = time.time()
        context = context or {}
        
        try:
            # Step 1: Check semantic cache
            cached = self.cache.get(query)
            if cached:
                logger.info(f"Cache hit for query: {query[:50]}...")
                return {
                    "response": cached,
                    "latency": time.time() - start_time,
                    "cached": True
                }
            
            # Step 2: Route the query
            route = self.router.route(query)
            if route == "simple":
                logger.info(f"Simple query routed to cheap model: {query[:50]}...")
                response = await self._call_cheap_model(query)
                self.cache.set(query, response, ttl_seconds=3600)
                return {
                    "response": response,
                    "latency": time.time() - start_time,
                    "cached": False,
                    "routed": "simple"
                }
            
            # Step 3: Complex query → agentic pipeline
            logger.info(f"Complex query routed to agentic pipeline: {query[:50]}...")
            
            # Step 3a: Identify tools needed
            # NOTE: In production, this would use an intent classification model
            # or the LLM itself (via function calling) to determine which tools are needed.
            # The implementation below is a simplified example.
            tools = self._identify_tools(query, context)
            
            # Step 3b: Execute tools in parallel
            if tools:
                logger.info(f"Executing {len(tools)} tools in parallel...")
                raw_results = await self._execute_tools_parallel(tools)
                
                # Step 3c: Compress all tool outputs
                compressed_results = {}
                for tool_name, raw_output in raw_results.items():
                    if "error" in raw_output:
                        compressed_results[tool_name] = raw_output
                    else:
                        compressed_results[tool_name] = self._compress_tool_output(tool_name, raw_output)
            else:
                compressed_results = {}
            
            # Step 3d: Generate response with compressed context
            response = await self._generate_response(query, compressed_results)
            
            # Step 3e: Cache the response
            self.cache.set(query, response, ttl_seconds=1800)  # 30 min for complex queries
            
            return {
                "response": response,
                "latency": time.time() - start_time,
                "cached": False,
                "routed": "complex",
                "tools_used": list(compressed_results.keys())
            }
            
        except Exception as e:
            logger.error(f"Pipeline failed: {e}")
            return {
                "response": "I encountered an error processing your request. Please try again.",
                "latency": time.time() - start_time,
                "error": str(e)
            }
    
    async def _call_cheap_model(self, query: str) -> str:
        """Call a cheap model for simple queries."""
        # Implement with GPT-4o-mini or similar
        await asyncio.sleep(0.3)  # Simulate
        return f"Simple response to: {query}"
    
    def _identify_tools(self, query: str, context: Dict) -> List[Dict]:
        """
        Identify which tools are needed for this query.
        
        NOTE: In production, this would be replaced with:
        - A fine-tuned intent classification model
        - LLM-based function calling (where the LLM decides which tools to call)
        - A dedicated router model trained on your specific tool set
        
        The simple keyword-based implementation below is for demonstration only.
        """
        tools = []
        if "package" in query.lower() or "tracking" in query.lower():
            tools.append({"name": "get_tracking", "args": {"package_id": context.get("package_id")}})
        if "weather" in query.lower():
            tools.append({"name": "get_weather", "args": {"city": context.get("city")}})
        return tools
    
    async def _generate_response(self, query: str, results: Dict) -> str:
        """Generate final response with compressed context."""
        # Implement with your LLM of choice
        await asyncio.sleep(0.5)  # Simulate
        return f"Response based on: {results}"

# Usage
async def main():
    pipeline = AgenticRAGPipeline()
    
    # Test queries
    queries = [
        "Where is my package #4412?",
        "Hello",
        "What's the weather in Chicago?",
    ]
    
    for query in queries:
        result = await pipeline.process_query(query)
        print(f"Query: {query}")
        print(f"Response: {result['response']}")
        print(f"Latency: {result['latency']:.2f}s")
        print(f"Cached: {result.get('cached', False)}")
        print("-" * 40)

if __name__ == "__main__":
    asyncio.run(main())