Skip to content

Add support for SSE (WagtailVectorIndexSSEConsumer) and async querying#28

Open
Morsey187 wants to merge 9 commits into
mainfrom
feature/event-stream-http-consumer
Open

Add support for SSE (WagtailVectorIndexSSEConsumer) and async querying#28
Morsey187 wants to merge 9 commits into
mainfrom
feature/event-stream-http-consumer

Conversation

@Morsey187
Copy link
Copy Markdown
Collaborator

@Morsey187 Morsey187 commented Dec 19, 2023

Changes

  1. Adds support for a new "aquery" method for asynchronous calls.
  2. Adds WagtailVectorIndexSSEConsumer a AsyncHttpConsumer which supports persistent HTTP connections using the event-stream format, where query responses can be streamed (resulting in that typewriter effect).

Testing

Requires following instructions in the added documentation including setting up an ASGI server.

Comment thread docs/quick-start.md
Comment thread src/wagtail_vector_index/consumers.py Outdated
Comment thread src/wagtail_vector_index/consumers.py Outdated
Comment thread docs/quick-start.md Outdated
Comment thread docs/quick-start.md Outdated
Comment thread src/wagtail_vector_index/consumers.py Outdated
Comment thread src/wagtail_vector_index/index/base.py Outdated
@Morsey187 Morsey187 changed the title Draft: Add support for SSE (WagtailVectorIndexSSEConsumer) and async querying Add support for SSE (WagtailVectorIndexSSEConsumer) and async querying Dec 21, 2023
Comment thread docs/quick-start.md Outdated
Comment thread docs/quick-start.md Outdated
Comment thread docs/quick-start.md Outdated
Comment thread docs/quick-start.md Outdated
Comment thread docs/quick-start.md
Comment thread docs/quick-start.md Outdated
Comment thread docs/quick-start.md
Comment thread docs/quick-start.md Outdated
Comment thread pyproject.toml Outdated
Comment thread src/wagtail_vector_index/consumers.py Outdated
# Finish the response
await self.send_body(b"")

async def process_prompt(self, query, vector_index):
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we've been type hinting from the onset in this package. Any chance you can continue the tradition consistently in your changes?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've had to update this to any, although I know it should be a VectorIndex type there are some import errors raise AppRegistryNotReady("Apps aren't loaded yet.") with using it. Theres also quite a few existing type errors which open another PR to resolve some of them.

Comment thread src/wagtail_vector_index/consumers.py Outdated
Note:
This consumer expects the following query parameters in the URL:
- 'query': The search query.
- 'page_type': The type of Wagtail page to search.
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: Vector indexes may not to be related to Wagtail pages - they could be Documents, other Django models, or some completely independent data source. If we can take the index name here and look it up in the registry, that would enable this endpoint to work across all defined indexes.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Had a look into this, however, could do with some pointers as I can't figure out why self.object_type.bulk_from_documents in aquery errors.

Comment thread docs/quick-start.md Outdated
@Morsey187 Morsey187 force-pushed the feature/event-stream-http-consumer branch from f1fa5e8 to 6b370e9 Compare March 11, 2024 09:38


@dataclass
class AsyncQueryResponse(Generic[VectorIndexableType]):
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: I have to return the whole response object to access it's iterator (response chunks). I’d prefer to avoid this and instead update QueryResponse to return the whole response instead of a string too.

Theres also the case that users might want to access the response object as a whole and for other attributes and methods like json().

If so maybe we can rename response as "llm_response"?

"""
Async version of the query method.
"""
if not self.chat_backend.can_stream():
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: I've added this here as I think there will likely be an issue with developers expecting streaming support when using async, especially since we're abstracting from the LLM package, so theres no obvious way for them to know without tracing errors back.

I don't think this is the correct solution if it is a problem, but have added it here to highlight the potential issue and get some thoughts?

Comment thread docs/quick-start.md
]
```

Next, you will need to define a new consumer inheriting from `WagtailVectorIndexSSEConsumer`, and assign a Wagtail page model for the vector index you'd like to use.
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Issue: I believe this is no longer true now we use a single consumer across all indexes?

form = WagtailVectorIndexQueryParamsForm(query_dict)
if not form.is_valid():
# Ignore "TRY301 Abstract `raise` to an inner function"
# So we can insure the event-stream is closed and no other code is executed
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nitpick (non-blocking): insure -> ensure

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants