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.

Sign up for GitHub

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

bugfix when initializing with async aoss vector store #17340

Merged
merged 7 commits into from
Dec 23, 2024
Merged

bugfix when initializing with async aoss vector store #17340

Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ff4215c
bugfix when initializing with async aoss
Aydin-ab Dec 20, 2024
841177e
Merge branch 'main' into bugfix-async-aoss-vector-stores
Aydin-ab Dec 20, 2024
6be8050
fix self.os_version error
Aydin-ab Dec 20, 2024
1c6f818
Merge branch 'bugfix-async-aoss-vector-stores' of https://github.com/…
Aydin-ab Dec 20, 2024
48add0f
lint
logan-markewich Dec 20, 2024
c380264
vbump
logan-markewich Dec 20, 2024
ada5dec
use asyncio_run
logan-markewich Dec 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions ..._stores/llama-index-vector-stores-opensearch/llama_index/vector_stores/opensearch/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import uuid
from datetime import datetime
from typing import Any, Dict, Iterable, List, Optional, Union, cast

import asyncio
from llama_index.core.bridge.pydantic import PrivateAttr

from llama_index.core.schema import BaseNode, MetadataMode, TextNode
Expand Down Expand Up @@ -123,14 +123,18 @@ def __init__(
self._os_async_client = os_async_client or self._get_async_opensearch_client(
Copy link
Collaborator

Choose a reason for hiding this comment

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

just curious, there's no way to create an sync client from an async client, or vice-versa? This would solve a lot of headaches if possible

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I guess you could if you dig into the kwargs and change it accordingly for each sync/async instantiation

from opensearchpy import AsyncHttpConnection, AWSV4SignerAsyncAuth, RequestsHttpConnection, AWSV4SignerAuth

# Sync parameters
kwargs['auth'] = AWSV4SignerAuth(credentials, region, 'aoss')
kwargs['connection_class'] = RequestsHttpConnection
self._os_client = os_client or self._get_opensearch_client(
    self._endpoint, **kwargs
)

# Async parameters
kwargs['auth'] = AWSV4SignerAsyncAuth(credentials, region, 'aoss')
kwargs['connection_class'] = AsyncHttpConnection
self._os_client = os_client or self._get_opensearch_client(
    self._endpoint, **kwargs
)

But you'd need the AWS credentials/region/service ('aoss') as additional inputs, seems too much

Copy link
Collaborator

Choose a reason for hiding this comment

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

ah yea, thats pretty messy. Too bad

self._endpoint, **kwargs
)
self._os_version = self._get_opensearch_version()
self._efficient_filtering_enabled = self._is_efficient_filtering_enabled(
self._os_version
)
self._efficient_filtering_enabled = self._is_efficient_filtering_enabled()
not_found_error = self._import_not_found_error()

try:
self._os_client.indices.get(index=self._index)
except TypeError:
# Probably using async so switch to async client
try:
asyncio.run(self._os_async_client.indices.get(index=self._index))
logan-markewich marked this conversation as resolved.
Show resolved Hide resolved
except not_found_error:
asyncio.run(self._os_async_client.indices.create(index=self._index, body=idx_conf))
asyncio.run(self._os_async_client.indices.refresh(index=self._index))
except not_found_error:
self._os_client.indices.create(index=self._index, body=idx_conf)
self._os_client.indices.refresh(index=self._index)
Expand Down Expand Up @@ -617,10 +621,18 @@ def _is_aoss_enabled(self, http_auth: Any) -> bool:
return True
return False

def _is_efficient_filtering_enabled(self, os_version: str) -> bool:
def _is_efficient_filtering_enabled(self) -> bool:
"""Check if kNN with efficient filtering is enabled."""
major, minor, patch = os_version.split(".")
return int(major) >= 2 and int(minor) >= 9
# Technically, AOSS supports efficient filtering,
# but we can't check the version number using .info(); AOSS doesn't support 'GET /'
# so we must skip and disable by default.
if self.is_aoss:
ef_enabled = False
else :
self._os_version = self._get_opensearch_version()
major, minor, patch = os_version.split(".")
ef_enabled = int(major) >= 2 and int(minor) >= 9
return ef_enabled

def index_results(self, nodes: List[BaseNode], **kwargs: Any) -> List[str]:
"""Store results in the index."""
Expand Down
Loading