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

Fix Nuclia Search Async #151

Merged
merged 3 commits into from
Feb 14, 2025
Merged

Fix Nuclia Search Async #151

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a59fe44
Fix Nuclia Search Async
bloodbare Jan 30, 2025
ffa68fc
Async generator support for async ask_stream
javitonino Feb 5, 2025
89f75f1
Reflect new context size in tests
javitonino Feb 5, 2025
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
33 changes: 32 additions & 1 deletion nuclia/decorators.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,35 @@ async def async_wrapper_checkout_kb(*args, **kwargs):
except ConnectError:
raise NucliaConnectionError(f"Could not connect to {ndb}")

@wraps(func)
async def async_generative_wrapper_checkout_kb(*args, **kwargs):
if "ndb" in kwargs:
async for value in func(*args, **kwargs):
yield value
else:
url = kwargs.get("url")
api_key = kwargs.get("api_key")
auth = get_async_auth()
if url is None:
# Get default KB
kbid = auth._config.get_default_kb()
if kbid is None:
raise NotDefinedDefault()
ndb = await get_async_client(kbid)
elif url.find(BASE_DOMAIN) >= 0:
region = url.split(".")[0].split("/")[-1]
ndb = AsyncNucliaDBClient(
environment=Environment.CLOUD,
url=url,
api_key=api_key,
region=region,
)
else:
ndb = AsyncNucliaDBClient(environment=Environment.OSS, url=url)
kwargs["ndb"] = ndb
async for value in func(*args, **kwargs):
yield value

@wraps(func)
def wrapper_checkout_kb(*args, **kwargs):
if "ndb" in kwargs:
Expand Down Expand Up @@ -89,8 +118,10 @@ def wrapper_checkout_kb(*args, **kwargs):
except ConnectError:
raise NucliaConnectionError(f"Could not connect to {ndb}")

if asyncio.iscoroutinefunction(func):
if inspect.iscoroutinefunction(func):
return async_wrapper_checkout_kb
elif inspect.isasyncgenfunction(func):
return async_generative_wrapper_checkout_kb
else:
return wrapper_checkout_kb

Expand Down
10 changes: 5 additions & 5 deletions nuclia/sdk/search.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ def search(

See https://docs.nuclia.dev/docs/api#tag/Search/operation/Search_Knowledge_Box_kb__kbid__search_post
"""
ndb: AsyncNucliaDBClient = kwargs["ndb"]
ndb: NucliaDBClient = kwargs["ndb"]
if isinstance(query, str):
req = SearchRequest(query=query, filters=(filters or []), **kwargs) # type: ignore
elif isinstance(query, SearchRequest):
Expand Down Expand Up @@ -325,7 +325,7 @@ async def search(

See https://docs.nuclia.dev/docs/api#tag/Search/operation/Search_Knowledge_Box_kb__kbid__search_post
"""
ndb: NucliaDBClient = kwargs["ndb"]
ndb: AsyncNucliaDBClient = kwargs["ndb"]
if isinstance(query, str):
req = SearchRequest(query=query, filters=(filters or []), **kwargs) # type: ignore
elif isinstance(query, SearchRequest):
Expand Down Expand Up @@ -393,7 +393,7 @@ async def ask(

See https://docs.nuclia.dev/docs/api#tag/Search/operation/Ask_Knowledge_Box_kb__kbid__ask_post
"""
ndb: NucliaDBClient = kwargs["ndb"]
ndb: AsyncNucliaDBClient = kwargs["ndb"]

if isinstance(query, str):
req = AskRequest(
Expand Down Expand Up @@ -479,7 +479,7 @@ async def ask_stream(

See https://docs.nuclia.dev/docs/api#tag/Search/operation/Ask_Knowledge_Box_kb__kbid__ask_post
"""
ndb: NucliaDBClient = kwargs["ndb"]
ndb: AsyncNucliaDBClient = kwargs["ndb"]
if isinstance(query, str):
req = AskRequest(
query=query,
Expand Down Expand Up @@ -519,7 +519,7 @@ async def ask_json(

See https://docs.nuclia.dev/docs/api#tag/Search/operation/Ask_Knowledge_Box_kb__kbid__ask_post
"""
ndb: NucliaDBClient = kwargs["ndb"]
ndb: AsyncNucliaDBClient = kwargs["ndb"]

if isinstance(query, str):
req = AskRequest(
Expand Down
2 changes: 1 addition & 1 deletion nuclia/tests/test_nua/test_predict.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ def test_predict_query(testing_config):
)
assert query.language == "en"
assert query.visual_llm is True
assert query.max_context == 128000
assert query.max_context == 127000
assert query.entities and query.entities.tokens[0].text == "Ramon"
assert query.sentence and len(query.sentence.data) == 1024

Expand Down