Skip to content

Commit 72b96ca

Browse files
committed
feat: add LLMWithGateway for enterprise OAuth support
Add LLMWithGateway subclass to support enterprise API gateways with OAuth 2.0 authentication (e.g., APIGEE, Azure API Management). Key features: - OAuth 2.0 token fetch and automatic refresh - Thread-safe token caching with TTL - Custom header injection for gateway-specific requirements - Template variable replacement for flexible configuration - Fully generic implementation (no vendor lock-in) Implementation approach: - Separate LLMWithGateway class (addresses PR #963 feedback from neubig) - Focused feature set for OAuth + custom headers (no over-engineering) - Comprehensive test coverage - Complete documentation with examples This replaces the previous approach of modifying the main LLM class, keeping the codebase cleaner and more maintainable. Example usage (Wells Fargo APIGEE + Tachyon): ```python llm = LLMWithGateway( model="gemini-1.5-flash", base_url=os.environ["TACHYON_API_URL"], gateway_auth_url=os.environ["APIGEE_TOKEN_URL"], gateway_auth_headers={ "X-Client-Id": os.environ["APIGEE_CLIENT_ID"], "X-Client-Secret": os.environ["APIGEE_CLIENT_SECRET"], }, gateway_auth_body={"grant_type": "client_credentials"}, custom_headers={"X-Tachyon-Key": os.environ["TACHYON_API_KEY"]}, ) ``` Files added: - openhands-sdk/openhands/sdk/llm/llm_with_gateway.py (new class) - tests/sdk/llm/test_llm_with_gateway.py (comprehensive tests) - openhands-sdk/docs/llm_with_gateway.md (API documentation) - examples/apigee_tachyon_example.py (working example)
1 parent b9860ce commit 72b96ca

File tree

5 files changed

+1164
-0
lines changed

5 files changed

+1164
-0
lines changed

examples/apigee_tachyon_example.py

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
"""Example: Using OpenHands SDK with APIGEE + Tachyon Gateway
2+
3+
This example demonstrates how to configure the OpenHands SDK to work with
4+
enterprise gateways that use OAuth 2.0 authentication, specifically APIGEE + Tachyon.
5+
6+
Setup:
7+
1. Set environment variables for your credentials
8+
2. Run this script to test the connection
9+
3. Use the same configuration in your OpenHands CLI or application
10+
11+
Requirements:
12+
- Access to APIGEE OAuth endpoint
13+
- APIGEE client credentials (key and secret)
14+
- Tachyon API access and key
15+
- Tachyon endpoint URL through APIGEE
16+
"""
17+
18+
import os
19+
import sys
20+
21+
from openhands.sdk.llm import LLMWithGateway
22+
23+
24+
def main():
25+
# Check required environment variables
26+
required_vars = [
27+
"APIGEE_TOKEN_URL",
28+
"APIGEE_CLIENT_ID",
29+
"APIGEE_CLIENT_SECRET",
30+
"TACHYON_API_URL",
31+
"TACHYON_API_KEY",
32+
]
33+
34+
missing_vars = [var for var in required_vars if var not in os.environ]
35+
if missing_vars:
36+
print("Error: Missing required environment variables:")
37+
for var in missing_vars:
38+
print(f" - {var}")
39+
print("\nPlease set these variables in your environment:")
40+
print(" export APIGEE_TOKEN_URL='https://apigee.example.com/oauth/token'")
41+
print(" export APIGEE_CLIENT_ID='your-client-id'")
42+
print(" export APIGEE_CLIENT_SECRET='your-client-secret'")
43+
print(" export TACHYON_API_URL='https://apigee.example.com/.../tachyon/v1'")
44+
print(" export TACHYON_API_KEY='your-tachyon-key'")
45+
sys.exit(1)
46+
47+
# Configure LLM with gateway support
48+
print("Configuring LLM with APIGEE + Tachyon gateway...")
49+
llm = LLMWithGateway(
50+
# Model configuration
51+
model="gemini-1.5-flash", # or "gemini-1.5-pro"
52+
base_url=os.environ["TACHYON_API_URL"],
53+
# APIGEE OAuth configuration
54+
gateway_auth_url=os.environ["APIGEE_TOKEN_URL"],
55+
gateway_auth_method="POST",
56+
gateway_auth_headers={
57+
"Content-Type": "application/json",
58+
"X-Client-Id": os.environ["APIGEE_CLIENT_ID"],
59+
"X-Client-Secret": os.environ["APIGEE_CLIENT_SECRET"],
60+
},
61+
gateway_auth_body={
62+
"grant_type": "client_credentials",
63+
# Add any additional fields required by your APIGEE setup:
64+
# "scope": "llm:access",
65+
# "audience": "tachyon-api",
66+
},
67+
gateway_auth_token_path="access_token", # Adjust if your response differs
68+
gateway_auth_token_ttl=86400, # 24 hours (adjust based on your token expiry)
69+
# Tachyon-specific configuration
70+
custom_headers={
71+
"X-Tachyon-Key": os.environ["TACHYON_API_KEY"],
72+
# Add any additional Tachyon headers here
73+
},
74+
# LLM configuration
75+
temperature=0.7,
76+
max_output_tokens=1000,
77+
usage_id="apigee-tachyon-example",
78+
)
79+
80+
print("✓ LLM configured successfully")
81+
print("\nTesting connection...")
82+
83+
try:
84+
# Test with a simple query
85+
messages = [{"role": "user", "content": "What is 2 + 2? Reply with just the number."}]
86+
87+
print("Sending test message to LLM...")
88+
response = llm.completion(messages)
89+
90+
print("✓ Connection successful!")
91+
print(f"\nResponse: {response.content[0].text}")
92+
print(f"\nMetrics:")
93+
print(f" - Model: {response.model}")
94+
print(f" - Input tokens: {llm.metrics.input_tokens}")
95+
print(f" - Output tokens: {llm.metrics.output_tokens}")
96+
print(f" - Total cost: ${llm.metrics.accumulated_cost:.4f}")
97+
98+
# Test with a more complex query
99+
print("\n" + "=" * 50)
100+
print("Testing with a more complex query...")
101+
102+
messages = [
103+
{
104+
"role": "user",
105+
"content": "Explain what enterprise API gateways like APIGEE do in one sentence.",
106+
}
107+
]
108+
109+
response = llm.completion(messages)
110+
print(f"\nResponse: {response.content[0].text}")
111+
112+
print("\n✓ All tests passed! Your APIGEE + Tachyon setup is working correctly.")
113+
114+
except Exception as e:
115+
print(f"\n✗ Error: {e}")
116+
print("\nTroubleshooting:")
117+
print("1. Verify your APIGEE credentials are correct")
118+
print("2. Check that the APIGEE OAuth URL is accessible")
119+
print("3. Ensure the Tachyon API URL is correct")
120+
print("4. Verify the Tachyon API key is valid")
121+
print("5. Check network connectivity to APIGEE endpoints")
122+
sys.exit(1)
123+
124+
125+
if __name__ == "__main__":
126+
main()

0 commit comments

Comments
 (0)