WIP

tools/models.py

@@ -41,6 +41,7 @@ class ToolOutput(BaseModel):
         "test_sample_needed",
         "more_tests_required",
         "refactor_analysis_complete",
+        "trace_complete",
         "resend_prompt",
         "continuation_available",
     ] = "success"
@@ -150,6 +151,129 @@ class ResendPromptRequest(BaseModel):
     metadata: dict[str, Any] = Field(default_factory=dict)
 
 
+class TraceEntryPoint(BaseModel):
+    """Entry point information for trace analysis"""
+
+    file: str = Field(..., description="Absolute path to the file")
+    class_or_struct: str = Field(..., description="Class or module name")
+    method: str = Field(..., description="Method or function name")
+    signature: str = Field(..., description="Full method signature")
+    parameters: Optional[dict[str, Any]] = Field(default_factory=dict, description="Parameter values used in analysis")
+
+
+class TraceTarget(BaseModel):
+    """Target information for dependency analysis"""
+
+    file: str = Field(..., description="Absolute path to the file")
+    class_or_struct: str = Field(..., description="Class or module name")
+    method: str = Field(..., description="Method or function name")
+    signature: str = Field(..., description="Full method signature")
+
+
+class CallPathStep(BaseModel):
+    """A single step in the call path trace"""
+
+    from_info: dict[str, Any] = Field(..., description="Source location information", alias="from")
+    to: dict[str, Any] = Field(..., description="Target location information")
+    reason: str = Field(..., description="Reason for the call or dependency")
+    condition: Optional[str] = Field(None, description="Conditional logic if applicable")
+    ambiguous: bool = Field(False, description="Whether this call is ambiguous")
+
+
+class BranchingPoint(BaseModel):
+    """A branching point in the execution flow"""
+
+    file: str = Field(..., description="File containing the branching point")
+    method: str = Field(..., description="Method containing the branching point")
+    line: int = Field(..., description="Line number of the branching point")
+    condition: str = Field(..., description="Branching condition")
+    branches: list[str] = Field(..., description="Possible execution branches")
+    ambiguous: bool = Field(False, description="Whether the branching is ambiguous")
+
+
+class SideEffect(BaseModel):
+    """A side effect detected in the trace"""
+
+    type: str = Field(..., description="Type of side effect")
+    description: str = Field(..., description="Description of the side effect")
+    file: str = Field(..., description="File where the side effect occurs")
+    method: str = Field(..., description="Method where the side effect occurs")
+    line: int = Field(..., description="Line number of the side effect")
+
+
+class UnresolvedDependency(BaseModel):
+    """An unresolved dependency in the trace"""
+
+    reason: str = Field(..., description="Reason why the dependency is unresolved")
+    affected_file: str = Field(..., description="File affected by the unresolved dependency")
+    line: int = Field(..., description="Line number of the unresolved dependency")
+
+
+class IncomingDependency(BaseModel):
+    """An incoming dependency (what calls this target)"""
+
+    from_file: str = Field(..., description="Source file of the dependency")
+    from_class: str = Field(..., description="Source class of the dependency")
+    from_method: str = Field(..., description="Source method of the dependency")
+    line: int = Field(..., description="Line number of the dependency")
+    type: str = Field(..., description="Type of dependency")
+
+
+class OutgoingDependency(BaseModel):
+    """An outgoing dependency (what this target calls)"""
+
+    to_file: str = Field(..., description="Target file of the dependency")
+    to_class: str = Field(..., description="Target class of the dependency")
+    to_method: str = Field(..., description="Target method of the dependency")
+    line: int = Field(..., description="Line number of the dependency")
+    type: str = Field(..., description="Type of dependency")
+
+
+class TypeDependency(BaseModel):
+    """A type-level dependency (inheritance, imports, etc.)"""
+
+    dependency_type: str = Field(..., description="Type of dependency")
+    source_file: str = Field(..., description="Source file of the dependency")
+    source_entity: str = Field(..., description="Source entity (class, module)")
+    target: str = Field(..., description="Target entity")
+
+
+class StateAccess(BaseModel):
+    """State access information"""
+
+    file: str = Field(..., description="File where state is accessed")
+    method: str = Field(..., description="Method accessing the state")
+    access_type: str = Field(..., description="Type of access (reads, writes, etc.)")
+    state_entity: str = Field(..., description="State entity being accessed")
+
+
+class TraceComplete(BaseModel):
+    """Complete trace analysis response"""
+
+    status: Literal["trace_complete"] = "trace_complete"
+    trace_type: Literal["precision", "dependencies"] = Field(..., description="Type of trace performed")
+
+    # Precision mode fields
+    entry_point: Optional[TraceEntryPoint] = Field(None, description="Entry point for precision trace")
+    call_path: Optional[list[CallPathStep]] = Field(default_factory=list, description="Call path for precision trace")
+    branching_points: Optional[list[BranchingPoint]] = Field(default_factory=list, description="Branching points")
+    side_effects: Optional[list[SideEffect]] = Field(default_factory=list, description="Side effects detected")
+    unresolved: Optional[list[UnresolvedDependency]] = Field(
+        default_factory=list, description="Unresolved dependencies"
+    )
+
+    # Dependencies mode fields
+    target: Optional[TraceTarget] = Field(None, description="Target for dependency analysis")
+    incoming_dependencies: Optional[list[IncomingDependency]] = Field(
+        default_factory=list, description="Incoming dependencies"
+    )
+    outgoing_dependencies: Optional[list[OutgoingDependency]] = Field(
+        default_factory=list, description="Outgoing dependencies"
+    )
+    type_dependencies: Optional[list[TypeDependency]] = Field(default_factory=list, description="Type dependencies")
+    state_access: Optional[list[StateAccess]] = Field(default_factory=list, description="State access information")
+
+
 # Registry mapping status strings to their corresponding Pydantic models
 SPECIAL_STATUS_MODELS = {
     "clarification_required": ClarificationRequest,
@@ -158,6 +282,7 @@ SPECIAL_STATUS_MODELS = {
     "test_sample_needed": TestSampleNeeded,
     "more_tests_required": MoreTestsRequired,
     "refactor_analysis_complete": RefactorAnalysisComplete,
+    "trace_complete": TraceComplete,
     "resend_prompt": ResendPromptRequest,
 }