mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
docs(custom_llm_server.md): add streaming example for custom llm call
Reference - https://github.com/BerriAI/litellm/issues/5147
This commit is contained in:
parent
30df7c7429
commit
50383cfe62
1 changed files with 50 additions and 0 deletions
|
|
@ -131,6 +131,56 @@ Expected Response
|
|||
}
|
||||
```
|
||||
|
||||
## Add Streaming Support
|
||||
|
||||
Here's a simple example of returning unix epoch seconds for both completion + streaming use-cases.
|
||||
|
||||
s/o [@Eloy Lafuente](https://github.com/stronk7) for this code example.
|
||||
|
||||
```python
|
||||
import time
|
||||
from typing import Iterator, AsyncIterator
|
||||
from litellm.types.utils import GenericStreamingChunk, ModelResponse
|
||||
from litellm import CustomLLM, completion, acompletion
|
||||
|
||||
class UnixTimeLLM(CustomLLM):
|
||||
def completion(self, *args, **kwargs) -> ModelResponse:
|
||||
return completion(
|
||||
model="test/unixtime",
|
||||
mock_response=str(int(time.time())),
|
||||
) # type: ignore
|
||||
|
||||
async def acompletion(self, *args, **kwargs) -> ModelResponse:
|
||||
return await acompletion(
|
||||
model="test/unixtime",
|
||||
mock_response=str(int(time.time())),
|
||||
) # type: ignore
|
||||
|
||||
def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
|
||||
generic_streaming_chunk: GenericStreamingChunk = {
|
||||
"finish_reason": "stop",
|
||||
"index": 0,
|
||||
"is_finished": True,
|
||||
"text": str(int(time.time())),
|
||||
"tool_use": None,
|
||||
"usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
|
||||
}
|
||||
return generic_streaming_chunk # type: ignore
|
||||
|
||||
async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
|
||||
generic_streaming_chunk: GenericStreamingChunk = {
|
||||
"finish_reason": "stop",
|
||||
"index": 0,
|
||||
"is_finished": True,
|
||||
"text": str(int(time.time())),
|
||||
"tool_use": None,
|
||||
"usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
|
||||
}
|
||||
yield generic_streaming_chunk # type: ignore
|
||||
|
||||
unixtime = UnixTimeLLM()
|
||||
```
|
||||
|
||||
## Custom Handler Spec
|
||||
|
||||
```python
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue