164 lines
4.4 KiB
Python
164 lines
4.4 KiB
Python
"""
|
|
Setup and testing utilities for PocketFlow tracing.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from typing import Optional
|
|
|
|
# Add parent directory to path for imports
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
|
|
|
try:
|
|
from langfuse import Langfuse
|
|
LANGFUSE_AVAILABLE = True
|
|
except ImportError:
|
|
LANGFUSE_AVAILABLE = False
|
|
|
|
from tracing import TracingConfig, LangfuseTracer
|
|
|
|
|
|
def setup_tracing(env_file: Optional[str] = None) -> TracingConfig:
|
|
"""
|
|
Set up tracing configuration and validate the setup.
|
|
|
|
Args:
|
|
env_file: Optional path to .env file. If None, uses default location.
|
|
|
|
Returns:
|
|
TracingConfig instance.
|
|
|
|
Raises:
|
|
RuntimeError: If setup fails.
|
|
"""
|
|
print("🔧 Setting up PocketFlow tracing...")
|
|
|
|
# Check if langfuse is installed
|
|
if not LANGFUSE_AVAILABLE:
|
|
raise RuntimeError(
|
|
"Langfuse package not installed. Install with: pip install langfuse"
|
|
)
|
|
|
|
# Load configuration
|
|
if env_file:
|
|
config = TracingConfig.from_env(env_file)
|
|
print(f"✓ Loaded configuration from: {env_file}")
|
|
else:
|
|
config = TracingConfig.from_env()
|
|
print("✓ Loaded configuration from environment")
|
|
|
|
# Validate configuration
|
|
if not config.validate():
|
|
raise RuntimeError(
|
|
"Invalid tracing configuration. Please check your environment variables:\n"
|
|
"- LANGFUSE_SECRET_KEY\n"
|
|
"- LANGFUSE_PUBLIC_KEY\n"
|
|
"- LANGFUSE_HOST"
|
|
)
|
|
|
|
print("✓ Configuration validated")
|
|
|
|
# Test connection
|
|
if test_langfuse_connection(config):
|
|
print("✓ Langfuse connection successful")
|
|
else:
|
|
raise RuntimeError("Failed to connect to Langfuse. Check your configuration and network.")
|
|
|
|
print("🎉 PocketFlow tracing setup complete!")
|
|
return config
|
|
|
|
|
|
def test_langfuse_connection(config: TracingConfig) -> bool:
|
|
"""
|
|
Test connection to Langfuse.
|
|
|
|
Args:
|
|
config: TracingConfig instance.
|
|
|
|
Returns:
|
|
True if connection successful, False otherwise.
|
|
"""
|
|
try:
|
|
# Create a test tracer
|
|
tracer = LangfuseTracer(config)
|
|
|
|
if not tracer.client:
|
|
return False
|
|
|
|
# Try to start and end a test trace
|
|
trace_id = tracer.start_trace("test_connection", {"test": True})
|
|
if trace_id:
|
|
tracer.end_trace({"test": "completed"}, "success")
|
|
tracer.flush()
|
|
return True
|
|
|
|
return False
|
|
|
|
except Exception as e:
|
|
if config.debug:
|
|
print(f"Connection test failed: {e}")
|
|
return False
|
|
|
|
|
|
def print_configuration_help():
|
|
"""Print help information for configuring tracing."""
|
|
print("""
|
|
🔧 PocketFlow Tracing Configuration Help
|
|
|
|
To use PocketFlow tracing, you need to configure Langfuse credentials.
|
|
|
|
1. Create or update your .env file with:
|
|
|
|
LANGFUSE_SECRET_KEY=your-secret-key
|
|
LANGFUSE_PUBLIC_KEY=your-public-key
|
|
LANGFUSE_HOST=your-langfuse-host
|
|
POCKETFLOW_TRACING_DEBUG=true
|
|
|
|
2. Optional configuration:
|
|
|
|
POCKETFLOW_TRACE_INPUTS=true
|
|
POCKETFLOW_TRACE_OUTPUTS=true
|
|
POCKETFLOW_TRACE_PREP=true
|
|
POCKETFLOW_TRACE_EXEC=true
|
|
POCKETFLOW_TRACE_POST=true
|
|
POCKETFLOW_TRACE_ERRORS=true
|
|
POCKETFLOW_SESSION_ID=your-session-id
|
|
POCKETFLOW_USER_ID=your-user-id
|
|
|
|
3. Install required packages:
|
|
|
|
pip install -r requirements.txt
|
|
|
|
4. Test your setup:
|
|
|
|
python -c "from utils import setup_tracing; setup_tracing()"
|
|
""")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
"""Command-line interface for setup and testing."""
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="PocketFlow Tracing Setup")
|
|
parser.add_argument("--test", action="store_true", help="Test Langfuse connection")
|
|
parser.add_argument("--help-config", action="store_true", help="Show configuration help")
|
|
parser.add_argument("--env-file", type=str, help="Path to .env file")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.help_config:
|
|
print_configuration_help()
|
|
sys.exit(0)
|
|
|
|
if args.test:
|
|
try:
|
|
config = setup_tracing(args.env_file)
|
|
print("\n✅ All tests passed! Your tracing setup is ready.")
|
|
except Exception as e:
|
|
print(f"\n❌ Setup failed: {e}")
|
|
print("\nFor help with configuration, run:")
|
|
print("python utils/setup.py --help-config")
|
|
sys.exit(1)
|
|
else:
|
|
print_configuration_help()
|