Bump up dependencies to the newest versions

allms/models/abstract.py

@@ -5,6 +5,7 @@
 import urllib
 from abc import ABC, abstractmethod
 from functools import partial
+from urllib.error import URLError
 
 import google
 import openai
@@ -77,8 +78,8 @@ def __init__(
 
         self._predict_example = create_base_retry_decorator(
             error_types=[
-                openai.error.RateLimitError, openai.error.APIError, openai.error.Timeout,
-                openai.error.APIConnectionError, openai.error.ServiceUnavailableError,
+                openai.RateLimitError, openai.APIError, openai.Timeout,
+                openai.APIConnectionError, openai.InternalServerError,
                 google.api_core.exceptions.ResourceExhausted, urllib.error.HTTPError
             ],
             max_retries=max_retries,
@@ -252,16 +253,16 @@ async def _predict_example(
                     model_response = await chain.arun({})
                 else:
                     model_response = await chain.arun(**input_data.input_mappings)
-        except openai.error.InvalidRequestError as invalid_request_error:
+        except openai.InternalServerError as invalid_request_error:
             logger.info(f"Error for id {input_data.id} has occurred. Message: {invalid_request_error} ")
-            if invalid_request_error.error.code == "content_filter":
+            if invalid_request_error.code == "content_filter":
                 model_response = None
                 error_message = f"{IODataConstants.CONTENT_FILTER_MESSAGE}: {invalid_request_error}"
             else:
                 model_response = None
                 error_message = f"{IODataConstants.ERROR_MESSAGE_STR}: {invalid_request_error}"
 
-        except (InvalidArgument, ValueError, TimeoutError, openai.error.Timeout, GCPInvalidRequestError) as other_error:
+        except (InvalidArgument, ValueError, TimeoutError, openai.APIError, GCPInvalidRequestError) as other_error:
             model_response = None
             logger.info(f"Error for id {input_data.id} has occurred. Message: {other_error} ")
             error_message = f"{type(other_error).__name__}: {other_error}"

allms/models/vertexai_base.py

@@ -1,9 +1,9 @@
 from typing import List, Optional, Any, Dict
 
 from google.cloud.aiplatform.models import Prediction
-from langchain_google_vertexai import VertexAI, VertexAIModelGarden
 from langchain_core.callbacks import AsyncCallbackManagerForLLMRun
 from langchain_core.outputs import LLMResult, Generation
+from langchain_google_vertexai import VertexAI, VertexAIModelGarden
 from pydash import chain
 
 from allms.constants.vertex_ai import VertexModelConstants

allms/models/vertexai_gemini.py

@@ -47,7 +47,7 @@ def __init__(
             event_loop=event_loop
         )
 
-    def _create_llm(self) -> VertexAI:
+    def _create_llm(self) -> CustomVertexAI:
         return CustomVertexAI(
             model_name=self._config.gemini_model_name,
             max_output_tokens=self._max_output_tokens,

allms/utils/response_parsing_utils.py

@@ -3,6 +3,7 @@
 
 from langchain.output_parsers import PydanticOutputParser
 from langchain.schema import OutputParserException
+from pydantic import ValidationError
 
 from allms.domain.response import ResponseData, ResponseParsingOutput
 
@@ -45,6 +46,15 @@ def _parse_response(
                     The exception message: {output_parser_exception}
                     """
             )
+        except ValidationError as validation_error:
+            return ResponseParsingOutput(
+                response=None,
+                error_message=f"""
+                    A ValidationError has occurred for the model response: {model_response_data.response}
+                    The exception message: {validation_error}
+                    """
+            )
+
 
     def parse_model_output(
         self, 

pyproject.toml

@@ -7,18 +7,19 @@ readme = "README.md"
 packages = [{include = "allms"}]
 
 [tool.poetry.dependencies]
-python = ">=3.8.1,<4.0"
+python = ">=3.9.0,<4.0"
 fsspec = "^2023.6.0"
-google-cloud-aiplatform = ">=1.57.0"
+google-cloud-aiplatform = "1.70.0"
 pydash = "^7.0.6"
 transformers = "^4.34.1"
-pydantic = "1.10.13"
-langchain = "^0.1.8"
-langchain-google-vertexai = "1.0.4"
-aioresponses = "^0.7.6"
-tiktoken = "^0.6.0"
-openai = "^0.27.8"
+pydantic = "2.7.4"
+langchain = "0.3.6"
+tiktoken = "^0.7.0"
+openai = "1.52.0"
 pytest-mock = "^3.14.0"
+respx = "^0.21.1"
+langchain-community = "^0.3.5"
+langchain-google-vertexai = "^2.0.7"
 sentencepiece = "^0.2.0"
 
 [tool.poetry.group.dev.dependencies]

tests/conftest.py

@@ -5,7 +5,6 @@
 from unittest.mock import patch
 
 import pytest
-from aioresponses import aioresponses
 from langchain_community.llms.fake import FakeListLLM
 
 from allms.domain.configuration import (
@@ -94,9 +93,3 @@ def models():
                     event_loop=event_loop
                 )
             }
-
-
-@pytest.fixture
-def mock_aioresponse():
-    with aioresponses() as http_mock:
-        yield http_mock

tests/test_end_to_end.py

@@ -1,6 +1,9 @@
 import re
 
 import pytest
+import httpx
+import respx
+from httpx import Response
 from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, PromptTemplate, \
     SystemMessagePromptTemplate
 
@@ -16,80 +19,81 @@ class TestEndToEnd:
 
     def test_model_is_queried_successfully(
             self,
-            mock_aioresponse,
             models
     ):
         # GIVEN
-        mock_aioresponse.post(
-            url=re.compile(f"^{AzureOpenAIEnv.OPENAI_API_BASE}.*$"),
-            payload={
-                "choices": [{
-                    "message": {
-                        "content": "{\"keywords\": [\"Indywidualna racja żywnościowa\", \"wojskowa\", \"S-R-9\", \"set nr 9\", \"Makaron po bolońsku\", \"Konserwa tyrolska\", \"Suchary\", \"Koncentrat napoju herbacianego instant o smaku owoców leśnych\", \"Dżem malinowy\", \"Baton zbożowo-owocowy o smaku figowym\"]}",
-                        "role": ""
-                    }
-                }],
-                "usage": {}
-            },
-            repeat=True
-        )
+        with respx.mock:
+            respx.post(
+                url=re.compile(f"^{AzureOpenAIEnv.OPENAI_API_BASE}.*$")).mock(
+                return_value=Response(status_code=200, json={
+                    "choices": [{
+                        "message": {
+                            "content": "{\"keywords\": [\"Indywidualna racja żywnościowa\", \"wojskowa\", \"S-R-9\", \"set nr 9\", \"Makaron po bolońsku\", \"Konserwa tyrolska\", \"Suchary\", \"Koncentrat napoju herbacianego instant o smaku owoców leśnych\", \"Dżem malinowy\", \"Baton zbożowo-owocowy o smaku figowym\"]}",
+                            "role": ""
+                        }
+                    }],
+                    "usage": {}
+                },
+                                      )
+            )
 
-        input_data = io_utils.load_csv_to_input_data(
-            limit=5,
-            path="./tests/resources/test_input_data.csv"
-        )
-        prompt_template_text = """Extract at most 10 keywords that could be used as features in a search index from this Polish product description.
+            input_data = io_utils.load_csv_to_input_data(
+                limit=5,
+                path="./tests/resources/test_input_data.csv"
+            )
+            prompt_template_text = """Extract at most 10 keywords that could be used as features in a search index from this Polish product description.
 
         {text}
         """
 
-        # WHEN
-        parsed_responses = models["azure_open_ai"].generate(
-            prompt=prompt_template_text,
-            input_data=input_data,
-            output_data_model_class=KeywordsOutputClass,
-            system_prompt="This is a system prompt."
-        )
-        parsed_responses = sorted(parsed_responses, key=lambda key: key.input_data.id)
+            # WHEN
+            parsed_responses = models["azure_open_ai"].generate(
+                prompt=prompt_template_text,
+                input_data=input_data,
+                output_data_model_class=KeywordsOutputClass,
+                system_prompt="This is a system prompt."
+            )
+            parsed_responses = sorted(parsed_responses, key=lambda key: key.input_data.id)
 
-        # THEN
-        expected_output = io_utils.load_csv("./tests/resources/test_end_to_end_expected_output.csv")
-        expected_output = sorted(expected_output, key=lambda example: example[IODataConstants.ID])
-        for idx in range(len(expected_output)):
-            expected_output[idx]["response"] = eval(expected_output[idx]["response"])
+            # THEN
+            expected_output = io_utils.load_csv("./tests/resources/test_end_to_end_expected_output.csv")
+            expected_output = sorted(expected_output, key=lambda example: example[IODataConstants.ID])
+            for idx in range(len(expected_output)):
+                expected_output[idx]["response"] = eval(expected_output[idx]["response"])
 
-        assert list(map(lambda output: output[IODataConstants.ID], expected_output)) == list(
-            map(lambda example: example.input_data.id, parsed_responses))
+            assert list(map(lambda output: output[IODataConstants.ID], expected_output)) == list(
+                map(lambda example: example.input_data.id, parsed_responses))
 
-        assert list(map(lambda output: output[IODataConstants.TEXT], expected_output)) == list(
-            map(lambda example: example.input_data.input_mappings["text"], parsed_responses))
+            assert list(map(lambda output: output[IODataConstants.TEXT], expected_output)) == list(
+                map(lambda example: example.input_data.input_mappings["text"], parsed_responses))
 
-        assert list(map(lambda output: output[IODataConstants.RESPONSE_STR_NAME], expected_output)) == list(
-            map(lambda example: example.response.keywords, parsed_responses))
+            assert list(map(lambda output: output[IODataConstants.RESPONSE_STR_NAME], expected_output)) == list(
+                map(lambda example: example.response.keywords, parsed_responses))
 
-        assert list(map(lambda output: int(output[IODataConstants.PROMPT_TOKENS_NUMBER]), expected_output)) == list(
-            map(lambda example: example.number_of_prompt_tokens, parsed_responses))
+            assert list(map(lambda output: int(output[IODataConstants.PROMPT_TOKENS_NUMBER]), expected_output)) == list(
+                map(lambda example: example.number_of_prompt_tokens, parsed_responses))
 
-        assert list(map(lambda output: int(output[IODataConstants.GENERATED_TOKENS_NUMBER]), expected_output)) == list(
-            map(lambda example: example.number_of_generated_tokens, parsed_responses))
-
-    def test_prompt_is_not_modified_for_open_source_models(self, mock_aioresponse, models, mocker):
+            assert list(
+                map(lambda output: int(output[IODataConstants.GENERATED_TOKENS_NUMBER]), expected_output)) == list(
+                map(lambda example: example.number_of_generated_tokens, parsed_responses))
+
+    def test_prompt_is_not_modified_for_open_source_models(self, models, mocker):
         # GIVEN
         open_source_models = ["azure_llama2", "azure_mistral", "vertex_gemma"]
 
-        mock_aioresponse.post(
-            url=re.compile(f"^https:\/\/dummy-endpoint.*$"),
-            payload={
-                "choices": [{
-                    "message": {
-                        "content": "{\"keywords\": [\"Indywidualna racja żywnościowa\", \"wojskowa\", \"S-R-9\", \"set nr 9\", \"Makaron po bolońsku\", \"Konserwa tyrolska\", \"Suchary\", \"Koncentrat napoju herbacianego instant o smaku owoców leśnych\", \"Dżem malinowy\", \"Baton zbożowo-owocowy o smaku figowym\"]}",
-                        "role": ""
-                    }
-                }],
-                "usage": {}
-            },
-            repeat=True
-        )
+        with respx.mock:
+            respx.post(
+                url=re.compile(f"^https:\/\/dummy-endpoint.*$")).mock(
+                return_value=Response(status_code=200, json={
+                    "choices": [{
+                        "message": {
+                            "content": "{\"keywords\": [\"Indywidualna racja żywnościowa\", \"wojskowa\", \"S-R-9\", \"set nr 9\", \"Makaron po bolońsku\", \"Konserwa tyrolska\", \"Suchary\", \"Koncentrat napoju herbacianego instant o smaku owoców leśnych\", \"Dżem malinowy\", \"Baton zbożowo-owocowy o smaku figowym\"]}",
+                            "role": ""
+                        }
+                    }],
+                    "usage": {}
+                },
+                                      ))
 
         input_data = io_utils.load_csv_to_input_data(
             limit=5,
@@ -114,42 +118,71 @@ def test_prompt_is_not_modified_for_open_source_models(self, mock_aioresponse, m
                 messages = [
                     HumanMessagePromptTemplate(
                         prompt=PromptTemplate(
-                            input_variables=["text"], 
+                            input_variables=["text"],
                             template=prompt_template_text
                         )
                     )
                 ]
                 if model_name != "azure_mistral":
                     messages = [
-                        SystemMessagePromptTemplate(
-                            prompt=PromptTemplate(
-                                input_variables=[], 
-                                template="This is a system prompt."
-                            )
-                        )
-                    ] + messages
+                                   SystemMessagePromptTemplate(
+                                       prompt=PromptTemplate(
+                                           input_variables=[],
+                                           template="This is a system prompt."
+                                       )
+                                   )
+                               ] + messages
                 prompt_template_spy.assert_called_with(messages)
             else:
                 prompt_template_spy.assert_called_with([
                     SystemMessagePromptTemplate(
                         prompt=PromptTemplate(
-                            input_variables=[], 
+                            input_variables=[],
                             template="This is a system prompt."
                         )
-                    ), 
+                    ),
                     HumanMessagePromptTemplate(
                         prompt=PromptTemplate(
-                            input_variables=["text"], 
+                            input_variables=["text"],
                             partial_variables={
-                                'output_data_model': 'The output should be formatted as a JSON instance that conforms to the JSON schema below.\n\nAs an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}\nthe object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.\n\nHere is the output schema:\n```\n{"properties": {"keywords": {"title": "Keywords", "description": "List of keywords", "type": "array", "items": {"type": "string"}}}, "required": ["keywords"]}\n```'
-                            }, 
+                                'output_data_model': 'The output should be formatted as a JSON instance that conforms to the JSON schema below.\n\nAs an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}\nthe object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.\n\nHere is the output schema:\n```\n{"properties": {"keywords": {"description": "List of keywords", "items": {"type": "string"}, "title": "Keywords", "type": "array"}}, "required": ["keywords"]}\n```'
+                            },
                             template=f"{prompt_template_text}\n\n{{output_data_model}}"
                         )
                     )
                 ])
 
-    def test_gemini_specific_args_are_passed_to_model(self):
+    def test_gemini_version_is_passed_to_model(self):
+        # GIVEN
+        model_config = VertexAIConfiguration(
+            cloud_project="dummy-project-id",
+            cloud_location="us-central1",
+            gemini_model_name="gemini-1.0-pro-001"
+        )
+
+        # WHEN
+        gemini_model = VertexAIGeminiModel(config=model_config)
+
+        # WHEN
+        gemini_model._llm.model_name == "gemini-model-name"
+
+    def test_model_times_out(
+            self,
+            models
+    ):
         # GIVEN
+        with respx.mock:
+            respx.post(re.compile(f"^https:\/\/dummy-endpoint.*$")).mock(
+                side_effect=httpx.TimeoutException("Request timed out")
+            )
+
+            # WHEN
+            responses = models["azure_open_ai"].generate("Some prompt")
+
+            # THEN
+            assert responses[0].response is None
+            assert "Request timed out" in responses[0].error
+    def test_gemini_specific_args_are_passed_to_model(self):
         gemini_model_name = "gemini-1.5-pro-001"
         gemini_safety_settings = {
             HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE,
@@ -164,7 +197,6 @@ def test_gemini_specific_args_are_passed_to_model(self):
             gemini_model_name=gemini_model_name,
             gemini_safety_settings=gemini_safety_settings
         )
-
         # WHEN
         gemini_model = VertexAIGeminiModel(config=model_config)
 
@@ -205,25 +237,3 @@ def test_incorrect_gemini_model_name_fail(self, model_name):
         # WHEN & THEN
         with pytest.raises(ValueError, match=f"Model {model_name} is not supported."):
             VertexAIGeminiModel(config=model_config)
-
-    def test_model_times_out(
-            self,
-            mock_aioresponse,
-            models
-    ):
-        # GIVEN
-        mock_aioresponse.post(
-            url=re.compile(f"^{AzureOpenAIEnv.OPENAI_API_BASE}.*$"),
-            exception=TimeoutError("Request timed out!"),
-            repeat=True
-        )
-
-        # WHEN
-        responses = models["azure_open_ai"].generate("Some prompt")
-
-        # THEN
-        assert responses[0].response is None
-        assert "Request timed out" in responses[0].error
-
-
-