Coverage for baibu/chat/tagging.py: 100%

1"""Optional tagging of conversations by topic, language and intent. 

2 

3Off unless ``CHAT_TAGGING_ENABLED``. Topics are defined by the deployment in 

4the admin (*Chat → Topics*); intents come from ``CHAT_TAGGING_INTENTS``. 

5A Celery beat task tags conversations that have been idle for 

6``CHAT_TAGGING_IDLE_MINUTES`` and changed since they were last tagged, but 

7only conversations whose consent tier allows evaluation. 

8 

9The tagger asks a language model through LiteLLM (``CHAT_TAGGING_MODEL``, 

10or the chat's default model) using the template 

11``chat/tagging_prompt.txt``. With the ``mock`` model it matches topic 

12labels and slugs as keywords instead, so it works without a model. 

13""" 

14 

15import json 

16import logging 

17import re 

18from dataclasses import dataclass 

19from dataclasses import field 

20from datetime import timedelta 

21 

22from django.conf import settings 

23from django.db import transaction 

24from django.template.loader import render_to_string 

25from django.utils import timezone 

26 

27from baibu.users.consent import TIER_RANK 

28 

29from . import llm 

30from .models import Conversation 

31from .models import ConversationTag 

32from .models import Message 

33from .models import Topic 

34 

35logger = logging.getLogger(__name__) 

36 

37_JSON_BLOCK = re.compile(r"\{.*\}", re.DOTALL) 

38 

39 

40class TaggingError(Exception): 

41 pass 

42 

43 

44@dataclass 

45class TagResult: 

46 topics: list[str] = field(default_factory=list) 

47 language: str = "" 

48 intent: str = "" 

49 

50 

51def transcript(conversation: Conversation) -> str: 

52 lines = [ 

53 f"{'User' if m.role == Message.Role.USER else 'Assistant'}: {m.content}" 

54 for m in conversation.messages.exclude(content="")[: settings.CHAT_TAGGING_MAX_MESSAGES] 

55 ] 

56 return "\n".join(lines)[: settings.CHAT_TAGGING_MAX_CHARACTERS] 

57 

58 

59def model_config() -> llm.ModelConfig: 

60 from .services import default_variant 

61 

62 if settings.CHAT_TAGGING_MODEL: 

63 return llm.ModelConfig(model=settings.CHAT_TAGGING_MODEL) 

64 return llm.ModelConfig.from_variant(default_variant()) 

65 

66 

67def keyword_tags(text: str, topics: list[Topic], language: str) -> TagResult: 

68 """The mock tagger: a topic matches if its slug or label appears as words.""" 

69 lowered = text.lower() 

70 found = [] 

71 for topic in topics: 

72 for term in {topic.slug.replace("-", " "), topic.label.lower()}: 

73 if re.search(rf"\b{re.escape(term.lower())}\b", lowered): 

74 found.append(topic.slug) 

75 break 

76 intent = "question" if "?" in text else "conversation" 

77 return TagResult(topics=found, language=language, intent=intent if intent in intents() else "") 

78 

79 

80def intents() -> list[str]: 

81 return list(settings.CHAT_TAGGING_INTENTS) 

82 

83 

84def model_tags(text: str, topics: list[Topic], config: llm.ModelConfig) -> TagResult: 

85 prompt = render_to_string( 

86 "chat/tagging_prompt.txt", 

87 {"topics": topics, "intents": intents(), "max_topics": settings.CHAT_TAGGING_MAX_TOPICS}, 

88 ) 

89 try: 

90 completion = llm.complete(config, [{"role": "system", "content": prompt}, {"role": "user", "content": text}]) 

91 except llm.ModelError as exc: 

92 raise TaggingError(str(exc)) from exc 

93 match = _JSON_BLOCK.search(completion.content or "") 

94 try: 

95 data = json.loads(match.group(0) if match else completion.content) 

96 except (json.JSONDecodeError, TypeError) as exc: 

97 msg = "Tagger answer is not JSON." 

98 raise TaggingError(msg) from exc 

99 if not isinstance(data, dict): 

100 msg = "Tagger answer is not an object." 

101 raise TaggingError(msg) 

102 known = {topic.slug for topic in topics} 

103 raw_topics = data.get("topics") if isinstance(data.get("topics"), list) else [] 

104 chosen = [slug for slug in dict.fromkeys(str(t) for t in raw_topics) if slug in known] 

105 intent = str(data.get("intent") or "") 

106 return TagResult( 

107 topics=chosen[: settings.CHAT_TAGGING_MAX_TOPICS], 

108 language=str(data.get("language") or "")[:24], 

109 intent=intent if intent in intents() else "", 

110 ) 

111 

112 

113def tag_conversation(conversation: Conversation) -> ConversationTag: 

114 topics = list(Topic.objects.filter(active=True)) 

115 activity = conversation.last_activity_at 

116 config = model_config() 

117 tag, _created = ConversationTag.objects.get_or_create( 

118 conversation=conversation, defaults={"tagged_activity_at": activity} 

119 ) 

120 text = transcript(conversation) 

121 try: 

122 if config.model == "mock": 

123 result = keyword_tags(text, topics, conversation.language_code) 

124 tagger = "keywords" 

125 else: 

126 result = model_tags(text, topics, config) 

127 tagger = config.model 

128 error = "" 

129 except TaggingError as exc: 

130 result, tagger, error = TagResult(), config.model, str(exc)[:500] 

131 logger.warning("Could not tag conversation %s: %s", conversation.pk, error) 

132 with transaction.atomic(): 

133 tag.language = result.language 

134 tag.intent = result.intent 

135 tag.tagger = tagger[:255] 

136 tag.error = error 

137 tag.tagged_activity_at = activity 

138 tag.save() 

139 tag.topics.set([t for t in topics if t.slug in result.topics]) 

140 return tag 

141 

142 

143def eligible_tiers() -> list[str]: 

144 minimum = TIER_RANK[settings.CHAT_TAGGING_MIN_TIER] 

145 return [tier for tier, rank in TIER_RANK.items() if rank >= minimum] 

146 

147 

148def due_conversations(now=None): 

149 """Idle conversations with consent, never tagged or changed since tagging.""" 

150 from django.db.models import F 

151 from django.db.models import Q 

152 

153 now = now or timezone.now() 

154 idle_since = now - timedelta(minutes=settings.CHAT_TAGGING_IDLE_MINUTES) 

155 return ( 

156 Conversation.objects.filter(last_activity_at__lte=idle_since, consent_tier__in=eligible_tiers()) 

157 .filter(messages__isnull=False) 

158 .filter(Q(tag__isnull=True) | Q(tag__tagged_activity_at__lt=F("last_activity_at"))) 

159 .distinct() 

160 .order_by("last_activity_at") 

161 ) 

162 

163 

164def tag_due(now=None) -> int: 

165 if not settings.CHAT_TAGGING_ENABLED: 

166 return 0 

167 count = 0 

168 for conversation in due_conversations(now)[: settings.CHAT_TAGGING_BATCH_SIZE]: 

169 tag_conversation(conversation) 

170 count += 1 

171 return count