From a406933990ea0e48504074a0339879f60610ebc3 Mon Sep 17 00:00:00 2001 From: GuanyiLi-Craig Date: Thu, 1 Jan 2026 16:14:14 +0000 Subject: [PATCH 1/4] improve code --- grafi/common/containers/container.py | 11 ++- .../event_stores/event_store_postgres.py | 77 ++++++++++++++++--- grafi/nodes/node_base.py | 6 +- .../function_calls/function_call_tool.py | 28 +++++-- grafi/topics/topic_base.py | 11 +++ grafi/workflows/workflow.py | 3 +- 6 files changed, 112 insertions(+), 24 deletions(-) diff --git a/grafi/common/containers/container.py b/grafi/common/containers/container.py index 3c703c7..2a4ba18 100644 --- a/grafi/common/containers/container.py +++ b/grafi/common/containers/container.py @@ -1,4 +1,5 @@ # container.py +import os import threading from typing import Any from typing import Optional @@ -68,11 +69,15 @@ def tracer(self) -> Tracer: # Slow path: initialize with lock (double-checked locking) with self._init_lock: if self._tracer is None: + # Use environment variables with sensible defaults + endpoint = os.getenv("OTEL_COLLECTOR_ENDPOINT", "localhost") + port = int(os.getenv("OTEL_COLLECTOR_PORT", "4317")) + project = os.getenv("GRAFI_PROJECT_NAME", "grafi-trace") self._tracer = setup_tracing( tracing_options=TracingOptions.AUTO, - collector_endpoint="localhost", - collector_port=4317, - project_name="grafi-trace", + collector_endpoint=endpoint, + collector_port=port, + project_name=project, ) return self._tracer diff --git a/grafi/common/event_stores/event_store_postgres.py b/grafi/common/event_stores/event_store_postgres.py index f5f6a3d..87985fb 100644 --- a/grafi/common/event_stores/event_store_postgres.py +++ b/grafi/common/event_stores/event_store_postgres.py @@ -6,6 +6,7 @@ from grafi.common.event_stores.event_store import EventStore from grafi.common.events.event import Event +from grafi.common.exceptions import EventPersistenceError try: @@ -14,6 +15,7 @@ from sqlalchemy import Integer from sqlalchemy import String from sqlalchemy import create_engine + from sqlalchemy import delete from sqlalchemy import select from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.ext.asyncio import AsyncSession @@ -55,13 +57,30 @@ class EventModel(Base): class EventStorePostgres(EventStore): """Postgres-backed implementation of the EventStore interface with async support.""" - def __init__(self, db_url: str): + def __init__( + self, + db_url: str, + pool_size: int = 5, + max_overflow: int = 10, + pool_timeout: int = 30, + ): """ Initialize the Postgres event store. - :param db_url: The SQLAlchemy database URL, e.g. 'postgresql://user:pass@host/dbname'. + + Args: + db_url: The SQLAlchemy database URL, e.g. 'postgresql://user:pass@host/dbname'. + pool_size: The number of connections to keep in the pool (default: 5). + max_overflow: Maximum overflow connections beyond pool_size (default: 10). + pool_timeout: Seconds to wait for a connection from the pool (default: 30). """ # Keep sync engine for initialization and sync methods - self.engine = create_engine(db_url, echo=False) + self.engine = create_engine( + db_url, + echo=False, + pool_size=pool_size, + max_overflow=max_overflow, + pool_timeout=pool_timeout, + ) Base.metadata.create_all(self.engine) # Create async engine and session for async methods @@ -69,18 +88,46 @@ def __init__(self, db_url: str): async_db_url = db_url.replace("postgresql://", "postgresql+asyncpg://") if "psycopg2" in async_db_url: async_db_url = async_db_url.replace("psycopg2", "asyncpg") - self.async_engine = create_async_engine(async_db_url, echo=False) + self.async_engine = create_async_engine( + async_db_url, + echo=False, + pool_size=pool_size, + max_overflow=max_overflow, + ) self.AsyncSession = async_sessionmaker( self.async_engine, class_=AsyncSession, expire_on_commit=False ) async def clear_events(self) -> None: """Clear all events from the database.""" - pass + async with self.AsyncSession() as session: + try: + await session.execute(delete(EventModel)) + await session.commit() + except Exception as e: + await session.rollback() + logger.error(f"Failed to clear events: {e}") + raise async def get_events(self) -> List[Event]: """Get all events from the database.""" - pass + async with self.AsyncSession() as session: + try: + result = await session.execute( + select(EventModel).order_by(EventModel.timestamp) + ) + rows = result.scalars().all() + + events: List[Event] = [] + for r in rows: + event = self._create_event_from_dict(r.event_data) + if event: + events.append(event) + + return events + except Exception as e: + logger.error(f"Failed to get events: {e}") + raise async def record_event(self, event: Event) -> None: """Record a single event into the database asynchronously.""" @@ -112,7 +159,10 @@ async def record_event(self, event: Event) -> None: except Exception as e: await session.rollback() logger.error(f"Failed to record event: {e}") - raise e + raise EventPersistenceError( + message=f"Failed to record event: {e}", + cause=e, + ) from e async def record_events(self, events: List[Event]) -> None: """Record multiple events into the database asynchronously.""" @@ -145,7 +195,10 @@ async def record_events(self, events: List[Event]) -> None: except Exception as e: await session.rollback() logger.error(f"Failed to record events: {e}") - raise e + raise EventPersistenceError( + message=f"Failed to record events: {e}", + cause=e, + ) from e async def get_event(self, event_id: str) -> Optional[Event]: """Get an event by ID asynchronously.""" @@ -162,7 +215,7 @@ async def get_event(self, event_id: str) -> Optional[Event]: return self._create_event_from_dict(row.event_data) except Exception as e: logger.error(f"Failed to get event {event_id}: {e}") - raise e + raise async def get_agent_events(self, assistant_request_id: str) -> List[Event]: """Get all events for a given assistant_request_id asynchronously.""" @@ -187,7 +240,7 @@ async def get_agent_events(self, assistant_request_id: str) -> List[Event]: return events except Exception as e: logger.error(f"Failed to get agent events {assistant_request_id}: {e}") - raise e + raise async def get_conversation_events(self, conversation_id: str) -> List[Event]: """Get all events for a given conversation ID asynchronously.""" @@ -214,7 +267,7 @@ async def get_conversation_events(self, conversation_id: str) -> List[Event]: logger.error( f"Failed to get conversation events {conversation_id}: {e}" ) - raise e + raise async def get_topic_events(self, name: str, offsets: List[int]) -> List[Event]: """Get all events for a given topic name and specific offsets asynchronously.""" @@ -260,7 +313,7 @@ async def get_topic_events(self, name: str, offsets: List[int]) -> List[Event]: return events except Exception as e: logger.error(f"Failed to get topic events for {name}: {e}") - raise e + raise async def initialize(self) -> None: """Initialize the database tables asynchronously.""" diff --git a/grafi/nodes/node_base.py b/grafi/nodes/node_base.py index a5ed37f..df228d0 100644 --- a/grafi/nodes/node_base.py +++ b/grafi/nodes/node_base.py @@ -39,10 +39,10 @@ class NodeBase(BaseModel): type: str = Field(default="Node") tool: Optional[Tool] = Field(default=None) oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.CHAIN - subscribed_expressions: List[SubExpr] = Field(default=[]) - publish_to: List[TopicBase] = Field(default=[]) + subscribed_expressions: List[SubExpr] = Field(default_factory=list) + publish_to: List[TopicBase] = Field(default_factory=list) - _subscribed_topics: Dict[str, TopicBase] = PrivateAttr(default={}) + _subscribed_topics: Dict[str, TopicBase] = PrivateAttr(default_factory=dict) _command: Optional[Command] = PrivateAttr(default=None) @property diff --git a/grafi/tools/function_calls/function_call_tool.py b/grafi/tools/function_calls/function_call_tool.py index 0426d30..e9cc39d 100644 --- a/grafi/tools/function_calls/function_call_tool.py +++ b/grafi/tools/function_calls/function_call_tool.py @@ -84,7 +84,8 @@ def __init_subclass__(cls, **kwargs: Any) -> None: function_spec: FunctionSpec = attr._function_spec cls.functions[function_spec.name] = attr cls.function_specs.append(function_spec) - else: + + if not cls.function_specs: logger.warning( f"{cls.__name__}: no method decorated with @llm_function found." ) @@ -117,13 +118,20 @@ async def invoke( Raises: ValueError: If the provided function_name doesn't match the registered function. """ - if len(input_data) > 0 and input_data[0].tool_calls is None: - logger.warning("Function call is None.") - raise ValueError("Function call is None.") + if not input_data: + logger.warning("No input data provided.") + yield [] + return + + tool_calls = input_data[0].tool_calls + if not tool_calls: + logger.warning("No tool calls found in input data.") + yield [] + return messages: Messages = [] - for tool_call in input_data[0].tool_calls if input_data[0].tool_calls else []: + for tool_call in tool_calls: if tool_call.function.name in self.functions: func = self.functions[tool_call.function.name] try: @@ -185,7 +193,17 @@ async def from_dict(cls, data: dict[str, Any]) -> "FunctionCallTool": Note: Functions are reconstructed from cloudpickle serialized data. + + Warning: + SECURITY: This method deserializes pickled code using cloudpickle. + Pickle deserialization can execute arbitrary code. Only use this + method with data from trusted sources. For production use with + external/untrusted data, consider using a safer serialization format. """ + logger.debug( + "Deserializing function call tool from pickle data. " + "Ensure data source is trusted." + ) function_call_tool_builder = ( cls.builder() diff --git a/grafi/topics/topic_base.py b/grafi/topics/topic_base.py index 88b3695..f48000e 100644 --- a/grafi/topics/topic_base.py +++ b/grafi/topics/topic_base.py @@ -181,6 +181,12 @@ async def from_dict(cls, data: dict[str, Any]) -> "TopicBase": Returns: TopicBase: A TopicBase instance created from the dictionary. + + Warning: + SECURITY: This method deserializes pickled code using cloudpickle. + Pickle deserialization can execute arbitrary code. Only use this + method with data from trusted sources. For production use with + external/untrusted data, consider using a safer serialization format. """ condition_data = data["condition"] if isinstance(condition_data, dict): @@ -188,6 +194,11 @@ async def from_dict(cls, data: dict[str, Any]) -> "TopicBase": else: encoded_condition = condition_data + logger.debug( + "Deserializing topic condition from pickle data. " + "Ensure data source is trusted." + ) + return cls( name=data["name"], type=data["type"], diff --git a/grafi/workflows/workflow.py b/grafi/workflows/workflow.py index a2f31d0..351ab26 100644 --- a/grafi/workflows/workflow.py +++ b/grafi/workflows/workflow.py @@ -7,6 +7,7 @@ from loguru import logger from openinference.semconv.trace import OpenInferenceSpanKindValues from pydantic import BaseModel +from pydantic import Field from pydantic import PrivateAttr from grafi.common.events.topic_events.consume_from_topic_event import ( @@ -26,7 +27,7 @@ class Workflow(BaseModel): workflow_id: str = default_id name: str = "Workflow" type: str = "Workflow" - nodes: Dict[str, NodeBase] = {} + nodes: Dict[str, NodeBase] = Field(default_factory=dict) # Stop flag to control workflow execution _stop_requested: bool = PrivateAttr(default=False) From 85ac7e531e47cf2231e63ee44d5b0117fc631b5a Mon Sep 17 00:00:00 2001 From: GuanyiLi-Craig Date: Thu, 1 Jan 2026 16:50:06 +0000 Subject: [PATCH 2/4] improve code --- grafi/assistants/assistant_base.py | 32 ++++++++++ grafi/common/models/base_builder.py | 2 - grafi/nodes/node_base.py | 58 +++++++++++++++++-- .../function_calls/function_call_tool.py | 24 +++++--- .../function_calls/impl/synthetic_tool.py | 48 +++++++++++++++ grafi/tools/functions/function_tool.py | 16 +++++ grafi/tools/llms/impl/claude_tool.py | 2 - grafi/tools/llms/impl/deepseek_tool.py | 2 - grafi/tools/llms/impl/gemini_tool.py | 2 - grafi/tools/llms/impl/openai_tool.py | 10 ++++ grafi/tools/llms/impl/openrouter_tool.py | 2 - grafi/tools/llms/llm.py | 33 +++++++++++ grafi/tools/tool.py | 24 ++++++++ .../expressions/subscription_builder.py | 22 +++++++ .../queue_impl/in_mem_topic_event_queue.py | 4 +- grafi/topics/topic_base.py | 27 +++++++++ grafi/workflows/workflow.py | 37 +++++++++++- 17 files changed, 319 insertions(+), 26 deletions(-) diff --git a/grafi/assistants/assistant_base.py b/grafi/assistants/assistant_base.py index 6dd0947..c58bf2b 100644 --- a/grafi/assistants/assistant_base.py +++ b/grafi/assistants/assistant_base.py @@ -98,17 +98,49 @@ class AssistantBaseBuilder(BaseBuilder[T_A]): """Inner builder class for Assistant construction.""" def oi_span_type(self, oi_span_type: OpenInferenceSpanKindValues) -> Self: + """Set the OpenInference span type for observability. + + Args: + oi_span_type: The span type for tracing (typically AGENT). + + Returns: + Self for method chaining. + """ self.kwargs["oi_span_type"] = oi_span_type return self def name(self, name: str) -> Self: + """Set the assistant's display name. + + Args: + name: Human-readable identifier for the assistant. + + Returns: + Self for method chaining. + """ self.kwargs["name"] = name return self def type(self, type_name: str) -> Self: + """Set the assistant's type identifier. + + Args: + type_name: Type classification for the assistant. + + Returns: + Self for method chaining. + """ self.kwargs["type"] = type_name return self def event_store(self, event_store: EventStore) -> Self: + """Register an event store for persistence. + + Args: + event_store: The event store implementation to use. + + Returns: + Self for method chaining. + """ container.register_event_store(event_store) return self diff --git a/grafi/common/models/base_builder.py b/grafi/common/models/base_builder.py index 99a9501..2123d33 100644 --- a/grafi/common/models/base_builder.py +++ b/grafi/common/models/base_builder.py @@ -1,6 +1,4 @@ # grafi/builder_core.py -from __future__ import annotations - from typing import Any from typing import Generic from typing import TypeVar diff --git a/grafi/nodes/node_base.py b/grafi/nodes/node_base.py index df228d0..f5ff20f 100644 --- a/grafi/nodes/node_base.py +++ b/grafi/nodes/node_base.py @@ -150,27 +150,67 @@ class NodeBaseBuilder(BaseBuilder[T_N]): """Inner builder class for workflow construction.""" def oi_span_type(self, oi_span_type: OpenInferenceSpanKindValues) -> Self: + """Set the OpenInference span type for observability. + + Args: + oi_span_type: The span type for tracing (e.g., CHAIN, TOOL). + + Returns: + Self for method chaining. + """ self.kwargs["oi_span_type"] = oi_span_type return self def name(self, name: str) -> Self: + """Set the node's display name. + + Args: + name: Human-readable identifier for the node. + + Returns: + Self for method chaining. + """ self.kwargs["name"] = name return self def type(self, type: str) -> Self: + """Set the node's type identifier. + + Args: + type: Type classification for the node. + + Returns: + Self for method chaining. + """ self.kwargs["type"] = type return self def tool(self, tool: Tool) -> Self: - """Set the tool for this node. Command will be auto-created.""" + """Set the tool for this node. Command will be auto-created. + + Args: + tool: The tool instance to execute when this node is invoked. + + Returns: + Self for method chaining. + """ self.kwargs["tool"] = tool return self def subscribe(self, subscribe_to: Union[TopicBase, SubExpr]) -> Self: - """ - Begin building a DSL expression. Returns a SubscriptionDSL.Builder, - which the user can chain with: - .subscribed_to(topicA).and_().subscribed_to(topicB).build() + """Subscribe this node to a topic or subscription expression. + + The node will be triggered when the subscription condition is met. + Can be chained with boolean expressions using SubscriptionBuilder. + + Args: + subscribe_to: A Topic or SubExpr (from SubscriptionBuilder). + + Returns: + Self for method chaining. + + Raises: + NodeExecutionError: If subscribe_to is not a Topic or SubExpr. """ if "subscribed_expressions" not in self.kwargs: self.kwargs["subscribed_expressions"] = [] @@ -187,6 +227,14 @@ def subscribe(self, subscribe_to: Union[TopicBase, SubExpr]) -> Self: return self def publish_to(self, topic: TopicBase) -> Self: + """Add a topic that this node will publish results to. + + Args: + topic: The topic to publish output events to. + + Returns: + Self for method chaining. + """ if "publish_to" not in self.kwargs: self.kwargs["publish_to"] = [] self.kwargs["publish_to"].append(topic) diff --git a/grafi/tools/function_calls/function_call_tool.py b/grafi/tools/function_calls/function_call_tool.py index e9cc39d..927367e 100644 --- a/grafi/tools/function_calls/function_call_tool.py +++ b/grafi/tools/function_calls/function_call_tool.py @@ -104,19 +104,16 @@ async def invoke( self, invoke_context: InvokeContext, input_data: Messages ) -> MsgsAGen: """ - Invoke the registered function with the given arguments. + Invoke the registered functions based on tool calls in input_data. This method is decorated with @record_tool_invoke to log its invoke. Args: - function_name (str): The name of the function to invoke. - arguments (Dict[str, Any]): The arguments to pass to the function. + invoke_context: Context for tracking the invocation lifecycle. + input_data: Messages containing tool_calls to execute. - Returns: - Any: The result of the function invoke. - - Raises: - ValueError: If the provided function_name doesn't match the registered function. + Yields: + Messages: Results from executing the tool calls. """ if not input_data: logger.warning("No input data provided.") @@ -231,6 +228,17 @@ class FunctionCallToolBuilder(ToolBuilder[T_F]): """ def function(self, function: Callable) -> Self: + """Add a function that can be called by the LLM. + + If the function is not already decorated with @llm_function, + it will be automatically wrapped. + + Args: + function: A callable to expose as a tool function. + + Returns: + Self for method chaining. + """ if not hasattr(function, "_function_spec"): function = llm_function(function) diff --git a/grafi/tools/function_calls/impl/synthetic_tool.py b/grafi/tools/function_calls/impl/synthetic_tool.py index 2349ca4..03b6faf 100644 --- a/grafi/tools/function_calls/impl/synthetic_tool.py +++ b/grafi/tools/function_calls/impl/synthetic_tool.py @@ -315,26 +315,74 @@ class SyntheticToolBuilder(FunctionCallToolBuilder[SyntheticTool]): """Builder for SyntheticTool instances.""" def tool_name(self, name: str) -> "SyntheticToolBuilder": + """Set the tool's name. + + Args: + name: Unique identifier for this synthetic tool. + + Returns: + Self for method chaining. + """ self.kwargs["tool_name"] = name self.kwargs["name"] = name return self def description(self, desc: str) -> "SyntheticToolBuilder": + """Set the tool's description for LLM context. + + Args: + desc: Human-readable description of what the tool does. + + Returns: + Self for method chaining. + """ self.kwargs["description"] = desc return self def input_model(self, model: type[BaseModel]) -> "SyntheticToolBuilder": + """Set the Pydantic model for input validation. + + Args: + model: A Pydantic BaseModel class defining the input schema. + + Returns: + Self for method chaining. + """ self.kwargs["input_model"] = model return self def output_model(self, model: type[BaseModel]) -> "SyntheticToolBuilder": + """Set the Pydantic model for output validation. + + Args: + model: A Pydantic BaseModel class defining the output schema. + + Returns: + Self for method chaining. + """ self.kwargs["output_model"] = model return self def model(self, model: str) -> "SyntheticToolBuilder": + """Set the LLM model to use for generation. + + Args: + model: Model identifier (e.g., 'gpt-4o-mini'). + + Returns: + Self for method chaining. + """ self.kwargs["model"] = model return self def openai_api_key(self, openai_api_key: str) -> "SyntheticToolBuilder": + """Set the OpenAI API key. + + Args: + openai_api_key: The API key for OpenAI authentication. + + Returns: + Self for method chaining. + """ self.kwargs["openai_api_key"] = openai_api_key return self diff --git a/grafi/tools/functions/function_tool.py b/grafi/tools/functions/function_tool.py index 9316e79..561a17e 100644 --- a/grafi/tools/functions/function_tool.py +++ b/grafi/tools/functions/function_tool.py @@ -138,9 +138,25 @@ class FunctionToolBuilder(ToolBuilder[T_FT]): """Builder for FunctionTool instances.""" def role(self, role: str) -> Self: + """Set the role for messages produced by this function. + + Args: + role: The message role (e.g., 'assistant', 'tool'). + + Returns: + Self for method chaining. + """ self.kwargs["role"] = role return self def function(self, function: Callable[[Messages], OutputType]) -> Self: + """Set the function to execute. + + Args: + function: A callable that takes Messages and returns output. + + Returns: + Self for method chaining. + """ self.kwargs["function"] = function return self diff --git a/grafi/tools/llms/impl/claude_tool.py b/grafi/tools/llms/impl/claude_tool.py index 28d9ad0..44c8165 100644 --- a/grafi/tools/llms/impl/claude_tool.py +++ b/grafi/tools/llms/impl/claude_tool.py @@ -2,8 +2,6 @@ ClaudeTool - Anthropic Claude implementation of grafi.tools.llms.llm.LLM """ -from __future__ import annotations - import asyncio import json import os diff --git a/grafi/tools/llms/impl/deepseek_tool.py b/grafi/tools/llms/impl/deepseek_tool.py index c676f34..0c35f46 100644 --- a/grafi/tools/llms/impl/deepseek_tool.py +++ b/grafi/tools/llms/impl/deepseek_tool.py @@ -9,8 +9,6 @@ setting `base_url="https://api.deepseek.com"`  :contentReference[oaicite:0]{index=0} """ -from __future__ import annotations - import asyncio import os from typing import Any diff --git a/grafi/tools/llms/impl/gemini_tool.py b/grafi/tools/llms/impl/gemini_tool.py index ee11320..1f27419 100644 --- a/grafi/tools/llms/impl/gemini_tool.py +++ b/grafi/tools/llms/impl/gemini_tool.py @@ -5,8 +5,6 @@ Docs & examples: https://ai.google.dev/gemini-api :contentReference[oaicite:0]{index=0} """ -from __future__ import annotations - import asyncio import json import os diff --git a/grafi/tools/llms/impl/openai_tool.py b/grafi/tools/llms/impl/openai_tool.py index 7be8aa3..d548f0b 100644 --- a/grafi/tools/llms/impl/openai_tool.py +++ b/grafi/tools/llms/impl/openai_tool.py @@ -228,6 +228,16 @@ async def from_dict(cls, data: Dict[str, Any]) -> "OpenAITool": class OpenAIToolBuilder(LLMBuilder[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 2703c88..41acfcc 100644 --- a/grafi/tools/llms/impl/openrouter_tool.py +++ b/grafi/tools/llms/impl/openrouter_tool.py @@ -2,8 +2,6 @@ OpenRouterTool - OpenRouter.ai implementation of grafi.tools.llms.llm.LLM """ -from __future__ import annotations - import asyncio import os from typing import Any diff --git a/grafi/tools/llms/llm.py b/grafi/tools/llms/llm.py index 57a5c6c..5a1b472 100644 --- a/grafi/tools/llms/llm.py +++ b/grafi/tools/llms/llm.py @@ -242,19 +242,52 @@ class LLMBuilder(ToolBuilder[T_L]): """Builder for LLM instances.""" def model(self, model: str) -> Self: + """Set the LLM model identifier. + + Args: + model: Model name (e.g., 'gpt-4', 'claude-3-opus'). + + Returns: + Self for method chaining. + """ self.kwargs["model"] = model return self def chat_params(self, params: Dict[str, Any]) -> Self: + """Set additional chat completion parameters. + + Args: + params: Dictionary of parameters (temperature, max_tokens, etc.). + If 'response_format' is included, structured_output is enabled. + + Returns: + Self for method chaining. + """ self.kwargs["chat_params"] = params if "response_format" in params: self.kwargs["structured_output"] = True return self def is_streaming(self, is_streaming: bool) -> Self: + """Enable or disable streaming responses. + + Args: + is_streaming: True to stream responses token-by-token. + + Returns: + Self for method chaining. + """ self.kwargs["is_streaming"] = is_streaming return self def system_message(self, system_message: Optional[str]) -> Self: + """Set the system message for the LLM. + + Args: + system_message: Instructions that guide the LLM's behavior. + + Returns: + Self for method chaining. + """ self.kwargs["system_message"] = system_message return self diff --git a/grafi/tools/tool.py b/grafi/tools/tool.py index 43bacee..29ef322 100644 --- a/grafi/tools/tool.py +++ b/grafi/tools/tool.py @@ -91,13 +91,37 @@ class ToolBuilder(BaseBuilder[T_T]): """Inner builder class for Tool construction.""" def name(self, name: str) -> Self: + """Set the tool's display name. + + Args: + name: Human-readable identifier for the tool. + + Returns: + Self for method chaining. + """ self.kwargs["name"] = name return self def type(self, type_name: str) -> Self: + """Set the tool's type identifier. + + Args: + type_name: Type classification for the tool. + + Returns: + Self for method chaining. + """ self.kwargs["type"] = type_name return self def oi_span_type(self, oi_span_type: OpenInferenceSpanKindValues) -> Self: + """Set the OpenInference span type for observability. + + Args: + oi_span_type: The span type for tracing (e.g., TOOL, LLM, CHAIN). + + Returns: + Self for method chaining. + """ self.kwargs["oi_span_type"] = oi_span_type return self diff --git a/grafi/topics/expressions/subscription_builder.py b/grafi/topics/expressions/subscription_builder.py index dad719d..1541eba 100644 --- a/grafi/topics/expressions/subscription_builder.py +++ b/grafi/topics/expressions/subscription_builder.py @@ -23,6 +23,18 @@ class SubscriptionBuilder(BaseModel): pending_op: Optional[LogicalOp] = None def subscribed_to(self, topic: TopicBase) -> "SubscriptionBuilder": + """Add a topic subscription to the expression. + + Args: + topic: The topic to subscribe to. + + Returns: + Self for method chaining. + + Raises: + ValueError: If topic is not a TopicBase instance. + ValueError: If chaining topics without an operator. + """ if not isinstance(topic, TopicBase): raise ValueError("subscribed_to(...) must receive a Topic object.") new_expr = TopicExpr(topic=topic) @@ -40,10 +52,20 @@ def subscribed_to(self, topic: TopicBase) -> "SubscriptionBuilder": return self def and_(self) -> "SubscriptionBuilder": + """Chain the next subscription with AND logic. + + Returns: + Self for method chaining. + """ self.pending_op = LogicalOp.AND return self def or_(self) -> "SubscriptionBuilder": + """Chain the next subscription with OR logic. + + Returns: + Self for method chaining. + """ self.pending_op = LogicalOp.OR return self 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 25e148e..e986b3d 100644 --- a/grafi/topics/queue_impl/in_mem_topic_event_queue.py +++ b/grafi/topics/queue_impl/in_mem_topic_event_queue.py @@ -15,8 +15,8 @@ class InMemTopicEventQueue(TopicEventQueue): """ In memory message queue where multiple publishers send events to all subscribers. - A publisher consists in any object who generates message. - A subscriber consists in any object who can consume messages. + A publisher consists of any object that generates messages. + A subscriber consists of any object that can consume messages. """ def __init__(self) -> None: diff --git a/grafi/topics/topic_base.py b/grafi/topics/topic_base.py index f48000e..7b7ee69 100644 --- a/grafi/topics/topic_base.py +++ b/grafi/topics/topic_base.py @@ -212,14 +212,41 @@ async def from_dict(cls, data: dict[str, Any]) -> "TopicBase": class TopicBaseBuilder(BaseBuilder[T_T]): + """Builder for TopicBase instances.""" + def name(self, name: str) -> Self: + """Set the topic's unique name. + + Args: + name: Unique identifier for the topic within a workflow. + + Returns: + Self for method chaining. + """ self.kwargs["name"] = name return self def type(self, type_name: str) -> Self: + """Set the topic's type identifier. + + Args: + type_name: Type classification for the topic. + + Returns: + Self for method chaining. + """ self.kwargs["type"] = type_name return self def condition(self, condition: Callable[[Messages], bool]) -> Self: + """Set a condition function that determines when this topic is satisfied. + + Args: + condition: A callable that takes Messages and returns True + when the topic's output condition is met. + + Returns: + Self for method chaining. + """ self.kwargs["condition"] = condition return self diff --git a/grafi/workflows/workflow.py b/grafi/workflows/workflow.py index 351ab26..464a148 100644 --- a/grafi/workflows/workflow.py +++ b/grafi/workflows/workflow.py @@ -74,25 +74,60 @@ async def from_dict(cls, data: dict[str, Any]) -> "Workflow": raise NotImplementedError("from_dict must be implemented in subclasses.") -T_W = TypeVar("T_W", bound="Workflow") # the Tool subclass +T_W = TypeVar("T_W", bound="Workflow") # the Workflow subclass class WorkflowBuilder(BaseBuilder[T_W]): """Inner builder class for Workflow construction.""" def oi_span_type(self, oi_span_type: OpenInferenceSpanKindValues) -> Self: + """Set the OpenInference span type for observability. + + Args: + oi_span_type: The span type for tracing (typically AGENT). + + Returns: + Self for method chaining. + """ self.kwargs["oi_span_type"] = oi_span_type return self def name(self, name: str) -> Self: + """Set the workflow's display name. + + Args: + name: Human-readable identifier for the workflow. + + Returns: + Self for method chaining. + """ self.kwargs["name"] = name return self def type(self, type_name: str) -> Self: + """Set the workflow's type identifier. + + Args: + type_name: Type classification for the workflow. + + Returns: + Self for method chaining. + """ self.kwargs["type"] = type_name return self def node(self, node: NodeBase) -> Self: + """Add a node to the workflow. + + Args: + node: The node to add to this workflow. + + Returns: + Self for method chaining. + + Raises: + DuplicateNodeError: If a node with the same name already exists. + """ if "nodes" not in self.kwargs: self.kwargs["nodes"] = {} if node.name in self.kwargs["nodes"]: From bc276af8f5bb9cbd24eaa54e6c082a869ffe73e7 Mon Sep 17 00:00:00 2001 From: GuanyiLi-Craig Date: Thu, 1 Jan 2026 16:56:45 +0000 Subject: [PATCH 3/4] update docs --- docs/docs/getting-started/quickstart.md | 4 +- .../docs/guide/connecting-to-an-mcp-server.md | 4 +- .../guide/getting-started-with-assistants.md | 2 +- docs/docs/index.md | 49 +++++++------------ 4 files changed, 23 insertions(+), 36 deletions(-) diff --git a/docs/docs/getting-started/quickstart.md b/docs/docs/getting-started/quickstart.md index b7952ae..3065048 100644 --- a/docs/docs/getting-started/quickstart.md +++ b/docs/docs/getting-started/quickstart.md @@ -77,11 +77,11 @@ You also need the following two dependencies for this guide. --- -## Use Build-in ReAct Agent +## Use Built-in ReAct Agent In graphite an agent is a specialized assistant that can handle events and perform actions based on the input it receives. We will create a ReAct agent that uses OpenAI's language model to process input, make function calls, and generate responses. -Create a file named `react_agent_app.py` and create a build-in react-agent: +Create a file named `react_agent_app.py` and create a built-in react-agent: ```python # react_agent_app.py diff --git a/docs/docs/guide/connecting-to-an-mcp-server.md b/docs/docs/guide/connecting-to-an-mcp-server.md index b81d5d9..d2ce3e2 100644 --- a/docs/docs/guide/connecting-to-an-mcp-server.md +++ b/docs/docs/guide/connecting-to-an-mcp-server.md @@ -19,7 +19,7 @@ Before getting started, make sure you have: - OpenAI API key - Basic understanding of Python and AI concepts - Basic understanding of what MCP Servers are -- Understand [Graphtie Assistants](../guide/getting-started-with-assistants.md) +- Understand [Graphite Assistants](../guide/getting-started-with-assistants.md) ## Code Walkthrough @@ -337,7 +337,7 @@ class StockAssistantBuilder( ``` -Graphtie is natively asychrnous, but you can chose to run syncrhous coroutines as well. For this case we are making a fully asynchronous workflow by overrding the `async def run()` method of the `Assistant` class. In order to run this create a `main.py` that will instantiate the assistant and execute it asynchrnously. +Graphite is natively asynchronous, but you can choose to run synchronous coroutines as well. For this case we are making a fully asynchronous workflow by overriding the `async def run()` method of the `Assistant` class. In order to run this create a `main.py` that will instantiate the assistant and execute it asynchronously. ```python #main.py diff --git a/docs/docs/guide/getting-started-with-assistants.md b/docs/docs/guide/getting-started-with-assistants.md index ba13ff0..0b42037 100644 --- a/docs/docs/guide/getting-started-with-assistants.md +++ b/docs/docs/guide/getting-started-with-assistants.md @@ -108,7 +108,7 @@ class FinanceAssistantBuilder(AssistantBaseBuilder[FinanceAssistant]): - Provides methods for setting API key, model, and system message - Returns `self` for method chaining -This class is used to set the fields from the `FinanceAssistant` the magic happens on the `builer` method up next. +This class is used to set the fields from the `FinanceAssistant`. The magic happens in the `builder` method up next. ### Builder Pattern Implementation diff --git a/docs/docs/index.md b/docs/docs/index.md index ec724a9..3e67509 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -8,50 +8,41 @@ Four critical capabilities—**observability, idempotency, auditability,** and * Overall, **Graphite** offers a powerful, extensible foundation for building AI solutions that scale, adapt to evolving compliance needs, and gracefully handle failures or user-driven pauses. By combining a robust workflow engine, well-structured nodes and tools, and a complete event model, Graphite enables teams to develop sophisticated conversational agents and automated pipelines with confidence. -## What is Graphite - -Graphite is an open-source platform that treats data as interconnected nodes and relationships, allowing you to: - -- **Process complex data relationships** with graph-based algorithms -- **Visualize data connections** through interactive network diagrams -- **Build analytical pipelines** that leverage graph structures -- **Scale efficiently** with distributed processing capabilities - -Whether you're analyzing social networks, tracking data lineage, exploring knowledge graphs, or building recommendation systems, Graphite provides the tools and abstractions you need to work effectively with connected data. - ## Key Features -**Graph-Native Processing**: Built from the ground up to handle graph data structures efficiently, with optimized algorithms for common graph operations like traversals, clustering, and pathfinding. +**Event-Driven Architecture**: Built on a pub/sub pattern where Topics manage message flow between Nodes, enabling loose coupling and flexible workflow composition. -**Visual Analytics**: Interactive visualization tools that help you explore and understand complex data relationships through customizable network diagrams and graph layouts. +**Modular Workflow Components**: Construct AI agents using composable layers - Assistants orchestrate Workflows, Workflows coordinate Nodes, and Nodes execute Tools. -**Flexible Data Integration**: Connect to various data sources including databases, APIs, and file formats, with built-in support for common graph data formats like GraphML, GEXF, and JSON. +**Multiple LLM Integrations**: Out-of-the-box support for OpenAI, Claude, Gemini, Ollama, DeepSeek, and OpenRouter, with a consistent interface across all providers. -**Extensible Architecture**: Plugin-based system that allows you to extend functionality with custom algorithms, data connectors, and visualization components. +**Function Calling Support**: Seamlessly integrate custom Python functions with LLMs through the FunctionCallTool, enabling agents to interact with external APIs and services. -**Performance Optimized**: Efficient memory management and parallel processing capabilities designed to handle large-scale graph datasets. +**MCP Server Integration**: Connect to Model Context Protocol servers for dynamic tool discovery and external data source access. + +**Production-Ready Features**: Built-in observability via OpenTelemetry, event sourcing for auditability, idempotent operations, and workflow restorability for fault tolerance. ## Who Should Use Graphite? -Graphite is designed for data scientists, analysts, researchers, and developers who work with interconnected data, including: +Graphite is designed for developers and teams building AI-powered applications, including: -- **Data Scientists** building recommendation engines or fraud detection systems -- **Business Analysts** exploring customer journey maps or organizational networks -- **Researchers** analyzing citation networks, protein interactions, or social structures -- **Developers** building applications that require graph-based computations +- **AI Engineers** building conversational agents with complex reasoning capabilities +- **Backend Developers** integrating LLM functionality into production systems +- **MLOps Teams** deploying observable, auditable AI workflows +- **Researchers** prototyping multi-step AI agents with tool use ## Getting Started This documentation will guide you through: 1. **Installation and Setup** - Get Graphite running in your environment -2. **Core Concepts** - Understand graphs, nodes, edges, and data models -3. **Data Import** - Load your data from various sources -4. **Processing and Analysis** - Apply algorithms and transformations -5. **Visualization** - Create interactive graph visualizations -6. **Advanced Topics** - Custom plugins, performance tuning, and deployment +2. **Core Concepts** - Understand the architecture: Assistants, Workflows, Nodes, and Tools +3. **Building Workflows** - Create event-driven AI pipelines +4. **Tool Integration** - Add LLMs, function calls, and MCP servers +5. **Observability** - Configure tracing with Arize and Phoenix +6. **Advanced Topics** - Event stores, workflow recovery, and custom tools -Ready to dive in? Start with our [Quick Start Guide](./getting-started/quickstart.md) to get Graphite up and running in minutes, or explore the [Core Concepts](./user-guide/architecture.md) to understand the fundamentals of graph-based data processing. +Ready to dive in? Start with our [Quick Start Guide](./getting-started/quickstart.md) to build your first AI agent, or explore the [Architecture](./user-guide/architecture.md) to understand how Graphite components work together. ## Community and Support @@ -61,7 +52,3 @@ Graphite is actively developed and maintained by the open-source community. Join - **Issues and Feature Requests**: Use GitHub Issues for bug reports and feature requests - **Discussions**: Join community discussions and get help from other users - **Contributing**: Check out our contribution guidelines to help improve Graphite - ---- - -*This documentation covers Graphite v0.0.x.* From a688016c399d4e6646cd543880661c9b373fd1e6 Mon Sep 17 00:00:00 2001 From: GuanyiLi-Craig Date: Thu, 1 Jan 2026 17:04:45 +0000 Subject: [PATCH 4/4] update docs --- docs/docs/guide/configuring-event-store.md | 4 ++-- docs/docs/guide/connecting-to-an-mcp-server.md | 6 +++--- docs/docs/guide/creating-a-simple-workflow.md | 4 ++-- docs/docs/guide/getting-started-with-assistants.md | 4 ++-- docs/docs/index.md | 2 +- docs/docs/user-guide/command.md | 8 ++++---- docs/docs/user-guide/event-interfaces.md | 6 +++--- docs/docs/user-guide/events/events.md | 6 +++--- docs/docs/user-guide/invoke-decorators.md | 2 +- docs/docs/user-guide/models.md | 6 +++--- docs/docs/user-guide/node.md | 2 +- docs/docs/user-guide/tools/function.md | 2 +- docs/docs/user-guide/tools/llm.md | 2 +- docs/docs/user-guide/tools/ollama.md | 4 ++-- docs/docs/user-guide/tools/openai.md | 4 ++-- docs/docs/user-guide/tools/tool.md | 10 +++++----- docs/docs/user-guide/topics/input_topics.md | 4 ++-- docs/docs/user-guide/topics/output_topics.md | 8 ++++---- docs/docs/user-guide/topics/topic.md | 4 ++-- 19 files changed, 44 insertions(+), 44 deletions(-) diff --git a/docs/docs/guide/configuring-event-store.md b/docs/docs/guide/configuring-event-store.md index 3fb17a8..b9e4f02 100644 --- a/docs/docs/guide/configuring-event-store.md +++ b/docs/docs/guide/configuring-event-store.md @@ -172,8 +172,8 @@ import uuid from grafi.agents.react_agent import create_react_agent from grafi.common.containers.container import container from grafi.common.event_stores.event_store_postgres import EventStorePostgres -from grafi.models.invoke_context import InvokeContext -from grafi.models.message import Message +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message postgres_event_store = EventStorePostgres( db_url="postgresql://postgres:postgres@localhost:5432/grafi_test_db", diff --git a/docs/docs/guide/connecting-to-an-mcp-server.md b/docs/docs/guide/connecting-to-an-mcp-server.md index d2ce3e2..2dbbbfe 100644 --- a/docs/docs/guide/connecting-to-an-mcp-server.md +++ b/docs/docs/guide/connecting-to-an-mcp-server.md @@ -348,9 +348,9 @@ from typing import Dict from grafi.common.containers.container import container from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent -from grafi.models.invoke_context import InvokeContext -from grafi.models.mcp_connections import StreamableHttpConnection -from grafi.models.message import Message +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.mcp_connections import StreamableHttpConnection +from grafi.common.models.message import Message from grafi.tools.function_calls.impl.mcp_tool import MCPTool from assistant import StockAssistant diff --git a/docs/docs/guide/creating-a-simple-workflow.md b/docs/docs/guide/creating-a-simple-workflow.md index 5c41511..fe5331b 100644 --- a/docs/docs/guide/creating-a-simple-workflow.md +++ b/docs/docs/guide/creating-a-simple-workflow.md @@ -43,8 +43,8 @@ The main function orchestrates the entire workflow. We start by defining a sampl ```python linenums="9" import uuid -from grafi.models.message import Message -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.common.models.invoke_context import InvokeContext def main(): user_input = "What is the capital of the United Kingdom" diff --git a/docs/docs/guide/getting-started-with-assistants.md b/docs/docs/guide/getting-started-with-assistants.md index 0b42037..1efab49 100644 --- a/docs/docs/guide/getting-started-with-assistants.md +++ b/docs/docs/guide/getting-started-with-assistants.md @@ -189,9 +189,9 @@ Prepare input data and context for workflow execution: ```python from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.invoke_context import InvokeContext from typing import Optional -from grafi.models.message import Message +from grafi.common.models.message import Message class FinanceAssistant(Assistant): diff --git a/docs/docs/index.md b/docs/docs/index.md index 3e67509..52bfa43 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -2,7 +2,7 @@ **Graphite** is an open-source framework for creating **domain-specific AI assistants** via composable, agentic workflows. It emphasizes loose coupling and well-defined interfaces, enabling developers to construct flexible, modular systems. Each major layer – **assistant, node, tool,** and **workflow** – has a clear role in orchestrating or executing tasks, with events serving as the single source of truth for every state change or data exchange. -This documentation details how **Graphite’s event-driven architecture** seamlessly supports complex business logic, from initial user requests through advanced tool integrations (e.g., LLM calls, function calls, RAG retrieval). Dedicated topics manage pub/sub operations, providing mechanisms for input, output, and human-in-the-loop interactions. Meanwhile, commands encapsulate invoke logic for each tool, allowing nodes to delegate work without tight coupling. +This documentation details how **Graphite's event-driven architecture** seamlessly supports complex business logic, from initial user requests through advanced tool integrations (e.g., LLM calls, function calls, MCP servers, and external APIs). Dedicated topics manage pub/sub operations, providing mechanisms for input, output, and human-in-the-loop interactions. Meanwhile, commands encapsulate invoke logic for each tool, allowing nodes to delegate work without tight coupling. Four critical capabilities—**observability, idempotency, auditability,** and **restorability**—underpin Graphite’s suitability for production AI environments. Observability is achieved via event sourcing and OpenTelemetry-based tracing, idempotency through carefully managed event stores and retry logic, auditability by logging every action and data flow, and restorability by maintaining offset-based consumption records that let workflows resume exactly where they left off. diff --git a/docs/docs/user-guide/command.md b/docs/docs/user-guide/command.md index f3af1b6..09f0ab8 100644 --- a/docs/docs/user-guide/command.md +++ b/docs/docs/user-guide/command.md @@ -258,7 +258,7 @@ Both commands enable a node to delegate specialized retrieval operations to thei Register custom commands for specific tool types: ```python -from grafi.models.command import use_command +from grafi.common.models.command import use_command @use_command(MyCustomCommand) class MySpecialTool(Tool): @@ -306,10 +306,10 @@ Create custom commands when you need: ```python from typing import List -from grafi.models.command import Command +from grafi.common.models.command import Command from grafi.common.events.topic_events.consume_from_topic_event import ConsumeFromTopicEvent -from grafi.models.invoke_context import InvokeContext -from grafi.models.message import Messages +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Messages class DatabaseQueryCommand(Command): """Command for database query tools with caching and optimization.""" diff --git a/docs/docs/user-guide/event-interfaces.md b/docs/docs/user-guide/event-interfaces.md index 4abe2bd..6f91b61 100644 --- a/docs/docs/user-guide/event-interfaces.md +++ b/docs/docs/user-guide/event-interfaces.md @@ -39,8 +39,8 @@ Published when a component sends data to a topic: ```python from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent -from grafi.models.invoke_context import InvokeContext -from grafi.models.message import Message +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message event = PublishToTopicEvent( publisher_name="ProcessorNode", @@ -94,7 +94,7 @@ class MyAssistant(Assistant): ```python from grafi.nodes.node import Node -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.invoke_context import InvokeContext 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 typing import List, AsyncGenerator diff --git a/docs/docs/user-guide/events/events.md b/docs/docs/user-guide/events/events.md index 70bf079..50aa315 100644 --- a/docs/docs/user-guide/events/events.md +++ b/docs/docs/user-guide/events/events.md @@ -80,7 +80,7 @@ Emitted when an assistant is invoked with input data. ```python from grafi.common.events.assistant_events.assistant_invoke_event import AssistantInvokeEvent -from grafi.models.message import Message +from grafi.common.models.message import Message event = AssistantInvokeEvent( invoke_context=context, @@ -307,8 +307,8 @@ The primary event used when consuming data from topics. This is the input format ```python from grafi.common.events.tool_events.tool_invoke_event import ToolInvokeEvent -from grafi.models.invoke_context import InvokeContext -from grafi.models.message import Message +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message # Create invoke context context = InvokeContext(assistant_request_id="req_123") diff --git a/docs/docs/user-guide/invoke-decorators.md b/docs/docs/user-guide/invoke-decorators.md index 250ad8a..b2b779c 100644 --- a/docs/docs/user-guide/invoke-decorators.md +++ b/docs/docs/user-guide/invoke-decorators.md @@ -116,7 +116,7 @@ Records synchronous node invocations in event-driven workflows. ```python from grafi.common.decorators.record_node_invoke import record_node_invoke -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.invoke_context import InvokeContext 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 typing import List, AsyncGenerator diff --git a/docs/docs/user-guide/models.md b/docs/docs/user-guide/models.md index 31fdab5..5121e58 100644 --- a/docs/docs/user-guide/models.md +++ b/docs/docs/user-guide/models.md @@ -32,7 +32,7 @@ In Graphite, various models provide the fundamental data structures that underpi ### Usage Example ```python -from grafi.models.message import Message +from grafi.common.models.message import Message # Creating a user message user_message = Message( @@ -64,7 +64,7 @@ assistant_message = Message( ### InvokeContext Usage Example ```python -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.invoke_context import InvokeContext context = InvokeContext( conversation_id="conv_123", @@ -96,7 +96,7 @@ context = InvokeContext( ### Command Usage Example ```python -from grafi.models.command import Command +from grafi.common.models.command import Command from grafi.tools.tool import Tool # Create command from tool diff --git a/docs/docs/user-guide/node.md b/docs/docs/user-guide/node.md index 8dbe5d7..188b6c6 100644 --- a/docs/docs/user-guide/node.md +++ b/docs/docs/user-guide/node.md @@ -93,7 +93,7 @@ node = Node.builder() .build() # Node invoke signature -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.invoke_context import InvokeContext from grafi.common.events.topic_events.consume_from_topic_event import ConsumeFromTopicEvent from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent diff --git a/docs/docs/user-guide/tools/function.md b/docs/docs/user-guide/tools/function.md index cae911c..9e3fef8 100644 --- a/docs/docs/user-guide/tools/function.md +++ b/docs/docs/user-guide/tools/function.md @@ -65,7 +65,7 @@ function_tool = ( ```python from grafi.tools.functions.function_tool import FunctionTool -from grafi.models.message import Messages +from grafi.common.models.message import Messages from pydantic import BaseModel class ProcessResult(BaseModel): diff --git a/docs/docs/user-guide/tools/llm.md b/docs/docs/user-guide/tools/llm.md index 3113901..3d34fe4 100644 --- a/docs/docs/user-guide/tools/llm.md +++ b/docs/docs/user-guide/tools/llm.md @@ -53,7 +53,7 @@ llm_tool = ( ### Function Calling Setup ```python -from grafi.models.function_spec import FunctionSpec +from grafi.common.models.function_spec import FunctionSpec # Create function specifications function_specs = [ diff --git a/docs/docs/user-guide/tools/ollama.md b/docs/docs/user-guide/tools/ollama.md index ba52d0b..f3c24b2 100644 --- a/docs/docs/user-guide/tools/ollama.md +++ b/docs/docs/user-guide/tools/ollama.md @@ -70,7 +70,7 @@ ollama_tool = ( ```python from grafi.tools.llms.impl.ollama_tool import OllamaTool -from grafi.models.message import Message +from grafi.common.models.message import Message # Create the tool ollama_tool = ( @@ -112,7 +112,7 @@ async def stream_example(): ### Function Calling ```python -from grafi.models.function_spec import FunctionSpec +from grafi.common.models.function_spec import FunctionSpec # Add function specifications function_spec = FunctionSpec( diff --git a/docs/docs/user-guide/tools/openai.md b/docs/docs/user-guide/tools/openai.md index 1b6e90f..88acdab 100644 --- a/docs/docs/user-guide/tools/openai.md +++ b/docs/docs/user-guide/tools/openai.md @@ -68,7 +68,7 @@ openai_tool = ( ```python from grafi.tools.llms.impl.openai_tool import OpenAITool -from grafi.models.message import Message +from grafi.common.models.message import Message # Create the tool openai_tool = ( @@ -110,7 +110,7 @@ async def stream_example(): ### Function Calling ```python -from grafi.models.function_spec import FunctionSpec +from grafi.common.models.function_spec import FunctionSpec # Add function specifications function_spec = FunctionSpec( diff --git a/docs/docs/user-guide/tools/tool.md b/docs/docs/user-guide/tools/tool.md index aec501a..5c6b4f0 100644 --- a/docs/docs/user-guide/tools/tool.md +++ b/docs/docs/user-guide/tools/tool.md @@ -132,8 +132,8 @@ from pydantic import Field from openinference.semconv.trace import OpenInferenceSpanKindValues from grafi.tools.tool import Tool, ToolBuilder -from grafi.models.invoke_context import InvokeContext -from grafi.models.message import Message, Messages, MsgsAGen +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message, Messages, MsgsAGen class TextProcessorTool(Tool): """A tool for processing text data.""" @@ -309,8 +309,8 @@ async def invoke(self, invoke_context: InvokeContext, input_data: Messages) -> M ```python import pytest -from grafi.models.message import Message -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.common.models.invoke_context import InvokeContext @pytest.mark.asyncio async def test_tool_invoke(): @@ -332,7 +332,7 @@ async def test_tool_invoke(): Tools integrate with Graphite's command system: ```python -from grafi.models.command import use_command +from grafi.common.models.command import use_command from grafi.tools.tool_command import ToolCommand @use_command(ToolCommand) diff --git a/docs/docs/user-guide/topics/input_topics.md b/docs/docs/user-guide/topics/input_topics.md index 6f954b2..c23ee8d 100644 --- a/docs/docs/user-guide/topics/input_topics.md +++ b/docs/docs/user-guide/topics/input_topics.md @@ -87,8 +87,8 @@ Builder for constructing InWorkflowInputTopic instances. ```python from grafi.topics.topic import InputTopic -from grafi.models.message import Message -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.common.models.invoke_context import InvokeContext # Create input topic input_topic = InputTopic(name="agent_input_topic") diff --git a/docs/docs/user-guide/topics/output_topics.md b/docs/docs/user-guide/topics/output_topics.md index 75eced0..2b89a87 100644 --- a/docs/docs/user-guide/topics/output_topics.md +++ b/docs/docs/user-guide/topics/output_topics.md @@ -106,8 +106,8 @@ Workflow topics are typically created as pairs for human-in-the-loop interaction ```python from grafi.topics.output_topic import OutputTopic, agent_output_topic -from grafi.models.message import Message -from grafi.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.common.models.invoke_context import InvokeContext # Create context and messages context = InvokeContext() @@ -131,7 +131,7 @@ if event: ```python import asyncio from typing import AsyncIterator -from grafi.models.message import Messages +from grafi.common.models.message import Messages async def streaming_response() -> AsyncIterator[Messages]: """Example async generator for streaming responses.""" @@ -208,7 +208,7 @@ async def managed_streaming(): ```python from grafi.topics.in_workflow_output_topic import InWorkflowOutputTopic -from grafi.models.message import Message +from grafi.common.models.message import Message # Create workflow output topic (paired with an input topic) workflow_output_topic = InWorkflowOutputTopic( diff --git a/docs/docs/user-guide/topics/topic.md b/docs/docs/user-guide/topics/topic.md index 23b348a..1821761 100644 --- a/docs/docs/user-guide/topics/topic.md +++ b/docs/docs/user-guide/topics/topic.md @@ -131,8 +131,8 @@ topic = (Topic.builder() ### Publishing Messages ```python -from grafi.models.invoke_context import InvokeContext -from grafi.models.message import Message +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message # Create context and messages context = InvokeContext()