File size: 6,108 Bytes
174e0f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
"""
API client for interacting with the Modius Agent Performance API.
"""
import requests
import logging
from typing import List, Dict, Any, Optional

from ..config.constants import API_BASE_URL, DATA_CONFIG
from .models import (
    AgentTypeResponse, AttributeDefinitionResponse, 
    AgentsListResponse, AttributeValuesResponse
)

logger = logging.getLogger(__name__)


class ModiusAPIClient:
    """Client for interacting with the Modius Agent Performance API."""
    
    def __init__(self, base_url: str = API_BASE_URL):
        self.base_url = base_url
        self.session = requests.Session()
        # Set default timeout and headers
        self.session.timeout = 30
        self.session.headers.update({
            'User-Agent': 'Modius-Agent-Performance/1.0'
        })
    
    def _make_request(self, endpoint: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Make a request to the API and return the response."""
        url = f"{self.base_url}{endpoint}"
        logger.debug(f"Making API request to: {url}")
        
        try:
            response = self.session.get(url, params=params)
            logger.debug(f"Response status: {response.status_code}")
            
            if response.status_code == 404:
                logger.warning(f"Resource not found: {url}")
                return {"error": "Not found", "status_code": 404}
            
            response.raise_for_status()
            return {"data": response.json(), "status_code": response.status_code}
            
        except requests.exceptions.RequestException as e:
            logger.error(f"API request error for {url}: {e}")
            return {"error": str(e), "status_code": getattr(e.response, 'status_code', 500)}
    
    def get_agent_type_by_name(self, type_name: str) -> AgentTypeResponse:
        """Get agent type by name."""
        endpoint = f"/api/agent-types/name/{type_name}"
        result = self._make_request(endpoint)
        
        if "error" in result:
            return AgentTypeResponse({}, result["status_code"])
        
        return AgentTypeResponse(result["data"], result["status_code"])
    
    def get_attribute_definition_by_name(self, attr_name: str) -> AttributeDefinitionResponse:
        """Get attribute definition by name."""
        endpoint = f"/api/attributes/name/{attr_name}"
        result = self._make_request(endpoint)
        
        if "error" in result:
            return AttributeDefinitionResponse({}, result["status_code"])
        
        return AttributeDefinitionResponse(result["data"], result["status_code"])
    
    def get_agents_by_type(self, type_id: int) -> AgentsListResponse:
        """Get all agents of a specific type."""
        endpoint = f"/api/agent-types/{type_id}/agents/"
        result = self._make_request(endpoint)
        
        if "error" in result:
            return AgentsListResponse([], result["status_code"])
        
        return AgentsListResponse(result["data"], result["status_code"])
    
    def get_agent_attributes(self, agent_id: int, limit: int = None) -> AttributeValuesResponse:
        """Get attributes for a specific agent."""
        endpoint = f"/api/agents/{agent_id}/attributes/"
        params = {"limit": limit or DATA_CONFIG["api_limit"]}
        result = self._make_request(endpoint, params)
        
        if "error" in result:
            return AttributeValuesResponse([], result["status_code"])
        
        return AttributeValuesResponse(result["data"], result["status_code"])
    
    def get_attribute_values_by_type_and_attr(
        self, 
        agents: List[Dict[str, Any]], 
        attr_def_id: int
    ) -> List[Dict[str, Any]]:
        """Get all attribute values for a specific attribute definition across all agents."""
        all_attributes = []
        logger.debug(f"Getting attributes for {len(agents)} agents with attr_def_id: {attr_def_id}")
        
        for agent in agents:
            agent_id = agent["agent_id"]
            
            # Get agent attributes
            response = self.get_agent_attributes(agent_id)
            
            if not response.is_success():
                logger.error(f"Failed to get attributes for agent ID {agent_id}")
                continue
            
            agent_attrs = response.get_attribute_values()
            logger.debug(f"Agent {agent_id} has {len(agent_attrs)} attributes")
            
            # Filter for the specific attribute definition ID
            filtered_attrs = [
                attr for attr in agent_attrs 
                if attr.get("attr_def_id") == attr_def_id
            ]
            logger.debug(f"Agent {agent_id} has {len(filtered_attrs)} matching attributes")
            
            if filtered_attrs:
                logger.debug(f"Sample attribute for agent {agent_id}: {filtered_attrs[0]}")
            
            all_attributes.extend(filtered_attrs)
        
        logger.info(f"Total matching attributes found across all agents: {len(all_attributes)}")
        return all_attributes
    
    def close(self):
        """Close the session."""
        self.session.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# Utility functions for backward compatibility
def get_agent_name(agent_id: int, agents: List[Dict[str, Any]]) -> str:
    """Get agent name from agent ID."""
    for agent in agents:
        if agent["agent_id"] == agent_id:
            return agent["agent_name"]
    return "Unknown"


# Global client instance (can be replaced with dependency injection later)
_client_instance = None


def get_api_client() -> ModiusAPIClient:
    """Get the global API client instance."""
    global _client_instance
    if _client_instance is None:
        _client_instance = ModiusAPIClient()
    return _client_instance


def close_api_client():
    """Close the global API client instance."""
    global _client_instance
    if _client_instance is not None:
        _client_instance.close()
        _client_instance = None