DavidL72Code/UMB_Sustainable_Chatbot
0
1"""Serializable conversation state transitions for corpus-grounded RAG chat."""2 3from __future__ import annotations4 5import re6from typing import Callable, Iterable, Optional7 8 9STATE_VERSION = 310 11 12def empty_state() -> dict:13 return {14 "version": STATE_VERSION,15 "mode": "idle",16 "active_subject": None,17 "candidate_subjects": [],18 "active_scope": None,19 "pending_query": None,20 "last_intent": None,21 "clarification_options": [],22 "subject_history": [],23 "last_query": None,24 }25 26 27def normalize_state(state: Optional[dict]) -> dict:28 normalized = empty_state()29 if not isinstance(state, dict):30 return normalized31 for key in normalized:32 if key in state:33 normalized[key] = state[key]34 normalized["version"] = STATE_VERSION35 normalized["candidate_subjects"] = list(normalized.get("candidate_subjects") or [])36 normalized["clarification_options"] = list(normalized.get("clarification_options") or [])37 normalized["subject_history"] = unique_subjects(normalized.get("subject_history") or [])38 if normalized["mode"] == "focused" and not normalized.get("active_subject"):39 normalized["mode"] = "idle"40 if normalized["mode"] == "comparing" and len(normalized["candidate_subjects"]) < 2:41 normalized["mode"] = "focused" if normalized.get("active_subject") else "idle"42 if normalized["mode"] == "awaiting_clarification" and not normalized.get("pending_query"):43 normalized["mode"] = "focused" if normalized.get("active_subject") else "idle"44 return normalized45 46 47def subject_key(subject: dict) -> str:48 return str(subject.get("unit_id") or subject.get("name") or "").strip().lower()49 50 51def unique_subjects(subjects: Iterable[dict]) -> list[dict]:52 result: list[dict] = []53 seen: set[str] = set()54 for subject in subjects:55 key = subject_key(subject)56 if not key or key in seen:57 continue58 seen.add(key)59 result.append(subject)60 return result61 62 63class ConversationStateMachine:64 """Resolve discourse references without retrieving or generating an answer."""65 66 def __init__(67 self,68 rewrite_callable: Optional[Callable[[str, dict], str]] = None,69 ) -> None:70 self.rewrite_callable = rewrite_callable71 72 PERSON_MARKERS = (73 "person", "who is", "background", "biography", "bio", "degree",74 "university", "role", "title", "expertise", "research", "career",75 )76 PROJECT_MARKERS = (77 "project", "initiative", "program", "launched",78 "participants", "benefit", "purpose", "goal", "caused it",79 "what does it do", "what is it about", "who leads it",80 )81 CONTEXT_MARKERS = re.compile(82 r"\b(it|its|they|them|their|he|him|his|she|her|hers|this|that|those|these)\b",83 re.IGNORECASE,84 )85 ELLIPSIS_MARKERS = re.compile(86 r"^(and\s+)?(what|which|who|when|where|why|how|does|did|is|are|was|were|can|could)\b",87 re.IGNORECASE,88 )89 CONTINUATION_MARKERS = re.compile(90 r"^\s*(and\b|also\b|actually\b|instead\b|going back\b|back to\b|what about\b|how about\b|"91 r"(?:no[, ]+)?i mean(?:t)?\b|rather\b|same (?:question|thing)\b)",92 re.IGNORECASE,93 )94 95 def classify_intent(self, message: str) -> str:96 lowered = message.lower()97 intent_terms = (98 ("cause", ("what caused", "caused", "why was", "why did", "motivat", "launched in response")),99 ("time", ("what year", "when", "how long", "timeframe")),100 ("funding", ("fund", "grant", "supported by", "sponsor")),101 ("leadership", ("who leads", "leader", "director", "chair")),102 ("education", ("degree", "university", "college", "studying", "candidate")),103 ("background", ("background", "biography", "bio", "career")),104 ("role", ("role", "position", "title", "what does she do", "what does he do")),105 ("research", ("research", "expertise", "focus", "work on")),106 ("audience", ("who is it for", "who does it serve", "participants", "audience")),107 ("count", ("how many", "count", "number of")),108 ("summary", ("tell me more", "what does it do", "what is it about", "overview")),109 )110 for intent, markers in intent_terms:111 if any(marker in lowered for marker in markers):112 return intent113 return "fact"114 115 def expected_subject_type(self, message: str, prior: dict) -> Optional[str]:116 lowered = message.lower()117 active = prior.get("active_subject") or {}118 if active.get("subject_type") == "publication" and any(119 marker in lowered for marker in ("title", "author", "publication", "paper", "study")120 ):121 return "publication"122 if re.search(r"\b(she|her|hers|he|him|his|that person|this person)\b", lowered):123 return "person"124 if any(marker in lowered for marker in self.PROJECT_MARKERS):125 return "project"126 if any(marker in lowered for marker in self.PERSON_MARKERS):127 return "person"128 if self.CONTEXT_MARKERS.search(message) or self.ELLIPSIS_MARKERS.search(message.strip()):129 return active.get("subject_type") or None130 return None131 132 def needs_context(self, message: str, prior: dict) -> bool:133 if self.CONTEXT_MARKERS.search(message):134 return True135 if self.CONTINUATION_MARKERS.search(message):136 return True137 if re.search(138 r"\b(?:[A-Z][\w'-]*\s+){1,5}(?:program|project|initiative|study|grant|funding)\b",139 message,140 ):141 return False142 if re.search(r"\b(the|that|this)\s+(project|initiative|program|person)(?:'s|\b)", message, re.IGNORECASE):143 return True144 if re.search(145 r"\b(students?|interns?|alumni|fellows?|staff|board members?|projects?|publications?|affiliates?)\b",146 message,147 re.IGNORECASE,148 ):149 return False150 if prior.get("mode") in {"focused", "comparing", "awaiting_clarification", "scoped"}:151 stripped = message.strip()152 return len(stripped.split()) <= 10 and bool(self.ELLIPSIS_MARKERS.search(stripped))153 return False154 155 @staticmethod156 def compatible(subject: dict, expected_type: Optional[str]) -> bool:157 return not expected_type or subject.get("subject_type") == expected_type158 159 def select_clarification_subject(self, message: str, candidates: list[dict]) -> Optional[dict]:160 lowered = re.sub(r"[^a-z0-9\s]", " ", message.lower())161 if re.search(r"\b(first|1st|former)\b", lowered) and candidates:162 return candidates[0]163 if re.search(r"\b(second|2nd|two|last|latter)\b", lowered) and len(candidates) > 1:164 return candidates[1]165 matches: list[dict] = []166 for subject in candidates:167 name = str(subject.get("name", "")).lower()168 tokens = [token for token in re.findall(r"[a-z0-9]+", name) if len(token) > 2]169 if name and (name in message.lower() or (tokens and all(token in lowered for token in tokens[-2:]))):170 matches.append(subject)171 return matches[0] if len(matches) == 1 else None172 173 def rewrite(self, message: str, subject: dict) -> str:174 if not self.rewrite_callable:175 return message176 rewritten = self.rewrite_callable(message, subject)177 return str(rewritten).strip() or message178 179 def clarify(self, message: str, candidates: list[dict], expected_type: Optional[str], prior: Optional[dict] = None) -> dict:180 candidates = unique_subjects(candidates)[:4]181 label = "project" if expected_type == "project" else "person" if expected_type == "person" else "subject"182 options = [f"{subject.get('name')} ({subject.get('subject_type', 'subject')})" for subject in candidates]183 state = normalize_state(prior)184 history = unique_subjects(185 list(state.get("subject_history") or [])186 + ([state["active_subject"]] if state.get("active_subject") else [])187 + list(state.get("candidate_subjects") or [])188 )189 state.update({190 "mode": "awaiting_clarification",191 "candidate_subjects": candidates,192 "pending_query": message,193 "last_intent": self.classify_intent(message),194 "clarification_options": options,195 "subject_history": history,196 })197 return {198 "resolved": False,199 "needs_clarification": True,200 "clarifying_question": f"Which {label} are you asking about?",201 "clarification_options": options,202 "state": state,203 }204 205 def resolve(self, message: str, prior_state: Optional[dict], explicit_subjects: list[dict]) -> dict:206 prior = normalize_state(prior_state)207 explicit = unique_subjects(explicit_subjects)208 intent = self.classify_intent(message)209 210 if prior.get("mode") == "awaiting_clarification":211 candidate_keys = {subject_key(subject) for subject in prior["candidate_subjects"]}212 explicit_choice = next(213 (subject for subject in explicit if subject_key(subject) in candidate_keys),214 None,215 )216 selected = explicit_choice or self.select_clarification_subject(message, prior["candidate_subjects"])217 if not selected and len(explicit) == 1 and self.CONTINUATION_MARKERS.search(message):218 selected = explicit[0]219 if selected:220 pending = str(prior.get("pending_query") or "Tell me more.")221 pending_intent = str(prior.get("last_intent") or self.classify_intent(pending))222 return self._resolved(223 self.rewrite(pending, selected), selected, pending_intent,224 used_context=True, prior=prior,225 )226 227 comparison_candidates = unique_subjects(prior.get("candidate_subjects") or [])228 if prior.get("mode") == "comparing" and len(comparison_candidates) > 1:229 lowered = message.lower()230 if re.search(r"\b(former|first one|first project|first person)\b", lowered):231 selected = comparison_candidates[0]232 return self._resolved(self.rewrite(message, selected), selected, intent, used_context=True, prior=prior)233 if re.search(r"\b(latter|second one|second project|second person)\b", lowered):234 selected = comparison_candidates[1]235 return self._resolved(self.rewrite(message, selected), selected, intent, used_context=True, prior=prior)236 collective = bool(re.search(r"\b(they|them|their|both|each|common|differ|difference|compare)\b", lowered))237 comparative_selection = bool(re.search(238 r"\bwhich(?:\s+(?:one|project|initiative|person))?\b.*\b"239 r"(first|earlier|later|newer|older|more|less|most|least|larger|smaller|broader)\b",240 lowered,241 ))242 if collective or comparative_selection:243 names = " and ".join(str(subject.get("name", "")) for subject in comparison_candidates)244 state = normalize_state(prior)245 state.update({"last_intent": intent, "last_query": message, "pending_query": None})246 rewritten = self.rewrite(message, {247 "name": names,248 "subject_type": "comparison",249 })250 return {251 "resolved": False, "needs_clarification": False,252 "rewritten_query": rewritten,253 "comparison_context": True, "used_context": True, "state": state,254 }255 256 if len(explicit) > 1:257 active = prior.get("active_subject") or {}258 correction = bool(re.match(r"^\s*(?:no\b|not\b|actually\b|i mean(?:t)?\b|rather\b|instead\b)", message, re.IGNORECASE))259 if active and correction and "compare" not in message.lower():260 selected = next(261 (subject for subject in reversed(explicit) if subject_key(subject) != subject_key(active)),262 None,263 )264 if selected:265 previous = str(prior.get("last_query") or message)266 rewritten = self.rewrite(previous, selected)267 return self._resolved(268 rewritten, selected,269 str(prior.get("last_intent") or intent), used_context=True, prior=prior,270 )271 state = normalize_state(prior)272 history = unique_subjects(273 list(state.get("subject_history") or [])274 + ([state["active_subject"]] if state.get("active_subject") else [])275 + explicit276 )[-12:]277 state.update({278 "mode": "comparing", "active_subject": None,279 "candidate_subjects": explicit, "pending_query": None,280 "last_intent": intent, "clarification_options": [],281 "subject_history": history, "last_query": message,282 })283 return {"resolved": False, "needs_clarification": False, "rewritten_query": message, "state": state}284 285 if len(explicit) == 1:286 subject = explicit[0]287 carry_previous = bool(288 prior.get("last_query")289 and self.CONTINUATION_MARKERS.search(message)290 and not re.search(r"\b(what|which|who|when|where|why|how|does|did|is|are|was|were|can|could)\b", message, re.IGNORECASE)291 )292 if carry_previous:293 previous = str(prior.get("last_query"))294 rewritten = self.rewrite(previous, subject)295 return self._resolved(296 rewritten, subject, str(prior.get("last_intent") or intent),297 used_context=True, prior=prior,298 )299 subject_name = str(subject.get("name", "")).strip()300 rewritten = (301 message302 if subject_name and subject_name.lower() in message.lower()303 else self.rewrite(message, subject)304 )305 return self._resolved(rewritten, subject, intent, used_context=False, prior=prior)306 307 if not self.needs_context(message, prior):308 state = normalize_state(prior)309 state["last_intent"] = intent310 state["last_query"] = message311 return {"resolved": False, "needs_clarification": False, "rewritten_query": message, "state": state}312 313 active_scope = prior.get("active_scope")314 if active_scope and not prior.get("active_subject") and not prior.get("candidate_subjects"):315 scope_name = str(active_scope.get("name", "")).strip()316 filter_text = str(active_scope.get("filter_text", "")).strip()317 rewritten = message318 if scope_name and self.rewrite_callable:319 rewritten = self.rewrite(message, {320 "name": scope_name,321 "subject_type": "scope",322 "filter_text": filter_text,323 })324 state = normalize_state(prior)325 state["last_intent"] = intent326 return {327 "resolved": False, "needs_clarification": False,328 "rewritten_query": rewritten, "scope_context": True, "state": state,329 }330 331 expected_type = self.expected_subject_type(message, prior)332 candidates = unique_subjects(prior.get("candidate_subjects") or [])333 active = prior.get("active_subject")334 compatible = [subject for subject in candidates if self.compatible(subject, expected_type)]335 historical = [336 subject for subject in reversed(prior.get("subject_history") or [])337 if self.compatible(subject, expected_type)338 ]339 340 if active and re.search(r"\b(?:the\s+)?(?:other|previous)\s+(?:one|project|initiative|person)\b", message, re.IGNORECASE):341 alternative = next(342 (subject for subject in historical if subject_key(subject) != subject_key(active)),343 None,344 )345 if alternative:346 return self._resolved(347 self.rewrite(message, alternative), alternative, intent,348 used_context=True, prior=prior,349 )350 351 if active and self.compatible(active, expected_type):352 return self._resolved(self.rewrite(message, active), active, intent, used_context=True, prior=prior)353 if len(compatible) == 1:354 return self._resolved(self.rewrite(message, compatible[0]), compatible[0], intent, used_context=True, prior=prior)355 if len(compatible) > 1:356 return self.clarify(message, compatible, expected_type, prior)357 if historical:358 return self._resolved(self.rewrite(message, historical[0]), historical[0], intent, used_context=True, prior=prior)359 if active:360 return self.clarify(message, [], expected_type, prior)361 if candidates:362 return self.clarify(message, candidates, expected_type, prior)363 return {"resolved": False, "needs_clarification": False, "rewritten_query": message, "state": prior}364 365 def _resolved(366 self, rewritten: str, subject: dict, intent: str, *, used_context: bool,367 prior: Optional[dict] = None,368 ) -> dict:369 state = normalize_state(prior)370 history = unique_subjects(371 list(state.get("subject_history") or [])372 + ([state["active_subject"]] if state.get("active_subject") else [])373 + list(state.get("candidate_subjects") or [])374 + [subject]375 )[-12:]376 state.update({377 "mode": "focused", "active_subject": subject,378 "candidate_subjects": [subject], "last_intent": intent,379 "pending_query": None, "clarification_options": [],380 "subject_history": history, "last_query": rewritten,381 })382 return {383 "resolved": True, "needs_clarification": False,384 "rewritten_query": rewritten, "active_subject": subject,385 "intent": intent, "used_context": used_context, "state": state,386 }387