calfkit 0.1.1__py3-none-any.whl

calfkit/__init__.py

@@ -0,0 +1,4 @@
+from importlib.metadata import version
+
+__version__ = version("calfkit")
+__all__ = ["__version__"]

calfkit/broker/__init__.py

@@ -0,0 +1,3 @@
+from calfkit.broker.broker import Broker
+
+__all__ = ["Broker"]

calfkit/broker/broker.py

@@ -0,0 +1,34 @@
+import os
+from collections.abc import Iterable
+from typing import Any, Literal
+
+from faststream import FastStream
+from faststream.kafka import KafkaBroker
+
+from calfkit.broker.deployable import Deployable
+from calfkit.broker.middleware import ContextInjectionMiddleware
+
+
+class Broker(KafkaBroker, Deployable):
+    """Lightweight client wrapper connecting to Calf brokers"""
+
+    mode = Literal["kafka_mode"]
+
+    def __init__(self, bootstrap_servers: str | Iterable[str] | None = None, **broker_kwargs: Any):
+        if not bootstrap_servers:
+            bootstrap_servers = os.getenv("CALF_HOST_URL")
+        super().__init__(
+            bootstrap_servers or "localhost",
+            middlewares=[ContextInjectionMiddleware],
+            **broker_kwargs,
+        )
+
+    @property
+    def app(self) -> FastStream:
+        return FastStream(self)
+
+    async def run_app(self) -> None:
+        await self.app.run()
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self, name)

calfkit/broker/deployable.py

@@ -0,0 +1,6 @@
+from abc import ABC, abstractmethod
+
+
+class Deployable(ABC):
+    @abstractmethod
+    async def run_app(self) -> None: ...

calfkit/broker/middleware.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+from faststream import BaseMiddleware
+from faststream.message import StreamMessage
+from faststream.types import AsyncFuncAny
+
+
+class ContextInjectionMiddleware(BaseMiddleware):
+    async def consume_scope(
+        self,
+        call_next: AsyncFuncAny,
+        msg: StreamMessage[Any],
+    ) -> Any:
+        with self.context.scope("correlation_id", msg.correlation_id):
+            return await super().consume_scope(call_next, msg)

calfkit/experimental/rpc_worker.py

@@ -0,0 +1,51 @@
+from asyncio import Future, wait_for
+from uuid import uuid4
+
+from faststream.kafka import KafkaBroker, KafkaMessage
+from faststream.types import SendableMessage
+
+
+class RPCWorker:
+    responses: dict[str, Future[bytes]]
+
+    def __init__(self, broker: KafkaBroker, reply_topic: str) -> None:
+        self.responses = {}
+        self.broker = broker
+        self.reply_topic = reply_topic
+
+        self.subscriber = broker.subscriber(reply_topic)
+        self.subscriber(self._handle_responses)
+
+    async def start(self) -> None:
+        await self.subscriber.start()
+
+    async def stop(self) -> None:
+        await self.subscriber.stop()
+
+    def _handle_responses(self, msg: KafkaMessage) -> None:
+        if future := self.responses.pop(msg.correlation_id, None):
+            future.set_result(msg.body)
+
+    async def request(
+        self,
+        data: SendableMessage,
+        topic: str,
+        timeout: float = 10.0,
+    ) -> bytes:
+        correlation_id = str(uuid4())
+        future = self.responses[correlation_id] = Future[bytes]()
+
+        await self.broker.publish(
+            data,
+            topic,
+            reply_to=self.reply_topic,
+            correlation_id=correlation_id,
+        )
+
+        try:
+            response: bytes = await wait_for(future, timeout=timeout)
+        except TimeoutError:
+            self.responses.pop(correlation_id, None)
+            raise
+        else:
+            return response

calfkit/messages/__init__.py

@@ -0,0 +1,9 @@
+"""Message utilities for calf SDK.
+
+This module provides utilities for working with pydantic_ai ModelMessage types,
+including message history manipulation and transformation.
+"""
+
+from .util import patch_system_prompts, validate_tool_call_pairs
+
+__all__ = ["patch_system_prompts", "validate_tool_call_pairs"]

calfkit/messages/util.py

