242 lines
8.2 KiB
YAML
242 lines
8.2 KiB
YAML
apiVersion: apps/v1
|
|
kind: Deployment
|
|
metadata:
|
|
name: agent
|
|
namespace: swarm-system
|
|
labels:
|
|
app: agent
|
|
component: worker
|
|
spec:
|
|
replicas: 1
|
|
selector:
|
|
matchLabels:
|
|
app: agent
|
|
template:
|
|
metadata:
|
|
labels:
|
|
app: agent
|
|
component: worker
|
|
spec:
|
|
serviceAccountName: orchestrator-sa
|
|
containers:
|
|
- name: agent
|
|
image: python:3.11-slim
|
|
command:
|
|
- /bin/bash
|
|
- -c
|
|
- |
|
|
set -e
|
|
echo "Installing dependencies..."
|
|
pip install --no-cache-dir \
|
|
websockets \
|
|
prometheus-client \
|
|
gitpython \
|
|
openai
|
|
|
|
echo "Setting up agent code..."
|
|
mkdir -p /tmp/agent
|
|
cd /tmp/agent
|
|
|
|
echo "Starting agent..."
|
|
python3 -c "
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from typing import Optional
|
|
import websockets
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class SimpleAgent:
|
|
def __init__(self, orchestrator_url: str, agent_id: str):
|
|
self.orchestrator_url = orchestrator_url
|
|
self.agent_id = agent_id
|
|
self.api_key = os.getenv('OPENAI_API_KEY')
|
|
self.api_base = os.getenv('OPENAI_API_BASE', 'https://api.openai.com/v1')
|
|
self.model = os.getenv('OPENAI_MODEL', 'gpt-4o-mini')
|
|
self.ws = None
|
|
self.running = True
|
|
|
|
async def connect(self):
|
|
logger.info(f'Connecting to orchestrator at {self.orchestrator_url}')
|
|
self.ws = await websockets.connect(self.orchestrator_url)
|
|
|
|
# Register with orchestrator
|
|
register_msg = {
|
|
'type': 'register',
|
|
'agent_id': self.agent_id,
|
|
'capabilities': ['code_generation', 'code_review'],
|
|
'status': 'idle'
|
|
}
|
|
await self.ws.send(json.dumps(register_msg))
|
|
logger.info(f'Agent {self.agent_id} registered')
|
|
|
|
async def send_heartbeat(self):
|
|
while self.running:
|
|
try:
|
|
if self.ws:
|
|
heartbeat_msg = {
|
|
'type': 'heartbeat',
|
|
'agent_id': self.agent_id,
|
|
'status': 'idle',
|
|
'timestamp': time.time()
|
|
}
|
|
await self.ws.send(json.dumps(heartbeat_msg))
|
|
logger.debug(f'Heartbeat sent')
|
|
except Exception as e:
|
|
logger.error(f'Heartbeat error: {e}')
|
|
await asyncio.sleep(15)
|
|
|
|
async def execute_task(self, task_data):
|
|
task_id = task_data.get('task_id')
|
|
description = task_data.get('description')
|
|
|
|
logger.info(f'Executing task {task_id}: {description}')
|
|
|
|
try:
|
|
# Use OpenAI-compatible API
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI(
|
|
api_key=self.api_key,
|
|
base_url=self.api_base
|
|
)
|
|
|
|
response = client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[
|
|
{'role': 'system', 'content': 'You are a helpful coding assistant. Generate clean, working code.'},
|
|
{'role': 'user', 'content': description}
|
|
],
|
|
max_tokens=2000
|
|
)
|
|
|
|
result = response.choices[0].message.content
|
|
|
|
# Send result back
|
|
result_msg = {
|
|
'type': 'task_result',
|
|
'task_id': task_id,
|
|
'agent_id': self.agent_id,
|
|
'status': 'completed',
|
|
'result': result,
|
|
'timestamp': time.time()
|
|
}
|
|
await self.ws.send(json.dumps(result_msg))
|
|
logger.info(f'Task {task_id} completed')
|
|
|
|
except Exception as e:
|
|
logger.error(f'Task execution failed: {e}')
|
|
error_msg = {
|
|
'type': 'task_result',
|
|
'task_id': task_id,
|
|
'agent_id': self.agent_id,
|
|
'status': 'failed',
|
|
'error': str(e),
|
|
'timestamp': time.time()
|
|
}
|
|
await self.ws.send(json.dumps(error_msg))
|
|
|
|
async def listen(self):
|
|
while self.running:
|
|
try:
|
|
message = await self.ws.recv()
|
|
data = json.loads(message)
|
|
msg_type = data.get('type')
|
|
|
|
logger.info(f'Received message: {msg_type}')
|
|
|
|
if msg_type == 'task_assignment':
|
|
await self.execute_task(data)
|
|
elif msg_type == 'ping':
|
|
pong_msg = {'type': 'pong', 'agent_id': self.agent_id}
|
|
await self.ws.send(json.dumps(pong_msg))
|
|
|
|
except ConnectionClosed:
|
|
logger.warning('Connection closed, reconnecting...')
|
|
await asyncio.sleep(5)
|
|
await self.connect()
|
|
except Exception as e:
|
|
logger.error(f'Listen error: {e}')
|
|
await asyncio.sleep(1)
|
|
|
|
async def run(self):
|
|
await self.connect()
|
|
|
|
# Start heartbeat task
|
|
heartbeat_task = asyncio.create_task(self.send_heartbeat())
|
|
|
|
# Start listening
|
|
await self.listen()
|
|
|
|
heartbeat_task.cancel()
|
|
|
|
async def main():
|
|
orchestrator_url = os.getenv('ORCHESTRATOR_URL', 'ws://orchestrator-service:8000/ws')
|
|
agent_id = os.getenv('AGENT_ID', f'agent-{uuid.uuid4().hex[:8]}')
|
|
|
|
logger.info(f'Starting agent {agent_id}')
|
|
logger.info(f'Orchestrator URL: {orchestrator_url}')
|
|
|
|
agent = SimpleAgent(orchestrator_url, agent_id)
|
|
await agent.run()
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main())
|
|
"
|
|
env:
|
|
- name: AGENT_ID
|
|
valueFrom:
|
|
fieldRef:
|
|
fieldPath: metadata.name
|
|
- name: ORCHESTRATOR_URL
|
|
value: "ws://orchestrator-service:8000/ws/$(AGENT_ID)"
|
|
- name: OPENAI_API_KEY
|
|
valueFrom:
|
|
secretKeyRef:
|
|
name: openai-api-key
|
|
key: api-key
|
|
- name: OPENAI_API_BASE
|
|
valueFrom:
|
|
secretKeyRef:
|
|
name: openai-api-key
|
|
key: api-base
|
|
optional: true
|
|
- name: OPENAI_MODEL
|
|
valueFrom:
|
|
secretKeyRef:
|
|
name: openai-api-key
|
|
key: model
|
|
optional: true
|
|
- name: PYTHONUNBUFFERED
|
|
value: "1"
|
|
resources:
|
|
requests:
|
|
cpu: 500m
|
|
memory: 512Mi
|
|
limits:
|
|
cpu: 1000m
|
|
memory: 1Gi
|
|
volumeMounts:
|
|
- name: agent-source
|
|
mountPath: /app
|
|
- name: workspace
|
|
mountPath: /workspace
|
|
volumes:
|
|
- name: agent-source
|
|
configMap:
|
|
name: agent-source
|
|
- name: workspace
|
|
emptyDir: {}
|
|
restartPolicy: Always
|