Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

VertexAI Streaming support #1354

Open
ADR-007 opened this issue Feb 21, 2025 · 0 comments
Open

VertexAI Streaming support #1354

ADR-007 opened this issue Feb 21, 2025 · 0 comments
Labels
bug Something isn't working enhancement New feature or request help wanted Extra attention is needed

Comments

@ADR-007
Copy link

ADR-007 commented Feb 21, 2025

Hello,

I see that streaming is not supported by instructor for VertexAI, while it is very simple to fix:

instructor/dsl/partial.py:

class PartialBase(Generic[T_Model]):
    # ...
    @staticmethod
    async def extract_json_async(
        completion: AsyncGenerator[Any, None], mode: Mode
    ) -> AsyncGenerator[str, None]:
        async for chunk in completion:
            try:
                # ...
                elif mode == Mode.VERTEXAI_JSON:
                    yield chunk.candidates[0].content.parts[0]
                # ...

Currently, I'm using this mixin with my models as a workaround:

class VertexAiStreamingModelMixin:
    @staticmethod
    async def extract_json_async(
        completion: AsyncGenerator[Any, None], mode: Mode
    ) -> AsyncGenerator[str, None]:
        # Gemini Streaming is not implemented on the instructor library yet
        # So, we need to override PartialBase.extract_json_async in
        # .venv/lib/python3.12/site-packages/instructor/dsl/partial.py
        if mode == Mode.VERTEXAI_JSON:
            async for chunk in completion:
                yield chunk.candidates[0].content.parts[0].text
        else:
            async for chunk in super().extract_json_async(completion, mode):
                yield chunk

I'm using VertexAI with Gemini Pro 1.5

It would be awesome to have it in the library!
Thanks!

@github-actions github-actions bot added bug Something isn't working enhancement New feature or request help wanted Extra attention is needed labels Feb 21, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working enhancement New feature or request help wanted Extra attention is needed
Projects
None yet
Development

No branches or pull requests

1 participant