Spaces:
Running on Zero
Running on Zero
| """Controlled Variable-level descent through direct GCMD children.""" | |
| from __future__ import annotations | |
| from pydantic import BaseModel, ConfigDict, Field | |
| from gcmd_classifier.classification.candidates import ( | |
| VariableCandidate, | |
| build_variable_candidates, | |
| validate_variable_candidate_relationship, | |
| ) | |
| from gcmd_classifier.classification.routers import ModelCallMetadata, TermBranchSeed | |
| from gcmd_classifier.config import ModelSettings | |
| from gcmd_classifier.errors import UnknownCandidateIDError | |
| from gcmd_classifier.llm.base import ( | |
| ModelClient, | |
| ModelRequest, | |
| ModelStage, | |
| RetryPolicy, | |
| generate_with_retries, | |
| ) | |
| from gcmd_classifier.llm.prompts import ParentContext, build_variable_prompt | |
| from gcmd_classifier.llm.schemas import CandidateDecision, VariableResponse | |
| from gcmd_classifier.models import ( | |
| ArticleRecord, | |
| HierarchyLevel, | |
| OutputError, | |
| OutputWarning, | |
| SupportType, | |
| ) | |
| from gcmd_classifier.vocabulary.index import VocabularyIndex | |
| _VARIABLE_PARENT_LEVELS = {"Term", "Variable_Level_1", "Variable_Level_2"} | |
| _VARIABLE_LEVELS = {"Variable_Level_1", "Variable_Level_2", "Variable_Level_3"} | |
| class EvidenceStep(BaseModel): | |
| """Evidence and routing metadata for one stage in a branch lineage.""" | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| branch_id: str = Field(min_length=1) | |
| UUID: str = Field(min_length=1) | |
| name: str = Field(min_length=1) | |
| level: str = Field(min_length=1) | |
| evidence: str = Field(min_length=1) | |
| support_type: SupportType | |
| confidence: float | None = Field(default=None, ge=0.0, le=1.0) | |
| reason: str | None = None | |
| candidate_id: str | None = None | |
| class VariableTerminalOutcome(BaseModel): | |
| """Terminal outcome for one independently descended branch.""" | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| branch_id: str = Field(min_length=1) | |
| parent_branch_id: str = Field(min_length=1) | |
| topic_uuid: str = Field(min_length=1) | |
| topic_name: str = Field(min_length=1) | |
| term_uuid: str = Field(min_length=1) | |
| term_name: str = Field(min_length=1) | |
| final_uuid: str = Field(min_length=1) | |
| final_name: str = Field(min_length=1) | |
| final_level: HierarchyLevel | |
| final_canonical_path: str = Field(min_length=1) | |
| path_components: tuple[str, ...] = Field(min_length=1) | |
| evidence: str = Field(min_length=1) | |
| evidence_trail: tuple[EvidenceStep, ...] = Field(default_factory=tuple) | |
| support_type: SupportType | |
| confidence: float | None = Field(default=None, ge=0.0, le=1.0) | |
| reason: str | None = None | |
| stop_reason: str | None = None | |
| candidate_id: str | None = None | |
| prompt_version: str | None = None | |
| model_provider: str | None = None | |
| model_name: str | None = None | |
| retry_count: int = Field(default=0, ge=0) | |
| class VariableBranchError(BaseModel): | |
| """Explicit branch-level failure that does not cancel sibling branches.""" | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| branch_id: str = Field(min_length=1) | |
| parent_branch_id: str = Field(min_length=1) | |
| parent_uuid: str = Field(min_length=1) | |
| parent_name: str = Field(min_length=1) | |
| parent_level: str = Field(min_length=1) | |
| code: str = Field(min_length=1) | |
| message: str = Field(min_length=1) | |
| class VariableDescentResult(BaseModel): | |
| """Result of controlled Variable-level descent for one Term branch.""" | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| terminals: tuple[VariableTerminalOutcome, ...] = Field(default_factory=tuple) | |
| errors: tuple[VariableBranchError, ...] = Field(default_factory=tuple) | |
| warnings: tuple[OutputWarning, ...] = Field(default_factory=tuple) | |
| diagnostics: tuple[OutputError, ...] = Field(default_factory=tuple) | |
| def terminal_count(self) -> int: | |
| """Number of successful terminal branch outcomes.""" | |
| return len(self.terminals) | |
| def has_errors(self) -> bool: | |
| """Whether any independent branch failed during descent.""" | |
| return bool(self.errors) | |
| def descend_variables( | |
| *, | |
| article: ArticleRecord, | |
| term_branch: TermBranchSeed, | |
| vocabulary: VocabularyIndex, | |
| model_client: ModelClient, | |
| settings: ModelSettings, | |
| retry_policy: RetryPolicy | None = None, | |
| ) -> VariableDescentResult: | |
| """Descend from a selected Term branch through supported direct Variable children.""" | |
| term = vocabulary.get(term_branch.term_uuid) | |
| trail = (_evidence_from_term_branch(term_branch),) | |
| context = _BranchContext( | |
| topic_uuid=term_branch.parent_topic_uuid, | |
| topic_name=term_branch.parent_topic_name, | |
| term_uuid=term_branch.term_uuid, | |
| term_name=term_branch.term_name, | |
| ) | |
| terminals, errors = _descend_parent( | |
| article=article, | |
| parent_uuid=term.UUID, | |
| parent_branch_id=term_branch.branch_id, | |
| branch_id=term_branch.branch_id, | |
| candidate_id=term_branch.candidate_id, | |
| parent_evidence=term_branch.evidence, | |
| parent_support_type=term_branch.support_type, | |
| parent_confidence=term_branch.confidence, | |
| parent_reason=term_branch.reason, | |
| trail=trail, | |
| context=context, | |
| vocabulary=vocabulary, | |
| model_client=model_client, | |
| settings=settings, | |
| retry_policy=retry_policy or RetryPolicy.from_settings(settings), | |
| parent_metadata=None, | |
| ) | |
| return VariableDescentResult(terminals=terminals, errors=errors) | |
| class _BranchContext(BaseModel): | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| topic_uuid: str | |
| topic_name: str | |
| term_uuid: str | |
| term_name: str | |
| def _descend_parent( | |
| *, | |
| article: ArticleRecord, | |
| parent_uuid: str, | |
| parent_branch_id: str, | |
| branch_id: str, | |
| candidate_id: str | None, | |
| parent_evidence: str, | |
| parent_support_type: SupportType, | |
| parent_confidence: float | None, | |
| parent_reason: str | None, | |
| trail: tuple[EvidenceStep, ...], | |
| context: _BranchContext, | |
| vocabulary: VocabularyIndex, | |
| model_client: ModelClient, | |
| settings: ModelSettings, | |
| retry_policy: RetryPolicy, | |
| parent_metadata: ModelCallMetadata | None, | |
| ) -> tuple[tuple[VariableTerminalOutcome, ...], tuple[VariableBranchError, ...]]: | |
| parent = vocabulary.get(parent_uuid) | |
| _validate_variable_parent(parent.level) | |
| if parent.level == "Variable_Level_3": | |
| return ( | |
| ( | |
| _terminal_outcome( | |
| parent=parent, | |
| branch_id=branch_id, | |
| parent_branch_id=parent_branch_id, | |
| context=context, | |
| evidence=parent_evidence, | |
| support_type=parent_support_type, | |
| confidence=parent_confidence, | |
| reason=parent_reason, | |
| stop_reason="Selected Variable_Level_3 is a leaf node.", | |
| candidate_id=candidate_id, | |
| trail=trail, | |
| metadata=parent_metadata, | |
| ), | |
| ), | |
| (), | |
| ) | |
| candidates = build_variable_candidates(vocabulary, parent_uuid=parent_uuid) | |
| if not candidates: | |
| return ( | |
| ( | |
| _terminal_outcome( | |
| parent=parent, | |
| branch_id=branch_id, | |
| parent_branch_id=parent_branch_id, | |
| context=context, | |
| evidence=parent_evidence, | |
| support_type=parent_support_type, | |
| confidence=parent_confidence, | |
| reason=parent_reason, | |
| stop_reason="No direct Variable children are available.", | |
| candidate_id=candidate_id, | |
| trail=trail, | |
| metadata=parent_metadata, | |
| ), | |
| ), | |
| (), | |
| ) | |
| candidates_by_id = {candidate.candidate_id: candidate for candidate in candidates} | |
| response = _call_variable_model( | |
| article=article, | |
| parent=parent, | |
| candidates=candidates, | |
| vocabulary=vocabulary, | |
| model_client=model_client, | |
| settings=settings, | |
| retry_policy=retry_policy, | |
| ) | |
| _validate_selected_candidate_ids(response.parsed.selected, candidates_by_id) | |
| metadata = _model_metadata(response) | |
| if response.parsed.stop_at_parent: | |
| return ( | |
| ( | |
| _terminal_outcome( | |
| parent=parent, | |
| branch_id=branch_id, | |
| parent_branch_id=parent_branch_id, | |
| context=context, | |
| evidence=parent_evidence, | |
| support_type=parent_support_type, | |
| confidence=parent_confidence, | |
| reason=parent_reason, | |
| stop_reason=response.parsed.stop_reason, | |
| candidate_id=candidate_id, | |
| trail=trail, | |
| metadata=metadata, | |
| ), | |
| ), | |
| (), | |
| ) | |
| terminals: list[VariableTerminalOutcome] = [] | |
| errors: list[VariableBranchError] = [] | |
| for decision in response.parsed.selected: | |
| candidate = candidates_by_id[decision.candidate_id] | |
| child = vocabulary.get(candidate.variable_uuid) | |
| child_branch_id = f"{branch_id}/variable:{candidate.candidate_id}" | |
| try: | |
| validate_variable_candidate_relationship( | |
| candidate, | |
| selected_parent_uuid=parent_uuid, | |
| index=vocabulary, | |
| ) | |
| child_trail = ( | |
| *trail, | |
| _evidence_from_child( | |
| child=child, | |
| branch_id=child_branch_id, | |
| decision=decision, | |
| ), | |
| ) | |
| child_terminals, child_errors = _descend_parent( | |
| article=article, | |
| parent_uuid=child.UUID, | |
| parent_branch_id=branch_id, | |
| branch_id=child_branch_id, | |
| candidate_id=decision.candidate_id, | |
| parent_evidence=decision.evidence, | |
| parent_support_type=decision.support_type, | |
| parent_confidence=decision.confidence, | |
| parent_reason=decision.reason, | |
| trail=child_trail, | |
| context=context, | |
| vocabulary=vocabulary, | |
| model_client=model_client, | |
| settings=settings, | |
| retry_policy=retry_policy, | |
| parent_metadata=metadata, | |
| ) | |
| terminals.extend(child_terminals) | |
| errors.extend(child_errors) | |
| except Exception as exc: | |
| errors.append( | |
| _branch_error( | |
| branch_id=child_branch_id, | |
| parent_branch_id=branch_id, | |
| parent=child, | |
| exc=exc, | |
| ) | |
| ) | |
| return tuple(terminals), tuple(errors) | |
| def _call_variable_model( | |
| *, | |
| article: ArticleRecord, | |
| parent, | |
| candidates: tuple[VariableCandidate, ...], | |
| vocabulary: VocabularyIndex, | |
| model_client: ModelClient, | |
| settings: ModelSettings, | |
| retry_policy: RetryPolicy, | |
| ): | |
| parent_context = ParentContext( | |
| candidate_id=parent.UUID, | |
| name=parent.name, | |
| level=parent.level, | |
| canonical_path=parent.canonical_path, | |
| ) | |
| prompt = build_variable_prompt( | |
| article=article, | |
| parent=parent_context, | |
| candidates=tuple(candidate.prompt_candidate for candidate in candidates), | |
| prompt_version=settings.prompt_version_variable, | |
| ) | |
| request = ModelRequest.from_settings( | |
| stage=ModelStage.VARIABLE, | |
| prompt=prompt, | |
| response_schema=VariableResponse, | |
| settings=settings, | |
| metadata={ | |
| "DOI": article.DOI, | |
| "parent_uuid": parent.UUID, | |
| "parent_level": parent.level, | |
| "candidate_ids": tuple(candidate.candidate_id for candidate in candidates), | |
| "candidate_uuids": tuple(candidate.variable_uuid for candidate in candidates), | |
| "vocabulary_version": vocabulary.vocabulary_version, | |
| }, | |
| ) | |
| return generate_with_retries(model_client, request, retry_policy) | |
| def _validate_selected_candidate_ids( | |
| selected: list[CandidateDecision], | |
| candidates_by_id: dict[str, VariableCandidate], | |
| ) -> None: | |
| seen: set[str] = set() | |
| for decision in selected: | |
| if decision.candidate_id not in candidates_by_id: | |
| raise UnknownCandidateIDError( | |
| f"Variable model selected unknown candidate_id {decision.candidate_id!r}." | |
| ) | |
| if decision.candidate_id in seen: | |
| raise UnknownCandidateIDError( | |
| f"Variable model selected duplicate candidate_id {decision.candidate_id!r}." | |
| ) | |
| seen.add(decision.candidate_id) | |
| def _terminal_outcome( | |
| *, | |
| parent, | |
| branch_id: str, | |
| parent_branch_id: str, | |
| context: _BranchContext, | |
| evidence: str, | |
| support_type: SupportType, | |
| confidence: float | None, | |
| reason: str | None, | |
| stop_reason: str | None, | |
| candidate_id: str | None, | |
| trail: tuple[EvidenceStep, ...], | |
| metadata: ModelCallMetadata | None, | |
| ) -> VariableTerminalOutcome: | |
| return VariableTerminalOutcome( | |
| branch_id=branch_id, | |
| parent_branch_id=parent_branch_id, | |
| topic_uuid=context.topic_uuid, | |
| topic_name=context.topic_name, | |
| term_uuid=context.term_uuid, | |
| term_name=context.term_name, | |
| final_uuid=parent.UUID, | |
| final_name=parent.name, | |
| final_level=parent.level, | |
| final_canonical_path=parent.canonical_path, | |
| path_components=parent.path_components, | |
| evidence=evidence, | |
| evidence_trail=trail, | |
| support_type=support_type, | |
| confidence=confidence, | |
| reason=reason, | |
| stop_reason=stop_reason, | |
| candidate_id=candidate_id, | |
| prompt_version=None if metadata is None else metadata.prompt_version, | |
| model_provider=None if metadata is None else metadata.provider, | |
| model_name=None if metadata is None else metadata.model_name, | |
| retry_count=0 if metadata is None else metadata.retry_count, | |
| ) | |
| def _evidence_from_term_branch(term_branch: TermBranchSeed) -> EvidenceStep: | |
| return EvidenceStep( | |
| branch_id=term_branch.branch_id, | |
| UUID=term_branch.term_uuid, | |
| name=term_branch.term_name, | |
| level=term_branch.term_level, | |
| evidence=term_branch.evidence, | |
| support_type=term_branch.support_type, | |
| confidence=term_branch.confidence, | |
| reason=term_branch.reason, | |
| candidate_id=term_branch.candidate_id, | |
| ) | |
| def _evidence_from_child(*, child, branch_id: str, decision: CandidateDecision) -> EvidenceStep: | |
| return EvidenceStep( | |
| branch_id=branch_id, | |
| UUID=child.UUID, | |
| name=child.name, | |
| level=child.level, | |
| evidence=decision.evidence, | |
| support_type=decision.support_type, | |
| confidence=decision.confidence, | |
| reason=decision.reason, | |
| candidate_id=decision.candidate_id, | |
| ) | |
| def _branch_error( | |
| *, | |
| branch_id: str, | |
| parent_branch_id: str, | |
| parent, | |
| exc: Exception, | |
| ) -> VariableBranchError: | |
| return VariableBranchError( | |
| branch_id=branch_id, | |
| parent_branch_id=parent_branch_id, | |
| parent_uuid=parent.UUID, | |
| parent_name=parent.name, | |
| parent_level=parent.level, | |
| code=exc.__class__.__name__, | |
| message=str(exc), | |
| ) | |
| def _validate_variable_parent(level: str) -> None: | |
| if level not in _VARIABLE_PARENT_LEVELS and level not in _VARIABLE_LEVELS: | |
| raise ValueError(f"Unsupported Variable descent parent level: {level!r}.") | |
| def _model_metadata(response) -> ModelCallMetadata: | |
| token_usage = response.token_usage | |
| return ModelCallMetadata( | |
| provider=response.provider, | |
| model_name=response.model_name, | |
| prompt_version=response.prompt_version, | |
| retry_count=response.retry_count, | |
| duration_seconds=response.duration_seconds, | |
| input_tokens=None if token_usage is None else token_usage.input_tokens, | |
| output_tokens=None if token_usage is None else token_usage.output_tokens, | |
| total_tokens=None if token_usage is None else token_usage.total_tokens, | |
| estimated_cost=response.estimated_cost, | |
| ) | |