Coverage for src / documint_mcp / embeddings.py: 0%

1"""Voyage AI embeddings and reranking for RAG. 

2 

3Optional enhancement layer — everything degrades gracefully to FTS-only 

4if the voyageai package is not installed or the API key is not set. 

5""" 

6from __future__ import annotations 

7 

8import logging 

9import os 

10 

11logger = logging.getLogger(__name__) 

12 

13_HAS_VOYAGE = False 

14_client = None 

15 

16try: 

17 import voyageai # noqa: PLC0415 

18 

19 _HAS_VOYAGE = True 

20except ImportError: 

21 pass 

22 

23 

24def get_client(): # noqa: ANN201 

25 """Return a cached Voyage AI client, or None if unavailable.""" 

26 global _client # noqa: PLW0603 

27 if not _HAS_VOYAGE: 

28 return None 

29 if _client is None: 

30 key = os.getenv("VOYAGE_API_KEY", "") 

31 if not key: 

32 return None 

33 _client = voyageai.Client(api_key=key) 

34 return _client 

35 

36 

37def embed( 

38 texts: list[str], model: str = "voyage-3-lite" 

39) -> list[list[float]] | None: 

40 """Embed a list of texts using Voyage AI. 

41 

42 Returns None if Voyage is unavailable or the call fails. 

43 """ 

44 client = get_client() 

45 if not client: 

46 return None 

47 try: 

48 result = client.embed(texts, model=model) 

49 return result.embeddings 

50 except Exception as exc: # noqa: BLE001 

51 logger.debug("Voyage embed failed: %s", exc) 

52 return None 

53 

54 

55def rerank( 

56 query: str, 

57 documents: list[str], 

58 model: str = "rerank-2", 

59 top_k: int = 3, 

60) -> list[str] | None: 

61 """Rerank documents for a query using Voyage AI. 

62 

63 Returns reranked document strings, or None if Voyage is unavailable. 

64 """ 

65 client = get_client() 

66 if not client: 

67 return None 

68 try: 

69 result = client.rerank(query, documents, model=model, top_k=top_k) 

70 return [documents[r.index] for r in result.results] 

71 except Exception as exc: # noqa: BLE001 

72 logger.debug("Voyage rerank failed: %s", exc) 

73 return None