How I Reduced AI Token Costs by 91% with Semantic Tool Selection and Redis

tool_embedding = {
    "description": "Send messages to Slack channels using webhooks",
    "parameters": [
        {"name": "channel_id", "description": "Target Slack channel identifier"},
        {"name": "message", "description": "Message content to send"},
        {"name": "webhook_url", "description": "Slack webhook endpoint"}
    ],
    "examples": [
        "send_slack_message(channel_id='#engineering', message='Deployment complete')",
        "Notify team about critical system updates via Slack"
    ],
    "returns": "Message delivery confirmation with timestamp"
}

def calculate_relevance_score(query_embedding, tool_embeddings):
    scores = {}
    
    # Weighted semantic matching
    scores['tool'] = cosine_similarity(query_embedding, tool_embeddings.description) * 0.50
    scores['parameters'] = max(cosine_similarity(query_embedding, param) for param in tool_embeddings.parameters) * 0.25
    scores['examples'] = max(cosine_similarity(query_embedding, example) for example in tool_embeddings.examples) * 0.15
    scores['returns'] = cosine_similarity(query_embedding, tool_embeddings.returns) * 0.10
    
    # Boost popular tools
    popularity_boost = calculate_popularity_boost(tool_id)
    
    # Apply confidence thresholding
    final_score = sum(scores.values()) + popularity_boost
    
    return final_score if final_score >= MIN_CONFIDENCE else 0

def adaptive_selection(scored_results, max_tools=5):
    scores = [result.score for result in scored_results]
    
    # Find natural score drop-off points
    score_drops = []
    for i in range(1, len(scores)):
        drop_percentage = (scores[i-1] - scores[i]) / scores[i-1]
        if drop_percentage > 0.3:  # 30% drop threshold
            score_drops.append(i)
    
    # Select tools until significant score drop or max limit
    if score_drops:
        selected_count = min(score_drops[0], max_tools)
    else:
        selected_count = max_tools
    
    return scored_results[:selected_count]

class VectorStore:
    def __init__(self, redis_client: Redis):
        self.redis = redis_client
        self.index_name = "tool_embeddings"
        self.vector_field = "embedding"
        
    async def create_hnsw_index(self):
        """Create HNSW index with optimal parameters for our use case"""
        await self.redis.ft().create_index(
            fields=[
                VectorField(self.vector_field, "FLAT", {
                    "TYPE": "FLOAT32",
                    "DIM": 1536,
                    "DISTANCE_METRIC": "COSINE",
                    "INITIAL_CAP": 1000,
                    "BLOCK_SIZE": 100
                }),
                TextField("tool_id"),
                TextField("component_type"),
                TextField("category"),
                NumericField("usage_count")
            ],
            definition=IndexDefinition(prefix=["tool:embedding:"], index_type=IndexType.HASH)
        )
    
    async def search_similar(self, query_embedding: List[float], 
                           k: int = 10, 
                           category: str = None) -> List[dict]:
        """Perform vector similarity search with optional filtering"""
        query = Query(f"*=>[KNN {k} @{self.vector_field} $query_vector]")
        
        # Add category filter if specified
        if category:
            query = Query(f"@category:{category}=>[KNN {k} @{self.vector_field} $query_vector]")
        
        query_params = {"query_vector": np.array(query_embedding).astype(np.float32).tobytes()}
        query.return_fields = ["tool_id", "component_type", "score", "category"]
        
        results = await self.redis.ft().search(query, query_params)
        return [dict(doc) for doc in results.docs]

class EmbeddingGenerator:
    def __init__(self, openai_client: AsyncOpenAI):
        self.client = openai_client
        self.batch_size = 100  # Optimize for OpenAI's rate limits
        
    async def generate_tool_embeddings(self, tool: Tool) -> ToolEmbedding:
        """Generate embeddings for all tool components"""
        
        # Collect all text components for batch processing
        texts_to_embed = [
            tool.description,
            *[f"{param.name}: {param.description}" for param in tool.parameters],
            *[example.description for example in tool.examples],
            tool.return_type.description if tool.return_type else ""
        ]
        
        # Batch generate embeddings for efficiency
        embeddings = await self.batch_generate_embeddings(texts_to_embed)
        
        # Map embeddings back to components
        tool_embedding = ToolEmbedding(
            tool_id=tool.id,
            description_embedding=embeddings[0],
            parameter_embeddings=embeddings[1:1+len(tool.parameters)],
            example_embeddings=embeddings[1+len(tool.parameters):-1],
            return_embedding=embeddings[-1] if tool.return_type else None
        )
        
        return tool_embedding
    
    async def batch_generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Efficient batch embedding generation with retry logic"""
        
        embeddings = []
        for i in range(0, len(texts), self.batch_size):
            batch = texts[i:i + self.batch_size]
            
            for attempt in range(3):  # Retry with exponential backoff
                try:
                    response = await self.client.embeddings.create(
                        model="text-embedding-3-small",
                        input=batch
                    )
                    batch_embeddings = [data.embedding for data in response.data]
                    embeddings.extend(batch_embeddings)
                    break
                    
                except RateLimitError:
                    wait_time = 2 ** attempt
                    await asyncio.sleep(wait_time)
                    continue
                except Exception as e:
                    logger.error(f"Embedding generation failed: {e}")
                    raise
        
        return embeddings

class SemanticSearch:
    def __init__(self, vector_store: VectorStore, 
                 embedding_generator: EmbeddingGenerator,
                 cache: QueryCache):
        self.vector_store = vector_store
        self.embedding_generator = embedding_generator
        self.cache = cache
        self.ranker = RelevanceRanker()
    
    async def search(self, query: str, max_results: int = 10) -> List[SearchResult]:
        """Main search method with caching and multi-component matching"""
        
        # Check cache first
        cache_key = self._generate_cache_key(query, max_results)
        cached_results = await self.cache.get(cache_key)
        if cached_results:
            return cached_results
        
        # Generate query embedding
        query_embedding = await self.embedding_generator.generate_query_embedding(query)
        
        # Search across all component types
        component_results = {}
        for component_type in ["description", "parameters", "examples", "returns"]:
            results = await self.vector_store.search_similar(
                query_embedding, 
                k=max_results * 2,  # Get more candidates for better ranking
                component_type=component_type
            )
            component_results[component_type] = results
        
        # Aggregate scores across components
        aggregated_scores = self._aggregate_component_scores(component_results)
        
        # Build search results with tool metadata
        search_results = await self._build_search_results(aggregated_scores, query)
        
        # Apply ranking and filtering
        ranked_results = await self.ranker.rank_results(search_results, query)
        
        # Cache results for 1 hour
        await self.cache.set(cache_key, ranked_results, ttl=3600)
        
        return ranked_results
    
    def _aggregate_component_scores(self, component_results: dict) -> dict:
        """Combine scores from different component types with weights"""
        
        weights = {
            "description": 0.50,
            "parameters": 0.25,
            "examples": 0.15,
            "returns": 0.10
        }
        
        tool_scores = defaultdict(float)
        
        for component_type, results in component_results.items():
            weight = weights[component_type]
            for result in results:
                tool_id = result["tool_id"]
                similarity_score = 1 - float(result["score"])  # Convert distance to similarity
                tool_scores[tool_id] += similarity_score * weight
        
        return dict(tool_scores)