fix(vertex_ai.py): add support for real async streaming + completion calls

This commit is contained in:
Krrish Dholakia 2023-12-13 11:53:55 -08:00
parent 07015843ac
commit 69c29f8f86
5 changed files with 134 additions and 49 deletions

View file

@ -19,6 +19,7 @@ import uuid
import aiohttp
import logging
import asyncio, httpx, inspect
from inspect import iscoroutine
import copy
from tokenizers import Tokenizer
from dataclasses import (
@ -5769,7 +5770,8 @@ class CustomStreamWrapper:
or self.custom_llm_provider == "azure"
or self.custom_llm_provider == "custom_openai"
or self.custom_llm_provider == "text-completion-openai"
or self.custom_llm_provider == "huggingface"):
or self.custom_llm_provider == "huggingface"
or self.custom_llm_provider == "vertex_ai"):
async for chunk in self.completion_stream:
if chunk == "None" or chunk is None:
raise Exception