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:
Krrish Dholakia 2024-08-21 15:45:54 -07:00
parent 30df7c7429
commit 50383cfe62

View file

@ -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