Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/clients/_cohere_mappers.py: 12%

1"""Mapping helpers for Cohere request payloads and tool schemas.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any 

6 

7from lexigram.ai.llm.clients._message_utils import serialize_text_only 

8from lexigram.ai.llm.clients._tools_utils import parse_json_arguments 

9from lexigram.ai.llm.http.client import ResilientHTTPClient 

10from lexigram.ai.llm.types import ChatMessage 

11 

12 

13def build_cohere_payload( 

14 *, 

15 client: ResilientHTTPClient, 

16 messages: list[ChatMessage] | list[dict[str, Any]], 

17 stream: bool, 

18 kwargs: dict[str, Any], 

19 default_model: str, 

20 logger: Any, 

21) -> tuple[ResilientHTTPClient, dict[str, Any], str]: 

22 """Build the Cohere API request payload from message history.""" 

23 request_kwargs = kwargs.copy() 

24 model = request_kwargs.pop("model", None) or default_model 

25 temperature = request_kwargs.pop("temperature", 0.7) 

26 max_tokens = request_kwargs.pop("max_tokens", None) 

27 tools = request_kwargs.pop("tools", None) 

28 documents = request_kwargs.pop("documents", None) 

29 

30 user_message = "" 

31 chat_history: list[dict[str, Any]] = [] 

32 preamble: str | None = None 

33 for msg in messages: 

34 if isinstance(msg, dict): 

35 role = msg.get("role", "user") 

36 content = msg.get("content", "") 

37 tool_call_id = msg.get("tool_call_id", "") 

38 tool_calls = msg.get("tool_calls", []) 

39 else: 

40 role = msg.role.value 

41 content = msg.content 

42 tool_call_id = msg.tool_call_id 

43 tool_calls = msg.tool_calls 

44 

45 text_content = serialize_text_only( 

46 content, 

47 logger=logger, 

48 client_name="cohere", 

49 ) 

50 if role == "system": 

51 preamble = text_content 

52 elif role == "tool": 

53 # Tool results are sent as a USER turn carrying tool_results 

54 chat_history.append( 

55 { 

56 "role": "USER", 

57 "message": text_content, 

58 "tool_results": [ 

59 { 

60 "call": {"name": tool_call_id or "", "parameters": {}}, 

61 "outputs": [{"tool_result": text_content}], 

62 } 

63 ], 

64 } 

65 ) 

66 elif role == "user": 

67 user_message = text_content 

68 elif role == "assistant": 

69 if chat_history and user_message: 

70 chat_history.append({"role": "USER", "message": user_message}) 

71 user_message = "" 

72 entry: dict[str, Any] = {"role": "CHATBOT", "message": text_content} 

73 serialized_calls = _cohere_assistant_tool_calls(tool_calls) 

74 if serialized_calls: 

75 entry["tool_calls"] = serialized_calls 

76 chat_history.append(entry) 

77 

78 payload: dict[str, Any] = { 

79 "model": model, 

80 "message": user_message, 

81 "temperature": temperature, 

82 "stream": stream, 

83 **request_kwargs, 

84 } 

85 if preamble is not None: 

86 payload["preamble"] = preamble 

87 if chat_history: 

88 payload["chat_history"] = chat_history 

89 if max_tokens: 

90 payload["max_tokens"] = max_tokens 

91 if tools: 

92 payload["tools"] = map_cohere_tools(tools) 

93 if documents: 

94 payload["documents"] = documents 

95 

96 return client, payload, model 

97 

98 

99def _cohere_assistant_tool_calls(tool_calls: Any) -> list[dict[str, Any]]: 

100 """Serialize assistant tool calls to Cohere ``tool_calls`` wire format. 

101 

102 Args: 

103 tool_calls: Framework tool calls from a prior assistant turn. 

104 

105 Returns: 

106 Cohere tool-call dicts (``{"name", "parameters"}``). 

107 """ 

108 serialized: list[dict[str, Any]] = [] 

109 for call in tool_calls or []: 

110 fn = getattr(call, "function", None) 

111 if fn is None or not getattr(fn, "name", None): 

112 continue 

113 serialized.append( 

114 { 

115 "name": fn.name, 

116 "parameters": parse_json_arguments(fn.arguments), 

117 } 

118 ) 

119 return serialized 

120 

121 

122def map_cohere_tools(tools: list[Any]) -> list[dict[str, Any]]: 

123 """Convert Lexigram tool definitions into Cohere tool schema format.""" 

124 cohere_tools: list[dict[str, Any]] = [] 

125 for tool in tools: 

126 if isinstance(tool, dict): 

127 cohere_tools.append(tool) 

128 continue 

129 fn = getattr(tool, "function", tool) 

130 name: str = getattr(fn, "name", "") or "" 

131 description: str = getattr(fn, "description", "") or "" 

132 schema: dict[str, Any] = getattr(fn, "parameters", {}) or {} 

133 properties: dict[str, Any] = schema.get("properties", {}) 

134 required_fields: list[str] = schema.get("required", []) 

135 param_defs: dict[str, Any] = {} 

136 for param_name, param_meta in properties.items(): 

137 param_defs[param_name] = { 

138 "description": param_meta.get("description", ""), 

139 "type": param_meta.get("type", "str"), 

140 "required": param_name in required_fields, 

141 } 

142 cohere_tools.append( 

143 { 

144 "name": name, 

145 "description": description, 

146 "parameter_definitions": param_defs, 

147 } 

148 ) 

149 return cohere_tools 

150 

151 

152COHERE_MODELS = { 

153 "command-r-plus": { 

154 "context_window": 128000, 

155 "supports_tools": True, 

156 "supports_rag": True, 

157 "description": "Most capable - best for RAG and complex tasks", 

158 }, 

159 "command-r": { 

160 "context_window": 128000, 

161 "supports_tools": True, 

162 "supports_rag": True, 

163 "description": "Balanced - RAG-optimized, cost-effective", 

164 }, 

165 "command": { 

166 "context_window": 4096, 

167 "supports_tools": False, 

168 "supports_rag": False, 

169 "description": "General purpose completion", 

170 }, 

171 "command-light": { 

172 "context_window": 4096, 

173 "supports_tools": False, 

174 "supports_rag": False, 

175 "description": "Fast, lightweight completion", 

176 }, 

177 "embed-english-v3.0": { 

178 "dimension": 1024, 

179 "description": "Best English embeddings", 

180 }, 

181 "embed-multilingual-v3.0": { 

182 "dimension": 1024, 

183 "description": "100+ languages supported", 

184 }, 

185 "embed-english-light-v3.0": { 

186 "dimension": 384, 

187 "description": "Lightweight English embeddings", 

188 }, 

189}