chore: import upstream snapshot with attribution

rag-sql-router/.streamlit/config.toml

@@ -0,0 +1,2 @@
+[browser]
+gatherUsageStats = false

rag-sql-router/README.md

@@ -0,0 +1,80 @@
+# RAG with SQL Router
+
+We are developing a system that will guide you in creating a custom agent. This agent can query either your Vector DB index for RAG-based retrieval or a separate SQL query engine. 
+
+## 🔍 **The Critical Component: Response Validation**
+
+**While everyone is trying to build agents, no one tells you how to ensure their outputs are reliable.**
+
+**[Cleanlab Codex](https://help.cleanlab.ai/codex/)**, developed by researchers from MIT, offers a platform to evaluate and monitor any RAG or agentic app you're building. This system integrates Cleanlab Codex for automatic response validation, ensuring your AI outputs are trustworthy and continuously improving.
+
+### **Why Cleanlab Codex is Essential:**
+
+- **🔍 Automatic Detection**: Detects inaccurate/unhelpful responses from your AI automatically
+- **📈 Continuous Improvement**: Allows Subject Matter Experts to directly improve responses without engineering intervention  
+- **🎯 Trust Scoring**: Provides reliability metrics for every response
+- **🔄 Real-time Validation**: Validates queries and responses in real-time
+- **📊 Analytics**: Track improvement rates and response quality over time
+
+### **How It Works in This System:**
+
+1. **Query Processing**: Your queries are automatically validated by Cleanlab Codex
+2. **Response Validation**: AI responses are scored for reliability and accuracy
+3. **SME Intervention**: Subject Matter Experts can improve responses through the Codex interface
+4. **Continuous Learning**: The system learns from validated responses for future queries
+
+We use:
+
+- [Llama_Index](https://docs.llamaindex.ai/en/stable/) for orchestration
+- [Docling](https://docling-project.github.io/docling) for simplifying document processing
+- [Milvus](https://milvus.io/) to self-host a VectorDB
+- **[Cleanlab Codex](https://help.cleanlab.ai/codex/)** for **response validation and reliability assurance** ⭐
+- [OpenRouterAI](https://openrouter.ai/docs/quick-start) to access Alibaba's Qwen model
+
+> **💡 Key Insight**: While most tutorials focus on building agents, **[Cleanlab Codex](https://help.cleanlab.ai/codex/)** addresses the critical gap of ensuring those agents produce reliable, trustworthy outputs.
+
+## Set Up
+
+Follow these steps one by one:
+
+### Setup Milvus VectorDB
+
+Milvus provides an installation script to install it as a docker container.
+
+To install Milvus in Docker, you can use the following command:
+
+```bash
+curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh -o standalone_embed.sh
+
+bash standalone_embed.sh start
+```
+
+### Install Dependencies
+
+```bash
+uv sync
+```
+
+## Run the Notebook
+
+You can run the `notebook.ipynb` file to test the functionality of the code in a Jupyter Notebook environment. This notebook will help you understand routing, tool calling, and validating responses.
+
+## Run the Application
+
+To run the Streamlit app, use the following command:
+
+```bash
+streamlit run app.py
+```
+
+Open your browser and navigate to `http://localhost:8501` to access the app.
+
+## 📬 Stay Updated with Our Newsletter!
+
+**Get a FREE Data Science eBook** 📖 with 150+ essential lessons in Data Science when you subscribe to our newsletter! Stay in the loop with the latest tutorials, insights, and exclusive resources. [Subscribe now!](https://join.dailydoseofds.com)
+
+[![Daily Dose of Data Science Newsletter](https://github.com/patchy631/ai-engineering/blob/main/resources/join_ddods.png)](https://join.dailydoseofds.com)
+
+## Contribution
+
+Contributions are welcome! Feel free to fork this repository and submit pull requests with your improvements.

rag-sql-router/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "rag-sql-router"
+version = "0.1.0"
+description = "Text2SQL with RAG, developing hybrid agentic workflow"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "cleanlab-codex>=1.0.26",
+    "llama-index>=0.12.52",
+    "llama-index-core>=0.12.52.post1",
+    "llama-index-embeddings-huggingface>=0.5.5",
+    "llama-index-llms-openai>=0.4.7",
+    "llama-index-llms-openrouter>=0.3.2",
+    "llama-index-node-parser-docling>=0.3.2",
+    "llama-index-readers-docling>=0.3.3",
+    "llama-index-readers-file>=0.4.11",
+    "llama-index-vector-stores-milvus>=0.8.7",
+    "nest-asyncio>=1.6.0",
+    "pymilvus>=2.5.14",
+    "python-dotenv>=1.1.1",
+    "sqlalchemy>=2.0.42",
+    "streamlit>=1.47.1",
+    "torch>=2.7.1",
+    "pandas>=2.0.0",
+    "plotly>=5.0.0",
+]

rag-sql-router/tools.py

@@ -0,0 +1,275 @@
+# Required imports
+import os
+import uuid
+from sqlalchemy import create_engine
+from llama_index.core import (
+    VectorStoreIndex,
+    SimpleDirectoryReader,
+    SQLDatabase,
+    PromptTemplate,
+    StorageContext,
+)
+from llama_index.core.query_engine import NLSQLTableQueryEngine
+from llama_index.core.tools import QueryEngineTool, FunctionTool
+from llama_index.core.node_parser import MarkdownNodeParser
+from llama_index.readers.docling import DoclingReader
+from llama_index.vector_stores.milvus import MilvusVectorStore
+from cleanlab_codex.project import Project
+from cleanlab_codex.client import Client
+
+
+#####################################
+# Define Tools for Router Agent
+#####################################
+def create_codex_project():
+    """Create a Codex project for document validation."""
+    try:
+        # Check if CODEX_API_KEY is available
+        if not os.environ.get("CODEX_API_KEY"):
+            print(
+                "Warning: CODEX_API_KEY not found. Codex validation will be disabled."
+            )
+            return None, None
+
+        # Create a unique identifier for the project
+        project_id = str(uuid.uuid4())[:8]  # Using first 8 chars for readability
+
+        codex_client = Client()
+        project = codex_client.create_project(name=f"RAG + SQL Router {project_id}")
+        access_key = project.create_access_key("default key")
+        project = Project.from_access_key(access_key)
+        return project, project_id
+    except Exception as e:
+        print(f"Error creating Codex project: {e}")
+        return None, None
+
+
+# Global variables for reuse - these will persist across function calls
+docs_query_engine = None
+codex_project = None
+current_session_id = None
+current_project_id = None
+
+
+def get_or_create_codex_project(session_id):
+    """Get existing Codex project or create a new one for the session."""
+    global codex_project, current_session_id, current_project_id
+    
+    # If we have a project and it's for the same session, reuse it
+    if codex_project is not None and current_session_id == session_id:
+        print(f"Reusing existing Codex project for session {session_id}")
+        return codex_project
+    
+    # Create a new project for this session
+    print(f"Creating new Codex project for session {session_id}")
+    codex_project, project_id = create_codex_project()
+    current_session_id = session_id
+    current_project_id = project_id
+    
+    return codex_project
+
+
+def get_codex_project_info():
+    """Get information about the current Codex project for debugging."""
+    global codex_project, current_session_id, current_project_id
+    
+    if codex_project is None:
+        return {
+            "status": "No project created",
+            "session_id": current_session_id,
+            "project_id": None
+        }
+    
+    try:
+        # Get the actual project name using the stored project ID
+        if current_project_id:
+            project_name = f"RAG + SQL Router {current_project_id}"
+        else:
+            project_name = "RAG + SQL Router Project"
+            
+        return {
+            "status": "Active",
+            "session_id": current_session_id,
+            "project_id": "Available",
+            "project_name": project_name
+        }
+    except Exception as e:
+        return {
+            "status": f"Error getting info: {str(e)}",
+            "session_id": current_session_id,
+            "project_id": "Unknown"
+        }
+
+
+def setup_sql_tool(db_path="city_database.sqlite", table_name="city_stats"):
+    """Setup SQL query tool for querying city database."""
+    # Validate database exists
+    if not os.path.exists(db_path):
+        raise FileNotFoundError(f"Database file not found: {db_path}")
+
+    try:
+        engine = create_engine(f"sqlite:///{db_path}")
+        sql_database = SQLDatabase(engine)
+    except Exception as e:
+        print(f"Error setting up SQL database: {e}")
+        raise
+
+    # Create SQL query engine
+    sql_query_engine = NLSQLTableQueryEngine(
+        sql_database=sql_database,
+        tables=[table_name],
+    )
+
+    # Create tool for SQL querying
+    sql_tool = QueryEngineTool.from_defaults(
+        query_engine=sql_query_engine,
+        name="sql_tool",
+        description=(
+            "Useful for translating a natural language query into a SQL query over"
+            " a table containing: city_stats, containing the population/state of"
+            " each city located in the USA."
+        ),
+    )
+
+    # Return the SQL tool
+    return sql_tool
+
+
+def setup_document_tool(file_dir, session_id=None, milvus_uri="http://localhost:19530"):
+    """Setup document query tool from uploaded documents with Codex validation."""
+    global docs_query_engine
+
+    # Create a reader and load the data
+    reader, node_parser = DoclingReader(), MarkdownNodeParser()
+    loader = SimpleDirectoryReader(
+        input_dir=file_dir,
+        file_extractor={
+            ".pdf": reader,
+            ".docx": reader,
+            ".pptx": reader,
+            ".txt": reader,
+        },
+    )
+    docs = loader.load_data()
+
+    # Creating a vector index over loaded data
+    unique_collection_id = uuid.uuid4().hex
+    collection_name = f"rag_with_sql_{unique_collection_id}"
+    vector_store = MilvusVectorStore(uri=milvus_uri, dim=384, overwrite=True, collection_name=collection_name)
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    vector_index = VectorStoreIndex.from_documents(
+        docs,
+        show_progress=True,
+        transformations=[node_parser],
+        storage_context=storage_context,
+    )
+
+    # Custom prompt template
+    template = (
+        "You are a meticulous and accurate document analyst. Your task is to answer the user's question based exclusively on the provided context. "
+        "Follow these rules strictly:\n"
+        "1. Your entire response must be grounded in the facts provided in the 'Context' section. Do not use any prior knowledge.\n"
+        "2. If multiple parts of the context are relevant, synthesize them into a single, coherent answer.\n"
+        "3. If the context does not contain the information needed to answer the question, you must state only: 'The provided context does not contain enough information to answer this question.'\n"
+        "-----------------------------------------\n"
+        "Context: {context_str}\n"
+        "-----------------------------------------\n"
+        "Question: {query_str}\n\n"
+        "Answer:"
+    )
+    qa_template = PromptTemplate(template)
+
+    # Create a query engine for the vector index
+    docs_query_engine = vector_index.as_query_engine(
+        text_qa_template=qa_template, similarity_top_k=3
+    )
+
+    # Get or create Codex project for this session
+    codex_project = get_or_create_codex_project(session_id)
+
+    # Define the document query function with Codex validation
+    def document_query_tool(query: str):
+        """Query documents with Codex validation for enhanced accuracy."""
+        # Step 1: Query the engine
+        response_obj = docs_query_engine.query(query)
+        initial_response = str(response_obj)
+
+        # Step 2: Gather source context
+        context = response_obj.source_nodes
+        context_str = "\n".join([n.node.text for n in context])
+
+        # Step 3: Prepare prompt for Codex validation
+        prompt_template = (
+            "You are a meticulous and accurate document analyst. Your task is to answer the user's question based exclusively on the provided context. "
+            "Follow these rules strictly:\n"
+            "1. Your entire response must be grounded in the facts provided in the 'Context' section. Do not use any prior knowledge.\n"
+            "2. If multiple parts of the context are relevant, synthesize them into a single, coherent answer.\n"
+            "3. If the context does not contain the information needed to answer the question, you must state only: 'The provided context does not contain enough information to answer this question.'\n"
+            "-----------------------------------------\n"
+            "Context: {context}\n"
+            "-----------------------------------------\n"
+            "Question: {query}\n\n"
+            "Answer:"
+        )
+        user_prompt = prompt_template.format(context=context_str, query=query)
+        messages = [{"role": "user", "content": user_prompt}]
+
+        # Step 4: Validate with Codex (if available)
+        if codex_project:
+            try:
+                print(f"Validating query with Codex: '{query[:50]}...'")
+                result = codex_project.validate(
+                    messages=messages,
+                    query=query,
+                    context=context_str,
+                    response=initial_response,
+                )
+                print("Codex validation completed successfully")
+
+                # Step 5: Final response selection
+                fallback_response = "I'm sorry, I couldn't find an answer — can I help with something else?"
+                final_response = (
+                    result.expert_answer
+                    if result.expert_answer and result.escalated_to_sme
+                    else (
+                        fallback_response
+                        if result.should_guardrail
+                        else initial_response
+                    )
+                )
+                trust_score = result.model_dump()["eval_scores"]["trustworthiness"]["score"]
+
+                # Return a dictionary to avoid tuple handling issues
+                return {
+                    "response": str(final_response),
+                    "trust_score": float(trust_score)
+                }
+            except Exception as e:
+                # If Codex validation fails, return the initial response
+                print(f"Codex validation failed: {e}")
+                return {
+                    "response": str(initial_response),
+                    "trust_score": None
+                }
+        else:
+            # If Codex is not available, return the initial response
+            print("Codex not available, using basic RAG response")
+            return {
+                "response": str(initial_response),
+                "trust_score": None
+            }
+
+    # Create tool for document querying using FunctionTool
+    docs_tool = FunctionTool.from_defaults(
+        document_query_tool,
+        name="document_tool",
+        description=(
+            "Useful for answering a natural language question by performing a semantic search over "
+            "a collection of documents. These documents may contain general knowledge, reports, "
+            "or domain-specific content. Returns the most relevant passages or synthesized answers. "
+            "If the user query does not relate to US city statistics (population and state), use this document search tool."
+        ),
+    )
+
+    # Return the document tool
+    return docs_tool

rag-sql-router/workflow.py

@@ -0,0 +1,223 @@
+# Required imports
+import asyncio
+from typing import Dict, List, Any, Optional
+from llama_index.core import Settings
+from llama_index.core.tools import BaseTool
+from llama_index.core.llms import ChatMessage
+from llama_index.core.llms.llm import ToolSelection, LLM
+from llama_index.core.workflow import (
+    Workflow,
+    Event,
+    StartEvent,
+    StopEvent,
+    step,
+    Context,
+)
+
+
+#####################################
+# Define Router Agent Workflow
+#####################################
+class InputEvent(Event):
+    """Input event."""
+
+
+class GatherToolsEvent(Event):
+    """Gather Tools Event"""
+
+    tool_calls: Any
+
+
+class ToolCallEvent(Event):
+    """Tool Call event"""
+
+    tool_call: ToolSelection
+
+
+class ToolCallEventResult(Event):
+    """Tool call event result."""
+
+    msg: ChatMessage
+
+
+class RouterOutputAgentWorkflow(Workflow):
+    """Custom router output agent workflow."""
+
+    def __init__(
+        self,
+        tools: List[BaseTool],
+        timeout: Optional[float] = 10.0,
+        disable_validation: bool = False,
+        verbose: bool = False,
+        llm: Optional[LLM] = None,
+        chat_history: Optional[List[ChatMessage]] = None,
+    ):
+        """Constructor."""
+        super().__init__(
+            timeout=timeout, disable_validation=disable_validation, verbose=verbose
+        )
+        self.tools: List[BaseTool] = tools
+        self.tools_dict: Optional[Dict[str, BaseTool]] = {
+            tool.metadata.name: tool for tool in self.tools
+        }
+        # Use provided LLM or fall back to Settings.llm
+        self.llm: LLM = llm or Settings.llm
+        if self.llm is None:
+            raise ValueError("No LLM provided and Settings.llm is not initialized")
+        self.chat_history: List[ChatMessage] = chat_history or []
+
+    def reset(self) -> None:
+        """Resets Chat History"""
+        self.chat_history = []
+
+    @step()
+    async def prepare_chat(self, ev: StartEvent) -> InputEvent:
+        message = ev.get("message")
+        if message is None:
+            raise ValueError("'message' field is required.")
+
+        # Add message to chat history
+        chat_history = self.chat_history
+        chat_history.append(ChatMessage(role="user", content=message))
+        return InputEvent()
+
+    @step()
+    async def chat(self, ev: InputEvent) -> GatherToolsEvent | StopEvent:
+        """Appends msg to chat history, then gets tool calls."""
+        try:
+            # Put message into LLM with tools included
+            chat_res = await self.llm.achat_with_tools(
+                self.tools,
+                chat_history=self.chat_history,
+                verbose=self._verbose,
+                allow_parallel_tool_calls=True,
+            )
+            tool_calls = self.llm.get_tool_calls_from_response(
+                chat_res, error_on_no_tool_call=False
+            )
+
+            ai_message = chat_res.message
+            self.chat_history.append(ai_message)
+            if self._verbose:
+                print(f"Chat message: {ai_message.content}")
+
+            # No tool calls, return chat message.
+            if not tool_calls:
+                return StopEvent(result=ai_message.content)
+
+            return GatherToolsEvent(tool_calls=tool_calls)
+        except asyncio.CancelledError:
+            print("Chat operation was cancelled")
+            return StopEvent(result="The operation was cancelled. Please try again.")
+        except Exception as e:
+            error_msg = f"Error during chat: {str(e)}"
+            print(error_msg)
+            return StopEvent(
+                result="I'm sorry, I encountered an issue processing your request. Could you try asking in a different way?"
+            )
+
+    @step(pass_context=True)
+    async def dispatch_calls(self, ctx: Context, ev: GatherToolsEvent) -> ToolCallEvent:
+        """Dispatches calls."""
+        tool_calls = ev.tool_calls
+        await ctx.set("num_tool_calls", len(tool_calls))
+
+        # Trigger tool call events
+        for tool_call in tool_calls:
+            ctx.send_event(ToolCallEvent(tool_call=tool_call))
+
+        return None
+
+    @step()
+    async def call_tool(self, ev: ToolCallEvent) -> ToolCallEventResult:
+        """Calls tool."""
+        try:
+            tool_call = ev.tool_call
+            # Get tool ID and function call
+            id_ = tool_call.tool_id
+
+            if self._verbose:
+                print(
+                    f"Calling function {tool_call.tool_name} with msg {tool_call.tool_kwargs}"
+                )
+
+            # Call function and put result into a chat message
+            tool = self.tools_dict[tool_call.tool_name]
+            output = await tool.acall(**tool_call.tool_kwargs)
+            
+            # Check if output is a dictionary (response, trust_score) for document tool
+            if isinstance(output, dict) and "response" in output:
+                response = output.get("response", "")
+                trust_score = output.get("trust_score")
+                # Ensure response is a string
+                content = str(response) if response is not None else ""
+                # Store additional metadata
+                additional_kwargs = {
+                    "tool_call_id": id_, 
+                    "name": tool_call.tool_name,
+                    "trust_score": trust_score,
+                    "tool_used": tool_call.tool_name
+                }
+                if self._verbose:
+                    print(f"Tool {tool_call.tool_name} returned dict: response='{content}', trust_score={trust_score}")
+            else:
+                content = str(output) if output is not None else ""
+                additional_kwargs = {
+                    "tool_call_id": id_, 
+                    "name": tool_call.tool_name,
+                    "tool_used": tool_call.tool_name
+                }
+                if self._verbose:
+                    print(f"Tool {tool_call.tool_name} returned: '{content}'")
+            
+            msg = ChatMessage(
+                name=tool_call.tool_name,
+                content=content,
+                role="tool",
+                additional_kwargs=additional_kwargs,
+            )
+
+            return ToolCallEventResult(msg=msg)
+            
+        except asyncio.CancelledError:
+            print(f"Tool call {tool_call.tool_name} was cancelled")
+            # Return a dummy result to avoid workflow breakdown
+            msg = ChatMessage(
+                name=tool_call.tool_name,
+                content="Tool execution was cancelled",
+                role="tool",
+                additional_kwargs={"tool_call_id": id_, "name": tool_call.tool_name, "tool_used": tool_call.tool_name},
+            )
+            return ToolCallEventResult(msg=msg)
+        except Exception as e:
+            print(f"Error in tool call {tool_call.tool_name}: {str(e)}")
+            # Return an error result instead of failing
+            msg = ChatMessage(
+                name=tool_call.tool_name,
+                content=f"Error executing tool: {str(e)}",
+                role="tool",
+                additional_kwargs={"tool_call_id": id_, "name": tool_call.tool_name, "tool_used": tool_call.tool_name},
+            )
+            return ToolCallEventResult(msg=msg)
+
+    @step(pass_context=True)
+    async def gather(self, ctx: Context, ev: ToolCallEventResult) -> StopEvent | None:
+        """Gathers tool calls."""
+        try:
+            # Wait for all tool call events to finish.
+            tool_events = ctx.collect_events(
+                ev, [ToolCallEventResult] * await ctx.get("num_tool_calls")
+            )
+            if not tool_events:
+                return None
+
+            for tool_event in tool_events:
+                # Append tool call chat messages to history
+                self.chat_history.append(tool_event.msg)
+
+            # After all tool calls finish, pass input event back, restart agent loop
+            return InputEvent()
+        except Exception as e:
+            print(f"Error in gather step: {str(e)}")
+            # Return a stop event instead of continuing the loop if there's an error
+            return StopEvent(result="I encountered an issue processing the tool responses. Please try again.")

rag-sql-router/workflow_all_flows.html

@@ -0,0 +1,155 @@
+<html>
+    <head>
+        <meta charset="utf-8">
+        
+            <script src="lib/bindings/utils.js"></script>
+            <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/dist/vis-network.min.css" integrity="sha512-WgxfT5LWjfszlPHXRmBWHkV2eceiWTOBvrKCNbdgDYTHrT2AeLCGbF4sZlZw3UMN3WtL0tGUoIAKsu8mllg/XA==" crossorigin="anonymous" referrerpolicy="no-referrer" />
+            <script src="https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/vis-network.min.js" integrity="sha512-LnvoEWDFrqGHlHmDD2101OrLcbsfkrzoSpvtSQtxK3RMnRV0eOkhhBN2dXHKRrUU8p2DGRTk35n4O8nWSVe1mQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+            
+        
+<center>
+<h1></h1>
+</center>
+
+<!-- <link rel="stylesheet" href="../node_modules/vis/dist/vis.min.css" type="text/css" />
+<script type="text/javascript" src="../node_modules/vis/dist/vis.js"> </script>-->
+        <link
+          href="https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/css/bootstrap.min.css"
+          rel="stylesheet"
+          integrity="sha384-eOJMYsd53ii+scO/bJGFsiCZc+5NDVN2yr8+0RDqr0Ql0h+rP48ckxlpbzKgwra6"
+          crossorigin="anonymous"
+        />
+        <script
+          src="https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/js/bootstrap.bundle.min.js"
+          integrity="sha384-JEW9xMcG8R+pH31jmWH6WWP0WintQrMb4s7ZOdauHnUtxwoG2vI5DkLtS3qm9Ekf"
+          crossorigin="anonymous"
+        ></script>
+
+
+        <center>
+          <h1></h1>
+        </center>
+        <style type="text/css">
+
+             #mynetwork {
+                 width: 100%;
+                 height: 750px;
+                 background-color: #ffffff;
+                 border: 1px solid lightgray;
+                 position: relative;
+                 float: left;
+             }
+
+             
+
+             
+
+             
+        </style>
+    </head>
+
+
+    <body>
+        <div class="card" style="width: 100%">
+            
+            
+            <div id="mynetwork" class="card-body"></div>
+        </div>
+
+        
+        
+
+        <script type="text/javascript">
+
+              // initialize global variables.
+              var edges;
+              var nodes;
+              var allNodes;
+              var allEdges;
+              var nodeColors;
+              var originalNodes;
+              var network;
+              var container;
+              var options, data;
+              var filter = {
+                  item : '',
+                  property : '',
+                  value : []
+              };
+
+              
+
+              
+
+              // This method is responsible for drawing the graph, returns the drawn network
+              function drawGraph() {
+                  var container = document.getElementById('mynetwork');
+
+                  
+
+                  // parsing and collecting nodes and edges from the python
+                  nodes = new vis.DataSet([{"color": "#ADD8E6", "id": "_done", "label": "_done", "shape": "box"}, {"color": "#FFA07A", "id": "StopEvent", "label": "StopEvent", "shape": "ellipse"}, {"color": "#ADD8E6", "id": "call_tool", "label": "call_tool", "shape": "box"}, {"color": "#90EE90", "id": "ToolCallEvent", "label": "ToolCallEvent", "shape": "ellipse"}, {"color": "#90EE90", "id": "ToolCallEventResult", "label": "ToolCallEventResult", "shape": "ellipse"}, {"color": "#ADD8E6", "id": "chat", "label": "chat", "shape": "box"}, {"color": "#90EE90", "id": "InputEvent", "label": "InputEvent", "shape": "ellipse"}, {"color": "#90EE90", "id": "GatherToolsEvent", "label": "GatherToolsEvent", "shape": "ellipse"}, {"color": "#ADD8E6", "id": "dispatch_calls", "label": "dispatch_calls", "shape": "box"}, {"color": "#ADD8E6", "id": "gather", "label": "gather", "shape": "box"}, {"color": "#ADD8E6", "id": "prepare_chat", "label": "prepare_chat", "shape": "box"}, {"color": "#E27AFF", "id": "StartEvent", "label": "StartEvent", "shape": "ellipse"}]);
+                  edges = new vis.DataSet([{"arrows": "to", "from": "StopEvent", "to": "_done"}, {"arrows": "to", "from": "call_tool", "to": "ToolCallEventResult"}, {"arrows": "to", "from": "ToolCallEvent", "to": "call_tool"}, {"arrows": "to", "from": "chat", "to": "GatherToolsEvent"}, {"arrows": "to", "from": "chat", "to": "StopEvent"}, {"arrows": "to", "from": "InputEvent", "to": "chat"}, {"arrows": "to", "from": "dispatch_calls", "to": "ToolCallEvent"}, {"arrows": "to", "from": "GatherToolsEvent", "to": "dispatch_calls"}, {"arrows": "to", "from": "gather", "to": "StopEvent"}, {"arrows": "to", "from": "ToolCallEventResult", "to": "gather"}, {"arrows": "to", "from": "prepare_chat", "to": "InputEvent"}, {"arrows": "to", "from": "StartEvent", "to": "prepare_chat"}]);
+
+                  nodeColors = {};
+                  allNodes = nodes.get({ returnType: "Object" });
+                  for (nodeId in allNodes) {
+                    nodeColors[nodeId] = allNodes[nodeId].color;
+                  }
+                  allEdges = edges.get({ returnType: "Object" });
+                  // adding nodes and edges to the graph
+                  data = {nodes: nodes, edges: edges};
+
+                  var options = {
+    "configure": {
+        "enabled": false
+    },
+    "edges": {
+        "color": {
+            "inherit": true
+        },
+        "smooth": {
+            "enabled": true,
+            "type": "dynamic"
+        }
+    },
+    "interaction": {
+        "dragNodes": true,
+        "hideEdgesOnDrag": false,
+        "hideNodesOnDrag": false
+    },
+    "physics": {
+        "enabled": true,
+        "stabilization": {
+            "enabled": true,
+            "fit": true,
+            "iterations": 1000,
+            "onlyDynamicEdges": false,
+            "updateInterval": 50
+        }
+    }
+};
+
+                  
+
+
+                  
+
+                  network = new vis.Network(container, data, options);
+
+                  
+
+                  
+
+                  
+
+
+                  
+
+                  return network;
+
+              }
+              drawGraph();
+        </script>
+    </body>
+</html>