diff --git a/src/codemodder/codetf/common.py b/src/codemodder/codetf/common.py index 7bfdee6d..f2cb5a9d 100644 --- a/src/codemodder/codetf/common.py +++ b/src/codemodder/codetf/common.py @@ -61,12 +61,14 @@ class DiffSide(CaseInsensitiveEnum): class PackageAction(BaseModel): + model_config = ConfigDict(frozen=True) action: Action result: PackageResult package: str class Change(BaseModel): + model_config = ConfigDict(frozen=True) lineNumber: int description: Optional[str] diffSide: DiffSide = DiffSide.RIGHT @@ -87,11 +89,13 @@ def validate_description(self): class Rating(BaseModel): + model_config = ConfigDict(frozen=True) score: int description: Optional[str] = None class FixQuality(BaseModel): + model_config = ConfigDict(frozen=True) safetyRating: Rating effectivenessRating: Rating cleanlinessRating: Rating diff --git a/src/codemodder/codetf/v3/codetf.py b/src/codemodder/codetf/v3/codetf.py index b4036472..d60e140c 100644 --- a/src/codemodder/codetf/v3/codetf.py +++ b/src/codemodder/codetf/v3/codetf.py @@ -3,9 +3,8 @@ from enum import Enum from typing import Optional -from pydantic import BaseModel, model_validator +from pydantic import BaseModel, ConfigDict, model_validator -from codemodder.llm import TokenUsage from codemodder.logging import logger from ..common import Change, CodeTFWriter, Finding, FixQuality @@ -39,6 +38,7 @@ class Run(BaseModel): class FixStatusType(str, Enum): """Status of a fix""" + model_config = ConfigDict(frozen=True) fixed = "fixed" skipped = "skipped" failed = "failed" @@ -48,12 +48,15 @@ class FixStatusType(str, Enum): class FixStatus(BaseModel): """Metadata describing fix outcome""" + model_config = ConfigDict(frozen=True) + status: FixStatusType reason: Optional[str] = None details: Optional[str] = None class ChangeSet(BaseModel): + model_config = ConfigDict(frozen=True) path: str diff: str changes: list[Change] = [] @@ -76,6 +79,7 @@ class Strategy(str, Enum): class AIMetadata(BaseModel): + model_config = ConfigDict(frozen=True) provider: Optional[str] = None models: Optional[list[str]] = None total_tokens: Optional[int] = None @@ -84,12 +88,14 @@ class AIMetadata(BaseModel): class GenerationMetadata(BaseModel): + model_config = ConfigDict(frozen=True) strategy: Strategy ai: Optional[AIMetadata] = None provisional: bool class FixMetadata(BaseModel): + model_config = ConfigDict(frozen=True) # Fix provider ID, corresponds to legacy codemod ID id: str # A brief summary of the fix @@ -103,6 +109,8 @@ class FixMetadata(BaseModel): class FixResult(BaseModel): """Result corresponding to a single finding""" + model_config = ConfigDict(frozen=True) + finding: Finding fixStatus: FixStatus changeSets: list[ChangeSet] = [] @@ -120,27 +128,6 @@ def validate_fixMetadata(self): raise ValueError("fixMetadata must be provided for fixed results") return self - def store_token_data(self, token_usage: TokenUsage): - if token_usage == TokenUsage() or not self.fixMetadata: - return - - if not self.fixMetadata.generation.ai: - self.fixMetadata.generation.ai = AIMetadata( - total_tokens=token_usage.total, - completion_tokens=token_usage.completion_tokens, - prompt_tokens=token_usage.prompt_tokens, - ) - else: - self.fixMetadata.generation.ai.completion_tokens = ( - self.fixMetadata.generation.ai.completion_tokens or 0 - ) + token_usage.completion_tokens - self.fixMetadata.generation.ai.prompt_tokens = ( - self.fixMetadata.generation.ai.prompt_tokens or 0 - ) + token_usage.prompt_tokens - self.fixMetadata.generation.ai.total_tokens = ( - self.fixMetadata.generation.ai.total_tokens or 0 - ) + token_usage.total - class CodeTF(CodeTFWriter, BaseModel): run: Run