File size: 4,153 Bytes
60ac2eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
#!/usr/bin/env python3
"""
Lineage Graph Extractor - Integration Example
This script demonstrates how to use the Lineage Graph Extractor agent
programmatically with the Anthropic API.
Usage:
python integration_example.py
"""
import os
from anthropic import Anthropic
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
def load_agent_config():
"""Load the agent configuration from memories/agent.md"""
config_path = os.path.join(os.path.dirname(__file__), "memories", "agent.md")
with open(config_path, "r") as f:
return f.read()
def extract_lineage(client, system_prompt, user_message):
"""
Send a lineage extraction request to the agent.
Args:
client: Anthropic client instance
system_prompt: Agent system prompt
user_message: User's lineage extraction request
Returns:
Agent's response text
"""
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=4000,
system=system_prompt,
messages=[{
"role": "user",
"content": user_message
}]
)
return response.content[0].text
def main():
"""Main function demonstrating agent usage"""
# Initialize Anthropic client
api_key = os.getenv("ANTHROPIC_API_KEY")
if not api_key:
print("Error: ANTHROPIC_API_KEY not found in environment variables.")
print("Please set it in your .env file.")
return
client = Anthropic(api_key=api_key)
# Load agent configuration
print("Loading agent configuration...")
system_prompt = load_agent_config()
print("✓ Agent configuration loaded\n")
# Example 1: Simple greeting to test agent
print("=" * 60)
print("Example 1: Testing agent connection")
print("=" * 60)
response = extract_lineage(
client,
system_prompt,
"Hello! What can you help me with?"
)
print(response)
print()
# Example 2: Extract lineage from sample metadata
print("=" * 60)
print("Example 2: Extract lineage from sample metadata")
print("=" * 60)
sample_metadata = """
{
"tables": [
{
"name": "raw_orders",
"type": "source",
"description": "Raw order data from API"
},
{
"name": "raw_customers",
"type": "source",
"description": "Raw customer data from database"
},
{
"name": "stg_orders",
"type": "staging",
"description": "Cleaned and standardized orders",
"depends_on": ["raw_orders"]
},
{
"name": "stg_customers",
"type": "staging",
"description": "Cleaned and standardized customers",
"depends_on": ["raw_customers"]
},
{
"name": "fct_orders",
"type": "fact",
"description": "Order facts with customer data",
"depends_on": ["stg_orders", "stg_customers"]
}
]
}
"""
response = extract_lineage(
client,
system_prompt,
f"Extract lineage from this metadata and create a Mermaid diagram:\n\n{sample_metadata}"
)
print(response)
print()
# Example 3: BigQuery extraction (requires credentials)
if os.getenv("GOOGLE_CLOUD_PROJECT"):
print("=" * 60)
print("Example 3: BigQuery lineage extraction")
print("=" * 60)
project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
response = extract_lineage(
client,
system_prompt,
f"Extract lineage from BigQuery project: {project_id}, dataset: analytics"
)
print(response)
else:
print("Skipping BigQuery example (GOOGLE_CLOUD_PROJECT not set)")
print("\n" + "=" * 60)
print("Examples complete!")
print("=" * 60)
if __name__ == "__main__":
main()
|