Skip to content

Commit 65fdc85

Browse files
authored
fix(llm): implement async chat methods for GoogleChat (Vertex AI Gemini) (#15994)
Fixed GoogleChat's async client bug where it incorrectly sent Vertex AI service account JSON as an OpenAI API key, by overriding _async_chat and _async_chat_streamly to use the async genai.Client methods for Gemini models.
1 parent ee75f53 commit 65fdc85

1 file changed

Lines changed: 96 additions & 0 deletions

File tree

rag/llm/chat_model.py

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1778,6 +1778,102 @@ def chat_streamly(self, system, history, gen_conf=None, **kwargs):
17781778

17791779
yield total_tokens
17801780

1781+
async def _async_chat(self, history, gen_conf, **kwargs):
1782+
if "claude" in self.model_name:
1783+
return await super()._async_chat(history, gen_conf, **kwargs)
1784+
1785+
gen_conf = dict(gen_conf or {})
1786+
system = history[0]["content"] if history and history[0]["role"] == "system" else ""
1787+
history = [h for h in history if h["role"] != "system"]
1788+
1789+
if "thinking_budget" not in gen_conf:
1790+
gen_conf["thinking_budget"] = 0
1791+
thinking_budget = gen_conf.pop("thinking_budget", 0)
1792+
gen_conf = self._clean_conf(gen_conf)
1793+
1794+
try:
1795+
from google.genai.types import Content, GenerateContentConfig, Part, ThinkingConfig
1796+
except ImportError as e:
1797+
logging.error(f"[GoogleChat] Failed to import google-genai: {e}. Please install: pip install google-genai>=1.41.0")
1798+
raise
1799+
1800+
config_dict = {}
1801+
if system:
1802+
config_dict["system_instruction"] = system
1803+
if "temperature" in gen_conf:
1804+
config_dict["temperature"] = gen_conf["temperature"]
1805+
if "top_p" in gen_conf:
1806+
config_dict["top_p"] = gen_conf["top_p"]
1807+
if "max_output_tokens" in gen_conf:
1808+
config_dict["max_output_tokens"] = gen_conf["max_output_tokens"]
1809+
config_dict["thinking_config"] = ThinkingConfig(thinking_budget=thinking_budget)
1810+
config = GenerateContentConfig(**config_dict)
1811+
1812+
contents = []
1813+
for item in history:
1814+
role = "model" if item["role"] == "assistant" else item["role"]
1815+
contents.append(Content(role=role, parts=[Part(text=item["content"])]))
1816+
1817+
response = await self.client.aio.models.generate_content(
1818+
model=self.model_name,
1819+
contents=contents,
1820+
config=config,
1821+
)
1822+
ans = response.text or ""
1823+
try:
1824+
total_tokens = response.usage_metadata.total_token_count
1825+
except Exception:
1826+
total_tokens = num_tokens_from_string(ans)
1827+
return ans, total_tokens
1828+
1829+
async def _async_chat_streamly(self, history, gen_conf, **kwargs):
1830+
if "claude" in self.model_name:
1831+
async for delta_ans, tol in super()._async_chat_streamly(history, gen_conf, **kwargs):
1832+
yield delta_ans, tol
1833+
return
1834+
1835+
gen_conf = dict(gen_conf or {})
1836+
system = history[0]["content"] if history and history[0]["role"] == "system" else ""
1837+
history = [h for h in history if h["role"] != "system"]
1838+
1839+
if "thinking_budget" not in gen_conf:
1840+
gen_conf["thinking_budget"] = 0
1841+
thinking_budget = gen_conf.pop("thinking_budget", 0)
1842+
gen_conf = self._clean_conf(gen_conf)
1843+
1844+
try:
1845+
from google.genai.types import Content, GenerateContentConfig, Part, ThinkingConfig
1846+
except ImportError as e:
1847+
logging.error(f"[GoogleChat] Failed to import google-genai: {e}. Please install: pip install google-genai>=1.41.0")
1848+
raise
1849+
1850+
config_dict = {}
1851+
if system:
1852+
config_dict["system_instruction"] = system
1853+
if "temperature" in gen_conf:
1854+
config_dict["temperature"] = gen_conf["temperature"]
1855+
if "top_p" in gen_conf:
1856+
config_dict["top_p"] = gen_conf["top_p"]
1857+
if "max_output_tokens" in gen_conf:
1858+
config_dict["max_output_tokens"] = gen_conf["max_output_tokens"]
1859+
config_dict["thinking_config"] = ThinkingConfig(thinking_budget=thinking_budget)
1860+
config = GenerateContentConfig(**config_dict)
1861+
1862+
contents = []
1863+
for item in history:
1864+
role = "model" if item["role"] == "assistant" else item["role"]
1865+
contents.append(Content(role=role, parts=[Part(text=item["content"])]))
1866+
1867+
stream = await self.client.aio.models.generate_content_stream(
1868+
model=self.model_name,
1869+
contents=contents,
1870+
config=config,
1871+
)
1872+
async for chunk in stream:
1873+
text = chunk.text
1874+
if text:
1875+
yield text, num_tokens_from_string(text)
1876+
17811877

17821878
class TokenPonyChat(Base):
17831879
_FACTORY_NAME = "TokenPony"

0 commit comments

Comments
 (0)