pixie-examples 0.1.1.dev3__py3-none-any.whl

examples/langchain/sql_agent.py

@@ -0,0 +1,176 @@
+"""
+SQL Agent (Multi-Turn & Multi-Step)
+
+This example demonstrates an agent that can answer questions about a SQL database.
+The agent can:
+1. Fetch available tables and schemas
+2. Decide which tables are relevant
+3. Generate SQL queries
+4. Execute queries and handle errors
+5. Formulate responses based on results
+
+Based on: https://docs.langchain.com/oss/python/langchain/sql-agent
+
+WARNING: Building Q&A systems of SQL databases requires executing model-generated SQL
+queries. Make sure database connection permissions are scoped as narrowly as possible.
+"""
+
+import pathlib
+import requests
+from langchain.agents import create_agent
+from langchain.chat_models import init_chat_model
+from langgraph.checkpoint.memory import InMemorySaver
+
+from langfuse.langchain import CallbackHandler
+import pixie
+from ..sql_utils import SQLDatabase, SQLDatabaseToolkit
+
+
+langfuse_handler = CallbackHandler()
+
+
+# System prompt for SQL agent
+SQL_AGENT_PROMPT = """
+You are an agent designed to interact with a SQL database.
+Given an input question, create a syntactically correct {dialect} query to run,
+then look at the results of the query and return the answer. Unless the user
+specifies a specific number of examples they wish to obtain, always limit your
+query to at most {top_k} results.
+
+You can order the results by a relevant column to return the most interesting
+examples in the database. Never query for all the columns from a specific table,
+only ask for the relevant columns given the question.
+
+You MUST double check your query before executing it. If you get an error while
+executing a query, rewrite the query and try again.
+
+DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the
+database.
+
+To start you should ALWAYS look at the tables in the database to see what you
+can query. Do NOT skip this step.
+
+Then you should query the schema of the most relevant tables.
+"""
+
+
+def setup_database():
+    """Download and setup the Chinook database if not already present."""
+    url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"
+    local_path = pathlib.Path("Chinook.db")
+
+    if local_path.exists():
+        print(f"{local_path} already exists, skipping download.")
+    else:
+        print("Downloading Chinook database...")
+        response = requests.get(url)
+        if response.status_code == 200:
+            local_path.write_bytes(response.content)
+            print(f"File downloaded and saved as {local_path}")
+        else:
+            raise Exception(
+                f"Failed to download the file. Status code: {response.status_code}"
+            )
+
+    return SQLDatabase.from_uri("sqlite:///Chinook.db")
+
+
+@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,
+    albums, tracks, customers, invoices, etc.
+
+    Args:
+        question: Natural language question about the database
+
+    Returns:
+        AI-generated answer based on SQL query results
+    """
+    # Setup database
+    db = setup_database()
+
+    # Initialize model
+    model = init_chat_model("gpt-4o-mini", temperature=0)
+
+    # Create SQL toolkit with tools for database interaction
+    toolkit = SQLDatabaseToolkit(db=db, llm=model)
+    tools = toolkit.get_tools()
+
+    # Format system prompt with database info
+    system_prompt = SQL_AGENT_PROMPT.format(dialect=db.dialect, top_k=5)
+
+    # Create agent
+    agent = create_agent(model, tools, system_prompt=system_prompt)
+
+    # Run the agent
+    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 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.
+
+    Yields:
+        AI responses to database queries
+    """
+    # Setup database
+    db = setup_database()
+
+    # Initialize model
+    model = init_chat_model("gpt-4o-mini", temperature=0)
+
+    # Create SQL toolkit
+    toolkit = SQLDatabaseToolkit(db=db, llm=model)
+    tools = toolkit.get_tools()
+
+    # Format system prompt
+    system_prompt = SQL_AGENT_PROMPT.format(dialect=db.dialect, top_k=5)
+
+    # Create agent with checkpointer for conversation memory
+    agent = create_agent(
+        model, tools, system_prompt=system_prompt, checkpointer=InMemorySaver()
+    )
+
+    # Send welcome message
+    yield f"""Welcome to the SQL Query Assistant!
+
+I can help you query the Chinook database, which contains information about:
+- Artists and Albums
+- Tracks and Genres
+- Customers and Invoices
+- Employees and more
+
+Available tables: {', '.join(db.get_usable_table_names())}
+
+Ask me any question about the data!"""
+
+    # Initialize conversation
+    thread_id = "sql_thread"
+    config = {"configurable": {"thread_id": thread_id}, "callbacks": [langfuse_handler]}
+
+    while True:
+        # Get user question
+        user_question = yield pixie.InputRequired(str)
+
+        # Check for exit
+        if user_question.lower() in {"exit", "quit", "bye"}:
+            yield "Goodbye! Feel free to come back if you have more questions about the database."
+            break
+
+        # Process with agent
+        result = agent.invoke(
+            {"messages": [{"role": "user", "content": user_question}]}, config  # type: ignore
+        )
+
+        # Yield the agent's response
+        yield result["messages"][-1].content

