@@ -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
17821878class TokenPonyChat (Base ):
17831879 _FACTORY_NAME = "TokenPony"
0 commit comments