diff --git a/docs/docs/getting-started/features.md b/docs/docs/getting-started/features.md index 5b465699..2cb6bb49 100644 --- a/docs/docs/getting-started/features.md +++ b/docs/docs/getting-started/features.md @@ -39,7 +39,9 @@ Graphite adopts an event-driven architecture where topics function as logical me Consumption events are only recorded once the entire node processing completes successfully. Until that point, the system treats partial or failed node invokes as if they never happened, preventing duplicated outputs or broken states. Should a node encounter an error (e.g., an LLM connection failure, external API issue, or function exception), Graphite detects the unconsumed events upon restoration and places the associated node(s) back into the invoke queue. This design ensures the node can safely retry from the same input without creating conflicting or duplicated consumption records. -By storing each event exactly once and withholding consumption records until success, Graphite guarantees idempotent behavior. Even if a node issues multiple invocations due to an error, the event logs and consumption rules still reconstruct a single, consistent path from invocation to response. This approach produces correct outcomes on retries while maintaining a complete, conflict-free audit trail. +By recording each event idempotently and withholding consumption records until success, Graphite makes its **event persistence and replay** idempotent: the event logs and consumption rules reconstruct a single, consistent path from invocation to response, maintaining a complete, conflict-free audit trail. + +This is not the same as exactly-once execution of external side effects. Node and tool invocation is **at-least-once**: on retry, a node runs again from the same input, so any external action it performs (an LLM call, an outbound API request, a write) can happen more than once unless the tool itself is idempotent. Make tools with external side effects idempotent (e.g. via an idempotency key) if you need exactly-once *effects*. ## Auditability diff --git a/docs/docs/user-guide/event-driven-workflow.md b/docs/docs/user-guide/event-driven-workflow.md index 305c7719..2c15ee78 100644 --- a/docs/docs/user-guide/event-driven-workflow.md +++ b/docs/docs/user-guide/event-driven-workflow.md @@ -100,14 +100,14 @@ The asynchronous execution model provides sophisticated event-driven processing 1. **Workflow Initialization**: Sets up initial state with `init_workflow` 2. **Concurrent Node Processing**: Spawns individual tasks for each node using `_invoke_node` 3. **Output Listening**: Creates listeners for each output topic to capture results -4. **Event Streaming**: Uses `MergeIdleQueue` to stream events as they become available +4. **Event Streaming**: Uses `AsyncOutputQueue` to stream output events as they become available 5. **Proper Termination**: Coordinates workflow completion using `AsyncNodeTracker` #### Key Components -- **AsyncNodeTracker**: Manages active node state and idle detection +- **AsyncNodeTracker**: Manages active node state and quiescence detection - **Output Listeners**: Monitor output topics for new events -- **MergeIdleQueue**: Coordinates between event availability and workflow idle state +- **AsyncOutputQueue**: Streams output-topic events and coordinates with workflow quiescence - **Offset Management**: Commits events immediately to prevent duplicates ### Event-Driven Node Triggering @@ -185,7 +185,9 @@ async for event in output_queue: The async workflow implements proper offset management to prevent duplicate data: ```python -async for event in MergeIdleQueue(queue, self._tracker): +output_queue = AsyncOutputQueue(output_topics, self.name, self._tracker) +await output_queue.start_listeners() +async for event in output_queue: consumed_output_event = ConsumeFromTopicEvent(...) # Commit BEFORE yielding to prevent duplicate data diff --git a/docs/docs/user-guide/tools/function.md b/docs/docs/user-guide/tools/function.md index 9e3fef80..7c044595 100644 --- a/docs/docs/user-guide/tools/function.md +++ b/docs/docs/user-guide/tools/function.md @@ -39,7 +39,7 @@ The tool automatically handles different response types in its `to_messages` met - **`BaseModel` instances**: Serialized to JSON using `model_dump_json()` - **Lists of `BaseModel` objects**: Converted to JSON arrays using `model_dump()` for each item - **String responses**: Used directly as message content -- **Other types**: Encoded using `jsonpickle` for complex object serialization +- **Other types**: Encoded with `json.dumps(..., default=str)` (no pickle) ### Async Support diff --git a/grafi/assistants/assistant.py b/grafi/assistants/assistant.py index ae55ccea..5d1884e9 100644 --- a/grafi/assistants/assistant.py +++ b/grafi/assistants/assistant.py @@ -69,6 +69,8 @@ def generate_manifest(self, output_dir: str = ".") -> str: with open(output_path, "w") as f: f.write(json.dumps(manifest_dict, indent=4)) + return output_path + @classmethod async def from_dict(cls, data: dict[str, Any]) -> "Assistant": """ diff --git a/grafi/assistants/assistant_base.py b/grafi/assistants/assistant_base.py index c58bf2ba..1596d1d2 100644 --- a/grafi/assistants/assistant_base.py +++ b/grafi/assistants/assistant_base.py @@ -44,8 +44,12 @@ def model_post_init(self, _context: Any) -> None: self._construct_workflow() def _construct_workflow(self) -> "AssistantBase": - """Construct the workflow for the assistant.""" - pass + """Construct the workflow for the assistant. + + Subclasses override this to build ``self.workflow`` and return ``self``. + The base is a no-op for a bare assistant with a default workflow. + """ + return self async def invoke( self, input_data: PublishToTopicEvent, is_sequential: bool = False diff --git a/grafi/common/decorators/llm_function.py b/grafi/common/decorators/llm_function.py index 7302f03d..ca6c6af5 100644 --- a/grafi/common/decorators/llm_function.py +++ b/grafi/common/decorators/llm_function.py @@ -149,7 +149,7 @@ def _type_to_schema(tp: Any) -> JsonSchema: f_schema = _type_to_schema(f_type) if f.default is not dataclasses.MISSING: f_schema.setdefault("default", f.default) - elif f.default_factory is not dataclasses.MISSING: # type: ignore[attr-defined] + elif f.default_factory is not dataclasses.MISSING: # we can't serialize the factory, just mark as optional pass else: @@ -164,7 +164,7 @@ def _type_to_schema(tp: Any) -> JsonSchema: import enum if isinstance(tp, type) and issubclass(tp, enum.Enum): - values = [m.value for m in tp] # type: ignore[arg-type] + values = [m.value for m in tp] # derive base type from first value base = _type_to_schema(type(values[0])) if values else {} base["enum"] = values @@ -242,11 +242,11 @@ def from_function( # Unwrap callable classes: use __call__ if not inspect.isroutine(fn) and hasattr(fn, "__call__"): - fn = fn.__call__ # type: ignore[assignment] + fn = fn.__call__ # Unwrap staticmethod if isinstance(fn, staticmethod): - fn = fn.__func__ # type: ignore[assignment] + fn = fn.__func__ sig = inspect.signature(fn) type_hints = get_type_hints(fn, include_extras=True) @@ -301,7 +301,7 @@ def from_function( return_schema: Optional[JsonSchema] = None return_ann = type_hints.get("return", sig.return_annotation) - if return_ann not in (inspect._empty, None, Any, ...): # type: ignore[attr-defined] + if return_ann not in (inspect._empty, None, Any, ...): # Convert return type to schema return_schema = _type_to_schema(return_ann) diff --git a/grafi/common/decorators/record_base.py b/grafi/common/decorators/record_base.py index c3adaf93..cb6cb32d 100644 --- a/grafi/common/decorators/record_base.py +++ b/grafi/common/decorators/record_base.py @@ -9,6 +9,7 @@ from typing import Callable from typing import Dict from typing import List +from typing import Optional from typing import Type from typing import TypeVar from typing import Union @@ -20,7 +21,6 @@ from pydantic import ConfigDict from pydantic_core import to_jsonable_python -from grafi.assistants.assistant_base import AssistantBase from grafi.common.containers.container import container from grafi.common.env import env_bool from grafi.common.events.component_base import ComponentEvent @@ -34,9 +34,6 @@ from grafi.common.models.default_id import default_id from grafi.common.models.invoke_context import InvokeContext from grafi.common.models.message import Message -from grafi.nodes.node_base import NodeBase -from grafi.tools.tool import Tool -from grafi.workflows.workflow import Workflow T = TypeVar("T") @@ -63,9 +60,10 @@ class ComponentConfig: event_types: Dict[ str, Type[ComponentEvent] ] # Maps 'invoke', 'respond', 'failed' to event classes - extract_metadata: Callable[ - [Union[AssistantBase, Workflow, NodeBase, Tool]], EventContext - ] # Extracts component-specific metadata + # Extracts component-specific metadata. Typed as ``Any`` so this common-layer + # module does not import the higher-layer component classes (Assistant / + # Workflow / Node / Tool) just for an annotation. + extract_metadata: Callable[[Any], EventContext] process_async_result: Callable[[List], Any] span_name_suffix: str = "invoke" # Suffix for span name @@ -237,27 +235,33 @@ def create_async_decorator(config: ComponentConfig) -> Callable: def decorator(func: Callable) -> Callable: @functools.wraps(func) async def wrapper( - self: Union[AssistantBase, Workflow, NodeBase, Tool], - *args, - **kwargs, + self: Any, + *args: Any, + **kwargs: Any, ) -> AsyncGenerator[Union[PublishToTopicEvent, List[Message]], None]: # Extract metadata using component-specific logic metadata = config.extract_metadata(self) - input_data: Union[ - List[ConsumeFromTopicEvent], List[Message], PublishToTopicEvent + input_data: Optional[ + Union[List[ConsumeFromTopicEvent], List[Message], PublishToTopicEvent] ] = None if isinstance(args[0], InvokeContext): invoke_context: InvokeContext = args[0] input_data = args[1] else: - # Assistant and workflow + # Assistant and workflow: args[0] is the input event (Any), which + # carries invoke_context. Read it off args[0] directly so mypy + # does not widen to the input_data union (which has no such attr). input_data = args[0] - invoke_context = input_data.invoke_context - - # Create invoke event - invoke_event = config.event_types["invoke"]( + invoke_context = args[0].invoke_context + + # Create invoke event. The factory maps each string key to the + # matching InvokeEvent/RespondEvent/FailedEvent subclass, so these + # per-event kwargs (input_data/output_data/error/...) are correct at + # runtime but unprovable to mypy through the Dict value's + # ComponentEvent base type -- hence the targeted call-arg ignores. + invoke_event = config.event_types["invoke"]( # type: ignore[call-arg] id=metadata.id, name=metadata.name, type=metadata.type, @@ -268,44 +272,52 @@ async def wrapper( # Execute with tracing output_data = None + error_details: Optional[Dict[str, Any]] = None try: with container.tracer.start_as_current_span( f"{metadata.name}.{config.span_name_suffix}" ) as span: - # Set span attributes - for key, value in metadata.model_dump().items(): - if value is not None: - span.set_attribute(key, value) - - span.set_attributes(invoke_context.model_dump()) - - # Set input (size-bounded; omitted if payloads are disabled) - input_payload = _span_payload(input_data) - if input_payload is not None: - span.set_attribute("input", input_payload) - - # Handle streaming - result_list: List = [] - - async for result in func(self, *args, **kwargs): - yield result - result_list.append(result) - - output_data = config.process_async_result(result_list) - - output_payload = _span_payload(output_data) - if output_payload is not None: - span.set_attribute("output", output_payload) + try: + # Set span attributes + for key, value in metadata.model_dump().items(): + if value is not None: + span.set_attribute(key, value) + + span.set_attributes(invoke_context.model_dump()) + + # Set input (size-bounded; omitted if payloads are disabled) + input_payload = _span_payload(input_data) + if input_payload is not None: + span.set_attribute("input", input_payload) + + # Handle streaming + result_list: List = [] + + async for result in func(self, *args, **kwargs): + yield result + result_list.append(result) + + output_data = config.process_async_result(result_list) + + output_payload = _span_payload(output_data) + if output_payload is not None: + span.set_attribute("output", output_payload) + except Exception as e: + # Build structured details once, while the traceback is + # still attached to the exception, and enrich the span + # WHILE it is still recording. Attributes set after the + # span context manager exits are dropped by the backend. + error_details = _build_error_details(e, metadata) + _record_span_error(span, e, error_details) + raise except Exception as e: - # Build structured details once, while the traceback is still - # attached to the exception. - error_details = _build_error_details(e, metadata) - - # Enrich the active span with structured error information. - if "span" in locals(): - _record_span_error(span, e, error_details) + # error_details is normally built above, inside the live span. + # Rebuild defensively only if the failure happened before the + # inner try (e.g. span creation itself). + if error_details is None: + error_details = _build_error_details(e, metadata) # Log a full traceback once at the layer closest to the failure; # outer layers log a concise summary (see _log_component_exception). @@ -313,7 +325,7 @@ async def wrapper( # Record failed event with both the human-readable string (kept # for backward compatibility) and the structured details. - failed_event = config.event_types["failed"]( + failed_event = config.event_types["failed"]( # type: ignore[call-arg] id=metadata.id, name=metadata.name, type=metadata.type, @@ -326,7 +338,7 @@ async def wrapper( raise else: # Record respond event - respond_event = config.event_types["respond"]( + respond_event = config.event_types["respond"]( # type: ignore[call-arg] id=metadata.id, name=metadata.name, type=metadata.type, diff --git a/grafi/common/event_stores/event_store.py b/grafi/common/event_stores/event_store.py index 24487be9..6b1701ea 100644 --- a/grafi/common/event_stores/event_store.py +++ b/grafi/common/event_stores/event_store.py @@ -7,32 +7,22 @@ from typing import List from typing import Optional from typing import Sequence -from typing import Type -from loguru import logger - -from grafi.common.events.component_events import AssistantFailedEvent -from grafi.common.events.component_events import AssistantInvokeEvent -from grafi.common.events.component_events import AssistantRespondEvent -from grafi.common.events.component_events import ConsumeFromTopicEvent -from grafi.common.events.component_events import NodeFailedEvent -from grafi.common.events.component_events import NodeInvokeEvent -from grafi.common.events.component_events import NodeRespondEvent -from grafi.common.events.component_events import ToolFailedEvent -from grafi.common.events.component_events import ToolInvokeEvent -from grafi.common.events.component_events import ToolRespondEvent -from grafi.common.events.component_events import WorkflowFailedEvent -from grafi.common.events.component_events import WorkflowInvokeEvent -from grafi.common.events.component_events import WorkflowRespondEvent from grafi.common.events.event import Event -from grafi.common.events.event import EventType -from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent -from grafi.common.events.topic_events.topic_event import TopicEvent +from grafi.common.events.event_codec import EventCodec +from grafi.common.events.event_codec import default_event_codec class EventStore(ABC): """Stores and manages events.""" + # Decoder for stored event dicts. Persistence backends delegate decoding to + # the codec instead of knowing concrete event classes. This is a shared, + # process-wide default registry: calling ``.register(...)`` on it affects all + # stores. For an isolated registry, assign a fresh ``EventCodec`` to the + # instance (``self._codec = EventCodec({...})``) in a subclass constructor. + _codec: EventCodec = default_event_codec + @abstractmethod async def record_event(self, event: Event) -> None: ... @@ -58,54 +48,9 @@ async def get_conversation_events(self, conversation_id: str) -> List[Event]: .. async def get_topic_events(self, name: str, offsets: List[int]) -> List[Event]: ... def _create_event_from_dict(self, event_dict: Dict[str, Any]) -> Optional[Event]: - """Create an event object from a dictionary. + """Decode a stored event dict into an :class:`Event` via the codec. - Returns ``None`` (with a logged warning) instead of raising when a single - stored event is malformed or of an unknown type. Retrieval methods skip - ``None`` results, so one corrupt row cannot abort an entire - ``get_agent_events`` / ``get_events`` call — which would otherwise make a - whole conversation (and any recovery that depends on it) unreadable. + Kept as a thin delegate for backends (and tests) that decode persisted + rows; the registry and error policy live in :class:`EventCodec`. """ - event_id = event_dict.get("event_id", "") - event_type: Any = event_dict.get("event_type") - if not isinstance(event_type, str): - logger.warning( - "Skipping event {} with missing/invalid event_type", event_id - ) - return None - - event_class = self._get_event_class(event_type) - if event_class is None: - logger.warning( - "Skipping event {} with unknown event type: {}", event_id, event_type - ) - return None - - try: - return event_class.from_dict(data=event_dict) - except Exception as e: - logger.error( - "Skipping event {} that failed to deserialize: {}", event_id, e - ) - return None - - def _get_event_class(self, event_type: str) -> Optional[Type[Event]]: - """Get the event class based on the event type string.""" - event_classes = { - EventType.NODE_FAILED.value: NodeFailedEvent, - EventType.NODE_INVOKE.value: NodeInvokeEvent, - EventType.NODE_RESPOND.value: NodeRespondEvent, - EventType.TOOL_FAILED.value: ToolFailedEvent, - EventType.TOOL_INVOKE.value: ToolInvokeEvent, - EventType.TOOL_RESPOND.value: ToolRespondEvent, - EventType.WORKFLOW_FAILED.value: WorkflowFailedEvent, - EventType.WORKFLOW_INVOKE.value: WorkflowInvokeEvent, - EventType.WORKFLOW_RESPOND.value: WorkflowRespondEvent, - EventType.ASSISTANT_FAILED.value: AssistantFailedEvent, - EventType.ASSISTANT_INVOKE.value: AssistantInvokeEvent, - EventType.ASSISTANT_RESPOND.value: AssistantRespondEvent, - EventType.TOPIC_EVENT.value: TopicEvent, - EventType.CONSUME_FROM_TOPIC.value: ConsumeFromTopicEvent, - EventType.PUBLISH_TO_TOPIC.value: PublishToTopicEvent, - } - return event_classes.get(event_type) + return self._codec.decode(event_dict) diff --git a/grafi/common/event_stores/event_store_postgres.py b/grafi/common/event_stores/event_store_postgres.py index 20503271..14ecfced 100644 --- a/grafi/common/event_stores/event_store_postgres.py +++ b/grafi/common/event_stores/event_store_postgres.py @@ -1,7 +1,10 @@ from datetime import datetime +from typing import Any +from typing import Dict from typing import List from typing import Optional from typing import Sequence +from typing import cast from loguru import logger @@ -110,6 +113,15 @@ async def clear_events(self) -> None: logger.error(f"Failed to clear events: {e}") raise + def _row_to_event(self, row: "EventModel") -> Optional[Event]: + """Decode a single ``EventModel`` row into an :class:`Event`. + + Centralizes the row-to-event conversion shared by every query method. + ``row.event_data`` is a JSONB column typed as ``Column[Any]`` at the + class level; at the instance level it holds the stored dict. + """ + return self._create_event_from_dict(cast(Dict[str, Any], row.event_data)) + async def get_events(self) -> List[Event]: """Get all events from the database.""" async with self.AsyncSession() as session: @@ -121,7 +133,7 @@ async def get_events(self) -> List[Event]: events: List[Event] = [] for r in rows: - event = self._create_event_from_dict(r.event_data) + event = self._row_to_event(r) if event: events.append(event) @@ -207,7 +219,7 @@ async def get_event(self, event_id: str) -> Optional[Event]: if not row: return None - return self._create_event_from_dict(row.event_data) + return self._row_to_event(row) except Exception as e: logger.error(f"Failed to get event {event_id}: {e}") raise @@ -228,7 +240,7 @@ async def get_agent_events(self, assistant_request_id: str) -> List[Event]: events: List[Event] = [] for r in rows: - event = self._create_event_from_dict(r.event_data) + event = self._row_to_event(r) if event: events.append(event) @@ -253,7 +265,7 @@ async def get_conversation_events(self, conversation_id: str) -> List[Event]: events: List[Event] = [] for r in rows: - event = self._create_event_from_dict(r.event_data) + event = self._row_to_event(r) if event: events.append(event) @@ -304,7 +316,7 @@ async def get_topic_events(self, name: str, offsets: List[int]) -> List[Event]: events: List[Event] = [] for r in rows: - event = self._create_event_from_dict(r.event_data) + event = self._row_to_event(r) if event: events.append(event) diff --git a/grafi/common/events/event_codec.py b/grafi/common/events/event_codec.py new file mode 100644 index 00000000..24e6b0ba --- /dev/null +++ b/grafi/common/events/event_codec.py @@ -0,0 +1,103 @@ +"""Registry-based decoding of stored event dicts into :class:`Event` objects. + +The codec owns event-type registration and decoding so that persistence +backends (``EventStore`` implementations) do not need to know concrete event +classes, and new event types can be registered without editing the store +(Open/Closed). Built-ins are registered here, at the composition boundary, +rather than imported into the persistence layer. +""" + +from typing import Any +from typing import Dict +from typing import Optional +from typing import Type + +from loguru import logger + +from grafi.common.events.component_events import AssistantFailedEvent +from grafi.common.events.component_events import AssistantInvokeEvent +from grafi.common.events.component_events import AssistantRespondEvent +from grafi.common.events.component_events import ConsumeFromTopicEvent +from grafi.common.events.component_events import NodeFailedEvent +from grafi.common.events.component_events import NodeInvokeEvent +from grafi.common.events.component_events import NodeRespondEvent +from grafi.common.events.component_events import ToolFailedEvent +from grafi.common.events.component_events import ToolInvokeEvent +from grafi.common.events.component_events import ToolRespondEvent +from grafi.common.events.component_events import WorkflowFailedEvent +from grafi.common.events.component_events import WorkflowInvokeEvent +from grafi.common.events.component_events import WorkflowRespondEvent +from grafi.common.events.event import Event +from grafi.common.events.event import EventType +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.events.topic_events.topic_event import TopicEvent + + +class EventCodec: + """Maps event-type strings to event classes and decodes stored dicts.""" + + def __init__(self, registry: Optional[Dict[str, Type[Event]]] = None) -> None: + self._registry: Dict[str, Type[Event]] = dict(registry) if registry else {} + + def register(self, event_type: str, event_class: Type[Event]) -> None: + """Register (or override) the class used to decode ``event_type``.""" + self._registry[event_type] = event_class + + def event_class(self, event_type: str) -> Optional[Type[Event]]: + """Return the registered class for ``event_type``, or ``None``.""" + return self._registry.get(event_type) + + def decode(self, event_dict: Dict[str, Any]) -> Optional[Event]: + """Decode one stored event dict into an :class:`Event`. + + Returns ``None`` (with a logged warning) instead of raising when an event + is malformed or of an unknown type. Retrieval methods skip ``None`` + results, so one corrupt row cannot abort an entire ``get_agent_events`` / + ``get_events`` call -- which would otherwise make a whole conversation + (and any recovery that depends on it) unreadable. + """ + event_id = event_dict.get("event_id", "") + event_type: Any = event_dict.get("event_type") + if not isinstance(event_type, str): + logger.warning( + "Skipping event {} with missing/invalid event_type", event_id + ) + return None + + event_class = self.event_class(event_type) + if event_class is None: + logger.warning( + "Skipping event {} with unknown event type: {}", event_id, event_type + ) + return None + + try: + return event_class.from_dict(data=event_dict) + except Exception as e: + logger.error( + "Skipping event {} that failed to deserialize: {}", event_id, e + ) + return None + + +# The default codec, pre-registered with the framework's built-in event types. +# Downstream code can register additional types without editing any EventStore. +default_event_codec = EventCodec( + { + EventType.NODE_FAILED.value: NodeFailedEvent, + EventType.NODE_INVOKE.value: NodeInvokeEvent, + EventType.NODE_RESPOND.value: NodeRespondEvent, + EventType.TOOL_FAILED.value: ToolFailedEvent, + EventType.TOOL_INVOKE.value: ToolInvokeEvent, + EventType.TOOL_RESPOND.value: ToolRespondEvent, + EventType.WORKFLOW_FAILED.value: WorkflowFailedEvent, + EventType.WORKFLOW_INVOKE.value: WorkflowInvokeEvent, + EventType.WORKFLOW_RESPOND.value: WorkflowRespondEvent, + EventType.ASSISTANT_FAILED.value: AssistantFailedEvent, + EventType.ASSISTANT_INVOKE.value: AssistantInvokeEvent, + EventType.ASSISTANT_RESPOND.value: AssistantRespondEvent, + EventType.TOPIC_EVENT.value: TopicEvent, + EventType.CONSUME_FROM_TOPIC.value: ConsumeFromTopicEvent, + EventType.PUBLISH_TO_TOPIC.value: PublishToTopicEvent, + } +) diff --git a/grafi/common/events/event_graph.py b/grafi/common/events/event_graph.py index c163da11..cb979607 100644 --- a/grafi/common/events/event_graph.py +++ b/grafi/common/events/event_graph.py @@ -19,8 +19,8 @@ class EventGraphNode(BaseModel): event_id: EventId event: TopicEvent - upstream_events: List[EventId] = Field(default=[]) - downstream_events: List[EventId] = Field(default=[]) + upstream_events: List[EventId] = Field(default_factory=list) + downstream_events: List[EventId] = Field(default_factory=list) def to_dict(self) -> dict[str, Any]: return { @@ -41,8 +41,8 @@ def from_dict(cls, data: dict) -> "EventGraphNode": class EventGraph(BaseModel): - nodes: Dict[EventId, EventGraphNode] = Field(default={}) - root_nodes: List[EventGraphNode] = Field(default=[]) + nodes: Dict[EventId, EventGraphNode] = Field(default_factory=dict) + root_nodes: List[EventGraphNode] = Field(default_factory=list) def _add_event(self, event: TopicEvent) -> EventGraphNode: """Add a new node to the graph if it doesn't exist""" diff --git a/grafi/common/models/async_result.py b/grafi/common/models/async_result.py index c51c43a2..34c6b6bf 100644 --- a/grafi/common/models/async_result.py +++ b/grafi/common/models/async_result.py @@ -1,7 +1,9 @@ import asyncio +from typing import Any from typing import AsyncGenerator +from typing import Generator from typing import Optional -from typing import TypeVar +from typing import cast from grafi.common.events.topic_events.consume_from_topic_event import ( ConsumeFromTopicEvent, @@ -10,9 +12,6 @@ _SENTINEL = object() -T = TypeVar("T") - - class AsyncResult: """ Wraps one of: @@ -37,7 +36,8 @@ class AsyncResult: def __init__(self, source: AsyncGenerator[ConsumeFromTopicEvent, None]): self._source = source - self._queue: asyncio.Queue[ConsumeFromTopicEvent] = asyncio.Queue() + # Holds stream items plus the end-of-stream sentinel, hence ``object``. + self._queue: asyncio.Queue[object] = asyncio.Queue() self._items: list[ConsumeFromTopicEvent] = [] self._done = asyncio.Event() self._started = False @@ -63,22 +63,22 @@ async def _producer(self) -> None: await self._queue.put(_SENTINEL) # --- async iteration support --- - def __aiter__(self): + def __aiter__(self) -> "AsyncResult": self._ensure_started() return self - async def __anext__(self) -> T: + async def __anext__(self) -> ConsumeFromTopicEvent: self._ensure_started() item = await self._queue.get() if item is _SENTINEL: if self._exc: raise self._exc raise StopAsyncIteration - return item + return cast(ConsumeFromTopicEvent, item) # --- awaitable support --- - def __await__(self): - async def _await_impl(): + def __await__(self) -> Generator[Any, None, list[ConsumeFromTopicEvent]]: + async def _await_impl() -> list[ConsumeFromTopicEvent]: self._ensure_started() await self._done.wait() if self._exc: @@ -111,7 +111,9 @@ async def aclose(self) -> None: pass -def async_func_wrapper(return_value) -> AsyncResult: +def async_func_wrapper( + return_value: AsyncGenerator[ConsumeFromTopicEvent, None], +) -> AsyncResult: """ Normalize a function's *return value* into an AsyncResult. diff --git a/grafi/common/models/function_spec.py b/grafi/common/models/function_spec.py index 53ca7602..01422c61 100644 --- a/grafi/common/models/function_spec.py +++ b/grafi/common/models/function_spec.py @@ -3,8 +3,6 @@ from typing import List from typing import Optional -from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam -from openai.types.shared_params.function_definition import FunctionDefinition from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field @@ -65,16 +63,9 @@ class FunctionSpec(BaseModel): description: Optional[str] parameters: ParametersSchema output_schema: Optional[JsonSchema] = None - - def to_openai_tool(self) -> ChatCompletionToolParam: - return ChatCompletionToolParam( - type="function", - function=FunctionDefinition( - name=self.name, - description=self.description, - parameters=self.parameters.model_dump(), - ), - ) + # Provider conversion (e.g. to OpenAI tool params) lives in provider + # adapters such as grafi.tools.llms.impl.openai_adapter, keeping this core + # model free of any vendor SDK. FunctionSpecs = List[FunctionSpec] diff --git a/grafi/tools/command.py b/grafi/tools/command.py index e18d8713..0f76c036 100644 --- a/grafi/tools/command.py +++ b/grafi/tools/command.py @@ -63,28 +63,6 @@ async def get_tool_input( def to_dict(self) -> dict[str, Any]: return {"class": self.__class__.__name__} - @classmethod - async def from_dict(cls, data: dict[str, Any]) -> "Command": - """ - Create a command instance from a dictionary representation. - - Args: - data (dict[str, Any]): A dictionary representation of the command. - - Returns: - Command: A command instance created from the dictionary. - - Note: - This base implementation returns a Command instance without a tool. - Subclasses should override this method if they need to reconstruct - the tool from the dictionary data. - """ - # Base Command doesn't serialize tool, so we can't reconstruct it - # Subclasses should override this if they need tool reconstruction - raise NotImplementedError( - "from_dict must be implemented by subclasses that need tool reconstruction" - ) - # Registry for tool types to command classes TOOL_COMMAND_REGISTRY: Dict[Type[Tool], Type[Command]] = {} diff --git a/grafi/tools/function_calls/function_call_tool.py b/grafi/tools/function_calls/function_call_tool.py index a994de86..e2dea42e 100644 --- a/grafi/tools/function_calls/function_call_tool.py +++ b/grafi/tools/function_calls/function_call_tool.py @@ -46,8 +46,8 @@ class FunctionCallTool(Tool): name: str = "FunctionCallTool" type: str = "FunctionCallTool" - function_specs: FunctionSpecs = Field(default=[]) - functions: Dict[str, Callable] = Field(default={}) + function_specs: FunctionSpecs = Field(default_factory=list) + functions: Dict[str, Callable] = Field(default_factory=dict) oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.TOOL @classmethod diff --git a/grafi/tools/function_calls/impl/agent_calling_tool.py b/grafi/tools/function_calls/impl/agent_calling_tool.py index 7acdcf9f..b3b07618 100644 --- a/grafi/tools/function_calls/impl/agent_calling_tool.py +++ b/grafi/tools/function_calls/impl/agent_calling_tool.py @@ -37,9 +37,14 @@ def model_post_init(self, _context: Any) -> None: description=self.agent_description, parameters=ParametersSchema( properties={ - "prompt": ParameterSchema( - type="string", - description=self.argument_description, + # ``type`` is an extra JSON-Schema key (ParameterSchema + # allows extras); validate from a dict so it is preserved + # without tripping the typed constructor signature. + "prompt": ParameterSchema.model_validate( + { + "type": "string", + "description": self.argument_description, + } ) }, required=["prompt"], diff --git a/grafi/tools/function_calls/impl/duckduckgo_tool.py b/grafi/tools/function_calls/impl/duckduckgo_tool.py index 27ffe514..6e5a6ebf 100644 --- a/grafi/tools/function_calls/impl/duckduckgo_tool.py +++ b/grafi/tools/function_calls/impl/duckduckgo_tool.py @@ -1,5 +1,6 @@ import json from typing import Any +from typing import Optional from typing import Self from grafi.common.decorators.llm_function import llm_function @@ -98,15 +99,15 @@ async def from_dict(cls, data: dict[str, Any]) -> "DuckDuckGoTool": class DuckDuckGoToolBuilder(FunctionCallToolBuilder[DuckDuckGoTool]): """Builder for DuckDuckGoTool instances.""" - def fixed_max_results(self, fixed_max_results: int) -> Self: + def fixed_max_results(self, fixed_max_results: Optional[int]) -> Self: self.kwargs["fixed_max_results"] = fixed_max_results return self - def headers(self, headers: dict[str, str]) -> Self: + def headers(self, headers: Optional[dict[str, str]]) -> Self: self.kwargs["headers"] = headers return self - def proxy(self, proxy: str) -> Self: + def proxy(self, proxy: Optional[str]) -> Self: self.kwargs["proxy"] = proxy return self diff --git a/grafi/tools/function_calls/impl/synthetic_tool.py b/grafi/tools/function_calls/impl/synthetic_tool.py index cc5ad851..def953f6 100644 --- a/grafi/tools/function_calls/impl/synthetic_tool.py +++ b/grafi/tools/function_calls/impl/synthetic_tool.py @@ -41,7 +41,7 @@ class SyntheticTool(FunctionCallTool): @field_validator("input_model", "output_model") @classmethod - def validate_pydantic_model_or_schema(cls, v: Any, info) -> Any: + def validate_pydantic_model_or_schema(cls, v: Any, info: Any) -> Any: """ Validate that input_model and output_model are either: - A Pydantic BaseModel class (not instance) - for type-safe Python usage @@ -231,8 +231,10 @@ async def _call_llm(self, prompt: str) -> str: }, } - # Use standard chat completion (not parse) - completion = await client.chat.completions.create( + # Use standard chat completion (not parse). The plain message and + # response_format dicts are valid at runtime but do not match the + # SDK's TypedDict overloads, which require literal-keyed construction. + completion = await client.chat.completions.create( # type: ignore[call-overload] model=self.model, messages=[{"role": "user", "content": prompt}], response_format=response_format, diff --git a/grafi/tools/function_calls/impl/tavily_tool.py b/grafi/tools/function_calls/impl/tavily_tool.py index d93bdf48..af00eb9c 100644 --- a/grafi/tools/function_calls/impl/tavily_tool.py +++ b/grafi/tools/function_calls/impl/tavily_tool.py @@ -3,6 +3,7 @@ from typing import Any from typing import Dict from typing import Literal +from typing import Optional from typing import Self from grafi.common.decorators.llm_function import llm_function @@ -114,7 +115,7 @@ async def from_dict(cls, data: dict[str, Any]) -> "TavilyTool": class TavilyToolBuilder(FunctionCallToolBuilder[TavilyTool]): """Builder for TavilyTool instances.""" - def api_key(self, api_key: str) -> Self: + def api_key(self, api_key: Optional[str]) -> Self: from tavily import TavilyClient self.kwargs["client"] = TavilyClient(api_key) diff --git a/grafi/tools/functions/impl/mcp_function_tool.py b/grafi/tools/functions/impl/mcp_function_tool.py index ae25c174..1fbe6a86 100644 --- a/grafi/tools/functions/impl/mcp_function_tool.py +++ b/grafi/tools/functions/impl/mcp_function_tool.py @@ -3,6 +3,8 @@ from typing import AsyncGenerator from typing import Callable from typing import Dict +from typing import Optional +from typing import cast from loguru import logger from pydantic import Field @@ -40,11 +42,17 @@ class MCPFunctionTool(FunctionTool): mcp_config: Dict[str, Any] = Field(default_factory=dict) - function: Callable[[Messages], AsyncGenerator[Messages, None]] = Field(default=None) + # MCP yields str chunks that FunctionTool.to_messages wraps. This return + # type intentionally differs from the base FunctionTool.function contract + # (Callable[..., OutputType]); reconciling the two is tracked for the tools + # refactor, so only the base-override mismatch is suppressed below. + function: Optional[Callable[[Messages], AsyncGenerator[str, None]]] = Field( # type: ignore[assignment] + default=None + ) function_name: str = Field(default="") - _function_spec: FunctionSpec = PrivateAttr(default=None) + _function_spec: Optional[FunctionSpec] = PrivateAttr(default=None) @classmethod async def initialize(cls, **kwargs: Any) -> "MCPFunctionTool": @@ -98,7 +106,7 @@ async def _get_function_spec(self) -> None: async def invoke_mcp_function( self, input_data: Messages, - ) -> AsyncGenerator[Messages, None]: + ) -> AsyncGenerator[str, None]: """ Invoke the MCPFunctionTool with the provided input data. @@ -112,7 +120,8 @@ async def invoke_mcp_function( """ input_message = input_data[-1] - kwargs = json.loads(input_message.content) + # The last message's content is the JSON-encoded tool-call arguments. + kwargs = json.loads(cast(str, input_message.content)) response_str = "" diff --git a/grafi/tools/llms/impl/claude_tool.py b/grafi/tools/llms/impl/claude_tool.py index cbec621b..54269f76 100644 --- a/grafi/tools/llms/impl/claude_tool.py +++ b/grafi/tools/llms/impl/claude_tool.py @@ -24,9 +24,7 @@ from grafi.tools.llms.llm import LLMBuilder try: - from anthropic import NOT_GIVEN from anthropic import AsyncAnthropic - from anthropic import NotGiven from anthropic import Omit from anthropic import omit from anthropic.types import Message as AnthropicMessage @@ -64,7 +62,7 @@ def builder(cls) -> "ClaudeToolBuilder": def prepare_api_input(self, input_data: Messages) -> tuple[ Union[str, Omit], List[MessageParam], - Union[List[ToolParam], NotGiven], + Union[List[ToolParam], Omit], ]: """grafi → Anthropic (system, message list, optional tools). @@ -137,7 +135,7 @@ def prepare_api_input(self, input_data: Messages) -> tuple[ } ) - return system, messages, tools or NOT_GIVEN + return system, messages, tools or omit @staticmethod def _content_to_text(content: Any) -> str: diff --git a/grafi/tools/llms/impl/deepseek_tool.py b/grafi/tools/llms/impl/deepseek_tool.py index 0c35f46b..8f8e690f 100644 --- a/grafi/tools/llms/impl/deepseek_tool.py +++ b/grafi/tools/llms/impl/deepseek_tool.py @@ -1,195 +1,48 @@ """ -DeepseekTool – DeepSeek implementation of grafi.tools.llms.llm.LLM +DeepseekTool – DeepSeek implementation of the OpenAI-compatible LLM tool. -DeepSeek’s HTTP interface is 100 % OpenAI‑compatible, so we reuse the -official `openai` Python SDK and simply change `base_url`. +DeepSeek's HTTP interface is OpenAI-compatible, so it reuses the shared +``OpenAICompatibleTool`` mechanics and only changes ``base_url`` / defaults. -Docs: https://api-docs.deepseek.com – see “Your First API Call”. -The page explicitly says you can call the API with the OpenAI SDK by -setting `base_url="https://api.deepseek.com"`  :contentReference[oaicite:0]{index=0} +Docs: https://api-docs.deepseek.com – call the API with the OpenAI SDK by +setting ``base_url="https://api.deepseek.com"``. """ -import asyncio import os from typing import Any +from typing import ClassVar from typing import Dict -from typing import List from typing import Optional -from typing import Self -from typing import Union -from typing import cast -from openai import AsyncClient -from openai import NotGiven -from openai import OpenAIError -from openai.types.chat import ChatCompletion -from openai.types.chat import ChatCompletionChunk -from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam -from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam from pydantic import Field -from grafi.common.decorators.record_decorators import record_tool_invoke -from grafi.common.exceptions import LLMToolException -from grafi.common.models.invoke_context import InvokeContext -from grafi.common.models.message import Message -from grafi.common.models.message import Messages -from grafi.common.models.message import MsgsAGen -from grafi.tools.llms.llm import LLM -from grafi.tools.llms.llm import LLMBuilder +from grafi.tools.llms.impl.openai_compatible import OpenAICompatibleTool +from grafi.tools.llms.impl.openai_compatible import OpenAICompatibleToolBuilder -class DeepseekTool(LLM): - """ - DeepseekTool – DeepSeek implementation of grafi.tools.llms.llm.LLM - """ +class DeepseekTool(OpenAICompatibleTool): + """DeepSeek implementation of the OpenAI-compatible chat-completions tool.""" name: str = Field(default="DeepseekTool") type: str = Field(default="DeepseekTool") api_key: Optional[str] = Field( default_factory=lambda: os.getenv("DEEPSEEK_API_KEY") ) - base_url: str = Field(default="https://api.deepseek.com") # SDK will append /v1 - model: str = Field(default="deepseek-chat") # or deepseek‑reasoner + base_url: Optional[str] = Field(default="https://api.deepseek.com") + model: str = Field(default="deepseek-chat") # or deepseek-reasoner + + _provider_label: ClassVar[str] = "DeepSeek" @classmethod def builder(cls) -> "DeepseekToolBuilder": - """ - Return a builder for DeepseekTool. - - This method allows for the construction of a DeepseekTool instance with specified parameters. - """ + """Return a builder for DeepseekTool.""" return DeepseekToolBuilder(cls) - # ------------------------------------------------------------------ # - # Shared helper to map grafi → SDK input # - # ------------------------------------------------------------------ # - def prepare_api_input( - self, input_data: Messages - ) -> tuple[ - List[ChatCompletionMessageParam], Union[List[ChatCompletionToolParam], NotGiven] - ]: - api_messages: List[ChatCompletionMessageParam] = ( - [ - cast( - ChatCompletionMessageParam, - {"role": "system", "content": self.system_message}, - ) - ] - if self.system_message - else [] - ) - - for m in input_data: - api_messages.append( - cast( - ChatCompletionMessageParam, - { - "name": m.name, - "role": m.role, - "content": m.content or "", - "tool_calls": m.tool_calls, - "tool_call_id": m.tool_call_id, - }, - ) - ) - - api_tools = [ - function_spec.to_openai_tool() - for function_spec in self.get_function_specs() - ] or None - - return api_messages, api_tools - - # ------------------------------------------------------------------ # - # Async call # - # ------------------------------------------------------------------ # - @record_tool_invoke - async def invoke( - self, - invoke_context: InvokeContext, - input_data: Messages, - ) -> MsgsAGen: - api_messages, api_tools = self.prepare_api_input(input_data) - try: - client = AsyncClient(api_key=self.api_key, base_url=self.base_url) - - if self.is_streaming: - async for chunk in await client.chat.completions.create( - model=self.model, - messages=api_messages, - tools=api_tools, - stream=True, - **self.chat_params, - ): - yield self.to_stream_messages(chunk) - else: - req_func = ( - client.chat.completions.create - if not self.structured_output - else client.beta.chat.completions.parse - ) - response: ChatCompletion = await req_func( - model=self.model, - messages=api_messages, - tools=api_tools, - **self.chat_params, - ) - - yield self.to_messages(response) - except asyncio.CancelledError: - raise # let caller handle - except OpenAIError as exc: - raise LLMToolException( - tool_name=self.name, - model=self.model, - message=f"DeepSeek API streaming failed: {exc}", - invoke_context=invoke_context, - cause=exc, - ) from exc - except Exception as exc: - raise LLMToolException( - tool_name=self.name, - model=self.model, - message=f"Unexpected error during DeepSeek streaming: {exc}", - invoke_context=invoke_context, - cause=exc, - ) from exc - - # ------------------------------------------------------------------ # - # Response converters # - # ------------------------------------------------------------------ # - def to_stream_messages(self, chunk: ChatCompletionChunk) -> Messages: - choice = chunk.choices[0] - delta = choice.delta - data = delta.model_dump() - if data.get("role") is None: - data["role"] = "assistant" - data["is_streaming"] = True - return [Message.model_validate(data)] - - def to_messages(self, resp: ChatCompletion) -> Messages: - return [Message.model_validate(resp.choices[0].message.model_dump())] - - # ------------------------------------------------------------------ # - # Serialisation helper # - # ------------------------------------------------------------------ # - def to_dict(self) -> Dict[str, Any]: - return { - **super().to_dict(), - "base_url": self.base_url, - } + # to_dict (with base_url) is inherited from OpenAICompatibleTool. @classmethod async def from_dict(cls, data: Dict[str, Any]) -> "DeepseekTool": - """ - Create a DeepseekTool instance from a dictionary representation. - - Args: - data (Dict[str, Any]): A dictionary representation of the DeepseekTool. - - Returns: - DeepseekTool: A DeepseekTool instance created from the dictionary. - """ + """Create a DeepseekTool instance from a dictionary representation.""" from openinference.semconv.trace import OpenInferenceSpanKindValues return ( @@ -207,13 +60,5 @@ async def from_dict(cls, data: Dict[str, Any]) -> "DeepseekTool": ) -class DeepseekToolBuilder(LLMBuilder[DeepseekTool]): +class DeepseekToolBuilder(OpenAICompatibleToolBuilder[DeepseekTool]): """Builder for DeepseekTool instances.""" - - def base_url(self, base_url: str) -> Self: - self.kwargs["base_url"] = base_url.rstrip("/") - return self - - def api_key(self, api_key: Optional[str]) -> Self: - self.kwargs["api_key"] = api_key - return self diff --git a/grafi/tools/llms/impl/ollama_tool.py b/grafi/tools/llms/impl/ollama_tool.py index 4d3761fc..fa8422bb 100644 --- a/grafi/tools/llms/impl/ollama_tool.py +++ b/grafi/tools/llms/impl/ollama_tool.py @@ -19,6 +19,7 @@ from grafi.common.models.message import Message from grafi.common.models.message import Messages from grafi.common.models.message import MsgsAGen +from grafi.tools.llms.impl.openai_adapter import to_openai_tool from grafi.tools.llms.llm import LLM from grafi.tools.llms.llm import LLMBuilder @@ -94,8 +95,7 @@ def prepare_api_input( # Extract function specifications from self.get_function_specs() api_functions = [ - function_spec.to_openai_tool() - for function_spec in self.get_function_specs() + to_openai_tool(function_spec) for function_spec in self.get_function_specs() ] or None return api_messages, api_functions diff --git a/grafi/tools/llms/impl/openai_adapter.py b/grafi/tools/llms/impl/openai_adapter.py new file mode 100644 index 00000000..16f13764 --- /dev/null +++ b/grafi/tools/llms/impl/openai_adapter.py @@ -0,0 +1,23 @@ +"""OpenAI adapter: convert provider-neutral Grafi models to OpenAI SDK types. + +Keeping vendor conversion here (rather than on the core ``FunctionSpec`` model) +lets ``grafi.common.models`` stay free of the OpenAI SDK, and concentrates +OpenAI-specific knowledge in the OpenAI tool layer (Open/Closed + DIP). +""" + +from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam +from openai.types.shared_params.function_definition import FunctionDefinition + +from grafi.common.models.function_spec import FunctionSpec + + +def to_openai_tool(spec: FunctionSpec) -> ChatCompletionToolParam: + """Convert a Grafi :class:`FunctionSpec` into an OpenAI tool parameter.""" + return ChatCompletionToolParam( + type="function", + function=FunctionDefinition( + name=spec.name, + description=spec.description, + parameters=spec.parameters.model_dump(), + ), + ) diff --git a/grafi/tools/llms/impl/openai_compatible.py b/grafi/tools/llms/impl/openai_compatible.py new file mode 100644 index 00000000..693387ac --- /dev/null +++ b/grafi/tools/llms/impl/openai_compatible.py @@ -0,0 +1,213 @@ +"""Shared base for OpenAI-compatible chat-completions providers. + +OpenAI, DeepSeek, and OpenRouter all speak the OpenAI chat-completions API via +the official ``openai`` SDK, differing only in endpoint/config (base URL, extra +headers, default model, API-key env var) and provider-specific error wording. +This base centralizes the identical request/response mechanics; concrete +providers specialize via a few small hooks. +""" + +import asyncio +from typing import Any +from typing import ClassVar +from typing import Dict +from typing import List +from typing import Optional +from typing import Self +from typing import TypeVar +from typing import Union +from typing import cast + +from openai import AsyncClient +from openai import AsyncStream +from openai import Omit +from openai import OpenAIError +from openai import omit +from openai.types.chat import ChatCompletion +from openai.types.chat import ChatCompletionChunk +from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam +from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam +from pydantic import Field + +from grafi.common.decorators.record_decorators import record_tool_invoke +from grafi.common.exceptions import LLMToolException +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.common.models.message import Messages +from grafi.common.models.message import MsgsAGen +from grafi.tools.llms.impl.openai_adapter import to_openai_tool +from grafi.tools.llms.llm import LLM +from grafi.tools.llms.llm import LLMBuilder + + +class OpenAICompatibleTool(LLM): + """Base for providers exposed through the OpenAI chat-completions API.""" + + # Human-readable provider name used in error messages. + _provider_label: ClassVar[str] = "OpenAI-compatible" + + # Whether structured-output requests may use the OpenAI-specific + # ``beta.chat.completions.parse`` endpoint. Providers that only emulate the + # standard chat-completions API (e.g. OpenRouter) set this False and rely on + # ``response_format`` in ``chat_params`` via the normal ``create`` call. + _supports_beta_parse: ClassVar[bool] = True + + # Optional endpoint override (``None`` uses the SDK's default OpenAI base). + base_url: Optional[str] = Field(default=None) + + def _extra_create_kwargs(self) -> Dict[str, Any]: + """Provider-specific keyword arguments merged into every request. + + Defaults to none; OpenRouter overrides this to attach extra headers. + """ + return {} + + def prepare_api_input( + self, input_data: Messages + ) -> tuple[ + List[ChatCompletionMessageParam], Union[List[ChatCompletionToolParam], Omit] + ]: + """Convert Grafi messages + function specs into SDK request parameters.""" + api_messages: List[ChatCompletionMessageParam] = ( + [ + cast( + ChatCompletionMessageParam, + {"role": "system", "content": self.system_message}, + ) + ] + if self.system_message + else [] + ) + + for message in input_data: + api_messages.append( + cast( + ChatCompletionMessageParam, + { + "name": message.name, + "role": message.role, + "content": message.content or "", + "tool_calls": message.tool_calls, + "tool_call_id": message.tool_call_id, + }, + ) + ) + + api_tools = [ + to_openai_tool(function_spec) for function_spec in self.get_function_specs() + ] or omit + + return api_messages, api_tools + + @record_tool_invoke + async def invoke( + self, + invoke_context: InvokeContext, + input_data: Messages, + ) -> MsgsAGen: + api_messages, api_tools = self.prepare_api_input(input_data) + + client_kwargs: Dict[str, Any] = {"api_key": self.api_key} + if self.base_url: + client_kwargs["base_url"] = self.base_url + + # ``name`` is Optional on the base Tool; concrete providers always set it, + # but fall back to the provider label so the error always has a name. + tool_name = self.name or self._provider_label + + # Base/provider kwargs merged with user chat_params (chat_params last, so + # a caller-supplied key overrides rather than collides with a duplicate + # keyword argument). + call_kwargs: Dict[str, Any] = { + "model": self.model, + "messages": api_messages, + "tools": api_tools, + **self._extra_create_kwargs(), + **self.chat_params, + } + + try: + async with AsyncClient(**client_kwargs) as client: + if self.is_streaming: + # ``**call_kwargs`` (Any) defeats overload resolution on + # ``stream=True``, so cast the result to the streaming type. + stream = cast( + AsyncStream[ChatCompletionChunk], + await client.chat.completions.create( + stream=True, **call_kwargs + ), + ) + async for chunk in stream: + yield self.to_stream_messages(chunk) + else: + use_parse = self.structured_output and self._supports_beta_parse + req_func = ( + client.beta.chat.completions.parse + if use_parse + else client.chat.completions.create + ) + response = cast( + ChatCompletion, + await req_func(**call_kwargs), + ) + yield self.to_messages(response) + except asyncio.CancelledError: + raise # let caller handle + except OpenAIError as exc: + raise LLMToolException( + tool_name=tool_name, + model=self.model, + message=f"{self._provider_label} API call failed: {exc}", + invoke_context=invoke_context, + cause=exc, + ) from exc + except Exception as exc: + raise LLMToolException( + tool_name=tool_name, + model=self.model, + message=f"Unexpected error during {self._provider_label} call: {exc}", + invoke_context=invoke_context, + cause=exc, + ) from exc + + def to_stream_messages(self, chunk: ChatCompletionChunk) -> Messages: + """Convert one streaming chunk into Grafi messages.""" + choice = chunk.choices[0] + data = choice.delta.model_dump() + if data.get("role") is None: + data["role"] = "assistant" + data["is_streaming"] = True + return [Message.model_validate(data)] + + def to_messages(self, response: ChatCompletion) -> Messages: + """Convert a non-streaming response into Grafi messages.""" + return [Message.model_validate(response.choices[0].message.model_dump())] + + def to_dict(self) -> Dict[str, Any]: + """Serialize the tool, including base_url when a provider sets one. + + OpenAI (no base_url) omits the key, matching its prior manifest shape; + DeepSeek/OpenRouter inherit base_url serialization here instead of + repeating it. + """ + data = super().to_dict() + if self.base_url: + data["base_url"] = self.base_url + return data + + +T_OAC = TypeVar("T_OAC", bound=OpenAICompatibleTool) + + +class OpenAICompatibleToolBuilder(LLMBuilder[T_OAC]): + """Builder for OpenAI-compatible tools.""" + + def api_key(self, api_key: Optional[str]) -> Self: + """Set the provider API key.""" + self.kwargs["api_key"] = api_key + return self + + def base_url(self, base_url: str) -> Self: + """Set the provider base URL (trailing slash trimmed).""" + self.kwargs["base_url"] = base_url.rstrip("/") + return self diff --git a/grafi/tools/llms/impl/openai_tool.py b/grafi/tools/llms/impl/openai_tool.py index d548f0b2..d15f1b4f 100644 --- a/grafi/tools/llms/impl/openai_tool.py +++ b/grafi/tools/llms/impl/openai_tool.py @@ -1,38 +1,18 @@ -import asyncio import os from typing import Any +from typing import ClassVar from typing import Dict -from typing import List from typing import Optional -from typing import Self -from typing import Union -from typing import cast -from openai import NOT_GIVEN -from openai import AsyncClient -from openai import NotGiven -from openai import OpenAIError -from openai.types.chat import ChatCompletion -from openai.types.chat import ChatCompletionChunk -from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam -from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam from pydantic import Field -from grafi.common.decorators.record_decorators import record_tool_invoke -from grafi.common.exceptions import LLMToolException -from grafi.common.models.invoke_context import InvokeContext -from grafi.common.models.message import Message -from grafi.common.models.message import Messages -from grafi.common.models.message import MsgsAGen -from grafi.tools.llms.llm import LLM -from grafi.tools.llms.llm import LLMBuilder +from grafi.tools.llms.impl.openai_compatible import OpenAICompatibleTool +from grafi.tools.llms.impl.openai_compatible import OpenAICompatibleToolBuilder -class OpenAITool(LLM): +class OpenAITool(OpenAICompatibleTool): """ - A class representing the OpenAI language model implementation. - - This class provides methods to interact with OpenAI's API for natural language processing tasks. + OpenAI implementation of the OpenAI-compatible chat-completions tool. Attributes: api_key (str): The API key for authenticating with OpenAI. @@ -44,173 +24,16 @@ class OpenAITool(LLM): api_key: Optional[str] = Field(default_factory=lambda: os.getenv("OPENAI_API_KEY")) model: str = Field(default="gpt-4o-mini") + _provider_label: ClassVar[str] = "OpenAI" + @classmethod def builder(cls) -> "OpenAIToolBuilder": - """ - Return a builder for OpenAITool. - - This method allows for the construction of an OpenAITool instance with specified parameters. - """ + """Return a builder for OpenAITool.""" return OpenAIToolBuilder(cls) - def prepare_api_input( - self, input_data: Messages - ) -> tuple[ - List[ChatCompletionMessageParam], Union[List[ChatCompletionToolParam], NotGiven] - ]: - """ - Prepare the input data for the OpenAI API. - - Args: - input_data (Messages): A list of Message objects. - - Returns: - tuple: A tuple containing: - - A list of dictionaries representing the messages for the API. - - A list of function specifications for the API, or None if no functions are present. - """ - api_messages = ( - [ - cast( - ChatCompletionMessageParam, - {"role": "system", "content": self.system_message}, - ) - ] - if self.system_message - else [] - ) - - for message in input_data: - api_message = { - "name": message.name, - "role": message.role, - "content": message.content or "", - "tool_calls": message.tool_calls, - "tool_call_id": message.tool_call_id, - } - api_messages.append(cast(ChatCompletionMessageParam, api_message)) - - # Extract function specifications if present in latest message - - api_tools = [ - function_spec.to_openai_tool() - for function_spec in self.get_function_specs() - ] or NOT_GIVEN - - return api_messages, api_tools - - @record_tool_invoke - async def invoke( - self, - invoke_context: InvokeContext, - input_data: Messages, - ) -> MsgsAGen: - api_messages, api_tools = self.prepare_api_input(input_data) - try: - async with AsyncClient(api_key=self.api_key) as client: - if self.is_streaming: - async for chunk in await client.chat.completions.create( - model=self.model, - messages=api_messages, - tools=api_tools, - stream=True, - **self.chat_params, - ): - yield self.to_stream_messages(chunk) - else: - req_func = ( - client.chat.completions.create - if not self.structured_output - else client.beta.chat.completions.parse - ) - response: ChatCompletion = await req_func( - model=self.model, - messages=api_messages, - tools=api_tools, - **self.chat_params, - ) - - yield self.to_messages(response) - except asyncio.CancelledError: - raise # let caller handle - except OpenAIError as exc: - raise LLMToolException( - tool_name=self.name, - model=self.model, - message=f"OpenAI API streaming failed: {exc}", - invoke_context=invoke_context, - cause=exc, - ) from exc - except Exception as e: - raise LLMToolException( - tool_name=self.name, - model=self.model, - message=f"Unexpected error during OpenAI streaming: {e}", - invoke_context=invoke_context, - cause=e, - ) from e - - def to_stream_messages(self, chunk: ChatCompletionChunk) -> Messages: - """ - Convert an OpenAI API response to a Message object. - - This method extracts relevant information from the API response and constructs a Message object. - - Args: - response (ChatCompletion): The response object from the OpenAI API. - - Returns: - Message: A Message object containing the extracted information from the API response. - """ - - # Extract the first choice - choice = chunk.choices[0] - message_data = choice.delta - data = message_data.model_dump() - if data.get("role") is None: - data["role"] = "assistant" - data["is_streaming"] = True - return [Message.model_validate(data)] - - def to_messages(self, response: ChatCompletion) -> Messages: - """ - Convert an OpenAI API response to a Message object. - - This method extracts relevant information from the API response and constructs a Message object. - - Args: - response (ChatCompletion): The response object from the OpenAI API. - - Returns: - Message: A Message object containing the extracted information from the API response. - """ - - # Extract the first choice - choice = response.choices[0] - return [Message.model_validate(choice.message.model_dump())] - - def to_dict(self) -> Dict[str, Any]: - """ - Convert the OpenAITool instance to a dictionary. - - Returns: - dict: A dictionary containing the attributes of the OpenAITool instance. - """ - return { - **super().to_dict(), - } - @classmethod async def from_dict(cls, data: Dict[str, Any]) -> "OpenAITool": - """ - Create an OpenAITool instance from a dictionary representation. - - Args: - data (Dict[str, Any]): A dictionary representation of the OpenAITool. - - Returns: - OpenAITool: An OpenAITool instance created from the dictionary. - """ + """Create an OpenAITool instance from a dictionary representation.""" from openinference.semconv.trace import OpenInferenceSpanKindValues return ( @@ -227,17 +50,5 @@ async def from_dict(cls, data: Dict[str, Any]) -> "OpenAITool": ) -class OpenAIToolBuilder(LLMBuilder[OpenAITool]): +class OpenAIToolBuilder(OpenAICompatibleToolBuilder[OpenAITool]): """Builder for OpenAITool instances.""" - - def api_key(self, api_key: Optional[str]) -> Self: - """Set the OpenAI API key. - - Args: - api_key: The API key for OpenAI authentication. - - Returns: - Self for method chaining. - """ - self.kwargs["api_key"] = api_key - return self diff --git a/grafi/tools/llms/impl/openrouter_tool.py b/grafi/tools/llms/impl/openrouter_tool.py index 41acfcc4..ebc93408 100644 --- a/grafi/tools/llms/impl/openrouter_tool.py +++ b/grafi/tools/llms/impl/openrouter_tool.py @@ -1,193 +1,59 @@ """ -OpenRouterTool - OpenRouter.ai implementation of grafi.tools.llms.llm.LLM +OpenRouterTool - OpenRouter.ai implementation of the OpenAI-compatible LLM tool. """ -import asyncio import os from typing import Any +from typing import ClassVar from typing import Dict -from typing import List from typing import Optional from typing import Self -from typing import Union -from typing import cast - -from openai import AsyncClient -from openai import NotGiven -from openai import OpenAIError -from openai.types.chat import ChatCompletion -from openai.types.chat import ChatCompletionChunk -from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam -from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam + from pydantic import Field -from grafi.common.decorators.record_decorators import record_tool_invoke -from grafi.common.exceptions import LLMToolException -from grafi.common.models.invoke_context import InvokeContext -from grafi.common.models.message import Message -from grafi.common.models.message import Messages -from grafi.common.models.message import MsgsAGen -from grafi.tools.llms.llm import LLM -from grafi.tools.llms.llm import LLMBuilder +from grafi.tools.llms.impl.openai_compatible import OpenAICompatibleTool +from grafi.tools.llms.impl.openai_compatible import OpenAICompatibleToolBuilder -class OpenRouterTool(LLM): - """ - OpenRouterTool - OpenRouter.ai implementation of grafi.tools.llms.llm.LLM - """ +class OpenRouterTool(OpenAICompatibleTool): + """OpenRouter.ai implementation of the OpenAI-compatible chat-completions tool.""" name: str = Field(default="OpenRouterTool") type: str = Field(default="OpenRouterTool") - api_key: Optional[str] = Field( default_factory=lambda: os.getenv("OPENROUTER_API_KEY") ) - base_url: str = Field(default="https://openrouter.ai/api/v1") + base_url: Optional[str] = Field(default="https://openrouter.ai/api/v1") model: str = Field(default="openrouter/auto") # Auto-router chooses best model # extra headers for leader-board visibility (optional) extra_headers: Dict[str, str] = Field(default_factory=dict) - chat_params: Dict[str, Any] = Field(default_factory=dict) + _provider_label: ClassVar[str] = "OpenRouter" + # OpenRouter emulates the chat-completions API but not OpenAI's beta + # parsed-completions endpoint; structured output goes through response_format + # on the standard create call. + _supports_beta_parse: ClassVar[bool] = False + + def _extra_create_kwargs(self) -> Dict[str, Any]: + # OpenRouter accepts optional attribution headers on each request. + return {"extra_headers": self.extra_headers or None} @classmethod def builder(cls) -> "OpenRouterToolBuilder": - """ - Return a builder for OpenRouterTool. - - This method allows for the construction of an OpenRouterTool instance with specified parameters. - """ + """Return a builder for OpenRouterTool.""" return OpenRouterToolBuilder(cls) - # ------------------------------------------------------------------ # - # Request conversion helper # - # ------------------------------------------------------------------ # - def prepare_api_input( - self, input_data: Messages - ) -> tuple[ - List[ChatCompletionMessageParam], Union[List[ChatCompletionToolParam], NotGiven] - ]: - api_messages: List[ChatCompletionMessageParam] = ( - [ - cast( - ChatCompletionMessageParam, - {"role": "system", "content": self.system_message}, - ) - ] - if self.system_message - else [] - ) - - for m in input_data: - api_messages.append( - cast( - ChatCompletionMessageParam, - { - "name": m.name, - "role": m.role, - "content": m.content or "", - "tool_calls": m.tool_calls, - "tool_call_id": m.tool_call_id, - }, - ) - ) - - # Extract function specifications from self.get_function_specs() - api_tools = [ - function_spec.to_openai_tool() - for function_spec in self.get_function_specs() - ] or None - - return api_messages, api_tools - - # ------------------------------------------------------------------ # - # Async call # - # ------------------------------------------------------------------ # - @record_tool_invoke - async def invoke( - self, - invoke_context: InvokeContext, - input_data: Messages, - ) -> MsgsAGen: - messages, tools = self.prepare_api_input(input_data) - - try: - client = AsyncClient(api_key=self.api_key, base_url=self.base_url) - - if self.is_streaming: - async for chunk in await client.chat.completions.create( - model=self.model, - messages=messages, - tools=tools, - stream=True, - extra_headers=self.extra_headers or None, - **self.chat_params, - ): - yield self.to_stream_messages(chunk) - else: - resp: ChatCompletion = await client.chat.completions.create( - model=self.model, - messages=messages, - tools=tools, - extra_headers=self.extra_headers or None, - **self.chat_params, - ) - yield self.to_messages(resp) - except asyncio.CancelledError: - raise - except OpenAIError as exc: - raise LLMToolException( - tool_name=self.name, - model=self.model, - message=f"OpenRouter async call failed: {exc}", - invoke_context=invoke_context, - cause=exc, - ) from exc - except Exception as exc: - raise LLMToolException( - tool_name=self.name, - model=self.model, - message=f"Unexpected error during OpenRouter async call: {exc}", - invoke_context=invoke_context, - cause=exc, - ) from exc - - # ------------------------------------------------------------------ # - # Response converters # - # ------------------------------------------------------------------ # - def to_stream_messages(self, chunk: ChatCompletionChunk) -> Messages: - choice = chunk.choices[0] - delta = choice.delta - data = delta.model_dump() - if data.get("role") is None: - data["role"] = "assistant" - data["is_streaming"] = True - return [Message.model_validate(data)] - - def to_messages(self, resp: ChatCompletion) -> Messages: - return [Message.model_validate(resp.choices[0].message.model_dump())] - - # ------------------------------------------------------------------ # - # Serialisation # - # ------------------------------------------------------------------ # def to_dict(self) -> Dict[str, Any]: + # base_url is serialized by OpenAICompatibleTool.to_dict. return { **super().to_dict(), - "base_url": self.base_url, "extra_headers": self.extra_headers, } @classmethod async def from_dict(cls, data: Dict[str, Any]) -> "OpenRouterTool": - """ - Create an OpenRouterTool instance from a dictionary representation. - - Args: - data (Dict[str, Any]): A dictionary representation of the OpenRouterTool. - - Returns: - OpenRouterTool: An OpenRouterTool instance created from the dictionary. - """ + """Create an OpenRouterTool instance from a dictionary representation.""" from openinference.semconv.trace import OpenInferenceSpanKindValues return ( @@ -199,25 +65,16 @@ async def from_dict(cls, data: Dict[str, Any]) -> "OpenRouterTool": .is_streaming(data.get("is_streaming", False)) .system_message(data.get("system_message", "")) .api_key(os.getenv("OPENROUTER_API_KEY")) + .model(data.get("model", "openrouter/auto")) .base_url(data.get("base_url", "https://openrouter.ai/api/v1")) .extra_headers(data.get("extra_headers", {})) .build() ) -class OpenRouterToolBuilder(LLMBuilder[OpenRouterTool]): - """ - Builder for OpenRouterTool. - """ - - def base_url(self, base_url: str) -> Self: - self.kwargs["base_url"] = base_url.rstrip("/") - return self +class OpenRouterToolBuilder(OpenAICompatibleToolBuilder[OpenRouterTool]): + """Builder for OpenRouterTool instances.""" def extra_headers(self, headers: Dict[str, str]) -> Self: self.kwargs["extra_headers"] = headers return self - - def api_key(self, api_key: Optional[str]) -> Self: - self.kwargs["api_key"] = api_key - return self diff --git a/grafi/tools/llms/llm.py b/grafi/tools/llms/llm.py index 05bb9a73..35860fe6 100644 --- a/grafi/tools/llms/llm.py +++ b/grafi/tools/llms/llm.py @@ -59,7 +59,7 @@ def _is_object_schema(node: Json) -> bool: SCHEMA_KEYS_ARRAY = ("allOf", "anyOf", "oneOf") # keys that can hold a schema or array of schemas - def _recurse(node: Any): + def _recurse(node: Any) -> None: if isinstance(node, dict): # Dive into $defs/definitions first (Pydantic v2 uses $defs) for defs_key in ("$defs", "definitions"): @@ -188,7 +188,7 @@ def _serialize_chat_params(self, params: Dict[str, Any]) -> Dict[str, Any]: Converts Pydantic v2 model instances and classes to their dict representation. """ - serialized_params = {} + serialized_params: Dict[str, Any] = {} for key, value in params.items(): if isinstance(value, BaseModel): # Use model_dump() for Pydantic v2 model instances diff --git a/grafi/tools/tool_factory.py b/grafi/tools/tool_factory.py index b7e95888..0512bde3 100644 --- a/grafi/tools/tool_factory.py +++ b/grafi/tools/tool_factory.py @@ -136,8 +136,9 @@ async def from_dict(cls, data: Dict[str, Any]) -> Tool: if tool_class is None: tool_class = cls._resolve_lazy(class_name) - if tool_class is None and data.get("base_class") is not None: - tool_class = cls._TOOL_REGISTRY.get(data.get("base_class")) + base_class = data.get("base_class") + if tool_class is None and base_class is not None: + tool_class = cls._TOOL_REGISTRY.get(base_class) if tool_class is None: raise ValueError( diff --git a/grafi/topics/expressions/subscription_builder.py b/grafi/topics/expressions/subscription_builder.py index 1541ebac..1d55ae32 100644 --- a/grafi/topics/expressions/subscription_builder.py +++ b/grafi/topics/expressions/subscription_builder.py @@ -71,8 +71,16 @@ def or_(self) -> "SubscriptionBuilder": def build(self) -> SubExpr: """ - Attach the final expression (root_expr) to the Node's subscribed_expressions, - then return the Node.Builder for further chaining. + Return the built subscription expression. + + Raises: + ValueError: If no topic was subscribed (nothing to build), so the + ``-> SubExpr`` contract always holds rather than silently + returning ``None``. """ - if self.root_expr: - return self.root_expr + if self.root_expr is None: + raise ValueError( + "Cannot build an empty subscription. " + "Call .subscribed_to(topic) before .build()." + ) + return self.root_expr diff --git a/grafi/topics/queue_impl/in_mem_topic_event_queue.py b/grafi/topics/queue_impl/in_mem_topic_event_queue.py index e3835d5d..6358d10d 100644 --- a/grafi/topics/queue_impl/in_mem_topic_event_queue.py +++ b/grafi/topics/queue_impl/in_mem_topic_event_queue.py @@ -102,6 +102,23 @@ async def commit_to(self, consumer_id: str, offset: int) -> int: return self._committed[consumer_id] + async def restore_consumer(self, consumer_id: str, committed_offset: int) -> None: + """Restore a consumer's cursors from a recorded commit point. + + See :meth:`TopicEventQueue.restore_consumer`. Idempotent and + order-independent: replaying the same or an earlier commit can only + advance a cursor, never rewind it, so events restored out of order still + reconstruct the correct position. + """ + async with self._cond: + # The next offset to read is one past the last committed offset. + self._consumed[consumer_id] = max( + self._consumed[consumer_id], committed_offset + 1 + ) + self._committed[consumer_id] = max( + self._committed[consumer_id], committed_offset + ) + async def reset(self) -> None: """ Reset the queue to its initial state asynchronously. diff --git a/grafi/topics/topic_base.py b/grafi/topics/topic_base.py index bffc81f0..3544640b 100644 --- a/grafi/topics/topic_base.py +++ b/grafi/topics/topic_base.py @@ -146,12 +146,12 @@ async def restore_topic(self, topic_event: TopicEvent) -> None: if isinstance(topic_event, PublishToTopicEvent): await self.event_queue.put(topic_event) elif isinstance(topic_event, ConsumeFromTopicEvent): - # Fetch the events for the consumer and commit the offset - await self.event_queue.fetch( - consumer_id=topic_event.consumer_name, offset=topic_event.offset + 1 - ) - await self.event_queue.commit_to( - topic_event.consumer_name, topic_event.offset + # Restore the consumer's cursor to the recorded commit point. The + # next event delivered is topic_event.offset + 1, so the offset + # immediately after the committed one is not skipped. + await self.event_queue.restore_consumer( + consumer_id=topic_event.consumer_name, + committed_offset=topic_event.offset, ) async def add_event(self, event: TopicEvent) -> Optional[TopicEvent]: diff --git a/grafi/topics/topic_event_queue.py b/grafi/topics/topic_event_queue.py index c8648b95..cabb247a 100644 --- a/grafi/topics/topic_event_queue.py +++ b/grafi/topics/topic_event_queue.py @@ -60,6 +60,27 @@ async def commit_to(self, consumer_id: str, offset: int) -> int: """ pass + @abstractmethod + async def restore_consumer(self, consumer_id: str, committed_offset: int) -> None: + """ + Restore a consumer's cursors during recovery, without consuming events. + + Marks the consumer as having consumed and committed every offset up to + and including ``committed_offset``. The next event delivered to this + consumer is ``committed_offset + 1``: nothing at or before + ``committed_offset`` is re-delivered, and nothing after it is skipped. + + Unlike :meth:`fetch`, this never blocks and never returns events. It is a + pure cursor restore used only by replay-based recovery, so its meaning is + unambiguous (``fetch`` over-advances the consumed cursor by one when + misused for restore). + + Args: + consumer_id: Unique identifier for the consumer + committed_offset: Highest offset the consumer had committed + """ + ... + @abstractmethod async def reset(self) -> None: """ diff --git a/grafi/workflows/impl/async_node_tracker.py b/grafi/workflows/impl/async_node_tracker.py index a3e2c6fd..c89dbf19 100644 --- a/grafi/workflows/impl/async_node_tracker.py +++ b/grafi/workflows/impl/async_node_tracker.py @@ -126,13 +126,33 @@ async def on_messages_published(self, count: int = 1, source: str = "") -> None: async def on_messages_committed(self, count: int = 1, source: str = "") -> None: """ - Called when messages are committed (consumed and acknowledged). - - Call site: _commit_events() in EventDrivenWorkflow + Called when deliveries are committed (consumed and acknowledged) or + released (a publication whose condition was not met). + + The pending counter tracks outstanding *deliveries* (one per consumer of + each published event), and deliveries are registered before an event + becomes consumable, so a commit can never outpace its registration. A + genuine underflow therefore signals a real accounting bug (e.g. a + double-commit). We log it loudly (rather than silently absorbing it as + the old ``max(0, ...)`` clamp did) and then clamp to 0 to keep the + counter valid; the parallel run is additionally backstopped by the output + queue's no-progress detection, so a miscount cannot hang the workflow. + + Call sites: _commit_events() in EventDrivenWorkflow, the output listener, + and publish_events() (releasing an unpublished, condition-filtered event). """ if count <= 0: return async with self._cond: + if count > self._uncommitted_messages: + # Tripwire: with deliveries registered before publication this + # should be unreachable; log loudly rather than silently absorb. + logger.error( + f"Tracker: delivery underflow committing {count} from " + f"{source} (pending={self._uncommitted_messages}); " + "clamping to 0 -- this indicates a double-commit or " + "accounting bug." + ) self._uncommitted_messages = max(0, self._uncommitted_messages - count) self._check_quiescence_unlocked() diff --git a/grafi/workflows/impl/async_output_queue.py b/grafi/workflows/impl/async_output_queue.py index 3c7452ee..c8a9eff2 100644 --- a/grafi/workflows/impl/async_output_queue.py +++ b/grafi/workflows/impl/async_output_queue.py @@ -1,4 +1,6 @@ import asyncio +from typing import Awaitable +from typing import Callable from typing import List from typing import Optional @@ -21,10 +23,23 @@ def __init__( output_topics: List[TopicBase], consumer_name: str, tracker: AsyncNodeTracker, + progress_possible: Optional[Callable[[], Awaitable[bool]]] = None, ): self.output_topics = output_topics self.consumer_name = consumer_name self.tracker = tracker + # Optional workflow callback: True while progress is still possible + # (a node is active, an event is still consumable, or some node can + # invoke). When this stays False while the tracker is non-quiescent, the + # outstanding deliveries are parked (e.g. an unsatisfied AND-subscription + # whose other branch never arrived) and iteration ends rather than hangs. + self._progress_possible = progress_possible + self._stuck_polls = 0 + # Tracker activity count at the last idle poll. Any change means a node + # processed between polls (real progress), which resets the stuck count + # even if no output-queue item appeared -- so termination never depends + # on the transient consume->enter window being unobservable. + self._last_activity = -1 self.queue: asyncio.Queue[TopicEvent] = asyncio.Queue() self._listener_tasks: List[asyncio.Task] = [] self._stopped = False @@ -112,7 +127,9 @@ async def __anext__(self) -> TopicEvent: # Fast path: queue has items if not self.queue.empty(): try: - return self.queue.get_nowait() + item = self.queue.get_nowait() + self._stuck_polls = 0 # progress made + return item except asyncio.QueueEmpty: pass @@ -146,7 +163,9 @@ async def __anext__(self) -> TopicEvent: # Got queue item if queue_task in done and not queue_task.cancelled(): try: - return queue_task.result() + item = queue_task.result() + self._stuck_polls = 0 # progress made + return item except asyncio.QueueEmpty: # Task was cancelled as part of normal cleanup; ignore. continue @@ -163,3 +182,28 @@ async def __anext__(self) -> TopicEvent: if self._listener_error is not None: raise self._listener_error raise StopAsyncIteration + + # Not quiescent and no queue item after the wait. If the workflow + # reports no further progress is possible, the outstanding tracker + # count is parked work that will never commit, so end iteration + # instead of looping forever. Terminate only after consecutive idle + # polls during which (a) no progress is possible AND (b) the tracker's + # activity count did not change -- i.e. no node processed in between. + # The activity-count guard means a node caught in the transient + # consume->enter window at one poll is detected as progress at the + # next (it will have entered), so a busy run is never mistaken for + # stuck regardless of scheduling. + if self._progress_possible is not None and self.queue.empty(): + activity = await self.tracker.get_activity_count() + progressed = activity != self._last_activity + self._last_activity = activity + if progressed or await self._progress_possible(): + self._stuck_polls = 0 + else: + self._stuck_polls += 1 + if self._stuck_polls >= 2: + logger.debug( + "Output queue: no progress possible; ending iteration " + "(parked deliveries)." + ) + raise StopAsyncIteration diff --git a/grafi/workflows/impl/event_driven_workflow.py b/grafi/workflows/impl/event_driven_workflow.py index 2f2c3972..8070bac1 100644 --- a/grafi/workflows/impl/event_driven_workflow.py +++ b/grafi/workflows/impl/event_driven_workflow.py @@ -55,14 +55,14 @@ class EventDrivenWorkflow(Workflow): oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.AGENT # Topics known to this workflow (e.g., "agent_input", "agent_stream_output") - _topics: Dict[str, TopicBase] = PrivateAttr(default={}) + _topics: Dict[str, TopicBase] = PrivateAttr(default_factory=dict) # Mapping of topic_name -> list of node_names that subscribe to that topic - _topic_nodes: Dict[str, List[str]] = PrivateAttr(default={}) + _topic_nodes: Dict[str, List[str]] = PrivateAttr(default_factory=dict) # Event graph for this workflow # Queue of nodes that are ready to invoke (in response to published events) - _invoke_queue: deque[NodeBase] = PrivateAttr(default=deque()) + _invoke_queue: deque[NodeBase] = PrivateAttr(default_factory=deque) _tracker: AsyncNodeTracker = PrivateAttr(default_factory=AsyncNodeTracker) @@ -241,27 +241,58 @@ async def _commit_events( len(topic_events), source=f"commit:{consumer_name}" ) + def _topic_consumers(self, topic_name: str) -> List[str]: + """Consumers that will commit each event published to ``topic_name``. + + These are the subscribing nodes, plus the workflow itself for output + topics (the output listener / ``_get_output_events`` drains them under + ``self.name``). This is the single source of truth for fan-out: one event + published here yields one delivery -- and one eventual commit -- per + consumer. Deduplicated, since a node consumes a topic once regardless of + how many of its subscription expressions reference it. + """ + consumers = list(dict.fromkeys(self._topic_nodes.get(topic_name, []))) + topic = self._topics.get(topic_name) + if topic is not None and topic.type in ( + TopicType.AGENT_OUTPUT_TOPIC_TYPE, + TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE, + ): + consumers.append(self.name) + return consumers + async def _count_pending_consumable(self) -> int: - """Count messages still awaiting consumption across all topics. + """Count deliveries still awaiting consumption across all topics. - For each topic this sums the unconsumed messages per consumer: subscriber - nodes for any topic, plus the workflow itself for output topics (which the - output listener drains under ``self.name``). The total equals the number - of commit operations the resumed run must perform to drain the restored - state, which is exactly what the quiescence tracker needs seeded. + For each topic this sums the unconsumed events per consumer. The total + equals the number of commit operations the resumed run must perform to + drain the restored state, which is exactly what the quiescence tracker + needs seeded on recovery. """ total = 0 for topic_name, topic in self._topics.items(): - consumer_names = list(self._topic_nodes.get(topic_name, [])) - if topic.type in ( - TopicType.AGENT_OUTPUT_TOPIC_TYPE, - TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE, - ): - consumer_names.append(self.name) - for consumer_name in consumer_names: + for consumer_name in self._topic_consumers(topic_name): total += await topic.unconsumed_count(consumer_name) return total + async def _progress_possible(self) -> bool: + """Whether the parallel run can still make progress. + + True while any node is actively processing or any event is still awaiting + consumption. A node can only be invoked when one of its subscribed topics + has an unconsumed event, so a zero pending-consumable count already + implies no node can invoke -- no separate can_invoke() sweep is needed. + + When this is False but the tracker is non-quiescent, the outstanding + deliveries are parked -- e.g. a node with an ``A AND B`` subscription that + consumed ``A`` but whose ``B`` will never arrive -- so the output queue + ends iteration instead of hanging on a commit that will never come. + (Per-consumer delivery accounting prevents premature termination; this + only catches the genuinely stuck case.) + """ + if not await self._tracker.is_idle(): + return True + return await self._count_pending_consumable() > 0 + async def _add_to_invoke_queue(self, event: TopicEvent) -> None: topic_name = event.name @@ -313,7 +344,12 @@ async def invoke_sequential( invoke_context, node_consumed_events ): published_events.extend( - await publish_events(node, result, self._tracker) + await publish_events( + node, + result, + self._tracker, + self._topic_consumers, + ) ) for event in published_events: @@ -369,7 +405,12 @@ async def invoke_parallel( ] # Create AsyncOutputQueue with output topics and tracker - output_queue = AsyncOutputQueue(output_topics, self.name, self._tracker) + output_queue = AsyncOutputQueue( + output_topics, + self.name, + self._tracker, + progress_possible=self._progress_possible, + ) await output_queue.start_listeners() consumed_output_events: List[ConsumeFromTopicEvent] = [] @@ -555,6 +596,7 @@ def _cancel_all_active_tasks() -> None: node=node, publish_event=event, tracker=self._tracker, + consumers_of=self._topic_consumers, ) ) @@ -662,6 +704,11 @@ async def init_workflow( ] events_to_record: List[Event] = [] + # One delivery per consumer of each seeded input topic (not one per + # topic), so a fan-out input feeding several nodes is not declared + # quiescent after only the first node commits. Accumulated here while + # ``event`` is narrowed to PublishToTopicEvent. + seeded_delivery_count = 0 for input_topic in input_topics: event = await input_topic.publish_data( input_data.model_copy( @@ -674,6 +721,7 @@ async def init_workflow( ) if event: events_to_record.append(event) + seeded_delivery_count += len(self._topic_consumers(event.name)) if is_sequential: await self._add_to_invoke_queue(event) @@ -681,16 +729,13 @@ async def init_workflow( f"init_workflow: events_to_record={len(events_to_record)}, input_topics={len(input_topics)}" ) if events_to_record: - # Track initial input messages for quiescence detection - if not is_sequential: + # Track initial input messages for quiescence detection. + if not is_sequential and seeded_delivery_count: logger.debug( - f"init_workflow: calling on_messages_published({len(events_to_record)})" + f"init_workflow: on_messages_published({seeded_delivery_count})" ) await self._tracker.on_messages_published( - len(events_to_record), source="init_workflow" - ) - logger.debug( - f"init_workflow: tracker after publish: {await self._tracker.get_metrics()}" + seeded_delivery_count, source="init_workflow" ) await container.event_store.record_events(events_to_record) else: @@ -756,11 +801,19 @@ async def init_workflow( ) ) if paired_event: - # Track the published message for quiescence detection + # Track one delivery per consumer of the paired + # input topic for quiescence detection. if not is_sequential: - await self._tracker.on_messages_published( - 1, source="restore_paired_input" + delivery_count = len( + self._topic_consumers( + paired_in_workflow_input_topic_name + ) ) + if delivery_count: + await self._tracker.on_messages_published( + delivery_count, + source="restore_paired_input", + ) if is_sequential: await self._add_to_invoke_queue(paired_event) await container.event_store.record_event(paired_event) diff --git a/grafi/workflows/impl/utils.py b/grafi/workflows/impl/utils.py index 93cbe31b..7a6a11f4 100644 --- a/grafi/workflows/impl/utils.py +++ b/grafi/workflows/impl/utils.py @@ -1,3 +1,4 @@ +from typing import Callable from typing import Dict from typing import List from typing import Sequence @@ -47,11 +48,13 @@ def get_async_output_events(events: Sequence[TopicEvent]) -> List[TopicEvent]: if streaming_events: base_event = streaming_events[0] - aggregated_content_parts = [] + aggregated_content_parts: List[str] = [] for event in streaming_events: messages = event.data if isinstance(event.data, list) else [event.data] for message in messages: - if message.content: + # Streaming text chunks carry str content; ignore non-str + # content (which ``str.join`` could not concatenate anyway). + if isinstance(message.content, str) and message.content: aggregated_content_parts.append(message.content) aggregated_content = "".join(aggregated_content_parts) @@ -84,25 +87,36 @@ async def publish_events( node: NodeBase, publish_event: PublishToTopicEvent, tracker: AsyncNodeTracker, + consumers_of: Callable[[str], Sequence[str]], ) -> List[PublishToTopicEvent]: """ Publish events to all topics the node publishes to. - CHANGE: Added optional tracker parameter. - When provided, notifies tracker of published messages. + The tracker counts outstanding *deliveries* -- one per consumer of each + published event -- not publications. ``consumers_of(topic_name)`` returns the + consumers the workflow topology will route a topic's events to (subscribing + nodes, plus the workflow itself for output topics). Each delivery is + registered *before* the event is enqueued, so a consumer can never commit a + delivery the tracker has not yet counted (which would strand the count and + prevent quiescence). """ published_events: List[PublishToTopicEvent] = [] for topic in node.publish_to: + delivery_count = len(consumers_of(topic.name)) + if delivery_count: + await tracker.on_messages_published( + delivery_count, source=f"node:{node.name}->{topic.name}" + ) event = await topic.publish_data(publish_event) if event: published_events.append(event) - - # NEW: Notify tracker of published messages - if tracker and published_events: - await tracker.on_messages_published( - len(published_events), source=f"node:{node.name}" - ) + elif delivery_count: + # Condition not met: the event was never enqueued, so no consumer + # will ever commit these deliveries. Release them immediately. + await tracker.on_messages_committed( + delivery_count, source=f"discard:{topic.name}" + ) return published_events diff --git a/pyproject.toml b/pyproject.toml index 08b7754a..6fc91eaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,6 @@ dependencies = [ "opentelemetry-exporter-otlp-proto-grpc>=1.42.1", "loguru>=0.7.3", "anyio>=4.14.0", - "cloudpickle>=3.1.2", ] [dependency-groups] @@ -88,14 +87,6 @@ ignore_missing_imports = true module = "googlesearch.*" ignore_missing_imports = true -[[tool.mypy.overrides]] -module = "jsonpickle.*" -ignore_missing_imports = true - -[[tool.mypy.overrides]] -module = "cloudpickle.*" -ignore_missing_imports = true - [[tool.mypy.overrides]] module = "tests.*" ignore_errors = true diff --git a/tests/assistants/test_assistant.py b/tests/assistants/test_assistant.py index e181c9a9..c1887ef7 100644 --- a/tests/assistants/test_assistant.py +++ b/tests/assistants/test_assistant.py @@ -244,11 +244,9 @@ def test_generate_manifest_success(self, mock_assistant, tmp_path): result_path = mock_assistant.generate_manifest(output_dir) - # Verify the method returns correct path + # Verify the method returns the path to the generated manifest. expected_path = os.path.join(output_dir, "test_assistant_manifest.json") - assert ( - result_path is None - ) # Method doesn't return path, but we can check file exists + assert result_path == expected_path # Verify file was created and contains correct data assert os.path.exists(expected_path) diff --git a/tests/common/decorators/test_record_decorators.py b/tests/common/decorators/test_record_decorators.py index 09294bf2..0fca8950 100644 --- a/tests/common/decorators/test_record_decorators.py +++ b/tests/common/decorators/test_record_decorators.py @@ -187,6 +187,92 @@ async def test_failed_event_carries_error_details(self, isolated_container): assert failed[0].error +class _RecordingSpan: + """Fake span that remembers whether it had already ended when each + operation was performed, so a test can prove enrichment happened while the + span was still recording.""" + + def __init__(self) -> None: + self.ended = False + self.attr_set_after_end: dict = {} # attribute key -> ended-state at set + self.exception_recorded_after_end: list = [] + self.status_set_after_end: list = [] + + def set_attribute(self, key, value): + self.attr_set_after_end[key] = self.ended + + def set_attributes(self, mapping): + for key in mapping: + self.attr_set_after_end[key] = self.ended + + def record_exception(self, exc): + self.exception_recorded_after_end.append(self.ended) + + def set_status(self, status): + self.status_set_after_end.append(self.ended) + + +class _RecordingSpanCM: + def __init__(self, span: _RecordingSpan) -> None: + self.span = span + + def __enter__(self) -> _RecordingSpan: + return self.span + + def __exit__(self, exc_type, exc, tb) -> bool: + # Ending the span (as OTel's context manager does) must happen AFTER + # error enrichment, never before. + self.span.ended = True + return False # do not suppress the exception + + +class _RecordingTracer: + def __init__(self) -> None: + self.span = _RecordingSpan() + + def start_as_current_span(self, _name): + return _RecordingSpanCM(self.span) + + +class TestSpanErrorEnrichment: + """Defect #4: failed-span error attributes must be set before span end.""" + + @pytest.mark.asyncio + async def test_error_attributes_recorded_before_span_ends(self): + prev_store = container._event_store + prev_tracer = container._tracer + tracer = _RecordingTracer() + container.register_event_store(EventStoreInMemory()) + container.register_tracer(tracer) + try: + tool = _ExplodingTool() + invoke_context = InvokeContext( + conversation_id="c", + invoke_id="i", + assistant_request_id="r", + ) + with pytest.raises(FunctionCallException): + async for _ in tool.invoke( + invoke_context, [Message(role="user", content="hi")] + ): + pass + finally: + container._event_store = prev_store + container._tracer = prev_tracer + + span = tracer.span + # The span was ended by the context manager exit. + assert span.ended is True + # Every error attribute was set while the span was still recording. + assert span.attr_set_after_end.get("error.type") is False + assert span.attr_set_after_end.get("error.message") is False + assert span.attr_set_after_end.get("error.stack") is False + assert span.attr_set_after_end.get("error.details") is False + # The exception and error status were also recorded before span end. + assert span.exception_recorded_after_end == [False] + assert span.status_set_after_end == [False] + + class TestTracebackDedup: """The full traceback is logged once even after the error is re-wrapped.""" diff --git a/tests/common/events/test_event_codec.py b/tests/common/events/test_event_codec.py new file mode 100644 index 00000000..5b4a09c8 --- /dev/null +++ b/tests/common/events/test_event_codec.py @@ -0,0 +1,67 @@ +"""Tests for EventCodec (Phase 6): registry-based event decoding.""" + +from datetime import datetime + +import pytest + +from grafi.common.events.event import Event +from grafi.common.events.event import EventType +from grafi.common.events.event_codec import EventCodec +from grafi.common.events.event_codec import default_event_codec +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message + + +def test_default_codec_knows_builtin_event_types(): + assert ( + default_event_codec.event_class(EventType.PUBLISH_TO_TOPIC.value) + is PublishToTopicEvent + ) + + +def test_decode_missing_event_type_returns_none(): + assert default_event_codec.decode({"event_id": "x"}) is None + + +def test_decode_unknown_event_type_returns_none(): + assert default_event_codec.decode({"event_id": "x", "event_type": "NOPE"}) is None + + +def test_decode_roundtrips_a_known_event(): + event = PublishToTopicEvent( + event_id="evt-1", + name="t", + offset=0, + publisher_name="p", + publisher_type="t", + consumed_event_ids=[], + invoke_context=InvokeContext( + conversation_id="c", invoke_id="i", assistant_request_id="r" + ), + data=[Message(role="user", content="hi")], + timestamp=datetime(2026, 1, 1), + ) + decoded = default_event_codec.decode(event.to_dict()) + assert isinstance(decoded, PublishToTopicEvent) + assert decoded.event_id == "evt-1" + + +def test_register_adds_a_new_event_type_without_touching_the_store(): + codec = EventCodec() + assert codec.event_class(EventType.TOPIC_EVENT.value) is None + + class _Custom(Event): + pass + + codec.register("custom_type", _Custom) + assert codec.event_class("custom_type") is _Custom + + +@pytest.mark.asyncio +async def test_store_delegates_decoding_to_codec(): + """EventStore._create_event_from_dict is a thin delegate over the codec.""" + from grafi.common.event_stores.event_store_in_memory import EventStoreInMemory + + store = EventStoreInMemory() + assert store._create_event_from_dict({"event_id": "x"}) is None diff --git a/tests/tools/llms/test_claude_tool.py b/tests/tools/llms/test_claude_tool.py index db5d9d17..557531fe 100644 --- a/tests/tools/llms/test_claude_tool.py +++ b/tests/tools/llms/test_claude_tool.py @@ -4,7 +4,7 @@ from unittest.mock import Mock import pytest -from anthropic import NOT_GIVEN +from anthropic import omit from anthropic.types.text_block import TextBlock from grafi.common.models.function_spec import FunctionSpec @@ -101,7 +101,7 @@ async def mock_aexit(self, *args): assert kwargs["messages"][0]["role"] == "user" assert kwargs["messages"][0]["content"] == "Say hello" # no tools in this call - assert kwargs["tools"] == NOT_GIVEN + assert kwargs["tools"] is omit # --------------------------------------------------------------------------- # @@ -206,7 +206,7 @@ def test_prepare_api_input(claude_instance): assert api_messages[0]["role"] == "user" assert api_messages[0]["content"] == "Hello!" assert api_messages[-1]["role"] == "assistant" - assert api_tools == NOT_GIVEN + assert api_tools is omit def test_prepare_api_input_tool_call_linkage(claude_instance): diff --git a/tests/tools/llms/test_deepseek_tool.py b/tests/tools/llms/test_deepseek_tool.py index 77619149..dd26a5a7 100644 --- a/tests/tools/llms/test_deepseek_tool.py +++ b/tests/tools/llms/test_deepseek_tool.py @@ -3,6 +3,7 @@ from unittest.mock import Mock import pytest +from openai import omit from openai.types.chat import ChatCompletion from openai.types.chat import ChatCompletionMessage @@ -58,8 +59,6 @@ def test_init(deepseek_instance): async def test_invoke_simple_response(monkeypatch, deepseek_instance, invoke_context): from unittest.mock import AsyncMock - import grafi.tools.llms.impl.deepseek_tool - mock_response = Mock(spec=ChatCompletion) mock_response.choices = [ Mock(message=ChatCompletionMessage(role="assistant", content="Hello, world!")) @@ -67,11 +66,13 @@ async def test_invoke_simple_response(monkeypatch, deepseek_instance, invoke_con mock_client = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) # Patch the OpenAI.Client constructor the tool uses mock_openai_cls = MagicMock(return_value=mock_client) monkeypatch.setattr( - grafi.tools.llms.impl.deepseek_tool, "AsyncClient", mock_openai_cls + "grafi.tools.llms.impl.openai_compatible.AsyncClient", mock_openai_cls ) input_data = [Message(role="user", content="Say hello")] @@ -102,7 +103,7 @@ async def test_invoke_simple_response(monkeypatch, deepseek_instance, invoke_con "tool_call_id": None, }, ] - assert call_args["tools"] is None + assert call_args["tools"] is omit # --------------------------------------------------------------------------- # @@ -112,8 +113,6 @@ async def test_invoke_simple_response(monkeypatch, deepseek_instance, invoke_con async def test_invoke_function_call(monkeypatch, deepseek_instance, invoke_context): from unittest.mock import AsyncMock - import grafi.tools.llms.impl.deepseek_tool - mock_response = Mock(spec=ChatCompletion) mock_response.choices = [ Mock( @@ -136,9 +135,10 @@ async def test_invoke_function_call(monkeypatch, deepseek_instance, invoke_conte mock_client = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) monkeypatch.setattr( - grafi.tools.llms.impl.deepseek_tool, - "AsyncClient", + "grafi.tools.llms.impl.openai_compatible.AsyncClient", MagicMock(return_value=mock_client), ) @@ -188,13 +188,12 @@ async def test_invoke_function_call(monkeypatch, deepseek_instance, invoke_conte # --------------------------------------------------------------------------- # @pytest.mark.asyncio async def test_invoke_api_error(monkeypatch, deepseek_instance, invoke_context): - import grafi.tools.llms.impl.deepseek_tool # Force constructor to raise – simulates any client error def _raise(*_a, **_kw): # noqa: D401 raise Exception("Error code") - monkeypatch.setattr(grafi.tools.llms.impl.deepseek_tool, "AsyncClient", _raise) + monkeypatch.setattr("grafi.tools.llms.impl.openai_compatible.AsyncClient", _raise) from grafi.common.exceptions import LLMToolException diff --git a/tests/tools/llms/test_openai_tool.py b/tests/tools/llms/test_openai_tool.py index 8e3f2052..c1754528 100644 --- a/tests/tools/llms/test_openai_tool.py +++ b/tests/tools/llms/test_openai_tool.py @@ -12,6 +12,7 @@ from grafi.common.models.function_spec import ParametersSchema from grafi.common.models.invoke_context import InvokeContext from grafi.common.models.message import Message +from grafi.tools.llms.impl.openai_adapter import to_openai_tool from grafi.tools.llms.impl.openai_tool import OpenAITool @@ -47,7 +48,6 @@ def test_init(openai_instance): @pytest.mark.asyncio async def test_invoke_simple_response(monkeypatch, openai_instance, invoke_context): - import grafi.tools.llms.impl.openai_tool mock_response = Mock(spec=ChatCompletion) mock_response.choices = [ @@ -75,7 +75,7 @@ async def mock_aexit(self, *args): # Mock the AsyncClient constructor to return our context manager mock_async_client_cls = MagicMock(return_value=mock_context_manager) monkeypatch.setattr( - grafi.tools.llms.impl.openai_tool, "AsyncClient", mock_async_client_cls + "grafi.tools.llms.impl.openai_compatible.AsyncClient", mock_async_client_cls ) input_data = [Message(role="user", content="Say hello")] @@ -97,7 +97,6 @@ async def mock_aexit(self, *args): @pytest.mark.asyncio async def test_invoke_function_call(monkeypatch, openai_instance, invoke_context): - import grafi.tools.llms.impl.openai_tool mock_response = Mock(spec=ChatCompletion) mock_response.choices = [ @@ -140,7 +139,7 @@ async def mock_aexit(self, *args): # Mock the AsyncClient constructor to return our context manager mock_async_client_cls = MagicMock(return_value=mock_context_manager) monkeypatch.setattr( - grafi.tools.llms.impl.openai_tool, "AsyncClient", mock_async_client_cls + "grafi.tools.llms.impl.openai_compatible.AsyncClient", mock_async_client_cls ) input_data = [Message(role="user", content="What's the weather in London?")] @@ -244,14 +243,16 @@ def test_prepare_api_input(openai_instance): role="user", content="What's the weather like?", tools=[ - FunctionSpec( - name="get_weather", - description="Get weather", - parameters=ParametersSchema( - type="object", - properties={"location": ParameterSchema(type="string")}, - ), - ).to_openai_tool() + to_openai_tool( + FunctionSpec( + name="get_weather", + description="Get weather", + parameters=ParametersSchema( + type="object", + properties={"location": ParameterSchema(type="string")}, + ), + ) + ) ], ), ] diff --git a/tests/tools/llms/test_openrouter_tool.py b/tests/tools/llms/test_openrouter_tool.py index 86793737..85daf582 100644 --- a/tests/tools/llms/test_openrouter_tool.py +++ b/tests/tools/llms/test_openrouter_tool.py @@ -4,6 +4,7 @@ from unittest.mock import Mock import pytest +from openai import omit from openai.types.chat import ChatCompletion from openai.types.chat import ChatCompletionMessage @@ -53,9 +54,6 @@ def test_init(openrouter_instance): # --------------------------------------------------------------------------- # @pytest.mark.asyncio async def test_invoke_simple_response(monkeypatch, openrouter_instance, invoke_context): - import grafi.tools.llms.impl.openrouter_tool - import grafi.tools.llms.impl.openrouter_tool as or_module - # Fake successful response mock_response = Mock(spec=ChatCompletion) mock_response.choices = [ @@ -64,12 +62,13 @@ async def test_invoke_simple_response(monkeypatch, openrouter_instance, invoke_c mock_client = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) - # Patch the OpenAI class used inside the tool + # Patch the OpenAI class used inside the shared base tool + mock_async_client_cls = MagicMock(return_value=mock_client) monkeypatch.setattr( - grafi.tools.llms.impl.openrouter_tool, - "AsyncClient", - MagicMock(return_value=mock_client), + "grafi.tools.llms.impl.openai_compatible.AsyncClient", mock_async_client_cls ) result_messages = [] @@ -84,7 +83,7 @@ async def test_invoke_simple_response(monkeypatch, openrouter_instance, invoke_c assert result_messages[0].content == "Hello, world!" # AsyncClient ctor must receive correct kwargs - or_module.AsyncClient.assert_called_once_with( + mock_async_client_cls.assert_called_once_with( api_key="test_api_key", base_url="https://openrouter.ai/api/v1" ) @@ -103,7 +102,6 @@ async def test_invoke_simple_response(monkeypatch, openrouter_instance, invoke_c async def test_invoke_with_extra_headers( monkeypatch, openrouter_instance, invoke_context ): - import grafi.tools.llms.impl.openrouter_tool openrouter_instance.extra_headers = { "HTTP-Referer": "https://my-app.example", @@ -117,9 +115,10 @@ async def test_invoke_with_extra_headers( mock_client = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) monkeypatch.setattr( - grafi.tools.llms.impl.openrouter_tool, - "AsyncClient", + "grafi.tools.llms.impl.openai_compatible.AsyncClient", MagicMock(return_value=mock_client), ) @@ -139,7 +138,6 @@ async def test_invoke_with_extra_headers( # --------------------------------------------------------------------------- # @pytest.mark.asyncio async def test_invoke_function_call(monkeypatch, openrouter_instance, invoke_context): - import grafi.tools.llms.impl.openrouter_tool mock_response = Mock(spec=ChatCompletion) mock_response.choices = [ @@ -163,9 +161,10 @@ async def test_invoke_function_call(monkeypatch, openrouter_instance, invoke_con mock_client = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) monkeypatch.setattr( - grafi.tools.llms.impl.openrouter_tool, - "AsyncClient", + "grafi.tools.llms.impl.openai_compatible.AsyncClient", MagicMock(return_value=mock_client), ) @@ -196,12 +195,11 @@ async def test_invoke_function_call(monkeypatch, openrouter_instance, invoke_con # --------------------------------------------------------------------------- # @pytest.mark.asyncio async def test_invoke_api_error(monkeypatch, openrouter_instance, invoke_context): - import grafi.tools.llms.impl.openrouter_tool def _raise(*_a, **_kw): # pragma: no cover raise Exception("Error code") - monkeypatch.setattr(grafi.tools.llms.impl.openrouter_tool, "AsyncClient", _raise) + monkeypatch.setattr("grafi.tools.llms.impl.openai_compatible.AsyncClient", _raise) from grafi.common.exceptions import LLMToolException @@ -223,7 +221,7 @@ def test_prepare_api_input(openrouter_instance): ] api_messages, api_tools = openrouter_instance.prepare_api_input(input_data) - assert api_tools is None + assert api_tools is omit assert api_messages[0]["content"] == "dummy system message" assert api_messages[-1]["role"] == "assistant" assert api_messages[-1]["content"] == "Hi there." diff --git a/tests/topics/test_topic_base.py b/tests/topics/test_topic_base.py index 64811f24..3c8da1f5 100644 --- a/tests/topics/test_topic_base.py +++ b/tests/topics/test_topic_base.py @@ -3,6 +3,9 @@ import pytest from grafi.common.callable_ref import CallableSerializationError +from grafi.common.events.topic_events.consume_from_topic_event import ( + ConsumeFromTopicEvent, +) from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent from grafi.common.models.invoke_context import InvokeContext from grafi.common.models.message import Message @@ -72,6 +75,45 @@ async def test_restore_topic(topic: TopicBase, invoke_context: InvokeContext): assert consumed_events[0].event_id == "event_1" +@pytest.mark.asyncio +async def test_restore_topic_consume_event_does_not_skip_next( + topic: TopicBase, invoke_context: InvokeContext +): + """Restoring a ConsumeFromTopicEvent advances the consumer's cursor to the + recorded commit point only, leaving the immediately following offset + available. Regression for the recovery off-by-one that skipped one event.""" + for i in range(3): + await topic.restore_topic( + PublishToTopicEvent( + event_id=f"event_{i}", + name="test_topic", + offset=i, + publisher_name="publisher1", + publisher_type="test", + invoke_context=invoke_context, + data=[Message(role="assistant", content=f"m{i}")], + timestamp=datetime(2023, 1, 1, 13, 0), + ) + ) + + # The consumer had committed through offset 0 before the crash. + await topic.restore_topic( + ConsumeFromTopicEvent( + event_id="consume_0", + name="test_topic", + offset=0, + consumer_name="c", + consumer_type="Node", + invoke_context=invoke_context, + data=[Message(role="assistant", content="m0")], + ) + ) + + # Offsets 1 and 2 remain available (offset 1 is NOT skipped). + remaining = await topic.consume("c") + assert [event.offset for event in remaining] == [1, 2] + + # Module-level condition used to exercise reference-based serialization. def module_level_condition(event: PublishToTopicEvent) -> bool: """A named, importable condition -> serializes as an import reference.""" diff --git a/tests/topics/test_topic_event_cache.py b/tests/topics/test_topic_event_cache.py index 9a614dab..cd7c2e10 100644 --- a/tests/topics/test_topic_event_cache.py +++ b/tests/topics/test_topic_event_cache.py @@ -212,6 +212,87 @@ async def test_commit_to(self, cache: TopicEventQueue): await cache.commit_to("consumer1", 5) assert cache._committed["consumer1"] == 5 + async def _put_n(self, cache: TopicEventQueue, n: int) -> list: + """Append ``n`` events (offsets 0..n-1) and return them.""" + events = [] + for i in range(n): + invoke_context = InvokeContext( + conversation_id="test-conversation", + invoke_id=f"test-invoke-{i}", + assistant_request_id="test-request", + ) + event = PublishToTopicEvent( + event_id=f"event-{i}", + name="test_topic", + offset=i, + publisher_name="test_publisher", + publisher_type="test_type", + consumed_event_ids=[], + invoke_context=invoke_context, + data=[Message(role="user", content=f"message {i}")], + timestamp=datetime.now(), + ) + events.append(event) + await cache.put(event) + return events + + @pytest.mark.asyncio + async def test_restore_consumer_does_not_skip_next_offset( + self, cache: TopicEventQueue + ): + """Regression for the recovery off-by-one: restoring a consume at offset + 0 must leave offsets 1..N available, not advance past offset 1.""" + events = await self._put_n(cache, 3) + + await cache.restore_consumer("consumer1", committed_offset=0) + + # Cursor sits exactly one past the committed offset. + assert cache._committed["consumer1"] == 0 + assert cache._consumed["consumer1"] == 1 + + # The very next offset (1) is still delivered, not skipped. + result = await cache.fetch("consumer1", timeout=0.1) + assert result == events[1:] + assert [e.offset for e in result] == [1, 2] + + @pytest.mark.asyncio + async def test_restore_consumer_is_idempotent_and_never_rewinds( + self, cache: TopicEventQueue + ): + """Replaying the same or an earlier commit only advances cursors.""" + await self._put_n(cache, 3) + + await cache.restore_consumer("consumer1", committed_offset=1) + assert cache._consumed["consumer1"] == 2 + assert cache._committed["consumer1"] == 1 + + # Same offset again: no change. + await cache.restore_consumer("consumer1", committed_offset=1) + assert cache._consumed["consumer1"] == 2 + assert cache._committed["consumer1"] == 1 + + # An earlier offset (out-of-order replay) must not rewind the cursor. + await cache.restore_consumer("consumer1", committed_offset=0) + assert cache._consumed["consumer1"] == 2 + assert cache._committed["consumer1"] == 1 + + @pytest.mark.asyncio + async def test_restore_consumer_is_independent_per_consumer( + self, cache: TopicEventQueue + ): + """Restoring one consumer's cursor does not affect another's.""" + events = await self._put_n(cache, 3) + + await cache.restore_consumer("consumer1", committed_offset=1) + + # consumer2 was never restored: it still sees every event. + result2 = await cache.fetch("consumer2", timeout=0.1) + assert result2 == events + + # consumer1 resumes from offset 2 only. + result1 = await cache.fetch("consumer1", timeout=0.1) + assert [e.offset for e in result1] == [2] + @pytest.mark.asyncio async def test_concurrent_producers(self, cache: TopicEventQueue): # Multiple producers adding events concurrently diff --git a/tests/workflow/test_fanout_quiescence.py b/tests/workflow/test_fanout_quiescence.py new file mode 100644 index 00000000..9e4bd698 --- /dev/null +++ b/tests/workflow/test_fanout_quiescence.py @@ -0,0 +1,164 @@ +"""Characterization tests for fan-out delivery accounting (Defect #2). + +The quiescence tracker counts outstanding *deliveries* (one per consumer of each +published event), not publications. Before the fix it counted publications, so a +single event fanned out to several subscribers was declared quiescent after only +the first subscriber committed -- silently dropping the rest of the work. +""" + +from unittest.mock import Mock +from unittest.mock import patch + +import pytest +from openinference.semconv.trace import OpenInferenceSpanKindValues + +from grafi.common.event_stores.event_store_in_memory import EventStoreInMemory +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.nodes.node import Node +from grafi.tools.tool import Tool +from grafi.topics.expressions.topic_expression import TopicExpr +from grafi.topics.topic_impl.input_topic import InputTopic +from grafi.topics.topic_impl.output_topic import OutputTopic +from grafi.topics.topic_types import TopicType +from grafi.workflows.impl.event_driven_workflow import EventDrivenWorkflow + + +class LabelTool(Tool): + """Echoes a per-node label so each subscriber's output is distinguishable.""" + + oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.TOOL + label: str = "?" + + async def invoke(self, invoke_context, input_data): + yield [Message(role="assistant", content=f"resp-{self.label}")] + + @classmethod + async def from_dict(cls, data): # pragma: no cover - not exercised here + return cls(oi_span_type=OpenInferenceSpanKindValues.TOOL) + + +def _input_event(request_id: str) -> PublishToTopicEvent: + return PublishToTopicEvent( + name="agent_input", + type=TopicType.AGENT_INPUT_TOPIC_TYPE, + publisher_name="EventDrivenWorkflow", + publisher_type="EventDrivenWorkflow", + offset=0, + invoke_context=InvokeContext( + conversation_id="c", + invoke_id="i", + assistant_request_id=request_id, + ), + data=[Message(role="user", content="hello")], + ) + + +def _build_fanout_workflow() -> EventDrivenWorkflow: + """One input topic feeding two independent subscribers, both publishing to + the single output topic.""" + input_topic = InputTopic(name="agent_input") + output_topic = OutputTopic(name="agent_output") + node_a = Node( + name="node_a", + type="Node", + tool=LabelTool(label="a"), + subscribed_expressions=[TopicExpr(topic=input_topic)], + publish_to=[output_topic], + ) + node_b = Node( + name="node_b", + type="Node", + tool=LabelTool(label="b"), + subscribed_expressions=[TopicExpr(topic=input_topic)], + publish_to=[output_topic], + ) + return EventDrivenWorkflow.builder().node(node_a).node(node_b).build() + + +def test_topic_consumers_counts_every_subscriber(): + """The topology helper reports both subscribers of the fan-out input topic + and the workflow itself as the consumer of the output topic.""" + workflow = _build_fanout_workflow() + + assert set(workflow._topic_consumers("agent_input")) == {"node_a", "node_b"} + assert workflow._topic_consumers("agent_output") == [workflow.name] + # Unknown topic -> no consumers (cannot hang execution). + assert workflow._topic_consumers("does_not_exist") == [] + + +@pytest.mark.asyncio +async def test_parallel_fanout_yields_every_subscriber_output(): + """Both subscribers must run: the workflow must not quiesce after only the + first commits. Regression for the fan-out under-count.""" + workflow = _build_fanout_workflow() + request_id = "fanout-parallel" + + fake_container = Mock() + fake_container.event_store = EventStoreInMemory() + + with patch("grafi.workflows.impl.event_driven_workflow.container", fake_container): + contents = [ + msg.content + async for event in workflow.invoke( + _input_event(request_id), is_sequential=False + ) + for msg in event.data + ] + + assert "resp-a" in contents + assert "resp-b" in contents + + +@pytest.mark.asyncio +async def test_parallel_does_not_hang_on_unsatisfied_and_subscription(): + """A topic feeding a firing node AND a node whose AND-subscription can never + be satisfied must still terminate (the parked delivery must not hang the + run). Regression for the per-consumer accounting over-count.""" + import asyncio + + from grafi.topics.expressions.subscription_builder import SubscriptionBuilder + from grafi.topics.topic_impl.topic import Topic + + input_topic = InputTopic(name="agent_input") + gate = Topic(name="gate_topic") # nothing ever publishes here + output_topic = OutputTopic(name="agent_output") + + firing = Node( + name="firing", + type="Node", + tool=LabelTool(label="ok"), + subscribed_expressions=[TopicExpr(topic=input_topic)], + publish_to=[output_topic], + ) + # Subscribes to (agent_input AND gate_topic); gate_topic never fires. + stuck = Node( + name="stuck", + type="Node", + tool=LabelTool(label="never"), + subscribed_expressions=[ + SubscriptionBuilder() + .subscribed_to(input_topic) + .and_() + .subscribed_to(gate) + .build() + ], + publish_to=[output_topic], + ) + workflow = EventDrivenWorkflow.builder().node(firing).node(stuck).build() + + fake_container = Mock() + fake_container.event_store = EventStoreInMemory() + + with patch("grafi.workflows.impl.event_driven_workflow.container", fake_container): + contents = [] + async with asyncio.timeout(10): # fails loudly if it hangs + async for event in workflow.invoke( + _input_event("and-sub"), is_sequential=False + ): + contents.extend(msg.content for msg in event.data) + + # The firing node's output is produced; the stuck AND-node never fires. + assert "resp-ok" in contents + assert "resp-never" not in contents diff --git a/tests/workflow/test_utils.py b/tests/workflow/test_utils.py index 681bd2e4..bb8a4565 100644 --- a/tests/workflow/test_utils.py +++ b/tests/workflow/test_utils.py @@ -1,5 +1,6 @@ from unittest.mock import AsyncMock from unittest.mock import MagicMock +from unittest.mock import call import pytest @@ -234,9 +235,17 @@ async def test_publish_events(self): ) mock_event2 = None # Test case where topic doesn't publish + mock_topic1.name = "topic1" + mock_topic2.name = "topic2" mock_topic1.publish_data.return_value = mock_event1 mock_topic2.publish_data.return_value = mock_event2 + # topic1 fans out to two consumers; topic2 to one. + consumers = {"topic1": ["A", "B"], "topic2": ["C"]} + + def consumers_of(topic_name: str): + return consumers.get(topic_name, []) + publish_to_event = PublishToTopicEvent( invoke_context=invoke_context, publisher_name=node.name, @@ -245,15 +254,24 @@ async def test_publish_events(self): consumed_event_ids=[event.event_id for event in consumed_events], ) - published_events = await publish_events(node, publish_to_event, tracker) + published_events = await publish_events( + node, publish_to_event, tracker, consumers_of + ) assert len(published_events) == 1 assert published_events[0] == mock_event1 # Verify topics were called correctly mock_topic1.publish_data.assert_called_once_with(publish_to_event) - tracker.on_messages_published.assert_awaited_once_with( - 1, source="node:test_node" + + # One delivery registered per consumer, per topic, before publishing. + assert tracker.on_messages_published.await_args_list == [ + call(2, source="node:test_node->topic1"), + call(1, source="node:test_node->topic2"), + ] + # topic2's condition filtered the event out, so its delivery is released. + tracker.on_messages_committed.assert_awaited_once_with( + 1, source="discard:topic2" ) diff --git a/uv.lock b/uv.lock index 114bd8ef..2ef548d7 100644 --- a/uv.lock +++ b/uv.lock @@ -763,15 +763,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, ] -[[package]] -name = "cloudpickle" -version = "3.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, -] - [[package]] name = "colorama" version = "0.4.6" @@ -1243,7 +1234,6 @@ version = "0.0.34" source = { editable = "." } dependencies = [ { name = "anyio" }, - { name = "cloudpickle" }, { name = "docstring-parser" }, { name = "loguru" }, { name = "openai" }, @@ -1287,7 +1277,6 @@ docs = [ [package.metadata] requires-dist = [ { name = "anyio", specifier = ">=4.14.0" }, - { name = "cloudpickle", specifier = ">=3.1.2" }, { name = "docstring-parser", specifier = ">=0.18.0" }, { name = "loguru", specifier = ">=0.7.3" }, { name = "openai", specifier = ">=2.42.0" },