diff --git a/pyproject.toml b/pyproject.toml index 24c3455..a764fe6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ ] [tool.uv.sources] -genai_utils = { git = "https://github.com/FullFact/genai-utils.git", tag = "5.0.0"} +genai_utils = { git = "https://github.com/FullFact/genai-utils.git", tag = "5.2.0"} [build-system] requires = ["setuptools >= 61.0", diff --git a/src/pastel/pastel.py b/src/pastel/pastel.py index 1c1d753..55dd1dc 100644 --- a/src/pastel/pastel.py +++ b/src/pastel/pastel.py @@ -10,7 +10,7 @@ import numpy as np import numpy.typing as npt import tenacity -from genai_utils.gemini import run_prompt +from genai_utils.gemini import run_prompt_async from google.api_core import exceptions as core_exceptions from pastel import pastel_functions @@ -207,7 +207,7 @@ async def _get_llm_answers_for_single_sentence( """Runs all genAI questions on the given sentence.""" sent_answers: dict[FEATURE_TYPE, float] = {} prompt = self.make_prompt(sentence) - raw_output = run_prompt(prompt) + raw_output = await run_prompt_async(prompt) raw_output = raw_output.strip().lower() if "question" in raw_output: output = raw_output[raw_output.index("0") :] diff --git a/tests/pastel/test_pastel.py b/tests/pastel/test_pastel.py index 4c3f4d1..ad82189 100644 --- a/tests/pastel/test_pastel.py +++ b/tests/pastel/test_pastel.py @@ -78,7 +78,7 @@ def test_quantify_answers(pastel_instance: Pastel) -> None: @patch( - "pastel.pastel.run_prompt", + "pastel.pastel.run_prompt_async", side_effect=ValueError("Gemini failed"), ) async def test_retries(mock_run_prompt: AsyncMock, pastel_instance: Pastel) -> None: diff --git a/uv.lock b/uv.lock index ff8b722..55b32c7 100644 --- a/uv.lock +++ b/uv.lock @@ -243,7 +243,7 @@ wheels = [ [[package]] name = "genai-utils" version = "0.0.1" -source = { git = "https://github.com/FullFact/genai-utils.git?tag=5.0.0#0a689e2a20bd32e09f1bb0d0543e2ba969392965" } +source = { git = "https://github.com/FullFact/genai-utils.git?tag=5.2.0#da1638fde3e86c770f51d2aec7405f1f17cb303b" } dependencies = [ { name = "google-cloud-core" }, { name = "google-genai" }, @@ -700,7 +700,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "genai-utils", git = "https://github.com/FullFact/genai-utils.git?tag=5.0.0" }, + { name = "genai-utils", git = "https://github.com/FullFact/genai-utils.git?tag=5.2.0" }, { name = "google-cloud-core", specifier = ">=2.4.3" }, { name = "pydantic", specifier = ">=2.11.7" }, { name = "scikit-learn", specifier = ">=1.7.1" },