Coverage for src / lexigram / contracts / ai / chat.py: 93%

1"""Chat History contracts for Lexigram. 

2 

3Defines chat message and history classes analogous to LangChain's chat history. 

4""" 

5 

6from __future__ import annotations 

7 

8from dataclasses import dataclass, field 

9from typing import Any 

10 

11 

12@dataclass(frozen=True) 

13class ChatMessage: 

14 """A chat message (like LangChain's BaseMessage). 

15 

16 Attributes: 

17 role: The role (user, assistant, system). 

18 content: The message content. 

19 """ 

20 

21 role: str 

22 content: str 

23 name: str | None = None 

24 additional_kwargs: dict[str, Any] = field(default_factory=dict) 

25 

26 def to_lc_dict(self) -> dict[str, Any]: 

27 """Convert to LangChain dict format.""" 

28 return { 

29 "type": "human" if self.role == "user" else "ai", 

30 "data": { 

31 "content": self.content, 

32 "additional_kwargs": self.additional_kwargs, 

33 }, 

34 } 

35 

36 

37class ChatHistory: 

38 """Chat history (like LangChain's ChatMessageHistory). 

39 

40 Manages chat messages and converts to LangChain format. 

41 """ 

42 

43 def __init__(self) -> None: 

44 self._messages: list[ChatMessage] = [] 

45 

46 def add_user_message(self, content: str) -> None: 

47 """Add a user message.""" 

48 self._messages.append(ChatMessage(role="user", content=content)) 

49 

50 def add_ai_message(self, content: str) -> None: 

51 """Add an AI message.""" 

52 self._messages.append(ChatMessage(role="assistant", content=content)) 

53 

54 def add_message(self, message: ChatMessage) -> None: 

55 """Add a message.""" 

56 self._messages.append(message) 

57 

58 def get_messages(self) -> list[ChatMessage]: 

59 """Get all messages.""" 

60 return list(self._messages) 

61 

62 def clear(self) -> None: 

63 """Clear all messages.""" 

64 self._messages.clear() 

65 

66 def to_lc_format(self) -> list[dict[str, Any]]: 

67 """Convert to LangChain message format.""" 

68 return [msg.to_lc_dict() for msg in self._messages] 

69 

70 

71__all__ = [ 

72 "ChatHistory", 

73 "ChatMessage", 

74]