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

1"""Calling a language model. 

2 

3Every call goes through :func:`complete`. Real models are reached through 

4LiteLLM, so a deployment can use any model and endpoint LiteLLM supports. 

5The model name ``mock`` answers locally (see :func:`_mock_completion`), so 

6the platform runs and can be tested without any model. 

7""" 

8 

9import json 

10import os 

11import time 

12from dataclasses import dataclass 

13from dataclasses import field 

14 

15from django.conf import settings 

16 

17 

18class ModelError(Exception): 

19 """The model call failed. ``code`` is safe to show and store.""" 

20 

21 def __init__(self, message: str, *, code: str = "model_error"): 

22 super().__init__(message) 

23 self.code = code 

24 

25 

26@dataclass 

27class ModelConfig: 

28 model: str 

29 api_base: str = "" 

30 api_key: str = "" 

31 parameters: dict = field(default_factory=dict) 

32 

33 @classmethod 

34 def from_variant(cls, variant) -> ModelConfig: 

35 if variant is None: 

36 return cls(model=settings.CHAT_MODEL) 

37 api_key = os.environ.get(variant.api_key_env, "") if variant.api_key_env else "" 

38 return cls( 

39 model=variant.model, api_base=variant.api_base, api_key=api_key, parameters=dict(variant.parameters) 

40 ) 

41 

42 

43@dataclass 

44class ToolCall: 

45 id: str 

46 name: str 

47 arguments: dict 

48 

49 

50@dataclass 

51class Completion: 

52 content: str 

53 model: str 

54 tool_calls: list[ToolCall] = field(default_factory=list) 

55 usage: dict = field(default_factory=dict) 

56 latency_ms: int = 0 

57 # The assistant turn as the model sent it, for replaying tool calls. 

58 raw_message: dict = field(default_factory=dict) 

59 

60 

61def complete(config: ModelConfig, messages: list[dict], *, tools: list[dict] | None = None) -> Completion: 

62 started = time.monotonic() 

63 if config.model == "mock": 

64 result = _mock_completion(messages, tools=tools) 

65 else: 

66 result = _litellm_completion(config, messages, tools=tools) 

67 result.latency_ms = int((time.monotonic() - started) * 1000) 

68 return result 

69 

70 

71def _litellm_completion(config: ModelConfig, messages, *, tools): 

72 import litellm 

73 

74 kwargs = { 

75 **config.parameters, 

76 "model": config.model, 

77 "messages": messages, 

78 "timeout": settings.CHAT_MODEL_TIMEOUT, 

79 } 

80 if config.api_base: 

81 kwargs["api_base"] = config.api_base 

82 if config.api_key: 

83 kwargs["api_key"] = config.api_key 

84 if tools: 

85 kwargs["tools"] = tools 

86 try: 

87 response = litellm.completion(**kwargs) 

88 except Exception as exc: 

89 code = "timeout" if "timeout" in type(exc).__name__.lower() else "model_error" 

90 msg = f"Model call failed: {type(exc).__name__}" 

91 raise ModelError(msg, code=code) from exc 

92 try: 

93 message = response.choices[0].message 

94 except (AttributeError, IndexError) as exc: 

95 msg = "Model returned no choices." 

96 raise ModelError(msg, code="empty_response") from exc 

97 tool_calls = [ 

98 ToolCall(id=call.id, name=call.function.name, arguments=_json_arguments(call.function.arguments)) 

99 for call in (getattr(message, "tool_calls", None) or []) 

100 ] 

101 usage = getattr(response, "usage", None) 

102 return Completion( 

103 content=message.content or "", 

104 model=getattr(response, "model", None) or config.model, 

105 tool_calls=tool_calls, 

106 usage={ 

107 "prompt_tokens": getattr(usage, "prompt_tokens", None), 

108 "completion_tokens": getattr(usage, "completion_tokens", None), 

109 } 

110 if usage 

111 else {}, 

112 raw_message={ 

113 "role": "assistant", 

114 "content": message.content or "", 

115 **( 

116 { 

117 "tool_calls": [ 

118 { 

119 "id": c.id, 

120 "type": "function", 

121 "function": {"name": c.name, "arguments": json.dumps(c.arguments)}, 

122 } 

123 for c in tool_calls 

124 ] 

125 } 

126 if tool_calls 

127 else {} 

128 ), 

129 }, 

130 ) 

131 

132 

133def _json_arguments(arguments) -> dict: 

134 if isinstance(arguments, dict): 

135 return arguments 

136 try: 

137 value = json.loads(arguments or "{}") 

138 except json.JSONDecodeError: 

139 return {} 

140 return value if isinstance(value, dict) else {} 

141 

142 

143MOCK_FAIL = "[mock-fail]" 

144MOCK_TOOL_PREFIX = "[mock-tool " 

145 

146 

147def _mock_completion(messages, *, tools): 

148 """Deterministic replies for development and tests. 

149 

150 - A user message containing ``[mock-fail]`` makes the call fail. 

151 - ``[mock-tool name] text`` asks for tool ``name`` with ``{"query": text}``, 

152 if that tool is offered; the reply after the tool result quotes it. 

153 - Otherwise the reply echoes the last user message. 

154 """ 

155 last = messages[-1] 

156 if last["role"] == "tool": 

157 reply = f"Here is what I found: {last['content'][:500]}" 

158 return Completion(content=reply, model="mock", raw_message={"role": "assistant", "content": reply}) 

159 text = last.get("content") or "" 

160 if MOCK_FAIL in text: 

161 msg = "Mock model failure." 

162 raise ModelError(msg, code="mock_failure") 

163 if text.startswith(MOCK_TOOL_PREFIX) and "]" in text: 

164 name, _, query = text[len(MOCK_TOOL_PREFIX) :].partition("]") 

165 if any(tool["function"]["name"] == name for tool in tools or []): 

166 call = ToolCall(id="mock-call-1", name=name, arguments={"query": query.strip()}) 

167 raw = { 

168 "role": "assistant", 

169 "content": "", 

170 "tool_calls": [ 

171 { 

172 "id": call.id, 

173 "type": "function", 

174 "function": {"name": name, "arguments": json.dumps(call.arguments)}, 

175 } 

176 ], 

177 } 

178 return Completion(content="", model="mock", tool_calls=[call], raw_message=raw) 

179 reply = f"You said: {text}" 

180 return Completion(content=reply, model="mock", raw_message={"role": "assistant", "content": reply})