@@ -0,0 +1,99 @@
+"""Message utility functions for calf SDK.
+
+This module provides utilities for working with pydantic_ai ModelMessage types,
+including message history manipulation and transformation.
+"""
+
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    RetryPromptPart,
+    SystemPromptPart,
+    ToolCallPart,
+    ToolReturnPart,
+)
+
+
+def patch_system_prompts(
+    base: list[ModelMessage],
+    incoming: list[ModelMessage],
+) -> list[ModelMessage]:
+    """Patch system prompts in message history.
+
+    If incoming messages contain system prompts, they replace any existing
+    system prompts in base. System prompts are consolidated and placed at
+    the front of the history.
+
+    If incoming has no system prompts, returns base unmodified.
+
+    Args:
+        base: The existing message history to patch.
+        incoming: The new messages that may contain replacement system prompts.
+
+    Returns:
+        A new message history with system prompts patched.
+
+    Examples:
+        >>> from pydantic_ai import ModelRequest, SystemPromptPart
+        >>> base = [ModelRequest(parts=[SystemPromptPart("old system")])]
+        >>> incoming = [ModelRequest(parts=[SystemPromptPart("new system")])]
+        >>> result = patch_system_prompts(base, incoming)
+        >>> len(result)
+        1
+        >>> result[0].parts[0].content
+        'new system'
+    """
+    incoming_system_parts: list[SystemPromptPart] = []
+    for msg in incoming:
+        if isinstance(msg, ModelResponse):
+            continue
+        for part in msg.parts:
+            if isinstance(part, SystemPromptPart):
+                incoming_system_parts.append(part)
+
+    if not incoming_system_parts:
+        return base
+
+    system_msg = ModelRequest(parts=incoming_system_parts)
+    result: list[ModelMessage] = []
+    for msg in base:
+        if isinstance(msg, ModelRequest):
+            non_system_parts = [p for p in msg.parts if not isinstance(p, SystemPromptPart)]
+            if non_system_parts:
+                result.append(ModelRequest(parts=non_system_parts))
+        else:
+            result.append(msg)
+
+    return [system_msg] + result
+
+
+def validate_tool_call_pairs(messages: list[ModelMessage]) -> bool:
+    """Validate that all tool calls have corresponding tool results.
+
+    Iterates through messages in reverse order to verify that every ToolCallPart
+    has a matching ToolReturnPart or RetryPromptPart with the same tool_call_id.
+
+    The first time a tool call is found without a matching result, the function
+    returns False immediately.
+
+    Args:
+        messages: List of ModelMessage to validate.
+
+    Returns:
+        True if all tool calls have matching results, False otherwise.
+    """
+    seen_result_ids: set[str] = set()
+
+    for message in reversed(messages):
+        if isinstance(message, ModelRequest):
+            for req_part in message.parts:
+                if isinstance(req_part, (ToolReturnPart, RetryPromptPart)):
+                    seen_result_ids.add(req_part.tool_call_id)
+        elif isinstance(message, ModelResponse):
+            for resp_part in message.parts:
+                if isinstance(resp_part, ToolCallPart):
+                    if resp_part.tool_call_id not in seen_result_ids:
+                        return False
+
+    return True

calfkit/models/event_envelope.py

@@ -0,0 +1,43 @@
+from collections.abc import Sequence
+from typing import Literal
+
+from pydantic_ai import ModelMessage, ModelRequest
+from pydantic_ai.models import ModelRequestParameters
+
+from calfkit.models.types import CompactBaseModel, SerializableModelSettings, ToolCallRequest
+
+
+class EventEnvelope(CompactBaseModel):
+    kind: Literal["tool_call_request", "user_prompt", "ai_response", "tool_result"]
+    text: str | None = None
+    trace_id: str | None = None
+
+    # Used to surface the tool call from latest message so tool call workers do not have to dig
+    tool_call_request: ToolCallRequest | None = None
+
+    # Optional inference-time patch in settings and parameters
+    patch_model_request_params: ModelRequestParameters | None = None
+    patch_model_settings: SerializableModelSettings | None = None
+
+    # Running message history
+    message_history: list[ModelMessage] = []
+
+    @property
+    def latest_message_in_history(self) -> ModelMessage | None:
+        return self.message_history[-1] if self.message_history else None
+
+    # The result message from a node
+    incoming_node_messages: Sequence[ModelMessage] = []
+
+    # thread id / conversation identifier
+    thread_id: str | None = None
+
+    # Allow client to dynamically patch system message at runtime
+    # Intentionally kept separate from message_history in order to simplify patch logic
+    system_message: ModelRequest | None = None
+
+    # Where the final response from AI should be published to
+    final_response_topic: str
+
+    # Whether the current message is the final response from the AI to the user
+    final_response: bool = False

calfkit/models/types.py

