Commit 72d001fa authored by Your Name's avatar Your Name

Completed refactoring of the providers module

parent d1079827
...@@ -44,13 +44,34 @@ AISBF is a modular proxy server for managing multiple AI provider integrations. ...@@ -44,13 +44,34 @@ AISBF is a modular proxy server for managing multiple AI provider integrations.
## Directory Structure ## Directory Structure
``` ```
geminiproxy/ aisbf/
├── aisbf/ # Main Python module ├── aisbf/ # Main Python module
│ ├── __init__.py # Module initialization with exports │ ├── __init__.py # Module initialization with exports
│ ├── config.py # Configuration management │ ├── config.py # Configuration management
│ ├── models.py # Pydantic models │ ├── models.py # Pydantic models
│ ├── providers.py # Provider handlers │ ├── handlers.py # Request handlers
│ └── handlers.py # Request handlers │ ├── auth/ # Authentication modules
│ │ ├── __init__.py # Auth module exports
│ │ ├── kiro.py # Kiro auth manager
│ │ ├── claude.py # Claude OAuth2 auth (was claude_auth.py)
│ │ └── kilo.py # Kilo OAuth2 auth (was kilo_oauth2.py)
│ └── providers/ # Provider handlers (split into individual modules)
│ ├── __init__.py # Re-exports, PROVIDER_HANDLERS, get_provider_handler()
│ ├── base.py # BaseProviderHandler, AnthropicFormatConverter, AdaptiveRateLimiter
│ ├── google.py # GoogleProviderHandler
│ ├── openai.py # OpenAIProviderHandler
│ ├── anthropic.py # AnthropicProviderHandler
│ ├── claude.py # ClaudeProviderHandler (OAuth2-based)
│ ├── kilo.py # KiloProviderHandler (OAuth2-based)
│ ├── ollama.py # OllamaProviderHandler
│ └── kiro/ # Kiro provider (subpackage)
│ ├── __init__.py
│ ├── handler.py
│ ├── converters.py
│ ├── converters_openai.py
│ ├── models.py
│ ├── parsers.py
│ └── utils.py
├── config/ # Configuration files directory ├── config/ # Configuration files directory
│ ├── providers.json # Default provider configurations │ ├── providers.json # Default provider configurations
│ └── rotations.json # Default rotation configurations │ └── rotations.json # Default rotation configurations
...@@ -63,7 +84,6 @@ geminiproxy/ ...@@ -63,7 +84,6 @@ geminiproxy/
├── start_proxy.sh # Development start script ├── start_proxy.sh # Development start script
├── aisbf.sh # Alternative start script ├── aisbf.sh # Alternative start script
├── requirements.txt # Python dependencies ├── requirements.txt # Python dependencies
├── INSTALL.md # Installation guide
├── PYPI.md # PyPI publishing guide ├── PYPI.md # PyPI publishing guide
├── DOCUMENTATION.md # Complete API documentation ├── DOCUMENTATION.md # Complete API documentation
└── README.md # Project documentation └── README.md # Project documentation
...@@ -95,14 +115,23 @@ Pydantic models for data validation: ...@@ -95,14 +115,23 @@ Pydantic models for data validation:
- `Provider` - Provider information - `Provider` - Provider information
- `ErrorTracking` - Error tracking data - `ErrorTracking` - Error tracking data
### aisbf/providers.py ### aisbf/providers/ (package)
Provider handler implementations: Provider handler implementations, split into individual modules:
- `BaseProviderHandler` - Base class with rate limiting and error tracking - `base.py` - `BaseProviderHandler`, `AnthropicFormatConverter`, `AdaptiveRateLimiter`, shared utilities
- `GoogleProviderHandler` - Google GenAI integration - `google.py` - `GoogleProviderHandler` - Google GenAI integration
- `OpenAIProviderHandler` - OpenAI API integration - `openai.py` - `OpenAIProviderHandler` - OpenAI API integration
- `AnthropicProviderHandler` - Anthropic API integration - `anthropic.py` - `AnthropicProviderHandler` - Anthropic API key integration
- `OllamaProviderHandler` - Ollama local integration - `claude.py` - `ClaudeProviderHandler` - Claude OAuth2 integration
- `get_provider_handler()` - Factory function to get appropriate handler - `kilo.py` - `KiloProviderHandler` - Kilo Gateway OAuth2 integration
- `ollama.py` - `OllamaProviderHandler` - Ollama local integration
- `kiro/` - `KiroProviderHandler` - Kiro/Amazon Q Developer integration (subpackage)
- `__init__.py` - Re-exports all handlers, `PROVIDER_HANDLERS` dict, `get_provider_handler()` factory
### aisbf/auth/ (package)
Authentication modules:
- `kiro.py` - `KiroAuthManager` for Kiro/Amazon Q Developer auth
- `claude.py` - `ClaudeAuth` for Claude OAuth2 PKCE flow (was `aisbf/claude_auth.py`)
- `kilo.py` - `KiloOAuth2` for Kilo Device Authorization Grant (was `aisbf/kilo_oauth2.py`)
### aisbf/handlers.py ### aisbf/handlers.py
Request handling logic: Request handling logic:
...@@ -327,8 +356,8 @@ When making changes: ...@@ -327,8 +356,8 @@ When making changes:
## Common Tasks ## Common Tasks
### Adding a New Provider ### Adding a New Provider
1. Create handler class in `aisbf/providers.py` inheriting from `BaseProviderHandler` 1. Create a new file `aisbf/providers/<provider_name>.py` with a handler class inheriting from `BaseProviderHandler`
2. Add to `PROVIDER_HANDLERS` dictionary 2. Import and add the handler to `PROVIDER_HANDLERS` in `aisbf/providers/__init__.py`
3. Add provider configuration to `config/providers.json` 3. Add provider configuration to `config/providers.json`
### Configuration Architecture ### Configuration Architecture
...@@ -538,8 +567,8 @@ Claude Code is an OAuth2-based authentication method for accessing Claude models ...@@ -538,8 +567,8 @@ Claude Code is an OAuth2-based authentication method for accessing Claude models
- Supports all Claude features: streaming, tools, vision, extended thinking - Supports all Claude features: streaming, tools, vision, extended thinking
**Integration Architecture:** **Integration Architecture:**
- [`ClaudeAuth`](aisbf/claude_auth.py) class handles OAuth2 PKCE flow - [`ClaudeAuth`](aisbf/auth/claude.py) class handles OAuth2 PKCE flow
- [`ClaudeProviderHandler`](aisbf/providers.py) in [`aisbf/providers.py`](aisbf/providers.py) manages API requests - [`ClaudeProviderHandler`](aisbf/providers/claude.py) in [`aisbf/providers/claude.py`](aisbf/providers/claude.py) manages API requests
- Supports all standard AISBF features: streaming, tools, rate limiting, error tracking - Supports all standard AISBF features: streaming, tools, rate limiting, error tracking
**Configuration:** **Configuration:**
...@@ -702,6 +731,18 @@ This AI.PROMPT file is automatically updated when significant changes are made t ...@@ -702,6 +731,18 @@ This AI.PROMPT file is automatically updated when significant changes are made t
### Recent Updates ### Recent Updates
**2026-04-03 - Provider Module Refactoring (Phase 2)**
- Split monolithic `aisbf/providers/__init__.py` (301K chars) into individual module files
- Created `aisbf/providers/base.py` with shared utilities: `BaseProviderHandler`, `AnthropicFormatConverter`, `AdaptiveRateLimiter`, `AISBF_DEBUG`
- Created individual handler files: `google.py`, `openai.py`, `anthropic.py`, `claude.py`, `kilo.py`, `ollama.py`
- Moved auth modules: `aisbf/claude_auth.py` → `aisbf/auth/claude.py`, `aisbf/kilo_oauth2.py` → `aisbf/auth/kilo.py`
- Updated `aisbf/auth/__init__.py` to export `ClaudeAuth` and `KiloOAuth2`
- Rewrote `aisbf/providers/__init__.py` as slim re-export module with `PROVIDER_HANDLERS` dict and `get_provider_handler()` factory
- Updated `aisbf/__init__.py` to import new handler classes (`ClaudeProviderHandler`, `KiloProviderHandler`)
- Updated `pyproject.toml` and `setup.py` with new file paths
- All existing `from aisbf.providers import X` imports continue to work (backward compatible)
- Deleted old `aisbf/claude_auth.py` and `aisbf/kilo_oauth2.py` (now in `aisbf/auth/`)
**2026-03-23 - TOR Hidden Service Support** **2026-03-23 - TOR Hidden Service Support**
- Added full TOR hidden service support for exposing AISBF over TOR network - Added full TOR hidden service support for exposing AISBF over TOR network
- Created aisbf/tor.py module with TorHiddenService class for managing TOR connections - Created aisbf/tor.py module with TorHiddenService class for managing TOR connections
......
...@@ -39,12 +39,16 @@ from .providers import ( ...@@ -39,12 +39,16 @@ from .providers import (
GoogleProviderHandler, GoogleProviderHandler,
OpenAIProviderHandler, OpenAIProviderHandler,
AnthropicProviderHandler, AnthropicProviderHandler,
ClaudeProviderHandler,
KiloProviderHandler,
OllamaProviderHandler, OllamaProviderHandler,
get_provider_handler, get_provider_handler,
PROVIDER_HANDLERS PROVIDER_HANDLERS
) )
from .providers.kiro import KiroProviderHandler from .providers.kiro import KiroProviderHandler
from .auth.kiro import KiroAuthManager from .auth.kiro import KiroAuthManager
from .auth.claude import ClaudeAuth
from .auth.kilo import KiloOAuth2
from .handlers import RequestHandler, RotationHandler, AutoselectHandler from .handlers import RequestHandler, RotationHandler, AutoselectHandler
from .utils import count_messages_tokens, split_messages_into_chunks, get_max_request_tokens_for_model from .utils import count_messages_tokens, split_messages_into_chunks, get_max_request_tokens_for_model
...@@ -73,11 +77,15 @@ __all__ = [ ...@@ -73,11 +77,15 @@ __all__ = [
"OpenAIProviderHandler", "OpenAIProviderHandler",
"AnthropicProviderHandler", "AnthropicProviderHandler",
"OllamaProviderHandler", "OllamaProviderHandler",
"ClaudeProviderHandler",
"KiloProviderHandler",
"KiroProviderHandler", "KiroProviderHandler",
"get_provider_handler", "get_provider_handler",
"PROVIDER_HANDLERS", "PROVIDER_HANDLERS",
# Auth # Auth
"KiroAuthManager", "KiroAuthManager",
"ClaudeAuth",
"KiloOAuth2",
# Handlers # Handlers
"RequestHandler", "RequestHandler",
"RotationHandler", "RotationHandler",
......
...@@ -22,8 +22,12 @@ Why did the programmer quit his job? Because he didn't get arrays! ...@@ -22,8 +22,12 @@ Why did the programmer quit his job? Because he didn't get arrays!
""" """
from .kiro import KiroAuthManager, AuthType from .kiro import KiroAuthManager, AuthType
from .claude import ClaudeAuth
from .kilo import KiloOAuth2
__all__ = [ __all__ = [
"KiroAuthManager", "KiroAuthManager",
"AuthType", "AuthType",
"ClaudeAuth",
"KiloOAuth2",
] ]
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
"""
Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
AISBF - AI Service Broker Framework || AI Should Be Free
Ollama provider handler.
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
Why did the programmer quit his job? Because he didn't get arrays!
"""
import httpx
import time
from typing import Dict, List, Optional, Union
from ..models import Model
from ..config import config
from .base import BaseProviderHandler, AISBF_DEBUG
class OllamaProviderHandler(BaseProviderHandler):
def __init__(self, provider_id: str, api_key: Optional[str] = None):
super().__init__(provider_id, api_key)
timeout = httpx.Timeout(
connect=60.0,
read=300.0,
write=60.0,
pool=60.0
)
self.client = httpx.AsyncClient(base_url=config.providers[provider_id].endpoint, timeout=timeout)
async def handle_request(self, model: str, messages: List[Dict], max_tokens: Optional[int] = None,
temperature: Optional[float] = 1.0, stream: Optional[bool] = False,
tools: Optional[List[Dict]] = None, tool_choice: Optional[Union[str, Dict]] = None) -> Dict:
"""
Handle request for Ollama provider.
Note: Ollama doesn't support tools/tool_choice, so these parameters are accepted but ignored.
"""
import logging
import json
logger = logging.getLogger(__name__)
logger.info(f"=== OllamaProviderHandler.handle_request START ===")
logger.info(f"Provider ID: {self.provider_id}")
logger.info(f"Endpoint: {self.client.base_url}")
logger.info(f"Model: {model}")
logger.info(f"Messages count: {len(messages)}")
logger.info(f"Max tokens: {max_tokens}")
logger.info(f"Temperature: {temperature}")
logger.info(f"Stream: {stream}")
logger.info(f"API key provided: {bool(self.api_key)}")
if self.is_rate_limited():
logger.error("Provider is rate limited")
raise Exception("Provider rate limited")
try:
logger.info("Testing Ollama connection...")
try:
health_response = await self.client.get("/api/tags", timeout=10.0)
logger.info(f"Ollama health check passed: {health_response.status_code}")
logger.info(f"Available models: {health_response.json().get('models', [])}")
except Exception as e:
logger.error(f"Ollama health check failed: {str(e)}")
logger.error(f"Cannot connect to Ollama at {self.client.base_url}")
logger.error(f"Please ensure Ollama is running and accessible")
raise Exception(f"Cannot connect to Ollama at {self.client.base_url}: {str(e)}")
logger.info("Applying rate limiting...")
await self.apply_rate_limit()
logger.info("Rate limiting applied")
prompt = "\n\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
logger.info(f"Prompt length: {len(prompt)} characters")
options = {"temperature": temperature}
if max_tokens is not None:
options["num_predict"] = max_tokens
request_data = {
"model": model,
"prompt": prompt,
"options": options,
"stream": False
}
headers = {}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
logger.info("API key added to request headers for Ollama cloud")
logger.info(f"Sending POST request to {self.client.base_url}/api/generate")
logger.info(f"Request data: {request_data}")
logger.info(f"Request headers: {headers}")
logger.info(f"Client timeout: {self.client.timeout}")
response = await self.client.post("/api/generate", json=request_data, headers=headers)
logger.info(f"Response status code: {response.status_code}")
logger.info(f"Response content type: {response.headers.get('content-type')}")
logger.info(f"Response content length: {len(response.content)} bytes")
logger.info(f"Raw response content (first 500 chars): {response.text[:500]}")
if response.status_code == 429:
try:
response_data = response.json()
except Exception:
response_data = response.text
self.handle_429_error(response_data, dict(response.headers))
response.raise_for_status()
response.raise_for_status()
content = response.text
logger.info(f"Attempting to parse response as JSON...")
try:
response_json = response.json()
logger.info(f"Response parsed as single JSON: {response_json}")
except json.JSONDecodeError as e:
logger.warning(f"Failed to parse as single JSON: {e}")
logger.info(f"Attempting to parse as multiple JSON objects...")
responses = []
for line in content.strip().split('\n'):
if line.strip():
try:
obj = json.loads(line)
responses.append(obj)
except json.JSONDecodeError as line_error:
logger.error(f"Failed to parse line: {line}")
logger.error(f"Error: {line_error}")
if not responses:
raise Exception("No valid JSON objects found in response")
response_json = responses[-1]
logger.info(f"Parsed {len(responses)} JSON objects, using last one: {response_json}")
logger.info(f"Final response: {response_json}")
self.record_success()
if AISBF_DEBUG:
logging.info(f"=== RAW OLLAMA RESPONSE ===")
logging.info(f"Raw response JSON: {response_json}")
logging.info(f"=== END RAW OLLAMA RESPONSE ===")
logger.info(f"=== OllamaProviderHandler.handle_request END ===")
openai_response = {
"id": f"ollama-{model}-{int(time.time())}",
"object": "chat.completion",
"created": int(time.time()),
"model": f"{self.provider_id}/{model}",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": response_json.get("response", "")
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": response_json.get("prompt_eval_count", 0),
"completion_tokens": response_json.get("eval_count", 0),
"total_tokens": response_json.get("prompt_eval_count", 0) + response_json.get("eval_count", 0)
}
}
if AISBF_DEBUG:
logging.info(f"=== FINAL OLLAMA RESPONSE DICT ===")
logging.info(f"Final response: {openai_response}")
logging.info(f"=== END FINAL OLLAMA RESPONSE DICT ===")
return openai_response
except Exception as e:
self.record_failure()
raise e
async def get_models(self) -> List[Model]:
await self.apply_rate_limit()
response = await self.client.get("/api/tags")
response.raise_for_status()
models = response.json().get('models', [])
return [Model(id=model, name=model, provider_id=self.provider_id) for model in models]
This diff is collapsed.
...@@ -49,6 +49,7 @@ Documentation = "https://git.nexlab.net/nexlab/aisbf.git" ...@@ -49,6 +49,7 @@ Documentation = "https://git.nexlab.net/nexlab/aisbf.git"
[tool.setuptools] [tool.setuptools]
packages = ["aisbf", "aisbf.auth", "aisbf.providers", "aisbf.providers.kiro"] packages = ["aisbf", "aisbf.auth", "aisbf.providers", "aisbf.providers.kiro"]
# Note: Provider handler modules (base, google, openai, anthropic, claude, kilo, ollama) are in aisbf.providers package
py-modules = ["cli"] py-modules = ["cli"]
[tool.setuptools.package-data] [tool.setuptools.package-data]
......
...@@ -99,6 +99,13 @@ setup( ...@@ -99,6 +99,13 @@ setup(
'aisbf/config.py', 'aisbf/config.py',
'aisbf/models.py', 'aisbf/models.py',
'aisbf/providers/__init__.py', 'aisbf/providers/__init__.py',
'aisbf/providers/base.py',
'aisbf/providers/google.py',
'aisbf/providers/openai.py',
'aisbf/providers/anthropic.py',
'aisbf/providers/claude.py',
'aisbf/providers/kilo.py',
'aisbf/providers/ollama.py',
'aisbf/handlers.py', 'aisbf/handlers.py',
'aisbf/context.py', 'aisbf/context.py',
'aisbf/utils.py', 'aisbf/utils.py',
...@@ -107,6 +114,8 @@ setup( ...@@ -107,6 +114,8 @@ setup(
'aisbf/tor.py', 'aisbf/tor.py',
'aisbf/auth/__init__.py', 'aisbf/auth/__init__.py',
'aisbf/auth/kiro.py', 'aisbf/auth/kiro.py',
'aisbf/auth/claude.py',
'aisbf/auth/kilo.py',
'aisbf/providers/kiro/__init__.py', 'aisbf/providers/kiro/__init__.py',
'aisbf/providers/kiro/handler.py', 'aisbf/providers/kiro/handler.py',
'aisbf/providers/kiro/converters.py', 'aisbf/providers/kiro/converters.py',
...@@ -114,7 +123,6 @@ setup( ...@@ -114,7 +123,6 @@ setup(
'aisbf/providers/kiro/models.py', 'aisbf/providers/kiro/models.py',
'aisbf/providers/kiro/parsers.py', 'aisbf/providers/kiro/parsers.py',
'aisbf/providers/kiro/utils.py', 'aisbf/providers/kiro/utils.py',
'aisbf/claude_auth.py',
'aisbf/semantic_classifier.py', 'aisbf/semantic_classifier.py',
'aisbf/batching.py', 'aisbf/batching.py',
'aisbf/cache.py', 'aisbf/cache.py',
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment