File size: 2,397 Bytes
100c46f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
"""
Base agent interface for all specialized agents.
"""
from typing import List, Dict, Any, Optional
from abc import ABC, abstractmethod
from smolagents import LiteLLMModel, CodeAgent, ToolCollection


class BaseAgent(ABC):
    """Abstract base class for all specialized agents."""
    
    def __init__(self, name: str, tools: List[Any], model: LiteLLMModel):
        self.name = name
        self.agent = CodeAgent(tools=tools, model=model)
        self.capabilities = self._define_capabilities()
    
    @abstractmethod
    def _define_capabilities(self) -> Dict[str, float]:
        """Define what this agent is capable of handling and with what confidence."""
        pass
    
    def can_handle(self, task: str) -> float:
        """Return confidence score for handling a task."""
        task = task.lower()
        max_confidence = 0.0
        
        # Check each word in the task against capabilities
        words = task.split()
        for word in words:
            for capability, confidence in self.capabilities.items():
                # Check if the capability appears as a whole word
                if capability in task and (
                    capability == word or  # Exact word match
                    f" {capability} " in f" {task} " or  # Word with spaces
                    capability.startswith(word) or  # Word is prefix
                    word.startswith(capability)  # Capability is prefix
                ):
                    max_confidence = max(max_confidence, confidence)
        
        return max_confidence
    
    @abstractmethod
    def execute(self, task: str) -> str:
        """Execute the task and return the result."""
        pass
    
    def _execute_through_agent(self, task: str) -> str:
        """Common method to execute task through the underlying agent."""
        try:
            result = self.agent.run(task)
            
            # If result is a dictionary with a 'result' key, extract it
            if isinstance(result, dict) and 'result' in result:
                return result['result']
            
            # If result is already a string, return it
            if isinstance(result, str):
                return result
                
            # Otherwise convert to string
            return str(result)
            
        except Exception as e:
            return f"❌ Error executing task: {str(e)}"