diff --git a/src/neo4j_graphrag/llm/anthropic_llm.py b/src/neo4j_graphrag/llm/anthropic_llm.py index 04d6555e..614c5b9c 100644 --- a/src/neo4j_graphrag/llm/anthropic_llm.py +++ b/src/neo4j_graphrag/llm/anthropic_llm.py @@ -13,7 +13,7 @@ # limitations under the License. from __future__ import annotations -from typing import Any, Iterable, Optional, TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, Iterable, Optional, cast from pydantic import ValidationError diff --git a/src/neo4j_graphrag/llm/cohere_llm.py b/src/neo4j_graphrag/llm/cohere_llm.py index 63e54aaa..024d578a 100644 --- a/src/neo4j_graphrag/llm/cohere_llm.py +++ b/src/neo4j_graphrag/llm/cohere_llm.py @@ -15,6 +15,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, Iterable, Optional, cast + from pydantic import ValidationError from neo4j_graphrag.exceptions import LLMGenerationError diff --git a/src/neo4j_graphrag/llm/mistralai_llm.py b/src/neo4j_graphrag/llm/mistralai_llm.py index a3c84759..761b9a62 100644 --- a/src/neo4j_graphrag/llm/mistralai_llm.py +++ b/src/neo4j_graphrag/llm/mistralai_llm.py @@ -16,6 +16,7 @@ import os from typing import Any, Iterable, Optional, cast + from pydantic import ValidationError from neo4j_graphrag.exceptions import LLMGenerationError @@ -30,7 +31,7 @@ ) try: - from mistralai import Mistral, Messages + from mistralai import Messages, Mistral from mistralai.models.sdkerror import SDKError except ImportError: Mistral = None # type: ignore diff --git a/src/neo4j_graphrag/llm/ollama_llm.py b/src/neo4j_graphrag/llm/ollama_llm.py index a36d34f9..08dfa655 100644 --- a/src/neo4j_graphrag/llm/ollama_llm.py +++ b/src/neo4j_graphrag/llm/ollama_llm.py @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations -from typing import Any, Iterable, Optional, Sequence, TYPE_CHECKING, cast + +from typing import TYPE_CHECKING, Any, Iterable, Optional, Sequence, cast from pydantic import ValidationError @@ -24,9 +25,9 @@ BaseMessage, LLMMessage, LLMResponse, + MessageList, SystemMessage, UserMessage, - MessageList, ) if TYPE_CHECKING: diff --git a/src/neo4j_graphrag/llm/openai_llm.py b/src/neo4j_graphrag/llm/openai_llm.py index b3d99411..353dd19d 100644 --- a/src/neo4j_graphrag/llm/openai_llm.py +++ b/src/neo4j_graphrag/llm/openai_llm.py @@ -25,9 +25,9 @@ BaseMessage, LLMMessage, LLMResponse, + MessageList, SystemMessage, UserMessage, - MessageList, ) if TYPE_CHECKING: diff --git a/src/neo4j_graphrag/llm/types.py b/src/neo4j_graphrag/llm/types.py index 77e89aef..2d293e93 100644 --- a/src/neo4j_graphrag/llm/types.py +++ b/src/neo4j_graphrag/llm/types.py @@ -1,6 +1,7 @@ -from pydantic import BaseModel from typing import Literal, TypedDict +from pydantic import BaseModel + class LLMResponse(BaseModel): content: str diff --git a/src/neo4j_graphrag/llm/vertexai_llm.py b/src/neo4j_graphrag/llm/vertexai_llm.py index 48acfc9f..f2cbe26a 100644 --- a/src/neo4j_graphrag/llm/vertexai_llm.py +++ b/src/neo4j_graphrag/llm/vertexai_llm.py @@ -23,10 +23,10 @@ try: from vertexai.generative_models import ( + Content, GenerativeModel, - ResponseValidationError, Part, - Content, + ResponseValidationError, ) except ImportError: GenerativeModel = None diff --git a/tests/e2e/test_simplekgpipeline_e2e.py b/tests/e2e/test_simplekgpipeline_e2e.py index 6b5ec04f..fad595b0 100644 --- a/tests/e2e/test_simplekgpipeline_e2e.py +++ b/tests/e2e/test_simplekgpipeline_e2e.py @@ -111,7 +111,6 @@ async def test_pipeline_builder_happy_path( await kg_builder_text.run_async(text=harry_potter_text) - @pytest.mark.asyncio @pytest.mark.usefixtures("setup_neo4j_for_kg_construction") async def test_pipeline_builder_two_documents( diff --git a/tests/unit/llm/test_anthropic_llm.py b/tests/unit/llm/test_anthropic_llm.py index fc8d2756..b1d583f6 100644 --- a/tests/unit/llm/test_anthropic_llm.py +++ b/tests/unit/llm/test_anthropic_llm.py @@ -18,7 +18,6 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import anthropic - import pytest from neo4j_graphrag.exceptions import LLMGenerationError from neo4j_graphrag.llm.anthropic_llm import AnthropicLLM