| | |
| |
|
| | from typing import Dict, Any, AsyncGenerator, Optional |
| | import httpx |
| | import random |
| | from abc import ABC, abstractmethod |
| | from app.config.config import settings |
| | from app.log.logger import get_api_client_logger |
| | from app.core.constants import DEFAULT_TIMEOUT |
| |
|
| | logger = get_api_client_logger() |
| |
|
| | class ApiClient(ABC): |
| | """API客户端基类""" |
| |
|
| | @abstractmethod |
| | async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]: |
| | pass |
| |
|
| | @abstractmethod |
| | async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]: |
| | pass |
| |
|
| |
|
| | class GeminiApiClient(ApiClient): |
| | """Gemini API客户端""" |
| |
|
| | def __init__(self, base_url: str, timeout: int = DEFAULT_TIMEOUT): |
| | self.base_url = base_url |
| | self.timeout = timeout |
| |
|
| | def _get_real_model(self, model: str) -> str: |
| | if model.endswith("-search"): |
| | model = model[:-7] |
| | if model.endswith("-image"): |
| | model = model[:-6] |
| | if model.endswith("-non-thinking"): |
| | model = model[:-13] |
| | if "-search" in model and "-non-thinking" in model: |
| | model = model[:-20] |
| | return model |
| |
|
| | def _prepare_headers(self) -> Dict[str, str]: |
| | headers = {} |
| | if settings.CUSTOM_HEADERS: |
| | headers.update(settings.CUSTOM_HEADERS) |
| | logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}") |
| | return headers |
| |
|
| | async def get_models(self, api_key: str) -> Optional[Dict[str, Any]]: |
| | """获取可用的 Gemini 模型列表""" |
| | timeout = httpx.Timeout(timeout=5) |
| | |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers() |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/models?key={api_key}&pageSize=1000" |
| | try: |
| | response = await client.get(url, headers=headers) |
| | response.raise_for_status() |
| | return response.json() |
| | except httpx.HTTPStatusError as e: |
| | logger.error(f"获取模型列表失败: {e.response.status_code}") |
| | logger.error(e.response.text) |
| | return None |
| | except httpx.RequestError as e: |
| | logger.error(f"请求模型列表失败: {e}") |
| | return None |
| | |
| | async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | model = self._get_real_model(model) |
| | |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| | |
| | headers = self._prepare_headers() |
| | |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/models/{model}:generateContent?key={api_key}" |
| | |
| | try: |
| | response = await client.post(url, json=payload, headers=headers) |
| | |
| | if response.status_code != 200: |
| | error_content = response.text |
| | logger.error(f"API call failed - Status: {response.status_code}, Content: {error_content}") |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | |
| | response_data = response.json() |
| | |
| | |
| | if not response_data.get("candidates"): |
| | logger.warning("No candidates found in API response") |
| | |
| | return response_data |
| | |
| | except httpx.TimeoutException as e: |
| | logger.error(f"Request timeout: {e}") |
| | raise Exception(f"Request timeout: {e}") |
| | except httpx.RequestError as e: |
| | logger.error(f"Request error: {e}") |
| | raise Exception(f"Request error: {e}") |
| | except Exception as e: |
| | logger.error(f"Unexpected error: {e}") |
| | raise |
| |
|
| | async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | model = self._get_real_model(model) |
| | |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers() |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse&key={api_key}" |
| | async with client.stream(method="POST", url=url, json=payload, headers=headers) as response: |
| | if response.status_code != 200: |
| | error_content = await response.aread() |
| | error_msg = error_content.decode("utf-8") |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_msg}") |
| | async for line in response.aiter_lines(): |
| | yield line |
| |
|
| | async def count_tokens(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | model = self._get_real_model(model) |
| |
|
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for counting tokens: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers() |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/models/{model}:countTokens?key={api_key}" |
| | response = await client.post(url, json=payload, headers=headers) |
| | if response.status_code != 200: |
| | error_content = response.text |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | return response.json() |
| |
|
| | async def embed_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]: |
| | """单一嵌入内容生成""" |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | model = self._get_real_model(model) |
| | |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for embedding: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers() |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/models/{model}:embedContent?key={api_key}" |
| | |
| | try: |
| | response = await client.post(url, json=payload, headers=headers) |
| | |
| | if response.status_code != 200: |
| | error_content = response.text |
| | logger.error(f"Embedding API call failed - Status: {response.status_code}, Content: {error_content}") |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | |
| | return response.json() |
| | |
| | except httpx.TimeoutException as e: |
| | logger.error(f"Embedding request timeout: {e}") |
| | raise Exception(f"Request timeout: {e}") |
| | except httpx.RequestError as e: |
| | logger.error(f"Embedding request error: {e}") |
| | raise Exception(f"Request error: {e}") |
| | except Exception as e: |
| | logger.error(f"Unexpected embedding error: {e}") |
| | raise |
| |
|
| | async def batch_embed_contents(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]: |
| | """批量嵌入内容生成""" |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | model = self._get_real_model(model) |
| | |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for batch embedding: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers() |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/models/{model}:batchEmbedContents?key={api_key}" |
| | |
| | try: |
| | response = await client.post(url, json=payload, headers=headers) |
| | |
| | if response.status_code != 200: |
| | error_content = response.text |
| | logger.error(f"Batch embedding API call failed - Status: {response.status_code}, Content: {error_content}") |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | |
| | return response.json() |
| | |
| | except httpx.TimeoutException as e: |
| | logger.error(f"Batch embedding request timeout: {e}") |
| | raise Exception(f"Request timeout: {e}") |
| | except httpx.RequestError as e: |
| | logger.error(f"Batch embedding request error: {e}") |
| | raise Exception(f"Request error: {e}") |
| | except Exception as e: |
| | logger.error(f"Unexpected batch embedding error: {e}") |
| | raise |
| |
|
| |
|
| | class OpenaiApiClient(ApiClient): |
| | """OpenAI API客户端""" |
| |
|
| | def __init__(self, base_url: str, timeout: int = DEFAULT_TIMEOUT): |
| | self.base_url = base_url |
| | self.timeout = timeout |
| | |
| | def _prepare_headers(self, api_key: str) -> Dict[str, str]: |
| | headers = {"Authorization": f"Bearer {api_key}"} |
| | if settings.CUSTOM_HEADERS: |
| | headers.update(settings.CUSTOM_HEADERS) |
| | logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}") |
| | return headers |
| |
|
| | async def get_models(self, api_key: str) -> Dict[str, Any]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| |
|
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers(api_key) |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/openai/models" |
| | response = await client.get(url, headers=headers) |
| | if response.status_code != 200: |
| | error_content = response.text |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | return response.json() |
| |
|
| | async def generate_content(self, payload: Dict[str, Any], api_key: str) -> Dict[str, Any]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | logger.info(f"settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: {settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY}") |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers(api_key) |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/openai/chat/completions" |
| | response = await client.post(url, json=payload, headers=headers) |
| | if response.status_code != 200: |
| | error_content = response.text |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | return response.json() |
| |
|
| | async def stream_generate_content(self, payload: Dict[str, Any], api_key: str) -> AsyncGenerator[str, None]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers(api_key) |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/openai/chat/completions" |
| | async with client.stream(method="POST", url=url, json=payload, headers=headers) as response: |
| | if response.status_code != 200: |
| | error_content = await response.aread() |
| | error_msg = error_content.decode("utf-8") |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_msg}") |
| | async for line in response.aiter_lines(): |
| | yield line |
| | |
| | async def create_embeddings(self, input: str, model: str, api_key: str) -> Dict[str, Any]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| | |
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers(api_key) |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/openai/embeddings" |
| | payload = { |
| | "input": input, |
| | "model": model, |
| | } |
| | response = await client.post(url, json=payload, headers=headers) |
| | if response.status_code != 200: |
| | error_content = response.text |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | return response.json() |
| | |
| | async def generate_images(self, payload: Dict[str, Any], api_key: str) -> Dict[str, Any]: |
| | timeout = httpx.Timeout(self.timeout, read=self.timeout) |
| |
|
| | proxy_to_use = None |
| | if settings.PROXIES: |
| | if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: |
| | proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)] |
| | else: |
| | proxy_to_use = random.choice(settings.PROXIES) |
| | logger.info(f"Using proxy for getting models: {proxy_to_use}") |
| |
|
| | headers = self._prepare_headers(api_key) |
| | async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: |
| | url = f"{self.base_url}/openai/images/generations" |
| | response = await client.post(url, json=payload, headers=headers) |
| | if response.status_code != 200: |
| | error_content = response.text |
| | raise Exception(f"API call failed with status code {response.status_code}, {error_content}") |
| | return response.json() |
