"""AIProvider interface + mock + registry.""" from __future__ import annotations from dataclasses import dataclass from typing import Protocol, Optional, Type from uuid import uuid4, UUID import structlog from .capabilities import ( CapabilityRequest, CapabilityResponse, ChatRequest, ChatResponse, CodeAnalysisRequest, CodeAnalysisResult, TestGenerationRequest, GeneratedTests, RequirementDecomposeRequest, RequirementDecomposition, SubTask, Finding, ) log = structlog.get_logger(__name__) @dataclass class ProviderCaps: supports_chat: bool = True supports_code_analysis: bool = True supports_test_generation: bool = True supports_decompose: bool = True max_input_chars: int = 200_000 class AIProvider(Protocol): """Low-level provider interface. The Gateway calls exactly one per request. Concrete implementations: MockProvider, OpenAICompatProvider, NVIDIADGXProvider, ... """ @property def name(self) -> str: ... @property def default_model(self) -> str: ... @property def capabilities(self) -> ProviderCaps: ... async def chat(self, req: ChatRequest) -> ChatResponse: ... async def analyze_code(self, req: CodeAnalysisRequest) -> CodeAnalysisResult: ... async def generate_tests(self, req: TestGenerationRequest) -> GeneratedTests: ... async def decompose_requirement( self, req: RequirementDecomposeRequest ) -> RequirementDecomposition: ... # ============================================================================== # MockProvider — deterministic, canned, latency-simulating responses # ============================================================================== class MockProvider: """Provider that always returns deterministic, predictable responses. Uses: * unit tests of code that depends on the AI Gateway * local dev when no API keys are set * UI demo / smoke tests """ name = "mock" default_model = "asf-mock-v1" capabilities = ProviderCaps() def __init__(self, *, response_latency_ms: int = 40) -> None: self._latency = max(0, response_latency_ms) async def _maybe_sleep(self) -> None: if self._latency <= 0: return import asyncio await asyncio.sleep(self._latency / 1000.0) async def chat(self, req: ChatRequest) -> ChatResponse: await self._maybe_sleep() # Mirror the last user message back as a demo reply — never misleadingly smart last_user = next( (m.content for m in reversed(req.messages) if m.role == "user"), "(no user message)", ) reply = ( "[ASF Mock AI] This is a deterministic response from MockProvider. " f"Your last message was: {last_user[:160]!r}. " "To enable real AI, configure a non-mock provider in the ASF AI Gateway." ) return ChatResponse( provider=self.name, model_used=self.default_model, execution_id=uuid4(), reply=reply, references=["mock://no-real-data"], input_tokens=len(last_user) // 4, output_tokens=len(reply) // 4, duration_ms=self._latency, estimated_cost=0.0, ) async def analyze_code(self, req: CodeAnalysisRequest) -> CodeAnalysisResult: await self._maybe_sleep() lines = [ln for ln in req.code_snippet.splitlines() if ln.strip()] findings: list[Finding] = [] if "TODO" in req.code_snippet: findings.append( Finding( severity="info", kind="todo", message="Mock analysis found TODO comment — remember to follow up.", ) ) if "except Exception" in req.code_snippet: findings.append( Finding( severity="medium", kind="broad_exception", message="Broad except clause detected — consider narrowing exception types.", ) ) return CodeAnalysisResult( provider=self.name, model_used=self.default_model, execution_id=uuid4(), summary=f"Mock: analyzed {len(lines)} non-empty lines; {len(findings)} informational finding(s).", findings=findings, input_tokens=max(10, len(req.code_snippet) // 4), output_tokens=100, duration_ms=self._latency, estimated_cost=0.0, ) async def generate_tests(self, req: TestGenerationRequest) -> GeneratedTests: await self._maybe_sleep() test_cases = [ "# Generated by MockProvider — placeholder skeleton\n" "def test_mock_happy_path():\n" " # TODO: replace with real assertions from a production AI provider\n" " assert True\n", "# Generated by MockProvider — placeholder skeleton\n" "def test_mock_failure_mode():\n" " # TODO: simulate and assert on failure behavior\n" " assert True\n", ] return GeneratedTests( provider=self.name, model_used=self.default_model, execution_id=uuid4(), test_cases=test_cases, fixtures_needed=[ "NOTE: Mock provider — fixtures listed are placeholders only.", ], coverage_notes="Mock response. Real AI provider will target specific branches.", input_tokens=max(10, len(req.code_snippet) // 4), output_tokens=300, duration_ms=self._latency, estimated_cost=0.0, ) async def decompose_requirement( self, req: RequirementDecomposeRequest ) -> RequirementDecomposition: await self._maybe_sleep() limit = max(1, min(5, req.max_subtasks)) subtasks = [ SubTask( title=f"(Mock) Clarify scope of requirement", description="Mock placeholder — confirm acceptance criteria with stakeholders.", estimated_complexity="s", ), SubTask( title=f"(Mock) Draft architecture", description="Mock placeholder — identify modules, data flow, contracts.", estimated_complexity="m", depends_on=[0], ), SubTask( title=f"(Mock) Implement feature", description="Mock placeholder — concrete coding work.", estimated_complexity="l", depends_on=[1], ), SubTask( title=f"(Mock) Add tests", description="Mock placeholder — unit + integration coverage.", estimated_complexity="m", depends_on=[2], ), SubTask( title=f"(Mock) Deploy + monitor", description="Mock placeholder — staging → production, observability.", estimated_complexity="s", depends_on=[3], ), ][:limit] return RequirementDecomposition( provider=self.name, model_used=self.default_model, execution_id=uuid4(), subtasks=subtasks, architecture_notes=["(Mock) Decomposition is deterministic placeholder data."], open_questions=[ "(Mock) Which non-functional requirements apply?", "(Mock) Any existing integration contracts to preserve?", ], input_tokens=max(10, len(req.requirement_text) // 4), output_tokens=400, duration_ms=self._latency, estimated_cost=0.0, ) # ============================================================================== # ProviderRegistry — maps provider name → AIProvider instance # ============================================================================== class ProviderRegistry: def __init__(self) -> None: self._providers: dict[str, AIProvider] = {} def register(self, name: str, provider: AIProvider) -> None: self._providers[name] = provider def get(self, name: str) -> AIProvider: if name not in self._providers: raise KeyError(f"Unknown AI provider: {name}. Registered: {list(self._providers)}") return self._providers[name] def names(self) -> list[str]: return list(self._providers.keys()) def default_registry() -> ProviderRegistry: r = ProviderRegistry() r.register("mock", MockProvider(response_latency_ms=40)) return r def build_default_gateway(): """Instantiate the canonical production AIGateway (mock by default). Imported lazily to avoid a cycle with ``app.gateway``. """ from .gateway import AIGateway, AIGatewayConfig return AIGateway(config=AIGatewayConfig(), registry=default_registry())