Add LlamaIndex Python Integration (#82)

logan-markewich · web-flow · commit 3761b09efd57 · 2025-06-18T11:17:14.000+02:00
* wip

* nit

* add gitignore

* wip

* wip

* remove non-working examples

* improve agentic generative UI example

* update to latest

* move and rename

* create llamaindex agent

* add LlamaIndexAgent package

* fix paths
diff --git a/typescript-sdk/apps/dojo/package.json b/typescript-sdk/apps/dojo/package.json
@@ -11,6 +11,7 @@
   "dependencies": {
     "@ag-ui/agno": "workspace:*",
     "@ag-ui/langgraph": "workspace:*",
+    "@ag-ui/llamaindex": "workspace:*",
     "@ag-ui/mastra": "workspace:*",
     "@ag-ui/middleware-starter": "workspace:*",
     "@ag-ui/server-starter": "workspace:*",
diff --git a/typescript-sdk/apps/dojo/src/agents.ts b/typescript-sdk/apps/dojo/src/agents.ts
@@ -8,6 +8,7 @@ import { VercelAISDKAgent } from "@ag-ui/vercel-ai-sdk";
 import { openai } from "@ai-sdk/openai";
 import { LangGraphAgent } from "@ag-ui/langgraph";
 import { AgnoAgent } from "@ag-ui/agno";
+import { LlamaIndexAgent } from "@ag-ui/llamaindex";
 
 export const agentsIntegrations: AgentIntegrationConfig[] = [
   {
@@ -112,4 +113,23 @@ export const agentsIntegrations: AgentIntegrationConfig[] = [
       };
     },
   },
+  {
+    id: "llama-index",
+    agents: async () => {
+      return {
+        agentic_chat: new LlamaIndexAgent({
+          url: "http://localhost:9000/agentic_chat/run",
+        }),
+        human_in_the_loop: new LlamaIndexAgent({
+          url: "http://localhost:9000/human_in_the_loop/run",
+        }),
+        agentic_generative_ui: new LlamaIndexAgent({
+          url: "http://localhost:9000/agentic_generative_ui/run",
+        }),
+        shared_state: new LlamaIndexAgent({
+          url: "http://localhost:9000/shared_state/run",
+        }),
+      };
+    },
+  }
 ];
diff --git a/typescript-sdk/apps/dojo/src/menu.ts b/typescript-sdk/apps/dojo/src/menu.ts
@@ -50,4 +50,14 @@ export const menuIntegrations: MenuIntegrationConfig[] = [
     name: "Agno",
     features: ["agentic_chat"],
   },
+  {
+    id: "llama-index",
+    name: "LlamaIndex",
+    features: [
+      "agentic_chat",
+      "human_in_the_loop",
+      "agentic_generative_ui",
+      "shared_state",
+    ],
+  },
 ];