@@ -0,0 +1,65 @@
+from typing import Any, TypeAlias
+
+from pydantic import BaseModel
+from pydantic_ai import ToolCallPart
+from typing_extensions import TypedDict
+
+ToolCallRequest: TypeAlias = ToolCallPart
+
+
+class SerializableModelSettings(TypedDict, total=False):
+    """Serializable version of pydantic_ai.ModelSettings.
+
+    This is a copy of ModelSettings with `timeout` narrowed
+    to `float` only to ensure JSON serializability.
+    """
+
+    max_tokens: int
+    """The maximum number of tokens to generate before stopping."""
+
+    temperature: float
+    """Amount of randomness injected into the response."""
+
+    top_p: float
+    """Nucleus sampling parameter."""
+
+    timeout: float
+    """Request timeout in seconds (float only, not httpx.Timeout)."""
+
+    parallel_tool_calls: bool
+    """Whether to allow parallel tool calls."""
+
+    seed: int
+    """Random seed for deterministic results."""
+
+    presence_penalty: float
+    """Penalize tokens based on whether they have appeared in the text so far."""
+
+    frequency_penalty: float
+    """Penalize tokens based on their existing frequency in the text so far."""
+
+    logit_bias: dict[str, int]
+    """Modify the likelihood of specified tokens appearing in the completion."""
+
+    stop_sequences: list[str]
+    """Sequences that will cause the model to stop generating."""
+
+    extra_headers: dict[str, str]
+    """Extra headers to send to the model."""
+
+    extra_body: dict[str, Any]
+    """Extra body to send to the model."""
+
+
+class CompactBaseModel(BaseModel):
+    """Base model that excludes unset and None values during serialization."""
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        kwargs.setdefault("exclude_unset", True)
+        kwargs.setdefault("exclude_none", True)
+        return super().model_dump(**kwargs)
+
+    def model_dump_json(self, **kwargs: Any) -> str:
+        kwargs.setdefault("exclude_unset", True)
+        kwargs.setdefault("exclude_none", True)
+        return super().model_dump_json(**kwargs)

calfkit/nodes/__init__.py

@@ -0,0 +1,16 @@
+from calfkit.nodes.agent_router_node import AgentRouterNode
+from calfkit.nodes.base_node import BaseNode, publish_to, subscribe_to
+from calfkit.nodes.base_tool_node import BaseToolNode, agent_tool
+from calfkit.nodes.chat_node import ChatNode
+from calfkit.nodes.registrator import Registrator
+
+__all__ = [
+    "AgentRouterNode",
+    "BaseNode",
+    "BaseToolNode",
+    "ChatNode",
+    "Registrator",
+    "agent_tool",
+    "publish_to",
+    "subscribe_to",
+]

calfkit/nodes/agent_router_node.py

