Skip to content

Commit 38606a6

Browse files
committed
update vertex2
1 parent a5ff969 commit 38606a6

1 file changed

Lines changed: 50 additions & 5 deletions

File tree

src/uipath_llamaindex/llms/vertex.py

Lines changed: 50 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -319,14 +319,59 @@ async def gen() -> CompletionResponseAsyncGen:
319319

320320
return gen()
321321

322+
def _create_sync_client(self) -> google.genai.Client:
323+
"""Create a fresh google.genai client for sync operations."""
324+
uipath_url = self._uipath_url
325+
headers = self._build_headers_static(self._uipath_token)
326+
327+
http_options = genai_types.HttpOptions(
328+
httpx_client=httpx.Client(
329+
transport=_UrlRewriteTransport(uipath_url),
330+
headers=headers,
331+
follow_redirects=True,
332+
),
333+
)
334+
335+
return google.genai.Client(
336+
api_key="uipath-gateway",
337+
http_options=http_options,
338+
)
339+
322340
@llm_chat_callback()
323341
def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
324-
"""Chat endpoint - delegates to parent's async chat via sync wrapper."""
325-
import asyncio
342+
"""Chat endpoint using sync client."""
343+
from llama_index.llms.google_genai.utils import chat_from_gemini_response
326344

327-
return asyncio.get_event_loop().run_until_complete(
328-
self.achat(messages, **kwargs)
329-
)
345+
generation_config = {
346+
**(self._generation_config or {}),
347+
**kwargs.pop("generation_config", {}),
348+
}
349+
350+
# Create a fresh sync client to avoid event loop issues
351+
sync_client = self._create_sync_client()
352+
353+
try:
354+
# Convert messages to Gemini format (sync version - no file uploads)
355+
contents = []
356+
for message in messages:
357+
content = genai_types.Content(
358+
role="user" if message.role.value == "user" else "model",
359+
parts=[genai_types.Part(text=message.content or "")],
360+
)
361+
contents.append(content)
362+
363+
# Use sync client to send message
364+
chat = sync_client.chats.create(
365+
model=self.model,
366+
config=generation_config,
367+
history=contents[:-1] if len(contents) > 1 else None,
368+
)
369+
response = chat.send_message(contents[-1].parts if contents else [])
370+
371+
return chat_from_gemini_response(response, [])
372+
finally:
373+
# Clean up the sync client
374+
pass
330375

331376
@llm_chat_callback()
332377
def stream_chat(

0 commit comments

Comments
 (0)