Add streaming support and refactor client interface

- Implement streaming functionality for all providers
- Update README with streaming examples and improved usage instructions
- Refactor LlmInterface to support both generate and stream methods
- Remove vLLM provider
- Improve type hints and error handling
- Update client initialization in examples

Author: brunneis

README.md

@@ -15,56 +15,127 @@
 pip install llmdk
 ```
 
-# Usage
+# Basic Usage
+
+## Client
+
+```python
+from llmdk import Llmdk, Providers
+
+# You can also set OPENAI_API_KEY
+client = Llmdk(
+    provider=Providers.OPENAI,
+    model_name='gpt-4o-mini',
+    # api_key='***',
+)
+```
+
+## Generate
+
+### Prompt
+
+```python
+output = client.generate(
+    'Who are you?',
+    # system='Write in Portuguese.',
+)
+```
+
+### List of messages
+
+```python
+output = client.generate(
+    messages=[
+        # {'role': 'system', 'content': 'Write in Portuguese.'},
+        {'role': 'user', 'content': 'Who are you?'},
+    ],
+)
+```
+
+## Stream
+
+### Prompt
+
+```python
+for chunk in client.stream(
+    'Who are you?',
+    # system='Write in Portuguese.',
+):
+    print(chunk, end='', flush=True)
+```
+
+### List of messages
+
+```python
+for chunk in client.stream([
+    # {'role': 'system', 'content': 'Write in Portuguese.'},
+    {'role': 'user', 'content': 'Who are you?'},
+]):
+    print(chunk, end='', flush=True)
+```
+
+# Supported Providers
 
 ## Anthropic
+
 ```python
 from llmdk import Llmdk, Providers
 
 # You can also set ANTHROPIC_API_KEY
-client = Llmdk(Providers.ANTHROPIC, 'claude-3-5-sonnet-20240620', api_key='***')
-output = client.generate('Who are you?')
+client = Llmdk(
+    provider=Providers.ANTHROPIC,
+    model_name='claude-3-5-sonnet-20240620',
+    # api_key='***',
+)
 ```
 
 ## Groq
+
 ```python
 from llmdk import Llmdk, Providers
 
 # You can also set GROQ_API_KEY
-client = Llmdk(Providers.GROQ, 'llama-3.1-70b-versatile', api_key='***')
-output = client.generate('Who are you?')
+client = Llmdk(
+    provider=Providers.GROQ,
+    model_name='llama-3.1-70b-versatile',
+    # api_key='***',
+)
 ```
 
 ## HuggingFace
+
 ```python
 from llmdk import Llmdk, Providers
 
 # You can also set HF_TOKEN
-client = Llmdk(Providers.HUGGINGFACE, 'meta-llama/Meta-Llama-3.1-70B-Instruct', api_key='***')
-output = client.generate('Who are you?')
+client = Llmdk(
+    provider=Providers.HUGGINGFACE,
+    model_name='meta-llama/Meta-Llama-3.1-70B-Instruct',
+    # api_key='***',
+)
 ```
 
 ## Ollama
+
 ```python
 from llmdk import Llmdk, Providers
 
-client = Llmdk(Providers.OLLAMA, 'llama3.1:8b', base_url='http://...')
-output = client.generate('Who are you?')
+client = Llmdk(
+    provider=Providers.OLLAMA,
+    model_name='llama3.2:1b',
+    # base_url='http://localhost:11434',
+)
 ```
 
 ## OpenAI
-```python
-from llmdk import Llmdk, Providers
 
-# You can also set OPENAI_API_KEY
-client = Llmdk(Providers.OPENAI, 'gpt-4o-2024-08-06', api_key='***')
-output = client.generate('Who are you?')
-```
-
-## vLLM
 ```python
 from llmdk import Llmdk, Providers
 
-client = Llmdk(Providers.VLLM, base_url='http://...')
-output = client.generate('Who are you?')
+# You can also set OPENAI_API_KEY
+client = Llmdk(
+    provider=Providers.OPENAI,
+    model_name='gpt-4o-mini',
+    # api_key='***',
+)
 ```

llmdk/llmdk.py

@@ -1,14 +1,14 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+from collections.abc import Iterator
 from enum import Enum
-
+from typing import Any, Dict, List, Optional
 from llmdk.providers.anthropic import AnthropicClient
 from llmdk.providers.groq import GroqClient
 from llmdk.providers.huggingface import HuggingFaceClient
 from llmdk.providers.ollama import OllamaClient
 from llmdk.providers.openai import OpenAiClient
-from llmdk.providers.vllm import VllmClient
 
 
 class Providers(Enum):
@@ -17,7 +17,6 @@ class Providers(Enum):
     HUGGINGFACE = 'huggingface'
     OLLAMA = 'ollama'
     OPENAI = 'openai'
-    VLLM = 'vllm'
 
 
 class Llmdk:
@@ -80,17 +79,22 @@ def __init__(
             )
             return
 
-        if (
-            provider == Providers.VLLM
-            or provider == Providers.VLLM.value
-        ):
-            self._client = VllmClient(
-                base_url=base_url,
-            )
-            return
-
         raise ValueError(f"Provider {provider} is not supported")
 
     # Fallback to the original client
     def __getattr__(self, name):
         return getattr(self._client, name)
+
+    def stream(
+        self,
+        prompt: str,
+        system: Optional[str] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
+        **kwargs: Any,
+    ) -> Iterator[str]:
+        return self._client.stream(
+            prompt,
+            system=system,
+            messages=messages,
+            **kwargs,
+        )

llmdk/providers/anthropic.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from os import environ as env
-from typing import Any, Optional, Dict
+from typing import Any, Iterator, Optional, Dict
 from anthropic import Anthropic
 from llmdk.providers.interface import LlmInterface
 
@@ -25,7 +25,30 @@ def __init__(
             base_url=base_url,
         )
 
+    def _prepare_payload(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        if 'max_tokens' not in payload:
+            payload['max_tokens'] = 4096  # Required by Anthropic
+
+        # Check for system message and move it to the system property
+        messages = payload.get('messages', [])
+        if messages and messages[0]['role'] == 'system':
+            payload['system'] = messages[0]['content']
+            payload['messages'] = messages[1:]
+
+        return payload
+
     def _execute_request(self, payload: Dict[str, Any]) -> str:
-        payload['max_tokens'] = 4096  # Required
-        completion = self._client.messages.create(**payload)
+        prepared_payload = self._prepare_payload(payload)
+        completion = self._client.messages.create(**prepared_payload)
         return completion.content[0].text
+
+    def _execute_stream_request(
+        self,
+        payload: Dict[str, Any],
+    ) -> Iterator[str]:
+        prepared_payload = self._prepare_payload(payload)
+
+        with self._client.messages.stream(**prepared_payload) as stream:
+            for message in stream:
+                if message.type == 'content_block_delta':
+                    yield message.delta.text

llmdk/providers/groq.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from os import environ as env
-from typing import Any, Optional, Dict
+from typing import Any, Iterator, Optional, Dict
 from groq import Groq
 from llmdk.providers.interface import LlmInterface
 
@@ -26,3 +26,12 @@ def __init__(
     def _execute_request(self, payload: Dict[str, Any]) -> str:
         completion = self._client.chat.completions.create(**payload)
         return completion.choices[0].message.content
+
+    def _execute_stream_request(
+        self,
+        payload: Dict[str, Any],
+    ) -> Iterator[str]:
+        payload['stream'] = True
+        for chunk in self._client.chat.completions.create(**payload):
+            if chunk.choices[0].delta.content is not None:
+                yield chunk.choices[0].delta.content

llmdk/providers/huggingface.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from os import environ as env
-from typing import Any, Optional, Dict
+from typing import Any, Iterator, Optional, Dict
 from huggingface_hub import InferenceClient
 from llmdk.providers.interface import LlmInterface
 
@@ -27,3 +27,12 @@ def __init__(
     def _execute_request(self, payload: Dict[str, Any]) -> str:
         completion = self._client.chat_completion(**payload)
         return completion.choices[0].message.content
+
+    def _execute_stream_request(
+        self,
+        payload: Dict[str, Any],
+    ) -> Iterator[str]:
+        payload['stream'] = True
+        output = self._client.chat.completions.create(**payload)
+        for chunk in output:
+            yield chunk.choices[0].delta.content