@@ -0,0 +1,194 @@
+from typing import Annotated, Any
+
+import uuid_utils
+from faststream import Context
+from faststream.kafka.annotations import (
+    KafkaBroker as BrokerAnnotation,
+)
+from pydantic_ai import ModelRequest, ModelResponse, SystemPromptPart
+from pydantic_ai.models import ModelRequestParameters
+
+from calfkit.broker.broker import Broker
+from calfkit.messages import patch_system_prompts, validate_tool_call_pairs
+from calfkit.models.event_envelope import EventEnvelope
+from calfkit.models.types import ToolCallRequest
+from calfkit.nodes.base_node import BaseNode, publish_to, subscribe_to
+from calfkit.nodes.base_tool_node import BaseToolNode
+from calfkit.stores.base import MessageHistoryStore
+
+
+class AgentRouterNode(BaseNode):
+    """Logic for the internal routing to operate agents"""
+
+    _router_sub_topic_name = "agent_router.input"
+    _router_pub_topic_name = "agent_router.output"
+
+    def __init__(
+        self,
+        chat_node: BaseNode,
+        tool_nodes: list[BaseToolNode],
+        system_prompt: str | None = None,
+        handoff_nodes: list[type[BaseNode]] = [],
+        message_history_store: MessageHistoryStore | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        self.chat = chat_node
+        self.tools = tool_nodes
+        self.handoffs = handoff_nodes
+        self.system_prompt = system_prompt
+        self.system_message = (
+            ModelRequest(parts=[SystemPromptPart(self.system_prompt)])
+            if self.system_prompt
+            else None
+        )
+        self.message_history_store = message_history_store
+
+        self.tools_topic_registry: dict[str, str] = {
+            tool.tool_schema().name: tool.subscribed_topic
+            for tool in tool_nodes
+            if tool.subscribed_topic is not None
+        }
+
+        self.tool_response_topics = [tool.publish_to_topic for tool in self.tools]
+
+        super().__init__(*args, **kwargs)
+
+    @subscribe_to(_router_sub_topic_name)
+    @publish_to(_router_pub_topic_name)
+    async def _router(
+        self,
+        ctx: EventEnvelope,
+        correlation_id: Annotated[str, Context()],
+        broker: BrokerAnnotation,
+    ) -> EventEnvelope:
+        if not ctx.incoming_node_messages:
+            raise RuntimeError("There is no response message to process")
+
+        # One central place where message history is updated
+        if self.message_history_store is not None and ctx.thread_id is not None:
+            await self.message_history_store.append_many(
+                thread_id=ctx.thread_id, messages=ctx.incoming_node_messages
+            )
+            ctx.message_history = await self.message_history_store.get(thread_id=ctx.thread_id)
+        else:
+            ctx.message_history.extend(ctx.incoming_node_messages)
+
+        # Apply system prompts with priority: incoming > self.system_message > existing history
+        # First, apply self.system_message as fallback (replaces existing history)
+        if ctx.system_message is not None:
+            ctx.message_history = patch_system_prompts(ctx.message_history, [ctx.system_message])
+        elif self.system_message is not None:
+            ctx.message_history = patch_system_prompts(
+                ctx.message_history,
+                [self.system_message],
+            )
+
+        if isinstance(ctx.latest_message_in_history, ModelResponse):
+            if (
+                ctx.latest_message_in_history.finish_reason == "tool_call"
+                or ctx.latest_message_in_history.tool_calls
+            ):
+                for tool_call in ctx.latest_message_in_history.tool_calls:
+                    await self._route_tool(ctx, tool_call, correlation_id, broker)
+            else:
+                # reply to sender here
+                await self._reply_to_sender(ctx, correlation_id, broker)
+        elif validate_tool_call_pairs(ctx.message_history):
+            await self._call_model(ctx, correlation_id, broker)
+        return ctx
+
+    async def _route_tool(
+        self,
+        event_envelope: EventEnvelope,
+        generated_tool_call: ToolCallRequest,
+        correlation_id: str,
+        broker: Any,
+    ) -> None:
+        tool_topic = self.tools_topic_registry.get(generated_tool_call.tool_name)
+        if tool_topic is None:
+            # TODO: implement a short circuit to respond with an
+            # error message for when provided tool does not exist.
+            return
+        event_envelope = event_envelope.model_copy(
+            update={"kind": "tool_call_request", "tool_call_request": generated_tool_call}
+        )
+        await broker.publish(
+            event_envelope,
+            topic=tool_topic,
+            correlation_id=correlation_id,
+            reply_to=self.subscribed_topic,
+        )
+
+    async def _reply_to_sender(
+        self, event_envelope: EventEnvelope, correlation_id: str, broker: Any
+    ) -> None:
+        event_envelope = event_envelope.model_copy(
+            update={"kind": "ai_response", "final_response": True}
+        )
+        await broker.publish(
+            event_envelope,
+            topic=event_envelope.final_response_topic,
+            correlation_id=correlation_id,
+        )
+
+    async def _call_model(
+        self,
+        event_envelope: EventEnvelope,
+        correlation_id: str,
+        broker: Any,
+    ) -> None:
+        patch_model_request_params = event_envelope.patch_model_request_params
+        if patch_model_request_params is None:
+            patch_model_request_params = ModelRequestParameters(
+                function_tools=[tool.tool_schema() for tool in self.tools]
+            )
+        event_envelope = event_envelope.model_copy(
+            update={"patch_model_request_params": patch_model_request_params}
+        )
+        await broker.publish(
+            event_envelope,
+            topic=self.chat.subscribed_topic,
+            correlation_id=correlation_id,
+            reply_to=self.subscribed_topic,
+        )
+
+    async def invoke(
+        self,
+        user_prompt: str,
+        broker: Broker,
+        final_response_topic: str,
+        thread_id: str | None = None,
+        correlation_id: str | None = None,
+    ) -> str:
+        """Invoke the agent
+
+        Args:
+            user_prompt (str): User prompt to request the model
+            broker (Broker): The broker to connect to
+            correlation_id (str | None, optional): Optionally provide a correlation ID
+            for this request. Defaults to None.
+
+        Returns:
+            str: The correlation ID for this request
+        """
+        patch_model_request_params = ModelRequestParameters(
+            function_tools=[tool.tool_schema() for tool in self.tools]
+        )
+        if correlation_id is None:
+            correlation_id = uuid_utils.uuid7().hex
+        new_node_messages = [ModelRequest.user_text_prompt(user_prompt)]
+        await broker.publish(
+            EventEnvelope(
+                kind="user_prompt",
+                trace_id=correlation_id,
+                patch_model_request_params=patch_model_request_params,
+                thread_id=thread_id,
+                incoming_node_messages=new_node_messages,
+                system_message=self.system_message,
+                final_response_topic=final_response_topic,
+            ),
+            topic=self.subscribed_topic or "",
+            correlation_id=correlation_id,
+        )
+        return correlation_id

calfkit/nodes/base_node.py

@@ -0,0 +1,62 @@
+from abc import ABC
+from collections.abc import Callable
+from functools import cached_property
+from typing import Any
+
+
+def subscribe_to(topic_name: str) -> Callable[[Any], Any]:
+    def decorator(fn: Any) -> Any:
+        fn._subscribe_to_topic_name = topic_name
+        return fn
+
+    return decorator
+
+
+def publish_to(topic_name: str) -> Callable[[Any], Any]:
+    def decorator(fn: Any) -> Any:
+        fn._publish_to_topic_name = topic_name
+        return fn
+
+    return decorator
+
+
+class BaseNode(ABC):
+    """Effectively a node is the data plane, defining the internal wiring and logic.
+    When provided to a NodeRunner, node logic can be deployed."""
+
+    _handler_registry: dict[Callable[..., Any], dict[str, str]] = {}
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self.bound_registry: dict[Callable[..., Any], dict[str, str]] = {
+            fn.__get__(self, type(self)): topics_dict
+            for fn, topics_dict in self._handler_registry.items()
+        }
+
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+
+        cls._handler_registry = {}
+
+        for attr in cls.__dict__.values():
+            publish_to_topic_name = getattr(attr, "_publish_to_topic_name", None)
+            subscribe_to_topic_name = getattr(attr, "_subscribe_to_topic_name", None)
+            if publish_to_topic_name:
+                cls._handler_registry[attr] = {"publish_topic": publish_to_topic_name}
+            if subscribe_to_topic_name:
+                cls._handler_registry[attr] = cls._handler_registry.get(attr, {}) | {
+                    "subscribe_topic": subscribe_to_topic_name
+                }
+
+    @cached_property
+    def subscribed_topic(self) -> str | None:
+        for topics_dict in self._handler_registry.values():
+            if "subscribe_topic" in topics_dict:
+                return topics_dict["subscribe_topic"]
+        return None
+
+    @cached_property
+    def publish_to_topic(self) -> str | None:
+        for topics_dict in self._handler_registry.values():
+            if "publish_topic" in topics_dict:
+                return topics_dict["publish_topic"]
+        return None

calfkit/nodes/base_tool_node.py

@@ -0,0 +1,55 @@
+import inspect
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable, Callable
+from typing import Any, cast
+
+from pydantic_ai import ModelRequest, Tool, ToolDefinition, ToolReturnPart
+
+from calfkit.models.event_envelope import EventEnvelope
+from calfkit.nodes.base_node import BaseNode, publish_to, subscribe_to
+
+
+class BaseToolNode(BaseNode, ABC):
+    @classmethod
+    @abstractmethod
+    def tool_schema(cls) -> ToolDefinition: ...
+
+
+def agent_tool(func: Callable[..., Any] | Callable[..., Awaitable[Any]]) -> BaseToolNode:
+    """Agent tool decorator to turn a function into a deployable node"""
+    tool = Tool(func)
+
+    class ToolNode(BaseToolNode):
+        @subscribe_to(f"tool_node.{func.__name__}.request")
+        @publish_to(f"tool_node.{func.__name__}.result")
+        async def on_enter(self, event_envelope: EventEnvelope) -> EventEnvelope:
+            if not event_envelope.tool_call_request:
+                raise RuntimeError("No tool call request found")
+            tool_cal_req = event_envelope.tool_call_request
+            kw_args = tool_cal_req.args_as_dict()
+            result = func(**kw_args)
+            if inspect.isawaitable(result):
+                result = await result
+            tool_result = ToolReturnPart(
+                tool_name=tool_cal_req.tool_name,
+                content=result,
+                tool_call_id=tool_cal_req.tool_call_id,
+            )
+            event_envelope = event_envelope.model_copy(
+                update={
+                    "kind": "tool_result",
+                    "incoming_node_messages": [ModelRequest(parts=[tool_result])],
+                }
+            )
+            return event_envelope
+
+        @classmethod
+        def tool_schema(cls) -> ToolDefinition:
+            return cast(ToolDefinition, tool.tool_def)
+
+    ToolNode.__name__ = func.__name__
+    ToolNode.__qualname__ = func.__qualname__
+    ToolNode.__doc__ = func.__doc__
+    ToolNode.__module__ = func.__module__
+
+    return ToolNode()