haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

haiku/rag/embeddings/voyageai.py

@@ -1,27 +1,182 @@
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Literal, cast
+
+from pydantic_ai.embeddings.base import EmbeddingModel
+from pydantic_ai.embeddings.result import EmbeddingResult, EmbedInputType
+from pydantic_ai.embeddings.settings import EmbeddingSettings
+from pydantic_ai.exceptions import ModelAPIError
+from pydantic_ai.usage import RequestUsage
+
 try:
-    from typing import overload
+    from voyageai.client_async import AsyncClient
+    from voyageai.error import VoyageError
+except ImportError as _import_error:
+    raise ImportError(
+        "Please install `voyageai` to use the VoyageAI embeddings model, "
+        "you can use — `pip install voyageai`"
+    ) from _import_error
+
+LatestVoyageAIEmbeddingModelNames = Literal[
+    "voyage-3-large",
+    "voyage-3.5",
+    "voyage-3.5-lite",
+    "voyage-code-3",
+    "voyage-finance-2",
+    "voyage-law-2",
+    "voyage-code-2",
+]
+"""Latest VoyageAI embedding models.
+
+See [VoyageAI Embeddings](https://docs.voyageai.com/docs/embeddings)
+for available models and their capabilities.
+"""
+
+VoyageAIEmbeddingModelName = str | LatestVoyageAIEmbeddingModelNames
+"""Possible VoyageAI embedding model names."""
+
+
+class VoyageAIEmbeddingSettings(EmbeddingSettings, total=False):
+    """Settings used for a VoyageAI embedding model request.
+
+    All fields from [`EmbeddingSettings`][pydantic_ai.embeddings.EmbeddingSettings] are supported,
+    plus VoyageAI-specific settings prefixed with `voyageai_`.
+    """
+
+    # ALL FIELDS MUST BE `voyageai_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+
+    voyageai_truncation: bool
+    """Whether to truncate inputs that exceed the model's context length.
+
+    Defaults to True. If False, an error is raised for inputs that are too long.
+    """
+
+    voyageai_output_dtype: Literal["float", "int8", "uint8", "binary", "ubinary"]
+    """The output data type for embeddings.
+
+    - `'float'` (default): 32-bit floats
+    - `'int8'`: Signed 8-bit integers (quantized)
+    - `'uint8'`: Unsigned 8-bit integers (quantized)
+    - `'binary'`: Binary embeddings
+    - `'ubinary'`: Unsigned binary embeddings
+    """
+
+
+_MAX_INPUT_TOKENS: dict[VoyageAIEmbeddingModelName, int] = {
+    "voyage-3-large": 32000,
+    "voyage-3.5": 32000,
+    "voyage-3.5-lite": 32000,
+    "voyage-code-3": 32000,
+    "voyage-finance-2": 32000,
+    "voyage-law-2": 16000,
+    "voyage-code-2": 16000,
+}
+
+
+@dataclass(init=False)
+class VoyageAIEmbeddingModel(EmbeddingModel):
+    """VoyageAI embedding model implementation.
+
+    VoyageAI provides state-of-the-art embedding models optimized for
+    retrieval, with specialized models for code, finance, and legal domains.
+
+    Example:
+    ```python
+    from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
+
+    model = VoyageAIEmbeddingModel('voyage-3.5')
+    ```
+    """
+
+    _model_name: VoyageAIEmbeddingModelName = field(repr=False)
+    _client: AsyncClient = field(repr=False)
+
+    def __init__(
+        self,
+        model_name: VoyageAIEmbeddingModelName,
+        *,
+        api_key: str | None = None,
+        max_retries: int = 0,
+        timeout: int | None = None,
+        settings: EmbeddingSettings | None = None,
+    ):
+        """Initialize a VoyageAI embedding model.
+
+        Args:
+            model_name: The name of the VoyageAI model to use.
+                See [VoyageAI models](https://docs.voyageai.com/docs/embeddings)
+                for available options.
+            api_key: The VoyageAI API key. If not provided, uses the
+                `VOYAGE_API_KEY` environment variable.
+            max_retries: Maximum number of retries for failed requests.
+            timeout: Request timeout in seconds.
+            settings: Model-specific [`EmbeddingSettings`][pydantic_ai.embeddings.EmbeddingSettings]
+                to use as defaults for this model.
+        """
+        self._model_name = model_name
+        self._client = AsyncClient(
+            api_key=api_key,
+            max_retries=max_retries,
+            timeout=timeout,
+        )
+
+        super().__init__(settings=settings)
+
+    @property
+    def model_name(self) -> VoyageAIEmbeddingModelName:
+        """The embedding model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The embedding model provider."""
+        return "voyageai"
+
+    async def embed(
+        self,
+        inputs: str | Sequence[str],
+        *,
+        input_type: EmbedInputType,
+        settings: EmbeddingSettings | None = None,
+    ) -> EmbeddingResult:
+        inputs, settings = self.prepare_embed(inputs, settings)
+        settings = cast(VoyageAIEmbeddingSettings, settings)
+
+        voyageai_input_type = "document" if input_type == "document" else "query"
 
