30 lines
745 B
Python
30 lines
745 B
Python
from typing import AsyncIterator
|
|
import litellm
|
|
from app.core.config import DEFAULT_MODEL
|
|
|
|
litellm.drop_params = True
|
|
|
|
|
|
async def stream_chat(
|
|
messages: list[dict],
|
|
model: str | None = None,
|
|
) -> AsyncIterator[str]:
|
|
response = await litellm.acompletion(
|
|
model=model or DEFAULT_MODEL,
|
|
messages=messages,
|
|
stream=True,
|
|
)
|
|
async for chunk in response:
|
|
delta = chunk.choices[0].delta.content
|
|
if delta:
|
|
yield delta
|
|
|
|
|
|
async def complete(messages: list[dict], model: str | None = None) -> str:
|
|
response = await litellm.acompletion(
|
|
model=model or DEFAULT_MODEL,
|
|
messages=messages,
|
|
stream=False,
|
|
)
|
|
return response.choices[0].message.content or ""
|