examples/langgraph/langgraph_rag.py

@@ -0,0 +1,241 @@
+"""
+LangGraph RAG Agent (Retrieval Augmented Generation)
+
+This example demonstrates building an agentic RAG system using LangGraph that can:
+1. Decide when to use retrieval vs. respond directly
+2. Grade retrieved documents for relevance
+3. Rewrite questions if documents aren't relevant
+4. Generate answers based on retrieved context
+
+Based on: https://docs.langchain.com/oss/python/langgraph/agentic-rag
+"""
+
+from pydantic import BaseModel, Field
+from typing import Literal
+from langchain.chat_models import init_chat_model
+from langchain.tools import tool
+from langchain.messages import HumanMessage
+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 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():
+    """Setup vectorstore with documents from Lilian Weng's blog."""
+    print("Loading documents from web...")
+
+    urls = [
+        "https://lilianweng.github.io/posts/2024-11-28-reward-hacking/",
+        "https://lilianweng.github.io/posts/2024-07-07-hallucination/",
+        "https://lilianweng.github.io/posts/2024-04-12-diffusion-video/",
+    ]
+
+    docs = [load_web_page(url) for url in urls]
+    docs_list = [item for sublist in docs for item in sublist]
+
+    print("Splitting documents...")
+    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+        chunk_size=100, chunk_overlap=50
+    )
+    doc_splits = text_splitter.split_documents(docs_list)
+
+    print("Creating vectorstore...")
+    vectorstore = InMemoryVectorStore.from_documents(
+        documents=doc_splits, embedding=OpenAIEmbeddings()
+    )
+
+    return vectorstore.as_retriever()
+
+
+def create_rag_graph(retriever, model):
+    """Create a LangGraph-based RAG agent."""
+
+    # Create retriever tool
+    @tool
+    def retrieve_blog_posts(query: str) -> str:
+        """Search and return information about Lilian Weng blog posts."""
+        docs = retriever.invoke(query)
+        return "\n\n".join([doc.page_content for doc in docs])
+
+    retriever_tool = retrieve_blog_posts
+
+    # 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"], config={"callbacks": [langfuse_handler]}
+        )
+        return {"messages": [response]}
+
+    # Grade documents schema
+    class GradeDocuments(BaseModel):
+        """Grade documents using a binary score for relevance check."""
+
+        binary_score: str = Field(
+            description="Relevance score: 'yes' if relevant, or 'no' if not relevant"
+        )
+
+    GRADE_PROMPT = (
+        "You are a grader assessing relevance of a retrieved document to a user question. \n "
+        "Here is the retrieved document: \n\n {context} \n\n"
+        "Here is the user question: {question} \n"
+        "If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n"
+        "Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."
+    )
+
+    grader_model = init_chat_model("gpt-4o", temperature=0)
+
+    # Conditional edge: Grade documents
+    def grade_documents(
+        state: MessagesState,
+    ) -> Literal["generate_answer", "rewrite_question"]:
+        """Determine whether the retrieved documents are relevant to the question."""
+        question = state["messages"][0].content
+        context = state["messages"][-1].content
+
+        prompt = GRADE_PROMPT.format(question=question, context=context)
+        response = grader_model.with_structured_output(GradeDocuments).invoke(
+            [{"role": "user", "content": prompt}],
+            config={"callbacks": [langfuse_handler]},
+        )
+
+        score = response.binary_score  # type: ignore
+        if score == "yes":
+            return "generate_answer"
+        else:
+            return "rewrite_question"
+
+    # Node: Rewrite question
+    REWRITE_PROMPT = (
+        "Look at the input and try to reason about the underlying semantic intent / meaning.\n"
+        "Here is the initial question:"
+        "\n ------- \n"
+        "{question}"
+        "\n ------- \n"
+        "Formulate an improved question:"
+    )
+
+    def rewrite_question(state: MessagesState):
+        """Rewrite the original user question."""
+        messages = state["messages"]
+        question = messages[0].content
+        prompt = REWRITE_PROMPT.format(question=question)
+        response = model.invoke(
+            [{"role": "user", "content": prompt}],
+            config={"callbacks": [langfuse_handler]},
+        )
+        return {"messages": [HumanMessage(content=response.content)]}
+
+    # Node: Generate answer
+    GENERATE_PROMPT = (
+        "You are an assistant for question-answering tasks. "
+        "Use the following pieces of retrieved context to answer the question. "
+        "If you don't know the answer, just say that you don't know. "
+        "Use three sentences maximum and keep the answer concise.\n"
+        "Question: {question} \n"
+        "Context: {context}"
+    )
+
+    def generate_answer(state: MessagesState):
+        """Generate an answer."""
+        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}],
+            config={"callbacks": [langfuse_handler]},
+        )
+        return {"messages": [response]}
+
+    # Build graph
+    workflow = StateGraph(MessagesState)
+
+    # Define nodes
+    workflow.add_node(generate_query_or_respond)
+    workflow.add_node("retrieve", ToolNode([retriever_tool]))
+    workflow.add_node(rewrite_question)
+    workflow.add_node(generate_answer)
+
+    # Define edges
+    workflow.add_edge(START, "generate_query_or_respond")
+
+    # Decide whether to retrieve
+    workflow.add_conditional_edges(
+        "generate_query_or_respond",
+        tools_condition,
+        {
+            "tools": "retrieve",
+            END: END,
+        },
+    )
+
+    # Grade documents after retrieval
+    workflow.add_conditional_edges("retrieve", grade_documents)
+    workflow.add_edge("generate_answer", END)
+    workflow.add_edge("rewrite_question", "generate_query_or_respond")
+
+    return workflow.compile()
+
+
+@pixie.app
+async def langgraph_rag_agent(question: str) -> str:
+    """Agentic RAG system that can answer questions about Lilian Weng's blog posts.
+
+    The agent:
+    1. Decides whether to retrieve or respond directly
+    2. Grades retrieved documents for relevance
+    3. Rewrites questions if needed
+    4. Generates answers based on context
+
+    Args:
+        question: Natural language question about the blog content
+
+    Returns:
+        AI-generated answer based on retrieved context
+    """
+    # Setup retriever (this will take a moment on first run)
+    retriever = setup_vectorstore()
+
+    # Initialize model
+    model = init_chat_model("gpt-4o-mini", temperature=0)
+
+    # Create graph
+    graph = create_rag_graph(retriever, model)
+
+    print(f"Processing question: {question}")
+
+    # Run the graph
+    result = graph.invoke(
+        {"messages": [{"role": "user", "content": question}]},  # type: ignore
+        config={"callbacks": [langfuse_handler]},
+    )
+
+    # Return the final answer
+    return result["messages"][-1].content

