Coverage for src/semware/models/schemas.py: 91%

1"""Pydantic schemas for data validation.""" 

2 

3from datetime import datetime 

4from typing import Any 

5 

6from pydantic import BaseModel, Field, field_validator 

7 

8 

9class TableSchema(BaseModel): 

10 """Schema definition for a table.""" 

11 

12 name: str = Field(..., description="Name of the table") 

13 columns: dict[str, str] = Field(..., description="Column name to data type mapping") 

14 id_column: str = Field(..., description="Name of the unique identifier column") 

15 embedding_column: str = Field( 

16 ..., description="Name of the column used for embeddings" 

17 ) 

18 

19 @field_validator("embedding_column") 

20 @classmethod 

21 def validate_embedding_column(cls, v: str, info) -> str: 

22 """Validate that embedding column exists and is string type.""" 

23 if hasattr(info, "data") and "columns" in info.data: 

24 columns = info.data["columns"] 

25 if v not in columns: 

26 raise ValueError(f"Embedding column '{v}' not found in columns") 

27 if columns[v] != "string": 

28 raise ValueError(f"Embedding column '{v}' must be of type 'string'") 

29 return v 

30 

31 @field_validator("id_column") 

32 @classmethod 

33 def validate_id_column(cls, v: str, info) -> str: 

34 """Validate that ID column exists in columns.""" 

35 if hasattr(info, "data") and "columns" in info.data: 

36 columns = info.data["columns"] 

37 if v not in columns: 

38 raise ValueError(f"ID column '{v}' not found in columns") 

39 return v 

40 

41 

42class TableInfo(BaseModel): 

43 """Information about a table.""" 

44 

45 name: str 

46 table_schema: TableSchema = Field(..., alias="schema") 

47 created_at: datetime 

48 record_count: int = 0 

49 

50 model_config = {"populate_by_name": True} 

51 

52 

53class DataRecord(BaseModel): 

54 """A data record to be inserted/updated.""" 

55 

56 data: dict[str, Any] = Field(..., description="The actual data record") 

57 

58 def get_id(self, id_column: str) -> str | int: 

59 """Get the ID value from the data record.""" 

60 if id_column not in self.data: 

61 raise ValueError(f"ID column '{id_column}' not found in data") 

62 return self.data[id_column] 

63 

64 def get_text_for_embedding(self, embedding_column: str) -> str: 

65 """Get the text content for embedding generation.""" 

66 if embedding_column not in self.data: 

67 raise ValueError(f"Embedding column '{embedding_column}' not found in data") 

68 

69 value = self.data[embedding_column] 

70 if not isinstance(value, str): 

71 raise ValueError( 

72 f"Embedding column '{embedding_column}' must contain string data" 

73 ) 

74 

75 return value 

76 

77 

78class SearchResult(BaseModel): 

79 """A search result record.""" 

80 

81 id: str | int 

82 data: dict[str, Any] 

83 similarity_score: float = Field( 

84 ..., ge=0.0, le=1.0, description="Similarity score between 0 and 1" 

85 ) 

86 

87 

88class ErrorResponse(BaseModel): 

89 """Standard error response.""" 

90 

91 error: str = Field(..., description="Error message") 

92 detail: str | None = Field(None, description="Detailed error information") 

93 error_code: str | None = Field( 

94 None, description="Error code for programmatic handling" 

95 ) 

96 

97 

98class SuccessResponse(BaseModel): 

99 """Standard success response.""" 

100 

101 message: str = Field(..., description="Success message") 

102 data: dict[str, Any] | None = Field(None, description="Additional response data")