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

1"""Cohere provider for enterprise-grade NLP with advanced embeddings and RAG. 

2 

3Cohere specializes in enterprise NLP with best-in-class embeddings, reranking, 

4and Retrieval-Augmented Generation (RAG) optimized models. 

5 

6Supported Models: 

7- command-r-plus: Most capable for RAG and complex tasks 

8- command-r: Balanced RAG-optimized model 

9- command: General purpose completion 

10- command-light: Fast, lightweight completion 

11 

12Embeddings: 

13- embed-english-v3.0: 1024-dim English embeddings 

14- embed-multilingual-v3.0: 1024-dim multilingual embeddings 

15- embed-english-light-v3.0: 384-dim lightweight embeddings 

16 

17Features: 

18- Best-in-class embeddings for semantic search 

19- Reranking for improved search results 

20- RAG-optimized models with grounded generation 

21- Multilingual support (100+ languages) 

22- Enterprise SLAs and support 

23 

24Example: 

25 >>> from lexigram.ai.llm import CohereClient 

26 >>> 

27 >>> async with CohereClient(api_key="...") as client: 

28 ... # Chat completion 

29 ... response = await client.complete( 

30 ... model="command-r-plus", 

31 ... messages=[{"role": "user", "content": "Explain RAG"}] 

32 ... ) 

33 ... 

34 ... # Embeddings 

35 ... embeddings = await client.embed( 

36 ... texts=["Document 1", "Document 2"] 

37 ... ) 

38 ... 

39 ... # Reranking 

40 ... ranked = await client.rerank( 

41 ... query="What is AI?", 

42 ... documents=["Doc about AI", "Unrelated doc"] 

43 ... ) 

44 

45API Documentation: https://docs.cohere.com 

46 

47""" 

48 

49from __future__ import annotations 

50 

51import asyncio 

52from collections.abc import AsyncIterator 

53from typing import Any, cast 

54 

55import aiohttp 

56 

57from lexigram.ai.llm.clients._cohere_mappers import ( 

58 COHERE_MODELS, 

59 build_cohere_payload, 

60) 

61from lexigram.ai.llm.clients.base import AbstractLLMClient 

62from lexigram.ai.llm.config import ClientConfig 

63from lexigram.ai.llm.exceptions import ( 

64 LLMAuthenticationError, 

65 LLMError, 

66 LLMModelNotFoundError, 

67 LLMQuotaExceededError, 

68 LLMRateLimitError, 

69) 

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

71from lexigram.ai.llm.types import ( 

72 AIError, 

73 ChatMessage, 

74 Completion, 

75 FunctionCall, 

76 StreamChunk, 

77 TokenUsage, 

78 ToolCall, 

79) 

80from lexigram.contracts.core.health import HealthCheckResult, HealthStatus 

81from lexigram.contracts.web.http_models import HttpStatusError 

82from lexigram.logging import ( 

83 get_logger, 

84) 

85from lexigram.result import Err, Ok, Result 

86from lexigram.serialization import loads 

87from lexigram.validation import SecretStr 

88 

89logger = get_logger(__name__) 

90 

91__all__ = ["COHERE_MODELS", "CohereClient"] 

92 

93 

94class CohereClient(AbstractLLMClient): 

95 """Client for Cohere's enterprise NLP API. 

96 

97 Conforms to: :class:`~lexigram.contracts.ai.LLMClientProtocol` protocol via structural typing. 

98 

99 Supports Chat, Embeddings, and Reranking with: 

100 - RAG-optimized models (Command R/R+) 

101 - High-performance embeddings 

102 - Native reranking support 

103 """ 

104 

105 def __init__(self, config: ClientConfig) -> None: 

106 """Initialize Cohere client. 

107 

108 Args: 

109 config: LLM configuration 

110 """ 

111 super().__init__(config=config) 

112 self._client: ResilientHTTPClient | None = None 

113 

114 @property 

115 def api_key(self) -> SecretStr: 

116 """Get API key from config.""" 

117 return self.config.api_key or SecretStr("") 

118 

119 @property 

120 def base_url(self) -> str: 