examples/langgraph/langgraph_sql_agent.py

@@ -0,0 +1,218 @@
+"""
+LangGraph SQL Agent (Custom Implementation)
+
+This example demonstrates building a SQL agent directly using LangGraph primitives
+for deeper customization. This gives more control over the agent's behavior compared
+to the higher-level LangChain agent.
+
+Based on: https://docs.langchain.com/oss/python/langgraph/sql-agent
+"""
+
+import pathlib
+import requests
+from typing import Literal
+from langchain.chat_models import init_chat_model
+from langchain.messages import AIMessage
+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 langfuse.langchain import CallbackHandler
+import pixie
+
+
+langfuse_handler = CallbackHandler()
+
+
+def setup_database():
+    """Download and setup the Chinook database if not already present."""
+    url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"
+    local_path = pathlib.Path("Chinook.db")
+
+    if local_path.exists():
+        print(f"{local_path} already exists, skipping download.")
+    else:
+        print("Downloading Chinook database...")
+        response = requests.get(url)
+        if response.status_code == 200:
+            local_path.write_bytes(response.content)
+            print(f"File downloaded and saved as {local_path}")
+        else:
+            raise Exception(
+                f"Failed to download the file. Status code: {response.status_code}"
+            )
+
+    return SQLDatabase.from_uri("sqlite:///Chinook.db")
+
+
+def create_sql_graph(db: SQLDatabase, model):
+    """Create a LangGraph-based SQL agent with custom workflow."""
+
+    # Get tools from toolkit
+    toolkit = SQLDatabaseToolkit(db=db, llm=model)
+    tools = toolkit.get_tools()
+
+    # Extract specific tools
+    get_schema_tool = next(tool for tool in tools if tool.name == "sql_db_schema")
+    get_schema_node = ToolNode([get_schema_tool], name="get_schema")
+
+    run_query_tool = next(tool for tool in tools if tool.name == "sql_db_query")
+    run_query_node = ToolNode([run_query_tool], name="run_query")
+
+    list_tables_tool = next(tool for tool in tools if tool.name == "sql_db_list_tables")
+
+    # Node: List tables (forced tool call)
+    def list_tables(state: MessagesState):
+        tool_call = {
+            "name": "sql_db_list_tables",
+            "args": {},
+            "id": "list_tables_call",
+            "type": "tool_call",
+        }
+        tool_call_message = AIMessage(content="", tool_calls=[tool_call])
+        tool_message = list_tables_tool.invoke(tool_call)
+        response = AIMessage(f"Available tables: {tool_message.content}")
+        return {"messages": [tool_call_message, tool_message, response]}
+
+    # 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"], config={"callbacks": [langfuse_handler]}
+        )
+        return {"messages": [response]}
+
+    # Node: Generate query
+    generate_query_prompt = f"""
+You are an agent designed to interact with a SQL database.
+Given an input question, create a syntactically correct {db.dialect} query to run,
+then look at the results of the query and return the answer. Unless the user
+specifies a specific number of examples they wish to obtain, always limit your
+query to at most 5 results.
+
+You can order the results by a relevant column to return the most interesting
+examples in the database. Never query for all the columns from a specific table,
+only ask for the relevant columns given the question.
+
+DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.
+"""
+
+    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"],
+            config={"callbacks": [langfuse_handler]},
+        )
+        return {"messages": [response]}
+
+    # Node: Check query
+    check_query_prompt = f"""
+You are a SQL expert with a strong attention to detail.
+Double check the {db.dialect} query for common mistakes, including:
+- Using NOT IN with NULL values
+- Using UNION when UNION ALL should have been used
+- Using BETWEEN for exclusive ranges
+- Data type mismatch in predicates
+- Properly quoting identifiers
+- Using the correct number of arguments for functions
+- Casting to the correct data type
+- Using the proper columns for joins
+
+If there are any of the above mistakes, rewrite the query. If there are no mistakes,
+just reproduce the original query.
+
+You will call the appropriate tool to execute the query after running this check.
+"""
+
+    def check_query(state: MessagesState):
+        from langchain.messages import AIMessage as AI
+
+        system_message = {"role": "system", "content": check_query_prompt}
+        last_message = state["messages"][-1]
+        # Only AIMessage has tool_calls
+        if isinstance(last_message, AI) and last_message.tool_calls:
+            tool_call = last_message.tool_calls[0]
+            user_message = {"role": "user", "content": tool_call["args"]["query"]}
+        else:
+            # 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], config={"callbacks": [langfuse_handler]}
+        )
+        if isinstance(last_message, AI):
+            response.id = last_message.id
+        return {"messages": [response]}
+
+    # Conditional edge: Continue or end
+    def should_continue(state: MessagesState) -> Literal["__end__", "check_query"]:
+        from langchain.messages import AIMessage as AI
+
+        messages = state["messages"]
+        last_message = messages[-1]
+        # Check if last message is AIMessage and has tool calls
+        if isinstance(last_message, AI) and last_message.tool_calls:
+            return "check_query"
+        else:
+            return "__end__"
+
+    # Build graph
+    builder = StateGraph(MessagesState)
+    builder.add_node("list_tables", list_tables)
+    builder.add_node("call_get_schema", call_get_schema)
+    builder.add_node("get_schema", get_schema_node)
+    builder.add_node("generate_query", generate_query)
+    builder.add_node("check_query", check_query)
+    builder.add_node("run_query", run_query_node)
+
+    builder.add_edge(START, "list_tables")
+    builder.add_edge("list_tables", "call_get_schema")
+    builder.add_edge("call_get_schema", "get_schema")
+    builder.add_edge("get_schema", "generate_query")
+    builder.add_conditional_edges("generate_query", should_continue)
+    builder.add_edge("check_query", "run_query")
+    builder.add_edge("run_query", "generate_query")
+
+    return builder.compile(checkpointer=InMemorySaver())
+
+
+@pixie.app
+async def langgraph_sql_agent(question: str) -> str:
+    """Custom SQL agent built with LangGraph primitives.
+
+    This agent has explicit control over the workflow:
+    1. Lists all tables
+    2. Gets schema for relevant tables
+    3. Generates SQL query
+    4. Checks query for errors
+    5. Executes query
+    6. Returns natural language answer
+
+    Args:
+        question: Natural language question about the database
+
+    Returns:
+        AI-generated answer based on SQL query results
+    """
+    # Setup database
+    db = setup_database()
+
+    # Initialize model
+    model = init_chat_model("gpt-4o-mini", temperature=0)
+
+    # Create graph
+    graph = create_sql_graph(db, model)
+
+    # Run the graph
+    result = graph.invoke(
+        {"messages": [{"role": "user", "content": question}]},  # type: ignore
+        {
+            "configurable": {"thread_id": "langgraph_sql"},
+            "callbacks": [langfuse_handler],
+        },
+    )
+
+    # Return the final message
+    return result["messages"][-1].content