tool_selector.py

"""
Tool selector module for finding the most appropriate tools based on user queries.
"""
from typing import List, Dict, Any, Tuple, Set
from abc import ABC, abstractmethod
from vector_store import EmbeddingInMemRepo
from embedding_model import EmbeddingModel
from tag_analyzer import TagAnalyzer
from utcp.shared.tool import Tool
from utcp.client.tool_search_strategy import ToolSearchStrategy

class ToolSelector(ToolSearchStrategy):
    """
    Class for selecting the most appropriate tools based on user queries.
    """
    
    def __init__(self, vector_store: EmbeddingInMemRepo, embedding_model: EmbeddingModel, tag_analyzer: TagAnalyzer):
        """
        Initialize the tool selector.
        
        Args:
            vector_store: The vector store for tool lookup
            embedding_model: The embedding model for generating embeddings
            tag_analyzer: The tag analyzer for extracting tags from queries
        """
        self.vector_store = vector_store
        self.embedding_model = embedding_model
        self.tag_analyzer = tag_analyzer
    
    async def search_tools(self, query: str, limit: int = 10) -> List[Tool]:
        """
        Search for tools relevant to the query.

        Args:
            query: The search query.
            limit: The maximum number of tools to return. 0 for no limit.

        Returns:
            A list of tools that match the search query.
        """
        # Get all available tags
        available_tags = self.vector_store.get_all_tags()
        
        # Strategy 1: Extract relevant tags from query
        relevant_tags = await self.tag_analyzer.extract_tags(query, available_tags)
        
        # Strategy 2: Generate embedding for the query
        query_embedding = await self.embedding_model.embed(query)
        
        # Perform searches using both strategies
        # Strategy 1 result: Search by tags
        tag_results = await self.vector_store.search_by_tags(relevant_tags, top_k=limit if limit > 0 else 50)
        
        # Strategy 2 result: Search by embedding
        embedding_results = await self.vector_store.search_by_embedding(query_embedding, top_k=limit if limit > 0 else 50)
        
        # Combine results from both strategies
        # Use a set to track unique tools
        unique_tools = set()
        combined_results = []
        
        # First add tools from embedding search with their scores
        for tool, score in embedding_results:
            if tool.name not in unique_tools:
                combined_results.append((tool, score, "embedding"))
                unique_tools.add(tool.name)
        
        # Then add tools from tag search if not already added
        for tool, score in tag_results:
            if tool.name not in unique_tools:
                combined_results.append((tool, score, "tag"))
                unique_tools.add(tool.name)
        
        # Sort combined results by score (descending)
        combined_results.sort(key=lambda x: x[1], reverse=True)
        
        # Return just the tools, limited to the specified limit
        if limit > 0:
            return [result[0] for result in combined_results[:limit]]
        else:
            return [result[0] for result in combined_results]
    
    async def find_tools(self, query: str, max_tools: int = 5) -> List[Tool]:
        """
        Legacy method for backward compatibility.
        Select the most appropriate tools for a user query using multiple strategies.
        
        Args:
            query: The user query
            max_tools: Maximum number of tools to return
            
        Returns:
            List of selected tools
        """
        return await self.search_tools(query, max_tools)
    
    def get_tool_summary(self, tools: List[Tool]) -> List[Dict[str, Any]]:
        """
        Generate a summary of tools for presenting to the agent.
        
        Args:
            tools: List of tools
            
        Returns:
            List of tool summary dictionaries
        """
        summaries = []
        
        for tool in tools:
            summary = {
                "name": tool.name,
                "description": tool.description,
                "tags": tool.tags,
                "provider": tool.tool_provider.name if tool.tool_provider else None,
                "provider_type": tool.tool_provider.provider_type if tool.tool_provider else None,
                # Include basic parameter information without full schema details
                "parameters": list(tool.inputs.properties.keys()) if tool.inputs and tool.inputs.properties else [],
                "returns": list(tool.outputs.properties.keys()) if tool.outputs and tool.outputs.properties else []
            }
            summaries.append(summary)
        
        return summaries