Add comprehensive development roadmap via GitHub Issues
Created 10 detailed GitHub issues covering: - Project activation and management UI (#1-2) - Worker node coordination and visualization (#3-4) - Automated GitHub repository scanning (#5) - Intelligent model-to-issue matching (#6) - Multi-model task execution system (#7) - N8N workflow integration (#8) - Hive-Bzzz P2P bridge (#9) - Peer assistance protocol (#10) Each issue includes detailed specifications, acceptance criteria, technical implementation notes, and dependency mapping. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
anthonyrawlins
312
backend/app/api/auto_agents.py
Normal file
312
backend/app/api/auto_agents.py
Normal file
@@ -0,0 +1,312 @@
|
||||
"""
|
||||
Auto-Discovery Agent Management Endpoints
|
||||
|
||||
This module provides API endpoints for automatic agent discovery and registration
|
||||
with dynamic capability detection based on installed models.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request, Depends, status
|
||||
from typing import List, Dict, Any, Optional
|
||||
from pydantic import BaseModel, Field
|
||||
from ..services.capability_detector import CapabilityDetector, detect_capabilities
|
||||
from ..core.unified_coordinator import Agent, AgentType
|
||||
from ..models.responses import (
|
||||
AgentListResponse,
|
||||
AgentRegistrationResponse,
|
||||
ErrorResponse,
|
||||
AgentModel,
|
||||
BaseResponse
|
||||
)
|
||||
from ..core.auth_deps import get_current_user_context
|
||||
from app.core.database import SessionLocal
|
||||
from app.models.agent import Agent as ORMAgent
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class AutoDiscoveryRequest(BaseModel):
|
||||
"""Request model for auto-discovery of agents"""
|
||||
endpoints: List[str] = Field(..., description="List of Ollama endpoints to scan")
|
||||
force_refresh: bool = Field(False, description="Force refresh of existing agents")
|
||||
|
||||
|
||||
class CapabilityReport(BaseModel):
|
||||
"""Model capability detection report"""
|
||||
endpoint: str
|
||||
models: List[str]
|
||||
model_count: int
|
||||
specialty: str
|
||||
capabilities: List[str]
|
||||
status: str
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
class AutoDiscoveryResponse(BaseResponse):
|
||||
"""Response for auto-discovery operations"""
|
||||
discovered_agents: List[CapabilityReport]
|
||||
registered_agents: List[str]
|
||||
failed_agents: List[str]
|
||||
total_discovered: int
|
||||
total_registered: int
|
||||
|
||||
|
||||
@router.post(
|
||||
"/auto-discovery",
|
||||
response_model=AutoDiscoveryResponse,
|
||||
status_code=status.HTTP_200_OK,
|
||||
summary="Auto-discover and register agents",
|
||||
description="""
|
||||
Automatically discover Ollama agents across the cluster and register them
|
||||
with dynamically detected capabilities based on installed models.
|
||||
|
||||
This endpoint:
|
||||
1. Scans provided endpoints for available models
|
||||
2. Analyzes model capabilities to determine agent specialization
|
||||
3. Registers agents with detected specializations
|
||||
4. Returns comprehensive discovery and registration report
|
||||
|
||||
**Dynamic Specializations:**
|
||||
- `advanced_coding`: Models like starcoder2, deepseek-coder-v2, devstral
|
||||
- `reasoning_analysis`: Models like phi4-reasoning, granite3-dense
|
||||
- `code_review_docs`: Models like codellama, qwen2.5-coder
|
||||
- `multimodal`: Models like llava with visual capabilities
|
||||
- `general_ai`: General purpose models and fallback category
|
||||
""",
|
||||
responses={
|
||||
200: {"description": "Auto-discovery completed successfully"},
|
||||
400: {"model": ErrorResponse, "description": "Invalid request parameters"},
|
||||
500: {"model": ErrorResponse, "description": "Discovery process failed"}
|
||||
}
|
||||
)
|
||||
async def auto_discover_agents(
|
||||
discovery_request: AutoDiscoveryRequest,
|
||||
request: Request,
|
||||
current_user: Dict[str, Any] = Depends(get_current_user_context)
|
||||
) -> AutoDiscoveryResponse:
|
||||
"""
|
||||
Auto-discover and register agents with dynamic capability detection.
|
||||
|
||||
Args:
|
||||
discovery_request: Discovery configuration and endpoints
|
||||
request: FastAPI request object
|
||||
current_user: Current authenticated user context
|
||||
|
||||
Returns:
|
||||
AutoDiscoveryResponse: Discovery results and registration status
|
||||
"""
|
||||
# Access coordinator
|
||||
hive_coordinator = getattr(request.app.state, 'hive_coordinator', None)
|
||||
if not hive_coordinator:
|
||||
from ..main import unified_coordinator
|
||||
hive_coordinator = unified_coordinator
|
||||
|
||||
if not hive_coordinator:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="Coordinator service unavailable"
|
||||
)
|
||||
|
||||
detector = CapabilityDetector()
|
||||
discovered_agents = []
|
||||
registered_agents = []
|
||||
failed_agents = []
|
||||
|
||||
try:
|
||||
# Scan all endpoints for capabilities
|
||||
capabilities_results = await detector.scan_cluster_capabilities(discovery_request.endpoints)
|
||||
|
||||
for endpoint, data in capabilities_results.items():
|
||||
# Create capability report
|
||||
report = CapabilityReport(
|
||||
endpoint=endpoint,
|
||||
models=data.get('models', []),
|
||||
model_count=data.get('model_count', 0),
|
||||
specialty=data.get('specialty', 'general_ai'),
|
||||
capabilities=data.get('capabilities', []),
|
||||
status=data.get('status', 'error'),
|
||||
error=data.get('error')
|
||||
)
|
||||
discovered_agents.append(report)
|
||||
|
||||
# Skip registration if offline or error
|
||||
if data['status'] != 'online' or not data['models']:
|
||||
failed_agents.append(endpoint)
|
||||
continue
|
||||
|
||||
# Generate agent ID from endpoint
|
||||
agent_id = endpoint.replace(':', '-').replace('.', '-')
|
||||
if agent_id.startswith('192-168-1-'):
|
||||
# Use hostname mapping for known cluster nodes
|
||||
hostname_map = {
|
||||
'192-168-1-27': 'walnut',
|
||||
'192-168-1-113': 'ironwood',
|
||||
'192-168-1-72': 'acacia',
|
||||
'192-168-1-132': 'rosewood',
|
||||
'192-168-1-106': 'forsteinet'
|
||||
}
|
||||
agent_id = hostname_map.get(agent_id.split('-11434')[0], agent_id)
|
||||
|
||||
# Select best model for the agent (prefer larger, more capable models)
|
||||
best_model = select_best_model(data['models'])
|
||||
|
||||
try:
|
||||
# Check if agent already exists
|
||||
with SessionLocal() as db:
|
||||
existing_agent = db.query(ORMAgent).filter(ORMAgent.id == agent_id).first()
|
||||
if existing_agent and not discovery_request.force_refresh:
|
||||
registered_agents.append(f"{agent_id} (already exists)")
|
||||
continue
|
||||
elif existing_agent and discovery_request.force_refresh:
|
||||
# Update existing agent
|
||||
existing_agent.specialty = data['specialty']
|
||||
existing_agent.model = best_model
|
||||
db.commit()
|
||||
registered_agents.append(f"{agent_id} (updated)")
|
||||
continue
|
||||
|
||||
# Map specialty to AgentType enum
|
||||
specialty_mapping = {
|
||||
'advanced_coding': AgentType.KERNEL_DEV,
|
||||
'reasoning_analysis': AgentType.REASONING,
|
||||
'code_review_docs': AgentType.DOCS_WRITER,
|
||||
'multimodal': AgentType.GENERAL_AI,
|
||||
'general_ai': AgentType.GENERAL_AI
|
||||
}
|
||||
agent_type = specialty_mapping.get(data['specialty'], AgentType.GENERAL_AI)
|
||||
|
||||
# Create and register agent
|
||||
agent = Agent(
|
||||
id=agent_id,
|
||||
endpoint=endpoint,
|
||||
model=best_model,
|
||||
specialty=agent_type,
|
||||
max_concurrent=2 # Default concurrent task limit
|
||||
)
|
||||
|
||||
# Add to coordinator
|
||||
hive_coordinator.add_agent(agent)
|
||||
registered_agents.append(agent_id)
|
||||
|
||||
except Exception as e:
|
||||
failed_agents.append(f"{endpoint}: {str(e)}")
|
||||
|
||||
return AutoDiscoveryResponse(
|
||||
status="success",
|
||||
message=f"Discovery completed: {len(registered_agents)} registered, {len(failed_agents)} failed",
|
||||
discovered_agents=discovered_agents,
|
||||
registered_agents=registered_agents,
|
||||
failed_agents=failed_agents,
|
||||
total_discovered=len(discovered_agents),
|
||||
total_registered=len(registered_agents)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Auto-discovery failed: {str(e)}"
|
||||
)
|
||||
finally:
|
||||
await detector.close()
|
||||
|
||||
|
||||
@router.get(
|
||||
"/cluster-capabilities",
|
||||
response_model=List[CapabilityReport],
|
||||
status_code=status.HTTP_200_OK,
|
||||
summary="Scan cluster capabilities without registration",
|
||||
description="""
|
||||
Scan the cluster for agent capabilities without registering them.
|
||||
|
||||
This endpoint provides a read-only view of what agents would be discovered
|
||||
and their detected capabilities, useful for:
|
||||
- Planning agent deployment strategies
|
||||
- Understanding cluster capacity
|
||||
- Debugging capability detection
|
||||
- Validating model installations
|
||||
""",
|
||||
responses={
|
||||
200: {"description": "Cluster capabilities scanned successfully"},
|
||||
500: {"model": ErrorResponse, "description": "Capability scan failed"}
|
||||
}
|
||||
)
|
||||
async def scan_cluster_capabilities(
|
||||
endpoints: List[str] = ["192.168.1.27:11434", "192.168.1.113:11434", "192.168.1.72:11434", "192.168.1.132:11434", "192.168.1.106:11434"],
|
||||
current_user: Dict[str, Any] = Depends(get_current_user_context)
|
||||
) -> List[CapabilityReport]:
|
||||
"""
|
||||
Scan cluster endpoints for model capabilities.
|
||||
|
||||
Args:
|
||||
endpoints: List of Ollama endpoints to scan
|
||||
current_user: Current authenticated user context
|
||||
|
||||
Returns:
|
||||
List[CapabilityReport]: Capability reports for each endpoint
|
||||
"""
|
||||
detector = CapabilityDetector()
|
||||
|
||||
try:
|
||||
capabilities_results = await detector.scan_cluster_capabilities(endpoints)
|
||||
|
||||
reports = []
|
||||
for endpoint, data in capabilities_results.items():
|
||||
report = CapabilityReport(
|
||||
endpoint=endpoint,
|
||||
models=data.get('models', []),
|
||||
model_count=data.get('model_count', 0),
|
||||
specialty=data.get('specialty', 'general_ai'),
|
||||
capabilities=data.get('capabilities', []),
|
||||
status=data.get('status', 'error'),
|
||||
error=data.get('error')
|
||||
)
|
||||
reports.append(report)
|
||||
|
||||
return reports
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Capability scan failed: {str(e)}"
|
||||
)
|
||||
finally:
|
||||
await detector.close()
|
||||
|
||||
|
||||
def select_best_model(models: List[str]) -> str:
|
||||
"""
|
||||
Select the best model from available models for agent registration.
|
||||
|
||||
Prioritizes models by capability and size:
|
||||
1. Advanced coding models (starcoder2, deepseek-coder-v2, devstral)
|
||||
2. Reasoning models (phi4, granite3-dense)
|
||||
3. Larger models over smaller ones
|
||||
4. Fallback to first available model
|
||||
"""
|
||||
if not models:
|
||||
return "unknown"
|
||||
|
||||
# Priority order for model selection
|
||||
priority_patterns = [
|
||||
"starcoder2:15b", "deepseek-coder-v2", "devstral",
|
||||
"phi4:14b", "phi4-reasoning", "qwen3:14b",
|
||||
"granite3-dense", "codellama", "qwen2.5-coder",
|
||||
"llama3.1:8b", "gemma3:12b", "mistral:7b"
|
||||
]
|
||||
|
||||
# Find highest priority model
|
||||
for pattern in priority_patterns:
|
||||
for model in models:
|
||||
if pattern in model.lower():
|
||||
return model
|
||||
|
||||
# Fallback: select largest model by parameter count
|
||||
def extract_size(model_name: str) -> int:
|
||||
"""Extract parameter count from model name"""
|
||||
import re
|
||||
size_match = re.search(r'(\d+)b', model_name.lower())
|
||||
if size_match:
|
||||
return int(size_match.group(1))
|
||||
return 0
|
||||
|
||||
largest_model = max(models, key=extract_size)
|
||||
return largest_model if extract_size(largest_model) > 0 else models[0]
|
||||
@@ -14,9 +14,16 @@ Key Features:
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, status
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.encoders import jsonable_encoder
|
||||
from typing import List, Dict, Any, Optional
|
||||
from datetime import datetime
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from ..core.auth_deps import get_current_user_context
|
||||
from ..core.unified_coordinator_refactored import UnifiedCoordinatorRefactored as UnifiedCoordinator
|
||||
from ..services.agent_service import AgentType
|
||||
from ..models.responses import (
|
||||
TaskListResponse,
|
||||
TaskCreationResponse,
|
||||
@@ -110,27 +117,54 @@ async def create_task(
|
||||
raise coordinator_unavailable_error()
|
||||
|
||||
try:
|
||||
# Convert Pydantic model to dict for coordinator
|
||||
task_dict = {
|
||||
"type": task_data.type,
|
||||
"priority": task_data.priority,
|
||||
"context": task_data.context,
|
||||
"preferred_agent": task_data.preferred_agent,
|
||||
"timeout": task_data.timeout,
|
||||
"user_id": current_user.get("user_id", "unknown")
|
||||
}
|
||||
# Convert task type string to AgentType enum
|
||||
try:
|
||||
agent_type = AgentType(task_data.type)
|
||||
except ValueError:
|
||||
raise validation_error("type", f"Invalid task type: {task_data.type}")
|
||||
|
||||
# Create task using coordinator
|
||||
task_id = await coordinator.submit_task(task_dict)
|
||||
try:
|
||||
task = coordinator.create_task(
|
||||
task_type=agent_type,
|
||||
context=task_data.context,
|
||||
priority=task_data.priority
|
||||
)
|
||||
logger.info(f"Task created successfully: {task.id}")
|
||||
except Exception as create_err:
|
||||
logger.error(f"Task creation failed: {create_err}")
|
||||
import traceback
|
||||
logger.error(f"Full traceback: {traceback.format_exc()}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Task creation failed: {str(create_err)}"
|
||||
)
|
||||
|
||||
# Get task details for response
|
||||
task_details = await coordinator.get_task_status(task_id)
|
||||
|
||||
return TaskCreationResponse(
|
||||
task_id=task_id,
|
||||
assigned_agent=task_details.get("assigned_agent") if task_details else task_data.preferred_agent,
|
||||
message=f"Task '{task_id}' created successfully with priority {task_data.priority}"
|
||||
)
|
||||
# Create simple dictionary response to avoid Pydantic datetime issues
|
||||
try:
|
||||
response_dict = {
|
||||
"status": "success",
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
"task_id": str(task.id),
|
||||
"assigned_agent": str(task.assigned_agent) if task.assigned_agent else None,
|
||||
"message": f"Task '{task.id}' created successfully with priority {task_data.priority}"
|
||||
}
|
||||
|
||||
return JSONResponse(
|
||||
status_code=201,
|
||||
content=response_dict
|
||||
)
|
||||
except Exception as response_err:
|
||||
logger.error(f"Response creation failed: {response_err}")
|
||||
# Return minimal safe response
|
||||
return JSONResponse(
|
||||
status_code=201,
|
||||
content={
|
||||
"status": "success",
|
||||
"task_id": str(task.id) if hasattr(task, 'id') else "unknown",
|
||||
"message": "Task created successfully"
|
||||
}
|
||||
)
|
||||
|
||||
except ValueError as e:
|
||||
raise validation_error("task_data", str(e))
|
||||
@@ -200,23 +234,23 @@ async def get_task(
|
||||
raise coordinator_unavailable_error()
|
||||
|
||||
try:
|
||||
task = await coordinator.get_task_status(task_id)
|
||||
task = coordinator.get_task_status(task_id)
|
||||
if not task:
|
||||
raise task_not_found_error(task_id)
|
||||
|
||||
# Convert coordinator task to response model
|
||||
return TaskModel(
|
||||
id=task.get("id", task_id),
|
||||
type=task.get("type", "unknown"),
|
||||
priority=task.get("priority", 3),
|
||||
status=task.get("status", "unknown"),
|
||||
context=task.get("context", {}),
|
||||
assigned_agent=task.get("assigned_agent"),
|
||||
result=task.get("result"),
|
||||
created_at=task.get("created_at"),
|
||||
started_at=task.get("started_at"),
|
||||
completed_at=task.get("completed_at"),
|
||||
error_message=task.get("error_message")
|
||||
id=task.id,
|
||||
type=task.type.value if hasattr(task.type, 'value') else str(task.type),
|
||||
priority=task.priority,
|
||||
status=task.status.value if hasattr(task.status, 'value') else str(task.status),
|
||||
context=task.context or {},
|
||||
assigned_agent=task.assigned_agent,
|
||||
result=task.result,
|
||||
created_at=task.created_at,
|
||||
started_at=getattr(task, 'started_at', None),
|
||||
completed_at=task.completed_at,
|
||||
error_message=getattr(task, 'error_message', None)
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
|
||||
BIN
backend/app/core/__pycache__/init_db.cpython-310.pyc
Normal file
BIN
backend/app/core/__pycache__/init_db.cpython-310.pyc
Normal file
Binary file not shown.
@@ -151,7 +151,21 @@ class UnifiedCoordinatorRefactored:
|
||||
|
||||
# Persist to database
|
||||
try:
|
||||
self.task_service.create_task(task)
|
||||
# Convert Task object to dictionary for database storage
|
||||
task_dict = {
|
||||
'id': task.id,
|
||||
'title': f"Task {task.type.value}",
|
||||
'description': f"Priority {task.priority} task",
|
||||
'priority': task.priority,
|
||||
'status': task.status.value,
|
||||
'assigned_agent': task.assigned_agent,
|
||||
'context': task.context,
|
||||
'payload': task.payload,
|
||||
'type': task.type.value,
|
||||
'created_at': task.created_at,
|
||||
'completed_at': task.completed_at
|
||||
}
|
||||
self.task_service.create_task(task_dict)
|
||||
logger.info(f"💾 Task {task_id} persisted to database")
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Failed to persist task {task_id} to database: {e}")
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from fastapi import FastAPI, Depends, HTTPException
|
||||
from fastapi import FastAPI, Depends, HTTPException, status
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.encoders import jsonable_encoder
|
||||
from contextlib import asynccontextmanager
|
||||
import json
|
||||
import asyncio
|
||||
@@ -49,6 +50,9 @@ async def lifespan(app: FastAPI):
|
||||
print("🤖 Initializing Unified Coordinator...")
|
||||
await unified_coordinator.start()
|
||||
|
||||
# Store coordinator in app state for endpoint access
|
||||
app.state.hive_coordinator = unified_coordinator
|
||||
app.state.unified_coordinator = unified_coordinator
|
||||
|
||||
startup_success = True
|
||||
print("✅ Hive Orchestrator with Unified Coordinator started successfully!")
|
||||
@@ -517,10 +521,7 @@ async def root():
|
||||
"timestamp": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check_internal():
|
||||
"""Internal health check endpoint for Docker and monitoring"""
|
||||
return {"status": "healthy", "timestamp": datetime.now().isoformat()}
|
||||
# Removed duplicate /health endpoint - using the enhanced one above
|
||||
|
||||
@app.get("/api/health")
|
||||
async def health_check():
|
||||
|
||||
@@ -48,6 +48,11 @@ class BaseResponse(BaseModel):
|
||||
status: StatusEnum = Field(..., description="Response status indicator")
|
||||
timestamp: datetime = Field(default_factory=datetime.utcnow, description="Response timestamp")
|
||||
message: Optional[str] = Field(None, description="Human-readable message")
|
||||
|
||||
class Config:
|
||||
json_encoders = {
|
||||
datetime: lambda v: v.isoformat() if v else None
|
||||
}
|
||||
|
||||
|
||||
class ErrorResponse(BaseResponse):
|
||||
@@ -173,12 +178,15 @@ class TaskModel(BaseModel):
|
||||
context: Dict[str, Any] = Field(..., description="Task context and parameters")
|
||||
assigned_agent: Optional[str] = Field(None, description="ID of assigned agent", example="walnut-codellama")
|
||||
result: Optional[Dict[str, Any]] = Field(None, description="Task execution results")
|
||||
created_at: datetime = Field(..., description="Task creation timestamp")
|
||||
created_at: Optional[datetime] = Field(None, description="Task creation timestamp")
|
||||
started_at: Optional[datetime] = Field(None, description="Task start timestamp")
|
||||
completed_at: Optional[datetime] = Field(None, description="Task completion timestamp")
|
||||
error_message: Optional[str] = Field(None, description="Error message if task failed")
|
||||
|
||||
class Config:
|
||||
json_encoders = {
|
||||
datetime: lambda v: v.isoformat() if v else None
|
||||
}
|
||||
schema_extra = {
|
||||
"example": {
|
||||
"id": "task-12345",
|
||||
@@ -228,7 +236,7 @@ class TaskCreationResponse(BaseResponse):
|
||||
status: StatusEnum = Field(StatusEnum.SUCCESS)
|
||||
task_id: str = Field(..., description="ID of the created task", example="task-12345")
|
||||
assigned_agent: Optional[str] = Field(None, description="ID of assigned agent", example="walnut-codellama")
|
||||
estimated_completion: Optional[datetime] = Field(None, description="Estimated completion time")
|
||||
estimated_completion: Optional[str] = Field(None, description="Estimated completion time (ISO format)")
|
||||
|
||||
class Config:
|
||||
schema_extra = {
|
||||
|
||||
BIN
backend/app/services/__pycache__/__init__.cpython-310.pyc
Normal file
BIN
backend/app/services/__pycache__/__init__.cpython-310.pyc
Normal file
Binary file not shown.
BIN
backend/app/services/__pycache__/task_service.cpython-310.pyc
Normal file
BIN
backend/app/services/__pycache__/task_service.cpython-310.pyc
Normal file
Binary file not shown.
259
backend/app/services/capability_detector.py
Normal file
259
backend/app/services/capability_detector.py
Normal file
@@ -0,0 +1,259 @@
|
||||
"""
|
||||
Capability Detection Service for Hive Agents
|
||||
|
||||
This service automatically detects agent capabilities and specializations based on
|
||||
the models installed on each Ollama endpoint. It replaces hardcoded specializations
|
||||
with dynamic detection based on actual model capabilities.
|
||||
"""
|
||||
|
||||
import httpx
|
||||
import asyncio
|
||||
from typing import Dict, List, Set, Optional, Tuple
|
||||
from enum import Enum
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ModelCapability(str, Enum):
|
||||
"""Model capability categories based on model characteristics"""
|
||||
CODE_GENERATION = "code_generation"
|
||||
CODE_REVIEW = "code_review"
|
||||
REASONING = "reasoning"
|
||||
DOCUMENTATION = "documentation"
|
||||
TESTING = "testing"
|
||||
VISUAL_ANALYSIS = "visual_analysis"
|
||||
GENERAL_AI = "general_ai"
|
||||
KERNEL_DEV = "kernel_dev"
|
||||
PYTORCH_DEV = "pytorch_dev"
|
||||
PROFILER = "profiler"
|
||||
|
||||
|
||||
class AgentSpecialty(str, Enum):
|
||||
"""Dynamic agent specializations based on model capabilities"""
|
||||
ADVANCED_CODING = "advanced_coding" # starcoder2, deepseek-coder-v2, devstral
|
||||
REASONING_ANALYSIS = "reasoning_analysis" # phi4-reasoning, granite3-dense
|
||||
CODE_REVIEW_DOCS = "code_review_docs" # codellama, qwen2.5-coder
|
||||
GENERAL_AI = "general_ai" # llama3, gemma, mistral
|
||||
MULTIMODAL = "multimodal" # llava, vision models
|
||||
LIGHTWEIGHT = "lightweight" # small models < 8B
|
||||
|
||||
|
||||
# Model capability mapping based on model names and characteristics
|
||||
MODEL_CAPABILITIES = {
|
||||
# Advanced coding models
|
||||
"starcoder2": [ModelCapability.CODE_GENERATION, ModelCapability.KERNEL_DEV],
|
||||
"deepseek-coder": [ModelCapability.CODE_GENERATION, ModelCapability.CODE_REVIEW],
|
||||
"devstral": [ModelCapability.CODE_GENERATION, ModelCapability.PROFILER],
|
||||
"codellama": [ModelCapability.CODE_GENERATION, ModelCapability.CODE_REVIEW],
|
||||
"qwen2.5-coder": [ModelCapability.CODE_GENERATION, ModelCapability.CODE_REVIEW],
|
||||
"qwen3": [ModelCapability.CODE_GENERATION, ModelCapability.REASONING],
|
||||
|
||||
# Reasoning and analysis models
|
||||
"phi4-reasoning": [ModelCapability.REASONING, ModelCapability.PROFILER],
|
||||
"phi4": [ModelCapability.REASONING, ModelCapability.GENERAL_AI],
|
||||
"granite3-dense": [ModelCapability.REASONING, ModelCapability.PYTORCH_DEV],
|
||||
"deepseek-r1": [ModelCapability.REASONING, ModelCapability.CODE_REVIEW],
|
||||
|
||||
# General purpose models
|
||||
"llama3": [ModelCapability.GENERAL_AI, ModelCapability.DOCUMENTATION],
|
||||
"gemma": [ModelCapability.GENERAL_AI, ModelCapability.TESTING],
|
||||
"mistral": [ModelCapability.GENERAL_AI, ModelCapability.DOCUMENTATION],
|
||||
"dolphin": [ModelCapability.GENERAL_AI, ModelCapability.REASONING],
|
||||
|
||||
# Multimodal models
|
||||
"llava": [ModelCapability.VISUAL_ANALYSIS, ModelCapability.DOCUMENTATION],
|
||||
|
||||
# Tool use models
|
||||
"llama3-groq-tool-use": [ModelCapability.CODE_GENERATION, ModelCapability.TESTING],
|
||||
}
|
||||
|
||||
# Specialization determination based on capabilities
|
||||
SPECIALTY_MAPPING = {
|
||||
frozenset([ModelCapability.CODE_GENERATION, ModelCapability.KERNEL_DEV]): AgentSpecialty.ADVANCED_CODING,
|
||||
frozenset([ModelCapability.CODE_GENERATION, ModelCapability.PROFILER]): AgentSpecialty.ADVANCED_CODING,
|
||||
frozenset([ModelCapability.REASONING, ModelCapability.PROFILER]): AgentSpecialty.REASONING_ANALYSIS,
|
||||
frozenset([ModelCapability.REASONING, ModelCapability.PYTORCH_DEV]): AgentSpecialty.REASONING_ANALYSIS,
|
||||
frozenset([ModelCapability.CODE_REVIEW, ModelCapability.DOCUMENTATION]): AgentSpecialty.CODE_REVIEW_DOCS,
|
||||
frozenset([ModelCapability.VISUAL_ANALYSIS]): AgentSpecialty.MULTIMODAL,
|
||||
}
|
||||
|
||||
|
||||
class CapabilityDetector:
|
||||
"""Detects agent capabilities by analyzing available models"""
|
||||
|
||||
def __init__(self, timeout: int = 10):
|
||||
self.timeout = timeout
|
||||
self.client = httpx.AsyncClient(timeout=timeout)
|
||||
|
||||
async def get_available_models(self, endpoint: str) -> List[Dict]:
|
||||
"""Get list of available models from Ollama endpoint"""
|
||||
try:
|
||||
# Handle endpoints with or without protocol
|
||||
if not endpoint.startswith(('http://', 'https://')):
|
||||
endpoint = f"http://{endpoint}"
|
||||
|
||||
# Ensure endpoint has port if not specified
|
||||
if ':' not in endpoint.split('//')[-1]:
|
||||
endpoint = f"{endpoint}:11434"
|
||||
|
||||
response = await self.client.get(f"{endpoint}/api/tags")
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data.get('models', [])
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get models from {endpoint}: {e}")
|
||||
return []
|
||||
|
||||
def analyze_model_capabilities(self, model_name: str) -> List[ModelCapability]:
|
||||
"""Analyze a single model to determine its capabilities"""
|
||||
capabilities = []
|
||||
|
||||
# Normalize model name for matching
|
||||
normalized_name = model_name.lower().split(':')[0] # Remove version tags
|
||||
|
||||
# Check for exact matches first
|
||||
for pattern, caps in MODEL_CAPABILITIES.items():
|
||||
if pattern in normalized_name:
|
||||
capabilities.extend(caps)
|
||||
break
|
||||
|
||||
# If no specific match, determine by model size and type
|
||||
if not capabilities:
|
||||
if any(size in normalized_name for size in ['3b', '7b']):
|
||||
capabilities.append(ModelCapability.LIGHTWEIGHT)
|
||||
capabilities.append(ModelCapability.GENERAL_AI)
|
||||
|
||||
return list(set(capabilities)) # Remove duplicates
|
||||
|
||||
def determine_agent_specialty(self, all_capabilities: List[ModelCapability]) -> AgentSpecialty:
|
||||
"""Determine agent specialty based on combined model capabilities"""
|
||||
capability_set = frozenset(all_capabilities)
|
||||
|
||||
# Check for exact specialty matches
|
||||
for caps, specialty in SPECIALTY_MAPPING.items():
|
||||
if caps.issubset(capability_set):
|
||||
return specialty
|
||||
|
||||
# Fallback logic based on dominant capabilities
|
||||
if ModelCapability.CODE_GENERATION in all_capabilities:
|
||||
if ModelCapability.REASONING in all_capabilities:
|
||||
return AgentSpecialty.ADVANCED_CODING
|
||||
elif ModelCapability.CODE_REVIEW in all_capabilities:
|
||||
return AgentSpecialty.CODE_REVIEW_DOCS
|
||||
else:
|
||||
return AgentSpecialty.ADVANCED_CODING
|
||||
|
||||
elif ModelCapability.REASONING in all_capabilities:
|
||||
return AgentSpecialty.REASONING_ANALYSIS
|
||||
|
||||
elif ModelCapability.VISUAL_ANALYSIS in all_capabilities:
|
||||
return AgentSpecialty.MULTIMODAL
|
||||
|
||||
else:
|
||||
return AgentSpecialty.GENERAL_AI
|
||||
|
||||
async def detect_agent_capabilities(self, endpoint: str) -> Tuple[List[str], AgentSpecialty, List[ModelCapability]]:
|
||||
"""
|
||||
Detect agent capabilities and determine specialty
|
||||
|
||||
Returns:
|
||||
Tuple of (model_names, specialty, capabilities)
|
||||
"""
|
||||
models = await self.get_available_models(endpoint)
|
||||
|
||||
if not models:
|
||||
return [], AgentSpecialty.GENERAL_AI, [ModelCapability.GENERAL_AI]
|
||||
|
||||
model_names = [model['name'] for model in models]
|
||||
all_capabilities = []
|
||||
|
||||
# Analyze each model
|
||||
for model in models:
|
||||
model_caps = self.analyze_model_capabilities(model['name'])
|
||||
all_capabilities.extend(model_caps)
|
||||
|
||||
# Remove duplicates and determine specialty
|
||||
unique_capabilities = list(set(all_capabilities))
|
||||
specialty = self.determine_agent_specialty(unique_capabilities)
|
||||
|
||||
return model_names, specialty, unique_capabilities
|
||||
|
||||
async def scan_cluster_capabilities(self, endpoints: List[str]) -> Dict[str, Dict]:
|
||||
"""Scan multiple endpoints and return capabilities for each"""
|
||||
results = {}
|
||||
|
||||
tasks = []
|
||||
for endpoint in endpoints:
|
||||
task = self.detect_agent_capabilities(endpoint)
|
||||
tasks.append((endpoint, task))
|
||||
|
||||
# Execute all scans concurrently
|
||||
for endpoint, task in tasks:
|
||||
try:
|
||||
models, specialty, capabilities = await task
|
||||
results[endpoint] = {
|
||||
'models': models,
|
||||
'model_count': len(models),
|
||||
'specialty': specialty,
|
||||
'capabilities': capabilities,
|
||||
'status': 'online' if models else 'offline'
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to scan {endpoint}: {e}")
|
||||
results[endpoint] = {
|
||||
'models': [],
|
||||
'model_count': 0,
|
||||
'specialty': AgentSpecialty.GENERAL_AI,
|
||||
'capabilities': [],
|
||||
'status': 'error',
|
||||
'error': str(e)
|
||||
}
|
||||
|
||||
return results
|
||||
|
||||
async def close(self):
|
||||
"""Close the HTTP client"""
|
||||
await self.client.aclose()
|
||||
|
||||
|
||||
# Convenience function for quick capability detection
|
||||
async def detect_capabilities(endpoint: str) -> Dict:
|
||||
"""Quick capability detection for a single endpoint"""
|
||||
detector = CapabilityDetector()
|
||||
try:
|
||||
models, specialty, capabilities = await detector.detect_agent_capabilities(endpoint)
|
||||
return {
|
||||
'endpoint': endpoint,
|
||||
'models': models,
|
||||
'model_count': len(models),
|
||||
'specialty': specialty.value,
|
||||
'capabilities': [cap.value for cap in capabilities],
|
||||
'status': 'online' if models else 'offline'
|
||||
}
|
||||
finally:
|
||||
await detector.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test the capability detector
|
||||
async def test_detection():
|
||||
endpoints = [
|
||||
"192.168.1.27:11434", # WALNUT
|
||||
"192.168.1.113:11434", # IRONWOOD
|
||||
"192.168.1.72:11434", # ACACIA
|
||||
]
|
||||
|
||||
detector = CapabilityDetector()
|
||||
try:
|
||||
results = await detector.scan_cluster_capabilities(endpoints)
|
||||
for endpoint, data in results.items():
|
||||
print(f"\n{endpoint}:")
|
||||
print(f" Models: {data['model_count']}")
|
||||
print(f" Specialty: {data['specialty']}")
|
||||
print(f" Capabilities: {data['capabilities']}")
|
||||
print(f" Status: {data['status']}")
|
||||
finally:
|
||||
await detector.close()
|
||||
|
||||
asyncio.run(test_detection())
|
||||
Reference in New Issue
Block a user