diff --git a/typescript-sdk/integrations/llamaindex/.gitignore b/typescript-sdk/integrations/llamaindex/.gitignore
@@ -0,0 +1,3 @@
+__pycache__
+*egg-info
+.venv
diff --git a/typescript-sdk/integrations/llamaindex/.npmignore b/typescript-sdk/integrations/llamaindex/.npmignore
@@ -0,0 +1,12 @@
+.turbo
+.DS_Store
+.git
+.gitignore
+.idea
+.vscode
+.env
+__tests__
+src
+tsup.config.ts
+tsconfig.json
+jest.config.js
diff --git a/typescript-sdk/integrations/llamaindex/jest.config.js b/typescript-sdk/integrations/llamaindex/jest.config.js
@@ -0,0 +1,9 @@
+/** @type {import('ts-jest').JestConfigWithTsJest} */
+module.exports = {
+  preset: "ts-jest",
+  testEnvironment: "node",
+  testMatch: ["**/*.test.ts"],
+  moduleNameMapper: {
+    "^@/(.*)$": "<rootDir>/src/$1",
+  },
+};
diff --git a/typescript-sdk/integrations/llamaindex/package.json b/typescript-sdk/integrations/llamaindex/package.json
@@ -0,0 +1,35 @@
+{
+  "name": "@ag-ui/llamaindex",
+  "author": "Logan Markewich <logan@runllama.ai>",
+  "version": "0.1.0",
+  "main": "./dist/index.js",
+  "module": "./dist/index.mjs",
+  "types": "./dist/index.d.ts",
+  "sideEffects": false,
+  "files": [
+    "dist/**"
+  ],
+  "scripts": {
+    "build": "tsup",
+    "dev": "tsup --watch",
+    "clean": "rm -rf dist",
+    "typecheck": "tsc --noEmit",
+    "test": "jest",
+    "link:global": "pnpm link --global",
+    "unlink:global": "pnpm unlink --global"
+  },
+  "dependencies": {
+    "@ag-ui/client": "workspace:*"
+  },
+  "peerDependencies": {
+    "rxjs": "7.8.1"
+  },
+  "devDependencies": {
+    "@types/jest": "^29.5.14",
+    "@types/node": "^20.11.19",
+    "jest": "^29.7.0",
+    "ts-jest": "^29.1.2",
+    "tsup": "^8.0.2",
+    "typescript": "^5.3.3"
+  }
+}
diff --git a/typescript-sdk/integrations/llamaindex/server-py/README.md b/typescript-sdk/integrations/llamaindex/server-py/README.md
@@ -0,0 +1,20 @@
+# LlamaIndex Python AG-UI Integration
+
+This package provides a FastAPI server that is bootstrapped with several AG-UI+LlamaIndex endpoints. It can be used to communicate with AG-UI compatible frameworks like [CopilotKit](https://docs.copilotkit.ai/).
+
+## Usage
+
+Launch the server with:
+
+```python
+uv sync
+uv run dev
+```
+
+Launch the frontend dojo with:
+
+```bash
+cd ../../
+pnpm install
+turbo run dev
+```
diff --git a/typescript-sdk/integrations/llamaindex/server-py/pyproject.toml b/typescript-sdk/integrations/llamaindex/server-py/pyproject.toml
@@ -0,0 +1,30 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "server"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.9, <3.14"
+dependencies = [
+    "llama-index-core>=0.12.41,<0.13",
+    "llama-index-agent-openai>=0.4.9,<0.5",
+    "llama-index-protocols-ag-ui>=0.1.2",
+    "jsonpatch>=1.33",
+    "uvicorn>=0.27.0",
+    "fastapi>=0.100.0",
+]
+authors = [
+    { name = "Logan Markewich", email = "logan@runllama.ai" },
+]
+
+[tool.hatch.build.targets.sdist]
+include = ["server/"]
+
+[tool.hatch.build.targets.wheel]
+include = ["server/"]
+
+[project.scripts]
+dev = "server:main"
diff --git a/typescript-sdk/integrations/llamaindex/server-py/server/__init__.py b/typescript-sdk/integrations/llamaindex/server-py/server/__init__.py
@@ -0,0 +1,23 @@
+import uvicorn
+from fastapi import FastAPI
+
+from .routers.agentic_chat import agentic_chat_router
+from .routers.human_in_the_loop import human_in_the_loop_router
+from .routers.agentic_generative_ui import agentic_generative_ui_router
+from .routers.shared_state import shared_state_router
+
+app = FastAPI(title="AG-UI Llama-Index Endpoint")
+
+app.include_router(agentic_chat_router, prefix="/agentic_chat")
+app.include_router(human_in_the_loop_router, prefix="/human_in_the_loop")
+app.include_router(agentic_generative_ui_router, prefix="/agentic_generative_ui")
+app.include_router(shared_state_router, prefix="/shared_state")
+
+def main():
+    """Main function to start the FastAPI server."""
+    uvicorn.run(app, host="0.0.0.0", port=9000)
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["main"]
diff --git a/typescript-sdk/integrations/llamaindex/server-py/server/routers/agentic_chat.py b/typescript-sdk/integrations/llamaindex/server-py/server/routers/agentic_chat.py
@@ -0,0 +1,16 @@
+from llama_index.llms.openai import OpenAI
+from llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router
+from typing import Annotated
+
+
+# This tool has a client-side version that is actually called to change the background
+def change_background(
+    background: Annotated[str, "The background. Prefer gradients."],
+) -> str:
+    """Change the background color of the chat. Can be anything that the CSS background attribute accepts. Regular colors, linear of radial gradients etc."""
+    return f"Changing background to {background}"
+
+agentic_chat_router = get_ag_ui_workflow_router(
+    llm=OpenAI(model="gpt-4.1"),
+    frontend_tools=[change_background],
+)
diff --git a/typescript-sdk/integrations/llamaindex/server-py/server/routers/agentic_generative_ui.py b/typescript-sdk/integrations/llamaindex/server-py/server/routers/agentic_generative_ui.py
@@ -0,0 +1,87 @@
+import asyncio
+import copy
+import jsonpatch
+from pydantic import BaseModel
+
+from llama_index.core.workflow import Context
+from llama_index.llms.openai import OpenAI
+from llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router
+from llama_index.protocols.ag_ui.events import StateDeltaWorkflowEvent, StateSnapshotWorkflowEvent
+
+class Step(BaseModel):
+    description: str
+
+class Task(BaseModel):
+    steps: list[Step]
+
+# Genrative UI demo
+async def run_task(
+    ctx: Context, task: Task,
+) -> str:
+    """Execute any list of steps needed to complete a task. Useful for anything the user wants to do."""
+    state = await ctx.get("state", default={})
+    task = Task.model_validate(task)
+
+    state = {
+        "steps": [
+            {
+                "description": step.description,
+                "status": "pending"
+            }
+            for step in task.steps
+        ]
+    }
+
+    # Send initial state snapshot
+    ctx.write_event_to_stream(
+        StateSnapshotWorkflowEvent(
+            snapshot=state
+        )
+    )
+
+    # Sleep for 1 second
+    await asyncio.sleep(1.0)
+
+    # Create a copy to track changes for JSON patches
+    previous_state = copy.deepcopy(state)
+
+    # Update each step and send deltas
+    for i, step in enumerate(state["steps"]):
+        step["status"] = "completed"
+        
+        # Generate JSON patch from previous state to current state
+        patch = jsonpatch.make_patch(previous_state, state)
+        
+        # Send state delta event
+        ctx.write_event_to_stream(
+            StateDeltaWorkflowEvent(
+                delta=patch.patch
+            )
+        )
+        
+        # Update previous state for next iteration
+        previous_state = copy.deepcopy(state)
+        
+        # Sleep for 1 second
+        await asyncio.sleep(1.0)
+
+    # Optionally send a final snapshot to the client
+    ctx.write_event_to_stream(
+        StateSnapshotWorkflowEvent(
+            snapshot=state
+        )
+    )
+
+    return "Done!"
+
+
+agentic_generative_ui_router = get_ag_ui_workflow_router(
+    llm=OpenAI(model="gpt-4.1"),
+    frontend_tools=[run_task],
+    initial_state={},
+    system_prompt=(
+        "You are a helpful assistant that can help the user with their task. "
+        "If the user asks you to do any task, use the run_task tool to do it. "
+        "Use your best judgement to describe the steps."
+    )
+)
diff --git a/typescript-sdk/integrations/llamaindex/server-py/server/routers/human_in_the_loop.py b/typescript-sdk/integrations/llamaindex/server-py/server/routers/human_in_the_loop.py
@@ -0,0 +1,21 @@
+from typing import Literal, List
+from pydantic import BaseModel
+
+from llama_index.llms.openai import OpenAI
+from llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router
+
+
+
+class Step(BaseModel):
+    description: str
+    status: Literal["enabled", "disabled", "executing"]
+
+
+def generate_task_steps(steps: List[Step]) -> str:
+    return f"Generated {len(steps)} steps"
+
+
+human_in_the_loop_router = get_ag_ui_workflow_router(
+    llm=OpenAI(model="gpt-4.1"),
+    frontend_tools=[generate_task_steps],
+)
diff --git a/typescript-sdk/integrations/llamaindex/server-py/server/routers/shared_state.py b/typescript-sdk/integrations/llamaindex/server-py/server/routers/shared_state.py
@@ -0,0 +1,52 @@
+from typing import Literal, List
+from pydantic import BaseModel
+
+from llama_index.core.workflow import Context
+from llama_index.llms.openai import OpenAI
+from llama_index.protocols.ag_ui.events import StateSnapshotWorkflowEvent
+from llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router
+
+
+class Ingredient(BaseModel):
+    icon: str
+    name: str
+    amount: str
+
+class Recipe(BaseModel):
+    skill_level: str
+    special_preferences: List[str]
+    cooking_time: str
+    ingredients: List[Ingredient]
+    instructions: List[str]
+
+
+async def update_recipe(ctx: Context, recipe: Recipe) -> str:
+    """Useful for recording a recipe to shared state."""
+    recipe = Recipe.model_validate(recipe)
+
+    state = await ctx.get("state")
+    if state is None:
+        state = {}
+
+    state["recipe"] = recipe.model_dump()
+
+    ctx.write_event_to_stream(
+        StateSnapshotWorkflowEvent(
+            snapshot=state
+        )
+    )
+
+    await ctx.set("state", state)
+
+    return "Recipe updated!"
+
+
+shared_state_router = get_ag_ui_workflow_router(
+    llm=OpenAI(model="gpt-4.1"),
+    frontend_tools=[update_recipe],
+    initial_state={
+        "recipe": None,
+    }
+)
+
+
diff --git a/typescript-sdk/integrations/llamaindex/server-py/uv.lock b/typescript-sdk/integrations/llamaindex/server-py/uv.lock
diff --git a/typescript-sdk/integrations/llamaindex/src/index.ts b/typescript-sdk/integrations/llamaindex/src/index.ts
diff --git a/typescript-sdk/integrations/llamaindex/tsconfig.json b/typescript-sdk/integrations/llamaindex/tsconfig.json
diff --git a/typescript-sdk/integrations/llamaindex/tsup.config.ts b/typescript-sdk/integrations/llamaindex/tsup.config.ts
diff --git a/typescript-sdk/integrations/server-starter-all-features/server/python/example_server/agentic_chat.py b/typescript-sdk/integrations/server-starter-all-features/server/python/example_server/agentic_chat.py
diff --git a/typescript-sdk/pnpm-lock.yaml b/typescript-sdk/pnpm-lock.yaml

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+__pycache__`
	`2`	`+*egg-info`
	`3`	`+.venv`