-    from voyageai.client import Client  # type: ignore
+        try:
+            response = await self._client.embed(
+                texts=list(inputs),
+                model=self.model_name,
+                input_type=voyageai_input_type,
+                truncation=settings.get("voyageai_truncation", True),
+                output_dtype=settings.get("voyageai_output_dtype", "float"),
+                output_dimension=settings.get("dimensions"),
+            )
+        except VoyageError as e:
+            raise ModelAPIError(model_name=self.model_name, message=str(e)) from e
 
-    from haiku.rag.embeddings.base import EmbedderBase
+        return EmbeddingResult(
+            embeddings=response.embeddings,
+            inputs=inputs,
+            input_type=input_type,
+            usage=_map_usage(response.total_tokens, self.model_name),
+            model_name=self.model_name,
+            provider_name=self.system,
+        )
 
-    class Embedder(EmbedderBase):
-        @overload
-        async def embed(self, text: str) -> list[float]: ...
+    async def max_input_tokens(self) -> int | None:
+        return _MAX_INPUT_TOKENS.get(self.model_name)
 
-        @overload
-        async def embed(self, text: list[str]) -> list[list[float]]: ...
 
-        async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
-            client = Client()
-            if not text:
-                return []
-            if isinstance(text, str):
-                res = client.embed([text], model=self._model, output_dtype="float")
-                return res.embeddings[0]  # type: ignore[return-value]
-            else:
-                res = client.embed(text, model=self._model, output_dtype="float")
-                return res.embeddings  # type: ignore[return-value]
+def _map_usage(total_tokens: int, model: str) -> RequestUsage:
+    usage_data = {"total_tokens": total_tokens}
+    response_data = {"model": model, "usage": usage_data}
 
-except ImportError:
-    pass
+    return RequestUsage.extract(
+        response_data,
+        provider="voyageai",
+        provider_url="https://api.voyageai.com",
+        provider_fallback="voyageai",
+        api_flavor="embeddings",
+    )

haiku/rag/graph/__init__.py

@@ -1,25 +1,14 @@
-"""Graph module for haiku.rag.
-
-This module contains all graph-related functionality including:
-- AG-UI protocol for graph streaming
-- Common graph utilities and models
-- Research graph implementation
-- Deep QA graph implementation
-"""
-
 from haiku.rag.graph.agui import (
     AGUIConsoleRenderer,
     AGUIEmitter,
     create_agui_server,
     stream_graph,
 )
