134 lines
4.0 KiB
Python
134 lines
4.0 KiB
Python
"""Typed capability requests / responses.
|
|
|
|
Business services construct these Pydantic models and call methods on the
|
|
AIGateway. They never know which concrete model or provider was chosen.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional, Sequence
|
|
from uuid import UUID
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class TraceContext(BaseModel):
|
|
"""Context propagated through every call for audit / logging."""
|
|
|
|
tenant_id: UUID
|
|
project_id: Optional[UUID] = None
|
|
user_id: Optional[UUID] = None
|
|
job_id: Optional[UUID] = None
|
|
correlation_id: Optional[str] = None
|
|
|
|
|
|
class CapabilityRequest(BaseModel):
|
|
"""Base marker for every capability request supported by the gateway."""
|
|
|
|
trace: TraceContext
|
|
preferred_model: Optional[str] = Field(
|
|
None, max_length=128, description="Hint only. Gateway's routing decision is authoritative."
|
|
)
|
|
priority: str = "normal"
|
|
|
|
|
|
class CapabilityResponse(BaseModel):
|
|
"""Every response carries at least trace + usage."""
|
|
|
|
provider: str
|
|
model_used: str
|
|
input_tokens: Optional[int] = None
|
|
output_tokens: Optional[int] = None
|
|
duration_ms: Optional[int] = None
|
|
estimated_cost: Optional[float] = None
|
|
execution_id: UUID
|
|
raw: Optional[dict[str, Any]] = Field(
|
|
default=None, description="Raw provider output, preserved for debugging."
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Individual capabilities
|
|
# ============================================================================
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
role: str = Field(..., pattern="^(system|user|assistant|tool)$")
|
|
content: str
|
|
|
|
|
|
class ChatRequest(CapabilityRequest):
|
|
messages: Sequence[ChatMessage]
|
|
temperature: float = Field(0.2, ge=0.0, le=2.0)
|
|
max_tokens: Optional[int] = Field(None, ge=1)
|
|
context_summary: Optional[str] = Field(
|
|
None, description="Short plaintext summary of the Context Engine assembly."
|
|
)
|
|
|
|
|
|
class ChatResponse(CapabilityResponse):
|
|
reply: str
|
|
references: list[str] = Field(default_factory=list)
|
|
|
|
|
|
class CodeAnalysisRequest(CapabilityRequest):
|
|
file_path: Optional[str] = None
|
|
language: Optional[str] = None
|
|
code_snippet: str = Field(..., min_length=1)
|
|
analysis_types: Sequence[str] = Field(
|
|
default_factory=lambda: ["bugs", "complexity", "security", "style"],
|
|
description="Which kinds of analyses to run.",
|
|
)
|
|
|
|
|
|
class Finding(BaseModel):
|
|
severity: str = Field(..., pattern="^(info|low|medium|high|critical)$")
|
|
kind: str
|
|
message: str
|
|
line_start: Optional[int] = None
|
|
line_end: Optional[int] = None
|
|
suggestion: Optional[str] = None
|
|
|
|
|
|
class CodeAnalysisResult(CapabilityResponse):
|
|
summary: str
|
|
findings: list[Finding] = Field(default_factory=list)
|
|
|
|
|
|
class TestGenerationRequest(CapabilityRequest):
|
|
file_path: Optional[str] = None
|
|
language: str = "python"
|
|
test_framework: str = "pytest"
|
|
code_snippet: str = Field(..., min_length=1)
|
|
goal: str = "Generate unit tests covering happy path + main failure modes"
|
|
existing_tests: Optional[str] = None
|
|
|
|
|
|
class GeneratedTests(CapabilityResponse):
|
|
test_cases: list[str] = Field(default_factory=list)
|
|
fixtures_needed: list[str] = Field(default_factory=list)
|
|
coverage_notes: str = ""
|
|
|
|
|
|
class RequirementDecomposeRequest(CapabilityRequest):
|
|
requirement_text: str = Field(..., min_length=1)
|
|
context: Optional[str] = None
|
|
existing_subtasks: Sequence[str] = Field(default_factory=list)
|
|
max_subtasks: int = Field(20, ge=1, le=100)
|
|
|
|
|
|
class SubTask(BaseModel):
|
|
title: str
|
|
description: str
|
|
estimated_complexity: str = Field(..., pattern="^(xs|s|m|l|xl)$")
|
|
depends_on: list[int] = Field(
|
|
default_factory=list,
|
|
description="Index (0-based) of sibling subtasks this one depends on.",
|
|
)
|
|
|
|
|
|
class RequirementDecomposition(CapabilityResponse):
|
|
subtasks: list[SubTask] = Field(default_factory=list)
|
|
architecture_notes: list[str] = Field(default_factory=list)
|
|
open_questions: list[str] = Field(default_factory=list)
|