Files
hive/backend/app/api/ai_models.py
anthonyrawlins 268214d971 Major WHOOSH system refactoring and feature enhancements
- Migrated from HIVE branding to WHOOSH across all components
- Enhanced backend API with new services: AI models, BZZZ integration, templates, members
- Added comprehensive testing suite with security, performance, and integration tests
- Improved frontend with new components for project setup, AI models, and team management
- Updated MCP server implementation with WHOOSH-specific tools and resources
- Enhanced deployment configurations with production-ready Docker setups
- Added comprehensive documentation and setup guides
- Implemented age encryption service and UCXL integration

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
2025-08-27 08:34:48 +10:00

350 lines
12 KiB
Python

"""
WHOOSH AI Models API - Phase 6.1
REST API endpoints for AI model management and usage
"""
from fastapi import APIRouter, HTTPException, Depends, BackgroundTasks
from typing import List, Dict, Any, Optional
from pydantic import BaseModel
import logging
from app.services.ai_model_service import ai_model_service, ModelCapability, AIModel
from app.core.auth_deps import get_current_user
from app.models.user import User
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/ai-models", tags=["AI Models"])
# Request/Response Models
class CompletionRequest(BaseModel):
prompt: str
model_name: Optional[str] = None
system_prompt: Optional[str] = None
max_tokens: int = 1000
temperature: float = 0.7
task_type: Optional[str] = None
context_requirements: int = 2048
class CompletionResponse(BaseModel):
success: bool
content: Optional[str] = None
model: str
response_time: Optional[float] = None
usage_stats: Optional[Dict[str, Any]] = None
error: Optional[str] = None
class ModelInfo(BaseModel):
name: str
node_url: str
capabilities: List[str]
context_length: int
parameter_count: str
specialization: Optional[str] = None
performance_score: float
availability: bool
usage_count: int
avg_response_time: float
class ClusterStatus(BaseModel):
total_nodes: int
healthy_nodes: int
total_models: int
models_by_capability: Dict[str, int]
cluster_load: float
model_usage_stats: Dict[str, Dict[str, Any]]
class ModelSelectionRequest(BaseModel):
task_type: str
context_requirements: int = 2048
prefer_specialized: bool = True
class CodeGenerationRequest(BaseModel):
description: str
language: str = "python"
context: Optional[str] = None
style: str = "clean" # clean, optimized, documented
max_tokens: int = 2000
class CodeReviewRequest(BaseModel):
code: str
language: str
focus_areas: List[str] = ["bugs", "performance", "security", "style"]
severity_level: str = "medium" # low, medium, high
@router.on_event("startup")
async def startup_ai_service():
"""Initialize AI model service on startup"""
try:
await ai_model_service.initialize()
logger.info("AI Model Service initialized successfully")
except Exception as e:
logger.error(f"Failed to initialize AI Model Service: {e}")
@router.on_event("shutdown")
async def shutdown_ai_service():
"""Cleanup AI model service on shutdown"""
await ai_model_service.cleanup()
@router.get("/status", response_model=ClusterStatus)
async def get_cluster_status(current_user: User = Depends(get_current_user)):
"""Get comprehensive cluster status"""
try:
status = await ai_model_service.get_cluster_status()
return ClusterStatus(**status)
except Exception as e:
logger.error(f"Error getting cluster status: {e}")
raise HTTPException(status_code=500, detail="Failed to get cluster status")
@router.get("/models", response_model=List[ModelInfo])
async def list_available_models(current_user: User = Depends(get_current_user)):
"""List all available AI models across the cluster"""
try:
models = []
for model in ai_model_service.models.values():
models.append(ModelInfo(
name=model.name,
node_url=model.node_url,
capabilities=[cap.value for cap in model.capabilities],
context_length=model.context_length,
parameter_count=model.parameter_count,
specialization=model.specialization,
performance_score=model.performance_score,
availability=model.availability,
usage_count=model.usage_count,
avg_response_time=model.avg_response_time
))
return sorted(models, key=lambda x: x.name)
except Exception as e:
logger.error(f"Error listing models: {e}")
raise HTTPException(status_code=500, detail="Failed to list models")
@router.post("/select-model", response_model=ModelInfo)
async def select_best_model(
request: ModelSelectionRequest,
current_user: User = Depends(get_current_user)
):
"""Select the best model for a specific task"""
try:
# Convert task_type string to enum
try:
task_capability = ModelCapability(request.task_type)
except ValueError:
raise HTTPException(
status_code=400,
detail=f"Invalid task type: {request.task_type}"
)
model = await ai_model_service.get_best_model_for_task(
task_type=task_capability,
context_requirements=request.context_requirements,
prefer_specialized=request.prefer_specialized
)
if not model:
raise HTTPException(
status_code=404,
detail="No suitable model found for the specified task"
)
return ModelInfo(
name=model.name,
node_url=model.node_url,
capabilities=[cap.value for cap in model.capabilities],
context_length=model.context_length,
parameter_count=model.parameter_count,
specialization=model.specialization,
performance_score=model.performance_score,
availability=model.availability,
usage_count=model.usage_count,
avg_response_time=model.avg_response_time
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error selecting model: {e}")
raise HTTPException(status_code=500, detail="Failed to select model")
@router.post("/generate", response_model=CompletionResponse)
async def generate_completion(
request: CompletionRequest,
current_user: User = Depends(get_current_user)
):
"""Generate completion using AI model"""
try:
model_name = request.model_name
# Auto-select model if not specified
if not model_name and request.task_type:
try:
task_capability = ModelCapability(request.task_type)
best_model = await ai_model_service.get_best_model_for_task(
task_type=task_capability,
context_requirements=request.context_requirements
)
if best_model:
model_name = best_model.name
except ValueError:
pass
if not model_name:
# Default to first available model
available_models = [m for m in ai_model_service.models.values() if m.availability]
if not available_models:
raise HTTPException(status_code=503, detail="No models available")
model_name = available_models[0].name
result = await ai_model_service.generate_completion(
model_name=model_name,
prompt=request.prompt,
system_prompt=request.system_prompt,
max_tokens=request.max_tokens,
temperature=request.temperature
)
return CompletionResponse(**result)
except Exception as e:
logger.error(f"Error generating completion: {e}")
raise HTTPException(status_code=500, detail=str(e))
@router.post("/code/generate", response_model=CompletionResponse)
async def generate_code(
request: CodeGenerationRequest,
current_user: User = Depends(get_current_user)
):
"""Generate code using AI models optimized for coding"""
try:
# Select best coding model
coding_model = await ai_model_service.get_best_model_for_task(
task_type=ModelCapability.CODE_GENERATION,
context_requirements=max(2048, len(request.description) * 4)
)
if not coding_model:
raise HTTPException(status_code=503, detail="No coding models available")
# Craft specialized prompt for code generation
system_prompt = f"""You are an expert {request.language} programmer. Generate clean, well-documented, and efficient code.
Style preferences: {request.style}
Language: {request.language}
Focus on: best practices, readability, and maintainability."""
prompt = f"""Generate {request.language} code for the following requirement:
Description: {request.description}
{f"Context: {request.context}" if request.context else ""}
Please provide:
1. Clean, well-structured code
2. Appropriate comments and documentation
3. Error handling where relevant
4. Following {request.language} best practices
Code:"""
result = await ai_model_service.generate_completion(
model_name=coding_model.name,
prompt=prompt,
system_prompt=system_prompt,
max_tokens=request.max_tokens,
temperature=0.3 # Lower temperature for more deterministic code
)
return CompletionResponse(**result)
except Exception as e:
logger.error(f"Error generating code: {e}")
raise HTTPException(status_code=500, detail=str(e))
@router.post("/code/review", response_model=CompletionResponse)
async def review_code(
request: CodeReviewRequest,
current_user: User = Depends(get_current_user)
):
"""Review code using AI models optimized for code analysis"""
try:
# Select best code review model
review_model = await ai_model_service.get_best_model_for_task(
task_type=ModelCapability.CODE_REVIEW,
context_requirements=max(4096, len(request.code) * 2)
)
if not review_model:
raise HTTPException(status_code=503, detail="No code review models available")
# Craft specialized prompt for code review
system_prompt = f"""You are an expert code reviewer specializing in {request.language}.
Provide constructive, actionable feedback focusing on: {', '.join(request.focus_areas)}.
Severity level: {request.severity_level}
Be specific about line numbers and provide concrete suggestions for improvement."""
focus_description = {
"bugs": "potential bugs and logic errors",
"performance": "performance optimizations and efficiency",
"security": "security vulnerabilities and best practices",
"style": "code style, formatting, and conventions",
"maintainability": "code maintainability and readability",
"testing": "test coverage and testability"
}
focus_details = [focus_description.get(area, area) for area in request.focus_areas]
prompt = f"""Please review this {request.language} code focusing on: {', '.join(focus_details)}
Code to review:
```{request.language}
{request.code}
```
Provide a detailed review including:
1. Overall assessment
2. Specific issues found (with line references if applicable)
3. Recommendations for improvement
4. Best practices that could be applied
5. Security considerations (if applicable)
Review:"""
result = await ai_model_service.generate_completion(
model_name=review_model.name,
prompt=prompt,
system_prompt=system_prompt,
max_tokens=2000,
temperature=0.5
)
return CompletionResponse(**result)
except Exception as e:
logger.error(f"Error reviewing code: {e}")
raise HTTPException(status_code=500, detail=str(e))
@router.post("/refresh-models")
async def refresh_model_discovery(
background_tasks: BackgroundTasks,
current_user: User = Depends(get_current_user)
):
"""Refresh model discovery across the cluster"""
try:
background_tasks.add_task(ai_model_service.discover_cluster_models)
return {"message": "Model discovery refresh initiated"}
except Exception as e:
logger.error(f"Error refreshing models: {e}")
raise HTTPException(status_code=500, detail="Failed to refresh models")
@router.get("/capabilities")
async def list_model_capabilities():
"""List all available model capabilities"""
return {
"capabilities": [
{
"name": cap.value,
"description": cap.value.replace("_", " ").title()
}
for cap in ModelCapability
]
}