diff --git a/python-agent/.gitignore b/python-agent/.gitignore new file mode 100644 index 00000000..d91333b0 --- /dev/null +++ b/python-agent/.gitignore @@ -0,0 +1,48 @@ +# Python cache files +__pycache__/ +*.py[cod] +*$py.class + +# Test databases +*.db + +# Virtual environments +venv/ +env/ +.env + +# IDE files +.vscode/ +.idea/ +*.swp +*.swo + +# Logs +*.log + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt \ No newline at end of file diff --git a/python-agent/README.md b/python-agent/README.md new file mode 100644 index 00000000..f8cfc432 --- /dev/null +++ b/python-agent/README.md @@ -0,0 +1,288 @@ +# Sentrius Python Agent + +This Python agent provides the same APIs and operations as the Java agent, enabling integration with the Sentrius platform for authentication, registration, heartbeat monitoring, and provenance event submission. + +## Features + +- **Keycloak Integration**: Full authentication support using Keycloak JWT tokens +- **Agent Registration**: Automatic registration with the Sentrius API server +- **Heartbeat Monitoring**: Continuous heartbeat mechanism to maintain connection +- **Provenance Events**: Submit detailed provenance events for audit trails +- **RSA Encryption**: Secure communication using ephemeral RSA keys +- **Configurable**: Support for both YAML configuration files and environment variables +- **Extensible**: Base agent framework for creating custom agents + +## Architecture + +The Python agent mirrors the Java agent architecture with these key components: + +### Services +- **KeycloakService**: Handles authentication and token management +- **AgentClientService**: Manages API communication with Sentrius server +- **EphemeralKeyGen**: RSA key generation and cryptographic operations +- **SentriusAgent**: Main agent framework coordinating all services + +### Agent Framework +- **BaseAgent**: Abstract base class for all agents + +## Configuration + +The Python agent supports multiple configuration approaches, similar to the Java agent: + +### Properties-based Configuration (Recommended) + +Similar to the Java agent's `application.properties`, you can use a properties file that references agent-specific YAML files: + +**application.properties:** +```properties +# Keycloak Configuration +keycloak.realm=sentrius +keycloak.base-url=${KEYCLOAK_BASE_URL:http://localhost:8180} +keycloak.client-id=${KEYCLOAK_CLIENT_ID:python-agents} +keycloak.client-secret=${KEYCLOAK_CLIENT_SECRET:your-client-secret} + +# Agent Configuration +agent.name.prefix=python-agent +agent.type=python +agent.callback.url=${AGENT_CALLBACK_URL:http://localhost:8093} +agent.api.url=${AGENT_API_URL:http://localhost:8080/} + +# Agent Definitions - these reference YAML files that define agent behavior +agent.chat.helper.config=chat-helper.yaml +agent.chat.helper.enabled=true + +agent.data.analyst.config=data-analyst.yaml +agent.data.analyst.enabled=false +``` + +**chat-helper.yaml:** +```yaml +description: "Agent that handles chat interactions and provides helpful responses via OpenAI integration." +context: | + You are a helpful agent that is responding to users via a chat interface. Please respond to the user in a friendly and helpful manner. + Return responses in the following format: + { + "previousOperation": "", + "nextOperation": "", + "terminalSummaryForLLM": "", + "responseForUser": "" + } +``` + +### Environment Variable Substitution + +Properties support environment variable substitution using the `${VARIABLE:default}` syntax: +- `${KEYCLOAK_BASE_URL:http://localhost:8180}` - Uses `KEYCLOAK_BASE_URL` env var or defaults to `http://localhost:8180` +- `${KEYCLOAK_CLIENT_ID:python-agents}` - Uses `KEYCLOAK_CLIENT_ID` env var or defaults to `python-agents` + +### Legacy YAML Configuration +```yaml +keycloak: + server_url: "http://localhost:8080" + realm: "sentrius" + client_id: "python-agent" + client_secret: "your-client-secret" + +agent: + name_prefix: "python-agent" + agent_type: "python" + callback_url: "http://localhost:8081" + api_url: "http://localhost:8080" + heartbeat_interval: 30 + +llm: + enabled: false + provider: "openai" + model: "gpt-3.5-turbo" + api_key: null + endpoint: null +``` + +### Environment Variables +```bash +KEYCLOAK_SERVER_URL=http://localhost:8080 +KEYCLOAK_REALM=sentrius +KEYCLOAK_CLIENT_ID=python-agent +KEYCLOAK_CLIENT_SECRET=your-client-secret +AGENT_NAME_PREFIX=python-agent +AGENT_API_URL=http://localhost:8080 +AGENT_CALLBACK_URL=http://localhost:8081 +AGENT_HEARTBEAT_INTERVAL=30 +``` + +## Usage + +### Running Agents + +#### With Properties Configuration +```bash +# Run chat helper agent using default application.properties +python main.py chat-helper + +# Run with custom configuration file +python main.py chat-helper --config my-app.properties + +# Run with task data +python main.py chat-helper --task-data '{"message": "Hello, how can you help?"}' + +# Test mode (no external services) +TEST_MODE=true python main.py chat-helper +``` + +#### Environment Variable Configuration +```bash +# Set environment variables +export KEYCLOAK_BASE_URL=http://localhost:8180 +export KEYCLOAK_CLIENT_ID=python-agents +export KEYCLOAK_CLIENT_SECRET=your-secret +export TEST_MODE=false + +# Run agent +python main.py chat-helper +``` + +### Agent Framework +The Python agent provides a framework for creating custom agents that integrate with the Sentrius platform. All agents interact through APIs using JWT authentication, working with DTOs from the API and the LLM proxy. + +### Creating Custom Agents +```python +from agents.base import BaseAgent + +class MyCustomAgent(BaseAgent): + def __init__(self, config_manager): + super().__init__(config_manager, name="my-custom-agent") + # Load agent-specific configuration + self.agent_definition = config_manager.get_agent_definition('my.custom') + + def execute_task(self, task_data=None): + # Your custom agent logic here + # Note: All data access is through Sentrius APIs, not direct database connections + self.submit_provenance( + event_type="CUSTOM_TASK", + details={"task": "custom_operation", "data": task_data} + ) + + # Return structured response + return { + "status": "completed", + "result": "Custom task executed successfully" + } +``` + +**my-custom.yaml:** +```yaml +description: "Custom agent that performs specialized tasks" +context: | + You are a custom agent designed to handle specific business logic. + Process requests according to your specialized capabilities. +``` + +**Add to application.properties:** +```properties +agent.my.custom.config=my-custom.yaml +agent.my.custom.enabled=true +``` + +### Running Custom Agents +```python +# Add to main.py AVAILABLE_AGENTS dict +from agents.my_custom.my_custom_agent import MyCustomAgent + +AVAILABLE_AGENTS = { + 'chat-helper': ChatHelperAgent, + 'my-custom': MyCustomAgent, # Add your agent here +} +``` + +```bash +# Run your custom agent +python main.py my-custom --task-data '{"operation": "process_data"}' +``` + +## API Operations + +The Python agent supports all the same API operations as the Java agent: + +### Agent Registration +- **Endpoint**: `POST /api/v1/agent/register` +- **Purpose**: Register the agent with the Sentrius API server +- **Authentication**: Keycloak JWT token required + +### Heartbeat +- **Endpoint**: `POST /api/v1/agent/heartbeat` +- **Purpose**: Send periodic status updates to maintain connection +- **Frequency**: Configurable (default: 30 seconds) + +### Provenance Submission +- **Endpoint**: `POST /api/v1/agent/provenance/submit` +- **Purpose**: Submit detailed provenance events for audit trails +- **Data**: Event type, timestamp, agent ID, and custom details + +## Dependencies + +- `requests`: HTTP client for API communication +- `PyJWT`: JWT token handling +- `cryptography`: RSA key generation and encryption +- `pyyaml`: YAML configuration parsing + +Note: The Python agent accesses data through Sentrius APIs using DTOs and the LLM proxy, not through direct database connections. + +## Installation + +```bash +pip install -r requirements.txt +``` + +## Testing + +Run the test suite: +```bash +python tests/test_services.py +``` + +## Security + +- Uses ephemeral RSA key pairs for secure communication +- Validates JWT tokens using Keycloak public keys +- Supports encrypted data exchange with the API server +- Maintains secure token management throughout agent lifecycle + +## Integration with Java Ecosystem + +This Python agent is designed to work seamlessly with the existing Java-based Sentrius infrastructure: + +- Compatible with the same API endpoints +- Uses identical authentication mechanisms +- Submits provenance events in the same format +- Supports the same agent lifecycle management +- Can be launched using the same agent launcher service + +## Example Provenance Events + +The agent automatically submits various provenance events: + +```json +{ + "event_type": "AGENT_REGISTRATION", + "timestamp": "2024-01-01T12:00:00.000Z", + "agent_id": "python-agent-abc123", + "details": { + "agent_id": "python-agent-abc123", + "callback_url": "http://localhost:8081", + "agent_type": "python" + } +} +``` + +```json +{ + "event_type": "SQL_QUERY_SUCCESS", + "timestamp": "2024-01-01T12:01:00.000Z", + "agent_id": "python-agent-abc123", + "details": { + "question_number": 1, + "question": "What are the top 5 customers by revenue?", + "response_length": 245 + } +} +``` \ No newline at end of file diff --git a/python-agent/agents/base.py b/python-agent/agents/base.py index 38961e2f..a3a80c37 100644 --- a/python-agent/agents/base.py +++ b/python-agent/agents/base.py @@ -1,12 +1,113 @@ from abc import ABC, abstractmethod +from typing import Dict, Any, Optional +import logging + +from services.sentrius_agent import SentriusAgent +from services.config import SentriusAgentConfig + +logger = logging.getLogger(__name__) + class BaseAgent(ABC): - """Abstract base class for all agents.""" - def __init__(self, name: str, config: dict): - self.name = name - self.config = config + """Abstract base class for all agents with Sentrius API integration.""" + + def __init__(self, config_manager, name: Optional[str] = None): + self.config_manager = config_manager + self.name = name or self.__class__.__name__.lower().replace('agent', '') + + # Check if we're in test mode + self.test_mode = config_manager.get_property('test.mode', 'false').lower() == 'true' + + if not self.test_mode: + # Load configuration for Sentrius integration + agent_config = config_manager.get_agent_config() + keycloak_config = config_manager.get_keycloak_config() + + # Create SentriusAgentConfig from the loaded configuration + self.config = SentriusAgentConfig( + keycloak_server_url=keycloak_config['server_url'], + keycloak_realm=keycloak_config['realm'], + keycloak_client_id=keycloak_config['client_id'], + keycloak_client_secret=keycloak_config['client_secret'], + agent_name_prefix=agent_config['name_prefix'], + agent_type=agent_config['agent_type'], + agent_callback_url=agent_config['callback_url'], + api_url=agent_config['api_url'], + heartbeat_interval=agent_config['heartbeat_interval'] + ) + + # Initialize Sentrius agent + self.sentrius_agent = SentriusAgent(self.config) + else: + logger.info("Running in test mode - external services disabled") + self.sentrius_agent = None + + logger.info(f"Initialized {self.__class__.__name__}: {config_manager}") @abstractmethod - def run(self): - """Method to execute the agent's task.""" + def execute_task(self, task_data: Optional[Dict[str, Any]] = None): + """Method to execute the agent's specific task.""" pass + + def run(self, task_data: Optional[Dict[str, Any]] = None): + """Main run method that handles agent lifecycle.""" + try: + if self.sentrius_agent and not self.test_mode: + with self.sentrius_agent: + logger.info(f"Starting {self.name} agent") + + # Submit start event + self.sentrius_agent.submit_provenance_event( + event_type="AGENT_START", + details={ + "agent_name": self.name, + "agent_class": self.__class__.__name__ + } + ) + + # Execute the specific task + result = self.execute_task(task_data) + + # Submit completion event + self.sentrius_agent.submit_provenance_event( + event_type="AGENT_COMPLETE", + details={ + "agent_name": self.name, + "status": "completed" + } + ) + + logger.info(f"{self.name} agent completed successfully") + return result + else: + # Test mode - just execute the task + logger.info(f"Starting {self.name} agent (test mode)") + result = self.execute_task(task_data) + logger.info(f"{self.name} agent completed successfully (test mode)") + return result + + except Exception as e: + logger.error(f"{self.name} agent failed: {e}") + + # Submit error event + if self.sentrius_agent and not self.test_mode: + try: + self.sentrius_agent.submit_provenance_event( + event_type="AGENT_ERROR", + details={ + "agent_name": self.name, + "error": str(e), + "error_type": type(e).__name__ + } + ) + except: + pass # Don't fail if we can't submit error event + + raise + + def submit_provenance(self, event_type: str, details: Dict[str, Any]): + """Submit a provenance event.""" + if self.sentrius_agent and not self.test_mode: + self.sentrius_agent.submit_provenance_event(event_type, details) + else: + logger.info(f"Test mode - would submit provenance: {event_type} - {details}") diff --git a/python-agent/agents/sql_agent/__init__.py b/python-agent/agents/chat_helper/__init__.py similarity index 100% rename from python-agent/agents/sql_agent/__init__.py rename to python-agent/agents/chat_helper/__init__.py diff --git a/python-agent/agents/chat_helper/chat_helper_agent.py b/python-agent/agents/chat_helper/chat_helper_agent.py new file mode 100644 index 00000000..83f559f0 --- /dev/null +++ b/python-agent/agents/chat_helper/chat_helper_agent.py @@ -0,0 +1,83 @@ +""" +Chat Helper Agent - Provides conversational AI assistance. +""" +import logging +from typing import Dict, Any, Optional +from agents.base import BaseAgent + +logger = logging.getLogger(__name__) + + +class ChatHelperAgent(BaseAgent): + """Agent that provides chat-based assistance using LLM integration.""" + + def __init__(self, config_manager): + super().__init__(config_manager) + self.agent_definition = config_manager.get_agent_definition('chat.helper') + if not self.agent_definition: + raise ValueError("Chat helper agent configuration not found") + + logger.info(f"Initialized ChatHelperAgent: {self.agent_definition.get('description', 'No description')}") + + def execute_task(self, task_data: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """Execute chat helper task.""" + try: + # Submit provenance for task start + self.submit_provenance("CHAT_TASK_START", { + "agent_type": "chat-helper", + "task_data": task_data + }) + + # Process the chat request (this would integrate with LLM) + response = self._process_chat_request(task_data) + + # Submit provenance for task completion + self.submit_provenance("CHAT_TASK_COMPLETE", { + "agent_type": "chat-helper", + "response": response + }) + + return response + + except Exception as e: + logger.error(f"Error executing chat helper task: {e}") + self.submit_provenance("CHAT_TASK_ERROR", { + "agent_type": "chat-helper", + "error": str(e) + }) + raise + + def _process_chat_request(self, task_data: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Process chat request using agent context and LLM.""" + if not task_data: + return { + "previousOperation": "initialization", + "nextOperation": "waiting_for_user_input", + "terminalSummaryForLLM": "Chat helper agent initialized and ready", + "responseForUser": "Hello! I'm your chat helper agent. How can I assist you today?" + } + + user_message = task_data.get('message', '') + context = self.agent_definition.get('context', '') + + # This would integrate with the LLM service + # For now, return a structured response based on the agent's context + return { + "previousOperation": "user_message_received", + "nextOperation": "generate_response", + "terminalSummaryForLLM": f"User asked: {user_message}", + "responseForUser": f"I received your message: '{user_message}'. I'm a helpful chat assistant ready to help!" + } + + def get_agent_info(self) -> Dict[str, Any]: + """Get information about this agent.""" + return { + "name": "chat-helper", + "type": "conversational", + "description": self.agent_definition.get('description', ''), + "capabilities": [ + "conversational_ai", + "user_assistance", + "structured_responses" + ] + } \ No newline at end of file diff --git a/python-agent/agents/sql_agent/sql_agent.py b/python-agent/agents/sql_agent/sql_agent.py deleted file mode 100644 index 359cf537..00000000 --- a/python-agent/agents/sql_agent/sql_agent.py +++ /dev/null @@ -1,33 +0,0 @@ -import yaml -from langchain.chat_models import ChatOpenAI -from langchain_experimental.sql import SQLDatabaseSequentialChain -from langchain.sql_database import SQLDatabase -from agents.base_agent import BaseAgent - -class SQLAgent(BaseAgent): - """SQL Agent using SQLDatabaseSequentialChain.""" - def __init__(self, config_path: str): - # Load configuration - with open(config_path, "r") as file: - config = yaml.safe_load(file) - super().__init__("SQLAgent", config) - - self.db_url = config.get("database_url") - self.questions_file = config.get("questions_file") - self.model_name = config.get("model_name", "gpt-4") - - # Initialize LangChain components - self.db = SQLDatabase.from_uri(self.db_url) - self.llm = ChatOpenAI(model=self.model_name) - self.chain = SQLDatabaseSequentialChain.from_llm(self.llm, self.db, verbose=True) - - def run(self): - """Executes the agent: loads questions and runs against the database.""" - with open(self.questions_file, "r") as file: - questions = yaml.safe_load(file) - - print(f"Running SQL Agent with {len(questions)} questions:") - for idx, question in enumerate(questions, start=1): - print(f"\nQuestion {idx}: {question}") - response = self.chain.run(question) - print(f"Answer: {response}") diff --git a/python-agent/application.properties b/python-agent/application.properties new file mode 100644 index 00000000..ddce1bc8 --- /dev/null +++ b/python-agent/application.properties @@ -0,0 +1,47 @@ +# Sentrius Python Agent Configuration + +# Test Mode - set to true to disable external service connections +test.mode=${TEST_MODE:false} + +# Keycloak Configuration +keycloak.realm=sentrius +keycloak.base-url=${KEYCLOAK_BASE_URL:http://localhost:8180} +keycloak.client-id=${KEYCLOAK_CLIENT_ID:python-agents} +keycloak.client-secret=${KEYCLOAK_CLIENT_SECRET:your-client-secret} + +# Agent Configuration +agent.name.prefix=python-agent +agent.type=python +agent.callback.url=${AGENT_CALLBACK_URL:http://localhost:8093} +agent.api.url=${AGENT_API_URL:http://localhost:8080/} +agent.heartbeat.interval=30 + +# LLM Configuration +agent.llm.endpoint=${LLM_ENDPOINT:http://localhost:8084/} +agent.llm.enabled=true + +# Agent Definitions - these reference YAML files that define agent behavior +agent.chat.helper.config=chat-helper.yaml +agent.chat.helper.enabled=true + +agent.data.analyst.config=data-analyst.yaml +agent.data.analyst.enabled=false + +agent.terminal.helper.config=terminal-helper.yaml +agent.terminal.helper.enabled=false + +# OpenTelemetry Configuration +otel.exporter.otlp.endpoint=${OTEL_EXPORTER_OTLP_ENDPOINT:http://localhost:4317} +otel.traces.exporter=otlp +otel.exporter.otlp.protocol=grpc +otel.metrics.exporter=none +otel.logs.exporter=none +otel.resource.attributes.service.name=python-agent +otel.traces.sampler=always_on +otel.exporter.otlp.timeout=10s + +# Provenance Configuration +provenance.kafka.topic=sentrius-provenance + +# Logging Configuration +logging.level=INFO \ No newline at end of file diff --git a/python-agent/chat-helper.yaml b/python-agent/chat-helper.yaml new file mode 100644 index 00000000..f8ecd04f --- /dev/null +++ b/python-agent/chat-helper.yaml @@ -0,0 +1,10 @@ +description: "Agent that handles chat interactions and provides helpful responses via OpenAI integration." +context: | + You are a helpful agent that is responding to users via a chat interface. Please respond to the user in a friendly and helpful manner. + Return responses in the following format: + { + "previousOperation": "", + "nextOperation": "", + "terminalSummaryForLLM": "", + "responseForUser": "" + } \ No newline at end of file diff --git a/python-agent/config.yaml b/python-agent/config.yaml new file mode 100644 index 00000000..e87ba419 --- /dev/null +++ b/python-agent/config.yaml @@ -0,0 +1,20 @@ +# Sentrius Python Agent Configuration +keycloak: + server_url: "http://localhost:8080" + realm: "sentrius" + client_id: "python-agent" + client_secret: "your-client-secret" + +agent: + name_prefix: "python-agent" + agent_type: "python" + callback_url: "http://localhost:8081" + api_url: "http://localhost:8080" + heartbeat_interval: 30 + +llm: + enabled: false + provider: "openai" + model: "gpt-3.5-turbo" + api_key: null + endpoint: null \ No newline at end of file diff --git a/python-agent/data-analyst.yaml b/python-agent/data-analyst.yaml new file mode 100644 index 00000000..cdf07292 --- /dev/null +++ b/python-agent/data-analyst.yaml @@ -0,0 +1,10 @@ +description: "Agent that provides data analysis and insights using AI capabilities." +context: | + You are a data analyst agent that helps users understand and analyze data. You can: + - Analyze datasets and provide insights + - Generate reports and summaries + - Suggest data visualization approaches + - Answer questions about data patterns and trends + + Always provide clear, actionable insights and suggest next steps for data exploration. + Return responses in a structured format that includes key findings and recommendations. \ No newline at end of file diff --git a/python-agent/main.py b/python-agent/main.py index 925e7078..f8de1a3a 100644 --- a/python-agent/main.py +++ b/python-agent/main.py @@ -1,21 +1,79 @@ import argparse -from agents.sql_agent.sql_agent import SQLAgent +import logging +import sys +from pathlib import Path + +# Add the current directory to the Python path +sys.path.append(str(Path(__file__).parent)) + +from utils.config_manager import ConfigManager +from agents.chat_helper.chat_helper_agent import ChatHelperAgent + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) + +AVAILABLE_AGENTS = { + 'chat-helper': ChatHelperAgent, +} + def main(): - parser = argparse.ArgumentParser(description="Run selected agent.") + parser = argparse.ArgumentParser(description="Run selected Sentrius Python agent.") parser.add_argument( "agent", - choices=["sql_agent"], + choices=list(AVAILABLE_AGENTS.keys()), help="Select the agent to run." ) + parser.add_argument( + "--config", + help="Path to agent configuration properties file", + default="application.properties" + ) + parser.add_argument( + "--task-data", + help="JSON string with task data for the agent", + default=None + ) + args = parser.parse_args() - if args.agent == "sql_agent": - print("Initializing SQL Agent...") - sql_agent = SQLAgent("agents/sql_agent/config.yaml") - sql_agent.run() - else: - print("Unknown agent. Exiting.") + try: + # Load configuration + config_manager = ConfigManager(args.config) + + # Check if the requested agent is enabled + if not config_manager.is_agent_enabled(args.agent.replace('-', '.')): + logger.error(f"Agent '{args.agent}' is not enabled in configuration") + return 1 + + # Initialize and run the agent + agent_class = AVAILABLE_AGENTS[args.agent] + agent = agent_class(config_manager) + + logger.info(f"Starting {args.agent} agent...") + + # Parse task data if provided + task_data = None + if args.task_data: + import json + task_data = json.loads(args.task_data) + + # Execute the agent task + result = agent.execute_task(task_data) + + logger.info(f"Agent execution completed successfully") + logger.info(f"Result: {result}") + + return 0 + + except Exception as e: + logger.error(f"Error running agent: {e}") + return 1 + if __name__ == "__main__": - main() + exit(main()) diff --git a/python-agent/requirements.txt b/python-agent/requirements.txt index 58e278eb..fd9ead33 100644 --- a/python-agent/requirements.txt +++ b/python-agent/requirements.txt @@ -1,6 +1,6 @@ argparse -langchain -sqlalchemy -openai pyyaml -psycopg2 \ No newline at end of file +requests>=2.25.0 +PyJWT>=2.4.0 +cryptography>=3.4.0 +dataclasses-json>=0.5.0 \ No newline at end of file diff --git a/python-agent/services/__init__.py b/python-agent/services/__init__.py new file mode 100644 index 00000000..5abec40a --- /dev/null +++ b/python-agent/services/__init__.py @@ -0,0 +1 @@ +# Services module for Sentrius Python Agent \ No newline at end of file diff --git a/python-agent/services/agent_client_service.py b/python-agent/services/agent_client_service.py new file mode 100644 index 00000000..f7577ead --- /dev/null +++ b/python-agent/services/agent_client_service.py @@ -0,0 +1,178 @@ +""" +Agent client service for API communication with Sentrius server. +Equivalent to Java AgentClientService class. +""" +import json +import requests +import logging +import time +from typing import Dict, Any, Optional +from dataclasses import dataclass, asdict + +logger = logging.getLogger(__name__) + + +@dataclass +class AgentRegistrationRequest: + """Agent registration request data.""" + agent_name: str + agent_callback_url: str + + +@dataclass +class AgentHeartbeat: + """Agent heartbeat data.""" + status: str = "ACTIVE" + last_activity: Optional[str] = None + message: Optional[str] = None + + +@dataclass +class ProvenanceEvent: + """Provenance event data.""" + event_type: str + timestamp: str + agent_id: str + details: Dict[str, Any] + + +@dataclass +class TokenDTO: + """Token data transfer object.""" + access_token: str + token_type: str = "Bearer" + + +class AgentClientService: + """Service for agent API communication with Sentrius server.""" + + def __init__(self, api_base_url: str, keycloak_service): + self.api_base_url = api_base_url.rstrip('/') + self.keycloak_service = keycloak_service + self.session = requests.Session() + + def _get_auth_headers(self) -> Dict[str, str]: + """Get authorization headers with Keycloak token.""" + token = self.keycloak_service.get_keycloak_token() + return { + 'Authorization': f'Bearer {token}', + 'Content-Type': 'application/json' + } + + def register_agent(self, agent_name: str, callback_url: str) -> Dict[str, Any]: + """ + Register agent with the Sentrius API server. + + Args: + agent_name: Name/ID of the agent + callback_url: Callback URL for the agent + + Returns: + Registration response data + """ + url = f"{self.api_base_url}/api/v1/agent/register" + headers = self._get_auth_headers() + + # Registration request doesn't need a body based on Java implementation + try: + response = self.session.post(url, headers=headers) + response.raise_for_status() + return response.json() + except requests.RequestException as e: + logger.error(f"Agent registration failed: {e}") + if hasattr(e, 'response') and e.response is not None: + logger.error(f"Response content: {e.response.text}") + raise + + def send_heartbeat(self, agent_id: str, status: str = "ACTIVE", + message: str = None) -> Dict[str, Any]: + """ + Send heartbeat to the Sentrius API server. + + Args: + agent_id: Agent identifier + status: Agent status + message: Optional status message + + Returns: + Heartbeat response data + """ + url = f"{self.api_base_url}/api/v1/agent/heartbeat" + headers = self._get_auth_headers() + + heartbeat_data = AgentHeartbeat( + status=status, + last_activity=time.strftime('%Y-%m-%dT%H:%M:%S.%fZ'), + message=message + ) + + try: + response = self.session.post( + url, + headers=headers, + json=asdict(heartbeat_data) + ) + response.raise_for_status() + return response.json() + except requests.RequestException as e: + logger.error(f"Heartbeat failed: {e}") + if hasattr(e, 'response') and e.response is not None: + logger.error(f"Response content: {e.response.text}") + raise + + def submit_provenance(self, event: ProvenanceEvent) -> Dict[str, Any]: + """ + Submit provenance event to the Sentrius API server. + + Args: + event: Provenance event data + + Returns: + Submission response data + """ + url = f"{self.api_base_url}/api/v1/agent/provenance/submit" + headers = self._get_auth_headers() + + try: + response = self.session.post( + url, + headers=headers, + json=asdict(event) + ) + response.raise_for_status() + return response.json() + except requests.RequestException as e: + logger.error(f"Provenance submission failed: {e}") + if hasattr(e, 'response') and e.response is not None: + logger.error(f"Response content: {e.response.text}") + raise + + def create_session(self, username: str, ip_address: str) -> Dict[str, Any]: + """ + Create a session log entry. + + Args: + username: Username for the session + ip_address: IP address of the client + + Returns: + Session creation response data + """ + # This appears to be a GET endpoint based on Java implementation + url = f"{self.api_base_url}/api/v1/agent/session" + headers = self._get_auth_headers() + + params = { + 'username': username, + 'ipAddress': ip_address + } + + try: + response = self.session.get(url, headers=headers, params=params) + response.raise_for_status() + return response.json() + except requests.RequestException as e: + logger.error(f"Session creation failed: {e}") + if hasattr(e, 'response') and e.response is not None: + logger.error(f"Response content: {e.response.text}") + raise \ No newline at end of file diff --git a/python-agent/services/config.py b/python-agent/services/config.py new file mode 100644 index 00000000..5027e637 --- /dev/null +++ b/python-agent/services/config.py @@ -0,0 +1,133 @@ +""" +Configuration management for Sentrius Python Agent. +""" +import yaml +import os +import logging +from dataclasses import dataclass +from typing import Optional, Dict, Any + +logger = logging.getLogger(__name__) + + +@dataclass +class KeycloakConfig: + """Keycloak configuration.""" + server_url: str + realm: str + client_id: str + client_secret: str + + +@dataclass +class AgentConfig: + """Agent configuration.""" + name_prefix: str + agent_type: str = "python" + callback_url: str = "http://localhost:8080" + api_url: str = "http://localhost:8080" + heartbeat_interval: int = 30 # seconds + + +@dataclass +class LLMConfig: + """LLM configuration.""" + enabled: bool = False + provider: str = "openai" + model: str = "gpt-3.5-turbo" + api_key: Optional[str] = None + endpoint: Optional[str] = None + + +@dataclass +class SentriusAgentConfig: + """Main configuration class for Sentrius Python Agent.""" + keycloak: KeycloakConfig + agent: AgentConfig + llm: LLMConfig + + @classmethod + def from_yaml(cls, config_path: str) -> 'SentriusAgentConfig': + """Load configuration from YAML file.""" + try: + with open(config_path, 'r') as f: + config_data = yaml.safe_load(f) + + return cls( + keycloak=KeycloakConfig(**config_data.get('keycloak', {})), + agent=AgentConfig(**config_data.get('agent', {})), + llm=LLMConfig(**config_data.get('llm', {})) + ) + except Exception as e: + logger.error(f"Failed to load configuration from {config_path}: {e}") + raise + + @classmethod + def from_env(cls) -> 'SentriusAgentConfig': + """Load configuration from environment variables.""" + try: + keycloak_config = KeycloakConfig( + server_url=os.getenv('KEYCLOAK_SERVER_URL', 'http://localhost:8080'), + realm=os.getenv('KEYCLOAK_REALM', 'sentrius'), + client_id=os.getenv('KEYCLOAK_CLIENT_ID', ''), + client_secret=os.getenv('KEYCLOAK_CLIENT_SECRET', '') + ) + + agent_config = AgentConfig( + name_prefix=os.getenv('AGENT_NAME_PREFIX', 'python-agent'), + agent_type=os.getenv('AGENT_TYPE', 'python'), + callback_url=os.getenv('AGENT_CALLBACK_URL', 'http://localhost:8080'), + api_url=os.getenv('AGENT_API_URL', 'http://localhost:8080'), + heartbeat_interval=int(os.getenv('AGENT_HEARTBEAT_INTERVAL', '30')) + ) + + llm_config = LLMConfig( + enabled=os.getenv('LLM_ENABLED', 'false').lower() == 'true', + provider=os.getenv('LLM_PROVIDER', 'openai'), + model=os.getenv('LLM_MODEL', 'gpt-3.5-turbo'), + api_key=os.getenv('LLM_API_KEY'), + endpoint=os.getenv('LLM_ENDPOINT') + ) + + return cls( + keycloak=keycloak_config, + agent=agent_config, + llm=llm_config + ) + except Exception as e: + logger.error(f"Failed to load configuration from environment: {e}") + raise + + def to_dict(self) -> Dict[str, Any]: + """Convert configuration to dictionary.""" + return { + 'keycloak': { + 'server_url': self.keycloak.server_url, + 'realm': self.keycloak.realm, + 'client_id': self.keycloak.client_id, + 'client_secret': self.keycloak.client_secret + }, + 'agent': { + 'name_prefix': self.agent.name_prefix, + 'agent_type': self.agent.agent_type, + 'callback_url': self.agent.callback_url, + 'api_url': self.agent.api_url, + 'heartbeat_interval': self.agent.heartbeat_interval + }, + 'llm': { + 'enabled': self.llm.enabled, + 'provider': self.llm.provider, + 'model': self.llm.model, + 'api_key': self.llm.api_key, + 'endpoint': self.llm.endpoint + } + } + + def save_to_yaml(self, config_path: str): + """Save configuration to YAML file.""" + try: + with open(config_path, 'w') as f: + yaml.dump(self.to_dict(), f, default_flow_style=False) + except Exception as e: + logger.error(f"Failed to save configuration to {config_path}: {e}") + raise \ No newline at end of file diff --git a/python-agent/services/key_service.py b/python-agent/services/key_service.py new file mode 100644 index 00000000..2b7c443f --- /dev/null +++ b/python-agent/services/key_service.py @@ -0,0 +1,107 @@ +""" +RSA key generation utilities for secure communication. +Equivalent to Java EphemeralKeyGen class. +""" +import base64 +import logging +from cryptography.hazmat.primitives import serialization, hashes +from cryptography.hazmat.primitives.asymmetric import rsa, padding +from cryptography.hazmat.backends import default_backend +from typing import Tuple + +logger = logging.getLogger(__name__) + + +class EphemeralKeyGen: + """Utility class for RSA key generation and cryptographic operations.""" + + @staticmethod + def generate_ephemeral_rsa_keypair(key_size: int = 2048) -> Tuple: + """ + Generate an ephemeral RSA key pair. + + Args: + key_size: Size of the RSA key (default 2048) + + Returns: + Tuple of (private_key, public_key) + """ + try: + private_key = rsa.generate_private_key( + public_exponent=65537, + key_size=key_size, + backend=default_backend() + ) + public_key = private_key.public_key() + return private_key, public_key + except Exception as e: + logger.error(f"Failed to generate RSA key pair: {e}") + raise + + @staticmethod + def get_base64_public_key(public_key) -> str: + """ + Convert public key to base64 encoded string. + + Args: + public_key: RSA public key object + + Returns: + Base64 encoded public key string + """ + try: + pem = public_key.public_key_pem() if hasattr(public_key, 'public_key_pem') else \ + public_key.public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo + ) + return base64.b64encode(pem).decode('utf-8') + except Exception as e: + logger.error(f"Failed to encode public key to base64: {e}") + raise + + @staticmethod + def decrypt_rsa_with_private_key(encrypted_secret: str, private_key) -> str: + """ + Decrypt RSA encrypted data using private key. + + Args: + encrypted_secret: Base64 encoded encrypted data + private_key: RSA private key object + + Returns: + Decrypted string + """ + try: + encrypted_data = base64.b64decode(encrypted_secret) + decrypted_bytes = private_key.decrypt( + encrypted_data, + padding.PKCS1v15() + ) + return decrypted_bytes.decode('utf-8') + except Exception as e: + logger.error(f"Failed to decrypt RSA data: {e}") + raise + + @staticmethod + def encrypt_rsa_with_public_key(data: str, public_key) -> str: + """ + Encrypt data using RSA public key. + + Args: + data: String data to encrypt + public_key: RSA public key object + + Returns: + Base64 encoded encrypted data + """ + try: + data_bytes = data.encode('utf-8') + encrypted_bytes = public_key.encrypt( + data_bytes, + padding.PKCS1v15() + ) + return base64.b64encode(encrypted_bytes).decode('utf-8') + except Exception as e: + logger.error(f"Failed to encrypt RSA data: {e}") + raise \ No newline at end of file diff --git a/python-agent/services/keycloak_service.py b/python-agent/services/keycloak_service.py new file mode 100644 index 00000000..ebb2913f --- /dev/null +++ b/python-agent/services/keycloak_service.py @@ -0,0 +1,138 @@ +""" +Keycloak service for handling authentication with Keycloak server. +Equivalent to Java KeycloakService class. +""" +import jwt +import requests +import logging +from typing import Optional, Dict, Any +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +import base64 + +logger = logging.getLogger(__name__) + + +class KeycloakService: + """Service for Keycloak authentication and token management.""" + + def __init__(self, server_url: str, realm: str, client_id: str, client_secret: str): + self.server_url = server_url.rstrip('/') + self.realm = realm + self.client_id = client_id + self.client_secret = client_secret + self.public_keys_cache = {} + + def get_keycloak_token(self) -> str: + """Get access token from Keycloak using client credentials.""" + token_url = f"{self.server_url}/realms/{self.realm}/protocol/openid-connect/token" + + data = { + 'grant_type': 'client_credentials', + 'client_id': self.client_id, + 'client_secret': self.client_secret + } + + try: + response = requests.post(token_url, data=data) + response.raise_for_status() + token_data = response.json() + return token_data['access_token'] + except requests.RequestException as e: + logger.error(f"Failed to get Keycloak token: {e}") + raise + + def validate_jwt(self, token: str) -> bool: + """Validate a JWT using Keycloak public key.""" + try: + # Extract kid from JWT header + kid = self._extract_kid(token) + if not kid: + logger.error("No 'kid' found in JWT header") + return False + + # Get public key for kid + public_key = self._get_public_key(kid) + if not public_key: + logger.error(f"No public key found for 'kid': {kid}") + return False + + # Validate JWT + jwt.decode(token, public_key, algorithms=['RS256']) + return True + except Exception as e: + logger.error(f"JWT validation failed: {e}") + return False + + def extract_agent_id(self, token: str) -> Optional[str]: + """Extract the client ID (agent identity) from a valid JWT.""" + try: + decoded = jwt.decode(token, options={"verify_signature": False}) + return decoded.get('azp') or decoded.get('client_id') + except Exception as e: + logger.error(f"Failed to extract agent ID: {e}") + return None + + def extract_username(self, token: str) -> Optional[str]: + """Extract username from JWT token.""" + try: + decoded = jwt.decode(token, options={"verify_signature": False}) + return decoded.get('preferred_username') or decoded.get('sub') + except Exception as e: + logger.error(f"Failed to extract username: {e}") + return None + + def _extract_kid(self, token: str) -> Optional[str]: + """Extract the 'kid' (Key ID) from JWT header.""" + try: + header = jwt.get_unverified_header(token) + return header.get('kid') + except Exception as e: + logger.error(f"Failed to extract kid: {e}") + return None + + def _get_public_key(self, kid: str): + """Get public key for the given kid.""" + if kid in self.public_keys_cache: + return self.public_keys_cache[kid] + + try: + # Fetch JWKS from Keycloak + jwks_url = f"{self.server_url}/realms/{self.realm}/protocol/openid-connect/certs" + response = requests.get(jwks_url) + response.raise_for_status() + jwks = response.json() + + # Find the key with matching kid + for key_data in jwks.get('keys', []): + if key_data.get('kid') == kid: + # Convert JWK to public key object + public_key = self._jwk_to_public_key(key_data) + self.public_keys_cache[kid] = public_key + return public_key + + except Exception as e: + logger.error(f"Failed to fetch public key: {e}") + + return None + + def _jwk_to_public_key(self, jwk_data: Dict[str, Any]): + """Convert JWK data to cryptography public key object.""" + try: + from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers + from cryptography.hazmat.backends import default_backend + + n = int.from_bytes( + base64.urlsafe_b64decode(jwk_data['n'] + '=='), + byteorder='big' + ) + e = int.from_bytes( + base64.urlsafe_b64decode(jwk_data['e'] + '=='), + byteorder='big' + ) + + public_numbers = RSAPublicNumbers(e, n) + return public_numbers.public_key(default_backend()) + except Exception as e: + logger.error(f"Failed to convert JWK to public key: {e}") + return None \ No newline at end of file diff --git a/python-agent/services/sentrius_agent.py b/python-agent/services/sentrius_agent.py new file mode 100644 index 00000000..617515d0 --- /dev/null +++ b/python-agent/services/sentrius_agent.py @@ -0,0 +1,163 @@ +""" +Main Sentrius Agent framework that integrates all services. +Equivalent to Java ChatAgent functionality. +""" +import time +import threading +import logging +import uuid +from typing import Optional, Dict, Any +from datetime import datetime, timezone + +from .keycloak_service import KeycloakService +from .agent_client_service import AgentClientService, ProvenanceEvent +from .key_service import EphemeralKeyGen +from .config import SentriusAgentConfig + +logger = logging.getLogger(__name__) + + +class SentriusAgent: + """Main Sentrius Agent framework class.""" + + def __init__(self, config: SentriusAgentConfig): + self.config = config + self.agent_id = f"{config.agent.name_prefix}-{uuid.uuid4().hex[:8]}" + self.running = False + self.heartbeat_thread: Optional[threading.Thread] = None + + # Initialize services + self.keycloak_service = KeycloakService( + server_url=config.keycloak.server_url, + realm=config.keycloak.realm, + client_id=config.keycloak.client_id, + client_secret=config.keycloak.client_secret + ) + + self.agent_client_service = AgentClientService( + api_base_url=config.agent.api_url, + keycloak_service=self.keycloak_service + ) + + # Generate ephemeral keys for secure communication + self.private_key, self.public_key = EphemeralKeyGen.generate_ephemeral_rsa_keypair() + + logger.info(f"Initialized Sentrius Agent: {self.agent_id}") + + def start(self): + """Start the agent and begin registration process.""" + logger.info(f"Starting Sentrius Agent: {self.agent_id}") + + try: + # Register with the API server + self._register_agent() + + # Start heartbeat mechanism + self._start_heartbeat() + + self.running = True + logger.info(f"Sentrius Agent {self.agent_id} started successfully") + + except Exception as e: + logger.error(f"Failed to start agent: {e}") + self.stop() + raise + + def stop(self): + """Stop the agent and cleanup resources.""" + logger.info(f"Stopping Sentrius Agent: {self.agent_id}") + + self.running = False + + # Stop heartbeat thread + if self.heartbeat_thread and self.heartbeat_thread.is_alive(): + self.heartbeat_thread.join(timeout=5) + + logger.info(f"Sentrius Agent {self.agent_id} stopped") + + def submit_provenance_event(self, event_type: str, details: Dict[str, Any]): + """Submit a provenance event to the API server.""" + try: + event = ProvenanceEvent( + event_type=event_type, + timestamp=datetime.now(timezone.utc).isoformat(), + agent_id=self.agent_id, + details=details + ) + + response = self.agent_client_service.submit_provenance(event) + logger.debug(f"Provenance event submitted: {response}") + + except Exception as e: + logger.error(f"Failed to submit provenance event: {e}") + raise + + def get_agent_id(self) -> str: + """Get the agent ID.""" + return self.agent_id + + def get_public_key_base64(self) -> str: + """Get the base64 encoded public key.""" + return EphemeralKeyGen.get_base64_public_key(self.public_key) + + def decrypt_with_private_key(self, encrypted_data: str) -> str: + """Decrypt data using the agent's private key.""" + return EphemeralKeyGen.decrypt_rsa_with_private_key(encrypted_data, self.private_key) + + def _register_agent(self): + """Register the agent with the API server.""" + try: + response = self.agent_client_service.register_agent( + agent_name=self.agent_id, + callback_url=self.config.agent.callback_url + ) + logger.info(f"Agent registration successful: {response}") + + # Submit registration provenance event + self.submit_provenance_event( + event_type="AGENT_REGISTRATION", + details={ + "agent_id": self.agent_id, + "callback_url": self.config.agent.callback_url, + "agent_type": self.config.agent.agent_type + } + ) + + except Exception as e: + logger.error(f"Agent registration failed: {e}") + raise + + def _start_heartbeat(self): + """Start the heartbeat mechanism.""" + if self.heartbeat_thread and self.heartbeat_thread.is_alive(): + return + + self.heartbeat_thread = threading.Thread(target=self._heartbeat_worker, daemon=True) + self.heartbeat_thread.start() + logger.info("Heartbeat mechanism started") + + def _heartbeat_worker(self): + """Heartbeat worker thread.""" + while self.running: + try: + response = self.agent_client_service.send_heartbeat( + agent_id=self.agent_id, + status="ACTIVE", + message="Agent running normally" + ) + logger.debug(f"Heartbeat sent: {response}") + + except Exception as e: + logger.error(f"Heartbeat failed: {e}") + + # Wait for next heartbeat interval + time.sleep(self.config.agent.heartbeat_interval) + + def __enter__(self): + """Context manager entry.""" + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.stop() \ No newline at end of file diff --git a/python-agent/terminal-helper.yaml b/python-agent/terminal-helper.yaml new file mode 100644 index 00000000..9a90c3c6 --- /dev/null +++ b/python-agent/terminal-helper.yaml @@ -0,0 +1,11 @@ +description: "Agent that assists with terminal operations and command-line tasks." +context: | + You are a terminal helper agent that assists users with command-line operations and system administration tasks. + You can help with: + - Explaining command-line tools and their usage + - Suggesting appropriate commands for specific tasks + - Troubleshooting system issues + - Automating repetitive terminal operations + + Always prioritize security and best practices when suggesting commands. + Provide clear explanations of what commands do before suggesting them. \ No newline at end of file diff --git a/python-agent/tests/__init__.py b/python-agent/tests/__init__.py new file mode 100644 index 00000000..18dae6c2 --- /dev/null +++ b/python-agent/tests/__init__.py @@ -0,0 +1 @@ +# Tests module for Sentrius Python Agent \ No newline at end of file diff --git a/python-agent/tests/test_config_manager.py b/python-agent/tests/test_config_manager.py new file mode 100644 index 00000000..74a6d756 --- /dev/null +++ b/python-agent/tests/test_config_manager.py @@ -0,0 +1,125 @@ +""" +Test configuration manager functionality. +""" +import unittest +import tempfile +import os +from pathlib import Path +from utils.config_manager import ConfigManager + + +class TestConfigManager(unittest.TestCase): + """Test the ConfigManager class.""" + + def setUp(self): + """Set up test configuration files.""" + # Create temporary directory for test files + self.test_dir = tempfile.mkdtemp() + self.properties_file = os.path.join(self.test_dir, "test.properties") + self.yaml_file = os.path.join(self.test_dir, "test-agent.yaml") + + # Create test properties file + properties_content = """ +# Test properties +test.mode=true +keycloak.realm=test-realm +keycloak.base-url=${KEYCLOAK_URL:http://localhost:8180} +agent.name.prefix=test-agent +agent.test.config=test-agent.yaml +agent.test.enabled=true +""" + with open(self.properties_file, 'w') as f: + f.write(properties_content) + + # Create test YAML file + yaml_content = """ +description: "Test agent for unit tests" +context: | + You are a test agent used for unit testing. + Always return structured responses for testing. +""" + with open(self.yaml_file, 'w') as f: + f.write(yaml_content) + + def tearDown(self): + """Clean up test files.""" + import shutil + shutil.rmtree(self.test_dir) + + def test_load_properties(self): + """Test loading properties file.""" + config_manager = ConfigManager(self.properties_file) + + self.assertEqual(config_manager.get_property('test.mode'), 'true') + self.assertEqual(config_manager.get_property('keycloak.realm'), 'test-realm') + self.assertEqual(config_manager.get_property('agent.name.prefix'), 'test-agent') + + def test_env_var_substitution(self): + """Test environment variable substitution.""" + # Set environment variable + os.environ['KEYCLOAK_URL'] = 'http://test.example.com:8080' + + config_manager = ConfigManager(self.properties_file) + + # Should use environment variable + self.assertEqual(config_manager.get_property('keycloak.base-url'), 'http://test.example.com:8080') + + # Clean up + del os.environ['KEYCLOAK_URL'] + + def test_env_var_default(self): + """Test environment variable default values.""" + config_manager = ConfigManager(self.properties_file) + + # Should use default value when env var is not set + self.assertEqual(config_manager.get_property('keycloak.base-url'), 'http://localhost:8180') + + def test_agent_config_loading(self): + """Test loading agent configuration.""" + # Change to test directory so relative paths work + original_cwd = os.getcwd() + os.chdir(self.test_dir) + + try: + config_manager = ConfigManager("test.properties") + + agent_config = config_manager.get_agent_definition('test') + self.assertIsNotNone(agent_config) + self.assertEqual(agent_config['description'], 'Test agent for unit tests') + self.assertIn('You are a test agent', agent_config['context']) + finally: + os.chdir(original_cwd) + + def test_enabled_agents(self): + """Test getting enabled agents.""" + config_manager = ConfigManager(self.properties_file) + + enabled_agents = config_manager.get_enabled_agents() + self.assertIn('test', enabled_agents) + + self.assertTrue(config_manager.is_agent_enabled('test')) + self.assertFalse(config_manager.is_agent_enabled('nonexistent')) + + def test_get_configs(self): + """Test getting different configuration sections.""" + config_manager = ConfigManager(self.properties_file) + + keycloak_config = config_manager.get_keycloak_config() + self.assertEqual(keycloak_config['realm'], 'test-realm') + self.assertEqual(keycloak_config['server_url'], 'http://localhost:8180') + + agent_config = config_manager.get_agent_config() + self.assertEqual(agent_config['name_prefix'], 'test-agent') + + def test_missing_files(self): + """Test handling of missing configuration files.""" + # Test with non-existent properties file + config_manager = ConfigManager('/nonexistent/file.properties') + + # Should not crash and should return defaults + self.assertEqual(len(config_manager.properties), 0) + self.assertEqual(config_manager.get_property('missing.key', 'default'), 'default') + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/python-agent/tests/test_integration.py b/python-agent/tests/test_integration.py new file mode 100644 index 00000000..4c1a5dc4 --- /dev/null +++ b/python-agent/tests/test_integration.py @@ -0,0 +1,85 @@ +""" +Integration test with mocked services to validate agent functionality. +""" +import unittest +import os +import sys +from unittest.mock import Mock, patch + +# Add parent directory to Python path for imports +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from services.config import SentriusAgentConfig, KeycloakConfig, AgentConfig, LLMConfig +from services.sentrius_agent import SentriusAgent + + +class TestSentriusAgentIntegration(unittest.TestCase): + """Test the complete agent integration with mocked services.""" + + def setUp(self): + """Set up test environment.""" + self.config = SentriusAgentConfig( + keycloak=KeycloakConfig( + server_url='http://localhost:8080', + realm='test-realm', + client_id='test-client', + client_secret='test-secret' + ), + agent=AgentConfig(name_prefix='test-agent'), + llm=LLMConfig() + ) + + @patch('services.sentrius_agent.KeycloakService') + @patch('services.sentrius_agent.AgentClientService') + def test_sentrius_agent_initialization(self, mock_agent_client, mock_keycloak): + """Test SentriusAgent initialization.""" + # Mock the services + mock_keycloak_instance = Mock() + mock_keycloak.return_value = mock_keycloak_instance + + mock_agent_client_instance = Mock() + mock_agent_client.return_value = mock_agent_client_instance + + # Create agent + agent = SentriusAgent(self.config) + + # Verify initialization + self.assertTrue(agent.agent_id.startswith('test-agent')) + self.assertFalse(agent.running) + + # Verify services were created + mock_keycloak.assert_called_once() + mock_agent_client.assert_called_once() + + @patch('services.sentrius_agent.KeycloakService') + @patch('services.sentrius_agent.AgentClientService') + def test_sentrius_agent_start_stop(self, mock_agent_client, mock_keycloak): + """Test agent start and stop functionality.""" + # Mock the services + mock_keycloak_instance = Mock() + mock_keycloak.return_value = mock_keycloak_instance + + mock_agent_client_instance = Mock() + mock_agent_client_instance.register_agent.return_value = {'status': 'success'} + mock_agent_client_instance.submit_provenance.return_value = {'status': 'success'} + mock_agent_client.return_value = mock_agent_client_instance + + # Create and start agent + agent = SentriusAgent(self.config) + agent.start() + + # Verify agent is running + self.assertTrue(agent.running) + + # Verify registration was called + mock_agent_client_instance.register_agent.assert_called_once() + + # Stop agent + agent.stop() + + # Verify agent is stopped + self.assertFalse(agent.running) + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/python-agent/tests/test_services.py b/python-agent/tests/test_services.py new file mode 100644 index 00000000..4437b97e --- /dev/null +++ b/python-agent/tests/test_services.py @@ -0,0 +1,136 @@ +""" +Basic tests for Sentrius Python Agent services. +""" +import unittest +import os +import sys +from unittest.mock import Mock, patch, MagicMock + +# Add parent directory to Python path for imports +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from services.config import SentriusAgentConfig, KeycloakConfig, AgentConfig, LLMConfig +from services.keycloak_service import KeycloakService +from services.key_service import EphemeralKeyGen + + +class TestSentriusAgentConfig(unittest.TestCase): + """Test configuration management.""" + + def test_config_from_env(self): + """Test loading configuration from environment variables.""" + with patch.dict(os.environ, { + 'KEYCLOAK_SERVER_URL': 'http://test:8080', + 'KEYCLOAK_REALM': 'test-realm', + 'KEYCLOAK_CLIENT_ID': 'test-client', + 'KEYCLOAK_CLIENT_SECRET': 'test-secret', + 'AGENT_NAME_PREFIX': 'test-agent' + }): + config = SentriusAgentConfig.from_env() + + self.assertEqual(config.keycloak.server_url, 'http://test:8080') + self.assertEqual(config.keycloak.realm, 'test-realm') + self.assertEqual(config.keycloak.client_id, 'test-client') + self.assertEqual(config.keycloak.client_secret, 'test-secret') + self.assertEqual(config.agent.name_prefix, 'test-agent') + + def test_config_to_dict(self): + """Test configuration conversion to dictionary.""" + config = SentriusAgentConfig( + keycloak=KeycloakConfig( + server_url='http://test:8080', + realm='test-realm', + client_id='test-client', + client_secret='test-secret' + ), + agent=AgentConfig(name_prefix='test-agent'), + llm=LLMConfig() + ) + + config_dict = config.to_dict() + + self.assertIn('keycloak', config_dict) + self.assertIn('agent', config_dict) + self.assertIn('llm', config_dict) + self.assertEqual(config_dict['keycloak']['server_url'], 'http://test:8080') + + +class TestEphemeralKeyGen(unittest.TestCase): + """Test RSA key generation utilities.""" + + def test_generate_keypair(self): + """Test RSA key pair generation.""" + private_key, public_key = EphemeralKeyGen.generate_ephemeral_rsa_keypair() + + self.assertIsNotNone(private_key) + self.assertIsNotNone(public_key) + + def test_base64_public_key(self): + """Test base64 encoding of public key.""" + private_key, public_key = EphemeralKeyGen.generate_ephemeral_rsa_keypair() + base64_key = EphemeralKeyGen.get_base64_public_key(public_key) + + self.assertIsInstance(base64_key, str) + self.assertTrue(len(base64_key) > 0) + + def test_encrypt_decrypt(self): + """Test RSA encryption and decryption.""" + private_key, public_key = EphemeralKeyGen.generate_ephemeral_rsa_keypair() + test_data = "Hello, World!" + + # Encrypt with public key + encrypted_data = EphemeralKeyGen.encrypt_rsa_with_public_key(test_data, public_key) + + # Decrypt with private key + decrypted_data = EphemeralKeyGen.decrypt_rsa_with_private_key(encrypted_data, private_key) + + self.assertEqual(test_data, decrypted_data) + + +class TestKeycloakService(unittest.TestCase): + """Test Keycloak service functionality.""" + + def setUp(self): + """Set up test environment.""" + self.keycloak_service = KeycloakService( + server_url='http://test:8080', + realm='test-realm', + client_id='test-client', + client_secret='test-secret' + ) + + @patch('requests.post') + def test_get_keycloak_token(self, mock_post): + """Test getting Keycloak token.""" + # Mock successful response + mock_response = Mock() + mock_response.json.return_value = {'access_token': 'test-token'} + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + token = self.keycloak_service.get_keycloak_token() + + self.assertEqual(token, 'test-token') + mock_post.assert_called_once() + + def test_extract_agent_id(self): + """Test extracting agent ID from JWT token.""" + # Create a mock JWT token (not signed, just for testing) + import base64 + import json + + header = {'typ': 'JWT', 'alg': 'RS256'} + payload = {'azp': 'test-agent-id', 'exp': 9999999999} + + header_b64 = base64.urlsafe_b64encode(json.dumps(header).encode()).decode().rstrip('=') + payload_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip('=') + + mock_token = f"{header_b64}.{payload_b64}.mock-signature" + + agent_id = self.keycloak_service.extract_agent_id(mock_token) + + self.assertEqual(agent_id, 'test-agent-id') + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/python-agent/utils/config_manager.py b/python-agent/utils/config_manager.py new file mode 100644 index 00000000..e04e32a1 --- /dev/null +++ b/python-agent/utils/config_manager.py @@ -0,0 +1,125 @@ +""" +Configuration management for the Sentrius Python Agent. +Supports loading from properties files and YAML files, similar to Java Spring configuration. +""" +import os +import yaml +import logging +from typing import Dict, Any, Optional, List +from pathlib import Path + +logger = logging.getLogger(__name__) + + +class ConfigManager: + """Manages configuration loading from properties and YAML files.""" + + def __init__(self, properties_file: str = "application.properties"): + self.properties_file = properties_file + self.properties = {} + self.agent_configs = {} + self._load_properties() + self._load_agent_configs() + + def _load_properties(self): + """Load configuration from properties file with environment variable substitution.""" + properties_path = Path(self.properties_file) + if not properties_path.exists(): + logger.warning(f"Properties file {self.properties_file} not found") + return + + try: + with open(properties_path, 'r') as f: + for line in f: + line = line.strip() + if line and not line.startswith('#'): + if '=' in line: + key, value = line.split('=', 1) + # Handle environment variable substitution + value = self._substitute_env_vars(value) + self.properties[key.strip()] = value.strip() + + logger.info(f"Loaded {len(self.properties)} properties from {self.properties_file}") + except Exception as e: + logger.error(f"Error loading properties file: {e}") + + def _substitute_env_vars(self, value: str) -> str: + """Substitute environment variables in property values.""" + # Handle ${VARIABLE:default} pattern + import re + pattern = r'\$\{([^:}]+):([^}]*)\}' + + def replacer(match): + env_var = match.group(1) + default_val = match.group(2) + return os.getenv(env_var, default_val) + + return re.sub(pattern, replacer, value) + + def _load_agent_configs(self): + """Load agent-specific configuration files referenced in properties.""" + agent_config_keys = [k for k in self.properties.keys() if k.endswith('.config')] + + for config_key in agent_config_keys: + yaml_file = self.properties[config_key] + agent_name = config_key.replace('.config', '').replace('agent.', '') + + try: + yaml_path = Path(yaml_file) + if yaml_path.exists(): + with open(yaml_path, 'r') as f: + config = yaml.safe_load(f) + self.agent_configs[agent_name] = config + logger.info(f"Loaded agent config for {agent_name} from {yaml_file}") + else: + logger.warning(f"Agent config file {yaml_file} not found for {agent_name}") + except Exception as e: + logger.error(f"Error loading agent config {yaml_file}: {e}") + + def get_property(self, key: str, default: Any = None) -> Any: + """Get a property value with optional default.""" + return self.properties.get(key, default) + + def get_keycloak_config(self) -> Dict[str, str]: + """Get Keycloak configuration.""" + return { + 'server_url': self.get_property('keycloak.base-url', 'http://localhost:8180'), + 'realm': self.get_property('keycloak.realm', 'sentrius'), + 'client_id': self.get_property('keycloak.client-id', 'python-agents'), + 'client_secret': self.get_property('keycloak.client-secret') + } + + def get_agent_config(self) -> Dict[str, Any]: + """Get agent configuration.""" + return { + 'name_prefix': self.get_property('agent.name.prefix', 'python-agent'), + 'agent_type': self.get_property('agent.type', 'python'), + 'callback_url': self.get_property('agent.callback.url', 'http://localhost:8093'), + 'api_url': self.get_property('agent.api.url', 'http://localhost:8080/'), + 'heartbeat_interval': int(self.get_property('agent.heartbeat.interval', '30')) + } + + def get_llm_config(self) -> Dict[str, Any]: + """Get LLM configuration.""" + return { + 'endpoint': self.get_property('agent.llm.endpoint', 'http://localhost:8084/'), + 'enabled': self.get_property('agent.llm.enabled', 'true').lower() == 'true' + } + + def get_agent_definition(self, agent_name: str) -> Optional[Dict[str, Any]]: + """Get agent definition from loaded YAML configs.""" + return self.agent_configs.get(agent_name) + + def get_enabled_agents(self) -> List[str]: + """Get list of enabled agents.""" + enabled_agents = [] + for key, value in self.properties.items(): + if key.endswith('.enabled') and value.lower() == 'true': + agent_name = key.replace('.enabled', '').replace('agent.', '') + enabled_agents.append(agent_name) + return enabled_agents + + def is_agent_enabled(self, agent_name: str) -> bool: + """Check if a specific agent is enabled.""" + enabled_key = f'agent.{agent_name}.enabled' + return self.get_property(enabled_key, 'false').lower() == 'true' \ No newline at end of file