swarms 7.8.4__py3-none-any.whl → 7.8.8__py3-none-any.whl

swarms/structs/swarm_matcher.py

@@ -1,8 +1,11 @@
 import json
-from typing import List, Optional, Tuple
+import os
+from typing import Dict, List, Literal, Optional, Tuple, Union
 
-from pydantic import BaseModel, Field
-from tenacity import retry, stop_after_attempt, wait_exponential
+import numpy as np
+from pydantic import BaseModel, Field, field_validator
+from pydantic.v1 import validator
+from litellm import embedding
 
 from swarms.utils.auto_download_check_packages import (
     auto_check_and_download_package,
@@ -20,18 +23,75 @@ class SwarmType(BaseModel):
     )
 
 
+api_key = os.getenv("OPENAI_API_KEY")
+
+
 class SwarmMatcherConfig(BaseModel):
-    model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
-    embedding_dim: int = (
-        512  # Dimension of the sentence-transformers model
+    backend: Literal["local", "openai"] = "local"
+    model_name: str = (
+        "sentence-transformers/all-MiniLM-L6-v2"  # For local embeddings
+    )
+    openai_model: str = (
+        "text-embedding-3-small"  # Default to newer OpenAI model
+    )
+    embedding_dim: int = 512  # For local embeddings
+    openai_dimensions: Optional[int] = (
+        None  # For OpenAI text-embedding-3-* models
     )
+    similarity_threshold: float = Field(default=0.5, ge=0.0, le=1.0)
+    cache_embeddings: bool = True
+    max_sequence_length: int = Field(default=512, ge=64, le=2048)
+    device: str = "cpu"  # Only used for local embeddings
+    batch_size: int = Field(default=32, ge=1)
+    openai_api_key: Optional[str] = os.getenv("OPENAI_API_KEY")
+    metadata: Optional[Dict] = Field(
+        default_factory=dict
+    )  # For OpenAI embedding calls
+
+    class Config:
+        validate_assignment = True
+
+    @validator("openai_dimensions")
+    def validate_dimensions(cls, v, values):
+        if values.get("backend") == "openai":
+            if (
+                values.get("openai_model", "").startswith(
+                    "text-embedding-3"
+                )
+                and v is None
+            ):
+                # Default to 1536 for text-embedding-3-small/large if not specified
+                return 1536
+        return v
+
+    @field_validator("openai_model")
+    def validate_model(cls, v, values):
+        if values.get("backend") == "openai":
+            valid_models = [
+                "text-embedding-3-small",
+                "text-embedding-3-large",
+                "text-embedding-ada-002",
+            ]
+            if v not in valid_models:
+                raise ValueError(
+                    f"OpenAI model must be one of: {', '.join(valid_models)}"
+                )
+        return v
 
 
 class SwarmMatcher:
     """
-    A class for matching tasks to swarm types based on their descriptions.
-    It utilizes a transformer model to generate embeddings for task and swarm type descriptions,
-    and then calculates the dot product to find the best match.
+    A class for matching tasks to swarm types based on their descriptions using semantic similarity.
+
+    This class uses transformer models to generate embeddings for both task descriptions and swarm type descriptions.
+    It then calculates similarity scores to find the most appropriate swarm type for a given task.
+
+    Features:
+    - Supports both local transformer models and OpenAI embeddings
+    - Implements embedding caching for improved performance
+    - Provides batch processing capabilities
+    - Includes retry mechanisms for API calls
+    - Supports saving/loading swarm type configurations
     """
 
     def __init__(self, config: SwarmMatcherConfig):
@@ -39,15 +99,35 @@ class SwarmMatcher:
         Initializes the SwarmMatcher with a configuration.
 
         Args:
-            config (SwarmMatcherConfig): The configuration for the SwarmMatcher.
+            config (SwarmMatcherConfig): Configuration object specifying model settings,
+                                        similarity thresholds, and other parameters.
+
+        Raises:
+            ImportError: If required dependencies (torch, transformers) are not available
+            Exception: If model initialization fails
         """
+        try:
+            self.config = config
+            if self.config.backend == "local":
+                transformers = self._setup_dependencies()
+                self._setup_model_and_tokenizer(transformers)
+            self._initialize_state()
+            self.initialize_swarm_types()
+            logger.debug("SwarmMatcher initialized successfully")
+        except Exception as e:
+            logger.error(f"Error initializing SwarmMatcher: {str(e)}")
+            raise
 
+    def _setup_dependencies(self):
+        """Set up required dependencies for the SwarmMatcher."""
         try:
+            import numpy as np
             import torch
         except ImportError:
             auto_check_and_download_package(
                 "torch", package_manager="pip", upgrade=True
             )
+            import numpy as np
             import torch
 
         try:
@@ -59,65 +139,279 @@ class SwarmMatcher:
             import transformers
 
         self.torch = torch
+        self.np = np
+        return transformers
+
+    def _setup_model_and_tokenizer(self, transformers):
+        """Initialize the model and tokenizer."""
+        self.device = self.torch.device(self.config.device)
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(
+            self.config.model_name
+        )
+        self.model = transformers.AutoModel.from_pretrained(
+            self.config.model_name
+        ).to(self.device)
+
+    def _initialize_state(self):
+        """Initialize internal state variables."""
+        self.swarm_types: List[SwarmType] = []
+        self._embedding_cache = (
+            {} if self.config.cache_embeddings else None
+        )
+
+    def _get_cached_embedding(
+        self, text: str
+    ) -> Optional[np.ndarray]:
+        """
+        Retrieves a cached embedding if available.
+
+        Args:
+            text (str): The text to look up in the cache
+
+        Returns:
+            Optional[np.ndarray]: The cached embedding if found, None otherwise
+        """
+        if self._embedding_cache is not None:
+            return self._embedding_cache.get(text)
+        return None
+
+    def _cache_embedding(self, text: str, embedding: np.ndarray):
+        """
+        Stores an embedding in the cache for future use.
+
+        Args:
+            text (str): The text associated with the embedding
+            embedding (np.ndarray): The embedding vector to cache
+        """
+        if self._embedding_cache is not None:
+            self._embedding_cache[text] = embedding
+
+    def _get_openai_embedding(self, text: str) -> np.ndarray:
+        """Get embedding using OpenAI's API via litellm."""
         try:
-            self.config = config
-            self.tokenizer = (
-                transformers.AutoTokenizer.from_pretrained(
-                    config.model_name
+            params = {
+                "model": self.config.openai_model,
+                "input": [text],
+            }
+
+            # Add dimensions parameter for text-embedding-3-* models
+            if (
+                self.config.openai_model.startswith(
+                    "text-embedding-3"
                 )
-            )
-            self.model = transformers.AutoModel.from_pretrained(
-                config.model_name
-            )
-            self.swarm_types: List[SwarmType] = []
-            logger.debug("SwarmMatcher initialized successfully")
+                and self.config.openai_dimensions
+            ):
+                params["dimensions"] = self.config.openai_dimensions
+
+            response = embedding(**params)
+            response = response.model_dump()
+
+            # Handle the response format
+            if "data" in response and len(response["data"]) > 0:
+                embedding_data = response["data"][0]["embedding"]
+            else:
+                raise ValueError(
+                    f"Unexpected response format from OpenAI API: {response}"
+                )
+
+            embedding_array = np.array(embedding_data)
+
+            # Log usage information if available
+            if "usage" in response:
+                logger.debug(
+                    f"OpenAI API usage - Prompt tokens: {response['usage'].get('prompt_tokens', 'N/A')}, "
+                    f"Total tokens: {response['usage'].get('total_tokens', 'N/A')}"
+                )
+
+            return embedding_array
         except Exception as e:
-            logger.error(f"Error initializing SwarmMatcher: {str(e)}")
+            logger.error(f"Error getting OpenAI embedding: {str(e)}")
             raise
 
-    @retry(
-        stop=stop_after_attempt(3),
-        wait=wait_exponential(multiplier=1, min=4, max=10),
-    )
-    def get_embedding(self, text: str):
+    def _get_openai_embeddings_batch(
+        self, texts: List[str]
+    ) -> np.ndarray:
+        """Get embeddings for a batch of texts using OpenAI's API via litellm."""
+        try:
+            params = {
+                "model": self.config.openai_model,
+                "input": texts,
+            }
+
+            # Add dimensions parameter for text-embedding-3-* models
+            if (
+                self.config.openai_model.startswith(
+                    "text-embedding-3"
+                )
+                and self.config.openai_dimensions
+            ):
+                params["dimensions"] = self.config.openai_dimensions
+
+            response = embedding(**params)
+            response = response.model_dump()
+
+            # Handle the response format
+            if "data" in response:
+                embeddings = [
+                    data["embedding"] for data in response["data"]
+                ]
+            else:
+                raise ValueError(
+                    f"Unexpected response format from OpenAI API: {response}"
+                )
+
+            # Log usage information if available
+            if "usage" in response:
+                logger.debug(
+                    f"Batch OpenAI API usage - Prompt tokens: {response['usage'].get('prompt_tokens', 'N/A')}, "
+                    f"Total tokens: {response['usage'].get('total_tokens', 'N/A')}"
+                )
+
+            return np.array(embeddings)
+        except Exception as e:
+            logger.error(
+                f"Error getting OpenAI embeddings batch: {str(e)}"
+            )
+            raise
+
+    def get_embedding(self, text: str) -> np.ndarray:
         """
         Generates an embedding for a given text using the configured model.
 
+        This method first checks the cache for an existing embedding. If not found,
+        it generates a new embedding using either the local transformer model or OpenAI API.
+
         Args:
-            text (str): The text for which to generate an embedding.
+            text (str): The text for which to generate an embedding
 
         Returns:
-            np.ndarray: The embedding vector for the text.
+            np.ndarray: The embedding vector for the text
+
+        Raises:
+            Exception: If embedding generation fails
         """
+        # Check cache first
+        cached_embedding = self._get_cached_embedding(text)
+        if cached_embedding is not None:
+            return cached_embedding
 
         logger.debug(f"Getting embedding for text: {text[:50]}...")
         try:
-            inputs = self.tokenizer(
-                text,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=512,
-            )
-            with self.torch.no_grad():
-                outputs = self.model(**inputs)
-            embedding = (
-                outputs.last_hidden_state.mean(dim=1)
-                .squeeze()
-                .numpy()
-            )
+            if self.config.backend == "openai":
+                embedding = self._get_openai_embedding(text)
+            else:
+                inputs = self.tokenizer(
+                    text,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=self.config.max_sequence_length,
+                )
+                # Move inputs to device
+                inputs = {
+                    k: v.to(self.device) for k, v in inputs.items()
+                }
+
+                with self.torch.no_grad():
+                    outputs = self.model(**inputs)
+                embedding = (
+                    outputs.last_hidden_state.mean(dim=1)
+                    .squeeze()
+                    .cpu()
+                    .numpy()
+                )
+
+            # Cache the embedding
+            self._cache_embedding(text, embedding)
+
             logger.debug("Embedding generated successfully")
             return embedding
         except Exception as e:
             logger.error(f"Error generating embedding: {str(e)}")
             raise
 
+    def get_embeddings_batch(self, texts: List[str]) -> np.ndarray:
+        """
+        Generate embeddings for multiple texts in batch for improved efficiency.
+
+        This method processes texts in batches, utilizing the cache where possible
+        and generating new embeddings only for uncached texts.
+
+        Args:
+            texts (List[str]): List of texts to generate embeddings for
+
+        Returns:
+            np.ndarray: Array of embeddings, one for each input text
+
+        Raises:
+            Exception: If batch processing fails
+        """
+        embeddings = []
+        batch_texts = []
+
+        for text in texts:
+            cached_embedding = self._get_cached_embedding(text)
+            if cached_embedding is not None:
+                embeddings.append(cached_embedding)
+            else:
+                batch_texts.append(text)
+
+        if batch_texts:
+            if self.config.backend == "openai":
+                batch_embeddings = self._get_openai_embeddings_batch(
+                    batch_texts
+                )
+                for text, embedding in zip(
+                    batch_texts, batch_embeddings
+                ):
+                    self._cache_embedding(text, embedding)
+                    embeddings.append(embedding)
+            else:
+                for i in range(
+                    0, len(batch_texts), self.config.batch_size
+                ):
+                    batch = batch_texts[
+                        i : i + self.config.batch_size
+                    ]
+                    inputs = self.tokenizer(
+                        batch,
+                        return_tensors="pt",
+                        padding=True,
+                        truncation=True,
+                        max_length=self.config.max_sequence_length,
+                    )
+                    inputs = {
+                        k: v.to(self.device)
+                        for k, v in inputs.items()
+                    }
+
+                    with self.torch.no_grad():
+                        outputs = self.model(**inputs)
+                    batch_embeddings = (
+                        outputs.last_hidden_state.mean(dim=1)
+                        .cpu()
+                        .numpy()
+                    )
+
+                    for text, embedding in zip(
+                        batch, batch_embeddings
+                    ):
+                        self._cache_embedding(text, embedding)
+                        embeddings.append(embedding)
+
+        return np.array(embeddings)
+
     def add_swarm_type(self, swarm_type: SwarmType):
         """
-        Adds a swarm type to the list of swarm types, generating an embedding for its description.
+        Adds a swarm type to the matcher's registry.
+
+        Generates and stores an embedding for the swarm type's description.
 
         Args:
-            swarm_type (SwarmType): The swarm type to add.
+            swarm_type (SwarmType): The swarm type to add
+
+        Raises:
+            Exception: If embedding generation or storage fails
         """
         logger.debug(f"Adding swarm type: {swarm_type.name}")
         try:
@@ -133,38 +427,104 @@ class SwarmMatcher:
 
     def find_best_match(self, task: str) -> Tuple[str, float]:
         """
-        Finds the best match for a given task among the registered swarm types.
+        Finds the best matching swarm type for a given task.
+
+        Uses semantic similarity to compare the task against all registered swarm types
+        and returns the best match along with its confidence score.
 
         Args:
-            task (str): The task for which to find the best match.
+            task (str): The task description to match
 
         Returns:
-            Tuple[str, float]: A tuple containing the name of the best matching swarm type and the score.
-        """
-        import numpy as np
+            Tuple[str, float]: A tuple containing:
+                - The name of the best matching swarm type
+                - The similarity score (between 0 and 1)
 
+        Raises:
+            Exception: If matching process fails
+        """
         logger.debug(f"Finding best match for task: {task[:50]}...")
         try:
             task_embedding = self.get_embedding(task)
             best_match = None
             best_score = -float("inf")
-            for swarm_type in self.swarm_types:
-                score = np.dot(
-                    task_embedding, np.array(swarm_type.embedding)
+
+            # Get all swarm type embeddings in batch
+            swarm_descriptions = [
+                st.description for st in self.swarm_types
+            ]
+            swarm_embeddings = self.get_embeddings_batch(
+                swarm_descriptions
+            )
+
+            # Calculate similarity scores in batch
+            scores = np.dot(task_embedding, swarm_embeddings.T)
+            best_idx = np.argmax(scores)
+            best_score = float(scores[best_idx])
+            best_match = self.swarm_types[best_idx]
+
+            if best_score < self.config.similarity_threshold:
+                logger.warning(
+                    f"Best match score {best_score} is below threshold {self.config.similarity_threshold}"
                 )
-                if score > best_score:
-                    best_score = score
-                    best_match = swarm_type
+
             logger.info(
                 f"Best match for task: {best_match.name} (score: {best_score})"
             )
-            return best_match.name, float(best_score)
+            return best_match.name, best_score
         except Exception as e:
             logger.error(
                 f"Error finding best match for task: {str(e)}"
             )
             raise
 
+    def find_top_k_matches(
+        self, task: str, k: int = 3
+    ) -> List[Tuple[str, float]]:
+        """
+        Finds the top k matching swarm types for a given task.
+
+        Returns all matches that exceed the similarity threshold, sorted by score.
+
+        Args:
+            task (str): The task for which to find matches.
+            k (int): Number of top matches to return.
+
+        Returns:
+            List[Tuple[str, float]]: List of tuples containing swarm names and their scores.
+        """
+        logger.debug(
+            f"Finding top {k} matches for task: {task[:50]}..."
+        )
+        try:
+            task_embedding = self.get_embedding(task)
+            swarm_descriptions = [
+                st.description for st in self.swarm_types
+            ]
+            swarm_embeddings = self.get_embeddings_batch(
+                swarm_descriptions
+            )
+
+            # Calculate similarity scores in batch
+            scores = np.dot(task_embedding, swarm_embeddings.T)
+            top_k_indices = np.argsort(scores)[-k:][::-1]
+
+            results = []
+            for idx in top_k_indices:
+                score = float(scores[idx])
+                if score >= self.config.similarity_threshold:
+                    results.append(
+                        (self.swarm_types[idx].name, score)
+                    )
+
+            logger.info(
+                f"Found {len(results)} matches above threshold"
+            )
+            return results
+        except Exception as e:
+            logger.error(f"Error finding top matches: {str(e)}")
+            raise
+
     def auto_select_swarm(self, task: str) -> str:
         """
         Automatically selects the best swarm type for a given task based on their descriptions.
@@ -226,56 +586,86 @@ class SwarmMatcher:
             logger.error(f"Error loading swarm types: {str(e)}")
             raise
 
+    def initialize_swarm_types(self):
+        logger.debug("Initializing swarm types")
+        swarm_types = [
+            SwarmType(
+                name="AgentRearrange",
+                description="Optimize agent order and rearrange flow for multi-step tasks, ensuring efficient task allocation and minimizing bottlenecks. Keywords: orchestration, coordination, pipeline optimization, task scheduling, resource allocation, workflow management, agent organization, process optimization",
+            ),
+            SwarmType(
+                name="MixtureOfAgents",
+                description="Combine diverse expert agents for comprehensive analysis, fostering a collaborative approach to problem-solving and leveraging individual strengths. Keywords: multi-agent system, expert collaboration, distributed intelligence, collective problem solving, agent specialization, team coordination, hybrid approaches, knowledge synthesis",
+            ),
+            SwarmType(
+                name="SpreadSheetSwarm",
+                description="Collaborative data processing and analysis in a spreadsheet-like environment, facilitating real-time data sharing and visualization. Keywords: data analysis, tabular processing, collaborative editing, data transformation, spreadsheet operations, data visualization, real-time collaboration, structured data",
+            ),
+            SwarmType(
+                name="SequentialWorkflow",
+                description="Execute tasks in a step-by-step, sequential process workflow, ensuring a logical and methodical approach to task execution. Keywords: linear processing, waterfall methodology, step-by-step execution, ordered tasks, sequential operations, process flow, systematic approach, staged execution",
+            ),
+            SwarmType(
+                name="ConcurrentWorkflow",
+                description="Process multiple tasks or data sources concurrently in parallel, maximizing productivity and reducing processing time. Keywords: parallel processing, multi-threading, asynchronous execution, distributed computing, concurrent operations, simultaneous tasks, parallel workflows, scalable processing",
+            ),
+            SwarmType(
+                name="HierarchicalSwarm",
+                description="Organize agents in a hierarchical structure with clear reporting lines and delegation of responsibilities. Keywords: management hierarchy, organizational structure, delegation, supervision, chain of command, tiered organization, structured coordination, leadership roles",
+            ),
+            SwarmType(
+                name="AdaptiveSwarm",
+                description="Dynamically adjust agent behavior and swarm configuration based on task requirements and performance feedback. Keywords: dynamic adaptation, self-optimization, feedback loops, learning systems, flexible configuration, responsive behavior, adaptive algorithms, real-time adjustment",
+            ),
+            SwarmType(
+                name="ConsensusSwarm",
+                description="Achieve group decisions through consensus mechanisms and voting protocols among multiple agents. Keywords: group decision making, voting systems, collective intelligence, agreement protocols, democratic processes, collaborative decisions, consensus building",
+            ),
+            SwarmType(
+                name="DeepResearchSwarm",
+                description="Conduct in-depth research and analysis by coordinating multiple agents to explore, synthesize, and validate information from various sources. Keywords: research methodology, information synthesis, data validation, comprehensive analysis, knowledge discovery, systematic investigation",
+            ),
+            SwarmType(
+                name="CouncilAsAJudge",
+                description="Evaluate and judge solutions or decisions through a council of expert agents acting as arbitrators. Keywords: evaluation, judgment, arbitration, expert assessment, quality control, decision validation, peer review, consensus building",
+            ),
+            SwarmType(
+                name="MALT",
+                description="Multi-Agent Language Tasks framework for coordinating language-based operations across multiple specialized agents. Keywords: language processing, task coordination, linguistic analysis, communication protocols, semantic understanding, natural language tasks",
+            ),
+            SwarmType(
+                name="GroupChat",
+                description="Enable dynamic multi-agent conversations and collaborative problem-solving through structured group discussions. Keywords: collaborative dialogue, group interaction, team communication, collective problem-solving, discussion facilitation, knowledge sharing",
+            ),
+            SwarmType(
+                name="MultiAgentRouter",
+                description="Intelligently route tasks and information between agents based on their specializations and current workload. Keywords: task distribution, load balancing, intelligent routing, agent specialization, workflow optimization, resource allocation",
+            ),
+            SwarmType(
+                name="MajorityVoting",
+                description="Make decisions through democratic voting mechanisms where multiple agents contribute their opinions and votes. Keywords: collective decision-making, democratic process, vote aggregation, opinion pooling, consensus building, collaborative choice",
+            ),
+        ]
+
+        try:
+            for swarm_type in swarm_types:
+                self.add_swarm_type(swarm_type)
+        except Exception as e:
+            logger.error(f"Error initializing swarm types: {str(e)}")
+            raise
+
 
-def initialize_swarm_types(matcher: SwarmMatcher):
-    logger.debug("Initializing swarm types")
-    swarm_types = [
-        SwarmType(
-            name="AgentRearrange",
-            description="Optimize agent order and rearrange flow for multi-step tasks, ensuring efficient task allocation and minimizing bottlenecks. Keywords: orchestration, coordination, pipeline optimization, task scheduling, resource allocation, workflow management, agent organization, process optimization",
-        ),
-        SwarmType(
-            name="MixtureOfAgents",
-            description="Combine diverse expert agents for comprehensive analysis, fostering a collaborative approach to problem-solving and leveraging individual strengths. Keywords: multi-agent system, expert collaboration, distributed intelligence, collective problem solving, agent specialization, team coordination, hybrid approaches, knowledge synthesis",
-        ),
-        SwarmType(
-            name="SpreadSheetSwarm",
-            description="Collaborative data processing and analysis in a spreadsheet-like environment, facilitating real-time data sharing and visualization. Keywords: data analysis, tabular processing, collaborative editing, data transformation, spreadsheet operations, data visualization, real-time collaboration, structured data",
-        ),
-        SwarmType(
-            name="SequentialWorkflow",
-            description="Execute tasks in a step-by-step, sequential process workflow, ensuring a logical and methodical approach to task execution. Keywords: linear processing, waterfall methodology, step-by-step execution, ordered tasks, sequential operations, process flow, systematic approach, staged execution",
-        ),
-        SwarmType(
-            name="ConcurrentWorkflow",
-            description="Process multiple tasks or data sources concurrently in parallel, maximizing productivity and reducing processing time. Keywords: parallel processing, multi-threading, asynchronous execution, distributed computing, concurrent operations, simultaneous tasks, parallel workflows, scalable processing",
-        ),
-        # SwarmType(
-        #     name="HierarchicalSwarm",
-        #     description="Organize agents in a hierarchical structure with clear reporting lines and delegation of responsibilities. Keywords: management hierarchy, organizational structure, delegation, supervision, chain of command, tiered organization, structured coordination",
-        # ),
-        # SwarmType(
-        #     name="AdaptiveSwarm",
-        #     description="Dynamically adjust agent behavior and swarm configuration based on task requirements and performance feedback. Keywords: dynamic adaptation, self-optimization, feedback loops, learning systems, flexible configuration, responsive behavior, adaptive algorithms",
-        # ),
-        # SwarmType(
-        #     name="ConsensusSwarm",
-        #     description="Achieve group decisions through consensus mechanisms and voting protocols among multiple agents. Keywords: group decision making, voting systems, collective intelligence, agreement protocols, democratic processes, collaborative decisions",
-        # ),
-    ]
-
-    for swarm_type in swarm_types:
-        matcher.add_swarm_type(swarm_type)
-    logger.debug("Swarm types initialized")
-
-
-def swarm_matcher(task: str, *args, **kwargs):
+def swarm_matcher(task: Union[str, List[str]], *args, **kwargs):
     """
     Runs the SwarmMatcher example with predefined tasks and swarm types.
     """
+    if isinstance(task, list):
+        task = "".join(task)
+    else:
+        task = task
+
     config = SwarmMatcherConfig()
     matcher = SwarmMatcher(config)
-    initialize_swarm_types(matcher)
 
     # matcher.save_swarm_types(f"swarm_logs/{uuid4().hex}.json")
 
@@ -286,319 +676,18 @@ def swarm_matcher(task: str, *args, **kwargs):
     return swarm_type
 
 
-# from typing import List, Tuple, Dict
-# from pydantic import BaseModel, Field
-# from loguru import logger
-# from uuid import uuid4
-# import chromadb
-# import json
-# from tenacity import retry, stop_after_attempt, wait_exponential
-
-
-# class SwarmType(BaseModel):
-#     """A swarm type with its name, description and optional metadata"""
-
-#     id: str = Field(default_factory=lambda: str(uuid4()))
-#     name: str
-#     description: str
-#     metadata: Dict = Field(default_factory=dict)
-
-
-# class SwarmMatcherConfig(BaseModel):
-#     """Configuration for the SwarmMatcher"""
-
-#     collection_name: str = "swarm_types"
-#     distance_metric: str = "cosine"  # or "l2" or "ip"
-#     embedding_function: str = (
-#         "sentence-transformers/all-mpnet-base-v2"  # Better model than MiniLM
-#     )
-#     persist_directory: str = "./chroma_db"
-
-
-# class SwarmMatcher:
-#     """
-#     An improved swarm matcher that uses ChromaDB for better vector similarity search.
-#     Features:
-#     - Persistent storage of embeddings
-#     - Better vector similarity search with multiple distance metrics
-#     - Improved embedding model
-#     - Metadata filtering capabilities
-#     - Batch operations support
-#     """
-
-#     def __init__(self, config: SwarmMatcherConfig):
-#         """Initialize the improved swarm matcher"""
-#         logger.add("swarm_matcher.log", rotation="100 MB")
-#         self.config = config
-
-#         # Initialize ChromaDB client with persistence
-#         self.chroma_client = chromadb.Client()
-
-#         # Get or create collection
-#         try:
-#             self.collection = self.chroma_client.get_collection(
-#                 name=config.collection_name,
-#             )
-#         except ValueError:
-#             self.collection = self.chroma_client.create_collection(
-#                 name=config.collection_name,
-#                 metadata={"hnsw:space": config.distance_metric},
-#             )
-
-#         logger.info(
-#             f"Initialized SwarmMatcher with collection '{config.collection_name}'"
-#         )
-
-#     def add_swarm_type(self, swarm_type: SwarmType) -> None:
-#         """Add a single swarm type to the collection"""
-#         try:
-#             self.collection.add(
-#                 ids=[swarm_type.id],
-#                 documents=[swarm_type.description],
-#                 metadatas=[
-#                     {"name": swarm_type.name, **swarm_type.metadata}
-#                 ],
-#             )
-#             logger.info(f"Added swarm type: {swarm_type.name}")
-#         except Exception as e:
-#             logger.error(
-#                 f"Error adding swarm type {swarm_type.name}: {str(e)}"
-#             )
-#             raise
-
-#     def add_swarm_types(self, swarm_types: List[SwarmType]) -> None:
-#         """Add multiple swarm types in batch"""
-#         try:
-#             self.collection.add(
-#                 ids=[st.id for st in swarm_types],
-#                 documents=[st.description for st in swarm_types],
-#                 metadatas=[
-#                     {"name": st.name, **st.metadata}
-#                     for st in swarm_types
-#                 ],
-#             )
-#             logger.info(f"Added {len(swarm_types)} swarm types")
-#         except Exception as e:
-#             logger.error(
-#                 f"Error adding swarm types in batch: {str(e)}"
-#             )
-#             raise
-
-#     @retry(
-#         stop=stop_after_attempt(3),
-#         wait=wait_exponential(multiplier=1, min=4, max=10),
-#     )
-#     def find_best_matches(
-#         self,
-#         task: str,
-#         n_results: int = 3,
-#         score_threshold: float = 0.7,
-#     ) -> List[Tuple[str, float]]:
-#         """
-#         Find the best matching swarm types for a given task
-#         Returns multiple matches with their scores
-#         """
-#         try:
-#             results = self.collection.query(
-#                 query_texts=[task],
-#                 n_results=n_results,
-#                 include=["metadatas", "distances"],
-#             )
-
-#             matches = []
-#             for metadata, distance in zip(
-#                 results["metadatas"][0], results["distances"][0]
-#             ):
-#                 # Convert distance to similarity score (1 - normalized_distance)
-#                 score = 1 - (
-#                     distance / 2
-#                 )  # Normalize cosine distance to [0,1]
-#                 if score >= score_threshold:
-#                     matches.append((metadata["name"], score))
-
-#             logger.info(f"Found {len(matches)} matches for task")
-#             return matches
-
-#         except Exception as e:
-#             logger.error(f"Error finding matches for task: {str(e)}")
-#             raise
-
-#     def auto_select_swarm(self, task: str) -> str:
-#         """
-#         Automatically select the best swarm type for a task
-#         Returns only the top match
-#         """
-#         matches = self.find_best_matches(task, n_results=1)
-#         if not matches:
-#             logger.warning("No suitable matches found for task")
-#             return "SequentialWorkflow"  # Default fallback
-
-#         best_match, score = matches[0]
-#         logger.info(
-#             f"Selected swarm type '{best_match}' with confidence {score:.3f}"
-#         )
-#         return best_match
-
-#     def run_multiple(self, tasks: List[str]) -> List[str]:
-#         """Process multiple tasks in batch"""
-#         return [self.auto_select_swarm(task) for task in tasks]
-
-#     def save_swarm_types(self, filename: str) -> None:
-#         """Export swarm types to JSON"""
-#         try:
-#             all_data = self.collection.get(
-#                 include=["metadatas", "documents"]
-#             )
-#             swarm_types = [
-#                 SwarmType(
-#                     id=id_,
-#                     name=metadata["name"],
-#                     description=document,
-#                     metadata={
-#                         k: v
-#                         for k, v in metadata.items()
-#                         if k != "name"
-#                     },
-#                 )
-#                 for id_, metadata, document in zip(
-#                     all_data["ids"],
-#                     all_data["metadatas"],
-#                     all_data["documents"],
-#                 )
-#             ]
-
-#             with open(filename, "w") as f:
-#                 json.dump(
-#                     [st.dict() for st in swarm_types], f, indent=2
-#                 )
-#             logger.info(f"Saved swarm types to {filename}")
-#         except Exception as e:
-#             logger.error(f"Error saving swarm types: {str(e)}")
-#             raise
-
-#     def load_swarm_types(self, filename: str) -> None:
-#         """Import swarm types from JSON"""
-#         try:
-#             with open(filename, "r") as f:
-#                 swarm_types_data = json.load(f)
-#             swarm_types = [SwarmType(**st) for st in swarm_types_data]
-#             self.add_swarm_types(swarm_types)
-#             logger.info(f"Loaded swarm types from {filename}")
-#         except Exception as e:
-#             logger.error(f"Error loading swarm types: {str(e)}")
-#             raise
-
-
-# def initialize_default_swarm_types(matcher: SwarmMatcher) -> None:
-#     """Initialize the matcher with default swarm types"""
-#     swarm_types = [
-#         SwarmType(
-#             name="AgentRearrange",
-#             description="""
-#             Optimize agent order and rearrange flow for multi-step tasks, ensuring efficient task allocation
-#             and minimizing bottlenecks. Specialized in orchestration, coordination, pipeline optimization,
-#             task scheduling, resource allocation, workflow management, agent organization, and process optimization.
-#             Best for tasks requiring complex agent interactions and workflow optimization.
-#             """,
-#             metadata={
-#                 "category": "optimization",
-#                 "complexity": "high",
-#             },
-#         ),
-#         SwarmType(
-#             name="MixtureOfAgents",
-#             description="""
-#             Combine diverse expert agents for comprehensive analysis, fostering a collaborative approach
-#             to problem-solving and leveraging individual strengths. Focuses on multi-agent systems,
-#             expert collaboration, distributed intelligence, collective problem solving, agent specialization,
-#             team coordination, hybrid approaches, and knowledge synthesis. Ideal for complex problems
-#             requiring multiple areas of expertise.
-#             """,
-#             metadata={
-#                 "category": "collaboration",
-#                 "complexity": "high",
-#             },
-#         ),
-#         SwarmType(
-#             name="SpreadSheetSwarm",
-#             description="""
-#             Collaborative data processing and analysis in a spreadsheet-like environment, facilitating
-#             real-time data sharing and visualization. Specializes in data analysis, tabular processing,
-#             collaborative editing, data transformation, spreadsheet operations, data visualization,
-#             real-time collaboration, and structured data handling. Perfect for data-intensive tasks
-#             requiring structured analysis.
-#             """,
-#             metadata={
-#                 "category": "data_processing",
-#                 "complexity": "medium",
-#             },
-#         ),
-#         SwarmType(
-#             name="SequentialWorkflow",
-#             description="""
-#             Execute tasks in a step-by-step, sequential process workflow, ensuring a logical and methodical
-#             approach to task execution. Focuses on linear processing, waterfall methodology, step-by-step
-#             execution, ordered tasks, sequential operations, process flow, systematic approach, and staged
-#             execution. Best for tasks requiring strict order and dependencies.
-#             """,
-#             metadata={"category": "workflow", "complexity": "low"},
-#         ),
-#         SwarmType(
-#             name="ConcurrentWorkflow",
-#             description="""
-#             Process multiple tasks or data sources concurrently in parallel, maximizing productivity
-#             and reducing processing time. Specializes in parallel processing, multi-threading,
-#             asynchronous execution, distributed computing, concurrent operations, simultaneous tasks,
-#             parallel workflows, and scalable processing. Ideal for independent tasks that can be
-#             processed simultaneously.
-#             """,
-#             metadata={"category": "workflow", "complexity": "medium"},
-#         ),
-#     ]
-
-#     matcher.add_swarm_types(swarm_types)
-#     logger.info("Initialized default swarm types")
-
-
-# def create_swarm_matcher(
-#     persist_dir: str = "./chroma_db",
-#     collection_name: str = "swarm_types",
-# ) -> SwarmMatcher:
-#     """Convenience function to create and initialize a swarm matcher"""
+# # Example usage
+# if __name__ == "__main__":
+#     # Create configuration
 #     config = SwarmMatcherConfig(
-#         persist_directory=persist_dir, collection_name=collection_name
+#         backend="openai",  # Using local embeddings for this example
+#         similarity_threshold=0.6,  # Increase threshold for more strict matching
+#         cache_embeddings=True,
 #     )
+
+#     # Initialize matcher
 #     matcher = SwarmMatcher(config)
-#     initialize_default_swarm_types(matcher)
-#     return matcher
 
+#     task = "I need to concurrently run 1000 tasks"
 
-# # Example usage
-# def swarm_matcher(task: str) -> str:
-#     # Create and initialize matcher
-#     matcher = create_swarm_matcher()
-
-#     swarm_type = matcher.auto_select_swarm(task)
-#     print(f"Task: {task}\nSelected Swarm: {swarm_type}\n")
-
-#     return swarm_type
-
-
-# # # Example usage
-# # if __name__ == "__main__":
-# #     # Create and initialize matcher
-# #     matcher = create_swarm_matcher()
-
-# #     # Example tasks
-# #     tasks = [
-# #         "Analyze this spreadsheet of sales data and create visualizations",
-# #         "Coordinate multiple AI agents to solve a complex problem",
-# #         "Process these tasks one after another in a specific order",
-# #         "Write multiple blog posts about the latest advancements in swarm intelligence all at once",
-# #         "Write a blog post about the latest advancements in swarm intelligence",
-# #     ]
-
-# #     # Process tasks
-# #     for task in tasks:
-# #         swarm_type = matcher.auto_select_swarm(task)
-# #         print(f"Task: {task}\nSelected Swarm: {swarm_type}\n")
+#     print(matcher.auto_select_swarm(task))