- 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>
350 lines
12 KiB
Python
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
|
|
]
|
|
} |