Coverage for src/dataknobs_fsm/llm/base.py: 0%

1"""Base LLM abstraction components.

2

3 This module provides the base abstractions for unified LLM operations.

4"""

5

6 from abc import ABC, abstractmethod

7 from dataclasses import dataclass, field

8 from enum import Enum

9 from typing import (

10 Any, Dict, List, Union, AsyncIterator, Iterator,

11 Callable, Protocol

12)

13 from datetime import datetime

14

15

16 class CompletionMode(Enum):

17 """LLM completion modes."""

18 CHAT = "chat" # Chat completion with message history

19 TEXT = "text" # Text completion

20 INSTRUCT = "instruct" # Instruction following

21 EMBEDDING = "embedding" # Generate embeddings

22 FUNCTION = "function" # Function calling

23

24

25 class ModelCapability(Enum):

26 """Model capabilities."""

27 TEXT_GENERATION = "text_generation"

28 CHAT = "chat"

29 EMBEDDINGS = "embeddings"

30 FUNCTION_CALLING = "function_calling"

31 VISION = "vision"

32 CODE = "code"

33 JSON_MODE = "json_mode"

34 STREAMING = "streaming"

35

36

37@dataclass

38 class LLMMessage:

39 """Represents a message in LLM conversation."""

40 role: str # 'system', 'user', 'assistant', 'function'

41 content: str

42 name: str | None = None # For function messages

43 function_call: Dict[str, Any] | None = None # For function calling

44 metadata: Dict[str, Any] = field(default_factory=dict)

45

46

47@dataclass

48 class LLMResponse:

49 """Response from LLM."""

50 content: str

51 model: str

52 finish_reason: str | None = None # 'stop', 'length', 'function_call'

53 usage: Dict[str, int] | None = None # tokens used

54 function_call: Dict[str, Any] | None = None

55 metadata: Dict[str, Any] = field(default_factory=dict)

56 created_at: datetime = field(default_factory=datetime.now)

57

58

59@dataclass

60 class LLMStreamResponse:

61 """Streaming response from LLM."""

62 delta: str # Incremental content

63 is_final: bool = False

64 finish_reason: str | None = None

65 usage: Dict[str, int] | None = None

66 metadata: Dict[str, Any] = field(default_factory=dict)

67

68

69@dataclass

70 class LLMConfig:

71 """Configuration for LLM operations."""

72 provider: str # 'openai', 'anthropic', 'ollama', etc.

73 model: str # Model name/identifier

74 api_key: str | None = None

75 api_base: str | None = None # Custom API endpoint

76

77 # Generation parameters

78 temperature: float = 0.7

79 max_tokens: int | None = None

80 top_p: float = 1.0

81 frequency_penalty: float = 0.0

82 presence_penalty: float = 0.0

83 stop_sequences: List[str] | None = None

84

85 # Mode settings

86 mode: CompletionMode = CompletionMode.CHAT

87 system_prompt: str | None = None

88 response_format: str | None = None # 'text' or 'json'

89

90 # Function calling

91 functions: List[Dict[str, Any]] | None = None

92 function_call: Union[str, Dict[str, str]] | None = None # 'auto', 'none', or specific function

93

94 # Streaming

95 stream: bool = False

96 stream_callback: Callable[[LLMStreamResponse], None] | None = None

97

98 # Rate limiting

99 rate_limit: int | None = None # Requests per minute

100 retry_count: int = 3

101 retry_delay: float = 1.0

102 timeout: float = 60.0

103

104 # Advanced settings

105 seed: int | None = None # For reproducibility

106 logit_bias: Dict[str, float] | None = None

107 user_id: str | None = None

108

109 # Provider-specific options

110 options: Dict[str, Any] = field(default_factory=dict)

111

112

113 class LLMProvider(ABC):

114 """Base LLM provider interface."""

115

116 def __init__(self, config: LLMConfig):

117 """Initialize provider with configuration."""

118 self.config = config

119 self._client = None

120 self._is_initialized = False

121

122 @abstractmethod

123 def initialize(self) -> None:

124 """Initialize the LLM client."""

125 pass

126

127 @abstractmethod

128 def close(self) -> None:

129 """Close the LLM client."""

130 pass

131

132 @abstractmethod

133 def validate_model(self) -> bool:

134 """Validate that the model is available."""

135 pass

136

137 @abstractmethod

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

139 """Get model capabilities."""

140 pass

141

142 @property

143 def is_initialized(self) -> bool:

144 """Check if provider is initialized."""

145 return self._is_initialized

146

147 def __enter__(self):

148 """Context manager entry."""

149 self.initialize()

150 return self

151

152 def __exit__(self, exc_type, exc_val, exc_tb):

153 """Context manager exit."""

154 self.close()

