BAML to Graph Data Example

app.py

import streamlit as st
from st_cytoscape import cytoscape
from baml_client import b
import json

# Set page config
st.set_page_config(layout="wide")

def process_graph_result(result):
    """Convert BAML graph result to cytoscape elements format"""
    
    elements = []

    # Add nodes
    if hasattr(result, 'nodes'):
        nodes = [n.model_dump() for n in result.nodes]
        elements.extend(nodes)
    
    # Add edges
    if hasattr(result, 'edges'):
        edges = [e.model_dump() for e in result.edges]
        elements.extend(edges)
    
    print(f'elements: {json.dumps(elements, indent=2)}')
    return elements

def main():
    st.title("Interactive Graph Visualization with Cytoscape")
    
    # Initialize session state for elements if it doesn't exist
    if 'elements' not in st.session_state:
        st.session_state.elements = []
    else:
        print(f'session state elements: {st.session_state.elements}')
    
    # Add text input
    user_input = st.text_area("Enter your text to extract a network graph:", 
                             height=150,
                             placeholder="Paste your text here...")
    
    # Add extract button
    if st.button("Extract Network Graph"):
        if user_input.strip():
            try:
                # Call the BAML function to process the text
                result = b.ExtractNetworkGraph(user_input)
                elements = process_graph_result(result)
                st.session_state.elements = elements
            except Exception as e:
                st.error(f"Error processing text: {e}")
        else:
            st.warning("Please enter some text to extract a graph.")

    if st.session_state.elements:
        stylesheet = [
            {
                'selector': 'node',
                'style': {
                    'content': 'data(label)',
                    'background-color': '#0074D9',
                    'color': '#fff',
                    'text-valign': 'center',
                    'text-halign': 'center',
                }
            },
            {
                'selector': 'edge',
                'style': {
                    'label': 'data(label)',
                    'width': 2,
                    'line-color': '#7FDBFF',
                    'target-arrow-color': '#7FDBFF',
                    'target-arrow-shape': 'triangle',
                    'curve-style': 'bezier'
                }
            }
        ]

        # Display the graph
        selected = cytoscape(
            st.session_state.elements,
            stylesheet=stylesheet,
            key="graph",
        )
        print(f'selected: {selected}')


if __name__ == "__main__":
    main()

graphs.baml

// Place this in the baml_src folder

class CytoscapeNode {
  data CytoscapeNodeData
}

class CytoscapeNodeData {
  id string @description("Required Unique identifier for the node")
  label string @description("Optional Display label for the node")
  type "person" | "organization" | "location" @description("Type of entity")
}

class CytoscapeEdge {
  data CytoscapeEdgeData
}

class CytoscapeEdgeData {
  id string @description("Required Unique identifier for the node")
  source string @description("Required Source node ID")
  target string @description("Required Target node ID")
  relationship string @description("Optional Type of relationship between nodes")
}

class CytoscapeGraph {
  nodes CytoscapeNode[]
  edges CytoscapeEdge[]
}

function ExtractNetworkGraph(text: string) -> CytoscapeGraph {
  client OllamaLlama3
  prompt #"
    Extract entities and their relationships from the text into a network graph format.
    Identify people, organizations, and locations as nodes, and their relationships as edges.
    Create unique IDs for each node and edge.

    {{ ctx.output_format }}

    {{ _.role("user") }} {{ text }}
  "#
}

test BasicExtraction {
  functions [ExtractNetworkGraph]
  args {
    text #"
      John Smith works at Apple Inc. in California. He reports to Lisa Jones, 
      who is the VP of Engineering at Apple Inc.
    "#
  }
}

test ComplexExtraction {
  functions [ExtractNetworkGraph]
  args {
    text #"
      Microsoft, based in Redmond, partnered with OpenAI in San Francisco. 
      Satya Nadella, CEO of Microsoft, announced the deal alongside Sam Altman from OpenAI.
    "#
  }
}