Coverage for astrocyte/providers/composite.py: 100%

1"""Composite LLMProvider — split completion and embedding across two backends. 

2 

3Delegates ``complete()`` to one provider and ``embed()`` to another. The 

4motivating composition is Claude-CLI completions (subscription auth, no 

5embeddings surface) + local sentence-transformers embeddings:: 

6 

7 from astrocyte.providers.claude_cli import ClaudeCliProvider 

8 from astrocyte.providers.composite import CompositeLLMProvider 

9 from astrocyte.providers.local_embeddings import LocalEmbeddingsProvider 

10 

11 provider = CompositeLLMProvider( 

12 completion_provider=ClaudeCliProvider(model="haiku"), 

13 embedding_provider=LocalEmbeddingsProvider(), 

14 ) 

15 

16Any pair of LLMProvider implementations composes the same way. 

17""" 

18 

19from __future__ import annotations 

20 

21from typing import Any, ClassVar 

22 

23from astrocyte.types import Completion, LLMCapabilities, Message, ToolDefinition 

24 

25 

26class CompositeLLMProvider: 

27 """LLMProvider that routes complete() and embed() to different backends.""" 

28 

29 SPI_VERSION: ClassVar[int] = 1 

30 

31 def __init__( 

32 self, 

33 *, 

34 completion_provider: Any, 

35 embedding_provider: Any, 

36 ) -> None: 

37 self._completion = completion_provider 

38 self._embedding = embedding_provider 

39 

40 def capabilities(self) -> LLMCapabilities: 

41 comp = self._completion.capabilities() 

42 emb = self._embedding.capabilities() 

43 return LLMCapabilities( 

44 supports_multimodal_completion=comp.supports_multimodal_completion, 

45 modalities_supported=comp.modalities_supported, 

46 supports_multimodal_embedding=emb.supports_multimodal_embedding, 

47 supports_batch_embed=emb.supports_batch_embed, 

48 ) 

49 

50 async def complete( 

51 self, 

52 messages: list[Message], 

53 model: str | None = None, 

54 max_tokens: int = 1024, 

55 temperature: float = 0.0, 

56 tools: list[ToolDefinition] | None = None, 

57 tool_choice: str | None = None, 

58 response_format: dict | None = None, 

59 ) -> Completion: 

60 return await self._completion.complete( 

61 messages, 

62 model=model, 

63 max_tokens=max_tokens, 

64 temperature=temperature, 

65 tools=tools, 

66 tool_choice=tool_choice, 

67 response_format=response_format, 

68 ) 

69 

70 async def embed( 

71 self, 

72 texts: list[str], 

73 model: str | None = None, 

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

75 return await self._embedding.embed(texts, model=model) 

76 

77 async def close(self) -> None: 

78 for backend in (self._completion, self._embedding): 

79 close = getattr(backend, "close", None) 

80 if close is not None: 

81 try: 

82 await close() 

83 except Exception: # noqa: BLE001 — best-effort teardown 

84 pass