Coverage for src/dataknobs_llm/llm/providers/huggingface.py: 20%

1"""HuggingFace Inference API provider implementation."""

2

3 import os

4 from typing import TYPE_CHECKING, Any, Dict, List, Union, AsyncIterator

5

6 from ..base import (

7 LLMConfig, LLMMessage, LLMResponse, LLMStreamResponse,

8 AsyncLLMProvider, ModelCapability,

9 normalize_llm_config

10)

11 from dataknobs_llm.prompts import AsyncPromptBuilder

12

13 if TYPE_CHECKING:

14 from dataknobs_config.config import Config

15

16

17 class HuggingFaceProvider(AsyncLLMProvider):

18 """HuggingFace Inference API provider."""

19

20 def __init__(

21 self,

22 config: Union[LLMConfig, "Config", Dict[str, Any]],

23 prompt_builder: AsyncPromptBuilder | None = None

24 ):

25 # Normalize config first

26 llm_config = normalize_llm_config(config)

27 super().__init__(llm_config, prompt_builder=prompt_builder)

28 self.base_url = llm_config.api_base or 'https://api-inference.huggingface.co/models'

29

30 async def initialize(self) -> None:

31 """Initialize HuggingFace client."""

32 try:

33 import aiohttp

34

35 api_key = self.config.api_key or os.environ.get('HUGGINGFACE_API_KEY')

36 if not api_key:

37 raise ValueError("HuggingFace API key not provided")

38

39 self._session = aiohttp.ClientSession(

40 headers={'Authorization': f'Bearer {api_key}'},

41 timeout=aiohttp.ClientTimeout(total=self.config.timeout)

42 )

43 self._is_initialized = True

44 except ImportError as e:

45 raise ImportError("aiohttp package not installed. Install with: pip install aiohttp") from e

46

47 async def close(self) -> None:

48 """Close HuggingFace client."""

49 if hasattr(self, '_session') and self._session:

50 await self._session.close()

51 self._is_initialized = False

52

53 async def validate_model(self) -> bool:

54 """Validate model availability."""

55 try:

56 url = f"{self.base_url}/{self.config.model}"

57 async with self._session.get(url) as response:

58 return response.status == 200

59 except Exception:

60 return False

61

62 def get_capabilities(self) -> List[ModelCapability]:

63 """Get HuggingFace model capabilities."""

64 # Basic capabilities for text generation models

65 return [

66 ModelCapability.TEXT_GENERATION,

67 ModelCapability.EMBEDDINGS if 'embedding' in self.config.model else None # type: ignore

68 ]

69

70 async def complete(

71 self,

72 messages: Union[str, List[LLMMessage]],

73 **kwargs

74 ) -> LLMResponse:

75 """Generate completion."""

76 if not self._is_initialized:

77 await self.initialize()

78

79 # Convert to prompt

80 if isinstance(messages, str):

81 prompt = messages

82 else:

83 prompt = self._build_prompt(messages)

84

85 # Make API call

86 url = f"{self.base_url}/{self.config.model}"

87 payload = {

88 'inputs': prompt,

89 'parameters': {

90 'temperature': self.config.temperature,

91 'top_p': self.config.top_p,

92 'max_new_tokens': self.config.max_tokens or 100,

93 'return_full_text': False

94 }

95 }

96

97 async with self._session.post(url, json=payload) as response:

98 response.raise_for_status()

99 data = await response.json()

100

101 # Parse response

102 if isinstance(data, list) and len(data) > 0:

103 text = data[0].get('generated_text', '')

104 else:

105 text = str(data)

106

107 return LLMResponse(

108 content=text,

109 model=self.config.model,

110 finish_reason='stop'

111 )

112

113 async def stream_complete(

114 self,

115 messages: Union[str, List[LLMMessage]],

116 **kwargs

117 ) -> AsyncIterator[LLMStreamResponse]:

118 """HuggingFace Inference API doesn't support streaming."""

119 # Simulate streaming by yielding complete response

120 response = await self.complete(messages, **kwargs)

121 yield LLMStreamResponse(

122 delta=response.content,

123 is_final=True,

124 finish_reason=response.finish_reason

125 )

126

127 async def embed(

128 self,

129 texts: Union[str, List[str]],

130 **kwargs

131 ) -> Union[List[float], List[List[float]]]:

132 """Generate embeddings."""

133 if not self._is_initialized:

134 await self.initialize()

135

136 if isinstance(texts, str):

137 texts = [texts]

138 single = True

139 else:

140 single = False

141

142 url = f"{self.base_url}/{self.config.model}"

143 payload = {'inputs': texts}

144

145 async with self._session.post(url, json=payload) as response:

146 response.raise_for_status()

147 embeddings = await response.json()

148

149 return embeddings[0] if single else embeddings

150

151 async def function_call(

152 self,

153 messages: List[LLMMessage],

154 functions: List[Dict[str, Any]],

155 **kwargs

156 ) -> LLMResponse:

157 """HuggingFace doesn't have native function calling."""

158 raise NotImplementedError("Function calling not supported for HuggingFace models")

159

160 def _build_prompt(self, messages: List[LLMMessage]) -> str:

161 """Build prompt from messages."""

162 prompt = ""

163 for msg in messages:

164 if msg.role == 'system':

165 prompt += f"{msg.content}\n\n"

166 elif msg.role == 'user':

167 prompt += f"User: {msg.content}\n"

168 elif msg.role == 'assistant':

169 prompt += f"Assistant: {msg.content}\n"

170 return prompt