Coverage for astrocyte/mip/intent.py: 76%

1"""MIP intent layer — LLM-based routing when mechanical rules can't resolve. 

2 

3Async — requires LLM call. Follows the pattern of pipeline/curated_retain.py. 

4""" 

5 

6from __future__ import annotations 

7 

8import json 

9import logging 

10from typing import TYPE_CHECKING 

11 

12from astrocyte.mip.schema import BankDefinition, IntentPolicy 

13from astrocyte.types import Message, RoutingDecision 

14 

15if TYPE_CHECKING: 

16 from astrocyte.mip.rule_engine import RuleEngineInput 

17 from astrocyte.provider import LLMProvider 

18 

19logger = logging.getLogger("astrocyte.mip") 

20 

21 

22async def resolve_intent( 

23 input_data: RuleEngineInput, 

24 intent_policy: IntentPolicy, 

25 available_banks: list[BankDefinition], 

26 llm_provider: LLMProvider, 

27 *, 

28 model: str | None = None, 

29) -> RoutingDecision: 

30 """Ask the LLM to route content when mechanical rules cannot resolve. 

31 

32 System message contains routing instructions (trusted). 

33 User message wraps untrusted content in XML delimiters. 

34 Falls back to passthrough on failure. 

35 """ 

36 max_tokens = 200 

37 if intent_policy.constraints and "max_tokens" in intent_policy.constraints: 

38 max_tokens = int(intent_policy.constraints["max_tokens"]) 

39 

40 system_msg = _build_system_message(intent_policy, available_banks) 

41 user_msg = _build_user_message(input_data) 

42 

43 try: 

44 completion = await llm_provider.complete( 

45 messages=[ 

46 Message(role="system", content=system_msg), 

47 Message(role="user", content=user_msg), 

48 ], 

49 max_tokens=max_tokens, 

50 temperature=0, 

51 ) 

52 return _parse_intent_response(completion.text) 

53 except Exception: 

54 logger.warning("MIP intent layer failed, falling back to passthrough") 

55 return RoutingDecision(resolved_by="passthrough", reasoning="Intent layer LLM call failed") 

56 

57 

58def _build_system_message( 

59 intent_policy: IntentPolicy, 

60 available_banks: list[BankDefinition], 

61) -> str: 

62 """Build the system message with routing instructions (trusted content only).""" 

63 banks_str = ", ".join(b.id for b in available_banks) if available_banks else "(none defined)" 

64 

65 if intent_policy.model_context: 

66 base = intent_policy.model_context.replace("{banks}", banks_str) 

67 else: 

68 base = ( 

69 f"You are a memory routing agent. Route content to the correct bank and apply tags.\n" 

70 f"Available banks: {banks_str}\n" 

71 f"Never override compliance rules." 

72 ) 

73 

74 return f"""{base} 

75 

76The user message contains the content to route inside <content> XML tags, along with metadata. 

77Respond with a JSON object: 

78{{"bank_id": "...", "tags": ["..."], "retain_policy": "default", "reasoning": "..."}}""" 

79 

80 

81def _build_user_message(input_data: RuleEngineInput) -> str: 

82 """Build the user message with untrusted content wrapped in XML delimiters.""" 

83 tags_str = ", ".join(input_data.tags) if input_data.tags else "(none)" 

84 return f"""<content> 

85{input_data.content[:500]} 

86</content> 

87 

88Content type: {input_data.content_type or "text"} 

89Source: {input_data.source or "unknown"} 

90Tags: {tags_str} 

91PII detected: {input_data.pii_detected}""" 

92 

93 

94def _parse_intent_response(response: str) -> RoutingDecision: 

95 """Parse LLM JSON response into RoutingDecision. Graceful fallback.""" 

96 try: 

97 text = response.strip() 

98 # Extract from code block if present 

99 if "```" in text: 

100 start = text.index("```") + 3 

101 if text[start:].startswith("json"): 

102 start += 4 

103 end = text.index("```", start) 

104 text = text[start:end].strip() 

105 

106 data = json.loads(text) 

107 return RoutingDecision( 

108 bank_id=data.get("bank_id"), 

109 tags=data.get("tags"), 

110 retain_policy=data.get("retain_policy"), 

111 resolved_by="intent", 

112 confidence=data.get("confidence", 0.8), 

113 reasoning=data.get("reasoning"), 

114 ) 

115 except (json.JSONDecodeError, ValueError, KeyError): 

116 logger.warning("Failed to parse MIP intent response, falling back to passthrough") 

117 return RoutingDecision(resolved_by="passthrough", reasoning="Failed to parse intent response")