121 """Get base URL from config.""" 

122 return self.config.api_base or "https://api.cohere.ai/v1" 

123 

124 def _get_client(self) -> ResilientHTTPClient: 

125 """Get or create HTTP client. 

126 

127 Returns: 

128 HTTP client instance. 

129 """ 

130 if self._client is None: 

131 self._client = ResilientHTTPClient( 

132 base_url=self.base_url, 

133 headers={ 

134 "Authorization": f"Bearer {self.api_key.get_secret_value()}", 

135 "Content-Type": "application/json", 

136 }, 

137 timeout=self.config.timeout, 

138 name="cohere-client", 

139 ) 

140 return self._client 

141 

142 def _build_payload( 

143 self, 

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

145 stream: bool, 

146 kwargs: dict[str, Any], 

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

148 """Build the Cohere API request payload from messages.""" 

149 return build_cohere_payload( 

150 client=self._get_client(), 

151 messages=messages, 

152 stream=stream, 

153 kwargs=kwargs, 

154 default_model=self.config.model, 

155 logger=logger, 

156 ) 

157 

158 async def _do_complete( 

159 self, 

160 messages: list[ChatMessage], 

161 **kwargs: Any, 

162 ) -> Result[Completion, LLMError]: 

163 """Generate non-streaming chat completion.""" 

164 try: 

165 client, payload, model = self._build_payload( 

166 messages, stream=False, kwargs=kwargs 

167 ) 

168 

169 # Non-streaming request 

170 response = await client.post("/chat", json=payload) 

171 response.raise_for_status() 

172 data = response.json() 

173 

174 # Parse response 

175 content = data.get("text", "") 

176 tool_calls = None 

177 

178 # Handle tool calls (normalize to our ToolCall schema) 

179 if "tool_calls" in data: 

180 tool_calls = [] 

181 for t in data["tool_calls"]: 

182 # Cohere may provide name/args or id/name/arguments 

183 name = t.get("name") or "" 

184 args = t.get("args") or t.get("arguments") or {} 

185 tid = t.get("id") or name or "" 

186 tool_calls.append( 

187 ToolCall( 

188 id=tid, 

189 type="function", 

190 function=FunctionCall(name=name, arguments=args), 

191 ), 

192 ) 

193 

194 # Handle citations (for RAG) 

195 citations = data.get("citations", []) 

196 

197 usage = {} 

198 if "meta" in data and "tokens" in data["meta"]: 

199 tokens = data["meta"]["tokens"] 

200 usage = { 

201 "prompt_tokens": tokens.get("input_tokens", 0), 

202 "completion_tokens": tokens.get("output_tokens", 0), 

203 "total_tokens": tokens.get("input_tokens", 0) 

204 + tokens.get("output_tokens", 0), 

205 } 

206 

207 return Ok( 

208 Completion( 

209 content=content, 

210 model=model, 

211 finish_reason=data.get("finish_reason", "stop"), 

212 usage=TokenUsage( 

213 prompt_tokens=usage.get("prompt_tokens", 0), 

214 completion_tokens=usage.get("completion_tokens", 0), 

215 total_tokens=usage.get("total_tokens", 0), 

216 ), 

217 tool_calls=tool_calls, 

218 metadata={"citations": citations} if citations else {}, 

219 ) 

220 ) 

221 

222 except (OSError, ConnectionError, TimeoutError, RuntimeError, ValueError) as e: 

223 return self._handle_error_as_result(e) 

224 

225 async def _do_stream_chat( 

226 self, 

227 messages: list[ChatMessage], 

228 **kwargs: Any, 

229 ) -> Result[AsyncIterator[StreamChunk], LLMError]: 

230 """Start a streaming completion.""" 

231 try: 

232 client, payload, _ = self._build_payload( 

233 messages, stream=True, kwargs=kwargs 

234 ) 

235 return Ok(self._stream_completion(client, payload)) 

236 except (OSError, ConnectionError, TimeoutError, RuntimeError, ValueError) as e: 

237 return self._handle_error_as_result(e) 

238 

239 async def _do_chat( 

240 self, 

241 messages: list[ChatMessage], 

242 tools: list[Any] | None = None, 

243 **kwargs: Any, 

244 ) -> Result[Completion, LLMError]: 

245 """Generate completion with optional tool/function calling. 

246 

247 Tool conversion happens in :func:`~lexigram.ai.llm.clients._cohere_mappers.build_cohere_payload` 

248 (via ``complete(..., tools=...)``); this method forwards the tool 

249 descriptors to keep the ``chat`` code path consistent. 

250 

251 Args: 

252 messages: Chat messages. 

253 tools: Optional tool definitions. 

254 **kwargs: Additional parameters forwarded to the provider. 

255 

256 Returns: 

257 ``Ok(Completion)`` on success or ``Err(LLMError)`` on failure. 

258 """ 

259 return await self._do_complete(messages, tools=tools, **kwargs) 

260 

261 async def _stream_completion( 

262 self, 

263 client: ResilientHTTPClient, 

264 payload: dict[str, Any], 

265 ) -> AsyncIterator[StreamChunk]: 

266 """Stream chat completion. 

267 

268 Args: 

269 client: HTTP client. 

270 payload: Request payload. 

271 

272 Yields: 

273 StreamChunk objects. 

274 

275 """ 

276 # Support clients that return a coroutine instead of an async context manager 

277 stream_ctx = client.stream("POST", "/chat", json=payload) 

278 

279 try: 

280 if asyncio.iscoroutine(stream_ctx): 

281 stream_ctx = await stream_ctx 

282 except aiohttp.ClientResponseError as e: 

283 # We want to catch specific HTTP errors during setup and convert them 

284 raise AIError(f"Cohere error: {e}") from e 

285 except (OSError, ValueError, RuntimeError) as e: 

286 raise AIError(f"Cohere error: {e}") from e 

287 

288 try: 

289 async with stream_ctx as response: 

290 response.raise_for_status() 

291 

292 iterator = response.aiter_lines() 

293 if asyncio.iscoroutine(iterator): 

294 iterator = await iterator 

295 

296 async for line in iterator: 

297 if not line: 

298 continue 

299 

300 try: 

301 data = loads(line) 

302 

303 # Cohere sends different event types 

304 event_type = data.get("event_type") 

305 

306 if event_type == "text-generation": 

307 content = data.get("text", "") 

308 if content: 

309 yield StreamChunk( 

310 delta=content, 

311 model=payload["model"], 

312 finish_reason=None, 

313 ) 

314 

315 elif event_type == "stream-end": 

316 # Final chunk with finish reason 

317 yield StreamChunk( 

318 delta="", 

319 model=payload["model"], 

320 finish_reason=data.get("finish_reason", "stop"), 

321 ) 

322 

323 except (ValueError, TypeError) as e: 

324 logger.debug("Error parsing stream chunk: %s", e) 

325 continue 

326 except (aiohttp.ClientError, OSError, ValueError, RuntimeError) as e: 

327 raise AIError(f"Cohere streaming error: {e}") from e 

328 

329 async def embed( 

330 self, 

331 texts: list[str] | str, 

332 model: str = "embed-english-v3.0", 

333 input_type: str = "search_document", 

334 **kwargs: Any, 

335 ) -> list[list[float]]: 

336 """Generate embeddings. 

337 

338 Args: 

339 texts: Text or list of texts to embed. 

340 model: Model ID (default: "embed-english-v3.0"). 

341 input_type: Type of input ("search_document", "search_query", "classification", "clustering"). 

342 **kwargs: Additional parameters. 

343 

344 Returns: 

345 List of embedding vectors. 

346 

347 Example: 

348 >>> # Embed documents 

349 >>> doc_embeddings = await client.embed( 

350 ... texts=["Doc 1", "Doc 2"], 

351 ... input_type="search_document" 

352 ... ) 

353 >>> 

354 >>> # Embed query 

355 >>> query_embedding = await client.embed( 

356 ... texts="What is AI?", 

357 ... input_type="search_query" 

358 ... ) 

359 

360 """ 

361 try: 

362 client = self._get_client() 

363 

364 # Ensure texts is list 

365 if isinstance(texts, str): 

366 texts = [texts] 

367 

368 payload = { 

369 "texts": texts, 

370 "model": model, 

371 "input_type": input_type, 

372 **kwargs, 

373 } 

374 

375 response = await client.post("/embed", json=payload) 

376 response.raise_for_status() 

377 data = response.json() 

378 

379 return cast("list[list[float]]", data["embeddings"]) 

380 except ( 

381 HttpStatusError, 

382 aiohttp.ClientError, 

383 OSError, 

384 KeyError, 

385 ValueError, 

386 ) as e: 

387 raise AIError(f"Cohere error: {e}") from e 

388 

389 async def rerank( 

390 self, 

391 query: str, 

392 documents: list[str] | list[dict[str, str]], 

393 model: str = "rerank-english-v3.0", 

394 top_n: int | None = None, 

395 **kwargs: Any, 

396 ) -> list[dict[str, Any]]: 

397 """Rerank documents for a query. 

398 

399 Args: 

400 query: Search query. 

401 documents: List of documents (strings or dicts with 'text' key). 

402 model: Reranking model (default: "rerank-english-v3.0"). 

403 top_n: Return top N results (default: all). 

404 **kwargs: Additional parameters. 

405 

406 Returns: 

407 List of ranked documents with scores. 

408 

409 Example: 

410 >>> results = await client.rerank( 

411 ... query="What is machine learning?", 

412 ... documents=[ 

413 ... "ML is a subset of AI...", 

414 ... "Unrelated document...", 

415 ... "Deep learning uses neural networks..." 

416 ... ], 

417 ... top_n=2 

418 ... ) 

419 >>> for result in results: 

420 ... print(f"Score: {result['relevance_score']:.3f} - {result['document']['text']}") 

421 

422 """ 

423 try: 

424 client = self._get_client() 

425 

426 payload = { 

427 "query": query, 

428 "documents": documents, 

429 "model": model, 

430 **kwargs, 

431 } 

432 

433 if top_n: 

434 payload["top_n"] = top_n 

435 

436 response = await client.post("/rerank", json=payload) 

437 response.raise_for_status() 

438 data = response.json() 

439 

440 return cast("list[dict[str, Any]]", data.get("results", [])) 

441 except ( 

442 HttpStatusError, 

443 aiohttp.ClientError, 

444 OSError, 

445 KeyError, 

446 ValueError, 

447 ) as e: 

448 raise AIError(f"Cohere error: {e}") from e 

449 

450 def _handle_error_as_result(self, error: Exception) -> Result[Any, LLMError]: 

451 """Map a caught exception to ``Err`` (recoverable) or re-raise (infra).""" 

452 if isinstance(error, aiohttp.ClientResponseError): 

453 if error.status == 401: 

454 raise LLMAuthenticationError( 

455 f"Cohere authentication failed: {error}" 

456 ) from error 

457 if error.status == 429: 

458 return Err(LLMRateLimitError(f"Cohere rate limit exceeded: {error}")) 

459 if error.status == 402: 

460 return Err(LLMQuotaExceededError(f"Cohere quota exceeded: {error}")) 

461 if error.status == 404: 

462 return Err(LLMModelNotFoundError(f"Cohere model not found: {error}")) 

463 raise AIError(f"Cohere infrastructure error: {error}") from error 

464 

465 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

466 """Perform a lightweight health check against the provider.""" 

467 try: 

468 client = self._get_client() 

469 # Simplest probe: list models 

470 response = await client.get("/models") 

471 response.raise_for_status() 

472 return HealthCheckResult(component="cohere", status=HealthStatus.HEALTHY) 

473 except (OSError, ConnectionError, TimeoutError, RuntimeError) as exc: 

474 return HealthCheckResult( 

475 component="cohere", 

476 status=HealthStatus.UNHEALTHY, 

477 error=str(exc), 

478 ) 

479 

480 async def close(self) -> None: 

481 """Close the HTTP client.""" 

482 if self._client: 

483 await self._client.close() 

484 self._client = None 

485 await super().close()