312 lines
9.1 KiB
Markdown
312 lines
9.1 KiB
Markdown
# API Integration
|
|
|
|
This document explains how Edison integrates with the OpenAI API to translate natural language into shell commands.
|
|
|
|
## Overview
|
|
|
|
```mermaid
|
|
flowchart LR
|
|
A[User Query] --> B[Prompt Construction]
|
|
B --> C[API Request]
|
|
C --> D[Response Processing]
|
|
D --> E[Command Extraction]
|
|
E --> F[Command Validation]
|
|
F --> G[Command Execution]
|
|
|
|
style A fill:#f9d5e5,stroke:#333,stroke-width:2px
|
|
style B fill:#eeeeee,stroke:#333,stroke-width:2px
|
|
style C fill:#d3f6db,stroke:#333,stroke-width:2px
|
|
style D fill:#d3f6f5,stroke:#333,stroke-width:2px
|
|
style E fill:#d5f6d5,stroke:#333,stroke-width:2px
|
|
style F fill:#f6f6d5,stroke:#333,stroke-width:2px
|
|
style G fill:#f5d5f5,stroke:#333,stroke-width:2px
|
|
```
|
|
|
|
## API Client Module
|
|
|
|
The `api_client.py` module is responsible for all OpenAI API interactions:
|
|
|
|
```python
|
|
def get_api_key(config):
|
|
"""Get the OpenAI API key from various sources."""
|
|
# Find API key from environment, file, or config
|
|
|
|
def create_client(config):
|
|
"""Create and initialize an OpenAI client."""
|
|
# Initialize client with API key
|
|
|
|
def call_api(client, config, query):
|
|
"""Call the OpenAI API with the given query."""
|
|
# Send request to API and extract response
|
|
|
|
def generate_command(client, config, query, max_retries=3):
|
|
"""Generate a command using the OpenAI API with retry logic."""
|
|
# Call API with retries for rate limits
|
|
```
|
|
|
|
## API Key Management
|
|
|
|
Edison supports multiple methods for supplying the OpenAI API key, processed in this order:
|
|
|
|
1. **Environment Variable**: `OPENAI_API_KEY`
|
|
2. **API Key File**: `~/.openai.apikey`
|
|
3. **Configuration File**: `openai_api_key` in `edison.yaml`
|
|
|
|
This implementation is in `get_api_key()`:
|
|
|
|
```python
|
|
def get_api_key(config):
|
|
"""Get the OpenAI API key from various sources."""
|
|
dotenv.load_dotenv()
|
|
|
|
# Method 1: Environment variable
|
|
api_key = os.getenv("OPENAI_API_KEY")
|
|
|
|
# Method 2: File in home directory
|
|
if not api_key:
|
|
home_path = os.path.expanduser("~")
|
|
api_key_path = os.path.join(home_path, ".openai.apikey")
|
|
if os.path.exists(api_key_path):
|
|
with open(api_key_path, 'r') as f:
|
|
api_key = f.read().strip()
|
|
|
|
# Method 3: Configuration file
|
|
if not api_key:
|
|
api_key = config.get("openai_api_key")
|
|
|
|
if not api_key:
|
|
raise ValueError("No OpenAI API key found")
|
|
|
|
return api_key
|
|
```
|
|
|
|
## Prompt Construction
|
|
|
|
Edison uses a template-based approach to construct effective prompts:
|
|
|
|
```mermaid
|
|
sequenceDiagram
|
|
participant User
|
|
participant PromptManager
|
|
participant Template
|
|
participant APIClient
|
|
participant OpenAI
|
|
|
|
User->>APIClient: Query: "list all files"
|
|
APIClient->>PromptManager: get_full_prompt(query, shell)
|
|
PromptManager->>Template: Load template
|
|
Template-->>PromptManager: Template content
|
|
PromptManager->>PromptManager: Format template with query
|
|
PromptManager-->>APIClient: Formatted prompt
|
|
APIClient->>OpenAI: Send API request
|
|
OpenAI-->>APIClient: Command response
|
|
```
|
|
|
|
The `prompt_manager.py` module handles this:
|
|
|
|
```python
|
|
def load_prompt_template(shell="bash"):
|
|
"""Load the prompt template with shell-specific considerations."""
|
|
# Load and return the appropriate template
|
|
|
|
def get_full_prompt(query, shell="bash"):
|
|
"""Get the full prompt for the given query and shell."""
|
|
# Format prompt with query and shell
|
|
```
|
|
|
|
### Prompt Template
|
|
|
|
Edison uses a prompt template (`edison.prompt`) to structure requests to the AI model. The template:
|
|
|
|
1. Provides context about the desired output format
|
|
2. Includes examples of good responses
|
|
3. Specifies the shell environment
|
|
4. Encourages safe commands
|
|
5. Includes the user's query
|
|
|
|
## API Request
|
|
|
|
Edison uses the OpenAI Python client library for API requests:
|
|
|
|
```python
|
|
def call_api(client, config, query):
|
|
"""Call the OpenAI API with the given query."""
|
|
prompt = prompt_manager.get_full_prompt(query, config.get("shell", "bash"))
|
|
system_prompt = prompt.split('\n')[0] if '\n' in prompt else prompt
|
|
|
|
response = client.chat.completions.create(
|
|
model=config.get("model", "gpt-3.5-turbo"),
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
temperature=config.get("temperature", 0),
|
|
max_tokens=config.get("max_tokens", 500),
|
|
)
|
|
|
|
return response.choices[0].message.content.strip()
|
|
```
|
|
|
|
### Streaming API Integration
|
|
|
|
Edison also supports streaming command generation, which provides a more responsive user experience:
|
|
|
|
```python
|
|
def call_api_streaming(client, config, query, callback):
|
|
"""Call the OpenAI API with streaming enabled.
|
|
|
|
Args:
|
|
client: The OpenAI client
|
|
config: The configuration dictionary
|
|
query: The user query
|
|
callback: Function to call with each token
|
|
"""
|
|
prompt = prompt_manager.get_full_prompt(query, config.get("shell", "bash"))
|
|
system_prompt = prompt.split('\n')[0] if '\n' in prompt else prompt
|
|
|
|
response = client.chat.completions.create(
|
|
model=config.get("model", "gpt-3.5-turbo"),
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
temperature=config.get("temperature", 0),
|
|
max_tokens=config.get("max_tokens", 500),
|
|
stream=True # Enable streaming
|
|
)
|
|
|
|
# Collect the full response while calling the callback for each chunk
|
|
full_response = ""
|
|
|
|
for chunk in response:
|
|
if chunk.choices and chunk.choices[0].delta.content:
|
|
content = chunk.choices[0].delta.content
|
|
full_response += content
|
|
callback(content) # Call the callback with each chunk
|
|
|
|
return full_response.strip()
|
|
```
|
|
|
|
This streaming functionality is integrated with the UI through a callback function that updates the display in real-time as tokens are received.
|
|
|
|
### Request Parameters
|
|
|
|
The API request includes these key parameters:
|
|
|
|
| Parameter | Description | Default |
|
|
|-----------|-------------|---------|
|
|
| model | The OpenAI model to use | gpt-3.5-turbo |
|
|
| temperature | Randomness of completions (0-1) | 0 |
|
|
| max_tokens | Maximum tokens in response | 500 |
|
|
|
|
## Error Handling and Retries
|
|
|
|
The `generate_command()` function implements retry logic to handle rate limiting:
|
|
|
|
```python
|
|
def generate_command(client, config, query, max_retries=3):
|
|
"""Generate a command using the OpenAI API with retry logic."""
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return call_api(client, config, query)
|
|
except Exception as e:
|
|
if "rate limit" in str(e).lower() and attempt < max_retries - 1:
|
|
wait_time = 2 ** attempt # Exponential backoff
|
|
logger.warning(f"Rate limited. Retrying in {wait_time} seconds...")
|
|
time.sleep(wait_time)
|
|
else:
|
|
logger.error(f"Error after {attempt+1} attempts: {str(e)}")
|
|
raise
|
|
```
|
|
|
|
Key features:
|
|
- Exponential backoff for rate limits
|
|
- Maximum retry attempts
|
|
- Detailed error logging
|
|
|
|
## Command Processing and Validation
|
|
|
|
After receiving the API response, Edison:
|
|
|
|
1. Extracts the command from the response
|
|
2. Validates the command for safety
|
|
3. Checks for markdown or other formatting issues
|
|
4. Prepares the command for execution
|
|
|
|
## Extending the API Integration
|
|
|
|
To support additional AI providers or models:
|
|
|
|
1. **Create a new client factory function**:
|
|
```python
|
|
def create_anthropic_client(config):
|
|
# Initialize Anthropic client
|
|
```
|
|
|
|
2. **Add a model selection mechanism**:
|
|
```python
|
|
def get_client_for_model(config):
|
|
model = config.get("model", "gpt-3.5-turbo")
|
|
if model.startswith("claude"):
|
|
return create_anthropic_client(config)
|
|
else:
|
|
return create_client(config)
|
|
```
|
|
|
|
3. **Implement provider-specific API call function**:
|
|
```python
|
|
def call_anthropic_api(client, config, query):
|
|
# Format request for Anthropic API
|
|
```
|
|
|
|
4. **Update the command generation logic**:
|
|
```python
|
|
def generate_command(client, config, query, max_retries=3):
|
|
model = config.get("model", "gpt-3.5-turbo")
|
|
if model.startswith("claude"):
|
|
return call_anthropic_api(client, config, query)
|
|
else:
|
|
return call_api(client, config, query)
|
|
```
|
|
|
|
## API Response Examples
|
|
|
|
### Successful Response
|
|
|
|
```json
|
|
{
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": "ls -la",
|
|
"role": "assistant"
|
|
},
|
|
"index": 0,
|
|
"finish_reason": "stop"
|
|
}
|
|
]
|
|
}
|
|
```
|
|
|
|
### Error Response
|
|
|
|
```json
|
|
{
|
|
"error": {
|
|
"message": "Rate limit exceeded",
|
|
"type": "rate_limit_error",
|
|
"param": null,
|
|
"code": null
|
|
}
|
|
}
|
|
```
|
|
|
|
## Performance Considerations
|
|
|
|
To optimize API usage and performance:
|
|
|
|
1. **Use Efficient Models**: Default to `gpt-3.5-turbo` for lower latency
|
|
2. **Limit Token Usage**: Keep max_tokens reasonable (default 500)
|
|
3. **Request Caching**: Consider implementing caching for common queries
|
|
4. **Concurrent Requests**: For batch processing, consider async requests
|
|
5. **Prompt Optimization**: Keep prompt templates concise but effective |