12 Commits
Author SHA1 Message Date
rune a47b1c95a2 bug fixes. Added model info screen with pricing info 2026-02-05 14:43:37 +01:00
rune 06a3c898d3 Lot's of changes. None breaking. v3.0.0-b3 2026-02-05 11:21:22 +01:00
rune ecc2489eef Fixed default model setting. Added ctrl+y to copy latest reply in markdown++ 2026-02-04 15:16:26 +01:00
rune 1191fa6d19 Updated README 2026-02-04 11:26:08 +01:00
rune 6298158d3c oAI version 3.0 beta 1 2026-02-04 11:22:53 +01:00
runeandrune b0cf88704e 2.1 (#2)
Final release of version 2.1.

Headlights:

### Core Features
- 🤖 Interactive chat with 300+ AI models via OpenRouter
- 🔍 Model selection with search and filtering
- 💾 Conversation save/load/export (Markdown, JSON, HTML)
- 📎 File attachments (images, PDFs, code files)
- 💰 Real-time cost tracking and credit monitoring
- 🎨 Rich terminal UI with syntax highlighting
- 📝 Persistent command history with search (Ctrl+R)
- 🌐 Online mode (web search capabilities)
- 🧠 Conversation memory toggle

### MCP Integration
- 🔧 **File Mode**: AI can read, search, and list local files
  - Automatic .gitignore filtering
  - Virtual environment exclusion
  - Large file handling (auto-truncates >50KB)

- ✍️ **Write Mode**: AI can modify files with permission
  - Create, edit, delete files
  - Move, copy, organize files
  - Always requires explicit opt-in

- 🗄️ **Database Mode**: AI can query SQLite databases
  - Read-only access (safe)
  - Schema inspection
  - Full SQL query support

Reviewed-on: #2
Co-authored-by: Rune Olsen <rune@rune.pm>
Co-committed-by: Rune Olsen <rune@rune.pm>
2026-02-03 09:02:44 +01:00
runeandrune 1ef7918291 1.9.6 (#1)
New functionality and bugfixes.

Reviewed-on: #1
Co-authored-by: Rune Olsen <rune@rune.pm>
Co-committed-by: Rune Olsen <rune@rune.pm>
2025-12-30 15:46:40 +01:00
rune a6f0edd9f3 New functionality++ Verision bump 2025-12-23 15:08:20 +01:00
rune d4e43e6cb2 Update README.md 2025-12-21 20:24:26 +01:00
rune 36a412138d More info in models using . To use ]7;file://localhost/ at start of query use ]7;file://localhost/. Plus some other changes. 2025-12-21 19:21:14 +01:00
rune 229ffdf51a Small changes. WIP. Use relases page to download app 2025-12-18 14:06:12 +01:00
rune 53b6ae3a76 Added some more function. E.g. use of models. 2025-12-17 14:43:47 +01:00
58 changed files with 16754 additions and 1665 deletions
+24 -2
View File
@@ -22,5 +22,27 @@ Pipfile.lock # Consider if you want to include or exclude
._*
*~.nib
*~.xib
README.md.old
oai.zip
# Claude Code local settings
.claude/
# Added by author
*.zip
.note
diagnose.py
*.log
*.xml
build*
*.spec
compiled/
images/oai-iOS-Default-1024x1024@1x.png
images/oai.icon/
b0.sh
*.bak
*.old
*.sh
*.back
requirements.txt
system_prompt.txt
CLAUDE*
SESSION*_COMPLETE.md
+368 -155
View File
@@ -1,228 +1,441 @@
# oAI - OpenRouter AI Chat
# oAI - Open AI Chat Client
A terminal-based chat interface for OpenRouter API with conversation management, cost tracking, and rich formatting.
## Description
oAI is a command-line chat application that provides an interactive interface to OpenRouter's AI models. It features conversation persistence, file attachments, export capabilities, and detailed session metrics.
A powerful, modern **Textual TUI** chat client with **multi-provider support** (OpenRouter, Anthropic, OpenAI, Ollama) and **MCP (Model Context Protocol)** integration, enabling AI to access local files and query SQLite databases.
## Features
- Interactive chat with multiple AI models via OpenRouter
- Model selection with search functionality
- Conversation save/load/export (Markdown, JSON, HTML)
- File attachment support (code files and images)
- Session cost tracking and credit monitoring
- Rich terminal formatting with syntax highlighting
- Persistent command history
- Configurable system prompts and token limits
- SQLite-based configuration and conversation storage
### Core Features
- 🖥️ **Modern Textual TUI** with async streaming and beautiful interface
- 🔄 **Multi-Provider Support** - OpenRouter, Anthropic (Claude), OpenAI (ChatGPT), Ollama (local)
- 🤖 Interactive chat with 300+ AI models across providers
- 🔍 Model selection with search, filtering, and capability icons
- 💾 Conversation save/load/export (Markdown, JSON, HTML)
- 📎 File attachments (images, PDFs, code files)
- 💰 Real-time cost tracking and credit monitoring (OpenRouter)
- 🎨 Dark theme with syntax highlighting and Markdown rendering
- 📝 Command history navigation (Up/Down arrows)
- 🌐 **Universal Online Mode** - Web search for ALL providers:
- **Anthropic Native** - Built-in search with automatic citations ($0.01/search)
- **DuckDuckGo** - Free web scraping (all providers)
- **Google Custom Search** - Premium search option
- 🧠 Conversation memory toggle
- ⌨️ Keyboard shortcuts (F1=Help, F2=Models, Ctrl+S=Stats)
### MCP Integration
- 🔧 **File Mode**: AI can read, search, and list local files
- Automatic .gitignore filtering
- Virtual environment exclusion
- Large file handling (auto-truncates >50KB)
- ✍️ **Write Mode**: AI can modify files with permission
- Create, edit, delete files
- Move, copy, organize files
- Always requires explicit opt-in
- 🗄️ **Database Mode**: AI can query SQLite databases
- Read-only access (safe)
- Schema inspection
- Full SQL query support
## Requirements
- Python 3.7 or higher
- OpenRouter API key (get one at https://openrouter.ai)
## Screenshot
[<img src="https://gitlab.pm/rune/oai/raw/branch/main/images/screenshot_01.png">](https://gitlab.pm/rune/oai/src/branch/main/README.md)
Screenshot of `/help` screen.
- Python 3.10-3.13
- API key for your chosen provider:
- **OpenRouter**: [openrouter.ai](https://openrouter.ai)
- **Anthropic**: [console.anthropic.com](https://console.anthropic.com)
- **OpenAI**: [platform.openai.com](https://platform.openai.com)
- **Ollama**: No API key needed (local server)
## Installation
### 1. Install Dependencies
### Option 1: Pre-built Binary (macOS/Linux) (Recommended)
Use the included `requirements.txt` file to install the dependencies:
Download from [Releases](https://gitlab.pm/rune/oai/releases):
- **macOS (Apple Silicon)**: `oai_v3.0.0_mac_arm64.zip`
- **Linux (x86_64)**: `oai_v3.0.0_linux_x86_64.zip`
```bash
pip install -r requirements.txt
```
### 2. Make the Script Executable
```bash
chmod +x oai.py
```
### 3. Copy to PATH
Copy the script to a directory in your `$PATH` environment variable. Common locations include:
```bash
# Option 1: System-wide (requires sudo)
sudo cp oai.py /usr/local/bin/oai
# Option 2: User-local (recommended)
# Extract and install
unzip oai_v3.0.0_*.zip
mkdir -p ~/.local/bin
cp oai.py ~/.local/bin/oai
mv oai ~/.local/bin/
# Add to PATH if not already (add to ~/.bashrc or ~/.zshrc)
# macOS only: Remove quarantine and approve
xattr -cr ~/.local/bin/oai
# Then right-click oai in Finder → Open With → Terminal → Click "Open"
```
### Add to PATH
```bash
# Add to ~/.zshrc or ~/.bashrc
export PATH="$HOME/.local/bin:$PATH"
```
### 4. Verify Installation
### Option 2: Install from Source
```bash
oai
# Clone the repository
git clone https://gitlab.pm/rune/oai.git
cd oai
# Install with pip
pip install -e .
```
### 5. Alternative Installation (for *nix systems)
If you have issues with the above method you can add an alias in your `.bashrc`, `.zshrc` etc.
## Quick Start
```bash
alias oai='python3 <path to your file>'
# Start oAI (launches TUI)
oai
# Start with specific provider
oai --provider anthropic
oai --provider openai
oai --provider ollama
# Or with options
oai --provider openrouter --model gpt-4o --online --mcp
# Show version
oai version
```
On first run, you will be prompted to enter your OpenRouter API key.
On first run, you'll be prompted for your API key. Configure additional providers anytime with `/config`.
## Usage
### Starting the Application
### Enable Web Search (All Providers)
```bash
oai
# Using Anthropic native search (best quality, automatic citations)
/config search_provider anthropic_native
/online on
# Using free DuckDuckGo (works with all providers)
/config search_provider duckduckgo
/online on
# Using Google Custom Search (requires API key)
/config search_provider google
/config google_api_key YOUR_KEY
/config google_search_engine_id YOUR_ID
/online on
```
### Basic Commands
```
/help Show all available commands
/model Select an AI model
/config api Set OpenRouter API key
exit Quit the application
```bash
# In the TUI interface:
/provider # Show current provider or switch
/provider anthropic # Switch to Anthropic (Claude)
/provider openai # Switch to OpenAI (ChatGPT)
/provider ollama # Switch to Ollama (local)
/model # Select AI model (or press F2)
/online on # Enable web search
/help # Show all commands (or press F1)
/mcp on # Enable file/database access
/stats # View session statistics (or press Ctrl+S)
/config # View configuration settings
/credits # Check account credits (shows API balance or console link)
Ctrl+Q # Quit
```
### Configuration
## Web Search
All configuration is stored in `~/.config/oai/`:
- `oai_config.db` - SQLite database for settings and conversations
- `oai.log` - Application log file
- `history.txt` - Command history
oAI provides universal web search capabilities for all AI providers with three options:
### Common Workflows
### Anthropic Native Search (Recommended for Anthropic)
**Select a Model:**
```
/model
Anthropic's built-in web search API with automatic citations:
```bash
/config search_provider anthropic_native
/online on
# Now ask questions requiring current information
What are the latest developments in quantum computing?
```
**Paste from clipboard:**
Paste and send content to model
```
/paste
**Features:**
-**Automatic citations** - Claude cites its sources
-**Smart searching** - Claude decides when to search
-**Progressive searches** - Multiple searches for complex queries
-**Best quality** - Professional-grade results
**Pricing:** $10 per 1,000 searches ($0.01 per search) + token costs
**Note:** Only works with Anthropic provider (Claude models)
### DuckDuckGo Search (Default - Free)
Free web scraping that works with ALL providers:
```bash
/config search_provider duckduckgo # Default
/online on
# Works with Anthropic, OpenAI, Ollama, and OpenRouter
```
Paste with prompt and send content to model
```
/paste Analyze this text
**Features:**
-**Free** - No API key or costs
-**Universal** - Works with all providers
-**Privacy-friendly** - Uses DuckDuckGo
### Google Custom Search (Premium Option)
Google's Custom Search API for high-quality results:
```bash
/config search_provider google
/config google_api_key YOUR_GOOGLE_API_KEY
/config google_search_engine_id YOUR_SEARCH_ENGINE_ID
/online on
```
**Start Chatting:**
```
You> Hello, how are you?
Get your API key: [Google Custom Search API](https://developers.google.com/custom-search/v1/overview)
## MCP (Model Context Protocol)
MCP allows the AI to interact with your local files and databases.
### File Access
```bash
/mcp on # Enable MCP
/mcp add ~/Projects # Grant access to folder
/mcp list # View allowed folders
# Now ask the AI:
"List all Python files in Projects"
"Read and explain main.py"
"Search for files containing 'TODO'"
```
**Attach Files:**
```
You> Debug this code @/path/to/script.py
You> Analyze this image @/path/to/image.png
### Write Mode
```bash
/mcp write on # Enable file modifications
# AI can now:
"Create a new file called utils.py"
"Edit config.json and update the API URL"
"Delete the old backup files" # Always asks for confirmation
```
**Save Conversation:**
```
/save my_conversation
```
### Database Mode
**Export to File:**
```
/export md notes.md
/export json backup.json
/export html report.html
```
**View Session Stats:**
```bash
/mcp add db ~/app/data.db # Add database
/mcp db 1 # Switch to database mode
# Ask the AI:
"Show all tables"
"Find users created this month"
"What's the schema for the orders table?"
```
/stats
/credits
```
**Prevous commands input:**
Use the up/down arrows to see earlier `/`commands and earlier input to model and `<enter>` to execute the same command or resend the same input.
## Command Reference
Use `/help` within the application for a complete command reference organized by category:
- Session Commands
- Model Commands
- Configuration
- Token & System
- Conversation Management
- Monitoring & Stats
- File Attachments
### Chat Commands
| Command | Description |
|---------|-------------|
| `/help [cmd]` | Show help |
| `/model [search]` | Select model |
| `/info [model]` | Model details |
| `/memory on\|off` | Toggle context |
| `/online on\|off` | Toggle web search |
| `/retry` | Resend last message |
| `/clear` | Clear screen |
## Configuration Options
### MCP Commands
| Command | Description |
|---------|-------------|
| `/mcp on\|off` | Enable/disable MCP |
| `/mcp status` | Show MCP status |
| `/mcp add <path>` | Add folder |
| `/mcp add db <path>` | Add database |
| `/mcp list` | List folders |
| `/mcp db list` | List databases |
| `/mcp db <n>` | Switch to database |
| `/mcp files` | Switch to file mode |
| `/mcp write on\|off` | Toggle write mode |
- API Key: `/config api`
- Base URL: `/config url`
- Streaming: `/config stream on|off`
- Default Model: `/config model`
- Cost Warning: `/config costwarning <amount>`
- Max Token Limit: `/config maxtoken <value>`
### Conversation Commands
| Command | Description |
|---------|-------------|
| `/save <name>` | Save conversation |
| `/load <name>` | Load conversation |
| `/list` | List saved conversations |
| `/delete <name>` | Delete conversation |
| `/export md\|json\|html <file>` | Export |
## File Support
### Configuration
| Command | Description |
|---------|-------------|
| `/config` | View settings |
| `/config provider <name>` | Set default provider |
| `/config openrouter_api_key` | Set OpenRouter API key |
| `/config anthropic_api_key` | Set Anthropic API key |
| `/config openai_api_key` | Set OpenAI API key |
| `/config ollama_base_url` | Set Ollama server URL |
| `/config search_provider <provider>` | Set search provider (anthropic_native/duckduckgo/google) |
| `/config google_api_key` | Set Google API key (for Google search) |
| `/config online on\|off` | Set default online mode |
| `/config model <id>` | Set default model |
| `/config stream on\|off` | Toggle streaming |
| `/stats` | Session statistics |
| `/credits` | Check credits (OpenRouter) |
**Supported Code Extensions:**
.py, .js, .ts, .cs, .java, .c, .cpp, .h, .hpp, .rb, .ruby, .php, .swift, .kt, .kts, .go, .sh, .bat, .ps1, .R, .scala, .pl, .lua, .dart, .elm, .xml, .json, .yaml, .yml, .md, .txt
## CLI Options
**Image Support:**
Any image format with proper MIME type (PNG, JPEG, GIF, etc.)
```bash
oai [OPTIONS]
## Data Storage
Options:
-p, --provider TEXT Provider to use (openrouter/anthropic/openai/ollama)
-m, --model TEXT Model ID to use
-s, --system TEXT System prompt
-o, --online Enable online mode (OpenRouter only)
--mcp Enable MCP server
-v, --version Show version
--help Show help
```
- Configuration: `~/.config/oai/oai_config.db`
- Logs: `~/.config/oai/oai.log`
- History: `~/.config/oai/history.txt`
Commands:
```bash
oai # Launch TUI (default)
oai version # Show version information
oai --help # Show help message
```
## Configuration
Configuration is stored in `~/.config/oai/`:
| File | Purpose |
|------|---------|
| `oai_config.db` | Settings, conversations, MCP config |
| `oai.log` | Application logs |
| `history.txt` | Command history |
## Project Structure
```
oai/
├── oai/
│ ├── __init__.py
│ ├── __main__.py # Entry point for python -m oai
│ ├── cli.py # Main CLI entry point
│ ├── constants.py # Configuration constants
│ ├── commands/ # Slash command handlers
│ ├── config/ # Settings and database
│ ├── core/ # Chat client and session
│ ├── mcp/ # MCP server and tools
│ ├── providers/ # AI provider abstraction
│ ├── tui/ # Textual TUI interface
│ │ ├── app.py # Main TUI application
│ │ ├── widgets/ # Custom widgets
│ │ ├── screens/ # Modal screens
│ │ └── styles.tcss # TUI styling
│ └── utils/ # Logging, export, etc.
├── pyproject.toml # Package configuration
├── build.sh # Binary build script
└── README.md
```
## Troubleshooting
### macOS Binary Issues
```bash
# Remove quarantine attribute
xattr -cr ~/.local/bin/oai
# Then in Finder: right-click oai → Open With → Terminal → Click "Open"
# After this, oai works from any terminal
```
### MCP Not Working
```bash
# Check if model supports function calling
/info # Look for "tools" in supported parameters
# Check MCP status
/mcp status
# View logs
tail -f ~/.config/oai/oai.log
```
### Import Errors
```bash
# Reinstall package
pip install -e . --force-reinstall
```
## Version History
### v3.0.0-b3 (Current - Beta 3)
- 🔄 **Multi-Provider Support** - OpenRouter, Anthropic (Claude), OpenAI (ChatGPT), Ollama (local)
- 🔌 **Provider Switching** - Switch between providers mid-session with `/provider`
- 🧠 **Provider Model Memory** - Remembers last used model per provider
- ⚙️ **Provider Configuration** - Separate API keys for each provider
- 🌐 **Universal Web Search** - Three search options for all providers:
- **Anthropic Native** - Built-in search with citations ($0.01/search)
- **DuckDuckGo** - Free web scraping (default)
- **Google Custom Search** - Premium option
- 🎨 **Enhanced Header** - Shows current provider and model
- 📊 **Updated Config Screen** - Shows provider-specific settings
- 🔧 **Command Dropdown** - Added `/provider` commands for discoverability
- 🎯 **Auto-Scrolling** - Chat window auto-scrolls during streaming responses
- 🎨 **Refined UI** - Thinner scrollbar, cleaner styling
- 🔧 **Updated Models** - Claude 4.5 (Sonnet, Haiku, Opus)
### v3.0.0 (Beta)
- 🎨 **Complete migration to Textual TUI** - Modern async terminal interface
- 🗑️ **Removed CLI interface** - TUI-only for cleaner codebase (11.6% smaller)
- 🖱️ **Modal screens** - Help, stats, config, credits, model selector
- ⌨️ **Keyboard shortcuts** - F1 (help), F2 (models), Ctrl+S (stats), etc.
- 🎯 **Capability indicators** - Visual icons for model features (vision, tools, online)
- 🎨 **Consistent dark theme** - Professional styling throughout
- 📊 **Enhanced model selector** - Search, filter, capability columns
- 🚀 **Default command** - Just run `oai` to launch TUI
- 🧹 **Code cleanup** - Removed 1,300+ lines of CLI code
### v2.1.0
- 🏗️ Complete codebase refactoring to modular package structure
- 🔌 Extensible provider architecture for adding new AI providers
- 📦 Proper Python packaging with pyproject.toml
- ✨ MCP integration (file access, write mode, database queries)
- 🔧 Command registry pattern for slash commands
- 📊 Improved cost tracking and session statistics
### v1.9.x
- Single-file implementation
- Core chat functionality
- File attachments
- Conversation management
## License
MIT License
Copyright (c) 2024 Rune Olsen
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
Full license: https://opensource.org/licenses/MIT
MIT License - See [LICENSE](LICENSE) for details.
## Author
**Rune Olsen**
- Project: https://iurl.no/oai
- Repository: https://gitlab.pm/rune/oai
Blog: https://blog.rune.pm
## Contributing
## Version
1. Fork the repository
2. Create a feature branch
3. Submit a pull request
1.0
---
## Support
For issues, questions, or contributions, visit https://iurl.no/oai and create an issue.
**⭐ Star this project if you find it useful!**
-1471
View File
File diff suppressed because it is too large Load Diff
+26
View File
@@ -0,0 +1,26 @@
"""
oAI - Open AI Chat Client
A feature-rich terminal-based chat application with multi-provider support
(OpenRouter, Anthropic, OpenAI, Ollama) and advanced Model Context Protocol (MCP)
integration for filesystem and database access.
Author: Rune
License: MIT
"""
__version__ = "3.0.0-b3"
__author__ = "Rune Olsen"
__license__ = "MIT"
# Lazy imports to avoid circular dependencies and improve startup time
# Full imports are available via submodules:
# from oai.config import Settings, Database
# from oai.providers import OpenRouterProvider, AIProvider
# from oai.mcp import MCPManager
__all__ = [
"__version__",
"__author__",
"__license__",
]
+8
View File
@@ -0,0 +1,8 @@
"""
Entry point for running oAI as a module: python -m oai
"""
from oai.cli import main
if __name__ == "__main__":
main()
+226
View File
@@ -0,0 +1,226 @@
"""
Main CLI entry point for oAI.
This module provides the command-line interface for the oAI TUI application.
"""
import sys
from typing import Optional
import typer
from oai import __version__
from oai.commands import register_all_commands
from oai.config.settings import Settings
from oai.constants import APP_URL, APP_VERSION
from oai.core.client import AIClient
from oai.core.session import ChatSession
from oai.mcp.manager import MCPManager
from oai.utils.logging import LoggingManager, get_logger
# Create Typer app
app = typer.Typer(
name="oai",
help=f"oAI - Open AI Chat Client (TUI)\n\nVersion: {APP_VERSION}",
add_completion=False,
epilog="For more information, visit: " + APP_URL,
)
@app.callback(invoke_without_command=True)
def main_callback(
ctx: typer.Context,
version_flag: bool = typer.Option(
False,
"--version",
"-v",
help="Show version information",
is_flag=True,
),
model: Optional[str] = typer.Option(
None,
"--model",
"-m",
help="Model ID to use",
),
system: Optional[str] = typer.Option(
None,
"--system",
"-s",
help="System prompt",
),
online: bool = typer.Option(
False,
"--online",
"-o",
help="Enable online mode",
),
mcp: bool = typer.Option(
False,
"--mcp",
help="Enable MCP server",
),
provider: Optional[str] = typer.Option(
None,
"--provider",
"-p",
help="AI provider to use (openrouter, anthropic, openai, ollama)",
),
) -> None:
"""Main callback - launches TUI by default."""
if version_flag:
typer.echo(f"oAI version {APP_VERSION}")
raise typer.Exit()
# If no subcommand provided, launch TUI
if ctx.invoked_subcommand is None:
_launch_tui(model, system, online, mcp, provider)
def _launch_tui(
model: Optional[str] = None,
system: Optional[str] = None,
online: bool = False,
mcp: bool = False,
provider: Optional[str] = None,
) -> None:
"""Launch the Textual TUI interface."""
from oai.constants import VALID_PROVIDERS
# Setup logging
logging_manager = LoggingManager()
logging_manager.setup()
logger = get_logger()
# Load settings
settings = Settings.load()
# Determine provider
selected_provider = provider or settings.default_provider
# Validate provider
if selected_provider not in VALID_PROVIDERS:
typer.echo(f"Error: Invalid provider: {selected_provider}", err=True)
typer.echo(f"Valid providers: {', '.join(VALID_PROVIDERS)}", err=True)
raise typer.Exit(1)
# Build provider API keys dict
provider_api_keys = {
"openrouter": settings.openrouter_api_key,
"anthropic": settings.anthropic_api_key,
"openai": settings.openai_api_key,
}
# Check if provider is configured (except Ollama which doesn't need API key)
if selected_provider != "ollama":
if not provider_api_keys.get(selected_provider):
typer.echo(f"Error: No API key configured for {selected_provider}", err=True)
typer.echo(f"Set it with: oai config {selected_provider}_api_key <key>", err=True)
raise typer.Exit(1)
# Initialize client
try:
client = AIClient(
provider_name=selected_provider,
provider_api_keys=provider_api_keys,
ollama_base_url=settings.ollama_base_url,
)
except Exception as e:
typer.echo(f"Error: Failed to initialize client: {e}", err=True)
raise typer.Exit(1)
# Register commands
register_all_commands()
# Initialize MCP manager (always create it, even if not enabled)
mcp_manager = MCPManager()
if mcp:
try:
result = mcp_manager.enable()
if result["success"]:
logger.info("MCP server enabled in files mode")
else:
logger.warning(f"MCP: {result.get('error', 'Failed to enable')}")
except Exception as e:
logger.warning(f"Failed to enable MCP: {e}")
# Create session with MCP manager
session = ChatSession(
client=client,
settings=settings,
mcp_manager=mcp_manager,
)
# Set system prompt if provided
if system:
session.set_system_prompt(system)
# Enable online mode if requested
if online:
session.online_enabled = True
# Set model if specified, otherwise use default
if model:
raw_model = client.get_raw_model(model)
if raw_model:
session.set_model(raw_model)
else:
logger.warning(f"Model '{model}' not found")
elif settings.default_model:
raw_model = client.get_raw_model(settings.default_model)
if raw_model:
session.set_model(raw_model)
else:
logger.warning(f"Default model '{settings.default_model}' not available")
# Run Textual app
from oai.tui.app import oAIChatApp
app_instance = oAIChatApp(session, settings, model)
app_instance.run()
@app.command()
def tui(
model: Optional[str] = typer.Option(
None,
"--model",
"-m",
help="Model ID to use",
),
system: Optional[str] = typer.Option(
None,
"--system",
"-s",
help="System prompt",
),
online: bool = typer.Option(
False,
"--online",
"-o",
help="Enable online mode",
),
mcp: bool = typer.Option(
False,
"--mcp",
help="Enable MCP server",
),
) -> None:
"""Start Textual TUI interface (alias for just running 'oai')."""
_launch_tui(model, system, online, mcp)
@app.command()
def version() -> None:
"""Show version information."""
typer.echo(f"oAI version {APP_VERSION}")
typer.echo(f"Visit {APP_URL} for more information")
def main() -> None:
"""Entry point for the CLI."""
app()
if __name__ == "__main__":
main()
+24
View File
@@ -0,0 +1,24 @@
"""
Command system for oAI.
This module provides a command registry and handler system
for processing slash commands in the chat interface.
"""
from oai.commands.registry import (
Command,
CommandRegistry,
CommandContext,
CommandResult,
registry,
)
from oai.commands.handlers import register_all_commands
__all__ = [
"Command",
"CommandRegistry",
"CommandContext",
"CommandResult",
"registry",
"register_all_commands",
]
File diff suppressed because it is too large Load Diff
+384
View File
@@ -0,0 +1,384 @@
"""
Command registry for oAI.
This module defines the command system infrastructure including
the Command base class, CommandContext for state, and CommandRegistry
for managing available commands.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING
from oai.utils.logging import get_logger
if TYPE_CHECKING:
from oai.config.settings import Settings
from oai.providers.base import AIProvider, ModelInfo
from oai.mcp.manager import MCPManager
class CommandStatus(str, Enum):
"""Status of command execution."""
SUCCESS = "success"
ERROR = "error"
CONTINUE = "continue" # Continue to next handler
EXIT = "exit" # Exit the application
@dataclass
class CommandResult:
"""
Result of a command execution.
Attributes:
status: Execution status
message: Optional message to display
data: Optional data payload
should_continue: Whether to continue the main loop
"""
status: CommandStatus = CommandStatus.SUCCESS
message: Optional[str] = None
data: Optional[Any] = None
should_continue: bool = True
@classmethod
def success(cls, message: Optional[str] = None, data: Any = None) -> "CommandResult":
"""Create a success result."""
return cls(status=CommandStatus.SUCCESS, message=message, data=data)
@classmethod
def error(cls, message: str) -> "CommandResult":
"""Create an error result."""
return cls(status=CommandStatus.ERROR, message=message)
@classmethod
def exit(cls, message: Optional[str] = None) -> "CommandResult":
"""Create an exit result."""
return cls(status=CommandStatus.EXIT, message=message, should_continue=False)
@dataclass
class CommandContext:
"""
Context object providing state to command handlers.
Contains all the session state needed by commands including
settings, provider, conversation history, and MCP manager.
Attributes:
settings: Application settings
provider: AI provider instance
mcp_manager: MCP manager instance
selected_model: Currently selected model
session_history: Conversation history
session_system_prompt: Current system prompt
memory_enabled: Whether memory is enabled
online_enabled: Whether online mode is enabled
session_tokens: Session token counts
session_cost: Session cost total
session: Reference to the ChatSession (for provider switching, etc.)
"""
settings: Optional["Settings"] = None
provider: Optional["AIProvider"] = None
mcp_manager: Optional["MCPManager"] = None
selected_model: Optional["ModelInfo"] = None
selected_model_raw: Optional[Dict[str, Any]] = None
session_history: List[Dict[str, Any]] = field(default_factory=list)
session_system_prompt: str = ""
memory_enabled: bool = True
memory_start_index: int = 0
online_enabled: bool = False
middle_out_enabled: bool = False
session_max_token: int = 0
total_input_tokens: int = 0
total_output_tokens: int = 0
total_cost: float = 0.0
message_count: int = 0
is_tui: bool = False # Flag for TUI mode
current_index: int = 0
session: Optional[Any] = None # Reference to ChatSession (avoid circular import)
@dataclass
class CommandHelp:
"""
Help information for a command.
Attributes:
description: Brief description
usage: Usage syntax
examples: List of (description, example) tuples
notes: Additional notes
aliases: Command aliases
"""
description: str
usage: str = ""
examples: List[tuple] = field(default_factory=list)
notes: str = ""
aliases: List[str] = field(default_factory=list)
class Command(ABC):
"""
Abstract base class for all commands.
Commands implement the execute method to handle their logic.
They can also provide help information and aliases.
"""
@property
@abstractmethod
def name(self) -> str:
"""Get the primary command name (e.g., '/help')."""
pass
@property
def aliases(self) -> List[str]:
"""Get command aliases (e.g., ['/h'] for help)."""
return []
@property
@abstractmethod
def help(self) -> CommandHelp:
"""Get command help information."""
pass
@abstractmethod
def execute(self, args: str, context: CommandContext) -> CommandResult:
"""
Execute the command.
Args:
args: Arguments passed to the command
context: Command execution context
Returns:
CommandResult indicating success/failure
"""
pass
def matches(self, input_text: str) -> bool:
"""
Check if this command matches the input.
Args:
input_text: User input text
Returns:
True if this command should handle the input
"""
input_lower = input_text.lower()
cmd_word = input_lower.split()[0] if input_lower.split() else ""
# Check primary name
if cmd_word == self.name.lower():
return True
# Check aliases
for alias in self.aliases:
if cmd_word == alias.lower():
return True
return False
def get_args(self, input_text: str) -> str:
"""
Extract arguments from the input text.
Args:
input_text: Full user input
Returns:
Arguments portion of the input
"""
parts = input_text.split(maxsplit=1)
return parts[1] if len(parts) > 1 else ""
class CommandRegistry:
"""
Registry for managing available commands.
Provides registration, lookup, and execution of commands.
"""
def __init__(self):
"""Initialize an empty command registry."""
self._commands: Dict[str, Command] = {}
self._aliases: Dict[str, str] = {}
self.logger = get_logger()
def register(self, command: Command) -> None:
"""
Register a command.
Args:
command: Command instance to register
Raises:
ValueError: If command name already registered
"""
name = command.name.lower()
if name in self._commands:
raise ValueError(f"Command '{name}' already registered")
self._commands[name] = command
# Register aliases
for alias in command.aliases:
alias_lower = alias.lower()
if alias_lower in self._aliases:
self.logger.warning(
f"Alias '{alias}' already registered, overwriting"
)
self._aliases[alias_lower] = name
self.logger.debug(f"Registered command: {name}")
def register_function(
self,
name: str,
handler: Callable[[str, CommandContext], CommandResult],
description: str,
usage: str = "",
aliases: Optional[List[str]] = None,
examples: Optional[List[tuple]] = None,
notes: str = "",
) -> None:
"""
Register a function-based command.
Convenience method for simple commands that don't need
a full Command class.
Args:
name: Command name (e.g., '/help')
handler: Function to execute
description: Help description
usage: Usage syntax
aliases: Command aliases
examples: Example usages
notes: Additional notes
"""
aliases = aliases or []
examples = examples or []
class FunctionCommand(Command):
@property
def name(self) -> str:
return name
@property
def aliases(self) -> List[str]:
return aliases
@property
def help(self) -> CommandHelp:
return CommandHelp(
description=description,
usage=usage,
examples=examples,
notes=notes,
aliases=aliases,
)
def execute(self, args: str, context: CommandContext) -> CommandResult:
return handler(args, context)
self.register(FunctionCommand())
def get(self, name: str) -> Optional[Command]:
"""
Get a command by name or alias.
Args:
name: Command name or alias
Returns:
Command instance or None if not found
"""
name_lower = name.lower()
# Check direct match
if name_lower in self._commands:
return self._commands[name_lower]
# Check aliases
if name_lower in self._aliases:
return self._commands[self._aliases[name_lower]]
return None
def find(self, input_text: str) -> Optional[Command]:
"""
Find a command that matches the input.
Args:
input_text: User input text
Returns:
Matching Command or None
"""
cmd_word = input_text.lower().split()[0] if input_text.split() else ""
return self.get(cmd_word)
def execute(self, input_text: str, context: CommandContext) -> Optional[CommandResult]:
"""
Execute a command matching the input.
Args:
input_text: User input text
context: Execution context
Returns:
CommandResult or None if no matching command
"""
command = self.find(input_text)
if command:
args = command.get_args(input_text)
self.logger.debug(f"Executing command: {command.name} with args: {args}")
return command.execute(args, context)
return None
def is_command(self, input_text: str) -> bool:
"""
Check if input is a valid command.
Args:
input_text: User input text
Returns:
True if input matches a registered command
"""
return self.find(input_text) is not None
def list_commands(self) -> List[Command]:
"""
Get all registered commands.
Returns:
List of Command instances
"""
return list(self._commands.values())
def get_all_names(self) -> List[str]:
"""
Get all command names and aliases.
Returns:
List of command names including aliases
"""
names = list(self._commands.keys())
names.extend(self._aliases.keys())
return sorted(set(names))
# Global registry instance
registry = CommandRegistry()
+11
View File
@@ -0,0 +1,11 @@
"""
Configuration management for oAI.
This package handles all configuration persistence, settings management,
and database operations for the application.
"""
from oai.config.settings import Settings
from oai.config.database import Database
__all__ = ["Settings", "Database"]
+472
View File
@@ -0,0 +1,472 @@
"""
Database persistence layer for oAI.
This module provides a clean abstraction for SQLite operations including
configuration storage, conversation persistence, and MCP statistics tracking.
All database operations are centralized here for maintainability.
"""
import sqlite3
import json
import datetime
from pathlib import Path
from typing import Optional, List, Dict, Any
from contextlib import contextmanager
from oai.constants import DATABASE_FILE, CONFIG_DIR
class Database:
"""
SQLite database manager for oAI.
Handles all database operations including:
- Configuration key-value storage
- Conversation session persistence
- MCP configuration and statistics
- Database registrations for MCP
Uses context managers for safe connection handling and supports
automatic table creation on first use.
"""
def __init__(self, db_path: Optional[Path] = None):
"""
Initialize the database manager.
Args:
db_path: Optional custom database path. Defaults to standard location.
"""
self.db_path = db_path or DATABASE_FILE
self._ensure_directories()
self._ensure_tables()
def _ensure_directories(self) -> None:
"""Ensure the configuration directory exists."""
CONFIG_DIR.mkdir(parents=True, exist_ok=True)
@contextmanager
def _connection(self):
"""
Context manager for database connections.
Yields:
sqlite3.Connection: Active database connection
Example:
with self._connection() as conn:
conn.execute("SELECT * FROM config")
"""
conn = sqlite3.connect(str(self.db_path))
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
finally:
conn.close()
def _ensure_tables(self) -> None:
"""Create all required tables if they don't exist."""
with self._connection() as conn:
# Main configuration table
conn.execute("""
CREATE TABLE IF NOT EXISTS config (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
)
""")
# Conversation sessions table
conn.execute("""
CREATE TABLE IF NOT EXISTS conversation_sessions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
timestamp TEXT NOT NULL,
data TEXT NOT NULL
)
""")
# MCP configuration table
conn.execute("""
CREATE TABLE IF NOT EXISTS mcp_config (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
)
""")
# MCP statistics table
conn.execute("""
CREATE TABLE IF NOT EXISTS mcp_stats (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp TEXT NOT NULL,
tool_name TEXT NOT NULL,
folder TEXT,
success INTEGER NOT NULL,
error_message TEXT
)
""")
# MCP databases table
conn.execute("""
CREATE TABLE IF NOT EXISTS mcp_databases (
id INTEGER PRIMARY KEY AUTOINCREMENT,
path TEXT NOT NULL UNIQUE,
name TEXT NOT NULL,
size INTEGER,
tables TEXT,
added_timestamp TEXT NOT NULL
)
""")
# =========================================================================
# CONFIGURATION METHODS
# =========================================================================
def get_config(self, key: str) -> Optional[str]:
"""
Retrieve a configuration value by key.
Args:
key: The configuration key to retrieve
Returns:
The configuration value, or None if not found
"""
with self._connection() as conn:
cursor = conn.execute(
"SELECT value FROM config WHERE key = ?",
(key,)
)
result = cursor.fetchone()
return result[0] if result else None
def set_config(self, key: str, value: str) -> None:
"""
Set a configuration value.
Args:
key: The configuration key
value: The value to store
"""
with self._connection() as conn:
conn.execute(
"INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)",
(key, value)
)
def delete_config(self, key: str) -> bool:
"""
Delete a configuration value.
Args:
key: The configuration key to delete
Returns:
True if a row was deleted, False otherwise
"""
with self._connection() as conn:
cursor = conn.execute(
"DELETE FROM config WHERE key = ?",
(key,)
)
return cursor.rowcount > 0
def get_all_config(self) -> Dict[str, str]:
"""
Retrieve all configuration values.
Returns:
Dictionary of all key-value pairs
"""
with self._connection() as conn:
cursor = conn.execute("SELECT key, value FROM config")
return dict(cursor.fetchall())
# =========================================================================
# MCP CONFIGURATION METHODS
# =========================================================================
def get_mcp_config(self, key: str) -> Optional[str]:
"""
Retrieve an MCP configuration value.
Args:
key: The MCP configuration key
Returns:
The configuration value, or None if not found
"""
with self._connection() as conn:
cursor = conn.execute(
"SELECT value FROM mcp_config WHERE key = ?",
(key,)
)
result = cursor.fetchone()
return result[0] if result else None
def set_mcp_config(self, key: str, value: str) -> None:
"""
Set an MCP configuration value.
Args:
key: The MCP configuration key
value: The value to store
"""
with self._connection() as conn:
conn.execute(
"INSERT OR REPLACE INTO mcp_config (key, value) VALUES (?, ?)",
(key, value)
)
# =========================================================================
# MCP STATISTICS METHODS
# =========================================================================
def log_mcp_stat(
self,
tool_name: str,
folder: Optional[str],
success: bool,
error_message: Optional[str] = None
) -> None:
"""
Log an MCP tool usage event.
Args:
tool_name: Name of the MCP tool that was called
folder: The folder path involved (if any)
success: Whether the call succeeded
error_message: Error message if the call failed
"""
timestamp = datetime.datetime.now().isoformat()
with self._connection() as conn:
conn.execute(
"""INSERT INTO mcp_stats
(timestamp, tool_name, folder, success, error_message)
VALUES (?, ?, ?, ?, ?)""",
(timestamp, tool_name, folder, 1 if success else 0, error_message)
)
def get_mcp_stats(self) -> Dict[str, Any]:
"""
Get aggregated MCP usage statistics.
Returns:
Dictionary containing usage statistics:
- total_calls: Total number of tool calls
- reads: Number of file reads
- lists: Number of directory listings
- searches: Number of file searches
- db_inspects: Number of database inspections
- db_searches: Number of database searches
- db_queries: Number of database queries
- last_used: Timestamp of last usage
"""
with self._connection() as conn:
cursor = conn.execute("""
SELECT
COUNT(*) as total_calls,
SUM(CASE WHEN tool_name = 'read_file' THEN 1 ELSE 0 END) as reads,
SUM(CASE WHEN tool_name = 'list_directory' THEN 1 ELSE 0 END) as lists,
SUM(CASE WHEN tool_name = 'search_files' THEN 1 ELSE 0 END) as searches,
SUM(CASE WHEN tool_name = 'inspect_database' THEN 1 ELSE 0 END) as db_inspects,
SUM(CASE WHEN tool_name = 'search_database' THEN 1 ELSE 0 END) as db_searches,
SUM(CASE WHEN tool_name = 'query_database' THEN 1 ELSE 0 END) as db_queries,
MAX(timestamp) as last_used
FROM mcp_stats
""")
row = cursor.fetchone()
return {
"total_calls": row[0] or 0,
"reads": row[1] or 0,
"lists": row[2] or 0,
"searches": row[3] or 0,
"db_inspects": row[4] or 0,
"db_searches": row[5] or 0,
"db_queries": row[6] or 0,
"last_used": row[7],
}
# =========================================================================
# MCP DATABASE REGISTRY METHODS
# =========================================================================
def add_mcp_database(self, db_info: Dict[str, Any]) -> int:
"""
Register a database for MCP access.
Args:
db_info: Dictionary containing:
- path: Database file path
- name: Display name
- size: File size in bytes
- tables: List of table names
- added: Timestamp when added
Returns:
The database ID
"""
with self._connection() as conn:
conn.execute(
"""INSERT INTO mcp_databases
(path, name, size, tables, added_timestamp)
VALUES (?, ?, ?, ?, ?)""",
(
db_info["path"],
db_info["name"],
db_info["size"],
json.dumps(db_info["tables"]),
db_info["added"]
)
)
cursor = conn.execute(
"SELECT id FROM mcp_databases WHERE path = ?",
(db_info["path"],)
)
return cursor.fetchone()[0]
def remove_mcp_database(self, db_path: str) -> bool:
"""
Remove a database from the MCP registry.
Args:
db_path: Path to the database file
Returns:
True if a row was deleted, False otherwise
"""
with self._connection() as conn:
cursor = conn.execute(
"DELETE FROM mcp_databases WHERE path = ?",
(db_path,)
)
return cursor.rowcount > 0
def get_mcp_databases(self) -> List[Dict[str, Any]]:
"""
Retrieve all registered MCP databases.
Returns:
List of database information dictionaries
"""
with self._connection() as conn:
cursor = conn.execute(
"""SELECT id, path, name, size, tables, added_timestamp
FROM mcp_databases ORDER BY id"""
)
databases = []
for row in cursor.fetchall():
tables_list = json.loads(row[4]) if row[4] else []
databases.append({
"id": row[0],
"path": row[1],
"name": row[2],
"size": row[3],
"tables": tables_list,
"added": row[5],
})
return databases
# =========================================================================
# CONVERSATION METHODS
# =========================================================================
def save_conversation(self, name: str, data: List[Dict[str, str]]) -> None:
"""
Save a conversation session.
Args:
name: Name/identifier for the conversation
data: List of message dictionaries with 'prompt' and 'response' keys
"""
timestamp = datetime.datetime.now().isoformat()
data_json = json.dumps(data)
with self._connection() as conn:
conn.execute(
"""INSERT INTO conversation_sessions
(name, timestamp, data) VALUES (?, ?, ?)""",
(name, timestamp, data_json)
)
def load_conversation(self, name: str) -> Optional[List[Dict[str, str]]]:
"""
Load a conversation by name.
Args:
name: Name of the conversation to load
Returns:
List of messages, or None if not found
"""
with self._connection() as conn:
cursor = conn.execute(
"""SELECT data FROM conversation_sessions
WHERE name = ?
ORDER BY timestamp DESC LIMIT 1""",
(name,)
)
result = cursor.fetchone()
if result:
return json.loads(result[0])
return None
def delete_conversation(self, name: str) -> int:
"""
Delete a conversation by name.
Args:
name: Name of the conversation to delete
Returns:
Number of rows deleted
"""
with self._connection() as conn:
cursor = conn.execute(
"DELETE FROM conversation_sessions WHERE name = ?",
(name,)
)
return cursor.rowcount
def list_conversations(self) -> List[Dict[str, Any]]:
"""
List all saved conversations.
Returns:
List of conversation summaries with name, timestamp, and message_count
"""
with self._connection() as conn:
cursor = conn.execute("""
SELECT name, MAX(timestamp) as last_saved, data
FROM conversation_sessions
GROUP BY name
ORDER BY last_saved DESC
""")
conversations = []
for row in cursor.fetchall():
name, timestamp, data_json = row
data = json.loads(data_json)
conversations.append({
"name": name,
"timestamp": timestamp,
"message_count": len(data),
})
return conversations
# Global database instance for convenience
_db: Optional[Database] = None
def get_database() -> Database:
"""
Get the global database instance.
Returns:
The shared Database instance
"""
global _db
if _db is None:
_db = Database()
return _db
+572
View File
@@ -0,0 +1,572 @@
"""
Settings management for oAI.
This module provides a centralized settings class that handles all application
configuration with type safety, validation, and persistence.
"""
from dataclasses import dataclass, field
from typing import Optional, Dict
from pathlib import Path
import json
from oai.constants import (
DEFAULT_BASE_URL,
DEFAULT_STREAM_ENABLED,
DEFAULT_MAX_TOKENS,
DEFAULT_ONLINE_MODE,
DEFAULT_COST_WARNING_THRESHOLD,
DEFAULT_LOG_MAX_SIZE_MB,
DEFAULT_LOG_BACKUP_COUNT,
DEFAULT_LOG_LEVEL,
DEFAULT_SYSTEM_PROMPT,
VALID_LOG_LEVELS,
DEFAULT_PROVIDER,
OLLAMA_DEFAULT_URL,
)
from oai.config.database import get_database
@dataclass
class Settings:
"""
Application settings with persistence support.
This class provides a clean interface for managing all configuration
options. Settings are automatically loaded from the database on
initialization and can be persisted back.
Attributes:
api_key: Legacy OpenRouter API key (deprecated, use openrouter_api_key)
base_url: API base URL
default_model: Default model ID to use
default_system_prompt: Custom system prompt (None = use hardcoded default, "" = blank)
stream_enabled: Whether to stream responses
max_tokens: Maximum tokens per request
cost_warning_threshold: Alert threshold for message cost
default_online_mode: Whether online mode is enabled by default
log_max_size_mb: Maximum log file size in MB
log_backup_count: Number of log file backups to keep
log_level: Logging level (debug/info/warning/error/critical)
# Provider-specific settings
default_provider: Default AI provider to use
openrouter_api_key: OpenRouter API key
anthropic_api_key: Anthropic API key
openai_api_key: OpenAI API key
ollama_base_url: Ollama server URL
"""
# Legacy field (kept for backward compatibility)
api_key: Optional[str] = None
# Provider configuration
default_provider: str = DEFAULT_PROVIDER
openrouter_api_key: Optional[str] = None
anthropic_api_key: Optional[str] = None
openai_api_key: Optional[str] = None
ollama_base_url: str = OLLAMA_DEFAULT_URL
provider_models: Dict[str, str] = field(default_factory=dict) # provider -> last_model_id
# Web search configuration (for online mode with non-OpenRouter providers)
search_provider: str = "duckduckgo" # "duckduckgo" or "google"
google_api_key: Optional[str] = None
google_search_engine_id: Optional[str] = None
search_num_results: int = 5
# General settings
base_url: str = DEFAULT_BASE_URL
default_model: Optional[str] = None
default_system_prompt: Optional[str] = None
stream_enabled: bool = DEFAULT_STREAM_ENABLED
max_tokens: int = DEFAULT_MAX_TOKENS
cost_warning_threshold: float = DEFAULT_COST_WARNING_THRESHOLD
default_online_mode: bool = DEFAULT_ONLINE_MODE
log_max_size_mb: int = DEFAULT_LOG_MAX_SIZE_MB
log_backup_count: int = DEFAULT_LOG_BACKUP_COUNT
log_level: str = DEFAULT_LOG_LEVEL
@property
def effective_system_prompt(self) -> str:
"""
Get the effective system prompt to use.
Returns:
The custom prompt if set, hardcoded default if None, or blank if explicitly set to ""
"""
if self.default_system_prompt is None:
return DEFAULT_SYSTEM_PROMPT
return self.default_system_prompt
def __post_init__(self):
"""Validate settings after initialization."""
self._validate()
def _validate(self) -> None:
"""Validate all settings values."""
# Validate log level
if self.log_level.lower() not in VALID_LOG_LEVELS:
raise ValueError(
f"Invalid log level: {self.log_level}. "
f"Must be one of: {', '.join(VALID_LOG_LEVELS.keys())}"
)
# Validate numeric bounds
if self.max_tokens < 1:
raise ValueError("max_tokens must be at least 1")
if self.cost_warning_threshold < 0:
raise ValueError("cost_warning_threshold must be non-negative")
if self.log_max_size_mb < 1:
raise ValueError("log_max_size_mb must be at least 1")
if self.log_backup_count < 0:
raise ValueError("log_backup_count must be non-negative")
@classmethod
def load(cls) -> "Settings":
"""
Load settings from the database.
Returns:
Settings instance with values from database
"""
db = get_database()
# Helper to safely parse boolean
def parse_bool(value: Optional[str], default: bool) -> bool:
if value is None:
return default
return value.lower() in ("on", "true", "1", "yes")
# Helper to safely parse int
def parse_int(value: Optional[str], default: int) -> int:
if value is None:
return default
try:
return int(value)
except ValueError:
return default
# Helper to safely parse float
def parse_float(value: Optional[str], default: float) -> float:
if value is None:
return default
try:
return float(value)
except ValueError:
return default
# Get system prompt from DB: None means not set (use default), "" means explicitly blank
system_prompt_value = db.get_config("default_system_prompt")
# Migration: copy legacy api_key to openrouter_api_key if not already set
legacy_api_key = db.get_config("api_key")
openrouter_key = db.get_config("openrouter_api_key")
if legacy_api_key and not openrouter_key:
db.set_config("openrouter_api_key", legacy_api_key)
openrouter_key = legacy_api_key
# Note: We keep the legacy api_key in DB for backward compatibility
# Load provider-model mapping
provider_models_json = db.get_config("provider_models")
provider_models = {}
if provider_models_json:
try:
provider_models = json.loads(provider_models_json)
except json.JSONDecodeError:
provider_models = {}
return cls(
# Legacy field
api_key=legacy_api_key,
# Provider configuration
default_provider=db.get_config("default_provider") or DEFAULT_PROVIDER,
openrouter_api_key=openrouter_key,
anthropic_api_key=db.get_config("anthropic_api_key"),
openai_api_key=db.get_config("openai_api_key"),
ollama_base_url=db.get_config("ollama_base_url") or OLLAMA_DEFAULT_URL,
provider_models=provider_models,
# Web search configuration
search_provider=db.get_config("search_provider") or "duckduckgo",
google_api_key=db.get_config("google_api_key"),
google_search_engine_id=db.get_config("google_search_engine_id"),
search_num_results=parse_int(db.get_config("search_num_results"), 5),
# General settings
base_url=db.get_config("base_url") or DEFAULT_BASE_URL,
default_model=db.get_config("default_model"),
default_system_prompt=system_prompt_value,
stream_enabled=parse_bool(
db.get_config("stream_enabled"),
DEFAULT_STREAM_ENABLED
),
max_tokens=parse_int(
db.get_config("max_token"),
DEFAULT_MAX_TOKENS
),
cost_warning_threshold=parse_float(
db.get_config("cost_warning_threshold"),
DEFAULT_COST_WARNING_THRESHOLD
),
default_online_mode=parse_bool(
db.get_config("default_online_mode"),
DEFAULT_ONLINE_MODE
),
log_max_size_mb=parse_int(
db.get_config("log_max_size_mb"),
DEFAULT_LOG_MAX_SIZE_MB
),
log_backup_count=parse_int(
db.get_config("log_backup_count"),
DEFAULT_LOG_BACKUP_COUNT
),
log_level=db.get_config("log_level") or DEFAULT_LOG_LEVEL,
)
def save(self) -> None:
"""Persist all settings to the database."""
db = get_database()
# Only save API key if it exists
if self.api_key:
db.set_config("api_key", self.api_key)
db.set_config("base_url", self.base_url)
if self.default_model:
db.set_config("default_model", self.default_model)
# Save system prompt: None means not set (don't save), otherwise save the value (even if "")
if self.default_system_prompt is not None:
db.set_config("default_system_prompt", self.default_system_prompt)
db.set_config("stream_enabled", "on" if self.stream_enabled else "off")
db.set_config("max_token", str(self.max_tokens))
db.set_config("cost_warning_threshold", str(self.cost_warning_threshold))
db.set_config("default_online_mode", "on" if self.default_online_mode else "off")
db.set_config("log_max_size_mb", str(self.log_max_size_mb))
db.set_config("log_backup_count", str(self.log_backup_count))
db.set_config("log_level", self.log_level)
def set_api_key(self, api_key: str) -> None:
"""
Set and persist the API key.
Args:
api_key: The new API key
"""
self.api_key = api_key.strip()
get_database().set_config("api_key", self.api_key)
def set_base_url(self, url: str) -> None:
"""
Set and persist the base URL.
Args:
url: The new base URL
"""
self.base_url = url.strip()
get_database().set_config("base_url", self.base_url)
def set_default_model(self, model_id: str) -> None:
"""
Set and persist the default model.
Args:
model_id: The model ID to set as default
"""
self.default_model = model_id
get_database().set_config("default_model", model_id)
def set_default_system_prompt(self, prompt: str) -> None:
"""
Set and persist the default system prompt.
Args:
prompt: The system prompt to use for all new sessions.
Empty string "" means blank prompt (no system message).
"""
self.default_system_prompt = prompt
get_database().set_config("default_system_prompt", prompt)
def clear_default_system_prompt(self) -> None:
"""
Clear the custom system prompt and revert to hardcoded default.
This removes the custom prompt from the database, causing the
application to use the built-in DEFAULT_SYSTEM_PROMPT.
"""
self.default_system_prompt = None
# Remove from database to indicate "not set"
db = get_database()
with db._connection() as conn:
conn.execute("DELETE FROM config WHERE key = ?", ("default_system_prompt",))
conn.commit()
def set_stream_enabled(self, enabled: bool) -> None:
"""
Set and persist the streaming preference.
Args:
enabled: Whether to enable streaming
"""
self.stream_enabled = enabled
get_database().set_config("stream_enabled", "on" if enabled else "off")
def set_max_tokens(self, max_tokens: int) -> None:
"""
Set and persist the maximum tokens.
Args:
max_tokens: Maximum number of tokens
Raises:
ValueError: If max_tokens is less than 1
"""
if max_tokens < 1:
raise ValueError("max_tokens must be at least 1")
self.max_tokens = max_tokens
get_database().set_config("max_token", str(max_tokens))
def set_cost_warning_threshold(self, threshold: float) -> None:
"""
Set and persist the cost warning threshold.
Args:
threshold: Cost threshold in USD
Raises:
ValueError: If threshold is negative
"""
if threshold < 0:
raise ValueError("cost_warning_threshold must be non-negative")
self.cost_warning_threshold = threshold
get_database().set_config("cost_warning_threshold", str(threshold))
def set_default_online_mode(self, enabled: bool) -> None:
"""
Set and persist the default online mode.
Args:
enabled: Whether online mode should be enabled by default
"""
self.default_online_mode = enabled
get_database().set_config("default_online_mode", "on" if enabled else "off")
def set_log_level(self, level: str) -> None:
"""
Set and persist the log level.
Args:
level: The log level (debug/info/warning/error/critical)
Raises:
ValueError: If level is not valid
"""
level_lower = level.lower()
if level_lower not in VALID_LOG_LEVELS:
raise ValueError(
f"Invalid log level: {level}. "
f"Must be one of: {', '.join(VALID_LOG_LEVELS.keys())}"
)
self.log_level = level_lower
get_database().set_config("log_level", level_lower)
def set_log_max_size(self, size_mb: int) -> None:
"""
Set and persist the maximum log file size.
Args:
size_mb: Maximum size in megabytes
Raises:
ValueError: If size_mb is less than 1
"""
if size_mb < 1:
raise ValueError("log_max_size_mb must be at least 1")
# Cap at 100 MB for safety
self.log_max_size_mb = min(size_mb, 100)
get_database().set_config("log_max_size_mb", str(self.log_max_size_mb))
def set_provider_api_key(self, provider: str, api_key: str) -> None:
"""
Set and persist an API key for a specific provider.
Args:
provider: Provider name ("openrouter", "anthropic", "openai")
api_key: The API key to set
Raises:
ValueError: If provider is invalid
"""
provider = provider.lower()
api_key = api_key.strip()
if provider == "openrouter":
self.openrouter_api_key = api_key
get_database().set_config("openrouter_api_key", api_key)
elif provider == "anthropic":
self.anthropic_api_key = api_key
get_database().set_config("anthropic_api_key", api_key)
elif provider == "openai":
self.openai_api_key = api_key
get_database().set_config("openai_api_key", api_key)
else:
raise ValueError(f"Invalid provider: {provider}")
def get_provider_api_key(self, provider: str) -> Optional[str]:
"""
Get the API key for a specific provider.
Args:
provider: Provider name ("openrouter", "anthropic", "openai", "ollama")
Returns:
API key or None if not set
Raises:
ValueError: If provider is invalid
"""
provider = provider.lower()
if provider == "openrouter":
return self.openrouter_api_key
elif provider == "anthropic":
return self.anthropic_api_key
elif provider == "openai":
return self.openai_api_key
elif provider == "ollama":
return "" # Ollama doesn't require an API key
else:
raise ValueError(f"Invalid provider: {provider}")
def set_default_provider(self, provider: str) -> None:
"""
Set and persist the default provider.
Args:
provider: Provider name
Raises:
ValueError: If provider is invalid
"""
from oai.constants import VALID_PROVIDERS
provider = provider.lower()
if provider not in VALID_PROVIDERS:
raise ValueError(
f"Invalid provider: {provider}. "
f"Valid providers: {', '.join(VALID_PROVIDERS)}"
)
self.default_provider = provider
get_database().set_config("default_provider", provider)
def set_ollama_base_url(self, url: str) -> None:
"""
Set and persist the Ollama base URL.
Args:
url: Ollama server URL
"""
self.ollama_base_url = url.strip()
get_database().set_config("ollama_base_url", self.ollama_base_url)
def set_search_provider(self, provider: str) -> None:
"""
Set and persist the web search provider.
Args:
provider: Search provider ("anthropic_native", "duckduckgo", "google")
Raises:
ValueError: If provider is invalid
"""
valid_providers = ["anthropic_native", "duckduckgo", "google"]
provider = provider.lower()
if provider not in valid_providers:
raise ValueError(
f"Invalid search provider: {provider}. "
f"Valid providers: {', '.join(valid_providers)}"
)
self.search_provider = provider
get_database().set_config("search_provider", provider)
def set_google_api_key(self, api_key: str) -> None:
"""
Set and persist the Google API key for Google Custom Search.
Args:
api_key: The Google API key
"""
self.google_api_key = api_key.strip()
get_database().set_config("google_api_key", self.google_api_key)
def set_google_search_engine_id(self, engine_id: str) -> None:
"""
Set and persist the Google Custom Search Engine ID.
Args:
engine_id: The Google Search Engine ID
"""
self.google_search_engine_id = engine_id.strip()
get_database().set_config("google_search_engine_id", self.google_search_engine_id)
def get_provider_model(self, provider: str) -> Optional[str]:
"""
Get the last used model for a provider.
Args:
provider: Provider name
Returns:
Model ID or None if not set
"""
return self.provider_models.get(provider)
def set_provider_model(self, provider: str, model_id: str) -> None:
"""
Set and persist the last used model for a provider.
Args:
provider: Provider name
model_id: Model ID to remember
"""
self.provider_models[provider] = model_id
# Save to database as JSON
get_database().set_config("provider_models", json.dumps(self.provider_models))
# Global settings instance
_settings: Optional[Settings] = None
def get_settings() -> Settings:
"""
Get the global settings instance.
Returns:
The shared Settings instance, loading from database if needed
"""
global _settings
if _settings is None:
_settings = Settings.load()
return _settings
def reload_settings() -> Settings:
"""
Force reload settings from the database.
Returns:
Fresh Settings instance
"""
global _settings
_settings = Settings.load()
return _settings
+471
View File
@@ -0,0 +1,471 @@
"""
Application-wide constants for oAI.
This module contains all configuration constants, default values, and static
definitions used throughout the application. Centralizing these values makes
the codebase easier to maintain and configure.
"""
from pathlib import Path
from typing import Set, Dict, Any
import logging
# Import version from single source of truth
from oai import __version__
# =============================================================================
# APPLICATION METADATA
# =============================================================================
APP_NAME = "oAI"
APP_VERSION = __version__ # Single source of truth in oai/__init__.py
APP_URL = "https://iurl.no/oai"
APP_DESCRIPTION = "Open AI Chat Client with Multi-Provider Support"
# =============================================================================
# FILE PATHS
# =============================================================================
HOME_DIR = Path.home()
CONFIG_DIR = HOME_DIR / ".config" / "oai"
CACHE_DIR = HOME_DIR / ".cache" / "oai"
HISTORY_FILE = CONFIG_DIR / "history.txt"
DATABASE_FILE = CONFIG_DIR / "oai_config.db"
LOG_FILE = CONFIG_DIR / "oai.log"
# =============================================================================
# API CONFIGURATION
# =============================================================================
DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_STREAM_ENABLED = True
DEFAULT_MAX_TOKENS = 100_000
DEFAULT_ONLINE_MODE = False
# =============================================================================
# PROVIDER CONFIGURATION
# =============================================================================
# Provider names
PROVIDER_OPENROUTER = "openrouter"
PROVIDER_ANTHROPIC = "anthropic"
PROVIDER_OPENAI = "openai"
PROVIDER_OLLAMA = "ollama"
VALID_PROVIDERS = [PROVIDER_OPENROUTER, PROVIDER_ANTHROPIC, PROVIDER_OPENAI, PROVIDER_OLLAMA]
# Provider base URLs
ANTHROPIC_BASE_URL = "https://api.anthropic.com"
OPENAI_BASE_URL = "https://api.openai.com/v1"
OLLAMA_DEFAULT_URL = "http://localhost:11434"
# Default provider
DEFAULT_PROVIDER = PROVIDER_OPENROUTER
# =============================================================================
# DEFAULT SYSTEM PROMPT
# =============================================================================
DEFAULT_SYSTEM_PROMPT = (
"You are a knowledgeable and helpful AI assistant. Provide clear, accurate, "
"and well-structured responses. Be concise yet thorough. When uncertain about "
"something, acknowledge your limitations. For technical topics, include relevant "
"details and examples when helpful."
)
# =============================================================================
# PRICING DEFAULTS (per million tokens)
# =============================================================================
DEFAULT_INPUT_PRICE = 3.0
DEFAULT_OUTPUT_PRICE = 15.0
MODEL_PRICING: Dict[str, float] = {
"input": DEFAULT_INPUT_PRICE,
"output": DEFAULT_OUTPUT_PRICE,
}
# =============================================================================
# CREDIT ALERTS
# =============================================================================
LOW_CREDIT_RATIO = 0.1 # Alert when credits < 10% of total
LOW_CREDIT_AMOUNT = 1.0 # Alert when credits < $1.00
DEFAULT_COST_WARNING_THRESHOLD = 0.01 # Alert when single message cost exceeds this
COST_WARNING_THRESHOLD = DEFAULT_COST_WARNING_THRESHOLD # Alias for convenience
# =============================================================================
# LOGGING CONFIGURATION
# =============================================================================
DEFAULT_LOG_MAX_SIZE_MB = 10
DEFAULT_LOG_BACKUP_COUNT = 2
DEFAULT_LOG_LEVEL = "info"
VALID_LOG_LEVELS: Dict[str, int] = {
"debug": logging.DEBUG,
"info": logging.INFO,
"warning": logging.WARNING,
"error": logging.ERROR,
"critical": logging.CRITICAL,
}
# =============================================================================
# FILE HANDLING
# =============================================================================
# Maximum file size for reading (10 MB)
MAX_FILE_SIZE = 10 * 1024 * 1024
# Content truncation threshold (50 KB)
CONTENT_TRUNCATION_THRESHOLD = 50 * 1024
# Maximum items in directory listing
MAX_LIST_ITEMS = 1000
# Supported code file extensions for syntax highlighting
SUPPORTED_CODE_EXTENSIONS: Set[str] = {
".py", ".js", ".ts", ".cs", ".java", ".c", ".cpp", ".h", ".hpp",
".rb", ".ruby", ".php", ".swift", ".kt", ".kts", ".go",
".sh", ".bat", ".ps1", ".r", ".scala", ".pl", ".lua", ".dart",
".elm", ".xml", ".json", ".yaml", ".yml", ".md", ".txt",
}
# All allowed file extensions for attachment
ALLOWED_FILE_EXTENSIONS: Set[str] = {
# Code files
".py", ".js", ".ts", ".jsx", ".tsx", ".vue", ".java", ".c", ".cpp", ".cc", ".cxx",
".h", ".hpp", ".hxx", ".rb", ".go", ".rs", ".swift", ".kt", ".kts", ".php",
".sh", ".bash", ".zsh", ".fish", ".bat", ".cmd", ".ps1",
# Data files
".json", ".csv", ".yaml", ".yml", ".toml", ".xml", ".sql", ".db", ".sqlite", ".sqlite3",
# Documents
".txt", ".md", ".log", ".conf", ".cfg", ".ini", ".env", ".properties",
# Images
".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".svg", ".ico",
# Archives
".zip", ".tar", ".gz", ".bz2", ".7z", ".rar", ".xz",
# Config files
".lock", ".gitignore", ".dockerignore", ".editorconfig", ".eslintrc",
".prettierrc", ".babelrc", ".nvmrc", ".npmrc",
# Binary/Compiled
".pyc", ".pyo", ".pyd", ".so", ".dll", ".dylib", ".exe", ".app",
".dmg", ".pkg", ".deb", ".rpm", ".apk", ".ipa",
# ML/AI
".pkl", ".pickle", ".joblib", ".npy", ".npz", ".safetensors", ".onnx",
".pt", ".pth", ".ckpt", ".pb", ".tflite", ".mlmodel", ".coreml", ".rknn",
# Data formats
".wasm", ".proto", ".graphql", ".graphqls", ".grpc", ".avro", ".parquet",
".orc", ".feather", ".arrow", ".hdf5", ".h5", ".mat", ".rdata", ".rds",
# Other
".pdf", ".class", ".jar", ".war",
}
# =============================================================================
# SECURITY CONFIGURATION
# =============================================================================
# System directories that should never be accessed
SYSTEM_DIRS_BLACKLIST: Set[str] = {
# macOS
"/System", "/Library", "/private", "/usr", "/bin", "/sbin",
# Linux
"/boot", "/dev", "/proc", "/sys", "/root",
# Windows
"C:\\Windows", "C:\\Program Files", "C:\\Program Files (x86)",
}
# Directories to skip during file operations
SKIP_DIRECTORIES: Set[str] = {
# Python virtual environments
".venv", "venv", "env", "virtualenv",
"site-packages", "dist-packages",
# Python caches
"__pycache__", ".pytest_cache", ".mypy_cache",
# JavaScript/Node
"node_modules",
# Version control
".git", ".svn",
# IDEs
".idea", ".vscode",
# Build directories
"build", "dist", "eggs", ".eggs",
}
# =============================================================================
# DATABASE QUERIES - SQL SAFETY
# =============================================================================
# Maximum query execution timeout (seconds)
MAX_QUERY_TIMEOUT = 5
# Maximum rows returned from queries
MAX_QUERY_RESULTS = 1000
# Default rows per query
DEFAULT_QUERY_LIMIT = 100
# Keywords that are blocked in database queries
DANGEROUS_SQL_KEYWORDS: Set[str] = {
"INSERT", "UPDATE", "DELETE", "DROP", "CREATE",
"ALTER", "TRUNCATE", "REPLACE", "ATTACH", "DETACH",
"PRAGMA", "VACUUM", "REINDEX",
}
# =============================================================================
# MCP CONFIGURATION
# =============================================================================
# Maximum tool call iterations per request
MAX_TOOL_LOOPS = 5
# =============================================================================
# VALID COMMANDS
# =============================================================================
VALID_COMMANDS: Set[str] = {
"/retry", "/online", "/memory", "/paste", "/export", "/save", "/load",
"/delete", "/list", "/prev", "/next", "/stats", "/middleout", "/reset",
"/info", "/model", "/maxtoken", "/system", "/config", "/credits", "/clear",
"/cl", "/help", "/mcp",
}
# =============================================================================
# COMMAND HELP DATABASE
# =============================================================================
COMMAND_HELP: Dict[str, Dict[str, Any]] = {
"/clear": {
"aliases": ["/cl"],
"description": "Clear the terminal screen for a clean interface.",
"usage": "/clear\n/cl",
"examples": [
("Clear screen", "/clear"),
("Using short alias", "/cl"),
],
"notes": "You can also use the keyboard shortcut Ctrl+L.",
},
"/help": {
"description": "Display help information for commands.",
"usage": "/help [command|topic]",
"examples": [
("Show all commands", "/help"),
("Get help for a specific command", "/help /model"),
("Get detailed MCP help", "/help mcp"),
],
"notes": "Use /help without arguments to see the full command list.",
},
"mcp": {
"description": "Complete guide to MCP (Model Context Protocol).",
"usage": "See detailed examples below",
"examples": [],
"notes": """
MCP (Model Context Protocol) gives your AI assistant direct access to:
• Local files and folders (read, search, list)
• SQLite databases (inspect, search, query)
FILE MODE (default):
/mcp on Start MCP server
/mcp add ~/Documents Grant access to folder
/mcp list View all allowed folders
DATABASE MODE:
/mcp add db ~/app/data.db Add specific database
/mcp db list View all databases
/mcp db 1 Work with database #1
/mcp files Switch back to file mode
WRITE MODE (optional):
/mcp write on Enable file modifications
/mcp write off Disable write mode (back to read-only)
For command-specific help: /help /mcp
""",
},
"/mcp": {
"description": "Manage MCP for local file access and SQLite database querying.",
"usage": "/mcp <command> [args]",
"examples": [
("Enable MCP server", "/mcp on"),
("Disable MCP server", "/mcp off"),
("Show MCP status", "/mcp status"),
("", ""),
("━━━ FILE MODE ━━━", ""),
("Add folder for file access", "/mcp add ~/Documents"),
("Remove folder", "/mcp remove ~/Desktop"),
("List allowed folders", "/mcp list"),
("Enable write mode", "/mcp write on"),
("", ""),
("━━━ DATABASE MODE ━━━", ""),
("Add SQLite database", "/mcp add db ~/app/data.db"),
("List all databases", "/mcp db list"),
("Switch to database #1", "/mcp db 1"),
("Switch back to file mode", "/mcp files"),
],
"notes": "MCP allows AI to read local files and query SQLite databases.",
},
"/memory": {
"description": "Toggle conversation memory.",
"usage": "/memory [on|off]",
"examples": [
("Check current memory status", "/memory"),
("Enable conversation memory", "/memory on"),
("Disable memory (save costs)", "/memory off"),
],
"notes": "Memory is ON by default. Disabling saves tokens.",
},
"/online": {
"description": "Enable or disable online mode (web search).",
"usage": "/online [on|off]",
"examples": [
("Check online mode status", "/online"),
("Enable web search", "/online on"),
("Disable web search", "/online off"),
],
"notes": "Not all models support online mode.",
},
"/paste": {
"description": "Paste plain text from clipboard and send to the AI.",
"usage": "/paste [prompt]",
"examples": [
("Paste clipboard content", "/paste"),
("Paste with a question", "/paste Explain this code"),
],
"notes": "Only plain text is supported.",
},
"/retry": {
"description": "Resend the last prompt from conversation history.",
"usage": "/retry",
"examples": [("Retry last message", "/retry")],
"notes": "Requires at least one message in history.",
},
"/next": {
"description": "View the next response in conversation history.",
"usage": "/next",
"examples": [("Navigate to next response", "/next")],
"notes": "Use /prev to go backward.",
},
"/prev": {
"description": "View the previous response in conversation history.",
"usage": "/prev",
"examples": [("Navigate to previous response", "/prev")],
"notes": "Use /next to go forward.",
},
"/reset": {
"description": "Clear conversation history and reset system prompt.",
"usage": "/reset",
"examples": [("Reset conversation", "/reset")],
"notes": "Requires confirmation.",
},
"/info": {
"description": "Display detailed information about a model.",
"usage": "/info [model_id]",
"examples": [
("Show current model info", "/info"),
("Show specific model info", "/info gpt-4o"),
],
"notes": "Shows pricing, capabilities, and context length.",
},
"/model": {
"description": "Select or change the AI model.",
"usage": "/model [search_term]",
"examples": [
("List all models", "/model"),
("Search for GPT models", "/model gpt"),
("Search for Claude models", "/model claude"),
],
"notes": "Models are numbered for easy selection.",
},
"/config": {
"description": "View or modify application configuration.",
"usage": "/config [setting] [value]",
"examples": [
("View all settings", "/config"),
("Set API key", "/config api"),
("Set default model", "/config model"),
("Set system prompt", "/config system You are a helpful assistant"),
("Enable streaming", "/config stream on"),
],
"notes": "Available: api, url, model, system, stream, costwarning, maxtoken, online, loglevel.",
},
"/maxtoken": {
"description": "Set a temporary session token limit.",
"usage": "/maxtoken [value]",
"examples": [
("View current session limit", "/maxtoken"),
("Set session limit to 2000", "/maxtoken 2000"),
],
"notes": "Cannot exceed stored max token limit.",
},
"/system": {
"description": "Set or clear the session-level system prompt.",
"usage": "/system [prompt|clear|default <prompt>]",
"examples": [
("View current system prompt", "/system"),
("Set as Python expert", "/system You are a Python expert"),
("Multiline with newlines", r"/system You are an expert.\nBe clear and concise."),
("Save as default", "/system default You are a helpful assistant"),
("Revert to default", "/system clear"),
("Blank prompt", '/system ""'),
],
"notes": r"Use \n for newlines. /system clear reverts to hardcoded default.",
},
"/save": {
"description": "Save the current conversation history.",
"usage": "/save <name>",
"examples": [("Save conversation", "/save my_chat")],
"notes": "Saved conversations can be loaded later with /load.",
},
"/load": {
"description": "Load a saved conversation.",
"usage": "/load <name|number>",
"examples": [
("Load by name", "/load my_chat"),
("Load by number from /list", "/load 3"),
],
"notes": "Use /list to see numbered conversations.",
},
"/delete": {
"description": "Delete a saved conversation.",
"usage": "/delete <name|number>",
"examples": [("Delete by name", "/delete my_chat")],
"notes": "Requires confirmation. Cannot be undone.",
},
"/list": {
"description": "List all saved conversations.",
"usage": "/list",
"examples": [("Show saved conversations", "/list")],
"notes": "Conversations are numbered for use with /load and /delete.",
},
"/export": {
"description": "Export the current conversation to a file.",
"usage": "/export <format> <filename>",
"examples": [
("Export as Markdown", "/export md notes.md"),
("Export as JSON", "/export json conversation.json"),
("Export as HTML", "/export html report.html"),
],
"notes": "Available formats: md, json, html.",
},
"/stats": {
"description": "Display session statistics.",
"usage": "/stats",
"examples": [("View session statistics", "/stats")],
"notes": "Shows tokens, costs, and credits.",
},
"/credits": {
"description": "Display your OpenRouter account credits.",
"usage": "/credits",
"examples": [("Check credits", "/credits")],
"notes": "Shows total, used, and remaining credits.",
},
"/middleout": {
"description": "Toggle middle-out transform for long prompts.",
"usage": "/middleout [on|off]",
"examples": [
("Check status", "/middleout"),
("Enable compression", "/middleout on"),
],
"notes": "Compresses prompts exceeding context size.",
},
}
+14
View File
@@ -0,0 +1,14 @@
"""
Core functionality for oAI.
This module provides the main session management and AI client
classes that power the chat application.
"""
from oai.core.session import ChatSession
from oai.core.client import AIClient
__all__ = [
"ChatSession",
"AIClient",
]
+502
View File
@@ -0,0 +1,502 @@
"""
AI Client for oAI.
This module provides a high-level client for interacting with AI models
through the provider abstraction layer.
"""
import asyncio
import json
from typing import Any, Callable, Dict, Iterator, List, Optional, Union
from oai.constants import APP_NAME, APP_URL, MODEL_PRICING, OLLAMA_DEFAULT_URL
from oai.providers.base import (
AIProvider,
ChatMessage,
ChatResponse,
ModelInfo,
StreamChunk,
ToolCall,
UsageStats,
)
from oai.utils.logging import get_logger
class AIClient:
"""
High-level AI client for chat interactions.
Provides a simplified interface for sending chat requests,
handling streaming, and managing tool calls.
Attributes:
provider: The underlying AI provider
provider_name: Name of the current provider
default_model: Default model ID to use
http_headers: Custom HTTP headers for requests
"""
def __init__(
self,
provider_name: str = "openrouter",
provider_api_keys: Optional[Dict[str, str]] = None,
ollama_base_url: str = OLLAMA_DEFAULT_URL,
app_name: str = APP_NAME,
app_url: str = APP_URL,
):
"""
Initialize the AI client with specified provider.
Args:
provider_name: Provider to use ("openrouter", "anthropic", "openai", "ollama")
provider_api_keys: Dict mapping provider names to API keys
ollama_base_url: Base URL for Ollama server
app_name: Application name for headers
app_url: Application URL for headers
Raises:
ValueError: If provider is invalid or not configured
"""
from oai.providers.registry import get_provider_class
self.provider_name = provider_name
self.provider_api_keys = provider_api_keys or {}
self.ollama_base_url = ollama_base_url
self.app_name = app_name
self.app_url = app_url
# Get provider class
provider_class = get_provider_class(provider_name)
if not provider_class:
raise ValueError(f"Unknown provider: {provider_name}")
# Get API key for this provider
api_key = self.provider_api_keys.get(provider_name, "")
# Initialize provider with appropriate parameters
if provider_name == "ollama":
self.provider: AIProvider = provider_class(
api_key=api_key,
base_url=ollama_base_url,
)
else:
self.provider: AIProvider = provider_class(
api_key=api_key,
app_name=app_name,
app_url=app_url,
)
self.default_model: Optional[str] = None
self.logger = get_logger()
self.logger.info(f"Initialized {provider_name} provider")
def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
"""
Get available models.
Args:
filter_text_only: Whether to exclude video-only models
Returns:
List of ModelInfo objects
"""
return self.provider.list_models(filter_text_only=filter_text_only)
def get_model(self, model_id: str) -> Optional[ModelInfo]:
"""
Get information about a specific model.
Args:
model_id: Model identifier
Returns:
ModelInfo or None if not found
"""
return self.provider.get_model(model_id)
def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
"""
Get raw model data for provider-specific fields.
Args:
model_id: Model identifier
Returns:
Raw model dictionary or None
"""
if hasattr(self.provider, "get_raw_model"):
return self.provider.get_raw_model(model_id)
return None
def chat(
self,
messages: List[Dict[str, Any]],
model: Optional[str] = None,
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
system_prompt: Optional[str] = None,
online: bool = False,
transforms: Optional[List[str]] = None,
enable_web_search: bool = False,
web_search_config: Optional[Dict[str, Any]] = None,
) -> Union[ChatResponse, Iterator[StreamChunk]]:
"""
Send a chat request.
Args:
messages: List of message dictionaries
model: Model ID (uses default if not specified)
stream: Whether to stream the response
max_tokens: Maximum tokens in response
temperature: Sampling temperature
tools: Tool definitions for function calling
tool_choice: Tool selection mode
system_prompt: System prompt to prepend
online: Whether to enable online mode
transforms: List of transforms (e.g., ["middle-out"])
enable_web_search: Enable native web search (Anthropic only)
web_search_config: Web search configuration (Anthropic only)
Returns:
ChatResponse for non-streaming, Iterator[StreamChunk] for streaming
Raises:
ValueError: If no model specified and no default set
"""
model_id = model or self.default_model
if not model_id:
raise ValueError("No model specified and no default set")
# Apply online mode suffix
if online and hasattr(self.provider, "get_effective_model_id"):
model_id = self.provider.get_effective_model_id(model_id, True)
# Convert dict messages to ChatMessage objects
chat_messages = []
# Add system prompt if provided
if system_prompt:
chat_messages.append(ChatMessage(role="system", content=system_prompt))
# Convert message dicts
for msg in messages:
# Convert tool_calls dicts to ToolCall objects if present
tool_calls_data = msg.get("tool_calls")
tool_calls_obj = None
if tool_calls_data:
from oai.providers.base import ToolCall, ToolFunction
tool_calls_obj = []
for tc in tool_calls_data:
# Handle both ToolCall objects and dicts
if isinstance(tc, ToolCall):
tool_calls_obj.append(tc)
elif isinstance(tc, dict):
func_data = tc.get("function", {})
tool_calls_obj.append(
ToolCall(
id=tc.get("id", ""),
type=tc.get("type", "function"),
function=ToolFunction(
name=func_data.get("name", ""),
arguments=func_data.get("arguments", "{}"),
),
)
)
chat_messages.append(
ChatMessage(
role=msg.get("role", "user"),
content=msg.get("content"),
tool_calls=tool_calls_obj,
tool_call_id=msg.get("tool_call_id"),
)
)
self.logger.debug(
f"Sending chat request: model={model_id}, "
f"messages={len(chat_messages)}, stream={stream}"
)
return self.provider.chat(
model=model_id,
messages=chat_messages,
stream=stream,
max_tokens=max_tokens,
temperature=temperature,
tools=tools,
tool_choice=tool_choice,
transforms=transforms,
enable_web_search=enable_web_search,
web_search_config=web_search_config or {},
)
def chat_with_tools(
self,
messages: List[Dict[str, Any]],
tools: List[Dict[str, Any]],
tool_executor: Callable[[str, Dict[str, Any]], Dict[str, Any]],
model: Optional[str] = None,
max_loops: int = 5,
max_tokens: Optional[int] = None,
system_prompt: Optional[str] = None,
on_tool_call: Optional[Callable[[ToolCall], None]] = None,
on_tool_result: Optional[Callable[[str, Dict[str, Any]], None]] = None,
) -> ChatResponse:
"""
Send a chat request with automatic tool call handling.
Executes tool calls returned by the model and continues
the conversation until no more tool calls are requested.
Args:
messages: Initial messages
tools: Tool definitions
tool_executor: Function to execute tool calls
model: Model ID
max_loops: Maximum tool call iterations
max_tokens: Maximum response tokens
system_prompt: System prompt
on_tool_call: Callback when tool is called
on_tool_result: Callback when tool returns result
Returns:
Final ChatResponse after all tool calls complete
"""
model_id = model or self.default_model
if not model_id:
raise ValueError("No model specified and no default set")
# Build initial messages
chat_messages = []
if system_prompt:
chat_messages.append({"role": "system", "content": system_prompt})
chat_messages.extend(messages)
loop_count = 0
current_response: Optional[ChatResponse] = None
while loop_count < max_loops:
# Send request
response = self.chat(
messages=chat_messages,
model=model_id,
stream=False,
max_tokens=max_tokens,
tools=tools,
tool_choice="auto",
)
if not isinstance(response, ChatResponse):
raise ValueError("Expected non-streaming response")
current_response = response
# Check for tool calls
tool_calls = response.tool_calls
if not tool_calls:
break
self.logger.info(f"Model requested {len(tool_calls)} tool call(s)")
# Process each tool call
tool_results = []
for tc in tool_calls:
if on_tool_call:
on_tool_call(tc)
try:
args = json.loads(tc.function.arguments)
except json.JSONDecodeError as e:
self.logger.error(f"Failed to parse tool arguments: {e}")
result = {"error": f"Invalid arguments: {e}"}
else:
result = tool_executor(tc.function.name, args)
if on_tool_result:
on_tool_result(tc.function.name, result)
tool_results.append({
"tool_call_id": tc.id,
"role": "tool",
"name": tc.function.name,
"content": json.dumps(result),
})
# Add assistant message with tool calls
assistant_msg = {
"role": "assistant",
"content": response.content,
"tool_calls": [
{
"id": tc.id,
"type": tc.type,
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in tool_calls
],
}
chat_messages.append(assistant_msg)
chat_messages.extend(tool_results)
loop_count += 1
if loop_count >= max_loops:
self.logger.warning(f"Reached max tool call loops ({max_loops})")
return current_response
def stream_chat(
self,
messages: List[Dict[str, Any]],
model: Optional[str] = None,
max_tokens: Optional[int] = None,
system_prompt: Optional[str] = None,
online: bool = False,
on_chunk: Optional[Callable[[StreamChunk], None]] = None,
) -> tuple[str, Optional[UsageStats]]:
"""
Stream a chat response and collect the full text.
Args:
messages: Chat messages
model: Model ID
max_tokens: Maximum tokens
system_prompt: System prompt
online: Online mode
on_chunk: Optional callback for each chunk
Returns:
Tuple of (full_response_text, usage_stats)
"""
response = self.chat(
messages=messages,
model=model,
stream=True,
max_tokens=max_tokens,
system_prompt=system_prompt,
online=online,
)
if isinstance(response, ChatResponse):
# Not actually streaming
return response.content or "", response.usage
full_text = ""
usage: Optional[UsageStats] = None
for chunk in response:
if chunk.error:
self.logger.error(f"Stream error: {chunk.error}")
break
if chunk.delta_content:
full_text += chunk.delta_content
if on_chunk:
on_chunk(chunk)
if chunk.usage:
usage = chunk.usage
return full_text, usage
def get_credits(self) -> Optional[Dict[str, Any]]:
"""
Get account credit information.
Returns:
Credit info dict or None if unavailable
"""
return self.provider.get_credits()
def estimate_cost(
self,
model_id: str,
input_tokens: int,
output_tokens: int,
) -> float:
"""
Estimate cost for a completion.
Args:
model_id: Model ID
input_tokens: Number of input tokens
output_tokens: Number of output tokens
Returns:
Estimated cost in USD
"""
if hasattr(self.provider, "estimate_cost"):
return self.provider.estimate_cost(model_id, input_tokens, output_tokens)
# Fallback to default pricing
input_cost = MODEL_PRICING["input"] * input_tokens / 1_000_000
output_cost = MODEL_PRICING["output"] * output_tokens / 1_000_000
return input_cost + output_cost
def set_default_model(self, model_id: str) -> None:
"""
Set the default model.
Args:
model_id: Model ID to use as default
"""
self.default_model = model_id
self.logger.info(f"Default model set to: {model_id}")
def clear_cache(self) -> None:
"""Clear the provider's model cache."""
if hasattr(self.provider, "clear_cache"):
self.provider.clear_cache()
def switch_provider(self, provider_name: str, ollama_base_url: Optional[str] = None) -> None:
"""
Switch to a different provider.
Args:
provider_name: Provider to switch to
ollama_base_url: Optional Ollama base URL (if switching to Ollama)
Raises:
ValueError: If provider is invalid or not configured
"""
from oai.providers.registry import get_provider_class
# Get provider class
provider_class = get_provider_class(provider_name)
if not provider_class:
raise ValueError(f"Unknown provider: {provider_name}")
# Get API key
api_key = self.provider_api_keys.get(provider_name, "")
# Check API key requirement
if provider_name != "ollama" and not api_key:
raise ValueError(f"No API key configured for {provider_name}")
# Initialize new provider
if provider_name == "ollama":
base_url = ollama_base_url or self.ollama_base_url
self.provider = provider_class(
api_key=api_key,
base_url=base_url,
)
self.ollama_base_url = base_url
else:
self.provider = provider_class(
api_key=api_key,
app_name=self.app_name,
app_url=self.app_url,
)
self.provider_name = provider_name
self.logger.info(f"Switched to {provider_name} provider")
# Clear model cache when switching providers
self.default_model = None
+984
View File
@@ -0,0 +1,984 @@
"""
Chat session management for oAI.
This module provides the ChatSession class that manages an interactive
chat session including history, state, and message handling.
"""
import asyncio
import json
import time
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Tuple
from oai.commands.registry import CommandContext, CommandResult, registry
from oai.config.database import Database
from oai.config.settings import Settings
from oai.constants import (
COST_WARNING_THRESHOLD,
LOW_CREDIT_AMOUNT,
LOW_CREDIT_RATIO,
)
from oai.core.client import AIClient
from oai.mcp.manager import MCPManager
from oai.providers.base import ChatResponse, StreamChunk, UsageStats
from oai.utils.logging import get_logger
from oai.utils.web_search import perform_web_search, format_search_results
logger = get_logger()
@dataclass
class SessionStats:
"""
Statistics for the current session.
Tracks tokens, costs, and message counts.
"""
total_input_tokens: int = 0
total_output_tokens: int = 0
total_cost: float = 0.0
message_count: int = 0
@property
def total_tokens(self) -> int:
"""Get total token count."""
return self.total_input_tokens + self.total_output_tokens
def add_usage(self, usage: Optional[UsageStats], cost: float = 0.0) -> None:
"""
Add usage stats from a response.
Args:
usage: Usage statistics
cost: Cost if not in usage
"""
if usage:
self.total_input_tokens += usage.prompt_tokens
self.total_output_tokens += usage.completion_tokens
if usage.total_cost_usd:
self.total_cost += usage.total_cost_usd
else:
self.total_cost += cost
else:
self.total_cost += cost
self.message_count += 1
@dataclass
class HistoryEntry:
"""
A single entry in the conversation history.
Stores the user prompt, assistant response, and metrics.
"""
prompt: str
response: str
prompt_tokens: int = 0
completion_tokens: int = 0
msg_cost: float = 0.0
timestamp: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary format."""
return {
"prompt": self.prompt,
"response": self.response,
"prompt_tokens": self.prompt_tokens,
"completion_tokens": self.completion_tokens,
"msg_cost": self.msg_cost,
}
class ChatSession:
"""
Manages an interactive chat session.
Handles conversation history, state management, command processing,
and communication with the AI client.
Attributes:
client: AI client for API requests
settings: Application settings
mcp_manager: MCP manager for file/database access
history: Conversation history
stats: Session statistics
"""
def __init__(
self,
client: AIClient,
settings: Settings,
mcp_manager: Optional[MCPManager] = None,
):
"""
Initialize a chat session.
Args:
client: AI client instance
settings: Application settings
mcp_manager: Optional MCP manager
"""
self.client = client
self.settings = settings
self.mcp_manager = mcp_manager
self.db = Database()
self.history: List[HistoryEntry] = []
self.stats = SessionStats()
# Session state
self.system_prompt: str = settings.effective_system_prompt
self.memory_enabled: bool = True
self.memory_start_index: int = 0
self.online_enabled: bool = settings.default_online_mode
self.middle_out_enabled: bool = False
self.session_max_token: int = 0
self.current_index: int = 0
# Selected model
self.selected_model: Optional[Dict[str, Any]] = None
self.logger = get_logger()
def get_context(self) -> CommandContext:
"""
Get the current command context.
Returns:
CommandContext with current session state
"""
return CommandContext(
settings=self.settings,
provider=self.client.provider,
mcp_manager=self.mcp_manager,
selected_model_raw=self.selected_model,
session_history=[e.to_dict() for e in self.history],
session_system_prompt=self.system_prompt,
memory_enabled=self.memory_enabled,
memory_start_index=self.memory_start_index,
online_enabled=self.online_enabled,
middle_out_enabled=self.middle_out_enabled,
session_max_token=self.session_max_token,
total_input_tokens=self.stats.total_input_tokens,
total_output_tokens=self.stats.total_output_tokens,
total_cost=self.stats.total_cost,
message_count=self.stats.message_count,
current_index=self.current_index,
session=self,
)
def set_model(self, model: Dict[str, Any]) -> None:
"""
Set the selected model.
Args:
model: Raw model dictionary
"""
self.selected_model = model
self.client.set_default_model(model["id"])
self.logger.info(f"Model selected: {model['id']}")
def build_api_messages(self, user_input: str) -> List[Dict[str, Any]]:
"""
Build the messages array for an API request.
Includes system prompt, history (if memory enabled), and current input.
Args:
user_input: Current user input
Returns:
List of message dictionaries
"""
messages = []
# Add system prompt
if self.system_prompt:
messages.append({"role": "system", "content": self.system_prompt})
# Add database context if in database mode
if self.mcp_manager and self.mcp_manager.enabled:
if self.mcp_manager.mode == "database" and self.mcp_manager.selected_db_index is not None:
db = self.mcp_manager.databases[self.mcp_manager.selected_db_index]
db_context = (
f"You are connected to SQLite database: {db['name']}\n"
f"Available tables: {', '.join(db['tables'])}\n\n"
"Use inspect_database, search_database, or query_database tools. "
"All queries are read-only."
)
messages.append({"role": "system", "content": db_context})
# Add history if memory enabled
if self.memory_enabled:
for i in range(self.memory_start_index, len(self.history)):
entry = self.history[i]
messages.append({"role": "user", "content": entry.prompt})
messages.append({"role": "assistant", "content": entry.response})
# Add current message
messages.append({"role": "user", "content": user_input})
return messages
def get_mcp_tools(self) -> Optional[List[Dict[str, Any]]]:
"""
Get MCP tool definitions if available.
Returns:
List of tool schemas or None
"""
if not self.mcp_manager or not self.mcp_manager.enabled:
return None
if not self.selected_model:
return None
# Check if model supports tools
supported_params = self.selected_model.get("supported_parameters", [])
if "tools" not in supported_params and "functions" not in supported_params:
return None
return self.mcp_manager.get_tools_schema()
async def execute_tool(
self,
tool_name: str,
tool_args: Dict[str, Any],
) -> Dict[str, Any]:
"""
Execute an MCP tool.
Args:
tool_name: Name of the tool
tool_args: Tool arguments
Returns:
Tool execution result
"""
if not self.mcp_manager:
return {"error": "MCP not available"}
return await self.mcp_manager.call_tool(tool_name, **tool_args)
def send_message(
self,
user_input: str,
stream: bool = True,
on_stream_chunk: Optional[Callable[[str], None]] = None,
) -> Tuple[str, Optional[UsageStats], float]:
"""
Send a message and get a response.
Args:
user_input: User's input text
stream: Whether to stream the response
on_stream_chunk: Callback for stream chunks
Returns:
Tuple of (response_text, usage_stats, response_time)
"""
if not self.selected_model:
raise ValueError("No model selected")
start_time = time.time()
messages = self.build_api_messages(user_input)
# Get MCP tools
tools = self.get_mcp_tools()
if tools:
# Disable streaming when tools are present
stream = False
# Build request parameters
model_id = self.selected_model["id"]
# Handle online mode
enable_web_search = False
web_search_config = {}
if self.online_enabled:
# OpenRouter handles online mode natively with :online suffix
if self.client.provider_name == "openrouter":
if hasattr(self.client.provider, "get_effective_model_id"):
model_id = self.client.provider.get_effective_model_id(model_id, True)
# Anthropic has native web search when search provider is set to anthropic_native
elif self.client.provider_name == "anthropic" and self.settings.search_provider == "anthropic_native":
enable_web_search = True
web_search_config = {
"max_uses": self.settings.search_num_results or 5
}
logger.info("Using Anthropic native web search")
else:
# For other providers, perform web search and inject results
logger.info(f"Performing web search for: {user_input}")
search_results = perform_web_search(
user_input,
num_results=self.settings.search_num_results,
provider=self.settings.search_provider,
google_api_key=self.settings.google_api_key,
google_search_engine_id=self.settings.google_search_engine_id
)
if search_results:
# Inject search results into messages
formatted_results = format_search_results(search_results)
search_context = f"\n\n{formatted_results}\n\nPlease use the above web search results to help answer the user's question."
# Add search results to the last user message
if messages and messages[-1]["role"] == "user":
messages[-1]["content"] += search_context
logger.info(f"Injected {len(search_results)} search results into context")
if self.online_enabled:
if hasattr(self.client.provider, "get_effective_model_id"):
model_id = self.client.provider.get_effective_model_id(model_id, True)
transforms = ["middle-out"] if self.middle_out_enabled else None
max_tokens = None
if self.session_max_token > 0:
max_tokens = self.session_max_token
if tools:
# Use tool handling flow
response = self._send_with_tools(
messages=messages,
model_id=model_id,
tools=tools,
max_tokens=max_tokens,
transforms=transforms,
)
response_time = time.time() - start_time
return response.content or "", response.usage, response_time
elif stream:
# Use streaming flow
full_text, usage = self._stream_response(
messages=messages,
model_id=model_id,
max_tokens=max_tokens,
transforms=transforms,
on_chunk=on_stream_chunk,
enable_web_search=enable_web_search,
web_search_config=web_search_config,
)
response_time = time.time() - start_time
return full_text, usage, response_time
else:
# Non-streaming request
response = self.client.chat(
messages=messages,
model=model_id,
stream=False,
max_tokens=max_tokens,
transforms=transforms,
)
response_time = time.time() - start_time
if isinstance(response, ChatResponse):
return response.content or "", response.usage, response_time
else:
return "", None, response_time
def _send_with_tools(
self,
messages: List[Dict[str, Any]],
model_id: str,
tools: List[Dict[str, Any]],
max_tokens: Optional[int] = None,
transforms: Optional[List[str]] = None,
) -> ChatResponse:
"""
Send a request with tool call handling.
Args:
messages: API messages
model_id: Model ID
tools: Tool definitions
max_tokens: Max tokens
transforms: Transforms list
Returns:
Final ChatResponse
"""
max_loops = 5
loop_count = 0
api_messages = list(messages)
while loop_count < max_loops:
response = self.client.chat(
messages=api_messages,
model=model_id,
stream=False,
max_tokens=max_tokens,
tools=tools,
tool_choice="auto",
transforms=transforms,
)
if not isinstance(response, ChatResponse):
raise ValueError("Expected ChatResponse")
tool_calls = response.tool_calls
if not tool_calls:
return response
# Tool calls requested by AI
tool_results = []
for tc in tool_calls:
try:
args = json.loads(tc.function.arguments)
except json.JSONDecodeError as e:
self.logger.error(f"Failed to parse tool arguments: {e}")
tool_results.append({
"tool_call_id": tc.id,
"role": "tool",
"name": tc.function.name,
"content": json.dumps({"error": f"Invalid arguments: {e}"}),
})
continue
# Display tool call
args_display = ", ".join(
f'{k}="{v}"' if isinstance(v, str) else f"{k}={v}"
for k, v in args.items()
)
# Executing tool: {tc.function.name}
# Execute tool
result = asyncio.run(self.execute_tool(tc.function.name, args))
if "error" in result:
# Tool execution error logged
pass
else:
# Tool execution successful
pass
tool_results.append({
"tool_call_id": tc.id,
"role": "tool",
"name": tc.function.name,
"content": json.dumps(result),
})
# Add assistant message with tool calls
api_messages.append({
"role": "assistant",
"content": response.content,
"tool_calls": [
{
"id": tc.id,
"type": tc.type,
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in tool_calls
],
})
api_messages.extend(tool_results)
# Processing tool results
loop_count += 1
self.logger.warning(f"Reached max tool loops ({max_loops})")
return response
def _stream_response(
self,
messages: List[Dict[str, Any]],
model_id: str,
max_tokens: Optional[int] = None,
transforms: Optional[List[str]] = None,
on_chunk: Optional[Callable[[str], None]] = None,
enable_web_search: bool = False,
web_search_config: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Optional[UsageStats]]:
"""
Stream a response with live display.
Args:
messages: API messages
model_id: Model ID
max_tokens: Max tokens
transforms: Transforms
on_chunk: Callback for chunks
enable_web_search: Whether to enable Anthropic native web search
web_search_config: Web search configuration
Returns:
Tuple of (full_text, usage)
"""
response = self.client.chat(
messages=messages,
model=model_id,
stream=True,
max_tokens=max_tokens,
transforms=transforms,
enable_web_search=enable_web_search,
web_search_config=web_search_config or {},
)
if isinstance(response, ChatResponse):
return response.content or "", response.usage
full_text = ""
usage: Optional[UsageStats] = None
try:
for chunk in response:
if chunk.error:
self.logger.error(f"Stream error: {chunk.error}")
break
if chunk.delta_content:
full_text += chunk.delta_content
if on_chunk:
on_chunk(chunk.delta_content)
if chunk.usage:
usage = chunk.usage
except KeyboardInterrupt:
self.logger.info("Streaming interrupted")
return "", None
return full_text, usage
# ========== ASYNC METHODS FOR TUI ==========
async def send_message_async(
self,
user_input: str,
stream: bool = True,
) -> AsyncIterator[StreamChunk]:
"""
Async version of send_message for Textual TUI.
Args:
user_input: User's input text
stream: Whether to stream the response
Yields:
StreamChunk objects for progressive display
"""
if not self.selected_model:
raise ValueError("No model selected")
messages = self.build_api_messages(user_input)
tools = self.get_mcp_tools()
if tools:
# Disable streaming when tools are present
stream = False
# Handle online mode
model_id = self.selected_model["id"]
enable_web_search = False
web_search_config = {}
if self.online_enabled:
# OpenRouter handles online mode natively with :online suffix
if self.client.provider_name == "openrouter":
if hasattr(self.client.provider, "get_effective_model_id"):
model_id = self.client.provider.get_effective_model_id(model_id, True)
# Anthropic has native web search when search provider is set to anthropic_native
elif self.client.provider_name == "anthropic" and self.settings.search_provider == "anthropic_native":
enable_web_search = True
web_search_config = {
"max_uses": self.settings.search_num_results or 5
}
logger.info("Using Anthropic native web search")
else:
# For other providers, perform web search and inject results
logger.info(f"Performing web search for: {user_input}")
search_results = await asyncio.to_thread(
perform_web_search,
user_input,
num_results=self.settings.search_num_results,
provider=self.settings.search_provider,
google_api_key=self.settings.google_api_key,
google_search_engine_id=self.settings.google_search_engine_id
)
if search_results:
# Inject search results into messages
formatted_results = format_search_results(search_results)
search_context = f"\n\n{formatted_results}\n\nPlease use the above web search results to help answer the user's question."
# Add search results to the last user message
if messages and messages[-1]["role"] == "user":
messages[-1]["content"] += search_context
logger.info(f"Injected {len(search_results)} search results into context")
transforms = ["middle-out"] if self.middle_out_enabled else None
max_tokens = None
if self.session_max_token > 0:
max_tokens = self.session_max_token
if tools:
# Use async tool handling flow
async for chunk in self._send_with_tools_async(
messages=messages,
model_id=model_id,
tools=tools,
max_tokens=max_tokens,
transforms=transforms,
):
yield chunk
elif stream:
# Use async streaming flow
async for chunk in self._stream_response_async(
messages=messages,
model_id=model_id,
max_tokens=max_tokens,
transforms=transforms,
enable_web_search=enable_web_search,
web_search_config=web_search_config,
):
yield chunk
else:
# Non-streaming request
response = self.client.chat(
messages=messages,
model=model_id,
stream=False,
max_tokens=max_tokens,
transforms=transforms,
)
if isinstance(response, ChatResponse):
# Yield single chunk with complete response
chunk = StreamChunk(
id="",
delta_content=response.content,
usage=response.usage,
error=None,
)
yield chunk
async def _send_with_tools_async(
self,
messages: List[Dict[str, Any]],
model_id: str,
tools: List[Dict[str, Any]],
max_tokens: Optional[int] = None,
transforms: Optional[List[str]] = None,
) -> AsyncIterator[StreamChunk]:
"""
Async version of _send_with_tools for TUI.
Args:
messages: API messages
model_id: Model ID
tools: Tool definitions
max_tokens: Max tokens
transforms: Transforms list
Yields:
StreamChunk objects including tool call notifications
"""
max_loops = 5
loop_count = 0
api_messages = list(messages)
while loop_count < max_loops:
response = self.client.chat(
messages=api_messages,
model=model_id,
stream=False,
max_tokens=max_tokens,
tools=tools,
tool_choice="auto",
transforms=transforms,
)
if not isinstance(response, ChatResponse):
raise ValueError("Expected ChatResponse")
tool_calls = response.tool_calls
if not tool_calls:
# Final response, yield it
chunk = StreamChunk(
id="",
delta_content=response.content,
usage=response.usage,
error=None,
)
yield chunk
return
# Yield notification about tool calls
tool_notification = f"\n🔧 AI requesting {len(tool_calls)} tool call(s)...\n"
yield StreamChunk(id="", delta_content=tool_notification, usage=None, error=None)
tool_results = []
for tc in tool_calls:
try:
args = json.loads(tc.function.arguments)
except json.JSONDecodeError as e:
self.logger.error(f"Failed to parse tool arguments: {e}")
tool_results.append({
"tool_call_id": tc.id,
"role": "tool",
"name": tc.function.name,
"content": json.dumps({"error": f"Invalid arguments: {e}"}),
})
continue
# Yield tool call display
args_display = ", ".join(
f'{k}="{v}"' if isinstance(v, str) else f"{k}={v}"
for k, v in args.items()
)
tool_display = f"{tc.function.name}({args_display})\n"
yield StreamChunk(id="", delta_content=tool_display, usage=None, error=None)
# Execute tool (await instead of asyncio.run)
result = await self.execute_tool(tc.function.name, args)
if "error" in result:
error_msg = f" ✗ Error: {result['error']}\n"
yield StreamChunk(id="", delta_content=error_msg, usage=None, error=None)
else:
success_msg = self._format_tool_success(tc.function.name, result)
yield StreamChunk(id="", delta_content=success_msg, usage=None, error=None)
tool_results.append({
"tool_call_id": tc.id,
"role": "tool",
"name": tc.function.name,
"content": json.dumps(result),
})
# Add assistant message with tool calls
api_messages.append({
"role": "assistant",
"content": response.content,
"tool_calls": [
{
"id": tc.id,
"type": tc.type,
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in tool_calls
],
})
# Add tool results
api_messages.extend(tool_results)
loop_count += 1
# Max loops reached
yield StreamChunk(
id="",
delta_content="\n⚠️ Maximum tool call loops reached\n",
usage=None,
error="Max loops reached"
)
def _format_tool_success(self, tool_name: str, result: Dict[str, Any]) -> str:
"""Format a success message for a tool call."""
if tool_name == "search_files":
count = result.get("count", 0)
return f" ✓ Found {count} file(s)\n"
elif tool_name == "read_file":
size = result.get("size", 0)
truncated = " (truncated)" if result.get("truncated") else ""
return f" ✓ Read {size} bytes{truncated}\n"
elif tool_name == "list_directory":
count = result.get("count", 0)
return f" ✓ Listed {count} item(s)\n"
elif tool_name == "inspect_database":
if "table" in result:
return f" ✓ Inspected table: {result['table']}\n"
else:
return f" ✓ Inspected database ({result.get('table_count', 0)} tables)\n"
elif tool_name == "search_database":
count = result.get("count", 0)
return f" ✓ Found {count} match(es)\n"
elif tool_name == "query_database":
count = result.get("count", 0)
return f" ✓ Query returned {count} row(s)\n"
else:
return " ✓ Success\n"
async def _stream_response_async(
self,
messages: List[Dict[str, Any]],
model_id: str,
max_tokens: Optional[int] = None,
transforms: Optional[List[str]] = None,
enable_web_search: bool = False,
web_search_config: Optional[Dict[str, Any]] = None,
) -> AsyncIterator[StreamChunk]:
"""
Async version of _stream_response for TUI.
Args:
messages: API messages
model_id: Model ID
max_tokens: Max tokens
transforms: Transforms
enable_web_search: Whether to enable Anthropic native web search
web_search_config: Web search configuration
Yields:
StreamChunk objects
"""
response = self.client.chat(
messages=messages,
model=model_id,
stream=True,
max_tokens=max_tokens,
transforms=transforms,
enable_web_search=enable_web_search,
web_search_config=web_search_config or {},
)
if isinstance(response, ChatResponse):
# Non-streaming response
chunk = StreamChunk(
id="",
delta_content=response.content,
usage=response.usage,
error=None,
)
yield chunk
return
# Stream the response
for chunk in response:
if chunk.error:
yield StreamChunk(id="", delta_content=None, usage=None, error=chunk.error)
break
yield chunk
# ========== END ASYNC METHODS ==========
def add_to_history(
self,
prompt: str,
response: str,
usage: Optional[UsageStats] = None,
cost: float = 0.0,
) -> None:
"""
Add an exchange to the history.
Args:
prompt: User prompt
response: Assistant response
usage: Usage statistics
cost: Cost if not in usage
"""
entry = HistoryEntry(
prompt=prompt,
response=response,
prompt_tokens=usage.prompt_tokens if usage else 0,
completion_tokens=usage.completion_tokens if usage else 0,
msg_cost=usage.total_cost_usd if usage and usage.total_cost_usd else cost,
timestamp=time.time(),
)
self.history.append(entry)
self.current_index = len(self.history) - 1
self.stats.add_usage(usage, cost)
def save_conversation(self, name: str) -> bool:
"""
Save the current conversation.
Args:
name: Name for the saved conversation
Returns:
True if saved successfully
"""
if not self.history:
return False
data = [e.to_dict() for e in self.history]
self.db.save_conversation(name, data)
self.logger.info(f"Saved conversation: {name}")
return True
def load_conversation(self, name: str) -> bool:
"""
Load a saved conversation.
Args:
name: Name of the conversation to load
Returns:
True if loaded successfully
"""
data = self.db.load_conversation(name)
if not data:
return False
self.history.clear()
for entry_dict in data:
self.history.append(HistoryEntry(
prompt=entry_dict.get("prompt", ""),
response=entry_dict.get("response", ""),
prompt_tokens=entry_dict.get("prompt_tokens", 0),
completion_tokens=entry_dict.get("completion_tokens", 0),
msg_cost=entry_dict.get("msg_cost", 0.0),
))
self.current_index = len(self.history) - 1
self.memory_start_index = 0
self.stats = SessionStats() # Reset stats for loaded conversation
self.logger.info(f"Loaded conversation: {name}")
return True
def reset(self) -> None:
"""Reset the session state."""
self.history.clear()
self.stats = SessionStats()
self.system_prompt = ""
self.memory_start_index = 0
self.current_index = 0
self.logger.info("Session reset")
def check_warnings(self) -> List[str]:
"""
Check for cost and credit warnings.
Returns:
List of warning messages
"""
warnings = []
# Check last message cost
if self.history:
last_cost = self.history[-1].msg_cost
threshold = self.settings.cost_warning_threshold
if last_cost > threshold:
warnings.append(
f"High cost: ${last_cost:.4f} exceeds threshold ${threshold:.4f}"
)
# Check credits
credits = self.client.get_credits()
if credits:
left = credits.get("credits_left", 0)
total = credits.get("total_credits", 0)
if left < LOW_CREDIT_AMOUNT:
warnings.append(f"Low credits: ${left:.2f} remaining!")
elif total > 0 and left < total * LOW_CREDIT_RATIO:
warnings.append(f"Credits low: less than 10% remaining (${left:.2f})")
return warnings
+28
View File
@@ -0,0 +1,28 @@
"""
Model Context Protocol (MCP) integration for oAI.
This package provides filesystem and database access capabilities
through the MCP standard, allowing AI models to interact with
local files and SQLite databases safely.
Key components:
- MCPManager: High-level manager for MCP operations
- MCPFilesystemServer: Filesystem and database access implementation
- GitignoreParser: Pattern matching for .gitignore support
- SQLiteQueryValidator: Query safety validation
- CrossPlatformMCPConfig: OS-specific configuration
"""
from oai.mcp.manager import MCPManager
from oai.mcp.server import MCPFilesystemServer
from oai.mcp.gitignore import GitignoreParser
from oai.mcp.validators import SQLiteQueryValidator
from oai.mcp.platform import CrossPlatformMCPConfig
__all__ = [
"MCPManager",
"MCPFilesystemServer",
"GitignoreParser",
"SQLiteQueryValidator",
"CrossPlatformMCPConfig",
]
+166
View File
@@ -0,0 +1,166 @@
"""
Gitignore pattern parsing for oAI MCP.
This module implements .gitignore pattern matching to filter files
during MCP filesystem operations.
"""
import fnmatch
from pathlib import Path
from typing import List, Tuple
from oai.utils.logging import get_logger
class GitignoreParser:
"""
Parse and apply .gitignore patterns.
Supports standard gitignore syntax including:
- Wildcards (*) and double wildcards (**)
- Directory-only patterns (ending with /)
- Negation patterns (starting with !)
- Comments (lines starting with #)
Patterns are applied relative to the directory containing
the .gitignore file.
"""
def __init__(self):
"""Initialize an empty pattern collection."""
# List of (pattern, is_negation, source_dir)
self.patterns: List[Tuple[str, bool, Path]] = []
def add_gitignore(self, gitignore_path: Path) -> None:
"""
Parse and add patterns from a .gitignore file.
Args:
gitignore_path: Path to the .gitignore file
"""
logger = get_logger()
if not gitignore_path.exists():
return
try:
source_dir = gitignore_path.parent
with open(gitignore_path, "r", encoding="utf-8") as f:
for line_num, line in enumerate(f, 1):
line = line.rstrip("\n\r")
# Skip empty lines and comments
if not line or line.startswith("#"):
continue
# Check for negation pattern
is_negation = line.startswith("!")
if is_negation:
line = line[1:]
# Remove leading slash (make relative to gitignore location)
if line.startswith("/"):
line = line[1:]
self.patterns.append((line, is_negation, source_dir))
logger.debug(
f"Loaded {len(self.patterns)} patterns from {gitignore_path}"
)
except Exception as e:
logger.warning(f"Error reading {gitignore_path}: {e}")
def should_ignore(self, path: Path) -> bool:
"""
Check if a path should be ignored based on gitignore patterns.
Patterns are evaluated in order, with later patterns overriding
earlier ones. Negation patterns (starting with !) un-ignore
previously matched paths.
Args:
path: Path to check
Returns:
True if the path should be ignored
"""
if not self.patterns:
return False
ignored = False
for pattern, is_negation, source_dir in self.patterns:
# Only apply pattern if path is under the source directory
try:
rel_path = path.relative_to(source_dir)
except ValueError:
# Path is not relative to this gitignore's directory
continue
rel_path_str = str(rel_path)
# Check if pattern matches
if self._match_pattern(pattern, rel_path_str, path.is_dir()):
if is_negation:
ignored = False # Negation patterns un-ignore
else:
ignored = True
return ignored
def _match_pattern(self, pattern: str, path: str, is_dir: bool) -> bool:
"""
Match a gitignore pattern against a path.
Args:
pattern: The gitignore pattern
path: The relative path string to match
is_dir: Whether the path is a directory
Returns:
True if the pattern matches
"""
# Directory-only pattern (ends with /)
if pattern.endswith("/"):
if not is_dir:
return False
pattern = pattern[:-1]
# Handle ** patterns (matches any number of directories)
if "**" in pattern:
pattern_parts = pattern.split("**")
if len(pattern_parts) == 2:
prefix, suffix = pattern_parts
# Match if path starts with prefix and ends with suffix
if prefix:
if not path.startswith(prefix.rstrip("/")):
return False
if suffix:
suffix = suffix.lstrip("/")
if not (path.endswith(suffix) or f"/{suffix}" in path):
return False
return True
# Direct match using fnmatch
if fnmatch.fnmatch(path, pattern):
return True
# Match as subdirectory pattern (pattern without / matches in any directory)
if "/" not in pattern:
parts = path.split("/")
if any(fnmatch.fnmatch(part, pattern) for part in parts):
return True
return False
def clear(self) -> None:
"""Clear all loaded patterns."""
self.patterns = []
@property
def pattern_count(self) -> int:
"""Get the number of loaded patterns."""
return len(self.patterns)
+1365
View File
File diff suppressed because it is too large Load Diff
+228
View File
@@ -0,0 +1,228 @@
"""
Cross-platform MCP configuration for oAI.
This module handles OS-specific configuration, path handling,
and security checks for the MCP filesystem server.
"""
import os
import platform
import subprocess
from pathlib import Path
from typing import List, Dict, Any, Optional
from oai.constants import SYSTEM_DIRS_BLACKLIST
from oai.utils.logging import get_logger
class CrossPlatformMCPConfig:
"""
Handle OS-specific MCP configuration.
Provides methods for path normalization, security validation,
and OS-specific default directories.
Attributes:
system: Operating system name
is_macos: Whether running on macOS
is_linux: Whether running on Linux
is_windows: Whether running on Windows
"""
def __init__(self):
"""Initialize platform detection."""
self.system = platform.system()
self.is_macos = self.system == "Darwin"
self.is_linux = self.system == "Linux"
self.is_windows = self.system == "Windows"
logger = get_logger()
logger.info(f"Detected OS: {self.system}")
def get_default_allowed_dirs(self) -> List[Path]:
"""
Get safe default directories for the current OS.
Returns:
List of default directories that are safe to access
"""
home = Path.home()
if self.is_macos:
return [
home / "Documents",
home / "Desktop",
home / "Downloads",
]
elif self.is_linux:
dirs = [home / "Documents"]
# Try to get XDG directories
try:
for xdg_dir in ["DOCUMENTS", "DESKTOP", "DOWNLOAD"]:
result = subprocess.run(
["xdg-user-dir", xdg_dir],
capture_output=True,
text=True,
timeout=1
)
if result.returncode == 0:
dir_path = Path(result.stdout.strip())
if dir_path.exists():
dirs.append(dir_path)
except (subprocess.TimeoutExpired, FileNotFoundError):
# Fallback to standard locations
dirs.extend([
home / "Desktop",
home / "Downloads",
])
return list(set(dirs))
elif self.is_windows:
return [
home / "Documents",
home / "Desktop",
home / "Downloads",
]
# Fallback for unknown OS
return [home]
def get_python_command(self) -> str:
"""
Get the Python executable path.
Returns:
Path to the Python executable
"""
import sys
return sys.executable
def get_filesystem_warning(self) -> str:
"""
Get OS-specific security warning message.
Returns:
Warning message for the current OS
"""
if self.is_macos:
return """
Note: macOS Security
The Filesystem MCP server needs access to your selected folder.
You may see a security prompt - click 'Allow' to proceed.
(System Settings > Privacy & Security > Files and Folders)
"""
elif self.is_linux:
return """
Note: Linux Security
The Filesystem MCP server will access your selected folder.
Ensure oAI has appropriate file permissions.
"""
elif self.is_windows:
return """
Note: Windows Security
The Filesystem MCP server will access your selected folder.
You may need to grant file access permissions.
"""
return ""
def normalize_path(self, path: str) -> Path:
"""
Normalize a path for the current OS.
Expands user directory (~) and resolves to absolute path.
Args:
path: Path string to normalize
Returns:
Normalized absolute Path
"""
return Path(os.path.expanduser(path)).resolve()
def is_system_directory(self, path: Path) -> bool:
"""
Check if a path is a protected system directory.
Args:
path: Path to check
Returns:
True if the path is a system directory
"""
path_str = str(path)
for blocked in SYSTEM_DIRS_BLACKLIST:
if path_str.startswith(blocked):
return True
return False
def is_safe_path(self, requested_path: Path, allowed_dirs: List[Path]) -> bool:
"""
Check if a path is within allowed directories.
Args:
requested_path: Path being requested
allowed_dirs: List of allowed parent directories
Returns:
True if the path is within an allowed directory
"""
try:
requested = requested_path.resolve()
for allowed in allowed_dirs:
try:
allowed_resolved = allowed.resolve()
requested.relative_to(allowed_resolved)
return True
except ValueError:
continue
return False
except Exception:
return False
def get_folder_stats(self, folder: Path) -> Dict[str, Any]:
"""
Get statistics for a folder.
Args:
folder: Path to the folder
Returns:
Dictionary with folder statistics:
- exists: Whether the folder exists
- file_count: Number of files (if exists)
- total_size: Total size in bytes (if exists)
- size_mb: Size in megabytes (if exists)
- error: Error message (if any)
"""
logger = get_logger()
try:
if not folder.exists() or not folder.is_dir():
return {"exists": False}
file_count = 0
total_size = 0
for item in folder.rglob("*"):
if item.is_file():
file_count += 1
try:
total_size += item.stat().st_size
except (OSError, PermissionError):
pass
return {
"exists": True,
"file_count": file_count,
"total_size": total_size,
"size_mb": total_size / (1024 * 1024),
}
except Exception as e:
logger.error(f"Error getting folder stats for {folder}: {e}")
return {"exists": False, "error": str(e)}
+1368
View File
File diff suppressed because it is too large Load Diff
+123
View File
@@ -0,0 +1,123 @@
"""
Query validation for oAI MCP database operations.
This module provides safety validation for SQL queries to ensure
only read-only operations are executed.
"""
import re
from typing import Tuple
from oai.constants import DANGEROUS_SQL_KEYWORDS
class SQLiteQueryValidator:
"""
Validate SQLite queries for read-only safety.
Ensures that only SELECT queries (including CTEs) are allowed
and blocks potentially dangerous operations like INSERT, UPDATE,
DELETE, DROP, etc.
"""
@staticmethod
def is_safe_query(query: str) -> Tuple[bool, str]:
"""
Validate that a query is a safe read-only SELECT.
The validation:
1. Checks that query starts with SELECT or WITH
2. Strips string literals before checking for dangerous keywords
3. Blocks any dangerous keywords outside of string literals
Args:
query: SQL query string to validate
Returns:
Tuple of (is_safe, error_message)
- is_safe: True if the query is safe to execute
- error_message: Description of why query is unsafe (empty if safe)
Examples:
>>> SQLiteQueryValidator.is_safe_query("SELECT * FROM users")
(True, "")
>>> SQLiteQueryValidator.is_safe_query("DELETE FROM users")
(False, "Only SELECT queries are allowed...")
>>> SQLiteQueryValidator.is_safe_query("SELECT 'DELETE' FROM users")
(True, "") # 'DELETE' is inside a string literal
"""
query_upper = query.strip().upper()
# Must start with SELECT or WITH (for CTEs)
if not (query_upper.startswith("SELECT") or query_upper.startswith("WITH")):
return False, "Only SELECT queries are allowed (including WITH/CTE)"
# Remove string literals before checking for dangerous keywords
# This prevents false positives when keywords appear in data
query_no_strings = re.sub(r"'[^']*'", "", query_upper)
query_no_strings = re.sub(r'"[^"]*"', "", query_no_strings)
# Check for dangerous keywords outside of quotes
for keyword in DANGEROUS_SQL_KEYWORDS:
if re.search(r"\b" + keyword + r"\b", query_no_strings):
return False, f"Keyword '{keyword}' not allowed in read-only mode"
return True, ""
@staticmethod
def sanitize_table_name(table_name: str) -> str:
"""
Sanitize a table name to prevent SQL injection.
Only allows alphanumeric characters and underscores.
Args:
table_name: Table name to sanitize
Returns:
Sanitized table name
Raises:
ValueError: If table name contains invalid characters
"""
# Remove any characters that aren't alphanumeric or underscore
sanitized = re.sub(r"[^\w]", "", table_name)
if not sanitized:
raise ValueError("Table name cannot be empty after sanitization")
if sanitized != table_name:
raise ValueError(
f"Table name contains invalid characters: {table_name}"
)
return sanitized
@staticmethod
def sanitize_column_name(column_name: str) -> str:
"""
Sanitize a column name to prevent SQL injection.
Only allows alphanumeric characters and underscores.
Args:
column_name: Column name to sanitize
Returns:
Sanitized column name
Raises:
ValueError: If column name contains invalid characters
"""
# Remove any characters that aren't alphanumeric or underscore
sanitized = re.sub(r"[^\w]", "", column_name)
if not sanitized:
raise ValueError("Column name cannot be empty after sanitization")
if sanitized != column_name:
raise ValueError(
f"Column name contains invalid characters: {column_name}"
)
return sanitized
+45
View File
@@ -0,0 +1,45 @@
"""
Provider abstraction for oAI.
This module provides a unified interface for AI model providers,
enabling easy extension to support additional providers beyond OpenRouter.
"""
from oai.providers.base import (
AIProvider,
ChatMessage,
ChatResponse,
ToolCall,
ToolFunction,
UsageStats,
ModelInfo,
ProviderCapabilities,
)
from oai.providers.openrouter import OpenRouterProvider
from oai.providers.anthropic import AnthropicProvider
from oai.providers.openai import OpenAIProvider
from oai.providers.ollama import OllamaProvider
from oai.providers.registry import register_provider
# Register all providers
register_provider("openrouter", OpenRouterProvider)
register_provider("anthropic", AnthropicProvider)
register_provider("openai", OpenAIProvider)
register_provider("ollama", OllamaProvider)
__all__ = [
# Base classes and types
"AIProvider",
"ChatMessage",
"ChatResponse",
"ToolCall",
"ToolFunction",
"UsageStats",
"ModelInfo",
"ProviderCapabilities",
# Provider implementations
"OpenRouterProvider",
"AnthropicProvider",
"OpenAIProvider",
"OllamaProvider",
]
+673
View File
@@ -0,0 +1,673 @@
"""
Anthropic provider for Claude models.
This provider connects to Anthropic's API for accessing Claude models.
"""
import json
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
import anthropic
from anthropic.types import Message, MessageStreamEvent
from oai.constants import ANTHROPIC_BASE_URL
from oai.providers.base import (
AIProvider,
ChatMessage,
ChatResponse,
ChatResponseChoice,
ModelInfo,
ProviderCapabilities,
StreamChunk,
ToolCall,
ToolFunction,
UsageStats,
)
from oai.utils.logging import get_logger
logger = get_logger()
# Model name aliases
MODEL_ALIASES = {
"claude-sonnet": "claude-sonnet-4-5-20250929",
"claude-haiku": "claude-haiku-4-5-20251001",
"claude-opus": "claude-opus-4-5-20251101",
# Legacy aliases
"claude-3-haiku": "claude-3-haiku-20240307",
"claude-3-7-sonnet": "claude-3-7-sonnet-20250219",
}
class AnthropicProvider(AIProvider):
"""
Anthropic API provider.
Provides access to Claude 3.5 Sonnet, Claude 3 Opus, and other Anthropic models.
"""
def __init__(
self,
api_key: str,
base_url: Optional[str] = None,
app_name: str = "oAI",
app_url: str = "",
**kwargs: Any,
):
"""
Initialize Anthropic provider.
Args:
api_key: Anthropic API key
base_url: Optional custom base URL
app_name: Application name (for headers)
app_url: Application URL (for headers)
**kwargs: Additional arguments
"""
super().__init__(api_key, base_url or ANTHROPIC_BASE_URL)
self.client = anthropic.Anthropic(api_key=api_key)
self.async_client = anthropic.AsyncAnthropic(api_key=api_key)
self._models_cache: Optional[List[ModelInfo]] = None
def _create_web_search_tool(self, config: Dict[str, Any]) -> Dict[str, Any]:
"""
Create Anthropic native web search tool definition.
Args:
config: Optional configuration for web search (max_uses, allowed_domains, etc.)
Returns:
Tool definition dict
"""
tool: Dict[str, Any] = {
"type": "web_search_20250305",
"name": "web_search",
}
# Add optional parameters if provided
if "max_uses" in config:
tool["max_uses"] = config["max_uses"]
else:
tool["max_uses"] = 5 # Default
if "allowed_domains" in config:
tool["allowed_domains"] = config["allowed_domains"]
if "blocked_domains" in config:
tool["blocked_domains"] = config["blocked_domains"]
if "user_location" in config:
tool["user_location"] = config["user_location"]
return tool
@property
def name(self) -> str:
"""Get provider name."""
return "Anthropic"
@property
def capabilities(self) -> ProviderCapabilities:
"""Get provider capabilities."""
return ProviderCapabilities(
streaming=True,
tools=True,
images=True,
online=True, # Web search via DuckDuckGo/Google
max_context=200000,
)
def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
"""
List available Anthropic models.
Args:
filter_text_only: Whether to filter for text models only
Returns:
List of ModelInfo objects
"""
if self._models_cache:
return self._models_cache
# Anthropic doesn't have a models list API, so we hardcode the available models
models = [
# Current Claude 4.5 models
ModelInfo(
id="claude-sonnet-4-5-20250929",
name="Claude Sonnet 4.5",
description="Smart model for complex agents and coding (recommended)",
context_length=200000,
pricing={"input": 3.0, "output": 15.0},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
ModelInfo(
id="claude-haiku-4-5-20251001",
name="Claude Haiku 4.5",
description="Fastest model with near-frontier intelligence",
context_length=200000,
pricing={"input": 1.0, "output": 5.0},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
ModelInfo(
id="claude-opus-4-5-20251101",
name="Claude Opus 4.5",
description="Premium model with maximum intelligence",
context_length=200000,
pricing={"input": 5.0, "output": 25.0},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
# Legacy models (still available)
ModelInfo(
id="claude-3-7-sonnet-20250219",
name="Claude Sonnet 3.7",
description="Legacy model - recommend migrating to 4.5",
context_length=200000,
pricing={"input": 3.0, "output": 15.0},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
ModelInfo(
id="claude-3-haiku-20240307",
name="Claude 3 Haiku",
description="Legacy fast model - recommend migrating to 4.5",
context_length=200000,
pricing={"input": 0.25, "output": 1.25},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
]
self._models_cache = models
logger.info(f"Loaded {len(models)} Anthropic models")
return models
def get_model(self, model_id: str) -> Optional[ModelInfo]:
"""
Get information about a specific model.
Args:
model_id: Model identifier
Returns:
ModelInfo or None
"""
# Resolve alias
resolved_id = MODEL_ALIASES.get(model_id, model_id)
models = self.list_models()
for model in models:
if model.id == resolved_id or model.id == model_id:
return model
return None
def chat(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, Iterator[StreamChunk]]:
"""
Send chat completion request to Anthropic.
Args:
model: Model ID
messages: Chat messages
stream: Whether to stream response
max_tokens: Maximum tokens
temperature: Sampling temperature
tools: Tool definitions
tool_choice: Tool selection mode
**kwargs: Additional parameters (including enable_web_search)
Returns:
ChatResponse or Iterator[StreamChunk]
"""
# Resolve model alias
model_id = MODEL_ALIASES.get(model, model)
# Extract system message (Anthropic requires it separate from messages)
system_prompt, anthropic_messages = self._convert_messages(messages)
# Build request parameters
params: Dict[str, Any] = {
"model": model_id,
"messages": anthropic_messages,
"max_tokens": max_tokens or 4096,
}
if system_prompt:
params["system"] = system_prompt
if temperature is not None:
params["temperature"] = temperature
# Prepare tools list
tools_list = []
# Add web search tool if requested via kwargs
if kwargs.get("enable_web_search", False):
web_search_config = kwargs.get("web_search_config", {})
tools_list.append(self._create_web_search_tool(web_search_config))
logger.info("Added Anthropic native web search tool")
# Add user-provided tools
if tools:
# Convert tools to Anthropic format
converted_tools = self._convert_tools(tools)
tools_list.extend(converted_tools)
if tools_list:
params["tools"] = tools_list
if tool_choice and tool_choice != "auto":
# Anthropic uses different tool_choice format
if tool_choice == "none":
pass # Don't include tools
elif tool_choice == "required":
params["tool_choice"] = {"type": "any"}
else:
params["tool_choice"] = {"type": "tool", "name": tool_choice}
logger.debug(f"Anthropic request: model={model_id}, messages={len(anthropic_messages)}")
try:
if stream:
return self._stream_chat(params)
else:
return self._sync_chat(params)
except Exception as e:
logger.error(f"Anthropic request failed: {e}")
return ChatResponse(
id="error",
choices=[
ChatResponseChoice(
index=0,
message=ChatMessage(role="assistant", content=f"Error: {str(e)}"),
finish_reason="error",
)
],
)
def _convert_messages(self, messages: List[ChatMessage]) -> tuple[str, List[Dict[str, Any]]]:
"""
Convert messages to Anthropic format.
Anthropic requires system messages to be separate from the conversation.
Args:
messages: List of ChatMessage objects
Returns:
Tuple of (system_prompt, anthropic_messages)
"""
system_prompt = ""
anthropic_messages = []
for msg in messages:
if msg.role == "system":
# Accumulate system messages
if system_prompt:
system_prompt += "\n\n"
system_prompt += msg.content or ""
else:
# Convert to Anthropic format
message_dict: Dict[str, Any] = {"role": msg.role}
# Handle content
if msg.content:
message_dict["content"] = msg.content
# Handle tool calls (assistant messages)
if msg.tool_calls:
# Anthropic format for tool use
content_blocks = []
if msg.content:
content_blocks.append({"type": "text", "text": msg.content})
for tc in msg.tool_calls:
content_blocks.append({
"type": "tool_use",
"id": tc.id,
"name": tc.function.name,
"input": json.loads(tc.function.arguments),
})
message_dict["content"] = content_blocks
# Handle tool results (tool messages)
if msg.role == "tool" and msg.tool_call_id:
# Convert to Anthropic's tool_result format
anthropic_messages.append({
"role": "user",
"content": [{
"type": "tool_result",
"tool_use_id": msg.tool_call_id,
"content": msg.content or "",
}]
})
continue
anthropic_messages.append(message_dict)
return system_prompt, anthropic_messages
def _convert_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Convert OpenAI-style tools to Anthropic format.
Args:
tools: OpenAI tool definitions
Returns:
Anthropic tool definitions
"""
anthropic_tools = []
for tool in tools:
if tool.get("type") == "function":
func = tool.get("function", {})
anthropic_tools.append({
"name": func.get("name"),
"description": func.get("description", ""),
"input_schema": func.get("parameters", {}),
})
return anthropic_tools
def _sync_chat(self, params: Dict[str, Any]) -> ChatResponse:
"""
Send synchronous chat request.
Args:
params: Request parameters
Returns:
ChatResponse
"""
message: Message = self.client.messages.create(**params)
# Extract content
content = ""
tool_calls = []
for block in message.content:
if block.type == "text":
content += block.text
elif block.type == "tool_use":
# Convert to ToolCall format
tool_calls.append(
ToolCall(
id=block.id,
type="function",
function=ToolFunction(
name=block.name,
arguments=json.dumps(block.input),
),
)
)
# Build ChatMessage
chat_message = ChatMessage(
role="assistant",
content=content if content else None,
tool_calls=tool_calls if tool_calls else None,
)
# Extract usage
usage = None
if message.usage:
usage = UsageStats(
prompt_tokens=message.usage.input_tokens,
completion_tokens=message.usage.output_tokens,
total_tokens=message.usage.input_tokens + message.usage.output_tokens,
)
return ChatResponse(
id=message.id,
choices=[
ChatResponseChoice(
index=0,
message=chat_message,
finish_reason=message.stop_reason,
)
],
usage=usage,
model=message.model,
)
def _stream_chat(self, params: Dict[str, Any]) -> Iterator[StreamChunk]:
"""
Stream chat response from Anthropic.
Args:
params: Request parameters
Yields:
StreamChunk objects
"""
stream = self.client.messages.stream(**params)
with stream as event_stream:
for event in event_stream:
event_data: MessageStreamEvent = event
# Handle different event types
if event_data.type == "content_block_delta":
delta = event_data.delta
if hasattr(delta, "text"):
yield StreamChunk(
id="stream",
delta_content=delta.text,
)
elif event_data.type == "message_stop":
# Final event with usage
pass
elif event_data.type == "message_delta":
# Contains stop reason and usage
usage = None
if hasattr(event_data, "usage"):
usage_data = event_data.usage
usage = UsageStats(
prompt_tokens=0,
completion_tokens=usage_data.output_tokens,
total_tokens=usage_data.output_tokens,
)
yield StreamChunk(
id="stream",
finish_reason=event_data.delta.stop_reason if hasattr(event_data.delta, "stop_reason") else None,
usage=usage,
)
async def chat_async(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
"""
Send async chat request to Anthropic.
Args:
model: Model ID
messages: Chat messages
stream: Whether to stream
max_tokens: Max tokens
temperature: Temperature
tools: Tool definitions
tool_choice: Tool choice
**kwargs: Additional args
Returns:
ChatResponse or AsyncIterator[StreamChunk]
"""
# Resolve model alias
model_id = MODEL_ALIASES.get(model, model)
# Convert messages
system_prompt, anthropic_messages = self._convert_messages(messages)
# Build params
params: Dict[str, Any] = {
"model": model_id,
"messages": anthropic_messages,
"max_tokens": max_tokens or 4096,
}
if system_prompt:
params["system"] = system_prompt
if temperature is not None:
params["temperature"] = temperature
if tools:
params["tools"] = self._convert_tools(tools)
if stream:
return self._stream_chat_async(params)
else:
message = await self.async_client.messages.create(**params)
return self._convert_message(message)
async def _stream_chat_async(self, params: Dict[str, Any]) -> AsyncIterator[StreamChunk]:
"""
Stream async chat response.
Args:
params: Request parameters
Yields:
StreamChunk objects
"""
stream = await self.async_client.messages.stream(**params)
async with stream as event_stream:
async for event in event_stream:
if event.type == "content_block_delta":
delta = event.delta
if hasattr(delta, "text"):
yield StreamChunk(
id="stream",
delta_content=delta.text,
)
def _convert_message(self, message: Message) -> ChatResponse:
"""Helper to convert Anthropic message to ChatResponse."""
content = ""
tool_calls = []
for block in message.content:
if block.type == "text":
content += block.text
elif block.type == "tool_use":
tool_calls.append(
ToolCall(
id=block.id,
type="function",
function=ToolFunction(
name=block.name,
arguments=json.dumps(block.input),
),
)
)
chat_message = ChatMessage(
role="assistant",
content=content if content else None,
tool_calls=tool_calls if tool_calls else None,
)
usage = None
if message.usage:
usage = UsageStats(
prompt_tokens=message.usage.input_tokens,
completion_tokens=message.usage.output_tokens,
total_tokens=message.usage.input_tokens + message.usage.output_tokens,
)
return ChatResponse(
id=message.id,
choices=[
ChatResponseChoice(
index=0,
message=chat_message,
finish_reason=message.stop_reason,
)
],
usage=usage,
model=message.model,
)
def get_credits(self) -> Optional[Dict[str, Any]]:
"""
Get account credits from Anthropic.
Note: Anthropic does not currently provide a public API endpoint
for checking account credits/balance. This information is only
available through the Anthropic Console web interface.
Returns:
None (credits API not available)
"""
# Anthropic doesn't provide a public credits API endpoint
# Users must check their balance at console.anthropic.com
return None
def clear_cache(self) -> None:
"""Clear model cache."""
self._models_cache = None
def get_raw_models(self) -> List[Dict[str, Any]]:
"""
Get raw model data as dictionaries.
Returns:
List of model dictionaries
"""
models = self.list_models()
return [
{
"id": model.id,
"name": model.name,
"description": model.description,
"context_length": model.context_length,
"pricing": model.pricing,
}
for model in models
]
def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
"""
Get raw model data for a specific model.
Args:
model_id: Model identifier
Returns:
Model dictionary or None
"""
model = self.get_model(model_id)
if model:
return {
"id": model.id,
"name": model.name,
"description": model.description,
"context_length": model.context_length,
"pricing": model.pricing,
}
return None
+413
View File
@@ -0,0 +1,413 @@
"""
Abstract base provider for AI model integration.
This module defines the interface that all AI providers must implement,
along with common data structures for requests and responses.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import (
Any,
AsyncIterator,
Dict,
Iterator,
List,
Optional,
Union,
)
class MessageRole(str, Enum):
"""Message roles in a conversation."""
SYSTEM = "system"
USER = "user"
ASSISTANT = "assistant"
TOOL = "tool"
@dataclass
class ToolFunction:
"""
Represents a function within a tool call.
Attributes:
name: The function name
arguments: JSON string of function arguments
"""
name: str
arguments: str
@dataclass
class ToolCall:
"""
Represents a tool/function call requested by the model.
Attributes:
id: Unique identifier for this tool call
type: Type of tool call (usually "function")
function: The function being called
"""
id: str
type: str
function: ToolFunction
@dataclass
class UsageStats:
"""
Token usage statistics from an API response.
Attributes:
prompt_tokens: Number of tokens in the prompt
completion_tokens: Number of tokens in the completion
total_tokens: Total tokens used
total_cost_usd: Cost in USD (if available from API)
"""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
total_cost_usd: Optional[float] = None
@property
def input_tokens(self) -> int:
"""Alias for prompt_tokens."""
return self.prompt_tokens
@property
def output_tokens(self) -> int:
"""Alias for completion_tokens."""
return self.completion_tokens
@dataclass
class ChatMessage:
"""
A single message in a chat conversation.
Attributes:
role: The role of the message sender
content: Message content (text or structured content blocks)
name: Optional name for the sender
tool_calls: List of tool calls (for assistant messages)
tool_call_id: Tool call ID this message responds to (for tool messages)
"""
role: str
content: Union[str, List[Dict[str, Any]], None] = None
name: Optional[str] = None
tool_calls: Optional[List[ToolCall]] = None
tool_call_id: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary format for API requests."""
result: Dict[str, Any] = {"role": self.role}
if self.content is not None:
result["content"] = self.content
if self.name:
result["name"] = self.name
if self.tool_calls:
result["tool_calls"] = [
{
"id": tc.id,
"type": tc.type,
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in self.tool_calls
]
if self.tool_call_id:
result["tool_call_id"] = self.tool_call_id
return result
@dataclass
class ChatResponseChoice:
"""
A single choice in a chat response.
Attributes:
index: Index of this choice
message: The response message
finish_reason: Why the response ended
"""
index: int
message: ChatMessage
finish_reason: Optional[str] = None
@dataclass
class ChatResponse:
"""
Response from a chat completion request.
Attributes:
id: Unique identifier for this response
choices: List of response choices
usage: Token usage statistics
model: Model that generated this response
created: Unix timestamp of creation
"""
id: str
choices: List[ChatResponseChoice]
usage: Optional[UsageStats] = None
model: Optional[str] = None
created: Optional[int] = None
@property
def message(self) -> Optional[ChatMessage]:
"""Get the first choice's message."""
if self.choices:
return self.choices[0].message
return None
@property
def content(self) -> Optional[str]:
"""Get the text content of the first choice."""
msg = self.message
if msg and isinstance(msg.content, str):
return msg.content
return None
@property
def tool_calls(self) -> Optional[List[ToolCall]]:
"""Get tool calls from the first choice."""
msg = self.message
if msg:
return msg.tool_calls
return None
@dataclass
class StreamChunk:
"""
A single chunk from a streaming response.
Attributes:
id: Response ID
delta_content: New content in this chunk
finish_reason: Finish reason (if this is the last chunk)
usage: Usage stats (usually in the last chunk)
error: Error message if something went wrong
"""
id: str
delta_content: Optional[str] = None
finish_reason: Optional[str] = None
usage: Optional[UsageStats] = None
error: Optional[str] = None
@dataclass
class ModelInfo:
"""
Information about an AI model.
Attributes:
id: Unique model identifier
name: Display name
description: Model description
context_length: Maximum context window size
pricing: Pricing info (input/output per million tokens)
supported_parameters: List of supported API parameters
input_modalities: Supported input types (text, image, etc.)
output_modalities: Supported output types
"""
id: str
name: str
description: str = ""
context_length: int = 0
pricing: Dict[str, float] = field(default_factory=dict)
supported_parameters: List[str] = field(default_factory=list)
input_modalities: List[str] = field(default_factory=lambda: ["text"])
output_modalities: List[str] = field(default_factory=lambda: ["text"])
def supports_images(self) -> bool:
"""Check if model supports image input."""
return "image" in self.input_modalities
def supports_tools(self) -> bool:
"""Check if model supports function calling/tools."""
return "tools" in self.supported_parameters or "functions" in self.supported_parameters
def supports_streaming(self) -> bool:
"""Check if model supports streaming responses."""
return "stream" in self.supported_parameters
def supports_online(self) -> bool:
"""Check if model supports web search (online mode)."""
return self.supports_tools()
@dataclass
class ProviderCapabilities:
"""
Capabilities supported by a provider.
Attributes:
streaming: Provider supports streaming responses
tools: Provider supports function calling
images: Provider supports image inputs
online: Provider supports web search
max_context: Maximum context length across all models
"""
streaming: bool = True
tools: bool = True
images: bool = True
online: bool = False
max_context: int = 128000
class AIProvider(ABC):
"""
Abstract base class for AI model providers.
All provider implementations must inherit from this class
and implement the required abstract methods.
"""
def __init__(self, api_key: str, base_url: Optional[str] = None):
"""
Initialize the provider.
Args:
api_key: API key for authentication
base_url: Optional custom base URL for the API
"""
self.api_key = api_key
self.base_url = base_url
@property
@abstractmethod
def name(self) -> str:
"""Get the provider name."""
pass
@property
@abstractmethod
def capabilities(self) -> ProviderCapabilities:
"""Get provider capabilities."""
pass
@abstractmethod
def list_models(self) -> List[ModelInfo]:
"""
Fetch available models from the provider.
Returns:
List of available models with their info
"""
pass
@abstractmethod
def get_model(self, model_id: str) -> Optional[ModelInfo]:
"""
Get information about a specific model.
Args:
model_id: The model identifier
Returns:
Model information or None if not found
"""
pass
@abstractmethod
def chat(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, Iterator[StreamChunk]]:
"""
Send a chat completion request.
Args:
model: Model ID to use
messages: List of chat messages
stream: Whether to stream the response
max_tokens: Maximum tokens in response
temperature: Sampling temperature
tools: List of tool definitions for function calling
tool_choice: How to handle tool selection ("auto", "none", etc.)
**kwargs: Additional provider-specific parameters
Returns:
ChatResponse for non-streaming, Iterator[StreamChunk] for streaming
"""
pass
@abstractmethod
async def chat_async(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
"""
Send an async chat completion request.
Args:
model: Model ID to use
messages: List of chat messages
stream: Whether to stream the response
max_tokens: Maximum tokens in response
temperature: Sampling temperature
tools: List of tool definitions for function calling
tool_choice: How to handle tool selection
**kwargs: Additional provider-specific parameters
Returns:
ChatResponse for non-streaming, AsyncIterator[StreamChunk] for streaming
"""
pass
@abstractmethod
def get_credits(self) -> Optional[Dict[str, Any]]:
"""
Get account credit/balance information.
Returns:
Dict with credit info or None if not supported
"""
pass
def validate_api_key(self) -> bool:
"""
Validate that the API key is valid.
Returns:
True if API key is valid
"""
try:
self.list_models()
return True
except Exception:
return False
+423
View File
@@ -0,0 +1,423 @@
"""
Ollama provider for local AI model serving.
This provider connects to a local Ollama server for running models
locally without API keys or external dependencies.
"""
import json
import time
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
import requests
from oai.constants import OLLAMA_DEFAULT_URL
from oai.providers.base import (
AIProvider,
ChatMessage,
ChatResponse,
ChatResponseChoice,
ModelInfo,
ProviderCapabilities,
StreamChunk,
UsageStats,
)
from oai.utils.logging import get_logger
logger = get_logger()
class OllamaProvider(AIProvider):
"""
Ollama local model provider.
Connects to a local Ollama server for running models locally.
No API key required.
"""
def __init__(
self,
api_key: str = "",
base_url: Optional[str] = None,
**kwargs: Any,
):
"""
Initialize Ollama provider.
Args:
api_key: Not used (Ollama doesn't require API keys)
base_url: Ollama server URL (default: http://localhost:11434)
**kwargs: Additional arguments (ignored)
"""
super().__init__(api_key or "", base_url)
self.base_url = base_url or OLLAMA_DEFAULT_URL
self._check_server_available()
def _check_server_available(self) -> bool:
"""
Check if Ollama server is accessible.
Returns:
True if server is accessible
"""
try:
response = requests.get(f"{self.base_url}/api/tags", timeout=2)
if response.ok:
logger.info(f"Ollama server accessible at {self.base_url}")
return True
else:
logger.warning(f"Ollama server returned status {response.status_code}")
return False
except requests.RequestException as e:
logger.warning(f"Ollama server not accessible at {self.base_url}: {e}")
return False
@property
def name(self) -> str:
"""Get provider name."""
return "Ollama"
@property
def capabilities(self) -> ProviderCapabilities:
"""Get provider capabilities."""
return ProviderCapabilities(
streaming=True,
tools=False, # Tool support varies by model
images=False, # Image support varies by model
online=True, # Web search via DuckDuckGo/Google
max_context=8192, # Varies by model
)
def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
"""
List models from local Ollama installation.
Args:
filter_text_only: Ignored for Ollama
Returns:
List of available models
"""
try:
response = requests.get(f"{self.base_url}/api/tags", timeout=5)
response.raise_for_status()
data = response.json()
models = []
for model_data in data.get("models", []):
models.append(self._parse_model(model_data))
logger.info(f"Found {len(models)} Ollama models")
return models
except requests.RequestException as e:
logger.error(f"Failed to list Ollama models: {e}")
return []
def _parse_model(self, model_data: Dict[str, Any]) -> ModelInfo:
"""
Parse Ollama model data into ModelInfo.
Args:
model_data: Raw model data from Ollama API
Returns:
ModelInfo object
"""
model_name = model_data.get("name", "unknown")
size_bytes = model_data.get("size", 0)
size_gb = size_bytes / (1024 ** 3) if size_bytes else 0
return ModelInfo(
id=model_name,
name=model_name,
description=f"Size: {size_gb:.1f}GB",
context_length=8192, # Default, varies by model
pricing={}, # Local models are free
supported_parameters=["stream", "temperature", "max_tokens"],
)
def get_model(self, model_id: str) -> Optional[ModelInfo]:
"""
Get information about a specific model.
Args:
model_id: Model identifier
Returns:
ModelInfo or None if not found
"""
models = self.list_models()
for model in models:
if model.id == model_id:
return model
return None
def chat(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, Iterator[StreamChunk]]:
"""
Send chat request to Ollama.
Args:
model: Model name
messages: Chat messages
stream: Whether to stream response
max_tokens: Maximum tokens (Ollama calls this num_predict)
temperature: Sampling temperature
tools: Not supported
tool_choice: Not supported
**kwargs: Additional parameters
Returns:
ChatResponse or Iterator[StreamChunk]
"""
# Convert messages to Ollama format
ollama_messages = []
for msg in messages:
ollama_messages.append({
"role": msg.role,
"content": msg.content or "",
})
# Build request payload
payload: Dict[str, Any] = {
"model": model,
"messages": ollama_messages,
"stream": stream,
}
# Add optional parameters
options = {}
if temperature is not None:
options["temperature"] = temperature
if max_tokens is not None:
options["num_predict"] = max_tokens
if options:
payload["options"] = options
logger.debug(f"Ollama request: model={model}, messages={len(ollama_messages)}")
try:
if stream:
return self._stream_chat(payload)
else:
return self._sync_chat(payload)
except requests.RequestException as e:
logger.error(f"Ollama request failed: {e}")
# Return error response
return ChatResponse(
id="error",
choices=[
ChatResponseChoice(
index=0,
message=ChatMessage(
role="assistant",
content=f"Error: {str(e)}",
),
finish_reason="error",
)
],
)
def _sync_chat(self, payload: Dict[str, Any]) -> ChatResponse:
"""
Send synchronous chat request.
Args:
payload: Request payload
Returns:
ChatResponse
"""
response = requests.post(
f"{self.base_url}/api/chat",
json=payload,
timeout=120,
)
response.raise_for_status()
data = response.json()
# Parse response
message_data = data.get("message", {})
content = message_data.get("content", "")
# Extract token usage if available
usage = None
if "prompt_eval_count" in data or "eval_count" in data:
usage = UsageStats(
prompt_tokens=data.get("prompt_eval_count", 0),
completion_tokens=data.get("eval_count", 0),
total_tokens=data.get("prompt_eval_count", 0) + data.get("eval_count", 0),
total_cost_usd=0.0, # Local models are free
)
return ChatResponse(
id=str(time.time()),
choices=[
ChatResponseChoice(
index=0,
message=ChatMessage(role="assistant", content=content),
finish_reason="stop",
)
],
usage=usage,
model=data.get("model"),
)
def _stream_chat(self, payload: Dict[str, Any]) -> Iterator[StreamChunk]:
"""
Stream chat response from Ollama.
Args:
payload: Request payload
Yields:
StreamChunk objects
"""
response = requests.post(
f"{self.base_url}/api/chat",
json=payload,
stream=True,
timeout=120,
)
response.raise_for_status()
total_prompt_tokens = 0
total_completion_tokens = 0
for line in response.iter_lines():
if not line:
continue
try:
data = json.loads(line)
# Extract content delta
message_data = data.get("message", {})
content = message_data.get("content", "")
# Check if done
done = data.get("done", False)
finish_reason = "stop" if done else None
# Extract usage if available
usage = None
if done and ("prompt_eval_count" in data or "eval_count" in data):
total_prompt_tokens = data.get("prompt_eval_count", 0)
total_completion_tokens = data.get("eval_count", 0)
usage = UsageStats(
prompt_tokens=total_prompt_tokens,
completion_tokens=total_completion_tokens,
total_tokens=total_prompt_tokens + total_completion_tokens,
total_cost_usd=0.0,
)
yield StreamChunk(
id=str(time.time()),
delta_content=content if content else None,
finish_reason=finish_reason,
usage=usage,
)
except json.JSONDecodeError as e:
logger.error(f"Failed to parse Ollama stream chunk: {e}")
yield StreamChunk(
id="error",
error=f"Parse error: {e}",
)
async def chat_async(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
"""
Async chat not implemented for Ollama.
Args:
model: Model name
messages: Chat messages
stream: Whether to stream
max_tokens: Max tokens
temperature: Temperature
tools: Tools (not supported)
tool_choice: Tool choice (not supported)
**kwargs: Additional args
Returns:
ChatResponse or AsyncIterator[StreamChunk]
Raises:
NotImplementedError: Async not implemented
"""
raise NotImplementedError("Async chat not implemented for Ollama provider")
def get_credits(self) -> Optional[Dict[str, Any]]:
"""
Get account credits.
Returns:
None (Ollama is local and free)
"""
return None
def clear_cache(self) -> None:
"""Clear model cache (no-op for Ollama)."""
pass
def get_raw_models(self) -> List[Dict[str, Any]]:
"""
Get raw model data as dictionaries.
Returns:
List of model dictionaries
"""
models = self.list_models()
return [
{
"id": model.id,
"name": model.name,
"description": model.description,
"context_length": model.context_length,
"pricing": model.pricing,
}
for model in models
]
def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
"""
Get raw model data for a specific model.
Args:
model_id: Model identifier
Returns:
Model dictionary or None
"""
model = self.get_model(model_id)
if model:
return {
"id": model.id,
"name": model.name,
"description": model.description,
"context_length": model.context_length,
"pricing": model.pricing,
}
return None
+630
View File
@@ -0,0 +1,630 @@
"""
OpenAI provider for GPT models.
This provider connects to OpenAI's API for accessing GPT-4, GPT-3.5, and other OpenAI models.
"""
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
from openai import OpenAI, AsyncOpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk
from oai.constants import OPENAI_BASE_URL
from oai.providers.base import (
AIProvider,
ChatMessage,
ChatResponse,
ChatResponseChoice,
ModelInfo,
ProviderCapabilities,
StreamChunk,
ToolCall,
ToolFunction,
UsageStats,
)
from oai.utils.logging import get_logger
logger = get_logger()
# Model aliases for convenience
MODEL_ALIASES = {
"gpt-4": "gpt-4-turbo",
"gpt-4-turbo": "gpt-4-turbo-2024-04-09",
"gpt-4o": "gpt-4o-2024-11-20",
"gpt-4o-mini": "gpt-4o-mini-2024-07-18",
"gpt-3.5": "gpt-3.5-turbo",
"gpt-3.5-turbo": "gpt-3.5-turbo-0125",
"o1": "o1-2024-12-17",
"o1-mini": "o1-mini-2024-09-12",
"o1-preview": "o1-preview-2024-09-12",
}
class OpenAIProvider(AIProvider):
"""
OpenAI API provider.
Provides access to GPT-4, GPT-3.5, o1, and other OpenAI models.
"""
def __init__(
self,
api_key: str,
base_url: Optional[str] = None,
app_name: str = "oAI",
app_url: str = "",
**kwargs: Any,
):
"""
Initialize OpenAI provider.
Args:
api_key: OpenAI API key
base_url: Optional custom base URL
app_name: Application name (for headers)
app_url: Application URL (for headers)
**kwargs: Additional arguments
"""
super().__init__(api_key, base_url or OPENAI_BASE_URL)
self.client = OpenAI(api_key=api_key, base_url=self.base_url)
self.async_client = AsyncOpenAI(api_key=api_key, base_url=self.base_url)
self._models_cache: Optional[List[ModelInfo]] = None
@property
def name(self) -> str:
"""Get provider name."""
return "OpenAI"
@property
def capabilities(self) -> ProviderCapabilities:
"""Get provider capabilities."""
return ProviderCapabilities(
streaming=True,
tools=True,
images=True,
online=True, # Web search via DuckDuckGo/Google
max_context=128000,
)
def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
"""
List available OpenAI models.
Args:
filter_text_only: Whether to filter for text models only
Returns:
List of ModelInfo objects
"""
if self._models_cache:
return self._models_cache
try:
models_response = self.client.models.list()
models = []
for model in models_response.data:
# Filter for chat models
if "gpt" in model.id or "o1" in model.id:
models.append(self._parse_model(model))
# Sort by name
models.sort(key=lambda m: m.name)
self._models_cache = models
logger.info(f"Loaded {len(models)} OpenAI models")
return models
except Exception as e:
logger.error(f"Failed to list OpenAI models: {e}")
return self._get_fallback_models()
def _get_fallback_models(self) -> List[ModelInfo]:
"""
Get fallback list of common OpenAI models.
Returns:
List of common models
"""
return [
ModelInfo(
id="gpt-4o",
name="GPT-4o",
description="Most capable GPT-4 model",
context_length=128000,
pricing={"input": 5.0, "output": 15.0},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
ModelInfo(
id="gpt-4o-mini",
name="GPT-4o Mini",
description="Affordable and fast GPT-4 class model",
context_length=128000,
pricing={"input": 0.15, "output": 0.6},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
ModelInfo(
id="gpt-4-turbo",
name="GPT-4 Turbo",
description="GPT-4 Turbo with vision",
context_length=128000,
pricing={"input": 10.0, "output": 30.0},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
input_modalities=["text", "image"],
),
ModelInfo(
id="gpt-3.5-turbo",
name="GPT-3.5 Turbo",
description="Fast and affordable model",
context_length=16384,
pricing={"input": 0.5, "output": 1.5},
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
),
ModelInfo(
id="o1",
name="o1",
description="Advanced reasoning model",
context_length=200000,
pricing={"input": 15.0, "output": 60.0},
supported_parameters=["max_tokens"],
),
ModelInfo(
id="o1-mini",
name="o1-mini",
description="Fast reasoning model",
context_length=128000,
pricing={"input": 3.0, "output": 12.0},
supported_parameters=["max_tokens"],
),
]
def _parse_model(self, model: Any) -> ModelInfo:
"""
Parse OpenAI model into ModelInfo.
Args:
model: OpenAI model object
Returns:
ModelInfo object
"""
model_id = model.id
# Determine context length
context_length = 8192 # Default
if "gpt-4o" in model_id or "gpt-4-turbo" in model_id:
context_length = 128000
elif "gpt-4" in model_id:
context_length = 8192
elif "gpt-3.5-turbo" in model_id:
context_length = 16384
elif "o1" in model_id:
context_length = 128000
# Determine pricing (approximate)
pricing = {}
if "gpt-4o-mini" in model_id:
pricing = {"input": 0.15, "output": 0.6}
elif "gpt-4o" in model_id:
pricing = {"input": 5.0, "output": 15.0}
elif "gpt-4-turbo" in model_id:
pricing = {"input": 10.0, "output": 30.0}
elif "gpt-4" in model_id:
pricing = {"input": 30.0, "output": 60.0}
elif "gpt-3.5" in model_id:
pricing = {"input": 0.5, "output": 1.5}
elif "o1" in model_id and "mini" not in model_id:
pricing = {"input": 15.0, "output": 60.0}
elif "o1-mini" in model_id:
pricing = {"input": 3.0, "output": 12.0}
return ModelInfo(
id=model_id,
name=model_id,
description="",
context_length=context_length,
pricing=pricing,
supported_parameters=["temperature", "max_tokens", "stream", "tools"],
)
def get_model(self, model_id: str) -> Optional[ModelInfo]:
"""
Get information about a specific model.
Args:
model_id: Model identifier
Returns:
ModelInfo or None
"""
# Resolve alias
resolved_id = MODEL_ALIASES.get(model_id, model_id)
models = self.list_models()
for model in models:
if model.id == resolved_id or model.id == model_id:
return model
# Try to fetch directly
try:
model = self.client.models.retrieve(resolved_id)
return self._parse_model(model)
except Exception:
return None
def chat(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, Iterator[StreamChunk]]:
"""
Send chat completion request to OpenAI.
Args:
model: Model ID
messages: Chat messages
stream: Whether to stream response
max_tokens: Maximum tokens
temperature: Sampling temperature
tools: Tool definitions
tool_choice: Tool selection mode
**kwargs: Additional parameters
Returns:
ChatResponse or Iterator[StreamChunk]
"""
# Resolve model alias
model_id = MODEL_ALIASES.get(model, model)
# Convert messages to OpenAI format
openai_messages = []
for msg in messages:
message_dict = {"role": msg.role, "content": msg.content or ""}
if msg.tool_calls:
message_dict["tool_calls"] = [
{
"id": tc.id,
"type": tc.type,
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in msg.tool_calls
]
if msg.tool_call_id:
message_dict["tool_call_id"] = msg.tool_call_id
openai_messages.append(message_dict)
# Build request parameters
params: Dict[str, Any] = {
"model": model_id,
"messages": openai_messages,
"stream": stream,
}
# Add optional parameters
if max_tokens is not None:
params["max_tokens"] = max_tokens
if temperature is not None and "o1" not in model_id:
# o1 models don't support temperature
params["temperature"] = temperature
if tools:
params["tools"] = tools
if tool_choice:
params["tool_choice"] = tool_choice
logger.debug(f"OpenAI request: model={model_id}, messages={len(openai_messages)}")
try:
if stream:
return self._stream_chat(params)
else:
return self._sync_chat(params)
except Exception as e:
logger.error(f"OpenAI request failed: {e}")
return ChatResponse(
id="error",
choices=[
ChatResponseChoice(
index=0,
message=ChatMessage(role="assistant", content=f"Error: {str(e)}"),
finish_reason="error",
)
],
)
def _sync_chat(self, params: Dict[str, Any]) -> ChatResponse:
"""
Send synchronous chat request.
Args:
params: Request parameters
Returns:
ChatResponse
"""
completion: ChatCompletion = self.client.chat.completions.create(**params)
# Convert to our format
choices = []
for choice in completion.choices:
# Convert tool calls if present
tool_calls = None
if choice.message.tool_calls:
tool_calls = [
ToolCall(
id=tc.id,
type=tc.type,
function=ToolFunction(
name=tc.function.name,
arguments=tc.function.arguments,
),
)
for tc in choice.message.tool_calls
]
choices.append(
ChatResponseChoice(
index=choice.index,
message=ChatMessage(
role=choice.message.role,
content=choice.message.content,
tool_calls=tool_calls,
),
finish_reason=choice.finish_reason,
)
)
# Convert usage
usage = None
if completion.usage:
usage = UsageStats(
prompt_tokens=completion.usage.prompt_tokens,
completion_tokens=completion.usage.completion_tokens,
total_tokens=completion.usage.total_tokens,
)
return ChatResponse(
id=completion.id,
choices=choices,
usage=usage,
model=completion.model,
created=completion.created,
)
def _stream_chat(self, params: Dict[str, Any]) -> Iterator[StreamChunk]:
"""
Stream chat response from OpenAI.
Args:
params: Request parameters
Yields:
StreamChunk objects
"""
stream = self.client.chat.completions.create(**params)
for chunk in stream:
chunk_data: ChatCompletionChunk = chunk
if not chunk_data.choices:
continue
choice = chunk_data.choices[0]
delta = choice.delta
# Extract content
content = delta.content if delta.content else None
# Extract finish reason
finish_reason = choice.finish_reason
# Extract usage (usually in last chunk)
usage = None
if hasattr(chunk_data, "usage") and chunk_data.usage:
usage = UsageStats(
prompt_tokens=chunk_data.usage.prompt_tokens,
completion_tokens=chunk_data.usage.completion_tokens,
total_tokens=chunk_data.usage.total_tokens,
)
yield StreamChunk(
id=chunk_data.id,
delta_content=content,
finish_reason=finish_reason,
usage=usage,
)
async def chat_async(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
"""
Send async chat request to OpenAI.
Args:
model: Model ID
messages: Chat messages
stream: Whether to stream
max_tokens: Max tokens
temperature: Temperature
tools: Tool definitions
tool_choice: Tool choice
**kwargs: Additional args
Returns:
ChatResponse or AsyncIterator[StreamChunk]
"""
# Resolve model alias
model_id = MODEL_ALIASES.get(model, model)
# Convert messages
openai_messages = [msg.to_dict() for msg in messages]
# Build params
params: Dict[str, Any] = {
"model": model_id,
"messages": openai_messages,
"stream": stream,
}
if max_tokens:
params["max_tokens"] = max_tokens
if temperature is not None and "o1" not in model_id:
params["temperature"] = temperature
if tools:
params["tools"] = tools
if tool_choice:
params["tool_choice"] = tool_choice
if stream:
return self._stream_chat_async(params)
else:
completion = await self.async_client.chat.completions.create(**params)
# Convert to ChatResponse (similar to _sync_chat)
return self._convert_completion(completion)
async def _stream_chat_async(self, params: Dict[str, Any]) -> AsyncIterator[StreamChunk]:
"""
Stream async chat response.
Args:
params: Request parameters
Yields:
StreamChunk objects
"""
stream = await self.async_client.chat.completions.create(**params)
async for chunk in stream:
if not chunk.choices:
continue
choice = chunk.choices[0]
delta = choice.delta
yield StreamChunk(
id=chunk.id,
delta_content=delta.content,
finish_reason=choice.finish_reason,
)
def _convert_completion(self, completion: ChatCompletion) -> ChatResponse:
"""Helper to convert OpenAI completion to ChatResponse."""
choices = []
for choice in completion.choices:
tool_calls = None
if choice.message.tool_calls:
tool_calls = [
ToolCall(
id=tc.id,
type=tc.type,
function=ToolFunction(
name=tc.function.name,
arguments=tc.function.arguments,
),
)
for tc in choice.message.tool_calls
]
choices.append(
ChatResponseChoice(
index=choice.index,
message=ChatMessage(
role=choice.message.role,
content=choice.message.content,
tool_calls=tool_calls,
),
finish_reason=choice.finish_reason,
)
)
usage = None
if completion.usage:
usage = UsageStats(
prompt_tokens=completion.usage.prompt_tokens,
completion_tokens=completion.usage.completion_tokens,
total_tokens=completion.usage.total_tokens,
)
return ChatResponse(
id=completion.id,
choices=choices,
usage=usage,
model=completion.model,
created=completion.created,
)
def get_credits(self) -> Optional[Dict[str, Any]]:
"""
Get account credits.
Returns:
None (OpenAI doesn't provide credit API)
"""
return None
def clear_cache(self) -> None:
"""Clear model cache."""
self._models_cache = None
def get_raw_models(self) -> List[Dict[str, Any]]:
"""
Get raw model data as dictionaries.
Returns:
List of model dictionaries
"""
models = self.list_models()
return [
{
"id": model.id,
"name": model.name,
"description": model.description,
"context_length": model.context_length,
"pricing": model.pricing,
}
for model in models
]
def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
"""
Get raw model data for a specific model.
Args:
model_id: Model identifier
Returns:
Model dictionary or None
"""
model = self.get_model(model_id)
if model:
return {
"id": model.id,
"name": model.name,
"description": model.description,
"context_length": model.context_length,
"pricing": model.pricing,
}
return None
+630
View File
@@ -0,0 +1,630 @@
"""
OpenRouter provider implementation.
This module implements the AIProvider interface for OpenRouter,
supporting chat completions, streaming, and function calling.
"""
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
import requests
from openrouter import OpenRouter
from oai.constants import APP_NAME, APP_URL, DEFAULT_BASE_URL
from oai.providers.base import (
AIProvider,
ChatMessage,
ChatResponse,
ChatResponseChoice,
ModelInfo,
ProviderCapabilities,
StreamChunk,
ToolCall,
ToolFunction,
UsageStats,
)
from oai.utils.logging import get_logger
class OpenRouterProvider(AIProvider):
"""
OpenRouter API provider implementation.
Provides access to multiple AI models through OpenRouter's unified API,
supporting chat completions, streaming responses, and function calling.
Attributes:
client: The underlying OpenRouter client
_models_cache: Cached list of available models
"""
def __init__(
self,
api_key: str,
base_url: Optional[str] = None,
app_name: str = APP_NAME,
app_url: str = APP_URL,
):
"""
Initialize the OpenRouter provider.
Args:
api_key: OpenRouter API key
base_url: Optional custom base URL
app_name: Application name for API headers
app_url: Application URL for API headers
"""
super().__init__(api_key, base_url or DEFAULT_BASE_URL)
self.app_name = app_name
self.app_url = app_url
self.client = OpenRouter(api_key=api_key)
self._models_cache: Optional[List[ModelInfo]] = None
self._raw_models_cache: Optional[List[Dict[str, Any]]] = None
self.logger = get_logger()
self.logger.info(f"OpenRouter provider initialized with base URL: {self.base_url}")
@property
def name(self) -> str:
"""Get the provider name."""
return "OpenRouter"
@property
def capabilities(self) -> ProviderCapabilities:
"""Get provider capabilities."""
return ProviderCapabilities(
streaming=True,
tools=True,
images=True,
online=True,
max_context=2000000, # Claude models support up to 200k
)
def _get_headers(self) -> Dict[str, str]:
"""Get standard HTTP headers for API requests."""
headers = {
"HTTP-Referer": self.app_url,
"X-Title": self.app_name,
}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
return headers
def _parse_model(self, model_data: Dict[str, Any]) -> ModelInfo:
"""
Parse raw model data into ModelInfo.
Args:
model_data: Raw model data from API
Returns:
Parsed ModelInfo object
"""
architecture = model_data.get("architecture", {})
pricing_data = model_data.get("pricing", {})
# Parse pricing (convert from string to float if needed)
pricing = {}
for key in ["prompt", "completion"]:
value = pricing_data.get(key)
if value is not None:
try:
# Convert from per-token to per-million-tokens
pricing[key] = float(value) * 1_000_000
except (ValueError, TypeError):
pricing[key] = 0.0
return ModelInfo(
id=model_data.get("id", ""),
name=model_data.get("name", model_data.get("id", "")),
description=model_data.get("description", ""),
context_length=model_data.get("context_length", 0),
pricing=pricing,
supported_parameters=model_data.get("supported_parameters", []),
input_modalities=architecture.get("input_modalities", ["text"]),
output_modalities=architecture.get("output_modalities", ["text"]),
)
def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
"""
Fetch available models from OpenRouter.
Args:
filter_text_only: If True, exclude video-only models
Returns:
List of available models
Raises:
Exception: If API request fails
"""
if self._models_cache is not None:
return self._models_cache
try:
response = requests.get(
f"{self.base_url}/models",
headers=self._get_headers(),
timeout=10,
)
response.raise_for_status()
raw_models = response.json().get("data", [])
self._raw_models_cache = raw_models
models = []
for model_data in raw_models:
# Optionally filter out video-only models
if filter_text_only:
modalities = model_data.get("modalities", [])
if modalities and "video" in modalities and "text" not in modalities:
continue
models.append(self._parse_model(model_data))
self._models_cache = models
self.logger.info(f"Fetched {len(models)} models from OpenRouter")
return models
except requests.RequestException as e:
self.logger.error(f"Failed to fetch models: {e}")
raise
def get_raw_models(self) -> List[Dict[str, Any]]:
"""
Get raw model data as returned by the API.
Useful for accessing provider-specific fields not in ModelInfo.
Returns:
List of raw model dictionaries
"""
if self._raw_models_cache is None:
self.list_models()
return self._raw_models_cache or []
def get_model(self, model_id: str) -> Optional[ModelInfo]:
"""
Get information about a specific model.
Args:
model_id: The model identifier
Returns:
Model information or None if not found
"""
models = self.list_models()
for model in models:
if model.id == model_id:
return model
return None
def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
"""
Get raw model data for a specific model.
Args:
model_id: The model identifier
Returns:
Raw model dictionary or None if not found
"""
raw_models = self.get_raw_models()
for model in raw_models:
if model.get("id") == model_id:
return model
return None
def _convert_messages(self, messages: List[ChatMessage]) -> List[Dict[str, Any]]:
"""
Convert ChatMessage objects to API format.
Args:
messages: List of ChatMessage objects
Returns:
List of message dictionaries for the API
"""
return [msg.to_dict() for msg in messages]
def _parse_usage(self, usage_data: Any) -> Optional[UsageStats]:
"""
Parse usage data from API response.
Args:
usage_data: Raw usage data from API
Returns:
Parsed UsageStats or None
"""
if not usage_data:
return None
# Handle both attribute and dict access
prompt_tokens = 0
completion_tokens = 0
total_cost = None
if hasattr(usage_data, "prompt_tokens"):
prompt_tokens = getattr(usage_data, "prompt_tokens", 0) or 0
elif isinstance(usage_data, dict):
prompt_tokens = usage_data.get("prompt_tokens", 0) or 0
if hasattr(usage_data, "completion_tokens"):
completion_tokens = getattr(usage_data, "completion_tokens", 0) or 0
elif isinstance(usage_data, dict):
completion_tokens = usage_data.get("completion_tokens", 0) or 0
# Try alternative naming (input_tokens/output_tokens)
if prompt_tokens == 0:
if hasattr(usage_data, "input_tokens"):
prompt_tokens = getattr(usage_data, "input_tokens", 0) or 0
elif isinstance(usage_data, dict):
prompt_tokens = usage_data.get("input_tokens", 0) or 0
if completion_tokens == 0:
if hasattr(usage_data, "output_tokens"):
completion_tokens = getattr(usage_data, "output_tokens", 0) or 0
elif isinstance(usage_data, dict):
completion_tokens = usage_data.get("output_tokens", 0) or 0
# Get cost if available
# OpenRouter returns cost in different places:
# 1. As 'total_cost_usd' in usage object (rare)
# 2. As 'usage' at root level (common - this is the dollar amount)
total_cost = None
if hasattr(usage_data, "total_cost_usd"):
total_cost = getattr(usage_data, "total_cost_usd", None)
elif hasattr(usage_data, "usage"):
# OpenRouter puts cost as 'usage' field (dollar amount)
total_cost = getattr(usage_data, "usage", None)
elif isinstance(usage_data, dict):
total_cost = usage_data.get("total_cost_usd") or usage_data.get("usage")
return UsageStats(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
total_cost_usd=float(total_cost) if total_cost else None,
)
def _parse_tool_calls(self, tool_calls_data: Any) -> Optional[List[ToolCall]]:
"""
Parse tool calls from API response.
Args:
tool_calls_data: Raw tool calls data
Returns:
List of ToolCall objects or None
"""
if not tool_calls_data:
return None
tool_calls = []
for tc in tool_calls_data:
# Handle both attribute and dict access
if hasattr(tc, "id"):
tc_id = tc.id
tc_type = getattr(tc, "type", "function")
func = tc.function
func_name = func.name
func_args = func.arguments
else:
tc_id = tc.get("id", "")
tc_type = tc.get("type", "function")
func = tc.get("function", {})
func_name = func.get("name", "")
func_args = func.get("arguments", "{}")
tool_calls.append(
ToolCall(
id=tc_id,
type=tc_type,
function=ToolFunction(name=func_name, arguments=func_args),
)
)
return tool_calls if tool_calls else None
def _parse_response(self, response: Any) -> ChatResponse:
"""
Parse API response into ChatResponse.
Args:
response: Raw API response
Returns:
Parsed ChatResponse
"""
choices = []
for choice in response.choices:
msg = choice.message
message = ChatMessage(
role=msg.role if hasattr(msg, "role") else "assistant",
content=msg.content if hasattr(msg, "content") else None,
tool_calls=self._parse_tool_calls(
getattr(msg, "tool_calls", None)
),
)
choices.append(
ChatResponseChoice(
index=choice.index if hasattr(choice, "index") else 0,
message=message,
finish_reason=getattr(choice, "finish_reason", None),
)
)
return ChatResponse(
id=response.id if hasattr(response, "id") else "",
choices=choices,
usage=self._parse_usage(getattr(response, "usage", None)),
model=getattr(response, "model", None),
created=getattr(response, "created", None),
)
def chat(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
transforms: Optional[List[str]] = None,
**kwargs: Any,
) -> Union[ChatResponse, Iterator[StreamChunk]]:
"""
Send a chat completion request to OpenRouter.
Args:
model: Model ID to use
messages: List of chat messages
stream: Whether to stream the response
max_tokens: Maximum tokens in response
temperature: Sampling temperature (0-2)
tools: List of tool definitions for function calling
tool_choice: How to handle tool selection ("auto", "none", etc.)
transforms: List of transforms (e.g., ["middle-out"])
**kwargs: Additional parameters
Returns:
ChatResponse for non-streaming, Iterator[StreamChunk] for streaming
"""
# Build request parameters
params: Dict[str, Any] = {
"model": model,
"messages": self._convert_messages(messages),
"stream": stream,
"http_headers": self._get_headers(),
}
# Request usage stats in streaming responses
if stream:
params["stream_options"] = {"include_usage": True}
if max_tokens is not None:
params["max_tokens"] = max_tokens
if temperature is not None:
params["temperature"] = temperature
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice or "auto"
if transforms:
params["transforms"] = transforms
# Add any additional parameters
params.update(kwargs)
self.logger.debug(f"Sending chat request to model {model}")
try:
response = self.client.chat.send(**params)
if stream:
return self._stream_response(response)
else:
return self._parse_response(response)
except Exception as e:
self.logger.error(f"Chat request failed: {e}")
raise
def _stream_response(self, response: Any) -> Iterator[StreamChunk]:
"""
Process a streaming response.
Args:
response: Streaming response from API
Yields:
StreamChunk objects
"""
last_usage = None
try:
for chunk in response:
# Check for errors
if hasattr(chunk, "error") and chunk.error:
yield StreamChunk(
id=getattr(chunk, "id", ""),
error=chunk.error.message if hasattr(chunk.error, "message") else str(chunk.error),
)
return
# Extract delta content
delta_content = None
finish_reason = None
if hasattr(chunk, "choices") and chunk.choices:
choice = chunk.choices[0]
if hasattr(choice, "delta"):
delta = choice.delta
if hasattr(delta, "content") and delta.content:
delta_content = delta.content
finish_reason = getattr(choice, "finish_reason", None)
# Track usage from last chunk
if hasattr(chunk, "usage") and chunk.usage:
last_usage = self._parse_usage(chunk.usage)
yield StreamChunk(
id=getattr(chunk, "id", ""),
delta_content=delta_content,
finish_reason=finish_reason,
usage=last_usage if finish_reason else None,
)
except Exception as e:
self.logger.error(f"Stream error: {e}")
yield StreamChunk(id="", error=str(e))
async def chat_async(
self,
model: str,
messages: List[ChatMessage],
stream: bool = False,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: Optional[str] = None,
**kwargs: Any,
) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
"""
Send an async chat completion request.
Note: Currently wraps the sync implementation.
TODO: Implement true async support when OpenRouter SDK supports it.
Args:
model: Model ID to use
messages: List of chat messages
stream: Whether to stream the response
max_tokens: Maximum tokens in response
temperature: Sampling temperature
tools: List of tool definitions
tool_choice: Tool selection mode
**kwargs: Additional parameters
Returns:
ChatResponse for non-streaming, AsyncIterator for streaming
"""
# For now, use sync implementation
# TODO: Add true async when SDK supports it
result = self.chat(
model=model,
messages=messages,
stream=stream,
max_tokens=max_tokens,
temperature=temperature,
tools=tools,
tool_choice=tool_choice,
**kwargs,
)
if stream and isinstance(result, Iterator):
# Convert sync iterator to async
async def async_iter() -> AsyncIterator[StreamChunk]:
for chunk in result:
yield chunk
return async_iter()
return result
def get_credits(self) -> Optional[Dict[str, Any]]:
"""
Get OpenRouter account credit information.
Returns:
Dict with credit info:
- total_credits: Total credits purchased
- used_credits: Credits used
- credits_left: Remaining credits
Raises:
Exception: If API request fails
"""
if not self.api_key:
return None
try:
response = requests.get(
f"{self.base_url}/credits",
headers=self._get_headers(),
timeout=10,
)
response.raise_for_status()
data = response.json().get("data", {})
total_credits = float(data.get("total_credits", 0))
total_usage = float(data.get("total_usage", 0))
credits_left = total_credits - total_usage
return {
"total_credits": total_credits,
"used_credits": total_usage,
"credits_left": credits_left,
"total_credits_formatted": f"${total_credits:.2f}",
"used_credits_formatted": f"${total_usage:.2f}",
"credits_left_formatted": f"${credits_left:.2f}",
}
except Exception as e:
self.logger.error(f"Failed to fetch credits: {e}")
return None
def clear_cache(self) -> None:
"""Clear the models cache to force a refresh."""
self._models_cache = None
self._raw_models_cache = None
self.logger.debug("Models cache cleared")
def get_effective_model_id(self, model_id: str, online_enabled: bool) -> str:
"""
Get the effective model ID with online suffix if needed.
Args:
model_id: Base model ID