add kwargs to invoke_context and add an example

Author: Guanyi Li

docs/docs/user-guide/models.md

@@ -58,7 +58,8 @@ assistant_message = Message(
 | `conversation_id`       | Unique identifier for a conversation between user and assistant.                       |
 | `invoke_id`             | Unique identifier for each conversation invoke - an invoke can involve multiple agents.|
 | `assistant_request_id`  | Created when an agent receives a request from the user.                                |
-| `user_id`               | Optional user identifier, defaults to empty string.                                   |
+| `user_id`               | Optional user identifier, defaults to empty string.                                    |
+| `kwargs`                | Optional additional keyword arguments and context for the workflow                     |
 
 ### InvokeContext Usage Example
 

grafi/common/models/invoke_context.py

@@ -1,3 +1,5 @@
+from typing import Any
+from typing import Dict
 from typing import Optional
 
 from pydantic import BaseModel
@@ -10,9 +12,22 @@ class InvokeContext(BaseModel):
     invoke_id: invoke id of each conversation, an invoke can involve multiple agents
     assistant_request_id: assistant_request_id is create when agent receive a request from the user
     user_id: user id
+    kwargs: optional field for any additional context or keyword arguments that need to be passed through the workflow
     """
 
-    conversation_id: str
-    invoke_id: str
-    assistant_request_id: str
-    user_id: Optional[str] = Field(default="")
+    conversation_id: str = Field(
+        description="Unique identifier for a conversation between user and assistant"
+    )
+    invoke_id: str = Field(
+        description="Unique identifier for each conversation invoke - an invoke can involve multiple agents"
+    )
+    assistant_request_id: str = Field(
+        description="Created when an agent receives a request from the user"
+    )
+    user_id: Optional[str] = Field(
+        default="", description="Optional user identifier, defaults to empty string"
+    )
+    kwargs: Optional[Dict[str, Any]] = Field(
+        default_factory=dict,
+        description="Additional keyword arguments and context for the workflow",
+    )

pyproject.toml

@@ -30,6 +30,7 @@ dev = [
     "llama-index-core>=0.12.48",
     "llama-index-embeddings-openai>=0.3.1",
     "llama-index-llms-openai>=0.4.7",
+    "markdown>=3.8.2",
     "mcp>=1.11.0",
     "mypy>=1.16.1",
     "ollama>=0.5.1",

tests_integration/invoke_kwargs/llm_prompt_template_command.py

@@ -0,0 +1,50 @@
+from typing import List
+
+from jinja2 import Template
+from loguru import logger
+
+from grafi.common.events.topic_events.consume_from_topic_event import (
+    ConsumeFromTopicEvent,
+)
+from grafi.common.models.invoke_context import InvokeContext
+from grafi.common.models.message import Message
+from grafi.common.models.message import Messages
+from grafi.tools.llms.llm_command import LLMCommand
+
+
+class LLMPromptTemplateCommand(LLMCommand):
+
+    def get_tool_input(
+        self,
+        invoke_context: InvokeContext,
+        node_input: List[ConsumeFromTopicEvent],
+    ) -> Messages:
+        """Prepare the input for the LLM command based on the node input and invoke context."""
+
+        messages = super().get_tool_input(invoke_context, node_input)
+
+        message = messages[-1] if messages else None
+
+        if "prompt_template" in invoke_context.kwargs:
+
+            template_str = invoke_context.kwargs["prompt_template"]
+            template: Template = Template(template_str)
+
+            if message and message.content:
+                # Render the Jinja template with the message content as input
+                try:
+                    rendered_prompt = template.render(input_text=message.content)
+
+                    # Create a new message with the rendered template
+                    new_message = Message(role="user", content=rendered_prompt)
+
+                    logger.info(
+                        f"Rendered prompt template with input: {message.content[:100]}..."
+                    )
+                    return [new_message]
+                except Exception as e:
+                    logger.error(f"Error rendering template: {e}")
+                    # Fallback to original message if template rendering fails
+                    return messages
+
+        return messages

tests_integration/invoke_kwargs/simple_llm_prompt_template_assistant.py

@@ -0,0 +1,92 @@
+import os
+from typing import Optional
+from typing import Self
+
+from openinference.semconv.trace import OpenInferenceSpanKindValues
+from pydantic import Field
+
+from grafi.assistants.assistant import Assistant
+from grafi.assistants.assistant_base import AssistantBaseBuilder
+from grafi.common.topics.output_topic import agent_output_topic
+from grafi.common.topics.topic import agent_input_topic
+from grafi.nodes.node import Node
+from grafi.tools.llms.impl.openai_tool import OpenAITool
+from grafi.workflows.impl.event_driven_workflow import EventDrivenWorkflow
+from tests_integration.invoke_kwargs.llm_prompt_template_command import (
+    LLMPromptTemplateCommand,
+)
+
+
+class SimpleLLMPromptTemplateAssistant(Assistant):
+    """
+    A simple assistant class that uses OpenAI's language model to process input and generate responses.
+
+    This class sets up a workflow with a single LLM node using OpenAI's API, and provides a method
+    to run input through this workflow.
+
+    Attributes:
+        api_key (str): The API key for OpenAI. If not provided, it tries to use the OPENAI_API_KEY environment variable.
+        model (str): The name of the OpenAI model to use.
+        event_store (EventStore): An instance of EventStore to record events during the assistant's operation.
+    """
+
+    oi_span_type: OpenInferenceSpanKindValues = Field(
+        default=OpenInferenceSpanKindValues.AGENT
+    )
+    name: str = Field(default="SimpleLLMPromptTemplateAssistant")
+    type: str = Field(default="SimpleLLMPromptTemplateAssistant")
+    api_key: Optional[str] = Field(default_factory=lambda: os.getenv("OPENAI_API_KEY"))
+    system_message: Optional[str] = Field(default=None)
+    model: str = Field(default="gpt-4o-mini")
+
+    @classmethod
+    def builder(cls) -> "SimpleLLMPromptTemplateAssistantBuilder":
+        """Return a builder for SimpleLLMPromptTemplateAssistant."""
+        return SimpleLLMPromptTemplateAssistantBuilder(cls)
+
+    def _construct_workflow(self) -> "SimpleLLMPromptTemplateAssistant":
+        # Create an LLM node
+        llm_node: Node = (
+            Node.builder()
+            .name("OpenAINode")
+            .subscribe(agent_input_topic)
+            .publish_to(agent_output_topic)
+            .build()
+        )
+
+        llm_node.command = LLMPromptTemplateCommand(
+            tool=OpenAITool.builder()
+            .name("OpenAITool")
+            .api_key(self.api_key)
+            .model(self.model)
+            .system_message(self.system_message)
+            .build()
+        )
+
+        # Create a workflow and add the LLM node
+        self.workflow = (
+            EventDrivenWorkflow.builder()
+            .name("SimpleLLMWorkflow")
+            .node(llm_node)
+            .build()
+        )
+
+        return self
+
+
+class SimpleLLMPromptTemplateAssistantBuilder(
+    AssistantBaseBuilder[SimpleLLMPromptTemplateAssistant]
+):
+    """Concrete builder for SimpleLLMPromptTemplateAssistant."""
+
+    def api_key(self, api_key: str) -> Self:
+        self.kwargs["api_key"] = api_key
+        return self
+
+    def system_message(self, system_message: str) -> Self:
+        self.kwargs["system_message"] = system_message
+        return self
+
+    def model(self, model: str) -> Self:
+        self.kwargs["model"] = model
+        return self

tests_integration/invoke_kwargs/simple_llm_prompt_template_assistant_example.py

@@ -0,0 +1,93 @@
+# We will test the SimpleLLMAssistant class in this file.
+
+import os
+import uuid
+
+import markdown
+
+from grafi.common.containers.container import container
+from grafi.common.instrumentations.tracing import TracingOptions
+from grafi.common.instrumentations.tracing import setup_tracing
+from grafi.common.models.invoke_context import InvokeContext
+from grafi.common.models.message import Message
+from tests_integration.invoke_kwargs.simple_llm_prompt_template_assistant import (
+    SimpleLLMPromptTemplateAssistant,
+)
+
+
+container.register_tracer(setup_tracing(tracing_options=TracingOptions.IN_MEMORY))
+event_store = container.event_store
+
+api_key = os.getenv("OPENAI_API_KEY", "")
+
+
+def get_invoke_context() -> InvokeContext:
+
+    return InvokeContext(
+        conversation_id="conversation_id",
+        invoke_id=uuid.uuid4().hex,
+        assistant_request_id=uuid.uuid4().hex,
+        kwargs={
+            "prompt_template": """You are a skilled poet and literary analyst. Your task is to analyze the given input and create a 14-line sonnet based on your analysis.
+
+## Input Analysis Instructions:
+1. Carefully read and understand the provided input
+2. Identify the main themes, emotions, imagery, and key concepts
+3. Determine the appropriate tone and mood for the sonnet
+4. Consider metaphors, symbolism, or literary devices that would enhance the poem
+
+## Sonnet Requirements:
+- Must be exactly 14 lines
+- Follow traditional sonnet structure (Shakespearean or Petrarchan)
+- Use appropriate rhyme scheme (ABAB CDCD EFEF GG for Shakespearean, or ABBAABBA CDECDE/CDCDCD for Petrarchan)
+- Maintain consistent meter (preferably iambic pentameter)
+- Include a clear thematic development with a turn (volta)
+- End with a powerful concluding couplet or tercet
+
+## Input to Analyze:
+{{ input_text }}
+
+## Analysis:
+Please first provide a brief analysis of the input, identifying:
+- Central themes: 
+- Emotional tone: 
+- Key imagery: 
+- Poetic approach: 
+
+## Generated Sonnet:
+Based on your analysis, create a 14-line sonnet that captures the essence of the input:
+
+```
+[Your 14-line sonnet here]
+```
+
+## Explanation:
+Briefly explain how your sonnet reflects the input and the literary devices used."""
+        },
+    )
+
+
+def test_simple_llm_assistant() -> None:
+    invoke_context = get_invoke_context()
+    assistant = (
+        SimpleLLMPromptTemplateAssistant.builder()
+        .name("SimpleLLMPromptTemplateAssistant")
+        .build()
+    )
+    event_store.clear_events()
+
+    input_data = [
+        Message(
+            content="Graphite is a event driven agentic AI platform, it offers real time observability, comprehensive auditing, and high performance workflow.",
+            role="user",
+        )
+    ]
+    output = assistant.invoke(invoke_context, input_data)
+
+    html = markdown.markdown(output[0].content)
+    print(html)
+    assert output is not None
+    assert len(event_store.get_events()) == 12
+
+
+test_simple_llm_assistant()