mirror of
https://github.com/mblanke/ThreatHunt.git
synced 2026-03-01 14:00:20 -05:00
75 lines
1.8 KiB
Python
75 lines
1.8 KiB
Python
from pydantic import BaseModel
|
|
from typing import Optional, List, Dict, Any
|
|
from datetime import datetime
|
|
|
|
|
|
class LLMRequestSchema(BaseModel):
|
|
"""Schema for LLM processing request"""
|
|
prompt: str
|
|
task_hints: Optional[List[str]] = []
|
|
requires_parallel: bool = False
|
|
requires_chaining: bool = False
|
|
batch_size: int = 1
|
|
operations: Optional[List[str]] = []
|
|
parameters: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
class RoutingDecision(BaseModel):
|
|
"""Schema for routing decision"""
|
|
task_type: str
|
|
model: str
|
|
endpoint: str
|
|
priority: int
|
|
description: str
|
|
requires_parallel: bool
|
|
requires_chaining: bool
|
|
|
|
|
|
class NodeInfo(BaseModel):
|
|
"""Schema for GPU node information"""
|
|
node_id: str
|
|
hostname: str
|
|
vram_total_gb: int
|
|
vram_used_gb: int
|
|
vram_available_gb: int
|
|
compute_utilization: float
|
|
status: str
|
|
models_loaded: List[str]
|
|
|
|
|
|
class SchedulingDecision(BaseModel):
|
|
"""Schema for job scheduling decision"""
|
|
job_id: str
|
|
execution_mode: str
|
|
nodes: Optional[List[Dict[str, str]]] = None
|
|
node: Optional[Dict[str, str]] = None
|
|
status: Optional[str] = None
|
|
queue_position: Optional[int] = None
|
|
|
|
|
|
class LLMResponseSchema(BaseModel):
|
|
"""Schema for LLM response"""
|
|
job_id: str
|
|
status: str
|
|
routing: Optional[RoutingDecision] = None
|
|
scheduling: Optional[SchedulingDecision] = None
|
|
result: Any
|
|
execution_mode: str
|
|
|
|
|
|
class ModelInfo(BaseModel):
|
|
"""Schema for model information"""
|
|
model_name: str
|
|
node_id: str
|
|
endpoint_url: str
|
|
is_available: bool
|
|
|
|
|
|
class MergedResult(BaseModel):
|
|
"""Schema for merged result"""
|
|
strategy: str
|
|
result: Any
|
|
confidence: Optional[float] = None
|
|
num_models: Optional[int] = None
|
|
all_results: Optional[List[Dict[str, Any]]] = None
|