llama-cpp-haystack 1.0.0__tar.gz → 1.2.0__tar.gz

{llama_cpp_haystack-1.0.0 → llama_cpp_haystack-1.2.0}/CHANGELOG.md

@@ -1,5 +1,37 @@
 # Changelog
 
+## [integrations/llama_cpp-v1.1.0] - 2025-06-19
+
+### 🐛 Bug Fixes
+
+- Fix llama.cpp types; add py.typed; Toolset support (#1973)
+
+### 🧪 Testing
+
+- Test llama.cpp with python 3.12 (#1601)
+
+### ⚙️ CI
+
+- Review testing workflows (#1541)
+
+### 🧹 Chores
+
+- Remove Python 3.8 support (#1421)
+- Use Haystack logging across integrations (#1484)
+- Update ChatGenerators with `deserialize_tools_or_toolset_inplace` (#1623)
+- Align core-integrations Hatch scripts (#1898)
+
+### 🌀 Miscellaneous
+
+- Chore: remove `jsonschema` dependency from `default` environment (#1368)
+
+## [integrations/llama_cpp-v1.0.0] - 2025-02-07
+
+### 🚀 Features
+
+- [**breaking**] Llama.cpp - unified support for tools + refactoring (#1357)
+
+
 ## [integrations/llama_cpp-v0.4.4] - 2025-01-16
 
 ### 🧹 Chores

{llama_cpp_haystack-1.0.0 → llama_cpp_haystack-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llama-cpp-haystack
-Version: 1.0.0
+Version: 1.2.0
 Summary: An integration between the llama.cpp LLM framework and Haystack
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/llama_cpp#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
@@ -12,15 +12,15 @@ License-File: LICENSE.txt
 Classifier: Development Status :: 4 - Beta
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
-Requires-Python: >=3.8
-Requires-Dist: haystack-ai>=2.9.0
+Requires-Python: >=3.9
+Requires-Dist: haystack-ai>=2.13.0
 Requires-Dist: llama-cpp-python>=0.2.87
 Description-Content-Type: text/markdown
 

{llama_cpp_haystack-1.0.0 → llama_cpp_haystack-1.2.0}/pyproject.toml

@@ -7,7 +7,7 @@ name = "llama-cpp-haystack"
 dynamic = ["version"]
 description = 'An integration between the llama.cpp LLM framework and Haystack'
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 license = "Apache-2.0"
 keywords = []
 authors = [
@@ -18,15 +18,15 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Development Status :: 4 - Beta",
     "Programming Language :: Python",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.9.0", "llama-cpp-python>=0.2.87"]
+dependencies = ["haystack-ai>=2.13.0", "llama-cpp-python>=0.2.87"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/llama_cpp#readme"
@@ -46,36 +46,35 @@ git_describe_command = 'git describe --tags --match="integrations/llama_cpp-v[0-
 
 [tool.hatch.envs.default]
 installer = "uv"
+dependencies = ["haystack-pydoc-tools", "ruff"]
+[tool.hatch.envs.default.scripts]
+docs = ["pydoc-markdown pydoc/config.yml"]
+fmt = "ruff check --fix {args} && ruff format {args}"
+fmt-check = "ruff check {args} && ruff format --check {args}"
+
+[tool.hatch.envs.test]
 dependencies = [
-    "coverage[toml]>=6.5",
     "pytest",
+    "pytest-asyncio",
+    "pytest-cov",
     "pytest-rerunfailures",
-    "haystack-pydoc-tools",
-    "transformers[sentencepiece]",
-    "jsonschema",  # needed for Tool
+    "mypy",
+    "pip",
+    "transformers[sentencepiece]"
 ]
-[tool.hatch.envs.default.scripts]
-test = "pytest {args:tests}"
-test-cov = "coverage run -m pytest {args:tests}"
-test-cov-retry = "test-cov --reruns 3 --reruns-delay 30 -x"
-cov-report = ["- coverage combine", "coverage report"]
-cov = ["test-cov", "cov-report"]
-cov-retry = ["test-cov-retry", "cov-report"]
-docs = ["pydoc-markdown pydoc/config.yml"]
-[[tool.hatch.envs.all.matrix]]
-python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
 
+[tool.hatch.envs.test.scripts]
+unit = 'pytest -m "not integration" {args:tests}'
+integration = 'pytest -m "integration" {args:tests}'
+all = 'pytest {args:tests}'
+cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
+types = "mypy -p haystack_integrations.components.generators.llama_cpp {args}"
 
-[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}"
-style = ["ruff check {args:.}", "black --check --diff {args:.}"]
-fmt = ["black {args:.}", "ruff check --fix {args:.}", "style"]
-all = ["style", "typing"]
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 [tool.hatch.metadata]
 allow-direct-references = true
@@ -83,10 +82,6 @@ allow-direct-references = true
 [tool.ruff.lint.isort]
 known-first-party = ["haystack_integrations"]
 
-[tool.black]
-target-version = ["py38"]
-line-length = 120
-skip-string-normalization = true
 
 [tool.ruff]
 target-version = "py38"
@@ -165,7 +160,3 @@ markers = [
     "integration: marks tests as slow (deselect with '-m \"not integration\"')",
 ]
 addopts = ["--import-mode=importlib"]
-
-[[tool.mypy.overrides]]
-module = ["haystack.*", "haystack_integrations.*", "pytest.*", "llama_cpp.*"]
-ignore_missing_imports = true

llama_cpp_haystack-1.2.0/src/haystack_integrations/components/generators/llama_cpp/chat/chat_generator.py

@@ -0,0 +1,436 @@
+import json
+from datetime import datetime, timezone
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from haystack import component, default_from_dict, default_to_dict, logging
+from haystack.components.generators.utils import _convert_streaming_chunks_to_chat_message
+from haystack.dataclasses import (
+    ChatMessage,
+    ComponentInfo,
+    StreamingCallbackT,
+    ToolCall,
+    ToolCallDelta,
+    select_streaming_callback,
+)
+from haystack.dataclasses.streaming_chunk import FinishReason, StreamingChunk, SyncStreamingCallbackT
+from haystack.tools import (
+    Tool,
+    Toolset,
+    _check_duplicate_tool_names,
+    deserialize_tools_or_toolset_inplace,
+    serialize_tools_or_toolset,
+)
+from haystack.utils import deserialize_callable, serialize_callable
+from llama_cpp import (
+    ChatCompletionMessageToolCall,
+    ChatCompletionRequestAssistantMessage,
+    ChatCompletionRequestMessage,
+    ChatCompletionResponseChoice,
+    ChatCompletionTool,
+    CreateChatCompletionResponse,
+    CreateChatCompletionStreamResponse,
+    Llama,
+)
+from llama_cpp.llama_tokenizer import LlamaHFTokenizer
+
+logger = logging.getLogger(__name__)
+
+FINISH_REASON_MAPPING: Dict[str, FinishReason] = {
+    "stop": "stop",
+    "length": "length",
+    "tool_calls": "tool_calls",
+    "function_call": "tool_calls",
+}
+
+
+def _convert_message_to_llamacpp_format(message: ChatMessage) -> ChatCompletionRequestMessage:
+    """
+    Convert a ChatMessage to the format expected by llama.cpp Chat API.
+    """
+    text_contents = message.texts
+    tool_calls = message.tool_calls
+    tool_call_results = message.tool_call_results
+
+    if not text_contents and not tool_calls and not tool_call_results:
+        msg = "A `ChatMessage` must contain at least one `TextContent`, `ToolCall`, or `ToolCallResult`."
+        raise ValueError(msg)
+    elif len(text_contents) + len(tool_call_results) > 1:
+        msg = "A `ChatMessage` can only contain one `TextContent` or one `ToolCallResult`."
+        raise ValueError(msg)
+
+    role = message._role.value
+
+    if role == "tool" and tool_call_results:
+        if tool_call_results[0].origin.id is None:
+            msg = "`ToolCall` must have a non-null `id` attribute to be used with llama.cpp."
+            raise ValueError(msg)
+        return {
+            "role": "function",
+            "content": tool_call_results[0].result,
+            "name": tool_call_results[0].origin.tool_name,
+        }
+
+    if role == "system":
+        content = text_contents[0] if text_contents else None
+        return {"role": "system", "content": content}
+
+    if role == "user":
+        content = text_contents[0] if text_contents else None
+        return {"role": "user", "content": content}
+
+    if role == "assistant":
+        result: ChatCompletionRequestAssistantMessage = {"role": "assistant"}
+
+        if text_contents:
+            result["content"] = text_contents[0]
+
+        if tool_calls:
+            llamacpp_tool_calls: List[ChatCompletionMessageToolCall] = []
+            for tc in tool_calls:
+                if tc.id is None:
+                    msg = "`ToolCall` must have a non-null `id` attribute to be used with llama.cpp."
+                    raise ValueError(msg)
+                llamacpp_tool_calls.append(
+                    {
+                        "id": tc.id,
+                        "type": "function",
+                        # We disable ensure_ascii so special chars like emojis are not converted
+                        "function": {"name": tc.tool_name, "arguments": json.dumps(tc.arguments, ensure_ascii=False)},
+                    }
+                )
+            result["tool_calls"] = llamacpp_tool_calls
+
+        return result
+
+    error_msg = f"Unknown role: {role}"
+    raise ValueError(error_msg)
+
+
+@component
+class LlamaCppChatGenerator:
+    """
+    Provides an interface to generate text using LLM via llama.cpp.
+
+    [llama.cpp](https://github.com/ggml-org/llama.cpp) is a project written in C/C++ for efficient inference of LLMs.
+    It employs the quantized GGUF format, suitable for running these models on standard machines (even without GPUs).
+
+    Usage example:
+    ```python
+    from haystack_integrations.components.generators.llama_cpp import LlamaCppChatGenerator
+    user_message = [ChatMessage.from_user("Who is the best American actor?")]
+    generator = LlamaCppGenerator(model="zephyr-7b-beta.Q4_0.gguf", n_ctx=2048, n_batch=512)
+
+    print(generator.run(user_message, generation_kwargs={"max_tokens": 128}))
+    # {"replies": [ChatMessage(content="John Cusack", role=<ChatRole.ASSISTANT: "assistant">, name=None, meta={...}]}
+    ```
+    """
+
+    def __init__(
+        self,
+        model: str,
+        n_ctx: Optional[int] = 0,
+        n_batch: Optional[int] = 512,
+        model_kwargs: Optional[Dict[str, Any]] = None,
+        generation_kwargs: Optional[Dict[str, Any]] = None,
+        *,
+        tools: Optional[Union[List[Tool], Toolset]] = None,
+        streaming_callback: Optional[StreamingCallbackT] = None,
+    ):
+        """
+        :param model: The path of a quantized model for text generation, for example, "zephyr-7b-beta.Q4_0.gguf".
+            If the model path is also specified in the `model_kwargs`, this parameter will be ignored.
+        :param n_ctx: The number of tokens in the context. When set to 0, the context will be taken from the model.
+        :param n_batch: Prompt processing maximum batch size.
+        :param model_kwargs: Dictionary containing keyword arguments used to initialize the LLM for text generation.
+            These keyword arguments provide fine-grained control over the model loading.
+            In case of duplication, these kwargs override `model`, `n_ctx`, and `n_batch` init parameters.
+            For more information on the available kwargs, see
+            [llama.cpp documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.__init__).
+        :param generation_kwargs:  A dictionary containing keyword arguments to customize text generation.
+            For more information on the available kwargs, see
+            [llama.cpp documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion).
+        :param tools:
+            A list of tools or a Toolset for which the model can prepare calls.
+            This parameter can accept either a list of `Tool` objects or a `Toolset` instance.
+        :param streaming_callback: A callback function that is called when a new token is received from the stream.
+        """
+
+        model_kwargs = model_kwargs or {}
+        generation_kwargs = generation_kwargs or {}
+
+        # check if the model_kwargs contain the essential parameters
+        # otherwise, populate them with values from init parameters
+        model_kwargs.setdefault("model_path", model)
+        model_kwargs.setdefault("n_ctx", n_ctx)
+        model_kwargs.setdefault("n_batch", n_batch)
+
+        _check_duplicate_tool_names(list(tools or []))
+
+        self.model_path = model
+        self.n_ctx = n_ctx
+        self.n_batch = n_batch
+        self.model_kwargs = model_kwargs
+        self.generation_kwargs = generation_kwargs
+        self._model: Optional[Llama] = None
+        self.tools = tools
+        self.streaming_callback = streaming_callback
+
+    def warm_up(self):
+        if "hf_tokenizer_path" in self.model_kwargs and "tokenizer" not in self.model_kwargs:
+            tokenizer = LlamaHFTokenizer.from_pretrained(self.model_kwargs["hf_tokenizer_path"])
+            self.model_kwargs["tokenizer"] = tokenizer
+
+        if self._model is None:
+            self._model = Llama(**self.model_kwargs)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+              Dictionary with serialized data.
+        """
+        callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None
+        return default_to_dict(
+            self,
+            model=self.model_path,
+            n_ctx=self.n_ctx,
+            n_batch=self.n_batch,
+            model_kwargs=self.model_kwargs,
+            generation_kwargs=self.generation_kwargs,
+            tools=serialize_tools_or_toolset(self.tools),
+            streaming_callback=callback_name,
+        )
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "LlamaCppChatGenerator":
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            Dictionary to deserialize from.
+        :returns:
+            Deserialized component.
+        """
+        deserialize_tools_or_toolset_inplace(data["init_parameters"], key="tools")
+        if (
+            "streaming_callback" in data["init_parameters"]
+            and data["init_parameters"]["streaming_callback"] is not None
+        ):
+            data["init_parameters"]["streaming_callback"] = deserialize_callable(
+                data["init_parameters"]["streaming_callback"]
+            )
+        return default_from_dict(cls, data)
+
+    @component.output_types(replies=List[ChatMessage])
+    def run(
+        self,
+        messages: List[ChatMessage],
+        generation_kwargs: Optional[Dict[str, Any]] = None,
+        *,
+        tools: Optional[Union[List[Tool], Toolset]] = None,
+        streaming_callback: Optional[StreamingCallbackT] = None,
+    ) -> Dict[str, List[ChatMessage]]:
+        """
+        Run the text generation model on the given list of ChatMessages.
+
+        :param messages:
+            A list of ChatMessage instances representing the input messages.
+        :param generation_kwargs:  A dictionary containing keyword arguments to customize text generation.
+            For more information on the available kwargs, see
+            [llama.cpp documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion).
+        :param tools:
+            A list of tools or a Toolset for which the model can prepare calls. If set, it will override the `tools`
+            parameter set during component initialization.
+        :param streaming_callback: A callback function that is called when a new token is received from the stream.
+            If set, it will override the `streaming_callback` parameter set during component initialization.
+        :returns: A dictionary with the following keys:
+            - `replies`: The responses from the model
+        """
+        if self._model is None:
+            error_msg = "The model has not been loaded. Please call warm_up() before running."
+            raise RuntimeError(error_msg)
+
+        if not messages:
+            return {"replies": []}
+
+        updated_generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
+        formatted_messages = [_convert_message_to_llamacpp_format(msg) for msg in messages]
+
+        tools = tools or self.tools
+        if isinstance(tools, Toolset):
+            tools = list(tools)
+        _check_duplicate_tool_names(tools)
+
+        llamacpp_tools: List[ChatCompletionTool] = []
+        if tools:
+            for t in tools:
+                llamacpp_tools.append(
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": t.tool_spec["name"],
+                            "description": t.tool_spec.get("description", ""),
+                            "parameters": t.tool_spec.get("parameters", {}),
+                        },
+                    }
+                )
+
+        streaming_callback = select_streaming_callback(
+            init_callback=self.streaming_callback,
+            runtime_callback=streaming_callback,
+            requires_async=False,
+        )
+
+        if streaming_callback:
+            response_stream = self._model.create_chat_completion(
+                messages=formatted_messages, tools=llamacpp_tools, **updated_generation_kwargs, stream=True
+            )
+            return self._handle_streaming_response(
+                response_stream=response_stream,  # type: ignore[arg-type]
+                streaming_callback=streaming_callback,
+                component_info=ComponentInfo.from_component(self),
+            )  # we know that response_stream is Iterator[CreateChatCompletionStreamResponse]
+            # because create_chat_completion was called with stream=True, but mypy doesn't know that
+
+        response = self._model.create_chat_completion(
+            messages=formatted_messages, tools=llamacpp_tools, **updated_generation_kwargs
+        )
+        replies = []
+        if not isinstance(response, dict):
+            msg = f"Expected a dictionary response, got a different object: {response}"
+            raise ValueError(msg)
+
+        for choice in response["choices"]:
+            chat_message = self._convert_chat_completion_choice_to_chat_message(choice, response)
+            replies.append(chat_message)
+        return {"replies": replies}
+
+    @staticmethod
+    def _handle_streaming_response(
+        response_stream: Iterator[CreateChatCompletionStreamResponse],
+        streaming_callback: SyncStreamingCallbackT,
+        component_info: ComponentInfo,
+    ) -> Dict[str, List[ChatMessage]]:
+        """
+        Take streaming responses from llama.cpp, convert to Haystack StreamingChunk objects, stream them,
+        and finally convert them to a ChatMessage.
+
+        :param response_stream: The streaming responses from llama.cpp.
+        :param streaming_callback: The callback function for streaming chunks.
+        :param component_info: The component info.
+        :returns: A dictionary with the replies.
+        """
+        streaming_chunks = []
+
+        seen_tool_call_ids = set()  # Track tool call IDs we've seen
+
+        for i, chunk in enumerate(response_stream):
+            content = ""
+            tool_calls = []
+            mapped_finish_reason = None
+
+            # Track new tool call IDs in this chunk.
+            # Considering tool call ID is the only reliable way to recognize tool calls in llama.cpp streaming.
+            # They are often spread across multiple chunks.
+            new_tool_call_ids = set()
+
+            if chunk.get("choices") and len(chunk["choices"]) > 0:
+                choice = chunk["choices"][0]
+                delta = choice.get("delta", {})
+
+                finish_reason = choice.get("finish_reason")
+                mapped_finish_reason = FINISH_REASON_MAPPING.get(finish_reason or "")
+
+                if content_raw := delta.get("content"):
+                    content = str(content_raw)
+
+                tool_calls_data = delta.get("tool_calls")
+                if tool_calls_data is not None and isinstance(tool_calls_data, list):
+                    for tool_call_chunk in tool_calls_data:
+                        tool_call_id = tool_call_chunk.get("id")
+                        is_new_tool_call = tool_call_id and tool_call_id not in seen_tool_call_ids
+
+                        if is_new_tool_call:
+                            new_tool_call_ids.add(tool_call_id)
+                            seen_tool_call_ids.add(tool_call_id)
+
+                        function_data = tool_call_chunk.get("function", {})
+
+                        # Only include tool_name if this is a new tool call
+                        tool_name = function_data.get("name", "") if is_new_tool_call else ""
+
+                        tool_calls.append(
+                            ToolCallDelta(
+                                index=tool_call_chunk.get("index"),
+                                id=tool_call_id,
+                                tool_name=tool_name,
+                                arguments=function_data.get("arguments"),
+                            )
+                        )
+
+            # start is True if it's the first chunk or if we have new tool call IDs
+            start = i == 0 or len(new_tool_call_ids) > 0
+
+            streaming_chunk = StreamingChunk(
+                content="" if tool_calls else content,  # prioritize tool calls over content when both are present
+                tool_calls=tool_calls,
+                component_info=component_info,
+                index=i,
+                start=start,
+                finish_reason=mapped_finish_reason,
+                meta={
+                    "model": chunk["model"],
+                    "received_at": datetime.fromtimestamp(chunk["created"], tz=timezone.utc).isoformat(),
+                },  # llama.cpp does not provide usage metadata during streaming
+            )
+
+            streaming_chunks.append(streaming_chunk)
+
+            # Stream the chunk
+            try:
+                streaming_callback(streaming_chunk)
+            except Exception as e:
+                logger.error(f"Error in streaming callback invocation: {e}")
+                continue
+
+        message = _convert_streaming_chunks_to_chat_message(streaming_chunks)
+        return {"replies": [message]}
+
+    @staticmethod
+    def _convert_chat_completion_choice_to_chat_message(
+        choice: ChatCompletionResponseChoice, response: CreateChatCompletionResponse
+    ) -> ChatMessage:
+        llamacpp_message = choice["message"]
+        text_content = llamacpp_message.get("content", "") or None
+        tool_calls = []
+
+        if llamacpp_tool_calls := llamacpp_message.get("tool_calls", []):
+            for llamacpp_tc in llamacpp_tool_calls:
+                arguments_str = llamacpp_tc["function"]["arguments"]
+                try:
+                    arguments = json.loads(arguments_str)
+                    tool_calls.append(
+                        ToolCall(id=llamacpp_tc["id"], tool_name=llamacpp_tc["function"]["name"], arguments=arguments)
+                    )
+                except json.JSONDecodeError:
+                    logger.warning(
+                        "Llama.cpp returned a malformed JSON string for tool call arguments. This tool call "
+                        "will be skipped. Tool call ID: {tc_id}, Tool name: {tc_name}, Arguments: {tc_args}",
+                        tc_id=llamacpp_tc["id"],
+                        tc_name=llamacpp_tc["function"]["name"],
+                        tc_args=arguments_str,
+                    )
+
+        finish_reason = choice.get("finish_reason")
+
+        meta = {
+            "response_id": response["id"],
+            "model": response["model"],
+            "created": response["created"],
+            "index": choice["index"],
+            "finish_reason": FINISH_REASON_MAPPING.get(finish_reason or ""),
+            "usage": response["usage"],
+        }
+
+        return ChatMessage.from_assistant(text=text_content, tool_calls=tool_calls, meta=meta)

{llama_cpp_haystack-1.0.0 → llama_cpp_haystack-1.2.0}/src/haystack_integrations/components/generators/llama_cpp/generator.py

@@ -1,7 +1,6 @@
-import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
-from haystack import component
+from haystack import component, logging
 
 from llama_cpp import Llama
 
@@ -13,7 +12,7 @@ class LlamaCppGenerator:
     """
     Provides an interface to generate text using LLM via llama.cpp.
 
-    [llama.cpp](https://github.com/ggerganov/llama.cpp) is a project written in C/C++ for efficient inference of LLMs.
+    [llama.cpp](https://github.com/ggml-org/llama.cpp) is a project written in C/C++ for efficient inference of LLMs.
     It employs the quantized GGUF format, suitable for running these models on standard machines (even without GPUs).
 
     Usage example:
@@ -63,14 +62,16 @@ class LlamaCppGenerator:
         self.n_batch = n_batch
         self.model_kwargs = model_kwargs
         self.generation_kwargs = generation_kwargs
-        self.model = None
+        self.model: Optional[Llama] = None
 
     def warm_up(self):
         if self.model is None:
             self.model = Llama(**self.model_kwargs)
 
     @component.output_types(replies=List[str], meta=List[Dict[str, Any]])
-    def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None):
+    def run(
+        self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Union[List[str], List[Dict[str, Any]]]]:
         """
         Run the text generation model on the given prompt.
 
@@ -93,6 +94,10 @@ class LlamaCppGenerator:
         updated_generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
 
         output = self.model.create_completion(prompt=prompt, **updated_generation_kwargs)
+        if not isinstance(output, dict):
+            msg = f"Expected a dictionary response, got a different object: {output}"
+            raise ValueError(msg)
+
         replies = [output["choices"][0]["text"]]
 
-        return {"replies": replies, "meta": [output]}
+        return {"replies": replies, "meta": [dict(output.items())]}