-from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
 from haiku.rag.graph.research.graph import build_research_graph
 
 __all__ = [
     "AGUIConsoleRenderer",
     "AGUIEmitter",
-    "build_deep_qa_graph",
     "build_research_graph",
     "create_agui_server",
     "stream_graph",

haiku/rag/graph/agui/__init__.py

@@ -1,8 +1,8 @@
 """Generic AG-UI protocol support for haiku.rag graphs."""
 
 from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
-from haiku.rag.graph.agui.emitter import AGUIEmitter
-from haiku.rag.graph.agui.events import (
+from haiku.rag.graph.agui.emitter import (
+    AGUIEmitter,
     AGUIEvent,
     emit_activity,
     emit_activity_delta,
@@ -17,6 +17,9 @@ from haiku.rag.graph.agui.events import (
     emit_text_message_content,
     emit_text_message_end,
     emit_text_message_start,
+    emit_tool_call_args,
+    emit_tool_call_end,
+    emit_tool_call_start,
 )
 from haiku.rag.graph.agui.server import (
     RunAgentInput,
@@ -48,6 +51,9 @@ __all__ = [
     "emit_text_message_content",
     "emit_text_message_end",
     "emit_text_message_start",
+    "emit_tool_call_args",
+    "emit_tool_call_end",
+    "emit_tool_call_start",
     "format_sse_event",
     "stream_graph",
 ]

haiku/rag/graph/agui/cli_renderer.py

@@ -5,7 +5,7 @@ from typing import Any
 
 from rich.console import Console
 
-from haiku.rag.graph.agui.events import AGUIEvent
+from haiku.rag.graph.agui.emitter import AGUIEvent
 
 
 class AGUIConsoleRenderer:

haiku/rag/graph/agui/emitter.py

@@ -2,23 +2,39 @@
 
 import asyncio
 import hashlib
+import json
 from collections.abc import AsyncIterator
+from typing import Any
 from uuid import uuid4
 
+from ag_ui.core import (
+    ActivitySnapshotEvent,
+    BaseEvent,
+    RunErrorEvent,
+    RunFinishedEvent,
+    RunStartedEvent,
+    StateDeltaEvent,
+    StateSnapshotEvent,
+    StepFinishedEvent,
+    StepStartedEvent,
+    TextMessageChunkEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
+)
 from pydantic import BaseModel
 
-from haiku.rag.graph.agui.events import (
-    AGUIEvent,
-    emit_activity,
-    emit_run_error,
-    emit_run_finished,
-    emit_run_started,
-    emit_state_delta,
-    emit_state_snapshot,
-    emit_step_finished,
-    emit_step_started,
-    emit_text_message,
-)
+from haiku.rag.graph.agui.state import compute_state_delta
+
+AGUIEvent = dict[str, Any]
+
+
+def _serialize_event(event: BaseEvent) -> AGUIEvent:
+    """Serialize an ag_ui event to a dict with camelCase keys."""
+    return event.model_dump(mode="json", by_alias=True, exclude_none=True)
 
 
 class AGUIEmitter[StateT: BaseModel, ResultT]:
@@ -54,7 +70,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         self._thread_id = thread_id or str(uuid4())
         self._run_id = run_id or str(uuid4())
         self._last_state: StateT | None = None
-        self._current_step: str | None = None
+        self._active_steps: set[str] = set()
         self._use_deltas = use_deltas
 
     @property
@@ -79,8 +95,14 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
             self._thread_id = self._generate_thread_id(state_json)
 
         # RunStarted (state snapshot follows immediately with full state)
-        self._emit(emit_run_started(self._thread_id, self._run_id))
-        self._emit(emit_state_snapshot(initial_state))
+        self.emit(
+            _serialize_event(
+                RunStartedEvent(thread_id=self._thread_id, run_id=self._run_id)
+            )
+        )
+        self.emit(
+            _serialize_event(StateSnapshotEvent(snapshot=initial_state.model_dump()))
+        )
         # Store a deep copy to detect future changes
         self._last_state = initial_state.model_copy(deep=True)
 
@@ -90,14 +112,17 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         Args:
             step_name: Name of the step being started
         """
-        self._current_step = step_name
-        self._emit(emit_step_started(step_name))
+        self._active_steps.add(step_name)
+        self.emit(_serialize_event(StepStartedEvent(step_name=step_name)))
 
-    def finish_step(self) -> None:
-        """Emit StepFinished event for the current step."""
-        if self._current_step:
-            self._emit(emit_step_finished(self._current_step))
-            self._current_step = None
+    def finish_step(self, step_name: str) -> None:
+        """Emit StepFinished event for the specified step.
+
+        Args:
+            step_name: Name of the step being finished
+        """
+        self._active_steps.discard(step_name)
+        self.emit(_serialize_event(StepFinishedEvent(step_name=step_name)))
 
     def log(self, message: str, role: str = "assistant") -> None:
         """Emit a text message event.
@@ -106,7 +131,16 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
             message: The message content
             role: The role of the sender (default: assistant)
         """
-        self._emit(emit_text_message(message, role))
+        message_id = str(uuid4())
+        self.emit(
+            _serialize_event(
+                TextMessageChunkEvent(
+                    message_id=message_id,
+                    role=role,  # type: ignore[arg-type]
+                    delta=message,
+                )
+            )
+        )
 
     def update_state(self, new_state: StateT) -> None:
         """Emit StateDelta or StateSnapshot for state change.
@@ -116,26 +150,40 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         """
         if self._use_deltas and self._last_state is not None:
             # Emit delta for incremental updates
-            self._emit(emit_state_delta(self._last_state, new_state))
+            delta = compute_state_delta(self._last_state, new_state)
+            self.emit(_serialize_event(StateDeltaEvent(delta=delta)))
         else:
             # Emit full snapshot for initial state or when deltas disabled
-            self._emit(emit_state_snapshot(new_state))
+            self.emit(
+                _serialize_event(StateSnapshotEvent(snapshot=new_state.model_dump()))
+            )
         # Store a deep copy to detect future changes
         self._last_state = new_state.model_copy(deep=True)
 
     def update_activity(
-        self, activity_type: str, content: str, message_id: str | None = None
+        self,
+        activity_type: str,
+        content: dict[str, Any],
+        message_id: str | None = None,
     ) -> None:
         """Emit ActivitySnapshot event.
 
         Args:
             activity_type: Type of activity (e.g., "planning", "searching")
-            content: Description of the activity
+            content: Structured payload representing the activity state
             message_id: Optional message ID to associate activity with (auto-generated if None)
         """
         if message_id is None:
             message_id = str(uuid4())
-        self._emit(emit_activity(message_id, activity_type, content))
+        self.emit(
+            _serialize_event(
+                ActivitySnapshotEvent(
+                    message_id=message_id,
+                    activity_type=activity_type,
+                    content=content,
+                )
+            )
+        )
 
     def finish_run(self, result: ResultT) -> None:
         """Emit RunFinished event.
@@ -143,7 +191,18 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         Args:
             result: The final result from the graph
         """
-        self._emit(emit_run_finished(self._thread_id, self._run_id, result))
+        # Convert result to dict if it's a Pydantic model
+        result_data: Any = result
+        if hasattr(result, "model_dump"):
+            result_data = result.model_dump()  # type: ignore[union-attr]
+
+        self.emit(
+            _serialize_event(
+                RunFinishedEvent(
+                    thread_id=self._thread_id, run_id=self._run_id, result=result_data
+                )
+            )
+        )
 
     def error(self, error: Exception, code: str | None = None) -> None:
         """Emit RunError event.
@@ -152,9 +211,9 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
             error: The exception that occurred
             code: Optional error code
         """
-        self._emit(emit_run_error(str(error), code))
+        self.emit(_serialize_event(RunErrorEvent(message=str(error), code=code)))
 
-    def _emit(self, event: AGUIEvent) -> None:
+    def emit(self, event: AGUIEvent) -> None:
         """Put event in queue.
 
         Args:
@@ -195,3 +254,132 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         # Use hash of input for deterministic thread ID
         hash_obj = hashlib.sha256(input_data.encode("utf-8"))
         return hash_obj.hexdigest()[:16]
+
+
+def emit_text_message_start(message_id: str, role: str = "assistant") -> AGUIEvent:
+    """Create a TextMessageStart event."""
+    return _serialize_event(
+        TextMessageStartEvent(message_id=message_id, role=role)  # type: ignore[arg-type]
+    )
+
+
+def emit_text_message_content(message_id: str, delta: str) -> AGUIEvent:
+    """Create a TextMessageContent event."""
+    return _serialize_event(TextMessageContentEvent(message_id=message_id, delta=delta))
+
+
+def emit_text_message_end(message_id: str) -> AGUIEvent:
+    """Create a TextMessageEnd event."""
+    return _serialize_event(TextMessageEndEvent(message_id=message_id))
+
+
+def emit_tool_call_start(
+    tool_call_id: str,
+    tool_name: str,
+    parent_message_id: str | None = None,
+) -> AGUIEvent:
+    """Create a ToolCallStart event."""
+    return _serialize_event(
+        ToolCallStartEvent(
+            tool_call_id=tool_call_id,
+            tool_call_name=tool_name,
+            parent_message_id=parent_message_id,
+        )
+    )
+
+
+def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> AGUIEvent:
+    """Create a ToolCallArgs event."""
+    return _serialize_event(
+        ToolCallArgsEvent(tool_call_id=tool_call_id, delta=json.dumps(args))
+    )
+
+
+def emit_tool_call_end(tool_call_id: str) -> AGUIEvent:
+    """Create a ToolCallEnd event."""
+    return _serialize_event(ToolCallEndEvent(tool_call_id=tool_call_id))
+
+
+def emit_run_started(thread_id: str, run_id: str) -> AGUIEvent:
+    """Create a RunStarted event."""
+    return _serialize_event(RunStartedEvent(thread_id=thread_id, run_id=run_id))
+
+
+def emit_run_finished(thread_id: str, run_id: str, result: Any) -> AGUIEvent:
+    """Create a RunFinished event."""
+    # Convert result to dict if it's a Pydantic model
+    if hasattr(result, "model_dump"):
+        result = result.model_dump()
+    return _serialize_event(
+        RunFinishedEvent(thread_id=thread_id, run_id=run_id, result=result)
+    )
+
+
+def emit_run_error(message: str, code: str | None = None) -> AGUIEvent:
+    """Create a RunError event."""
+    return _serialize_event(RunErrorEvent(message=message, code=code))
+
+
+def emit_step_started(step_name: str) -> AGUIEvent:
+    """Create a StepStarted event."""
+    return _serialize_event(StepStartedEvent(step_name=step_name))
+
+
+def emit_step_finished(step_name: str) -> AGUIEvent:
+    """Create a StepFinished event."""
+    return _serialize_event(StepFinishedEvent(step_name=step_name))
+
+
+def emit_text_message(content: str, role: str = "assistant") -> AGUIEvent:
+    """Create a TextMessageChunk event (convenience wrapper)."""
+    message_id = str(uuid4())
+    return _serialize_event(
+        TextMessageChunkEvent(
+            message_id=message_id,
+            role=role,  # type: ignore[arg-type]
+            delta=content,
+        )
+    )
+
+
+def emit_state_snapshot(state: BaseModel) -> AGUIEvent:
+    """Create a StateSnapshot event."""
+    return _serialize_event(StateSnapshotEvent(snapshot=state.model_dump()))
+
+
+def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> AGUIEvent:
+    """Create a StateDelta event with JSON Patch operations."""
+    delta = compute_state_delta(old_state, new_state)
+    return _serialize_event(StateDeltaEvent(delta=delta))
+
+
+def emit_activity(
+    message_id: str,
+    activity_type: str,
+    content: dict[str, Any],
+) -> AGUIEvent:
+    """Create an ActivitySnapshot event."""
+    return _serialize_event(
+        ActivitySnapshotEvent(
+            message_id=message_id,
+            activity_type=activity_type,
+            content=content,
+        )
+    )
+
+
+def emit_activity_delta(
+    message_id: str,
+    activity_type: str,
+    patch: list[dict[str, Any]],
+) -> AGUIEvent:
+    """Create an ActivityDelta event with JSON Patch operations."""
+    from ag_ui.core import ActivityDeltaEvent
+
+    return _serialize_event(
+        ActivityDeltaEvent(
+            message_id=message_id,
+            activity_type=activity_type,
+            patch=patch,
+        )
+    )