From f12548a9b79e24f960606fc459d5ebd2ed2dd839 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Thu, 28 Aug 2025 18:04:35 -0500 Subject: [PATCH 1/8] Added embedding search feature for utcp 1.0 --- core/EMBEDDING_SEARCH_README.md | 278 +++++++++++++++ core/examples/embedding_search_example.py | 262 ++++++++++++++ core/pyproject.toml | 5 + core/src/utcp/implementations/__init__.py | 2 + .../utcp/implementations/embedding_search.py | 229 ++++++++++++ core/src/utcp/plugins/plugin_loader.py | 2 + .../implementations/test_embedding_search.py | 337 ++++++++++++++++++ plugins/tool_search/embedding/README.md | 19 + plugins/tool_search/embedding/pyproject.toml | 39 ++ .../src/utcp_embedding_search/__init__.py | 7 + 10 files changed, 1180 insertions(+) create mode 100644 core/EMBEDDING_SEARCH_README.md create mode 100644 core/examples/embedding_search_example.py create mode 100644 core/src/utcp/implementations/embedding_search.py create mode 100644 core/tests/implementations/test_embedding_search.py create mode 100644 plugins/tool_search/embedding/README.md create mode 100644 plugins/tool_search/embedding/pyproject.toml create mode 100644 plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py diff --git a/core/EMBEDDING_SEARCH_README.md b/core/EMBEDDING_SEARCH_README.md new file mode 100644 index 0000000..55327af --- /dev/null +++ b/core/EMBEDDING_SEARCH_README.md @@ -0,0 +1,278 @@ +# UTCP Embedding Search Plugin + +This document describes the new embedding-based semantic search strategy for UTCP tools, which provides intelligent tool discovery based on meaning similarity rather than just keyword matching. + +## Overview + +The `EmbeddingSearchStrategy` is a plugin that implements semantic search for UTCP tools using sentence embeddings. It converts tool descriptions and search queries into numerical vectors and finds the most semantically similar tools using cosine similarity. + +## Features + +- **Semantic Understanding**: Finds tools based on meaning, not just exact keyword matches +- **Configurable Similarity Threshold**: Adjustable threshold for result quality +- **Automatic Fallback**: Falls back to simple text similarity if sentence-transformers is unavailable +- **Embedding Caching**: Caches tool embeddings for improved performance +- **Tag Filtering**: Supports filtering results by required tags +- **Async Support**: Fully asynchronous implementation for non-blocking operations +- **Context Manager**: Proper resource management with async context manager + +## Installation + +### Basic Installation + +The core functionality is available with the basic dependencies: + +```bash +pip install utcp +``` + +### Enhanced Semantic Search + +For the best semantic search experience, install the optional embedding dependencies: + +```bash +pip install utcp[embedding] +``` + +This installs: +- `sentence-transformers>=2.2.0` - For high-quality sentence embeddings +- `torch>=1.9.0` - PyTorch backend for sentence-transformers + +## Quick Start + +### Basic Usage + +```python +import asyncio +from utcp.implementations.embedding_search import EmbeddingSearchStrategy +from utcp.implementations.in_mem_tool_repository import InMemToolRepository + +async def main(): + # Create a tool repository with some tools + tool_repo = InMemToolRepository() + # ... add tools to repository ... + + # Create the embedding search strategy + strategy = EmbeddingSearchStrategy( + model_name="all-MiniLM-L6-v2", + similarity_threshold=0.3, + max_workers=4, + cache_embeddings=True + ) + + # Search for tools + results = await strategy.search_tools( + tool_repo, + "I need to process some data", + limit=5 + ) + + for tool in results: + print(f"Found: {tool.name} - {tool.description}") + +asyncio.run(main()) +``` + +### Using as a Context Manager + +```python +async with EmbeddingSearchStrategy() as strategy: + results = await strategy.search_tools(tool_repo, "cooking tools", limit=3) + # Strategy automatically manages resources +``` + +## Configuration Options + +### EmbeddingSearchStrategy Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `model_name` | str | "all-MiniLM-L6-v2" | Sentence transformer model to use | +| `similarity_threshold` | float | 0.3 | Minimum similarity score (0.0-1.0) | +| `max_workers` | int | 4 | Maximum worker threads for embedding generation | +| `cache_embeddings` | bool | True | Whether to cache tool embeddings | + +### Recommended Similarity Thresholds + +- **0.2-0.3**: Broad matches, more results +- **0.4-0.5**: Balanced precision and recall +- **0.6-0.7**: High precision, fewer results +- **0.8+**: Very high precision, may miss relevant tools + +## How It Works + +### 1. Text to Embedding Conversion + +The strategy converts text into numerical embeddings: + +```python +# Query: "I need to analyze data" +# Gets converted to: [0.1, -0.3, 0.8, ...] (384-dimensional vector) + +# Tool description: "Analyze datasets and generate insights" +# Gets converted to: [0.2, -0.1, 0.9, ...] (384-dimensional vector) +``` + +### 2. Similarity Calculation + +Uses cosine similarity to measure how similar two embeddings are: + +```python +similarity = dot_product(embedding1, embedding2) / (norm(embedding1) * norm(embedding2)) +``` + +### 3. Result Ranking + +Tools are ranked by similarity score and filtered by the threshold. + +## Advanced Usage + +### Tag Filtering + +```python +# Only return tools with specific tags +results = await strategy.search_tools( + tool_repo, + "data processing", + limit=5, + any_of_tags_required=["data", "analysis"] +) +``` + +### Custom Model Selection + +```python +# Use a different sentence transformer model +strategy = EmbeddingSearchStrategy( + model_name="paraphrase-multilingual-MiniLM-L12-v2", # Multilingual support + similarity_threshold=0.4 +) +``` + +### Performance Tuning + +```python +# Optimize for your use case +strategy = EmbeddingSearchStrategy( + max_workers=8, # More workers for faster processing + cache_embeddings=True, # Cache for repeated searches + similarity_threshold=0.5 # Higher threshold for quality +) +``` + +## Fallback Behavior + +If `sentence-transformers` is not available, the strategy automatically falls back to a simple text similarity approach: + +- Uses character frequency-based embeddings +- Maintains the same API +- Provides reasonable results for basic use cases + +## Integration with UTCP + +The embedding search strategy integrates seamlessly with the UTCP plugin system: + +```python +# The strategy is automatically registered when the module is imported +from utcp.implementations.embedding_search import EmbeddingSearchStrategy + +# Use it in your UTCP client configuration +config = UtcpClientConfig( + tool_search_strategy=EmbeddingSearchStrategy( + similarity_threshold=0.4 + ) +) +``` + +## Performance Considerations + +### Memory Usage + +- Each tool embedding uses ~1.5KB of memory (384 dimensions Ɨ 4 bytes) +- With caching enabled, memory usage grows with the number of tools +- Consider disabling caching for very large tool repositories + +### Processing Speed + +- First search: Slower due to model loading and embedding generation +- Subsequent searches: Faster due to caching +- More workers = faster embedding generation but higher memory usage + +### Model Loading + +- Models are downloaded on first use (~80MB for all-MiniLM-L6-v2) +- Consider pre-downloading models in production environments + +## Troubleshooting + +### Common Issues + +1. **Import Error for sentence-transformers** + - Install with: `pip install sentence-transformers` + - The strategy will fall back to simple text similarity + +2. **Slow First Search** + - This is normal - the model needs to load + - Subsequent searches will be faster + +3. **Memory Issues** + - Reduce `max_workers` + - Disable `cache_embeddings` + - Use a smaller model (e.g., "all-MiniLM-L6-v2" instead of larger models) + +4. **Low Quality Results** + - Adjust `similarity_threshold` + - Ensure tool descriptions are detailed and meaningful + - Consider using more specific queries + +### Debug Mode + +Enable debug logging to see what's happening: + +```python +import logging +logging.basicConfig(level=logging.DEBUG) + +# The strategy will log detailed information about: +# - Model loading +# - Embedding generation +# - Similarity calculations +# - Search results +``` + +## Examples + +See `core/examples/embedding_search_example.py` for comprehensive examples demonstrating: + +- Basic search functionality +- Tag filtering +- Similarity threshold adjustment +- Context manager usage +- Comparison with tag-based search + +## Contributing + +To contribute to the embedding search plugin: + +1. Fork the repository +2. Create a feature branch +3. Add tests for new functionality +4. Ensure all tests pass +5. Submit a pull request + +## License + +This plugin is part of the UTCP project and follows the same license terms. + +## Support + +For issues and questions: + +1. Check the troubleshooting section above +2. Review the example code +3. Open an issue on the GitHub repository +4. Check the UTCP documentation + +--- + +**Note**: The embedding search plugin requires Python 3.10+ and is designed to work seamlessly with the existing UTCP ecosystem while providing enhanced semantic search capabilities. diff --git a/core/examples/embedding_search_example.py b/core/examples/embedding_search_example.py new file mode 100644 index 0000000..6183e61 --- /dev/null +++ b/core/examples/embedding_search_example.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +""" +Example demonstrating the embedding search strategy for UTCP tools. + +This example shows how to: +1. Create tools with descriptions and tags +2. Use the embedding search strategy to find semantically similar tools +3. Compare results with the traditional tag-based search +""" + +import asyncio +import sys +import os + +# Add the src directory to the path so we can import utcp modules +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) + +from utcp.data.tool import Tool, JsonSchema +from utcp.data.call_template import CallTemplate +from utcp.implementations.embedding_search import EmbeddingSearchStrategy +from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategy +from utcp.implementations.in_mem_tool_repository import InMemToolRepository + + +async def create_sample_tools(): + """Create a collection of sample tools for demonstration.""" + tools = [] + + # Cooking tools + tools.append(Tool( + name="cooking.spatula", + description="A kitchen utensil used for flipping and turning food while cooking", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["cooking", "kitchen", "utensil"], + tool_call_template=CallTemplate(name="cooking.spatula", description="Spatula tool") + )) + + tools.append(Tool( + name="cooking.whisk", + description="A kitchen tool for mixing and aerating ingredients", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["cooking", "kitchen", "mixing"], + tool_call_template=CallTemplate(name="cooking.whisk", description="Whisk tool") + )) + + tools.append(Tool( + name="cooking.knife", + description="A sharp blade for cutting and chopping food ingredients", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["cooking", "kitchen", "cutting"], + tool_call_template=CallTemplate(name="cooking.knife", description="Knife tool") + )) + + # Programming tools + tools.append(Tool( + name="dev.code_review", + description="Review and analyze source code for quality and best practices", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["programming", "development", "code"], + tool_call_template=CallTemplate(name="dev.code_review", description="Code review tool") + )) + + tools.append(Tool( + name="dev.debug", + description="Find and fix bugs in software code", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["programming", "development", "debugging"], + tool_call_template=CallTemplate(name="dev.debug", description="Debugging tool") + )) + + tools.append(Tool( + name="dev.test", + description="Run automated tests to verify code functionality", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["programming", "development", "testing"], + tool_call_template=CallTemplate(name="dev.test", description="Testing tool") + )) + + # Data analysis tools + tools.append(Tool( + name="data.analyze", + description="Analyze datasets and generate insights from data", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["data", "analysis", "insights"], + tool_call_template=CallTemplate(name="data.analyze", description="Data analysis tool") + )) + + tools.append(Tool( + name="data.visualize", + description="Create charts and graphs to represent data visually", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["data", "visualization", "charts"], + tool_call_template=CallTemplate(name="data.visualize", description="Data visualization tool") + )) + + tools.append(Tool( + name="data.clean", + description="Clean and preprocess raw data for analysis", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["data", "cleaning", "preprocessing"], + tool_call_template=CallTemplate(name="data.clean", description="Data cleaning tool") + )) + + return tools + + +async def demonstrate_search_strategies(): + """Demonstrate both search strategies with example queries.""" + + # Create tools and repository + tools = await create_sample_tools() + tool_repo = InMemToolRepository() + + # Add tools to repository + for tool in tools: + await tool_repo.save_tool(tool) + + # Create search strategies + embedding_strategy = EmbeddingSearchStrategy( + model_name="all-MiniLM-L6-v2", + similarity_threshold=0.3, + max_workers=2, + cache_embeddings=True + ) + + tag_strategy = TagAndDescriptionWordMatchStrategy() + + # Example queries to test + test_queries = [ + "I need to cook something", + "Help me write better code", + "I have data to analyze", + "Kitchen equipment", + "Software development", + "Data science tasks" + ] + + print("šŸ” UTCP Embedding Search Strategy Demo") + print("=" * 50) + print() + + for query in test_queries: + print(f"Query: '{query}'") + print("-" * 30) + + # Search with embedding strategy + print("šŸ“Š Embedding Search Results:") + try: + embedding_results = await embedding_strategy.search_tools( + tool_repo, query, limit=3 + ) + for i, tool in enumerate(embedding_results, 1): + print(f" {i}. {tool.name} (tags: {', '.join(tool.tags)})") + print(f" {tool.description}") + except Exception as e: + print(f" Error: {e}") + + print() + + # Search with tag strategy + print("šŸ·ļø Tag-based Search Results:") + try: + tag_results = await tag_strategy.search_tools( + tool_repo, query, limit=3 + ) + for i, tool in enumerate(tag_results, 1): + print(f" {i}. {tool.name} (tags: {', '.join(tool.tags)})") + print(f" {tool.description}") + except Exception as e: + print(f" Error: {e}") + + print("\n" + "=" * 50 + "\n") + + +async def demonstrate_advanced_features(): + """Demonstrate advanced features of the embedding search strategy.""" + + print("šŸš€ Advanced Embedding Search Features") + print("=" * 50) + print() + + # Create tools and repository + tools = await create_sample_tools() + tool_repo = InMemToolRepository() + + for tool in tools: + await tool_repo.save_tool(tool) + + # Create strategy with custom configuration + strategy = EmbeddingSearchStrategy( + model_name="all-MiniLM-L6-v2", + similarity_threshold=0.5, # Higher threshold for more precise matches + max_workers=4, + cache_embeddings=True + ) + + # Test tag filtering + print("1. Tag Filtering Example:") + print("Query: 'cooking tools' with required tags: ['cooking']") + results = await strategy.search_tools( + tool_repo, + "cooking tools", + limit=5, + any_of_tags_required=["cooking"] + ) + + for i, tool in enumerate(results, 1): + print(f" {i}. {tool.name} (tags: {', '.join(tool.tags)})") + + print() + + # Test different similarity thresholds + print("2. Similarity Threshold Comparison:") + thresholds = [0.2, 0.4, 0.6, 0.8] + query = "food preparation" + + for threshold in thresholds: + strategy.similarity_threshold = threshold + results = await strategy.search_tools(tool_repo, query, limit=5) + print(f" Threshold {threshold}: {len(results)} results") + + print() + + # Test context manager usage + print("3. Context Manager Usage:") + async with strategy as ctx_strategy: + results = await ctx_strategy.search_tools(tool_repo, "software", limit=3) + print(f" Found {len(results)} software-related tools") + + print("\n" + "=" * 50) + + +async def main(): + """Main function to run the demonstration.""" + try: + await demonstrate_search_strategies() + await demonstrate_advanced_features() + + print("āœ… Demo completed successfully!") + print("\nšŸ’” Tips:") + print("- Install sentence-transformers for better semantic search: pip install sentence-transformers") + print("- Adjust similarity_threshold based on your needs (0.3-0.7 recommended)") + print("- Use tag filtering to narrow down results by category") + print("- The strategy automatically falls back to simple text similarity if sentence-transformers is not available") + + except Exception as e: + print(f"āŒ Error during demo: {e}") + import traceback + traceback.print_exc() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/core/pyproject.toml b/core/pyproject.toml index aa17d4d..130c8ee 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "pydantic>=2.0", "python-dotenv>=1.0", "tomli>=2.0", + "numpy>=1.21.0", ] classifiers = [ "Development Status :: 4 - Beta", @@ -33,6 +34,10 @@ dev = [ "coverage", "twine", ] +embedding = [ + "sentence-transformers>=2.2.0", + "torch>=1.9.0", +] [project.urls] Homepage = "https://utcp.io" diff --git a/core/src/utcp/implementations/__init__.py b/core/src/utcp/implementations/__init__.py index 12fb6d6..a362425 100644 --- a/core/src/utcp/implementations/__init__.py +++ b/core/src/utcp/implementations/__init__.py @@ -1,7 +1,9 @@ from utcp.implementations.in_mem_tool_repository import InMemToolRepository from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategy +from utcp.implementations.embedding_search import EmbeddingSearchStrategy __all__ = [ "InMemToolRepository", "TagAndDescriptionWordMatchStrategy", + "EmbeddingSearchStrategy", ] diff --git a/core/src/utcp/implementations/embedding_search.py b/core/src/utcp/implementations/embedding_search.py new file mode 100644 index 0000000..4e6b244 --- /dev/null +++ b/core/src/utcp/implementations/embedding_search.py @@ -0,0 +1,229 @@ +"""Embedding-based semantic search strategy for UTCP tools. + +This module provides a semantic search implementation that uses sentence embeddings +to find tools based on meaning similarity rather than just keyword matching. +""" + +import asyncio +import logging +from typing import List, Tuple, Optional, Literal, Dict, Any +from concurrent.futures import ThreadPoolExecutor +import numpy as np +from pydantic import BaseModel, Field + +from utcp.interfaces.tool_search_strategy import ToolSearchStrategy +from utcp.data.tool import Tool +from utcp.interfaces.concurrent_tool_repository import ConcurrentToolRepository +from utcp.interfaces.serializer import Serializer + +logger = logging.getLogger(__name__) + +class EmbeddingSearchStrategy(ToolSearchStrategy): + """Semantic search strategy using sentence embeddings. + + This strategy converts tool descriptions and search queries into numerical + embeddings and finds the most semantically similar tools using cosine similarity. + """ + + tool_search_strategy_type: Literal["embedding_search"] = "embedding_search" + + # Configuration parameters + model_name: str = Field(default="all-MiniLM-L6-v2", description="Sentence transformer model to use") + similarity_threshold: float = Field(default=0.3, description="Minimum similarity score to consider a match") + max_workers: int = Field(default=4, description="Maximum number of worker threads for embedding generation") + cache_embeddings: bool = Field(default=True, description="Whether to cache tool embeddings for performance") + + def __init__(self, **data): + super().__init__(**data) + self._embedding_model = None + self._tool_embeddings_cache: Dict[str, np.ndarray] = {} + self._executor = ThreadPoolExecutor(max_workers=self.max_workers) + self._model_loaded = False + + async def _ensure_model_loaded(self): + """Ensure the embedding model is loaded.""" + if self._model_loaded: + return + + try: + # Import sentence-transformers here to avoid dependency issues + from sentence_transformers import SentenceTransformer + + # Load the model in a thread to avoid blocking + loop = asyncio.get_event_loop() + self._embedding_model = await loop.run_in_executor( + self._executor, + SentenceTransformer, + self.model_name + ) + self._model_loaded = True + logger.info(f"Loaded embedding model: {self.model_name}") + + except ImportError: + logger.warning("sentence-transformers not available, falling back to simple text similarity") + self._embedding_model = None + self._model_loaded = True + except Exception as e: + logger.error(f"Failed to load embedding model: {e}") + self._embedding_model = None + self._model_loaded = True + + async def _get_text_embedding(self, text: str) -> np.ndarray: + """Generate embedding for given text.""" + if not text: + return np.zeros(384) # Default dimension for all-MiniLM-L6-v2 + + if self._embedding_model is None: + # Fallback to simple text similarity + return self._simple_text_embedding(text) + + try: + loop = asyncio.get_event_loop() + embedding = await loop.run_in_executor( + self._executor, + self._embedding_model.encode, + text + ) + return embedding + except Exception as e: + logger.warning(f"Failed to generate embedding for text: {e}") + return self._simple_text_embedding(text) + + def _simple_text_embedding(self, text: str) -> np.ndarray: + """Simple fallback embedding using character frequency.""" + # Create a simple embedding based on character frequency + # This is a fallback when sentence-transformers is not available + embedding = np.zeros(384) + text_lower = text.lower() + + # Simple character frequency-based embedding + for i, char in enumerate(text_lower): + if i < 384: + embedding[i % 384] += ord(char) / 1000.0 + + # Normalize + norm = np.linalg.norm(embedding) + if norm > 0: + embedding = embedding / norm + + return embedding + + async def _get_tool_embedding(self, tool: Tool) -> np.ndarray: + """Get or generate embedding for a tool.""" + if not self.cache_embeddings or tool.name not in self._tool_embeddings_cache: + # Create text representation of the tool + tool_text = f"{tool.name} {tool.description} {' '.join(tool.tags)}" + embedding = await self._get_text_embedding(tool_text) + + if self.cache_embeddings: + self._tool_embeddings_cache[tool.name] = embedding + + return embedding + + return self._tool_embeddings_cache[tool.name] + + def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float: + """Calculate cosine similarity between two vectors.""" + try: + dot_product = np.dot(a, b) + norm_a = np.linalg.norm(a) + norm_b = np.linalg.norm(b) + + if norm_a == 0 or norm_b == 0: + return 0.0 + + return dot_product / (norm_a * norm_b) + except Exception as e: + logger.warning(f"Error calculating cosine similarity: {e}") + return 0.0 + + async def search_tools( + self, + tool_repository: ConcurrentToolRepository, + query: str, + limit: int = 10, + any_of_tags_required: Optional[List[str]] = None + ) -> List[Tool]: + """Search for tools using semantic similarity. + + Args: + tool_repository: The tool repository to search within. + query: The search query string. + limit: Maximum number of tools to return. + any_of_tags_required: Optional list of tags where one of them must be present. + + Returns: + List of Tool objects ranked by semantic similarity. + """ + if limit < 0: + raise ValueError("limit must be non-negative") + + # Ensure the embedding model is loaded + await self._ensure_model_loaded() + + # Get all tools + tools: List[Tool] = await tool_repository.get_tools() + + # Filter by required tags if specified + if any_of_tags_required and len(any_of_tags_required) > 0: + any_of_tags_required = [tag.lower() for tag in any_of_tags_required] + tools = [ + tool for tool in tools + if any(tag.lower() in any_of_tags_required for tag in tool.tags) + ] + + if not tools: + return [] + + # Generate query embedding + query_embedding = await self._get_text_embedding(query) + + # Calculate similarity scores for all tools + tool_scores: List[Tuple[Tool, float]] = [] + + for tool in tools: + try: + tool_embedding = await self._get_tool_embedding(tool) + similarity = self._cosine_similarity(query_embedding, tool_embedding) + + if similarity >= self.similarity_threshold: + tool_scores.append((tool, similarity)) + + except Exception as e: + logger.warning(f"Error processing tool {tool.name}: {e}") + continue + + # Sort by similarity score (descending) + sorted_tools = [ + tool for tool, score in sorted( + tool_scores, + key=lambda x: x[1], + reverse=True + ) + ] + + # Return up to 'limit' tools + return sorted_tools[:limit] if limit > 0 else sorted_tools + + async def __aenter__(self): + """Async context manager entry.""" + await self._ensure_model_loaded() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Async context manager exit.""" + if self._executor: + self._executor.shutdown(wait=False) + + +class EmbeddingSearchStrategyConfigSerializer(Serializer[EmbeddingSearchStrategy]): + """Serializer for EmbeddingSearchStrategy configuration.""" + + def to_dict(self, obj: EmbeddingSearchStrategy) -> dict: + return obj.model_dump() + + def validate_dict(self, data: dict) -> EmbeddingSearchStrategy: + try: + return EmbeddingSearchStrategy.model_validate(data) + except Exception as e: + raise ValueError(f"Invalid configuration: {e}") from e diff --git a/core/src/utcp/plugins/plugin_loader.py b/core/src/utcp/plugins/plugin_loader.py index 6fa2cf6..100000b 100644 --- a/core/src/utcp/plugins/plugin_loader.py +++ b/core/src/utcp/plugins/plugin_loader.py @@ -6,6 +6,7 @@ def _load_plugins(): from utcp.interfaces.tool_search_strategy import ToolSearchStrategyConfigSerializer from utcp.implementations.in_mem_tool_repository import InMemToolRepositoryConfigSerializer from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategyConfigSerializer + from utcp.implementations.embedding_search import EmbeddingSearchStrategyConfigSerializer from utcp.data.auth_implementations import OAuth2AuthSerializer, BasicAuthSerializer, ApiKeyAuthSerializer from utcp.data.variable_loader_implementations import DotEnvVariableLoaderSerializer from utcp.implementations.post_processors import FilterDictPostProcessorConfigSerializer, LimitStringsPostProcessorConfigSerializer @@ -19,6 +20,7 @@ def _load_plugins(): register_tool_repository(ConcurrentToolRepositoryConfigSerializer.default_repository, InMemToolRepositoryConfigSerializer()) register_tool_search_strategy(ToolSearchStrategyConfigSerializer.default_strategy, TagAndDescriptionWordMatchStrategyConfigSerializer()) + register_tool_search_strategy("embedding_search", EmbeddingSearchStrategyConfigSerializer()) register_tool_post_processor("filter_dict", FilterDictPostProcessorConfigSerializer()) register_tool_post_processor("limit_strings", LimitStringsPostProcessorConfigSerializer()) diff --git a/core/tests/implementations/test_embedding_search.py b/core/tests/implementations/test_embedding_search.py new file mode 100644 index 0000000..334100e --- /dev/null +++ b/core/tests/implementations/test_embedding_search.py @@ -0,0 +1,337 @@ +"""Tests for the EmbeddingSearchStrategy implementation.""" + +import pytest +import numpy as np +from unittest.mock import patch +from typing import List + +from utcp.implementations.embedding_search import EmbeddingSearchStrategy +from utcp.data.tool import Tool, JsonSchema +from utcp.data.call_template import CallTemplate + + +class MockToolRepository: + """Simplified mock repository for testing.""" + + def __init__(self, tools: List[Tool]): + self.tools = tools + + async def get_tools(self) -> List[Tool]: + return self.tools + + +@pytest.fixture +def sample_tools(): + """Create sample tools for testing.""" + tools = [] + + # Tool 1: Cooking related + tool1 = Tool( + name="cooking.spatula", + description="A kitchen utensil used for flipping and turning food while cooking", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["cooking", "kitchen", "utensil"], + tool_call_template=CallTemplate( + name="cooking.spatula", + description="Spatula tool", + call_template_type="default" + ) + ) + tools.append(tool1) + + # Tool 2: Programming related + tool2 = Tool( + name="dev.code_review", + description="Review and analyze source code for quality and best practices", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["programming", "development", "code"], + tool_call_template=CallTemplate( + name="dev.code_review", + description="Code review tool", + call_template_type="default" + ) + ) + tools.append(tool2) + + # Tool 3: Data analysis + tool3 = Tool( + name="data.analyze", + description="Analyze datasets and generate insights from data", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["data", "analysis", "insights"], + tool_call_template=CallTemplate( + name="data.analyze", + description="Data analysis tool", + call_template_type="default" + ) + ) + tools.append(tool3) + + return tools + + +@pytest.fixture +def embedding_strategy(): + """Create an embedding search strategy instance.""" + return EmbeddingSearchStrategy( + model_name="all-MiniLM-L6-v2", + similarity_threshold=0.3, + max_workers=2, + cache_embeddings=True + ) + + +# --- Your existing tests remain unchanged below this line --- + + +@pytest.mark.asyncio + +async def test_embedding_strategy_initialization(embedding_strategy): + """Test that the embedding strategy initializes correctly.""" + assert embedding_strategy.tool_search_strategy_type == "embedding_search" + assert embedding_strategy.model_name == "all-MiniLM-L6-v2" + assert embedding_strategy.similarity_threshold == 0.3 + assert embedding_strategy.max_workers == 2 + assert embedding_strategy.cache_embeddings is True + + +@pytest.mark.asyncio +async def test_simple_text_embedding_fallback(embedding_strategy): + """Test the fallback text embedding when sentence-transformers is not available.""" + # Mock the embedding model to be None to trigger fallback + embedding_strategy._embedding_model = None + embedding_strategy._model_loaded = True + + text = "test text" + embedding = await embedding_strategy._get_text_embedding(text) + + assert isinstance(embedding, np.ndarray) + assert embedding.shape == (384,) + assert np.linalg.norm(embedding) > 0 + + +@pytest.mark.asyncio +async def test_cosine_similarity_calculation(embedding_strategy): + """Test cosine similarity calculation.""" + # Test with identical vectors + vec1 = np.array([1.0, 0.0, 0.0]) + vec2 = np.array([1.0, 0.0, 0.0]) + similarity = embedding_strategy._cosine_similarity(vec1, vec2) + assert similarity == pytest.approx(1.0) + + # Test with orthogonal vectors + vec3 = np.array([0.0, 1.0, 0.0]) + similarity = embedding_strategy._cosine_similarity(vec1, vec3) + assert similarity == pytest.approx(0.0) + + # Test with zero vectors + vec4 = np.zeros(3) + similarity = embedding_strategy._cosine_similarity(vec1, vec4) + assert similarity == 0.0 + + +@pytest.mark.asyncio +async def test_tool_embedding_generation(embedding_strategy, sample_tools): + """Test that tool embeddings are generated and cached correctly.""" + tool = sample_tools[0] + + # Mock the text embedding method + with patch.object(embedding_strategy, '_get_text_embedding') as mock_embed: + mock_embed.return_value = np.random.rand(384) + + # First call should generate and cache + embedding1 = await embedding_strategy._get_tool_embedding(tool) + assert tool.name in embedding_strategy._tool_embeddings_cache + + # Second call should use cache + embedding2 = await embedding_strategy._get_tool_embedding(tool) + assert np.array_equal(embedding1, embedding2) + + # Verify the mock was called only once + mock_embed.assert_called_once() + + +@pytest.mark.asyncio +async def test_search_tools_basic(embedding_strategy, sample_tools): + """Test basic search functionality.""" + tool_repo = MockToolRepository(sample_tools) + + # Mock the embedding methods + with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed: + + # Create mock embeddings + query_embedding = np.random.rand(384) + tool_embeddings = [np.random.rand(384) for _ in sample_tools] + + mock_query_embed.return_value = query_embedding + mock_tool_embed.side_effect = tool_embeddings + + # Mock cosine similarity to return high scores + with patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: + mock_sim.return_value = 0.8 # High similarity + + results = await embedding_strategy.search_tools(tool_repo, "cooking", limit=2) + + assert len(results) == 2 + assert all(isinstance(tool, Tool) for tool in results) + + +@pytest.mark.asyncio +async def test_search_tools_with_tag_filtering(embedding_strategy, sample_tools): + """Test search with tag filtering.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed, \ + patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: + + mock_query_embed.return_value = np.random.rand(384) + mock_tool_embed.return_value = np.random.rand(384) + mock_sim.return_value = 0.8 + + # Search with required tags + results = await embedding_strategy.search_tools( + tool_repo, + "cooking", + limit=10, + any_of_tags_required=["cooking", "kitchen"] + ) + + # Should only return tools with cooking or kitchen tags + assert all( + any(tag in ["cooking", "kitchen"] for tag in tool.tags) + for tool in results + ) + + +@pytest.mark.asyncio +async def test_search_tools_with_similarity_threshold(embedding_strategy, sample_tools): + """Test that similarity threshold filtering works correctly.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed, \ + patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: + + mock_query_embed.return_value = np.random.rand(384) + mock_tool_embed.return_value = np.random.rand(384) + + # Set threshold to 0.5 and return scores below and above + embedding_strategy.similarity_threshold = 0.5 + mock_sim.side_effect = [0.3, 0.7, 0.2] # Only second tool should pass + + results = await embedding_strategy.search_tools(tool_repo, "test", limit=10) + + assert len(results) == 1 # Only one tool above threshold + + +@pytest.mark.asyncio +async def test_search_tools_limit_respected(embedding_strategy, sample_tools): + """Test that the limit parameter is respected.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed, \ + patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: + + mock_query_embed.return_value = np.random.rand(384) + mock_tool_embed.return_value = np.random.rand(384) + mock_sim.return_value = 0.8 + + # Test with limit 1 + results = await embedding_strategy.search_tools(tool_repo, "test", limit=1) + assert len(results) == 1 + + # Test with limit 0 (no limit) + results = await embedding_strategy.search_tools(tool_repo, "test", limit=0) + assert len(results) == 3 # All tools + + +@pytest.mark.asyncio +async def test_search_tools_empty_repository(embedding_strategy): + """Test search behavior with empty tool repository.""" + tool_repo = MockToolRepository([]) + + results = await embedding_strategy.search_tools(tool_repo, "test", limit=10) + assert results == [] + + +@pytest.mark.asyncio +async def test_search_tools_invalid_limit(embedding_strategy, sample_tools): + """Test that invalid limit values raise appropriate errors.""" + tool_repo = MockToolRepository(sample_tools) + + with pytest.raises(ValueError, match="limit must be non-negative"): + await embedding_strategy.search_tools(tool_repo, "test", limit=-1) + + +@pytest.mark.asyncio +async def test_context_manager_behavior(embedding_strategy): + """Test async context manager behavior.""" + async with embedding_strategy as strategy: + assert strategy._model_loaded is True + + # Executor should be shut down + assert strategy._executor._shutdown is True + + +@pytest.mark.asyncio +async def test_error_handling_in_search(embedding_strategy, sample_tools): + """Test that errors in search are handled gracefully.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed: + + mock_query_embed.return_value = np.random.rand(384) + + # Make the second tool fail + def mock_tool_embed_side_effect(tool): + if tool.name == "dev.code_review": + raise Exception("Simulated error") + return np.random.rand(384) + + mock_tool_embed.side_effect = mock_tool_embed_side_effect + + # Mock cosine similarity + with patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: + mock_sim.return_value = 0.8 + + # Should not crash, just skip the problematic tool + results = await embedding_strategy.search_tools(tool_repo, "test", limit=10) + + # Should return tools that didn't fail + assert len(results) == 2 # One tool failed, so only 2 results + + +@pytest.mark.asyncio +async def test_embedding_strategy_config_serializer(): + """Test the configuration serializer.""" + from utcp.implementations.embedding_search import EmbeddingSearchStrategyConfigSerializer + + serializer = EmbeddingSearchStrategyConfigSerializer() + + # Test serialization + strategy = EmbeddingSearchStrategy( + model_name="test-model", + similarity_threshold=0.5, + max_workers=8, + cache_embeddings=False + ) + + config_dict = serializer.to_dict(strategy) + assert config_dict["model_name"] == "test-model" + assert config_dict["similarity_threshold"] == 0.5 + assert config_dict["max_workers"] == 8 + assert config_dict["cache_embeddings"] is False + + # Test deserialization + restored_strategy = serializer.validate_dict(config_dict) + assert restored_strategy.model_name == "test-model" + assert restored_strategy.similarity_threshold == 0.5 + assert restored_strategy.max_workers == 8 + assert restored_strategy.cache_embeddings is False diff --git a/plugins/tool_search/embedding/README.md b/plugins/tool_search/embedding/README.md new file mode 100644 index 0000000..7a55d82 --- /dev/null +++ b/plugins/tool_search/embedding/README.md @@ -0,0 +1,19 @@ +# UTCP Embedding Search Plugin + +This plugin registers the embedding-based semantic search strategy with UTCP 1.0 via entry points. + +## Installation + +```bash +pip install utcp-embedding-search +``` + +Optionally, for high-quality embeddings: + +```bash +pip install utcp-embedding-search[embedding] +``` + +## How it works + +When installed, this package exposes an entry point under `utcp.plugins` so the UTCP core can auto-discover and register the `embedding_search` strategy. diff --git a/plugins/tool_search/embedding/pyproject.toml b/plugins/tool_search/embedding/pyproject.toml new file mode 100644 index 0000000..a8cabed --- /dev/null +++ b/plugins/tool_search/embedding/pyproject.toml @@ -0,0 +1,39 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "utcp-embedding-search" +version = "1.0.0" +authors = [ + { name = "UTCP Contributors" }, +] +description = "UTCP plugin providing embedding-based semantic tool search." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "utcp>=1.0", +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] +license = "MPL-2.0" + +[project.optional-dependencies] +embedding = [ + "sentence-transformers>=2.2.0", + "torch>=1.9.0", +] + +authors = [ { name = "UTCP Contributors" } ] + +[project.urls] +Homepage = "https://utcp.io" +Source = "https://github.com/universal-tool-calling-protocol/python-utcp" +Issues = "https://github.com/universal-tool-calling-protocol/python-utcp/issues" + +[project.entry-points."utcp.plugins"] +embedding_search = "utcp_embedding_search:register" diff --git a/plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py b/plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py new file mode 100644 index 0000000..a20fc1d --- /dev/null +++ b/plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py @@ -0,0 +1,7 @@ +from utcp.plugins.discovery import register_tool_search_strategy +from utcp.implementations.embedding_search import EmbeddingSearchStrategyConfigSerializer + + +def register(): + """Entry point function to register the embedding search strategy.""" + register_tool_search_strategy("embedding_search", EmbeddingSearchStrategyConfigSerializer()) From c8a4777b4491ecb2c539fc5497015af96b617f97 Mon Sep 17 00:00:00 2001 From: Thuraabtech <97426541+Thuraabtech@users.noreply.github.com> Date: Thu, 28 Aug 2025 18:58:36 -0500 Subject: [PATCH 2/8] Update plugins/tool_search/embedding/pyproject.toml Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- plugins/tool_search/embedding/pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/plugins/tool_search/embedding/pyproject.toml b/plugins/tool_search/embedding/pyproject.toml index a8cabed..0a4c90e 100644 --- a/plugins/tool_search/embedding/pyproject.toml +++ b/plugins/tool_search/embedding/pyproject.toml @@ -28,7 +28,6 @@ embedding = [ "torch>=1.9.0", ] -authors = [ { name = "UTCP Contributors" } ] [project.urls] Homepage = "https://utcp.io" From 2dd5752b3d7ad4733da15aa76d388c23a4f72a6b Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Tue, 2 Sep 2025 14:42:33 +0200 Subject: [PATCH 3/8] Update plugins/tool_search/embedding/README.md Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com> --- plugins/tool_search/embedding/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/tool_search/embedding/README.md b/plugins/tool_search/embedding/README.md index 7a55d82..ca0b906 100644 --- a/plugins/tool_search/embedding/README.md +++ b/plugins/tool_search/embedding/README.md @@ -11,7 +11,7 @@ pip install utcp-embedding-search Optionally, for high-quality embeddings: ```bash -pip install utcp-embedding-search[embedding] +pip install "utcp-embedding-search[embedding]" ``` ## How it works From 8dc2fe6bdf5b4a827a94cad7cec52fe87a3712fd Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Thu, 4 Sep 2025 12:28:25 -0500 Subject: [PATCH 4/8] To be resolve --- .../tool_search/in_mem_embeddings/README.md | 21 ++ .../in_mem_embeddings/pyproject.toml | 37 ++ .../src/utcp_in_mem_embeddings/__init__.py | 7 + .../in_mem_embeddings_search.py | 231 ++++++++++++ .../in_mem_embeddings/test_integration.py | 87 +++++ .../in_mem_embeddings/test_performance.py | 101 ++++++ .../in_mem_embeddings/test_plugin.py | 105 ++++++ .../tests/test_in_mem_embeddings_search.py | 343 ++++++++++++++++++ 8 files changed, 932 insertions(+) create mode 100644 plugins/tool_search/in_mem_embeddings/README.md create mode 100644 plugins/tool_search/in_mem_embeddings/pyproject.toml create mode 100644 plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/__init__.py create mode 100644 plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py create mode 100644 plugins/tool_search/in_mem_embeddings/test_integration.py create mode 100644 plugins/tool_search/in_mem_embeddings/test_performance.py create mode 100644 plugins/tool_search/in_mem_embeddings/test_plugin.py create mode 100644 plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py diff --git a/plugins/tool_search/in_mem_embeddings/README.md b/plugins/tool_search/in_mem_embeddings/README.md new file mode 100644 index 0000000..4e435e4 --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/README.md @@ -0,0 +1,21 @@ +# UTCP In-Memory Embeddings Search Plugin + +This plugin registers the in-memory embedding-based semantic search strategy with UTCP 1.0 via entry points. + +## Installation + +```bash +pip install utcp-in-mem-embeddings +``` + +Optionally, for high-quality embeddings: + +```bash +pip install utcp-in-mem-embeddings[embedding] +``` + +## How it works + +When installed, this package exposes an entry point under `utcp.plugins` so the UTCP core can auto-discover and register the `in_mem_embeddings` strategy. + +The embeddings are cached in memory for improved performance during repeated searches. diff --git a/plugins/tool_search/in_mem_embeddings/pyproject.toml b/plugins/tool_search/in_mem_embeddings/pyproject.toml new file mode 100644 index 0000000..0bdccd8 --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "utcp-in-mem-embeddings" +version = "1.0.0" +authors = [ + { name = "UTCP Contributors" }, +] +description = "UTCP plugin providing in-memory embedding-based semantic tool search." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "utcp>=1.0", +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] +license = "MPL-2.0" + +[project.optional-dependencies] +embedding = [ + "sentence-transformers>=2.2.0", + "torch>=1.9.0", +] + +[project.urls] +Homepage = "https://utcp.io" +Source = "https://github.com/universal-tool-calling-protocol/python-utcp" +Issues = "https://github.com/universal-tool-calling-protocol/python-utcp/issues" + +[project.entry-points."utcp.plugins"] +in_mem_embeddings = "utcp_in_mem_embeddings:register" diff --git a/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/__init__.py b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/__init__.py new file mode 100644 index 0000000..53da84d --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/__init__.py @@ -0,0 +1,7 @@ +from utcp.plugins.discovery import register_tool_search_strategy +from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategyConfigSerializer + + +def register(): + """Entry point function to register the in-memory embeddings search strategy.""" + register_tool_search_strategy("in_mem_embeddings", InMemEmbeddingsSearchStrategyConfigSerializer()) diff --git a/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py new file mode 100644 index 0000000..83cff80 --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py @@ -0,0 +1,231 @@ +"""In-memory embedding-based semantic search strategy for UTCP tools. + +This module provides a semantic search implementation that uses sentence embeddings +to find tools based on meaning similarity rather than just keyword matching. +Embeddings are cached in memory for improved performance. +""" + +import asyncio +import logging +from typing import List, Tuple, Optional, Literal, Dict, Any +from concurrent.futures import ThreadPoolExecutor +import numpy as np +from pydantic import BaseModel, Field + +from utcp.interfaces.tool_search_strategy import ToolSearchStrategy +from utcp.data.tool import Tool +from utcp.interfaces.concurrent_tool_repository import ConcurrentToolRepository +from utcp.interfaces.serializer import Serializer + +logger = logging.getLogger(__name__) + +class InMemEmbeddingsSearchStrategy(ToolSearchStrategy): + """In-memory semantic search strategy using sentence embeddings. + + This strategy converts tool descriptions and search queries into numerical + embeddings and finds the most semantically similar tools using cosine similarity. + Embeddings are cached in memory for improved performance during repeated searches. + """ + + tool_search_strategy_type: Literal["in_mem_embeddings"] = "in_mem_embeddings" + + # Configuration parameters + model_name: str = Field(default="all-MiniLM-L6-v2", description="Sentence transformer model to use") + similarity_threshold: float = Field(default=0.3, description="Minimum similarity score to consider a match") + max_workers: int = Field(default=4, description="Maximum number of worker threads for embedding generation") + cache_embeddings: bool = Field(default=True, description="Whether to cache tool embeddings for performance") + + def __init__(self, **data): + super().__init__(**data) + self._embedding_model = None + self._tool_embeddings_cache: Dict[str, np.ndarray] = {} + self._executor = ThreadPoolExecutor(max_workers=self.max_workers) + self._model_loaded = False + + async def _ensure_model_loaded(self): + """Ensure the embedding model is loaded.""" + if self._model_loaded: + return + + try: + # Import sentence-transformers here to avoid dependency issues + from sentence_transformers import SentenceTransformer + + # Load the model in a thread to avoid blocking + loop = asyncio.get_event_loop() + self._embedding_model = await loop.run_in_executor( + self._executor, + SentenceTransformer, + self.model_name + ) + self._model_loaded = True + logger.info(f"Loaded embedding model: {self.model_name}") + + except ImportError: + logger.warning("sentence-transformers not available, falling back to simple text similarity") + self._embedding_model = None + self._model_loaded = True + except Exception as e: + logger.error(f"Failed to load embedding model: {e}") + self._embedding_model = None + self._model_loaded = True + + async def _get_text_embedding(self, text: str) -> np.ndarray: + """Generate embedding for given text.""" + if not text: + return np.zeros(384) # Default dimension for all-MiniLM-L6-v2 + + if self._embedding_model is None: + # Fallback to simple text similarity + return self._simple_text_embedding(text) + + try: + loop = asyncio.get_event_loop() + embedding = await loop.run_in_executor( + self._executor, + self._embedding_model.encode, + text + ) + return embedding + except Exception as e: + logger.warning(f"Failed to generate embedding for text: {e}") + return self._simple_text_embedding(text) + + def _simple_text_embedding(self, text: str) -> np.ndarray: + """Simple fallback embedding using character frequency.""" + # Create a simple embedding based on character frequency + # This is a fallback when sentence-transformers is not available + embedding = np.zeros(384) + text_lower = text.lower() + + # Simple character frequency-based embedding + for i, char in enumerate(text_lower): + if i < 384: + embedding[i % 384] += ord(char) / 1000.0 + + # Normalize + norm = np.linalg.norm(embedding) + if norm > 0: + embedding = embedding / norm + + return embedding + + async def _get_tool_embedding(self, tool: Tool) -> np.ndarray: + """Get or generate embedding for a tool.""" + if not self.cache_embeddings or tool.name not in self._tool_embeddings_cache: + # Create text representation of the tool + tool_text = f"{tool.name} {tool.description} {' '.join(tool.tags)}" + embedding = await self._get_text_embedding(tool_text) + + if self.cache_embeddings: + self._tool_embeddings_cache[tool.name] = embedding + + return embedding + + return self._tool_embeddings_cache[tool.name] + + def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float: + """Calculate cosine similarity between two vectors.""" + try: + dot_product = np.dot(a, b) + norm_a = np.linalg.norm(a) + norm_b = np.linalg.norm(b) + + if norm_a == 0 or norm_b == 0: + return 0.0 + + return dot_product / (norm_a * norm_b) + except Exception as e: + logger.warning(f"Error calculating cosine similarity: {e}") + return 0.0 + + async def search_tools( + self, + tool_repository: ConcurrentToolRepository, + query: str, + limit: int = 10, + any_of_tags_required: Optional[List[str]] = None + ) -> List[Tool]: + """Search for tools using semantic similarity. + + Args: + tool_repository: The tool repository to search within. + query: The search query string. + limit: Maximum number of tools to return. + any_of_tags_required: Optional list of tags where one of them must be present. + + Returns: + List of Tool objects ranked by semantic similarity. + """ + if limit < 0: + raise ValueError("limit must be non-negative") + + # Ensure the embedding model is loaded + await self._ensure_model_loaded() + + # Get all tools + tools: List[Tool] = await tool_repository.get_tools() + + # Filter by required tags if specified + if any_of_tags_required and len(any_of_tags_required) > 0: + any_of_tags_required = [tag.lower() for tag in any_of_tags_required] + tools = [ + tool for tool in tools + if any(tag.lower() in any_of_tags_required for tag in tool.tags) + ] + + if not tools: + return [] + + # Generate query embedding + query_embedding = await self._get_text_embedding(query) + + # Calculate similarity scores for all tools + tool_scores: List[Tuple[Tool, float]] = [] + + for tool in tools: + try: + tool_embedding = await self._get_tool_embedding(tool) + similarity = self._cosine_similarity(query_embedding, tool_embedding) + + if similarity >= self.similarity_threshold: + tool_scores.append((tool, similarity)) + + except Exception as e: + logger.warning(f"Error processing tool {tool.name}: {e}") + continue + + # Sort by similarity score (descending) + sorted_tools = [ + tool for tool, score in sorted( + tool_scores, + key=lambda x: x[1], + reverse=True + ) + ] + + # Return up to 'limit' tools + return sorted_tools[:limit] if limit > 0 else sorted_tools + + async def __aenter__(self): + """Async context manager entry.""" + await self._ensure_model_loaded() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Async context manager exit.""" + if self._executor: + self._executor.shutdown(wait=False) + + +class InMemEmbeddingsSearchStrategyConfigSerializer(Serializer[InMemEmbeddingsSearchStrategy]): + """Serializer for InMemEmbeddingsSearchStrategy configuration.""" + + def to_dict(self, obj: InMemEmbeddingsSearchStrategy) -> dict: + return obj.model_dump() + + def validate_dict(self, data: dict) -> InMemEmbeddingsSearchStrategy: + try: + return InMemEmbeddingsSearchStrategy.model_validate(data) + except Exception as e: + raise ValueError(f"Invalid configuration: {e}") from e diff --git a/plugins/tool_search/in_mem_embeddings/test_integration.py b/plugins/tool_search/in_mem_embeddings/test_integration.py new file mode 100644 index 0000000..aad11ec --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/test_integration.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Integration test to verify the plugin works with the core UTCP system.""" + +import sys +import asyncio +from pathlib import Path + +# Add paths +plugin_src = Path(__file__).parent / "src" +core_src = Path(__file__).parent.parent.parent.parent / "core" / "src" +sys.path.insert(0, str(plugin_src)) +sys.path.insert(0, str(core_src)) + +async def test_integration(): + """Test plugin integration with core system.""" + print("šŸ”— Testing Integration with Core UTCP System...") + + try: + # Test 1: Plugin registration + print("1. Testing plugin registration...") + from utcp_in_mem_embeddings import register + register() + print(" āœ… Plugin registered successfully") + + # Test 2: Core system can discover the plugin + print("2. Testing plugin discovery...") + from utcp.interfaces.tool_search_strategy import ToolSearchStrategyConfigSerializer + strategies = ToolSearchStrategyConfigSerializer.tool_search_strategy_implementations + assert "in_mem_embeddings" in strategies + print(" āœ… Plugin discovered by core system") + + # Test 3: Create strategy through core system + print("3. Testing strategy creation through core...") + from utcp.interfaces.tool_search_strategy import ToolSearchStrategyConfigSerializer + serializer = ToolSearchStrategyConfigSerializer() + + # This should work if the plugin is properly registered + strategy_config = { + "tool_search_strategy_type": "in_mem_embeddings", + "model_name": "all-MiniLM-L6-v2", + "similarity_threshold": 0.3 + } + + strategy = serializer.validate_dict(strategy_config) + print(f" āœ… Strategy created: {strategy.tool_search_strategy_type}") + + # Test 4: Basic functionality test + print("4. Testing basic search functionality...") + from utcp.data.tool import Tool, JsonSchema + from utcp.data.call_template import CallTemplate + from utcp.implementations.in_mem_tool_repository import InMemToolRepository + + # Create sample tools + tools = [ + Tool( + name="test.tool1", + description="A test tool for cooking", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["cooking", "test"], + tool_call_template=CallTemplate( + name="test.tool1", + description="Test tool", + call_template_type="default" + ) + ) + ] + + # Create repository + repo = InMemToolRepository(tools) + + # Test search + results = await strategy.search_tools(repo, "cooking", limit=1) + print(f" āœ… Search completed, found {len(results)} results") + + print("\nšŸŽ‰ Integration test passed! Plugin works with core system.") + return True + + except Exception as e: + print(f"āŒ Integration test failed: {e}") + import traceback + traceback.print_exc() + return False + +if __name__ == "__main__": + success = asyncio.run(test_integration()) + sys.exit(0 if success else 1) diff --git a/plugins/tool_search/in_mem_embeddings/test_performance.py b/plugins/tool_search/in_mem_embeddings/test_performance.py new file mode 100644 index 0000000..d5189f6 --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/test_performance.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Performance test for the in-memory embeddings plugin.""" + +import sys +import asyncio +import time +from pathlib import Path + +# Add paths +plugin_src = Path(__file__).parent / "src" +core_src = Path(__file__).parent.parent.parent.parent / "core" / "src" +sys.path.insert(0, str(plugin_src)) +sys.path.insert(0, str(core_src)) + +async def test_performance(): + """Test plugin performance with multiple tools and searches.""" + print("āš” Testing Performance...") + + try: + from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategy + from utcp.data.tool import Tool, JsonSchema + from utcp.data.call_template import CallTemplate + + # Create strategy + strategy = InMemEmbeddingsSearchStrategy( + model_name="all-MiniLM-L6-v2", + similarity_threshold=0.3, + max_workers=2, + cache_embeddings=True + ) + + # Create many tools + print("1. Creating 100 test tools...") + tools = [] + for i in range(100): + tool = Tool( + name=f"test.tool{i}", + description=f"Test tool {i} for various purposes like cooking, coding, data analysis", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["test", f"category{i%5}"], + tool_call_template=CallTemplate( + name=f"test.tool{i}", + description=f"Test tool {i}", + call_template_type="default" + ) + ) + tools.append(tool) + + # Mock repository + class MockRepo: + def __init__(self, tools): + self.tools = tools + async def get_tools(self): + return self.tools + + repo = MockRepo(tools) + + # Test 1: First search (cold start) + print("2. Testing cold start performance...") + start_time = time.time() + results1 = await strategy.search_tools(repo, "cooking tools", limit=10) + cold_time = time.time() - start_time + print(f" ā±ļø Cold start: {cold_time:.3f}s, found {len(results1)} results") + + # Test 2: Second search (warm cache) + print("3. Testing warm cache performance...") + start_time = time.time() + results2 = await strategy.search_tools(repo, "coding tools", limit=10) + warm_time = time.time() - start_time + print(f" ā±ļø Warm cache: {warm_time:.3f}s, found {len(results2)} results") + + # Test 3: Multiple searches + print("4. Testing multiple searches...") + queries = ["cooking", "programming", "data analysis", "testing", "utilities"] + start_time = time.time() + + for query in queries: + await strategy.search_tools(repo, query, limit=5) + + total_time = time.time() - start_time + avg_time = total_time / len(queries) + print(f" ā±ļø Average per search: {avg_time:.3f}s") + + # Performance assertions + assert cold_time < 5.0, f"Cold start too slow: {cold_time}s" + assert warm_time < 1.0, f"Warm cache too slow: {warm_time}s" + assert avg_time < 0.5, f"Average search too slow: {avg_time}s" + + print("\nšŸŽ‰ Performance test passed!") + return True + + except Exception as e: + print(f"āŒ Performance test failed: {e}") + import traceback + traceback.print_exc() + return False + +if __name__ == "__main__": + success = asyncio.run(test_performance()) + sys.exit(0 if success else 1) diff --git a/plugins/tool_search/in_mem_embeddings/test_plugin.py b/plugins/tool_search/in_mem_embeddings/test_plugin.py new file mode 100644 index 0000000..95f11cf --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/test_plugin.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Simple test script to verify the in-memory embeddings plugin works.""" + +import sys +import os +import asyncio +from pathlib import Path + +# Add the plugin source to Python path +plugin_src = Path(__file__).parent / "src" +sys.path.insert(0, str(plugin_src)) + +# Add core to path for imports +core_src = Path(__file__).parent.parent.parent.parent / "core" / "src" +sys.path.insert(0, str(core_src)) + +async def test_plugin(): + """Test the plugin functionality.""" + print("šŸ§Ŗ Testing In-Memory Embeddings Plugin...") + + try: + # Test 1: Import the plugin + print("1. Testing imports...") + from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategy + from utcp_in_mem_embeddings import register + print(" āœ… Imports successful") + + # Test 2: Create strategy instance + print("2. Testing strategy creation...") + strategy = InMemEmbeddingsSearchStrategy( + model_name="all-MiniLM-L6-v2", + similarity_threshold=0.3, + max_workers=2, + cache_embeddings=True + ) + print(f" āœ… Strategy created: {strategy.tool_search_strategy_type}") + + # Test 3: Test registration function + print("3. Testing registration...") + register() + print(" āœ… Registration function works") + + # Test 4: Test basic functionality + print("4. Testing basic functionality...") + + # Create mock tools + from utcp.data.tool import Tool, JsonSchema + from utcp.data.call_template import CallTemplate + + tools = [ + Tool( + name="cooking.spatula", + description="A kitchen utensil for flipping food", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["cooking", "kitchen"], + tool_call_template=CallTemplate( + name="cooking.spatula", + description="Spatula tool", + call_template_type="default" + ) + ), + Tool( + name="dev.code_review", + description="Review source code for quality", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["programming", "development"], + tool_call_template=CallTemplate( + name="dev.code_review", + description="Code review tool", + call_template_type="default" + ) + ) + ] + + # Create mock repository + class MockRepo: + def __init__(self, tools): + self.tools = tools + async def get_tools(self): + return self.tools + + repo = MockRepo(tools) + + # Test search + results = await strategy.search_tools(repo, "cooking utensils", limit=2) + print(f" āœ… Search completed, found {len(results)} results") + + if results: + print(f" šŸ“‹ Top result: {results[0].name}") + + print("\nšŸŽ‰ All tests passed! Plugin is working correctly.") + + except Exception as e: + print(f"āŒ Test failed: {e}") + import traceback + traceback.print_exc() + return False + + return True + +if __name__ == "__main__": + success = asyncio.run(test_plugin()) + sys.exit(0 if success else 1) diff --git a/plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py b/plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py new file mode 100644 index 0000000..5f37c61 --- /dev/null +++ b/plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py @@ -0,0 +1,343 @@ +"""Tests for the InMemEmbeddingsSearchStrategy implementation.""" +"""just test""" +import pytest +import numpy as np +import sys +from pathlib import Path +from unittest.mock import patch +from typing import List + +# Add plugin source to path +plugin_src = Path(__file__).parent.parent / "src" +sys.path.insert(0, str(plugin_src)) + +# Add core to path +core_src = Path(__file__).parent.parent.parent.parent / "core" / "src" +sys.path.insert(0, str(core_src)) + +from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategy +from utcp.data.tool import Tool, JsonSchema +from utcp.data.call_template import CallTemplate + + +class MockToolRepository: + """Simplified mock repository for testing.""" + + def __init__(self, tools: List[Tool]): + self.tools = tools + + async def get_tools(self) -> List[Tool]: + return self.tools + + +@pytest.fixture +def sample_tools(): + """Create sample tools for testing.""" + tools = [] + + # Tool 1: Cooking related + tool1 = Tool( + name="cooking.spatula", + description="A kitchen utensil used for flipping and turning food while cooking", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["cooking", "kitchen", "utensil"], + tool_call_template=CallTemplate( + name="cooking.spatula", + description="Spatula tool", + call_template_type="default" + ) + ) + tools.append(tool1) + + # Tool 2: Programming related + tool2 = Tool( + name="dev.code_review", + description="Review and analyze source code for quality and best practices", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["programming", "development", "code"], + tool_call_template=CallTemplate( + name="dev.code_review", + description="Code review tool", + call_template_type="default" + ) + ) + tools.append(tool2) + + # Tool 3: Data analysis + tool3 = Tool( + name="data.analyze", + description="Analyze datasets and generate insights from data", + inputs=JsonSchema(), + outputs=JsonSchema(), + tags=["data", "analysis", "insights"], + tool_call_template=CallTemplate( + name="data.analyze", + description="Data analysis tool", + call_template_type="default" + ) + ) + tools.append(tool3) + + return tools + + +@pytest.fixture +def in_mem_embeddings_strategy(): + """Create an in-memory embeddings search strategy instance.""" + return InMemEmbeddingsSearchStrategy( + model_name="all-MiniLM-L6-v2", + similarity_threshold=0.3, + max_workers=2, + cache_embeddings=True + ) + + +@pytest.mark.asyncio +async def test_in_mem_embeddings_strategy_initialization(in_mem_embeddings_strategy): + """Test that the in-memory embeddings strategy initializes correctly.""" + assert in_mem_embeddings_strategy.tool_search_strategy_type == "in_mem_embeddings" + assert in_mem_embeddings_strategy.model_name == "all-MiniLM-L6-v2" + assert in_mem_embeddings_strategy.similarity_threshold == 0.3 + assert in_mem_embeddings_strategy.max_workers == 2 + assert in_mem_embeddings_strategy.cache_embeddings is True + + +@pytest.mark.asyncio +async def test_simple_text_embedding_fallback(in_mem_embeddings_strategy): + """Test the fallback text embedding when sentence-transformers is not available.""" + # Mock the embedding model to be None to trigger fallback + in_mem_embeddings_strategy._embedding_model = None + in_mem_embeddings_strategy._model_loaded = True + + text = "test text" + embedding = await in_mem_embeddings_strategy._get_text_embedding(text) + + assert isinstance(embedding, np.ndarray) + assert embedding.shape == (384,) + assert np.linalg.norm(embedding) > 0 + + +@pytest.mark.asyncio +async def test_cosine_similarity_calculation(in_mem_embeddings_strategy): + """Test cosine similarity calculation.""" + # Test with identical vectors + vec1 = np.array([1.0, 0.0, 0.0]) + vec2 = np.array([1.0, 0.0, 0.0]) + similarity = in_mem_embeddings_strategy._cosine_similarity(vec1, vec2) + assert similarity == pytest.approx(1.0) + + # Test with orthogonal vectors + vec3 = np.array([0.0, 1.0, 0.0]) + similarity = in_mem_embeddings_strategy._cosine_similarity(vec1, vec3) + assert similarity == pytest.approx(0.0) + + # Test with zero vectors + vec4 = np.zeros(3) + similarity = in_mem_embeddings_strategy._cosine_similarity(vec1, vec4) + assert similarity == 0.0 + + +@pytest.mark.asyncio +async def test_tool_embedding_generation(in_mem_embeddings_strategy, sample_tools): + """Test that tool embeddings are generated and cached correctly.""" + tool = sample_tools[0] + + # Mock the text embedding method + with patch.object(in_mem_embeddings_strategy, '_get_text_embedding') as mock_embed: + mock_embed.return_value = np.random.rand(384) + + # First call should generate and cache + embedding1 = await in_mem_embeddings_strategy._get_tool_embedding(tool) + assert tool.name in in_mem_embeddings_strategy._tool_embeddings_cache + + # Second call should use cache + embedding2 = await in_mem_embeddings_strategy._get_tool_embedding(tool) + assert np.array_equal(embedding1, embedding2) + + # Verify the mock was called only once + mock_embed.assert_called_once() + + +@pytest.mark.asyncio +async def test_search_tools_basic(in_mem_embeddings_strategy, sample_tools): + """Test basic search functionality.""" + tool_repo = MockToolRepository(sample_tools) + + # Mock the embedding methods + with patch.object(in_mem_embeddings_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(in_mem_embeddings_strategy, '_get_tool_embedding') as mock_tool_embed: + + # Create mock embeddings + query_embedding = np.random.rand(384) + tool_embeddings = [np.random.rand(384) for _ in sample_tools] + + mock_query_embed.return_value = query_embedding + mock_tool_embed.side_effect = tool_embeddings + + # Mock cosine similarity to return high scores + with patch.object(in_mem_embeddings_strategy, '_cosine_similarity') as mock_sim: + mock_sim.return_value = 0.8 # High similarity + + results = await in_mem_embeddings_strategy.search_tools(tool_repo, "cooking", limit=2) + + assert len(results) == 2 + assert all(isinstance(tool, Tool) for tool in results) + + +@pytest.mark.asyncio +async def test_search_tools_with_tag_filtering(in_mem_embeddings_strategy, sample_tools): + """Test search with tag filtering.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(in_mem_embeddings_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(in_mem_embeddings_strategy, '_get_tool_embedding') as mock_tool_embed, \ + patch.object(in_mem_embeddings_strategy, '_cosine_similarity') as mock_sim: + + mock_query_embed.return_value = np.random.rand(384) + mock_tool_embed.return_value = np.random.rand(384) + mock_sim.return_value = 0.8 + + # Search with required tags + results = await in_mem_embeddings_strategy.search_tools( + tool_repo, + "cooking", + limit=10, + any_of_tags_required=["cooking", "kitchen"] + ) + + # Should only return tools with cooking or kitchen tags + assert all( + any(tag in ["cooking", "kitchen"] for tag in tool.tags) + for tool in results + ) + + +@pytest.mark.asyncio +async def test_search_tools_with_similarity_threshold(in_mem_embeddings_strategy, sample_tools): + """Test that similarity threshold filtering works correctly.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(in_mem_embeddings_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(in_mem_embeddings_strategy, '_get_tool_embedding') as mock_tool_embed, \ + patch.object(in_mem_embeddings_strategy, '_cosine_similarity') as mock_sim: + + mock_query_embed.return_value = np.random.rand(384) + mock_tool_embed.return_value = np.random.rand(384) + + # Set threshold to 0.5 and return scores below and above + in_mem_embeddings_strategy.similarity_threshold = 0.5 + mock_sim.side_effect = [0.3, 0.7, 0.2] # Only second tool should pass + + results = await in_mem_embeddings_strategy.search_tools(tool_repo, "test", limit=10) + + assert len(results) == 1 # Only one tool above threshold + + +@pytest.mark.asyncio +async def test_search_tools_limit_respected(in_mem_embeddings_strategy, sample_tools): + """Test that the limit parameter is respected.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(in_mem_embeddings_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(in_mem_embeddings_strategy, '_get_tool_embedding') as mock_tool_embed, \ + patch.object(in_mem_embeddings_strategy, '_cosine_similarity') as mock_sim: + + mock_query_embed.return_value = np.random.rand(384) + mock_tool_embed.return_value = np.random.rand(384) + mock_sim.return_value = 0.8 + + # Test with limit 1 + results = await in_mem_embeddings_strategy.search_tools(tool_repo, "test", limit=1) + assert len(results) == 1 + + # Test with limit 0 (no limit) + results = await in_mem_embeddings_strategy.search_tools(tool_repo, "test", limit=0) + assert len(results) == 3 # All tools + + +@pytest.mark.asyncio +async def test_search_tools_empty_repository(in_mem_embeddings_strategy): + """Test search behavior with empty tool repository.""" + tool_repo = MockToolRepository([]) + + results = await in_mem_embeddings_strategy.search_tools(tool_repo, "test", limit=10) + assert results == [] + + +@pytest.mark.asyncio +async def test_search_tools_invalid_limit(in_mem_embeddings_strategy, sample_tools): + """Test that invalid limit values raise appropriate errors.""" + tool_repo = MockToolRepository(sample_tools) + + with pytest.raises(ValueError, match="limit must be non-negative"): + await in_mem_embeddings_strategy.search_tools(tool_repo, "test", limit=-1) + + +@pytest.mark.asyncio +async def test_context_manager_behavior(in_mem_embeddings_strategy): + """Test async context manager behavior.""" + async with in_mem_embeddings_strategy as strategy: + assert strategy._model_loaded is True + + # Executor should be shut down + assert strategy._executor._shutdown is True + + +@pytest.mark.asyncio +async def test_error_handling_in_search(in_mem_embeddings_strategy, sample_tools): + """Test that errors in search are handled gracefully.""" + tool_repo = MockToolRepository(sample_tools) + + with patch.object(in_mem_embeddings_strategy, '_get_text_embedding') as mock_query_embed, \ + patch.object(in_mem_embeddings_strategy, '_get_tool_embedding') as mock_tool_embed: + + mock_query_embed.return_value = np.random.rand(384) + + # Make the second tool fail + def mock_tool_embed_side_effect(tool): + if tool.name == "dev.code_review": + raise Exception("Simulated error") + return np.random.rand(384) + + mock_tool_embed.side_effect = mock_tool_embed_side_effect + + # Mock cosine similarity + with patch.object(in_mem_embeddings_strategy, '_cosine_similarity') as mock_sim: + mock_sim.return_value = 0.8 + + # Should not crash, just skip the problematic tool + results = await in_mem_embeddings_strategy.search_tools(tool_repo, "test", limit=10) + + # Should return tools that didn't fail + assert len(results) == 2 # One tool failed, so only 2 results + + +@pytest.mark.asyncio +async def test_in_mem_embeddings_strategy_config_serializer(): + """Test the configuration serializer.""" + from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategyConfigSerializer + + serializer = InMemEmbeddingsSearchStrategyConfigSerializer() + + # Test serialization + strategy = InMemEmbeddingsSearchStrategy( + model_name="test-model", + similarity_threshold=0.5, + max_workers=8, + cache_embeddings=False + ) + + config_dict = serializer.to_dict(strategy) + assert config_dict["model_name"] == "test-model" + assert config_dict["similarity_threshold"] == 0.5 + assert config_dict["max_workers"] == 8 + assert config_dict["cache_embeddings"] is False + + # Test deserialization + restored_strategy = serializer.validate_dict(config_dict) + assert restored_strategy.model_name == "test-model" + assert restored_strategy.similarity_threshold == 0.5 + assert restored_strategy.max_workers == 8 + assert restored_strategy.cache_embeddings is False From ba3e1fcde5fcff452d2b06457d8fedf8bcfe1eb1 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Thu, 4 Sep 2025 12:36:10 -0500 Subject: [PATCH 5/8] folder structure to be resolved --- core/EMBEDDING_SEARCH_README.md | 278 --------------- core/examples/embedding_search_example.py | 262 -------------- core/src/utcp/implementations/__init__.py | 2 - .../utcp/implementations/embedding_search.py | 229 ------------ core/src/utcp/plugins/plugin_loader.py | 4 +- .../implementations/test_embedding_search.py | 337 ------------------ plugins/tool_search/embedding/README.md | 19 - plugins/tool_search/embedding/pyproject.toml | 39 -- .../src/utcp_embedding_search/__init__.py | 7 - 9 files changed, 2 insertions(+), 1175 deletions(-) delete mode 100644 core/EMBEDDING_SEARCH_README.md delete mode 100644 core/examples/embedding_search_example.py delete mode 100644 core/src/utcp/implementations/embedding_search.py delete mode 100644 core/tests/implementations/test_embedding_search.py delete mode 100644 plugins/tool_search/embedding/README.md delete mode 100644 plugins/tool_search/embedding/pyproject.toml delete mode 100644 plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py diff --git a/core/EMBEDDING_SEARCH_README.md b/core/EMBEDDING_SEARCH_README.md deleted file mode 100644 index 55327af..0000000 --- a/core/EMBEDDING_SEARCH_README.md +++ /dev/null @@ -1,278 +0,0 @@ -# UTCP Embedding Search Plugin - -This document describes the new embedding-based semantic search strategy for UTCP tools, which provides intelligent tool discovery based on meaning similarity rather than just keyword matching. - -## Overview - -The `EmbeddingSearchStrategy` is a plugin that implements semantic search for UTCP tools using sentence embeddings. It converts tool descriptions and search queries into numerical vectors and finds the most semantically similar tools using cosine similarity. - -## Features - -- **Semantic Understanding**: Finds tools based on meaning, not just exact keyword matches -- **Configurable Similarity Threshold**: Adjustable threshold for result quality -- **Automatic Fallback**: Falls back to simple text similarity if sentence-transformers is unavailable -- **Embedding Caching**: Caches tool embeddings for improved performance -- **Tag Filtering**: Supports filtering results by required tags -- **Async Support**: Fully asynchronous implementation for non-blocking operations -- **Context Manager**: Proper resource management with async context manager - -## Installation - -### Basic Installation - -The core functionality is available with the basic dependencies: - -```bash -pip install utcp -``` - -### Enhanced Semantic Search - -For the best semantic search experience, install the optional embedding dependencies: - -```bash -pip install utcp[embedding] -``` - -This installs: -- `sentence-transformers>=2.2.0` - For high-quality sentence embeddings -- `torch>=1.9.0` - PyTorch backend for sentence-transformers - -## Quick Start - -### Basic Usage - -```python -import asyncio -from utcp.implementations.embedding_search import EmbeddingSearchStrategy -from utcp.implementations.in_mem_tool_repository import InMemToolRepository - -async def main(): - # Create a tool repository with some tools - tool_repo = InMemToolRepository() - # ... add tools to repository ... - - # Create the embedding search strategy - strategy = EmbeddingSearchStrategy( - model_name="all-MiniLM-L6-v2", - similarity_threshold=0.3, - max_workers=4, - cache_embeddings=True - ) - - # Search for tools - results = await strategy.search_tools( - tool_repo, - "I need to process some data", - limit=5 - ) - - for tool in results: - print(f"Found: {tool.name} - {tool.description}") - -asyncio.run(main()) -``` - -### Using as a Context Manager - -```python -async with EmbeddingSearchStrategy() as strategy: - results = await strategy.search_tools(tool_repo, "cooking tools", limit=3) - # Strategy automatically manages resources -``` - -## Configuration Options - -### EmbeddingSearchStrategy Parameters - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `model_name` | str | "all-MiniLM-L6-v2" | Sentence transformer model to use | -| `similarity_threshold` | float | 0.3 | Minimum similarity score (0.0-1.0) | -| `max_workers` | int | 4 | Maximum worker threads for embedding generation | -| `cache_embeddings` | bool | True | Whether to cache tool embeddings | - -### Recommended Similarity Thresholds - -- **0.2-0.3**: Broad matches, more results -- **0.4-0.5**: Balanced precision and recall -- **0.6-0.7**: High precision, fewer results -- **0.8+**: Very high precision, may miss relevant tools - -## How It Works - -### 1. Text to Embedding Conversion - -The strategy converts text into numerical embeddings: - -```python -# Query: "I need to analyze data" -# Gets converted to: [0.1, -0.3, 0.8, ...] (384-dimensional vector) - -# Tool description: "Analyze datasets and generate insights" -# Gets converted to: [0.2, -0.1, 0.9, ...] (384-dimensional vector) -``` - -### 2. Similarity Calculation - -Uses cosine similarity to measure how similar two embeddings are: - -```python -similarity = dot_product(embedding1, embedding2) / (norm(embedding1) * norm(embedding2)) -``` - -### 3. Result Ranking - -Tools are ranked by similarity score and filtered by the threshold. - -## Advanced Usage - -### Tag Filtering - -```python -# Only return tools with specific tags -results = await strategy.search_tools( - tool_repo, - "data processing", - limit=5, - any_of_tags_required=["data", "analysis"] -) -``` - -### Custom Model Selection - -```python -# Use a different sentence transformer model -strategy = EmbeddingSearchStrategy( - model_name="paraphrase-multilingual-MiniLM-L12-v2", # Multilingual support - similarity_threshold=0.4 -) -``` - -### Performance Tuning - -```python -# Optimize for your use case -strategy = EmbeddingSearchStrategy( - max_workers=8, # More workers for faster processing - cache_embeddings=True, # Cache for repeated searches - similarity_threshold=0.5 # Higher threshold for quality -) -``` - -## Fallback Behavior - -If `sentence-transformers` is not available, the strategy automatically falls back to a simple text similarity approach: - -- Uses character frequency-based embeddings -- Maintains the same API -- Provides reasonable results for basic use cases - -## Integration with UTCP - -The embedding search strategy integrates seamlessly with the UTCP plugin system: - -```python -# The strategy is automatically registered when the module is imported -from utcp.implementations.embedding_search import EmbeddingSearchStrategy - -# Use it in your UTCP client configuration -config = UtcpClientConfig( - tool_search_strategy=EmbeddingSearchStrategy( - similarity_threshold=0.4 - ) -) -``` - -## Performance Considerations - -### Memory Usage - -- Each tool embedding uses ~1.5KB of memory (384 dimensions Ɨ 4 bytes) -- With caching enabled, memory usage grows with the number of tools -- Consider disabling caching for very large tool repositories - -### Processing Speed - -- First search: Slower due to model loading and embedding generation -- Subsequent searches: Faster due to caching -- More workers = faster embedding generation but higher memory usage - -### Model Loading - -- Models are downloaded on first use (~80MB for all-MiniLM-L6-v2) -- Consider pre-downloading models in production environments - -## Troubleshooting - -### Common Issues - -1. **Import Error for sentence-transformers** - - Install with: `pip install sentence-transformers` - - The strategy will fall back to simple text similarity - -2. **Slow First Search** - - This is normal - the model needs to load - - Subsequent searches will be faster - -3. **Memory Issues** - - Reduce `max_workers` - - Disable `cache_embeddings` - - Use a smaller model (e.g., "all-MiniLM-L6-v2" instead of larger models) - -4. **Low Quality Results** - - Adjust `similarity_threshold` - - Ensure tool descriptions are detailed and meaningful - - Consider using more specific queries - -### Debug Mode - -Enable debug logging to see what's happening: - -```python -import logging -logging.basicConfig(level=logging.DEBUG) - -# The strategy will log detailed information about: -# - Model loading -# - Embedding generation -# - Similarity calculations -# - Search results -``` - -## Examples - -See `core/examples/embedding_search_example.py` for comprehensive examples demonstrating: - -- Basic search functionality -- Tag filtering -- Similarity threshold adjustment -- Context manager usage -- Comparison with tag-based search - -## Contributing - -To contribute to the embedding search plugin: - -1. Fork the repository -2. Create a feature branch -3. Add tests for new functionality -4. Ensure all tests pass -5. Submit a pull request - -## License - -This plugin is part of the UTCP project and follows the same license terms. - -## Support - -For issues and questions: - -1. Check the troubleshooting section above -2. Review the example code -3. Open an issue on the GitHub repository -4. Check the UTCP documentation - ---- - -**Note**: The embedding search plugin requires Python 3.10+ and is designed to work seamlessly with the existing UTCP ecosystem while providing enhanced semantic search capabilities. diff --git a/core/examples/embedding_search_example.py b/core/examples/embedding_search_example.py deleted file mode 100644 index 6183e61..0000000 --- a/core/examples/embedding_search_example.py +++ /dev/null @@ -1,262 +0,0 @@ -#!/usr/bin/env python3 -""" -Example demonstrating the embedding search strategy for UTCP tools. - -This example shows how to: -1. Create tools with descriptions and tags -2. Use the embedding search strategy to find semantically similar tools -3. Compare results with the traditional tag-based search -""" - -import asyncio -import sys -import os - -# Add the src directory to the path so we can import utcp modules -sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) - -from utcp.data.tool import Tool, JsonSchema -from utcp.data.call_template import CallTemplate -from utcp.implementations.embedding_search import EmbeddingSearchStrategy -from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategy -from utcp.implementations.in_mem_tool_repository import InMemToolRepository - - -async def create_sample_tools(): - """Create a collection of sample tools for demonstration.""" - tools = [] - - # Cooking tools - tools.append(Tool( - name="cooking.spatula", - description="A kitchen utensil used for flipping and turning food while cooking", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["cooking", "kitchen", "utensil"], - tool_call_template=CallTemplate(name="cooking.spatula", description="Spatula tool") - )) - - tools.append(Tool( - name="cooking.whisk", - description="A kitchen tool for mixing and aerating ingredients", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["cooking", "kitchen", "mixing"], - tool_call_template=CallTemplate(name="cooking.whisk", description="Whisk tool") - )) - - tools.append(Tool( - name="cooking.knife", - description="A sharp blade for cutting and chopping food ingredients", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["cooking", "kitchen", "cutting"], - tool_call_template=CallTemplate(name="cooking.knife", description="Knife tool") - )) - - # Programming tools - tools.append(Tool( - name="dev.code_review", - description="Review and analyze source code for quality and best practices", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["programming", "development", "code"], - tool_call_template=CallTemplate(name="dev.code_review", description="Code review tool") - )) - - tools.append(Tool( - name="dev.debug", - description="Find and fix bugs in software code", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["programming", "development", "debugging"], - tool_call_template=CallTemplate(name="dev.debug", description="Debugging tool") - )) - - tools.append(Tool( - name="dev.test", - description="Run automated tests to verify code functionality", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["programming", "development", "testing"], - tool_call_template=CallTemplate(name="dev.test", description="Testing tool") - )) - - # Data analysis tools - tools.append(Tool( - name="data.analyze", - description="Analyze datasets and generate insights from data", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["data", "analysis", "insights"], - tool_call_template=CallTemplate(name="data.analyze", description="Data analysis tool") - )) - - tools.append(Tool( - name="data.visualize", - description="Create charts and graphs to represent data visually", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["data", "visualization", "charts"], - tool_call_template=CallTemplate(name="data.visualize", description="Data visualization tool") - )) - - tools.append(Tool( - name="data.clean", - description="Clean and preprocess raw data for analysis", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["data", "cleaning", "preprocessing"], - tool_call_template=CallTemplate(name="data.clean", description="Data cleaning tool") - )) - - return tools - - -async def demonstrate_search_strategies(): - """Demonstrate both search strategies with example queries.""" - - # Create tools and repository - tools = await create_sample_tools() - tool_repo = InMemToolRepository() - - # Add tools to repository - for tool in tools: - await tool_repo.save_tool(tool) - - # Create search strategies - embedding_strategy = EmbeddingSearchStrategy( - model_name="all-MiniLM-L6-v2", - similarity_threshold=0.3, - max_workers=2, - cache_embeddings=True - ) - - tag_strategy = TagAndDescriptionWordMatchStrategy() - - # Example queries to test - test_queries = [ - "I need to cook something", - "Help me write better code", - "I have data to analyze", - "Kitchen equipment", - "Software development", - "Data science tasks" - ] - - print("šŸ” UTCP Embedding Search Strategy Demo") - print("=" * 50) - print() - - for query in test_queries: - print(f"Query: '{query}'") - print("-" * 30) - - # Search with embedding strategy - print("šŸ“Š Embedding Search Results:") - try: - embedding_results = await embedding_strategy.search_tools( - tool_repo, query, limit=3 - ) - for i, tool in enumerate(embedding_results, 1): - print(f" {i}. {tool.name} (tags: {', '.join(tool.tags)})") - print(f" {tool.description}") - except Exception as e: - print(f" Error: {e}") - - print() - - # Search with tag strategy - print("šŸ·ļø Tag-based Search Results:") - try: - tag_results = await tag_strategy.search_tools( - tool_repo, query, limit=3 - ) - for i, tool in enumerate(tag_results, 1): - print(f" {i}. {tool.name} (tags: {', '.join(tool.tags)})") - print(f" {tool.description}") - except Exception as e: - print(f" Error: {e}") - - print("\n" + "=" * 50 + "\n") - - -async def demonstrate_advanced_features(): - """Demonstrate advanced features of the embedding search strategy.""" - - print("šŸš€ Advanced Embedding Search Features") - print("=" * 50) - print() - - # Create tools and repository - tools = await create_sample_tools() - tool_repo = InMemToolRepository() - - for tool in tools: - await tool_repo.save_tool(tool) - - # Create strategy with custom configuration - strategy = EmbeddingSearchStrategy( - model_name="all-MiniLM-L6-v2", - similarity_threshold=0.5, # Higher threshold for more precise matches - max_workers=4, - cache_embeddings=True - ) - - # Test tag filtering - print("1. Tag Filtering Example:") - print("Query: 'cooking tools' with required tags: ['cooking']") - results = await strategy.search_tools( - tool_repo, - "cooking tools", - limit=5, - any_of_tags_required=["cooking"] - ) - - for i, tool in enumerate(results, 1): - print(f" {i}. {tool.name} (tags: {', '.join(tool.tags)})") - - print() - - # Test different similarity thresholds - print("2. Similarity Threshold Comparison:") - thresholds = [0.2, 0.4, 0.6, 0.8] - query = "food preparation" - - for threshold in thresholds: - strategy.similarity_threshold = threshold - results = await strategy.search_tools(tool_repo, query, limit=5) - print(f" Threshold {threshold}: {len(results)} results") - - print() - - # Test context manager usage - print("3. Context Manager Usage:") - async with strategy as ctx_strategy: - results = await ctx_strategy.search_tools(tool_repo, "software", limit=3) - print(f" Found {len(results)} software-related tools") - - print("\n" + "=" * 50) - - -async def main(): - """Main function to run the demonstration.""" - try: - await demonstrate_search_strategies() - await demonstrate_advanced_features() - - print("āœ… Demo completed successfully!") - print("\nšŸ’” Tips:") - print("- Install sentence-transformers for better semantic search: pip install sentence-transformers") - print("- Adjust similarity_threshold based on your needs (0.3-0.7 recommended)") - print("- Use tag filtering to narrow down results by category") - print("- The strategy automatically falls back to simple text similarity if sentence-transformers is not available") - - except Exception as e: - print(f"āŒ Error during demo: {e}") - import traceback - traceback.print_exc() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/core/src/utcp/implementations/__init__.py b/core/src/utcp/implementations/__init__.py index a362425..12fb6d6 100644 --- a/core/src/utcp/implementations/__init__.py +++ b/core/src/utcp/implementations/__init__.py @@ -1,9 +1,7 @@ from utcp.implementations.in_mem_tool_repository import InMemToolRepository from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategy -from utcp.implementations.embedding_search import EmbeddingSearchStrategy __all__ = [ "InMemToolRepository", "TagAndDescriptionWordMatchStrategy", - "EmbeddingSearchStrategy", ] diff --git a/core/src/utcp/implementations/embedding_search.py b/core/src/utcp/implementations/embedding_search.py deleted file mode 100644 index 4e6b244..0000000 --- a/core/src/utcp/implementations/embedding_search.py +++ /dev/null @@ -1,229 +0,0 @@ -"""Embedding-based semantic search strategy for UTCP tools. - -This module provides a semantic search implementation that uses sentence embeddings -to find tools based on meaning similarity rather than just keyword matching. -""" - -import asyncio -import logging -from typing import List, Tuple, Optional, Literal, Dict, Any -from concurrent.futures import ThreadPoolExecutor -import numpy as np -from pydantic import BaseModel, Field - -from utcp.interfaces.tool_search_strategy import ToolSearchStrategy -from utcp.data.tool import Tool -from utcp.interfaces.concurrent_tool_repository import ConcurrentToolRepository -from utcp.interfaces.serializer import Serializer - -logger = logging.getLogger(__name__) - -class EmbeddingSearchStrategy(ToolSearchStrategy): - """Semantic search strategy using sentence embeddings. - - This strategy converts tool descriptions and search queries into numerical - embeddings and finds the most semantically similar tools using cosine similarity. - """ - - tool_search_strategy_type: Literal["embedding_search"] = "embedding_search" - - # Configuration parameters - model_name: str = Field(default="all-MiniLM-L6-v2", description="Sentence transformer model to use") - similarity_threshold: float = Field(default=0.3, description="Minimum similarity score to consider a match") - max_workers: int = Field(default=4, description="Maximum number of worker threads for embedding generation") - cache_embeddings: bool = Field(default=True, description="Whether to cache tool embeddings for performance") - - def __init__(self, **data): - super().__init__(**data) - self._embedding_model = None - self._tool_embeddings_cache: Dict[str, np.ndarray] = {} - self._executor = ThreadPoolExecutor(max_workers=self.max_workers) - self._model_loaded = False - - async def _ensure_model_loaded(self): - """Ensure the embedding model is loaded.""" - if self._model_loaded: - return - - try: - # Import sentence-transformers here to avoid dependency issues - from sentence_transformers import SentenceTransformer - - # Load the model in a thread to avoid blocking - loop = asyncio.get_event_loop() - self._embedding_model = await loop.run_in_executor( - self._executor, - SentenceTransformer, - self.model_name - ) - self._model_loaded = True - logger.info(f"Loaded embedding model: {self.model_name}") - - except ImportError: - logger.warning("sentence-transformers not available, falling back to simple text similarity") - self._embedding_model = None - self._model_loaded = True - except Exception as e: - logger.error(f"Failed to load embedding model: {e}") - self._embedding_model = None - self._model_loaded = True - - async def _get_text_embedding(self, text: str) -> np.ndarray: - """Generate embedding for given text.""" - if not text: - return np.zeros(384) # Default dimension for all-MiniLM-L6-v2 - - if self._embedding_model is None: - # Fallback to simple text similarity - return self._simple_text_embedding(text) - - try: - loop = asyncio.get_event_loop() - embedding = await loop.run_in_executor( - self._executor, - self._embedding_model.encode, - text - ) - return embedding - except Exception as e: - logger.warning(f"Failed to generate embedding for text: {e}") - return self._simple_text_embedding(text) - - def _simple_text_embedding(self, text: str) -> np.ndarray: - """Simple fallback embedding using character frequency.""" - # Create a simple embedding based on character frequency - # This is a fallback when sentence-transformers is not available - embedding = np.zeros(384) - text_lower = text.lower() - - # Simple character frequency-based embedding - for i, char in enumerate(text_lower): - if i < 384: - embedding[i % 384] += ord(char) / 1000.0 - - # Normalize - norm = np.linalg.norm(embedding) - if norm > 0: - embedding = embedding / norm - - return embedding - - async def _get_tool_embedding(self, tool: Tool) -> np.ndarray: - """Get or generate embedding for a tool.""" - if not self.cache_embeddings or tool.name not in self._tool_embeddings_cache: - # Create text representation of the tool - tool_text = f"{tool.name} {tool.description} {' '.join(tool.tags)}" - embedding = await self._get_text_embedding(tool_text) - - if self.cache_embeddings: - self._tool_embeddings_cache[tool.name] = embedding - - return embedding - - return self._tool_embeddings_cache[tool.name] - - def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float: - """Calculate cosine similarity between two vectors.""" - try: - dot_product = np.dot(a, b) - norm_a = np.linalg.norm(a) - norm_b = np.linalg.norm(b) - - if norm_a == 0 or norm_b == 0: - return 0.0 - - return dot_product / (norm_a * norm_b) - except Exception as e: - logger.warning(f"Error calculating cosine similarity: {e}") - return 0.0 - - async def search_tools( - self, - tool_repository: ConcurrentToolRepository, - query: str, - limit: int = 10, - any_of_tags_required: Optional[List[str]] = None - ) -> List[Tool]: - """Search for tools using semantic similarity. - - Args: - tool_repository: The tool repository to search within. - query: The search query string. - limit: Maximum number of tools to return. - any_of_tags_required: Optional list of tags where one of them must be present. - - Returns: - List of Tool objects ranked by semantic similarity. - """ - if limit < 0: - raise ValueError("limit must be non-negative") - - # Ensure the embedding model is loaded - await self._ensure_model_loaded() - - # Get all tools - tools: List[Tool] = await tool_repository.get_tools() - - # Filter by required tags if specified - if any_of_tags_required and len(any_of_tags_required) > 0: - any_of_tags_required = [tag.lower() for tag in any_of_tags_required] - tools = [ - tool for tool in tools - if any(tag.lower() in any_of_tags_required for tag in tool.tags) - ] - - if not tools: - return [] - - # Generate query embedding - query_embedding = await self._get_text_embedding(query) - - # Calculate similarity scores for all tools - tool_scores: List[Tuple[Tool, float]] = [] - - for tool in tools: - try: - tool_embedding = await self._get_tool_embedding(tool) - similarity = self._cosine_similarity(query_embedding, tool_embedding) - - if similarity >= self.similarity_threshold: - tool_scores.append((tool, similarity)) - - except Exception as e: - logger.warning(f"Error processing tool {tool.name}: {e}") - continue - - # Sort by similarity score (descending) - sorted_tools = [ - tool for tool, score in sorted( - tool_scores, - key=lambda x: x[1], - reverse=True - ) - ] - - # Return up to 'limit' tools - return sorted_tools[:limit] if limit > 0 else sorted_tools - - async def __aenter__(self): - """Async context manager entry.""" - await self._ensure_model_loaded() - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - """Async context manager exit.""" - if self._executor: - self._executor.shutdown(wait=False) - - -class EmbeddingSearchStrategyConfigSerializer(Serializer[EmbeddingSearchStrategy]): - """Serializer for EmbeddingSearchStrategy configuration.""" - - def to_dict(self, obj: EmbeddingSearchStrategy) -> dict: - return obj.model_dump() - - def validate_dict(self, data: dict) -> EmbeddingSearchStrategy: - try: - return EmbeddingSearchStrategy.model_validate(data) - except Exception as e: - raise ValueError(f"Invalid configuration: {e}") from e diff --git a/core/src/utcp/plugins/plugin_loader.py b/core/src/utcp/plugins/plugin_loader.py index 1695f55..d13b02f 100644 --- a/core/src/utcp/plugins/plugin_loader.py +++ b/core/src/utcp/plugins/plugin_loader.py @@ -6,7 +6,7 @@ def _load_plugins(): from utcp.interfaces.tool_search_strategy import ToolSearchStrategyConfigSerializer from utcp.implementations.in_mem_tool_repository import InMemToolRepositoryConfigSerializer from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategyConfigSerializer - from utcp.implementations.embedding_search import EmbeddingSearchStrategyConfigSerializer + from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategyConfigSerializer from utcp.data.auth_implementations import OAuth2AuthSerializer, BasicAuthSerializer, ApiKeyAuthSerializer from utcp.data.variable_loader_implementations import DotEnvVariableLoaderSerializer from utcp.implementations.post_processors import FilterDictPostProcessorConfigSerializer, LimitStringsPostProcessorConfigSerializer @@ -20,7 +20,7 @@ def _load_plugins(): register_tool_repository(ConcurrentToolRepositoryConfigSerializer.default_repository, InMemToolRepositoryConfigSerializer()) register_tool_search_strategy(ToolSearchStrategyConfigSerializer.default_strategy, TagAndDescriptionWordMatchStrategyConfigSerializer()) - register_tool_search_strategy("embedding_search", EmbeddingSearchStrategyConfigSerializer()) + register_tool_search_strategy("in_mem_embeddings", InMemEmbeddingsSearchStrategyConfigSerializer()) register_tool_post_processor("filter_dict", FilterDictPostProcessorConfigSerializer()) register_tool_post_processor("limit_strings", LimitStringsPostProcessorConfigSerializer()) diff --git a/core/tests/implementations/test_embedding_search.py b/core/tests/implementations/test_embedding_search.py deleted file mode 100644 index 334100e..0000000 --- a/core/tests/implementations/test_embedding_search.py +++ /dev/null @@ -1,337 +0,0 @@ -"""Tests for the EmbeddingSearchStrategy implementation.""" - -import pytest -import numpy as np -from unittest.mock import patch -from typing import List - -from utcp.implementations.embedding_search import EmbeddingSearchStrategy -from utcp.data.tool import Tool, JsonSchema -from utcp.data.call_template import CallTemplate - - -class MockToolRepository: - """Simplified mock repository for testing.""" - - def __init__(self, tools: List[Tool]): - self.tools = tools - - async def get_tools(self) -> List[Tool]: - return self.tools - - -@pytest.fixture -def sample_tools(): - """Create sample tools for testing.""" - tools = [] - - # Tool 1: Cooking related - tool1 = Tool( - name="cooking.spatula", - description="A kitchen utensil used for flipping and turning food while cooking", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["cooking", "kitchen", "utensil"], - tool_call_template=CallTemplate( - name="cooking.spatula", - description="Spatula tool", - call_template_type="default" - ) - ) - tools.append(tool1) - - # Tool 2: Programming related - tool2 = Tool( - name="dev.code_review", - description="Review and analyze source code for quality and best practices", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["programming", "development", "code"], - tool_call_template=CallTemplate( - name="dev.code_review", - description="Code review tool", - call_template_type="default" - ) - ) - tools.append(tool2) - - # Tool 3: Data analysis - tool3 = Tool( - name="data.analyze", - description="Analyze datasets and generate insights from data", - inputs=JsonSchema(), - outputs=JsonSchema(), - tags=["data", "analysis", "insights"], - tool_call_template=CallTemplate( - name="data.analyze", - description="Data analysis tool", - call_template_type="default" - ) - ) - tools.append(tool3) - - return tools - - -@pytest.fixture -def embedding_strategy(): - """Create an embedding search strategy instance.""" - return EmbeddingSearchStrategy( - model_name="all-MiniLM-L6-v2", - similarity_threshold=0.3, - max_workers=2, - cache_embeddings=True - ) - - -# --- Your existing tests remain unchanged below this line --- - - -@pytest.mark.asyncio - -async def test_embedding_strategy_initialization(embedding_strategy): - """Test that the embedding strategy initializes correctly.""" - assert embedding_strategy.tool_search_strategy_type == "embedding_search" - assert embedding_strategy.model_name == "all-MiniLM-L6-v2" - assert embedding_strategy.similarity_threshold == 0.3 - assert embedding_strategy.max_workers == 2 - assert embedding_strategy.cache_embeddings is True - - -@pytest.mark.asyncio -async def test_simple_text_embedding_fallback(embedding_strategy): - """Test the fallback text embedding when sentence-transformers is not available.""" - # Mock the embedding model to be None to trigger fallback - embedding_strategy._embedding_model = None - embedding_strategy._model_loaded = True - - text = "test text" - embedding = await embedding_strategy._get_text_embedding(text) - - assert isinstance(embedding, np.ndarray) - assert embedding.shape == (384,) - assert np.linalg.norm(embedding) > 0 - - -@pytest.mark.asyncio -async def test_cosine_similarity_calculation(embedding_strategy): - """Test cosine similarity calculation.""" - # Test with identical vectors - vec1 = np.array([1.0, 0.0, 0.0]) - vec2 = np.array([1.0, 0.0, 0.0]) - similarity = embedding_strategy._cosine_similarity(vec1, vec2) - assert similarity == pytest.approx(1.0) - - # Test with orthogonal vectors - vec3 = np.array([0.0, 1.0, 0.0]) - similarity = embedding_strategy._cosine_similarity(vec1, vec3) - assert similarity == pytest.approx(0.0) - - # Test with zero vectors - vec4 = np.zeros(3) - similarity = embedding_strategy._cosine_similarity(vec1, vec4) - assert similarity == 0.0 - - -@pytest.mark.asyncio -async def test_tool_embedding_generation(embedding_strategy, sample_tools): - """Test that tool embeddings are generated and cached correctly.""" - tool = sample_tools[0] - - # Mock the text embedding method - with patch.object(embedding_strategy, '_get_text_embedding') as mock_embed: - mock_embed.return_value = np.random.rand(384) - - # First call should generate and cache - embedding1 = await embedding_strategy._get_tool_embedding(tool) - assert tool.name in embedding_strategy._tool_embeddings_cache - - # Second call should use cache - embedding2 = await embedding_strategy._get_tool_embedding(tool) - assert np.array_equal(embedding1, embedding2) - - # Verify the mock was called only once - mock_embed.assert_called_once() - - -@pytest.mark.asyncio -async def test_search_tools_basic(embedding_strategy, sample_tools): - """Test basic search functionality.""" - tool_repo = MockToolRepository(sample_tools) - - # Mock the embedding methods - with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ - patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed: - - # Create mock embeddings - query_embedding = np.random.rand(384) - tool_embeddings = [np.random.rand(384) for _ in sample_tools] - - mock_query_embed.return_value = query_embedding - mock_tool_embed.side_effect = tool_embeddings - - # Mock cosine similarity to return high scores - with patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: - mock_sim.return_value = 0.8 # High similarity - - results = await embedding_strategy.search_tools(tool_repo, "cooking", limit=2) - - assert len(results) == 2 - assert all(isinstance(tool, Tool) for tool in results) - - -@pytest.mark.asyncio -async def test_search_tools_with_tag_filtering(embedding_strategy, sample_tools): - """Test search with tag filtering.""" - tool_repo = MockToolRepository(sample_tools) - - with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ - patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed, \ - patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: - - mock_query_embed.return_value = np.random.rand(384) - mock_tool_embed.return_value = np.random.rand(384) - mock_sim.return_value = 0.8 - - # Search with required tags - results = await embedding_strategy.search_tools( - tool_repo, - "cooking", - limit=10, - any_of_tags_required=["cooking", "kitchen"] - ) - - # Should only return tools with cooking or kitchen tags - assert all( - any(tag in ["cooking", "kitchen"] for tag in tool.tags) - for tool in results - ) - - -@pytest.mark.asyncio -async def test_search_tools_with_similarity_threshold(embedding_strategy, sample_tools): - """Test that similarity threshold filtering works correctly.""" - tool_repo = MockToolRepository(sample_tools) - - with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ - patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed, \ - patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: - - mock_query_embed.return_value = np.random.rand(384) - mock_tool_embed.return_value = np.random.rand(384) - - # Set threshold to 0.5 and return scores below and above - embedding_strategy.similarity_threshold = 0.5 - mock_sim.side_effect = [0.3, 0.7, 0.2] # Only second tool should pass - - results = await embedding_strategy.search_tools(tool_repo, "test", limit=10) - - assert len(results) == 1 # Only one tool above threshold - - -@pytest.mark.asyncio -async def test_search_tools_limit_respected(embedding_strategy, sample_tools): - """Test that the limit parameter is respected.""" - tool_repo = MockToolRepository(sample_tools) - - with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ - patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed, \ - patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: - - mock_query_embed.return_value = np.random.rand(384) - mock_tool_embed.return_value = np.random.rand(384) - mock_sim.return_value = 0.8 - - # Test with limit 1 - results = await embedding_strategy.search_tools(tool_repo, "test", limit=1) - assert len(results) == 1 - - # Test with limit 0 (no limit) - results = await embedding_strategy.search_tools(tool_repo, "test", limit=0) - assert len(results) == 3 # All tools - - -@pytest.mark.asyncio -async def test_search_tools_empty_repository(embedding_strategy): - """Test search behavior with empty tool repository.""" - tool_repo = MockToolRepository([]) - - results = await embedding_strategy.search_tools(tool_repo, "test", limit=10) - assert results == [] - - -@pytest.mark.asyncio -async def test_search_tools_invalid_limit(embedding_strategy, sample_tools): - """Test that invalid limit values raise appropriate errors.""" - tool_repo = MockToolRepository(sample_tools) - - with pytest.raises(ValueError, match="limit must be non-negative"): - await embedding_strategy.search_tools(tool_repo, "test", limit=-1) - - -@pytest.mark.asyncio -async def test_context_manager_behavior(embedding_strategy): - """Test async context manager behavior.""" - async with embedding_strategy as strategy: - assert strategy._model_loaded is True - - # Executor should be shut down - assert strategy._executor._shutdown is True - - -@pytest.mark.asyncio -async def test_error_handling_in_search(embedding_strategy, sample_tools): - """Test that errors in search are handled gracefully.""" - tool_repo = MockToolRepository(sample_tools) - - with patch.object(embedding_strategy, '_get_text_embedding') as mock_query_embed, \ - patch.object(embedding_strategy, '_get_tool_embedding') as mock_tool_embed: - - mock_query_embed.return_value = np.random.rand(384) - - # Make the second tool fail - def mock_tool_embed_side_effect(tool): - if tool.name == "dev.code_review": - raise Exception("Simulated error") - return np.random.rand(384) - - mock_tool_embed.side_effect = mock_tool_embed_side_effect - - # Mock cosine similarity - with patch.object(embedding_strategy, '_cosine_similarity') as mock_sim: - mock_sim.return_value = 0.8 - - # Should not crash, just skip the problematic tool - results = await embedding_strategy.search_tools(tool_repo, "test", limit=10) - - # Should return tools that didn't fail - assert len(results) == 2 # One tool failed, so only 2 results - - -@pytest.mark.asyncio -async def test_embedding_strategy_config_serializer(): - """Test the configuration serializer.""" - from utcp.implementations.embedding_search import EmbeddingSearchStrategyConfigSerializer - - serializer = EmbeddingSearchStrategyConfigSerializer() - - # Test serialization - strategy = EmbeddingSearchStrategy( - model_name="test-model", - similarity_threshold=0.5, - max_workers=8, - cache_embeddings=False - ) - - config_dict = serializer.to_dict(strategy) - assert config_dict["model_name"] == "test-model" - assert config_dict["similarity_threshold"] == 0.5 - assert config_dict["max_workers"] == 8 - assert config_dict["cache_embeddings"] is False - - # Test deserialization - restored_strategy = serializer.validate_dict(config_dict) - assert restored_strategy.model_name == "test-model" - assert restored_strategy.similarity_threshold == 0.5 - assert restored_strategy.max_workers == 8 - assert restored_strategy.cache_embeddings is False diff --git a/plugins/tool_search/embedding/README.md b/plugins/tool_search/embedding/README.md deleted file mode 100644 index 7a55d82..0000000 --- a/plugins/tool_search/embedding/README.md +++ /dev/null @@ -1,19 +0,0 @@ -# UTCP Embedding Search Plugin - -This plugin registers the embedding-based semantic search strategy with UTCP 1.0 via entry points. - -## Installation - -```bash -pip install utcp-embedding-search -``` - -Optionally, for high-quality embeddings: - -```bash -pip install utcp-embedding-search[embedding] -``` - -## How it works - -When installed, this package exposes an entry point under `utcp.plugins` so the UTCP core can auto-discover and register the `embedding_search` strategy. diff --git a/plugins/tool_search/embedding/pyproject.toml b/plugins/tool_search/embedding/pyproject.toml deleted file mode 100644 index a8cabed..0000000 --- a/plugins/tool_search/embedding/pyproject.toml +++ /dev/null @@ -1,39 +0,0 @@ -[build-system] -requires = ["setuptools>=61.0"] -build-backend = "setuptools.build_meta" - -[project] -name = "utcp-embedding-search" -version = "1.0.0" -authors = [ - { name = "UTCP Contributors" }, -] -description = "UTCP plugin providing embedding-based semantic tool search." -readme = "README.md" -requires-python = ">=3.10" -dependencies = [ - "utcp>=1.0", -] -classifiers = [ - "Development Status :: 4 - Beta", - "Intended Audience :: Developers", - "Programming Language :: Python :: 3", - "Operating System :: OS Independent", -] -license = "MPL-2.0" - -[project.optional-dependencies] -embedding = [ - "sentence-transformers>=2.2.0", - "torch>=1.9.0", -] - -authors = [ { name = "UTCP Contributors" } ] - -[project.urls] -Homepage = "https://utcp.io" -Source = "https://github.com/universal-tool-calling-protocol/python-utcp" -Issues = "https://github.com/universal-tool-calling-protocol/python-utcp/issues" - -[project.entry-points."utcp.plugins"] -embedding_search = "utcp_embedding_search:register" diff --git a/plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py b/plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py deleted file mode 100644 index a20fc1d..0000000 --- a/plugins/tool_search/embedding/src/utcp_embedding_search/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from utcp.plugins.discovery import register_tool_search_strategy -from utcp.implementations.embedding_search import EmbeddingSearchStrategyConfigSerializer - - -def register(): - """Entry point function to register the embedding search strategy.""" - register_tool_search_strategy("embedding_search", EmbeddingSearchStrategyConfigSerializer()) From 235d49041a2ce4e11f600c7eff95c58175afa9a1 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Sat, 6 Sep 2025 14:29:34 -0500 Subject: [PATCH 6/8] Correct folder placement done. --- .../tool_search/in_mem_embeddings/README.md | 18 ++++++++++++++++++ .../in_mem_embeddings/test_integration.py | 9 ++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/plugins/tool_search/in_mem_embeddings/README.md b/plugins/tool_search/in_mem_embeddings/README.md index 4e435e4..0a0eb66 100644 --- a/plugins/tool_search/in_mem_embeddings/README.md +++ b/plugins/tool_search/in_mem_embeddings/README.md @@ -14,6 +14,24 @@ Optionally, for high-quality embeddings: pip install utcp-in-mem-embeddings[embedding] ``` +Or install the required dependencies directly: + +```bash +pip install "sentence-transformers>=2.2.0" "torch>=1.9.0" +``` + +## Why are sentence-transformers and torch needed? + +While the plugin works without these packages (using a simple character frequency-based fallback), installing them provides significant benefits: + +- **Enhanced Semantic Understanding**: The `sentence-transformers` package provides pre-trained models that convert text into high-quality vector embeddings, capturing the semantic meaning of text rather than just keywords. + +- **Better Search Results**: With these packages installed, the search can understand conceptual similarity between queries and tools, even when they don't share exact keywords. + +- **Performance**: The default model (all-MiniLM-L6-v2) offers a good balance between quality and performance for semantic search applications. + +- **Fallback Mechanism**: Without these packages, the plugin automatically falls back to a simpler text similarity method, which works but with reduced accuracy. + ## How it works When installed, this package exposes an entry point under `utcp.plugins` so the UTCP core can auto-discover and register the `in_mem_embeddings` strategy. diff --git a/plugins/tool_search/in_mem_embeddings/test_integration.py b/plugins/tool_search/in_mem_embeddings/test_integration.py index aad11ec..0094403 100644 --- a/plugins/tool_search/in_mem_embeddings/test_integration.py +++ b/plugins/tool_search/in_mem_embeddings/test_integration.py @@ -67,8 +67,15 @@ async def test_integration(): ] # Create repository - repo = InMemToolRepository(tools) + repo = InMemToolRepository() + # Create a manual and add it to the repository + from utcp.data.utcp_manual import UtcpManual + manual = UtcpManual(tools=tools) + manual_call_template = CallTemplate(name="test_manual", description="Test manual", call_template_type="default") + await repo.save_manual(manual_call_template, manual) + + # Test search results = await strategy.search_tools(repo, "cooking", limit=1) print(f" āœ… Search completed, found {len(results)} results") From eef4809d48bd68a8195b3579c5f52ce14b977876 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Sat, 6 Sep 2025 14:50:23 -0500 Subject: [PATCH 7/8] Description for values accepted by model_name --- .../utcp_in_mem_embeddings/in_mem_embeddings_search.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py index 83cff80..2e6d7e9 100644 --- a/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py +++ b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py @@ -30,7 +30,15 @@ class InMemEmbeddingsSearchStrategy(ToolSearchStrategy): tool_search_strategy_type: Literal["in_mem_embeddings"] = "in_mem_embeddings" # Configuration parameters - model_name: str = Field(default="all-MiniLM-L6-v2", description="Sentence transformer model to use") + model_name: str = Field( + default="all-MiniLM-L6-v2", + description="Sentence transformer model name to use for embeddings. " + "Accepts any model from Hugging Face sentence-transformers library. " + "Popular options: 'all-MiniLM-L6-v2' (fast, good quality), " + "'all-mpnet-base-v2' (slower, higher quality), " + "'paraphrase-MiniLM-L6-v2' (paraphrase detection). " + "See https://huggingface.co/sentence-transformers for full list." + ) similarity_threshold: float = Field(default=0.3, description="Minimum similarity score to consider a match") max_workers: int = Field(default=4, description="Maximum number of worker threads for embedding generation") cache_embeddings: bool = Field(default=True, description="Whether to cache tool embeddings for performance") From 21b548a945698f8696b2ece1d44a8dede540fa81 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Mon, 8 Sep 2025 21:00:09 -0500 Subject: [PATCH 8/8] Resolved cubic suggestions --- core/pyproject.toml | 5 ----- core/src/utcp/plugins/plugin_loader.py | 13 +++++++++++-- plugins/tool_search/embedding/README.md | 4 ++-- .../tool_search/in_mem_embeddings/README.md | 2 +- .../in_mem_embeddings/pyproject.toml | 1 - .../in_mem_embeddings_search.py | 16 +++++++++------- .../in_mem_embeddings/test_integration.py | 10 ++++++---- .../in_mem_embeddings/test_performance.py | 18 +++++++++--------- .../tests/test_in_mem_embeddings_search.py | 2 +- 9 files changed, 39 insertions(+), 32 deletions(-) diff --git a/core/pyproject.toml b/core/pyproject.toml index 130c8ee..aa17d4d 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -15,7 +15,6 @@ dependencies = [ "pydantic>=2.0", "python-dotenv>=1.0", "tomli>=2.0", - "numpy>=1.21.0", ] classifiers = [ "Development Status :: 4 - Beta", @@ -34,10 +33,6 @@ dev = [ "coverage", "twine", ] -embedding = [ - "sentence-transformers>=2.2.0", - "torch>=1.9.0", -] [project.urls] Homepage = "https://utcp.io" diff --git a/core/src/utcp/plugins/plugin_loader.py b/core/src/utcp/plugins/plugin_loader.py index d13b02f..0eb5668 100644 --- a/core/src/utcp/plugins/plugin_loader.py +++ b/core/src/utcp/plugins/plugin_loader.py @@ -6,10 +6,16 @@ def _load_plugins(): from utcp.interfaces.tool_search_strategy import ToolSearchStrategyConfigSerializer from utcp.implementations.in_mem_tool_repository import InMemToolRepositoryConfigSerializer from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategyConfigSerializer - from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategyConfigSerializer from utcp.data.auth_implementations import OAuth2AuthSerializer, BasicAuthSerializer, ApiKeyAuthSerializer from utcp.data.variable_loader_implementations import DotEnvVariableLoaderSerializer from utcp.implementations.post_processors import FilterDictPostProcessorConfigSerializer, LimitStringsPostProcessorConfigSerializer + + # Try to import optional plugin, skip if not installed + try: + from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategyConfigSerializer + in_mem_embeddings_available = True + except ImportError: + in_mem_embeddings_available = False register_auth("oauth2", OAuth2AuthSerializer()) register_auth("basic", BasicAuthSerializer()) @@ -20,7 +26,10 @@ def _load_plugins(): register_tool_repository(ConcurrentToolRepositoryConfigSerializer.default_repository, InMemToolRepositoryConfigSerializer()) register_tool_search_strategy(ToolSearchStrategyConfigSerializer.default_strategy, TagAndDescriptionWordMatchStrategyConfigSerializer()) - register_tool_search_strategy("in_mem_embeddings", InMemEmbeddingsSearchStrategyConfigSerializer()) + + # Register optional plugin only if available + if in_mem_embeddings_available: + register_tool_search_strategy("in_mem_embeddings", InMemEmbeddingsSearchStrategyConfigSerializer()) register_tool_post_processor("filter_dict", FilterDictPostProcessorConfigSerializer()) register_tool_post_processor("limit_strings", LimitStringsPostProcessorConfigSerializer()) diff --git a/plugins/tool_search/embedding/README.md b/plugins/tool_search/embedding/README.md index ca0b906..71cfa9e 100644 --- a/plugins/tool_search/embedding/README.md +++ b/plugins/tool_search/embedding/README.md @@ -11,9 +11,9 @@ pip install utcp-embedding-search Optionally, for high-quality embeddings: ```bash -pip install "utcp-embedding-search[embedding]" +pip install "utcp-in-mem-embeddings[embedding]" ``` ## How it works -When installed, this package exposes an entry point under `utcp.plugins` so the UTCP core can auto-discover and register the `embedding_search` strategy. +When installed, this package exposes an entry point under `utcp.plugins` so the UTCP core can auto-discover and register the `in_mem_embeddings` strategy. diff --git a/plugins/tool_search/in_mem_embeddings/README.md b/plugins/tool_search/in_mem_embeddings/README.md index 0a0eb66..5a844a6 100644 --- a/plugins/tool_search/in_mem_embeddings/README.md +++ b/plugins/tool_search/in_mem_embeddings/README.md @@ -11,7 +11,7 @@ pip install utcp-in-mem-embeddings Optionally, for high-quality embeddings: ```bash -pip install utcp-in-mem-embeddings[embedding] +pip install "utcp-in-mem-embeddings[embedding]" ``` Or install the required dependencies directly: diff --git a/plugins/tool_search/in_mem_embeddings/pyproject.toml b/plugins/tool_search/in_mem_embeddings/pyproject.toml index 2e7e389..e77ba4e 100644 --- a/plugins/tool_search/in_mem_embeddings/pyproject.toml +++ b/plugins/tool_search/in_mem_embeddings/pyproject.toml @@ -28,7 +28,6 @@ embedding = [ "torch>=1.9.0", ] -authors = [ { name = "UTCP Contributors" } ] [project.urls] Homepage = "https://utcp.io" diff --git a/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py index 2e6d7e9..669748d 100644 --- a/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py +++ b/plugins/tool_search/in_mem_embeddings/src/utcp_in_mem_embeddings/in_mem_embeddings_search.py @@ -10,7 +10,7 @@ from typing import List, Tuple, Optional, Literal, Dict, Any from concurrent.futures import ThreadPoolExecutor import numpy as np -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, PrivateAttr from utcp.interfaces.tool_search_strategy import ToolSearchStrategy from utcp.data.tool import Tool @@ -43,12 +43,15 @@ class InMemEmbeddingsSearchStrategy(ToolSearchStrategy): max_workers: int = Field(default=4, description="Maximum number of worker threads for embedding generation") cache_embeddings: bool = Field(default=True, description="Whether to cache tool embeddings for performance") + # Private attributes + _embedding_model: Optional[Any] = PrivateAttr(default=None) + _tool_embeddings_cache: Dict[str, np.ndarray] = PrivateAttr(default_factory=dict) + _executor: Optional[ThreadPoolExecutor] = PrivateAttr(default=None) + _model_loaded: bool = PrivateAttr(default=False) + def __init__(self, **data): super().__init__(**data) - self._embedding_model = None - self._tool_embeddings_cache: Dict[str, np.ndarray] = {} self._executor = ThreadPoolExecutor(max_workers=self.max_workers) - self._model_loaded = False async def _ensure_model_loaded(self): """Ensure the embedding model is loaded.""" @@ -60,7 +63,7 @@ async def _ensure_model_loaded(self): from sentence_transformers import SentenceTransformer # Load the model in a thread to avoid blocking - loop = asyncio.get_event_loop() + loop = asyncio.get_running_loop() self._embedding_model = await loop.run_in_executor( self._executor, SentenceTransformer, @@ -108,8 +111,7 @@ def _simple_text_embedding(self, text: str) -> np.ndarray: # Simple character frequency-based embedding for i, char in enumerate(text_lower): - if i < 384: - embedding[i % 384] += ord(char) / 1000.0 + embedding[i % 384] += ord(char) / 1000.0 # Normalize norm = np.linalg.norm(embedding) diff --git a/plugins/tool_search/in_mem_embeddings/test_integration.py b/plugins/tool_search/in_mem_embeddings/test_integration.py index 0094403..908be2b 100644 --- a/plugins/tool_search/in_mem_embeddings/test_integration.py +++ b/plugins/tool_search/in_mem_embeddings/test_integration.py @@ -6,8 +6,8 @@ from pathlib import Path # Add paths -plugin_src = Path(__file__).parent / "src" -core_src = Path(__file__).parent.parent.parent.parent / "core" / "src" +plugin_src = (Path(__file__).parent / "src").resolve() +core_src = (Path(__file__).parent.parent.parent.parent / "core" / "src").resolve() sys.path.insert(0, str(plugin_src)) sys.path.insert(0, str(core_src)) @@ -60,7 +60,6 @@ async def test_integration(): tags=["cooking", "test"], tool_call_template=CallTemplate( name="test.tool1", - description="Test tool", call_template_type="default" ) ) @@ -72,7 +71,7 @@ async def test_integration(): # Create a manual and add it to the repository from utcp.data.utcp_manual import UtcpManual manual = UtcpManual(tools=tools) - manual_call_template = CallTemplate(name="test_manual", description="Test manual", call_template_type="default") + manual_call_template = CallTemplate(name="test_manual", call_template_type="default") await repo.save_manual(manual_call_template, manual) @@ -80,6 +79,9 @@ async def test_integration(): results = await strategy.search_tools(repo, "cooking", limit=1) print(f" āœ… Search completed, found {len(results)} results") + # Validate search results + assert len(results) > 0, "Search should return at least one result for 'cooking' query" + print("\nšŸŽ‰ Integration test passed! Plugin works with core system.") return True diff --git a/plugins/tool_search/in_mem_embeddings/test_performance.py b/plugins/tool_search/in_mem_embeddings/test_performance.py index d5189f6..b447a35 100644 --- a/plugins/tool_search/in_mem_embeddings/test_performance.py +++ b/plugins/tool_search/in_mem_embeddings/test_performance.py @@ -34,13 +34,13 @@ async def test_performance(): tools = [] for i in range(100): tool = Tool( - name=f"test.tool{i}", + name=f"test_tool{i}", description=f"Test tool {i} for various purposes like cooking, coding, data analysis", inputs=JsonSchema(), outputs=JsonSchema(), tags=["test", f"category{i%5}"], tool_call_template=CallTemplate( - name=f"test.tool{i}", + name=f"test_tool{i}", description=f"Test tool {i}", call_template_type="default" ) @@ -58,32 +58,32 @@ async def get_tools(self): # Test 1: First search (cold start) print("2. Testing cold start performance...") - start_time = time.time() + start_time = time.perf_counter() results1 = await strategy.search_tools(repo, "cooking tools", limit=10) - cold_time = time.time() - start_time + cold_time = time.perf_counter() - start_time print(f" ā±ļø Cold start: {cold_time:.3f}s, found {len(results1)} results") # Test 2: Second search (warm cache) print("3. Testing warm cache performance...") - start_time = time.time() + start_time = time.perf_counter() results2 = await strategy.search_tools(repo, "coding tools", limit=10) - warm_time = time.time() - start_time + warm_time = time.perf_counter() - start_time print(f" ā±ļø Warm cache: {warm_time:.3f}s, found {len(results2)} results") # Test 3: Multiple searches print("4. Testing multiple searches...") queries = ["cooking", "programming", "data analysis", "testing", "utilities"] - start_time = time.time() + start_time = time.perf_counter() for query in queries: await strategy.search_tools(repo, query, limit=5) - total_time = time.time() - start_time + total_time = time.perf_counter() - start_time avg_time = total_time / len(queries) print(f" ā±ļø Average per search: {avg_time:.3f}s") # Performance assertions - assert cold_time < 5.0, f"Cold start too slow: {cold_time}s" + assert cold_time < 10.0, f"Cold start too slow: {cold_time}s" # Allow more time for model loading assert warm_time < 1.0, f"Warm cache too slow: {warm_time}s" assert avg_time < 0.5, f"Average search too slow: {avg_time}s" diff --git a/plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py b/plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py index 5f37c61..d27773e 100644 --- a/plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py +++ b/plugins/tool_search/in_mem_embeddings/tests/test_in_mem_embeddings_search.py @@ -12,7 +12,7 @@ sys.path.insert(0, str(plugin_src)) # Add core to path -core_src = Path(__file__).parent.parent.parent.parent / "core" / "src" +core_src = Path(__file__).parent.parent.parent.parent.parent / "core" / "src" sys.path.insert(0, str(core_src)) from utcp_in_mem_embeddings.in_mem_embeddings_search import InMemEmbeddingsSearchStrategy