155

156

157 class AsyncLLMProvider(LLMProvider):

158 """Async LLM provider interface."""

159

160 @abstractmethod

161 async def complete(

162 self,

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

164 **kwargs

165 ) -> LLMResponse:

166 """Generate completion asynchronously.

167

168 Args:

169 messages: Input messages or prompt

170 **kwargs: Additional parameters

171

172 Returns:

173 LLM response

174 """

175 pass

176

177 @abstractmethod

178 async def stream_complete(

179 self,

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

181 **kwargs

182 ) -> AsyncIterator[LLMStreamResponse]:

183 """Generate streaming completion asynchronously.

184

185 Args:

186 messages: Input messages or prompt

187 **kwargs: Additional parameters

188

189 Yields:

190 Streaming response chunks

191 """

192 pass

193

194 @abstractmethod

195 async def embed(

196 self,

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

198 **kwargs

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

200 """Generate embeddings asynchronously.

201

202 Args:

203 texts: Input text(s)

204 **kwargs: Additional parameters

205

206 Returns:

207 Embedding vector(s)

208 """

209 pass

210

211 @abstractmethod

212 async def function_call(

213 self,

214 messages: List[LLMMessage],

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

216 **kwargs

217 ) -> LLMResponse:

218 """Execute function calling asynchronously.

219

220 Args:

221 messages: Conversation messages

222 functions: Available functions

223 **kwargs: Additional parameters

224

225 Returns:

226 Response with function call

227 """

228 pass

229

230 async def initialize(self) -> None:

231 """Initialize the async LLM client."""

232 self._is_initialized = True

233

234 async def close(self) -> None:

235 """Close the async LLM client."""

236 self._is_initialized = False

237

238 async def __aenter__(self):

239 """Async context manager entry."""

240 await self.initialize()

241 return self

242

243 async def __aexit__(self, exc_type, exc_val, exc_tb):

244 """Async context manager exit."""

245 await self.close()

246

247

248 class SyncLLMProvider(LLMProvider):

249 """Synchronous LLM provider interface."""

250

251 @abstractmethod

252 def complete(

253 self,

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

255 **kwargs

256 ) -> LLMResponse:

257 """Generate completion synchronously.

258

259 Args:

260 messages: Input messages or prompt

261 **kwargs: Additional parameters

262

263 Returns:

264 LLM response

265 """

266 pass

267

268 @abstractmethod

269 def stream_complete(

270 self,

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

272 **kwargs

273 ) -> Iterator[LLMStreamResponse]:

274 """Generate streaming completion synchronously.

275

276 Args:

277 messages: Input messages or prompt

278 **kwargs: Additional parameters

279

280 Yields:

281 Streaming response chunks

282 """

283 pass

284

285 @abstractmethod

286 def embed(

287 self,

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

289 **kwargs

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

291 """Generate embeddings synchronously.

292

293 Args:

294 texts: Input text(s)

295 **kwargs: Additional parameters

296

297 Returns:

298 Embedding vector(s)

299 """

300 pass

301

302 @abstractmethod

303 def function_call(

304 self,

305 messages: List[LLMMessage],

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

307 **kwargs

308 ) -> LLMResponse:

309 """Execute function calling synchronously.

310

311 Args:

312 messages: Conversation messages

313 functions: Available functions

314 **kwargs: Additional parameters

315

316 Returns:

317 Response with function call

318 """

319 pass

320

321 def initialize(self) -> None:

322 """Initialize the sync LLM client."""

323 self._is_initialized = True

324

325 def close(self) -> None:

326 """Close the sync LLM client."""

327 self._is_initialized = False

328

329

330 class LLMAdapter(ABC):

331 """Base adapter for converting between different LLM formats."""

332

333 @abstractmethod

334 def adapt_messages(

335 self,

336 messages: List[LLMMessage]

337 ) -> Any:

338 """Adapt messages to provider format."""

339 pass

340

341 @abstractmethod

342 def adapt_response(

343 self,

344 response: Any

345 ) -> LLMResponse:

346 """Adapt provider response to standard format."""

347 pass

348

349 @abstractmethod

350 def adapt_config(

351 self,

352 config: LLMConfig

353 ) -> Dict[str, Any]:

354 """Adapt configuration to provider format."""

355 pass

356

357

358 class LLMMiddleware(Protocol):

359 """Protocol for LLM middleware."""

360

361 async def process_request(

362 self,

363 messages: List[LLMMessage],

364 config: LLMConfig

365 ) -> List[LLMMessage]:

366 """Process request before sending to LLM."""

367 ...

368

369 async def process_response(

370 self,

371 response: LLMResponse,

372 config: LLMConfig

373 ) -> LLMResponse:

374 """Process response from LLM."""

375 ...