fix: fix Ollama types + add py.typed (#1922)

* fix: fix Ollama types + add py.typed

* Update integrations/ollama/src/haystack_integrations/components/embedders/ollama/document_embedder.py

Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com>

* better typing for embedders

---------

Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com>

Author: anakin87

.github/workflows/ollama.yml

@@ -84,11 +84,9 @@ jobs:
       - name: Install Hatch
         run: pip install --upgrade hatch
 
-    # TODO: Once this integration is properly typed, use hatch run test:types
-    # https://github.com/deepset-ai/haystack-core-integrations/issues/1771
       - name: Lint
         if: matrix.python-version == '3.9' && runner.os == 'Linux'
-        run: hatch run fmt-check && hatch run lint:typing
+        run: hatch run fmt-check && hatch run test:types
 
       - name: Generate docs
         if: matrix.python-version == '3.9' && runner.os == 'Linux'

integrations/ollama/pyproject.toml

@@ -71,18 +71,14 @@ integration = 'pytest -m "integration" {args:tests}'
 all = 'pytest {args:tests}'
 cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
 
-types = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = """mypy -p haystack_integrations.components.embedders.ollama \
+-p haystack_integrations.components.generators.ollama {args}"""
 
-# TODO: remove lint environment once this integration is properly typed
-# test environment should be used instead
-# https://github.com/deepset-ai/haystack-core-integrations/issues/1771
-[tool.hatch.envs.lint]
-installer = "uv"
-detached = true
-dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
-
-[tool.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 [tool.hatch.metadata]
 allow-direct-references = true
@@ -173,7 +169,3 @@ log_cli = true
 addopts = ["--import-mode=importlib"]
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "class"
-
-[[tool.mypy.overrides]]
-module = ["haystack.*", "haystack_integrations.*", "pytest.*", "ollama.*", "pydantic.*"]
-ignore_missing_imports = true

integrations/ollama/src/haystack_integrations/components/embedders/ollama/document_embedder.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from haystack import Document, component
 from tqdm import tqdm
@@ -114,7 +114,7 @@ def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
 
     def _embed_batch(
         self, texts_to_embed: List[str], batch_size: int, generation_kwargs: Optional[Dict[str, Any]] = None
-    ):
+    ) -> List[List[float]]:
         """
         Internal method to embed a batch of texts.
         """
@@ -132,7 +132,7 @@ def _embed_batch(
 
     async def _embed_batch_async(
         self, texts_to_embed: List[str], batch_size: int, generation_kwargs: Optional[Dict[str, Any]] = None
-    ):
+    ) -> List[List[float]]:
         """
         Internal method to embed a batch of texts asynchronously.
         """
@@ -160,7 +160,9 @@ async def _embed_batch_async(
         return all_embeddings
 
     @component.output_types(documents=List[Document], meta=Dict[str, Any])
-    def run(self, documents: List[Document], generation_kwargs: Optional[Dict[str, Any]] = None):
+    def run(
+        self, documents: List[Document], generation_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Union[List[Document], Dict[str, Any]]]:
         """
         Runs an Ollama Model to compute embeddings of the provided documents.
 
@@ -193,7 +195,9 @@ def run(self, documents: List[Document], generation_kwargs: Optional[Dict[str, A
         return {"documents": documents, "meta": {"model": self.model}}
 
     @component.output_types(documents=List[Document], meta=Dict[str, Any])
-    async def run_async(self, documents: List[Document], generation_kwargs: Optional[Dict[str, Any]] = None):
+    async def run_async(
+        self, documents: List[Document], generation_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Union[List[Document], Dict[str, Any]]]:
         """
         Asynchronously run an Ollama Model to compute embeddings of the provided documents.
 

integrations/ollama/src/haystack_integrations/components/embedders/ollama/text_embedder.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from haystack import component
 
@@ -49,7 +49,9 @@ def __init__(
         self._async_client = AsyncClient(host=self.url, timeout=self.timeout)
 
     @component.output_types(embedding=List[float], meta=Dict[str, Any])
-    def run(self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None):
+    def run(
+        self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Union[List[float], Dict[str, Any]]]:
         """
         Runs an Ollama Model to compute embeddings of the provided text.
 
@@ -69,7 +71,9 @@ def run(self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None):
         return result
 
     @component.output_types(embedding=List[float], meta=Dict[str, Any])
-    async def run_async(self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None):
+    async def run_async(
+        self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Union[List[float], Dict[str, Any]]]:
         """
         Asynchronously run an Ollama Model to compute embeddings of the provided text.
 

integrations/ollama/src/haystack_integrations/components/embedders/py.typed

@@ -220,7 +220,7 @@ def __init__(
             - JSON Schema: The response is formatted as a JSON object
                 that adheres to the specified JSON Schema. (needs Ollama ≥ 0.1.34)
         """
-        _check_duplicate_tool_names(tools)
+        _check_duplicate_tool_names(list(tools or []))
 
         self.model = model
         self.url = url
@@ -361,10 +361,11 @@ def _handle_streaming_response(
         # Compose final reply
         text = "".join(c.content for c in chunks)
 
-        tool_calls = [
-            ToolCall(tool_name=name_by_id[tool_call_id], arguments=arg_by_id.get(tool_call_id))
-            for tool_call_id in id_order
-        ]
+        tool_calls = []
+        for tool_call_id in id_order:
+            arguments = arg_by_id.get(tool_call_id, {})
+            assert isinstance(arguments, dict)  # final arguments are a dictionary  # noqa: S101
+            tool_calls.append(ToolCall(tool_name=name_by_id[tool_call_id], arguments=arguments))
 
         reply = ChatMessage.from_assistant(
             text=text,
@@ -381,7 +382,7 @@ def run(
         tools: Optional[Union[List[Tool], Toolset]] = None,
         *,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
-    ):
+    ) -> Dict[str, List[ChatMessage]]:
         """
         Runs an Ollama Model on a given chat history.
 
@@ -405,7 +406,7 @@ def run(
         """
         generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
         tools = tools or self.tools
-        _check_duplicate_tool_names(tools)
+        _check_duplicate_tool_names(list(tools or []))
 
         # Convert Toolset → list[Tool] for JSON serialization
         if isinstance(tools, Toolset):
@@ -421,7 +422,7 @@ def run(
             model=self.model,
             messages=ollama_messages,
             tools=ollama_tools,
-            stream=is_stream,
+            stream=is_stream,  # type: ignore[call-overload]  # Ollama expects Literal[True] or Literal[False], not bool
             keep_alive=self.keep_alive,
             options=generation_kwargs,
             format=self.response_format,

integrations/ollama/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py

@@ -239,7 +239,7 @@ def run(
         generation_kwargs: Optional[Dict[str, Any]] = None,
         *,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
-    ):
+    ) -> Dict[str, List[Any]]:
         """
         Runs an Ollama Model on the given prompt.
 
@@ -263,7 +263,7 @@ def run(
         response = self._client.generate(
             model=self.model,
             prompt=prompt,
-            stream=stream,
+            stream=stream,  # type: ignore[call-overload]  # Ollama expects Literal[True] or Literal[False], not bool
             keep_alive=self.keep_alive,
             options=generation_kwargs,
         )