Coverage for src/lexigram/ai/types.py: 100%

1"""Shared type definitions for Lexigram AI. 

2 

3This module contains type definitions that are used across multiple 

4sub-modules in intelligence package. 

5""" 

6 

7from __future__ import annotations 

8 

9from collections.abc import Callable 

10from dataclasses import dataclass 

11from datetime import UTC, datetime 

12from typing import Any, TypeVar 

13 

14from lexigram.contracts.ai.types import EmbeddingModel, ModelProvider, VectorProvider 

15from lexigram.contracts.core import JSON, Metadata 

16from lexigram.domain import DomainModel 

17from lexigram.validation import ConfigDict, Field 

18 

19# Note: model_serializer not yet implemented in lexigram.contracts 

20# Using a stub for now 

21DecoratedMethod = TypeVar("DecoratedMethod", bound=Callable[..., Any]) 

22 

23 

24def model_serializer(mode: Any) -> Callable[[DecoratedMethod], DecoratedMethod]: 

25 def decorator(func: DecoratedMethod) -> DecoratedMethod: 

26 return func 

27 

28 return decorator 

29 

30 

31@dataclass(init=False) 

32class AIBaseEvent(DomainModel): 

33 """Base class for AI intelligence domain events.""" 

34 

35 timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC)) 

36 metadata: Metadata = Field(default_factory=dict) 

37 

38 model_config = ConfigDict(arbitrary_types_allowed=True) 

39 

40 @model_serializer(mode="wrap") 

41 def serialize_model(self, handler: Callable[..., Any]) -> Any: 

42 """Custom serializer to handle datetime objects.""" 

43 data = handler(self) 

44 # Convert datetime objects to ISO format strings 

45 if isinstance(data.get("timestamp"), datetime): 

46 data["timestamp"] = data["timestamp"].isoformat() 

47 return data 

48 

49 

50__all__ = [ 

51 "JSON", 

52 "AIBaseEvent", 

53 "EmbeddingModel", 

54 "Metadata", 

55 "ModelProvider", 

56 "VectorProvider", 

57]