pixie-examples 0.1.0__tar.gz → 0.1.1.dev5__tar.gz

pixie_examples-0.1.1.dev5/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Finto Technologies
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pixie_examples-0.1.1.dev5/PKG-INFO

@@ -0,0 +1,118 @@
+Metadata-Version: 2.4
+Name: pixie-examples
+Version: 0.1.1.dev5
+Summary: examples for using Pixie
+License: MIT
+License-File: LICENSE
+Author: Yiou Li
+Author-email: yol@gopixie.ai
+Requires-Python: >=3.11,<3.14
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: beautifulsoup4 (>=4.14.3,<5.0.0)
+Requires-Dist: langchain (>=1.2.3,<2.0.0)
+Requires-Dist: langchain-openai (>=1.1.7,<2.0.0)
+Requires-Dist: langchain-text-splitters (>=1.1.0,<2.0.0)
+Requires-Dist: langgraph (>=1.0.5,<2.0.0)
+Requires-Dist: lxml (>=6.0.2,<7.0.0)
+Requires-Dist: numpy (>=2.4.1,<3.0.0)
+Requires-Dist: openai-agents (>=0.6.5,<0.7.0)
+Requires-Dist: openinference-instrumentation-crewai (>=0.1.17,<0.2.0)
+Requires-Dist: openinference-instrumentation-dspy (>=0.1.33,<0.2.0)
+Requires-Dist: openinference-instrumentation-google-adk (>=0.1.8,<0.2.0)
+Requires-Dist: openinference-instrumentation-openai-agents (>=1.4.0,<2.0.0)
+Requires-Dist: pixie-sdk (>=0.1.1.dev13,<0.2.0)
+Requires-Dist: pydantic (>=2.7.4,<3.0.0)
+Requires-Dist: pydantic-ai-slim (>=1.39.0,<2.0.0)
+Requires-Dist: pymarkdownlnt (>=0.9.34,<0.10.0)
+Requires-Dist: requests (>=2.32.5,<3.0.0)
+Requires-Dist: sqlalchemy (>=2.0.45,<3.0.0)
+Project-URL: Changelog, https://github.com/yiouli/pixie-examples/commits/main/
+Project-URL: Documentation, https://yiouli.github.io/pixie-sdk-py/
+Project-URL: Homepage, https://gopixie.ai/?url=https%3A%2F%2Fdemo.yiouli.us%2Fgraphql
+Project-URL: Issues, https://github.com/yiouli/pixie-examples/issues
+Project-URL: Repository, https://github.com/yiouli/pixie-examples
+Description-Content-Type: text/markdown
+
+# Pixie Examples
+
+[![MIT License](https://img.shields.io/badge/License-MIT-red.svg?style=flat-square)](https://opensource.org/licenses/MIT)
+[![Python Version](https://img.shields.io/badge/python-3.11%2B-blue?style=flat-square)](https://www.python.org/downloads/)
+[![Discord](https://img.shields.io/discord/1459772566528069715?style=flat-square&logo=Discord&logoColor=white&label=Discord&color=%23434EE4)](https://discord.gg/YMNYu6Z3)
+
+This repository contains a collection of example applications integrated with [**Pixie SDK**](https://github.com/yiouli/pixie-sdk-py) for interactive debugging.
+
+## Get Started
+
+> You can play with the demo site [here](https://gopixie.ai/?url=https://demo.yiouli.us/graphql) withou any setup.
+
+### 1. Setup
+
+Clone this repository:
+
+Install `pixie-examples` python package:
+
+```bash
+pip install pixie-examples
+```
+
+Create a `.env` file with your API keys:
+
+```ini
+# .env
+OPENAI_API_KEY=...
+# Add other API keys as needed for specific examples
+```
+
+Start the Pixie server:
+
+```bash
+pixie
+```
+
+### 2. Debug with Web UI
+
+Visit [gopixie.ai](https://gopixie.ai) to interact with and debug your applications through the web interface.
+
+## Important Links
+
+- [**Pixie SDK**](https://github.com/yiouli/pixie-sdk-py)
+- [**Documentation**](https://yiouli.github.io/pixie-sdk-py/)
+- [**Discord**](https://discord.gg/YMNYu6Z3)
+
+## Examples Catelog
+
+### Quickstart
+
+- **Basic Example**: Simple hello world application to get started with Pixie SDK
+
+### Pydantic AI Examples
+
+- **Bank Support**: Multi-turn chatbot for banking customer support
+- **Flight Booking**: Multi-agent system for flight booking
+- **Question Graph**: Graph-based question answering system
+- **SQL Generation**: Multi-step workflow for generating SQL queries
+- **Structured Output**: Examples of structured data handling
+
+### OpenAI Agents SDK Examples
+
+- **Customer Service**: Multi-agent customer service system
+- **Financial Research Agent**: Multi-step financial research workflow
+- **LLM-as-a-Judge**: Evaluation and judging patterns
+- **Routing**: Agent routing and handoffs
+
+### LangChain Examples
+
+- **Basic Agent**: Simple LangChain agent integration
+- **Customer Support**: Customer support chatbot
+- **Personal Assistant**: Multi-agent personal assistant
+- **SQL Agent**: SQL query generation with LangChain
+
+### LangGraph Examples
+
+- **RAG System**: Retrieval-augmented generation with LangGraph
+- **SQL Agent**: SQL agent built with LangGraph state machines
+

pixie_examples-0.1.1.dev5/README.md

@@ -0,0 +1,78 @@
+# Pixie Examples
+
+[![MIT License](https://img.shields.io/badge/License-MIT-red.svg?style=flat-square)](https://opensource.org/licenses/MIT)
+[![Python Version](https://img.shields.io/badge/python-3.11%2B-blue?style=flat-square)](https://www.python.org/downloads/)
+[![Discord](https://img.shields.io/discord/1459772566528069715?style=flat-square&logo=Discord&logoColor=white&label=Discord&color=%23434EE4)](https://discord.gg/YMNYu6Z3)
+
+This repository contains a collection of example applications integrated with [**Pixie SDK**](https://github.com/yiouli/pixie-sdk-py) for interactive debugging.
+
+## Get Started
+
+> You can play with the demo site [here](https://gopixie.ai/?url=https://demo.yiouli.us/graphql) withou any setup.
+
+### 1. Setup
+
+Clone this repository:
+
+Install `pixie-examples` python package:
+
+```bash
+pip install pixie-examples
+```
+
+Create a `.env` file with your API keys:
+
+```ini
+# .env
+OPENAI_API_KEY=...
+# Add other API keys as needed for specific examples
+```
+
+Start the Pixie server:
+
+```bash
+pixie
+```
+
+### 2. Debug with Web UI
+
+Visit [gopixie.ai](https://gopixie.ai) to interact with and debug your applications through the web interface.
+
+## Important Links
+
+- [**Pixie SDK**](https://github.com/yiouli/pixie-sdk-py)
+- [**Documentation**](https://yiouli.github.io/pixie-sdk-py/)
+- [**Discord**](https://discord.gg/YMNYu6Z3)
+
+## Examples Catelog
+
+### Quickstart
+
+- **Basic Example**: Simple hello world application to get started with Pixie SDK
+
+### Pydantic AI Examples
+
+- **Bank Support**: Multi-turn chatbot for banking customer support
+- **Flight Booking**: Multi-agent system for flight booking
+- **Question Graph**: Graph-based question answering system
+- **SQL Generation**: Multi-step workflow for generating SQL queries
+- **Structured Output**: Examples of structured data handling
+
+### OpenAI Agents SDK Examples
+
+- **Customer Service**: Multi-agent customer service system
+- **Financial Research Agent**: Multi-step financial research workflow
+- **LLM-as-a-Judge**: Evaluation and judging patterns
+- **Routing**: Agent routing and handoffs
+
+### LangChain Examples
+
+- **Basic Agent**: Simple LangChain agent integration
+- **Customer Support**: Customer support chatbot
+- **Personal Assistant**: Multi-agent personal assistant
+- **SQL Agent**: SQL query generation with LangChain
+
+### LangGraph Examples
+
+- **RAG System**: Retrieval-augmented generation with LangGraph
+- **SQL Agent**: SQL agent built with LangGraph state machines

{pixie_examples-0.1.0 → pixie_examples-0.1.1.dev5}/examples/langchain/basic_agent.py

@@ -5,10 +5,14 @@ This example demonstrates a simple agent that can answer questions and call tool
 Based on: https://docs.langchain.com/oss/python/langchain/quickstart
 """
 
-from types import NoneType
 from langchain.agents import create_agent
 from langchain.chat_models import init_chat_model
-from pixie import pixie_app, PixieGenerator, UserInputRequirement
+
+from langfuse.langchain import CallbackHandler
+import pixie
+
+
+langfuse_handler = CallbackHandler()
 
 
 def get_weather(city: str) -> str:
@@ -16,8 +20,8 @@ def get_weather(city: str) -> str:
     return f"It's always sunny in {city}!"
 
 
-@pixie_app
-async def basic_weather_agent(query: str) -> str:
+@pixie.app
+async def langchain_basic_weather_agent(query: str) -> str:
     """A simple weather agent that can answer questions using tools.
 
     Args:
@@ -37,21 +41,21 @@ async def basic_weather_agent(query: str) -> str:
     )
 
     # Run the agent
-    result = agent.invoke({"messages": [{"role": "user", "content": query}]})
+    result = agent.invoke(
+        {"messages": [{"role": "user", "content": query}]},
+        config={"callbacks": [langfuse_handler]},
+    )
 
     # Return the final response
     return result["messages"][-1].content
 
 
-@pixie_app
-async def interactive_weather_agent(_: NoneType) -> PixieGenerator[str, str]:
+@pixie.app
+async def langchain_interactive_weather_agent() -> pixie.PixieGenerator[str, str]:
     """An interactive weather chatbot that maintains conversation.
 
     This agent can have multi-turn conversations with the user.
 
-    Args:
-        _: No initial input required
-
     Yields:
         AI responses to user questions
     """
@@ -73,7 +77,7 @@ async def interactive_weather_agent(_: NoneType) -> PixieGenerator[str, str]:
 
     while True:
         # Get user input
-        user_query = yield UserInputRequirement(str)
+        user_query = yield pixie.InputRequired(str)
 
         # Check for exit commands
         if user_query.lower() in {"exit", "quit", "bye", "goodbye"}:
@@ -84,7 +88,9 @@ async def interactive_weather_agent(_: NoneType) -> PixieGenerator[str, str]:
         messages.append({"role": "user", "content": user_query})
 
         # Run agent with full conversation history
-        result = agent.invoke({"messages": messages})
+        result = agent.invoke(
+            {"messages": messages}, config={"callbacks": [langfuse_handler]}
+        )
 
         # Update history with AI response
         messages = result["messages"]

{pixie_examples-0.1.0 → pixie_examples-0.1.1.dev5}/examples/langchain/customer_support.py

@@ -7,7 +7,6 @@ as it moves through different states of a workflow.
 Based on: https://docs.langchain.com/oss/python/langchain/multi-agent/handoffs-customer-support
 """
 
-from types import NoneType
 from typing import Literal, NotRequired
 from langchain.agents import create_agent, AgentState
 from langchain.chat_models import init_chat_model
@@ -16,7 +15,12 @@ from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResp
 from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.types import Command
 from typing import Callable
-from pixie import pixie_app, PixieGenerator, UserInputRequirement
+
+from langfuse.langchain import CallbackHandler
+import pixie
+
+
+langfuse_handler = CallbackHandler()
 
 
 # Define the possible workflow steps
@@ -168,8 +172,8 @@ def apply_step_config(
     return handler(request)
 
 
-@pixie_app
-async def customer_support(_: NoneType) -> PixieGenerator[str, str]:
+@pixie.app
+async def langchain_customer_support() -> pixie.PixieGenerator[str, str]:
     """Customer support agent with state machine workflow.
 
     The agent progresses through three stages:
@@ -177,9 +181,6 @@ async def customer_support(_: NoneType) -> PixieGenerator[str, str]:
     2. Issue classification (hardware/software)
     3. Resolution (solution or escalation)
 
-    Args:
-        _: No initial input required
-
     Yields:
         AI responses guiding the support workflow
     """
@@ -208,11 +209,11 @@ async def customer_support(_: NoneType) -> PixieGenerator[str, str]:
 
     # Initialize conversation
     thread_id = "support_thread"
-    config = {"configurable": {"thread_id": thread_id}}
+    config = {"configurable": {"thread_id": thread_id}, "callbacks": [langfuse_handler]}
 
     while True:
         # Get user input
-        user_message = yield UserInputRequirement(str)
+        user_message = yield pixie.InputRequired(str)
 
         # Check for exit
         if user_message.lower() in {"exit", "quit", "bye"}:

{pixie_examples-0.1.0 → pixie_examples-0.1.1.dev5}/examples/langchain/personal_assistant.py

@@ -7,12 +7,16 @@ coordinates specialized worker agents (calendar and email agents).
 Based on: https://docs.langchain.com/oss/python/langchain/multi-agent/subagents-personal-assistant
 """
 
-from types import NoneType
 from langchain.agents import create_agent
 from langchain.chat_models import init_chat_model
 from langchain.tools import tool
 from langgraph.checkpoint.memory import InMemorySaver
-from pixie import pixie_app, PixieGenerator, UserInputRequirement
+
+from langfuse.langchain import CallbackHandler
+import pixie
+
+
+langfuse_handler = CallbackHandler()
 
 
 # Define calendar tools (stubs for demonstration)
@@ -68,17 +72,14 @@ SUPERVISOR_PROMPT = (
 )
 
 
-@pixie_app
-async def personal_assistant(_: NoneType) -> PixieGenerator[str, str]:
+@pixie.app
+async def langchain_personal_assistant() -> pixie.PixieGenerator[str, str]:
     """Multi-agent personal assistant with calendar and email subagents.
 
     The supervisor coordinates specialized worker agents:
     - Calendar agent: handles scheduling and availability
     - Email agent: manages communication and drafts
 
-    Args:
-        _: No initial input required
-
     Yields:
         AI responses to user requests
     """
@@ -108,7 +109,8 @@ async def personal_assistant(_: NoneType) -> PixieGenerator[str, str]:
         Handles date/time parsing, availability checking, and event creation.
         """
         result = calendar_agent.invoke(
-            {"messages": [{"role": "user", "content": request}]}
+            {"messages": [{"role": "user", "content": request}]},
+            config={"callbacks": [langfuse_handler]},
         )
         return result["messages"][-1].content
 
@@ -120,7 +122,8 @@ async def personal_assistant(_: NoneType) -> PixieGenerator[str, str]:
         communication. Handles recipient extraction, subject generation, and email composition.
         """
         result = email_agent.invoke(
-            {"messages": [{"role": "user", "content": request}]}
+            {"messages": [{"role": "user", "content": request}]},
+            config={"callbacks": [langfuse_handler]},
         )
         return result["messages"][-1].content
 
@@ -140,11 +143,11 @@ async def personal_assistant(_: NoneType) -> PixieGenerator[str, str]:
 
     # Initialize conversation
     thread_id = "personal_assistant_thread"
-    config = {"configurable": {"thread_id": thread_id}}
+    config = {"configurable": {"thread_id": thread_id}, "callbacks": [langfuse_handler]}
 
     while True:
         # Get user request
-        user_request = yield UserInputRequirement(str)
+        user_request = yield pixie.InputRequired(str)
 
         # Check for exit
         if user_request.lower() in {"exit", "quit", "bye", "goodbye"}:

{pixie_examples-0.1.0 → pixie_examples-0.1.1.dev5}/examples/langchain/sql_agent.py

@@ -17,13 +17,16 @@ queries. Make sure database connection permissions are scoped as narrowly as pos
 
 import pathlib
 import requests
-from types import NoneType
 from langchain.agents import create_agent
 from langchain.chat_models import init_chat_model
-from langchain_community.utilities import SQLDatabase
-from langchain_community.agent_toolkits import SQLDatabaseToolkit
 from langgraph.checkpoint.memory import InMemorySaver
-from pixie import pixie_app, PixieGenerator, UserInputRequirement
+
+from langfuse.langchain import CallbackHandler
+import pixie
+from ..sql_utils import SQLDatabase, SQLDatabaseToolkit
+
+
+langfuse_handler = CallbackHandler()
 
 
 # System prompt for SQL agent
@@ -72,8 +75,8 @@ def setup_database():
     return SQLDatabase.from_uri("sqlite:///Chinook.db")
 
 
-@pixie_app
-async def sql_query_agent(question: str) -> str:
+@pixie.app
+async def langchain_sql_query_agent(question: str) -> str:
     """SQL database query agent that can answer questions about the Chinook database.
 
     The Chinook database represents a digital media store with tables for artists,
@@ -102,21 +105,21 @@ async def sql_query_agent(question: str) -> str:
     agent = create_agent(model, tools, system_prompt=system_prompt)
 
     # Run the agent
-    result = agent.invoke({"messages": [{"role": "user", "content": question}]})
+    result = agent.invoke(
+        {"messages": [{"role": "user", "content": question}]},
+        config={"callbacks": [langfuse_handler]},
+    )
 
     # Return the final answer
     return result["messages"][-1].content
 
 
-@pixie_app
-async def interactive_sql_agent(_: NoneType) -> PixieGenerator[str, str]:
+@pixie.app
+async def langchain_interactive_sql_agent() -> pixie.PixieGenerator[str, str]:
     """Interactive SQL database query agent with multi-turn conversation.
 
     This agent maintains conversation history and can handle follow-up questions.
 
-    Args:
-        _: No initial input required
-
     Yields:
         AI responses to database queries
     """
@@ -153,11 +156,11 @@ Ask me any question about the data!"""
 
     # Initialize conversation
     thread_id = "sql_thread"
-    config = {"configurable": {"thread_id": thread_id}}
+    config = {"configurable": {"thread_id": thread_id}, "callbacks": [langfuse_handler]}
 
     while True:
         # Get user question
-        user_question = yield UserInputRequirement(str)
+        user_question = yield pixie.InputRequired(str)
 
         # Check for exit
         if user_question.lower() in {"exit", "quit", "bye"}:

{pixie_examples-0.1.0/examples/langchain → pixie_examples-0.1.1.dev5/examples/langgraph}/langgraph_rag.py

@@ -15,13 +15,36 @@ from typing import Literal
 from langchain.chat_models import init_chat_model
 from langchain.tools import tool
 from langchain.messages import HumanMessage
-from langchain_community.document_loaders import WebBaseLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.vectorstores import InMemoryVectorStore
+from langchain_core.documents import Document
 from langchain_openai import OpenAIEmbeddings
 from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
-from pixie import pixie_app
+
+from langfuse.langchain import CallbackHandler
+import pixie
+import requests
+from bs4 import BeautifulSoup
+
+
+langfuse_handler = CallbackHandler()
+
+
+def load_web_page(url: str) -> list[Document]:
+    """Simple web page loader using requests and BeautifulSoup.
+
+    Replaces langchain_community.document_loaders.WebBaseLoader
+    to avoid the langchain-community dependency.
+    """
+    response = requests.get(url)
+    response.raise_for_status()
+    soup = BeautifulSoup(response.content, "html.parser")
+
+    # Extract text from the page
+    text = soup.get_text(separator="\n", strip=True)
+
+    return [Document(page_content=text, metadata={"source": url})]
 
 
 def setup_vectorstore():
@@ -34,7 +57,7 @@ def setup_vectorstore():
         "https://lilianweng.github.io/posts/2024-04-12-diffusion-video/",
     ]
 
-    docs = [WebBaseLoader(url).load() for url in urls]
+    docs = [load_web_page(url) for url in urls]
     docs_list = [item for sublist in docs for item in sublist]
 
     print("Splitting documents...")
@@ -66,7 +89,9 @@ def create_rag_graph(retriever, model):
     # Node: Generate query or respond
     def generate_query_or_respond(state: MessagesState):
         """Call the model to generate a response or use retrieval tool."""
-        response = model.bind_tools([retriever_tool]).invoke(state["messages"])
+        response = model.bind_tools([retriever_tool]).invoke(
+            state["messages"], config={"callbacks": [langfuse_handler]}
+        )
         return {"messages": [response]}
 
     # Grade documents schema
@@ -97,7 +122,8 @@ def create_rag_graph(retriever, model):
 
         prompt = GRADE_PROMPT.format(question=question, context=context)
         response = grader_model.with_structured_output(GradeDocuments).invoke(
-            [{"role": "user", "content": prompt}]
+            [{"role": "user", "content": prompt}],
+            config={"callbacks": [langfuse_handler]},
         )
 
         score = response.binary_score  # type: ignore
@@ -121,7 +147,10 @@ def create_rag_graph(retriever, model):
         messages = state["messages"]
         question = messages[0].content
         prompt = REWRITE_PROMPT.format(question=question)
-        response = model.invoke([{"role": "user", "content": prompt}])
+        response = model.invoke(
+            [{"role": "user", "content": prompt}],
+            config={"callbacks": [langfuse_handler]},
+        )
         return {"messages": [HumanMessage(content=response.content)]}
 
     # Node: Generate answer
@@ -139,7 +168,10 @@ def create_rag_graph(retriever, model):
         question = state["messages"][0].content
         context = state["messages"][-1].content
         prompt = GENERATE_PROMPT.format(question=question, context=context)
-        response = model.invoke([{"role": "user", "content": prompt}])
+        response = model.invoke(
+            [{"role": "user", "content": prompt}],
+            config={"callbacks": [langfuse_handler]},
+        )
         return {"messages": [response]}
 
     # Build graph
@@ -172,7 +204,7 @@ def create_rag_graph(retriever, model):
     return workflow.compile()
 
 
-@pixie_app
+@pixie.app
 async def langgraph_rag_agent(question: str) -> str:
     """Agentic RAG system that can answer questions about Lilian Weng's blog posts.
 
@@ -200,7 +232,10 @@ async def langgraph_rag_agent(question: str) -> str:
     print(f"Processing question: {question}")
 
     # Run the graph
-    result = graph.invoke({"messages": [{"role": "user", "content": question}]})  # type: ignore
+    result = graph.invoke(
+        {"messages": [{"role": "user", "content": question}]},  # type: ignore
+        config={"callbacks": [langfuse_handler]},
+    )
 
     # Return the final answer
     return result["messages"][-1].content

{pixie_examples-0.1.0/examples/langchain → pixie_examples-0.1.1.dev5/examples/langgraph}/langgraph_sql_agent.py

@@ -13,12 +13,16 @@ import requests
 from typing import Literal
 from langchain.chat_models import init_chat_model
 from langchain.messages import AIMessage
-from langchain_community.utilities import SQLDatabase
-from langchain_community.agent_toolkits import SQLDatabaseToolkit
 from langgraph.graph import START, MessagesState, StateGraph
+from ..sql_utils import SQLDatabase, SQLDatabaseToolkit
 from langgraph.prebuilt import ToolNode
 from langgraph.checkpoint.memory import InMemorySaver
-from pixie import pixie_app
+
+from langfuse.langchain import CallbackHandler
+import pixie
+
+
+langfuse_handler = CallbackHandler()
 
 
 def setup_database():
@@ -74,7 +78,9 @@ def create_sql_graph(db: SQLDatabase, model):
     # Node: Force model to call get_schema
     def call_get_schema(state: MessagesState):
         llm_with_tools = model.bind_tools([get_schema_tool], tool_choice="any")
-        response = llm_with_tools.invoke(state["messages"])
+        response = llm_with_tools.invoke(
+            state["messages"], config={"callbacks": [langfuse_handler]}
+        )
         return {"messages": [response]}
 
     # Node: Generate query
@@ -95,7 +101,10 @@ DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the databa
     def generate_query(state: MessagesState):
         system_message = {"role": "system", "content": generate_query_prompt}
         llm_with_tools = model.bind_tools([run_query_tool])
-        response = llm_with_tools.invoke([system_message] + state["messages"])
+        response = llm_with_tools.invoke(
+            [system_message] + state["messages"],
+            config={"callbacks": [langfuse_handler]},
+        )
         return {"messages": [response]}
 
     # Node: Check query
@@ -130,7 +139,9 @@ You will call the appropriate tool to execute the query after running this check
             # Fallback if no tool calls
             user_message = {"role": "user", "content": "Please check the query"}
         llm_with_tools = model.bind_tools([run_query_tool], tool_choice="any")
-        response = llm_with_tools.invoke([system_message, user_message])
+        response = llm_with_tools.invoke(
+            [system_message, user_message], config={"callbacks": [langfuse_handler]}
+        )
         if isinstance(last_message, AI):
             response.id = last_message.id
         return {"messages": [response]}
@@ -167,7 +178,7 @@ You will call the appropriate tool to execute the query after running this check
     return builder.compile(checkpointer=InMemorySaver())
 
 
-@pixie_app
+@pixie.app
 async def langgraph_sql_agent(question: str) -> str:
     """Custom SQL agent built with LangGraph primitives.
 
@@ -197,7 +208,10 @@ async def langgraph_sql_agent(question: str) -> str:
     # Run the graph
     result = graph.invoke(
         {"messages": [{"role": "user", "content": question}]},  # type: ignore
-        {"configurable": {"thread_id": "langgraph_sql"}},
+        {
+            "configurable": {"thread_id": "langgraph_sql"},
+            "callbacks": [langfuse_handler],
+        },
     )
 
     # Return the final message