Merge pull request #33 from syedfakher27/adk-middleware

Adk middleware

Author: contextablemark

typescript-sdk/apps/dojo/src/agents.ts

@@ -67,6 +67,7 @@ export const agentsIntegrations: AgentIntegrationConfig[] = [
         tool_based_generative_ui: new ServerStarterAgent({ url: "http://localhost:8000/adk-tool-based-generative-ui" }),
         human_in_the_loop: new ServerStarterAgent({ url: "http://localhost:8000/adk-human-in-loop-agent" }),
         shared_state: new ServerStarterAgent({ url: "http://localhost:8000/adk-shared-state-agent" }),
+        predictive_state_updates: new ServerStarterAgent({ url: "http://localhost:8000/adk-predictive-state-agent" }),
       };
     },
   },

typescript-sdk/apps/dojo/src/menu.ts

@@ -26,7 +26,7 @@ export const menuIntegrations: MenuIntegrationConfig[] = [
   {
     id: "adk-middleware",
     name: "ADK Middleware",
-    features: ["agentic_chat","tool_based_generative_ui","human_in_the_loop","shared_state"],
+    features: ["agentic_chat","tool_based_generative_ui","human_in_the_loop","shared_state","predictive_state_updates"],
   },
   {
     id: "server-starter-all-features",

typescript-sdk/integrations/adk-middleware/CHANGELOG.md

@@ -10,12 +10,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - **NEW**: SystemMessage support for ADK agents (issue #22) - SystemMessages as first message are now appended to agent instructions
 - **NEW**: Comprehensive tests for SystemMessage functionality including edge cases
+- **NEW**: Long running tools can be defined in backend side as well
+- **NEW**: Predictive state demo is added in dojo App
 
 ### Fixed  
 - **FIXED**: Race condition in tool result processing causing "No pending tool calls found" warnings
 - **FIXED**: Tool call removal now happens after pending check to prevent race conditions
 - **IMPROVED**: Better handling of empty tool result content with graceful JSON parsing fallback
 - **FIXED**: Pending tool call state management now uses SessionManager methods (issue #25)
+- **FIXED**: Pending tools issue for normal backend tools is now fixed (issue #32)
+- **FIXED**: TestEventTranslatorComprehensive unit test cases fixed
 
 ### Enhanced
 - **LOGGING**: Added debug logging for tool result processing to aid in troubleshooting

typescript-sdk/integrations/adk-middleware/examples/fastapi_server.py

@@ -12,6 +12,7 @@
 from .tool_based_generative_ui.agent import haiku_generator_agent
 from .human_in_the_loop.agent import human_in_loop_agent
 from .shared_state.agent import shared_state_agent
+from .predictive_state_updates.agent import predictive_state_updates_agent
 
 # Basic logging configuration
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -33,14 +34,15 @@
     sample_agent = LlmAgent(
         name="assistant",
         model="gemini-2.0-flash",
-        instruction="You are a helpful assistant.",
+        instruction="You are a helpful assistant. Help users by answering their questions and assisting with their needs.",
         tools=[adk_tools.preload_memory_tool.PreloadMemoryTool()]
     )
     # Register the agent
     registry.set_default_agent(sample_agent)
     registry.register_agent('adk-tool-based-generative-ui', haiku_generator_agent)
     registry.register_agent('adk-human-in-loop-agent', human_in_loop_agent)
     registry.register_agent('adk-shared-state-agent', shared_state_agent)
+    registry.register_agent('adk-predictive-state-agent', predictive_state_updates_agent)
     # Create ADK middleware agent
     adk_agent = ADKAgent(
         app_name="demo_app",
@@ -70,6 +72,13 @@
         use_in_memory_services=True
     )
 
+    adk_predictive_state_agent = ADKAgent(
+        app_name="demo_app",
+        user_id="demo_user",
+        session_timeout_seconds=3600,
+        use_in_memory_services=True
+    )
+    
     # Create FastAPI app
     app = FastAPI(title="ADK Middleware Demo")
 
@@ -78,6 +87,7 @@
     add_adk_fastapi_endpoint(app, adk_agent_haiku_generator, path="/adk-tool-based-generative-ui")
     add_adk_fastapi_endpoint(app, adk_human_in_loop_agent, path="/adk-human-in-loop-agent")
     add_adk_fastapi_endpoint(app, adk_shared_state_agent, path="/adk-shared-state-agent")
+    add_adk_fastapi_endpoint(app, adk_predictive_state_agent, path="/adk-predictive-state-agent")
 
     @app.get("/")
     async def root():

typescript-sdk/integrations/adk-middleware/examples/human_in_the_loop/agent.py

@@ -37,7 +37,7 @@
 
 
 human_in_loop_agent = Agent(
-    model='gemini-1.5-flash',
+    model='gemini-2.5-flash',
     name='human_in_loop_agent',
     instruction=f"""
         You are a human-in-the-loop task planning assistant that helps break down complex tasks into manageable steps with human oversight and approval.

typescript-sdk/integrations/adk-middleware/examples/predictive_state_updates/agent.py

@@ -0,0 +1,132 @@
+"""
+A demo of predictive state updates using Google ADK.
+"""
+
+from dotenv import load_dotenv
+load_dotenv()
+
+import json
+import uuid
+from typing import Dict, List, Any, Optional
+from google.adk.agents import LlmAgent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.sessions import InMemorySessionService, Session
+from google.adk.runners import Runner
+from google.adk.events import Event, EventActions
+from google.adk.tools import FunctionTool, ToolContext
+from google.genai.types import Content, Part, FunctionDeclaration
+from google.adk.models import LlmResponse, LlmRequest
+from google.genai import types
+
+
+def write_document(
+    tool_context: ToolContext,
+    document: str
+) -> Dict[str, str]:
+    """
+    Write a document. Use markdown formatting to format the document.
+    It's good to format the document extensively so it's easy to read.
+    You can use all kinds of markdown.
+    However, do not use italic or strike-through formatting, it's reserved for another purpose.
+    You MUST write the full document, even when changing only a few words.
+    When making edits to the document, try to make them minimal - do not change every word.
+    Keep stories SHORT!
+    
+    Args:
+        document: The document content to write in markdown format
+        
+    Returns:
+        Dict indicating success status and message
+    """
+    try:
+        # Update the session state with the new document
+        tool_context.state["document"] = document
+        
+        return {"status": "success", "message": "Document written successfully"}
+        
+    except Exception as e:
+        return {"status": "error", "message": f"Error writing document: {str(e)}"}
+
+
+def on_before_agent(callback_context: CallbackContext):
+    """
+    Initialize document state if it doesn't exist.
+    """
+    if "document" not in callback_context.state:
+        # Initialize with empty document
+        callback_context.state["document"] = None
+    
+    return None
+
+
+def before_model_modifier(
+    callback_context: CallbackContext, llm_request: LlmRequest
+) -> Optional[LlmResponse]:
+    """
+    Modifies the LLM request to include the current document state.
+    This enables predictive state updates by providing context about the current document.
+    """
+    agent_name = callback_context.agent_name
+    if agent_name == "DocumentAgent":
+        current_document = "No document yet"
+        if "document" in callback_context.state and callback_context.state["document"] is not None:
+            try:
+                current_document = callback_context.state["document"]
+            except Exception as e:
+                current_document = f"Error retrieving document: {str(e)}"
+        
+        # Modify the system instruction to include current document state
+        original_instruction = llm_request.config.system_instruction or types.Content(role="system", parts=[])
+        prefix = f"""You are a helpful assistant for writing documents. 
+        To write the document, you MUST use the write_document tool.
+        You MUST write the full document, even when changing only a few words.
+        When you wrote the document, DO NOT repeat it as a message. 
+        Just briefly summarize the changes you made. 2 sentences max.
+        This is the current state of the document: ----
+        {current_document}
+        -----"""
+        
+        # Ensure system_instruction is Content and parts list exists
+        if not isinstance(original_instruction, types.Content):
+            original_instruction = types.Content(role="system", parts=[types.Part(text=str(original_instruction))])
+        if not original_instruction.parts:
+            original_instruction.parts.append(types.Part(text=""))
+
+        # Modify the text of the first part
+        modified_text = prefix + (original_instruction.parts[0].text or "")
+        original_instruction.parts[0].text = modified_text
+        llm_request.config.system_instruction = original_instruction
+
+    return None
+
+
+# Create the predictive state updates agent
+predictive_state_updates_agent = LlmAgent(
+    name="DocumentAgent",
+    model="gemini-2.5-pro",
+    instruction="""
+    You are a helpful assistant for writing documents. 
+    To write the document, you MUST use the write_document tool.
+    You MUST write the full document, even when changing only a few words.
+    When you wrote the document, DO NOT repeat it as a message. 
+    Just briefly summarize the changes you made. 2 sentences max.
+
+    IMPORTANT RULES:
+    1. Always use the write_document tool for any document writing or editing requests
+    2. Write complete documents, not fragments
+    3. Use markdown formatting for better readability
+    4. Keep stories SHORT and engaging
+    5. After using the tool, provide a brief summary of what you created or changed
+    6. Do not use italic or strike-through formatting
+
+    Examples of when to use the tool:
+    - "Write a story about..." → Use tool with complete story in markdown
+    - "Edit the document to..." → Use tool with the full edited document
+    - "Add a paragraph about..." → Use tool with the complete updated document
+
+    Always provide complete, well-formatted documents that users can read and use.
+    """,
+    tools=[write_document],
+    before_agent_callback=on_before_agent,
+    before_model_callback=before_model_modifier
+)

typescript-sdk/integrations/adk-middleware/examples/shared_state/agent.py

@@ -69,7 +69,8 @@ class Recipe(BaseModel):
 def generate_recipe(
     tool_context: ToolContext,
     skill_level: str,
-    special_preferences: str = "",
+    title: str,
+    special_preferences: List[str] = [],
     cooking_time: str = "",
     ingredients: List[dict] = [],
     instructions: List[str] = [],
@@ -79,14 +80,20 @@ def generate_recipe(
     Generate or update a recipe using the provided recipe data.
     
     Args:
+        "title": {
+            "type": "string",
+            "description": "**REQUIRED** - The title of the recipe."
+        },
         "skill_level": {
             "type": "string",
             "enum": ["Beginner","Intermediate","Advanced"],
             "description": "**REQUIRED** - The skill level required for the recipe. Must be one of the predefined skill levels (Beginner, Intermediate, Advanced)."
         },
         "special_preferences": {
-            "type": "string",
-            "description": "**OPTIONAL** - Special dietary preferences for the recipe as comma-separated values. Example: 'High Protein, Low Carb, Gluten Free'. Leave empty or omit if no special preferences."
+            "type": "array",
+            "items": {"type": "string"},
+            "enum": ["High Protein","Low Carb","Spicy","Budget-Friendly","One-Pot Meal","Vegetarian","Vegan"],
+            "description": "**OPTIONAL** - Special dietary preferences for the recipe as comma-separated values. Example: 'High Protein, Low Carb, Gluten Free'. Leave empty array if no special preferences."
         },
         "cooking_time": {
             "type": "string",
@@ -123,6 +130,7 @@ def generate_recipe(
 
         # Create RecipeData object to validate structure
         recipe = {
+            "title": title,
             "skill_level": skill_level,
             "special_preferences": special_preferences ,
             "cooking_time": cooking_time ,
@@ -143,8 +151,7 @@ def generate_recipe(
 
         tool_context.state["recipe"] = current_recipe
 
-        # Log the update
-        print(f"Recipe updated: {recipe.get('change')}")
+
 
         return {"status": "success", "message": "Recipe generated successfully"}
 
@@ -158,18 +165,19 @@ def on_before_agent(callback_context: CallbackContext):
     """
     Initialize recipe state if it doesn't exist.
     """
-    print('recipe state ==>',callback_context.state.get('recipe'))
+
     if "recipe" not in callback_context.state:
         # Initialize with default recipe
         default_recipe =     {
+            "title": "Make Your Recipe",
             "skill_level": "Beginner",
             "special_preferences": [],
             "cooking_time": '15 min',
             "ingredients": [{"icon": "🍴", "name": "Sample Ingredient", "amount": "1 unit"}],
             "instructions": ["First step instruction"]
         }
         callback_context.state["recipe"] = default_recipe
-        print("Initialized default recipe state")
+ 
 
     return None
 
@@ -181,7 +189,6 @@ def before_model_modifier(
 ) -> Optional[LlmResponse]:
     """Inspects/modifies the LLM request or skips the call."""
     agent_name = callback_context.agent_name
-    print(f"[Callback] Before model call for agent: {agent_name}")
     if agent_name == "RecipeAgent":
         recipe_json = "No recipe yet"
         if "recipe" in callback_context.state and callback_context.state["recipe"] is not None:
@@ -212,6 +219,28 @@ def before_model_modifier(
     return None
 
 
+# --- Define the Callback Function ---
+def simple_after_model_modifier(
+    callback_context: CallbackContext, llm_response: LlmResponse
+) -> Optional[LlmResponse]:
+    """Stop the consecutive tool calling of the agent"""
+    agent_name = callback_context.agent_name
+    # --- Inspection ---
+    if agent_name == "RecipeAgent":
+        original_text = ""
+        if llm_response.content and llm_response.content.parts:
+            # Assuming simple text response for this example
+            if  llm_response.content.role=='model' and llm_response.content.parts[0].text:
+                original_text = llm_response.content.parts[0].text
+                callback_context._invocation_context.end_invocation = True
+        
+        elif llm_response.error_message:
+            return None
+        else:
+            return None # Nothing to modify
+    return None
+
+
 shared_state_agent = LlmAgent(
         name="RecipeAgent",
         model="gemini-2.5-pro",
@@ -224,16 +253,18 @@ def before_model_modifier(
         3. When modifying an existing recipe, include the changes parameter to describe what was modified
         4. Be creative and helpful in generating complete, practical recipes
         5. After using the tool, provide a brief summary of what you created or changed
+        6. If user ask to improve the recipe then add more ingredients and make it healthier
+        7. When you see the 'Recipe generated successfully' confirmation message, wish the user well with their cooking by telling them to enjoy their dish.
 
         Examples of when to use the tool:
         - "Create a pasta recipe" → Use tool with skill_level, ingredients, instructions
-        - "Make it vegetarian" → Use tool with special_preferences="vegetarian" and changes describing the modification
+        - "Make it vegetarian" → Use tool with special_preferences=["vegetarian"] and changes describing the modification
         - "Add some herbs" → Use tool with updated ingredients and changes describing the addition
 
         Always provide complete, practical recipes that users can actually cook.
         """,
         tools=[generate_recipe],
         before_agent_callback=on_before_agent,
-        before_model_callback=before_model_modifier
-    )
-
+        before_model_callback=before_model_modifier,
+        after_model_callback = simple_after_model_modifier
+    )