mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-08-11 19:57:32 +00:00
openai lib
This commit is contained in:
9qeklajc
@@ -20,7 +20,6 @@ dependencies = [
|
||||
"nostr>=0.0.2",
|
||||
"mdurl==0.1.2",
|
||||
"pillow>=10",
|
||||
"google-generativeai>=0.8.5",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
|
||||
@@ -1,39 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any, AsyncGenerator, Awaitable, Callable
|
||||
|
||||
import google.generativeai as genai
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from .base import BaseAPIClient
|
||||
|
||||
|
||||
class GeminiClient(BaseAPIClient):
|
||||
"""Native Gemini API client using Google's official package."""
|
||||
"""Gemini API client using OpenAI compatibility layer."""
|
||||
|
||||
def __init__(self, api_key: str, base_url: str | None = None):
|
||||
super().__init__(api_key, base_url)
|
||||
genai.configure(api_key=api_key) # type: ignore
|
||||
|
||||
def _validate_model_name(self, model: str) -> str:
|
||||
return model
|
||||
|
||||
def _convert_openai_to_gemini_messages(self, messages: list[dict[str, Any]]) -> str:
|
||||
"""Convert OpenAI messages to a simple string for Gemini."""
|
||||
combined_content = []
|
||||
|
||||
for message in messages:
|
||||
role = message.get("role", "user")
|
||||
content = message.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
combined_content.append(f"System: {content}")
|
||||
elif role == "user":
|
||||
combined_content.append(f"User: {content}")
|
||||
elif role == "assistant":
|
||||
combined_content.append(f"Assistant: {content}")
|
||||
|
||||
return "\n".join(combined_content)
|
||||
self.client = AsyncOpenAI(
|
||||
api_key=api_key,
|
||||
base_url=base_url
|
||||
or "https://generativelanguage.googleapis.com/v1beta/openai/",
|
||||
)
|
||||
|
||||
async def generate_content(
|
||||
self,
|
||||
@@ -43,53 +26,22 @@ class GeminiClient(BaseAPIClient):
|
||||
max_tokens: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Generate content using Gemini API (non-streaming)."""
|
||||
model = self._validate_model_name(model)
|
||||
prompt = self._convert_openai_to_gemini_messages(messages)
|
||||
|
||||
generation_config = {}
|
||||
if temperature is not None:
|
||||
generation_config["temperature"] = temperature
|
||||
if max_tokens is not None:
|
||||
generation_config["max_output_tokens"] = max_tokens
|
||||
if "top_p" in kwargs:
|
||||
generation_config["top_p"] = kwargs["top_p"]
|
||||
|
||||
model_instance = genai.GenerativeModel(model) # type: ignore
|
||||
response = await model_instance.generate_content_async( # type: ignore
|
||||
prompt,
|
||||
generation_config=generation_config if generation_config else None, # type: ignore
|
||||
)
|
||||
|
||||
return {
|
||||
"id": f"chatcmpl-{hash(str(response))}"[1:16],
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"""Generate content using Gemini API via OpenAI SDK (non-streaming)."""
|
||||
args = {
|
||||
"model": model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": response.text,
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": getattr(
|
||||
response.usage_metadata, "prompt_token_count", 0
|
||||
),
|
||||
"completion_tokens": getattr(
|
||||
response.usage_metadata, "candidates_token_count", 0
|
||||
),
|
||||
"total_tokens": getattr(
|
||||
response.usage_metadata, "total_token_count", 0
|
||||
),
|
||||
},
|
||||
"messages": messages,
|
||||
}
|
||||
if temperature is not None:
|
||||
args["temperature"] = temperature
|
||||
if max_tokens is not None:
|
||||
args["max_tokens"] = max_tokens
|
||||
if "top_p" in kwargs:
|
||||
args["top_p"] = kwargs["top_p"]
|
||||
|
||||
async def generate_content_stream( # type: ignore[override]
|
||||
response = await self.client.chat.completions.create(**args)
|
||||
return response.model_dump()
|
||||
|
||||
async def generate_content_stream(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
@@ -100,107 +52,42 @@ class GeminiClient(BaseAPIClient):
|
||||
| None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncGenerator[dict[str, Any], None]:
|
||||
"""Generate content using Gemini API (streaming)."""
|
||||
model = self._validate_model_name(model)
|
||||
prompt = self._convert_openai_to_gemini_messages(messages)
|
||||
|
||||
stream_id = f"chatcmpl-{abs(hash(prompt + str(time.time())))}"[:28]
|
||||
created_time = int(time.time())
|
||||
generation_config = {}
|
||||
"""Generate content using Gemini API via OpenAI SDK (streaming)."""
|
||||
args = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
if temperature is not None:
|
||||
generation_config["temperature"] = temperature
|
||||
args["temperature"] = temperature
|
||||
if max_tokens is not None:
|
||||
generation_config["max_output_tokens"] = max_tokens
|
||||
args["max_tokens"] = max_tokens
|
||||
if "top_p" in kwargs:
|
||||
generation_config["top_p"] = kwargs["top_p"]
|
||||
args["top_p"] = kwargs["top_p"]
|
||||
|
||||
model_instance = genai.GenerativeModel(model) # type: ignore
|
||||
response_stream = await model_instance.generate_content_async( # type: ignore
|
||||
prompt,
|
||||
generation_config=generation_config if generation_config else None, # type: ignore
|
||||
stream=True,
|
||||
)
|
||||
stream = await self.client.chat.completions.create(**args)
|
||||
|
||||
async for chunk in response_stream:
|
||||
finish_reason = None
|
||||
content = ""
|
||||
final_usage = None
|
||||
|
||||
if hasattr(chunk, "text") and chunk.text:
|
||||
content = chunk.text
|
||||
if hasattr(chunk, "candidates") and chunk.candidates:
|
||||
candidate = chunk.candidates[0]
|
||||
finish_reason_raw = getattr(candidate, "finish_reason", None)
|
||||
async for chunk in stream:
|
||||
chunk_data = chunk.model_dump()
|
||||
|
||||
if finish_reason_raw == 1:
|
||||
finish_reason = "stop"
|
||||
elif finish_reason_raw == 2:
|
||||
finish_reason = "length"
|
||||
elif finish_reason_raw == 3:
|
||||
finish_reason = "content_filter"
|
||||
elif finish_reason_raw == 4:
|
||||
finish_reason = "content_filter"
|
||||
usage_data = None
|
||||
if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
|
||||
usage_metadata = chunk.usage_metadata
|
||||
usage_data = {
|
||||
"prompt_tokens": getattr(usage_metadata, "prompt_token_count", 0),
|
||||
"completion_tokens": getattr(
|
||||
usage_metadata, "candidates_token_count", 0
|
||||
),
|
||||
"total_tokens": getattr(usage_metadata, "total_token_count", 0),
|
||||
}
|
||||
|
||||
if hasattr(usage_metadata, "cached_content_token_count"):
|
||||
cached_tokens = getattr(
|
||||
usage_metadata, "cached_content_token_count", 0
|
||||
)
|
||||
if cached_tokens > 0:
|
||||
usage_data["prompt_tokens_details"] = {
|
||||
"cached_tokens": cached_tokens
|
||||
}
|
||||
chunk_data = {
|
||||
"id": stream_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created_time,
|
||||
"model": model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": content} if content else {},
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
if usage_data:
|
||||
chunk_data["usage"] = usage_data
|
||||
if chunk.usage:
|
||||
final_usage = chunk.usage.model_dump()
|
||||
if usage_callback:
|
||||
usage_callback(usage_data)
|
||||
usage_callback(final_usage)
|
||||
|
||||
yield chunk_data
|
||||
|
||||
if finish_reason == "stop":
|
||||
if completion_callback:
|
||||
await completion_callback(model, usage_data)
|
||||
break
|
||||
if completion_callback:
|
||||
await completion_callback(model, final_usage)
|
||||
|
||||
async def list_models(self) -> list[dict[str, Any]]:
|
||||
"""List available Gemini models."""
|
||||
try:
|
||||
models = genai.list_models() # type: ignore
|
||||
return [
|
||||
{
|
||||
"name": model.name,
|
||||
"display_name": getattr(model, "display_name", model.name),
|
||||
"description": getattr(model, "description", ""),
|
||||
"supported_generation_methods": getattr(
|
||||
model, "supported_generation_methods", ["generateContent"]
|
||||
),
|
||||
"input_token_limit": getattr(model, "input_token_limit", 32768),
|
||||
"output_token_limit": getattr(model, "output_token_limit", 8192),
|
||||
}
|
||||
for model in models
|
||||
]
|
||||
response = await self.client.models.list()
|
||||
return [model.model_dump() for model in response.data]
|
||||
except Exception as e:
|
||||
from ...core.logging import get_logger
|
||||
|
||||
|
||||
+6
-144
@@ -64,75 +64,10 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
"platform_url": cls.platform_url,
|
||||
}
|
||||
|
||||
def prepare_headers(self, request_headers: dict) -> dict:
|
||||
headers = dict(request_headers)
|
||||
removed_headers = []
|
||||
|
||||
for header in [
|
||||
"host",
|
||||
"content-length",
|
||||
"refund-lnurl",
|
||||
"key-expiry-time",
|
||||
"x-cashu",
|
||||
"authorization",
|
||||
]:
|
||||
if headers.pop(header, None) is not None:
|
||||
removed_headers.append(header)
|
||||
|
||||
if self.api_key:
|
||||
headers["x-goog-api-key"] = self.api_key
|
||||
|
||||
return headers
|
||||
|
||||
def transform_model_name(self, model_id: str) -> str:
|
||||
return model_id.removeprefix("gemini/")
|
||||
|
||||
def prepare_request_body(
|
||||
self, body: bytes | None, model_obj: Model
|
||||
) -> bytes | None:
|
||||
if not body:
|
||||
return body
|
||||
|
||||
try:
|
||||
openai_data = json.loads(body)
|
||||
if not isinstance(openai_data, dict):
|
||||
return body
|
||||
|
||||
gemini_data = {}
|
||||
|
||||
if "messages" in openai_data:
|
||||
contents = []
|
||||
for message in openai_data["messages"]:
|
||||
role = message.get("role", "user")
|
||||
content = message.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
continue
|
||||
elif role == "assistant":
|
||||
role = "model"
|
||||
|
||||
contents.append({"role": role, "parts": [{"text": content}]})
|
||||
gemini_data["contents"] = contents
|
||||
|
||||
generation_config = {}
|
||||
if "temperature" in openai_data:
|
||||
generation_config["temperature"] = openai_data["temperature"]
|
||||
if "max_tokens" in openai_data:
|
||||
generation_config["maxOutputTokens"] = openai_data["max_tokens"]
|
||||
if "top_p" in openai_data:
|
||||
generation_config["topP"] = openai_data["top_p"]
|
||||
|
||||
if generation_config:
|
||||
gemini_data["generationConfig"] = generation_config # type: ignore
|
||||
|
||||
return json.dumps(gemini_data).encode()
|
||||
|
||||
except (json.JSONDecodeError, KeyError, TypeError) as e:
|
||||
logger.warning(
|
||||
f"Failed to transform request body for Gemini: {e}",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
return body
|
||||
|
||||
async def forward_request(
|
||||
self,
|
||||
@@ -324,61 +259,7 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
model_obj,
|
||||
)
|
||||
|
||||
def _transform_gemini_to_openai(self, gemini_response: dict, model_id: str) -> dict:
|
||||
"""Transform Gemini API response to OpenAI format."""
|
||||
candidates = gemini_response.get("candidates", [])
|
||||
if not candidates:
|
||||
return {
|
||||
"id": f"chatcmpl-{hash(str(gemini_response))}"[1:16],
|
||||
"object": "chat.completion",
|
||||
"created": int(__import__("time").time()),
|
||||
"model": model_id,
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
},
|
||||
}
|
||||
|
||||
first_candidate = candidates[0]
|
||||
content = ""
|
||||
if "content" in first_candidate and "parts" in first_candidate["content"]:
|
||||
parts = first_candidate["content"]["parts"]
|
||||
if parts and "text" in parts[0]:
|
||||
content = parts[0]["text"]
|
||||
|
||||
finish_reason = first_candidate.get("finishReason", "stop")
|
||||
if finish_reason == "STOP":
|
||||
finish_reason = "stop"
|
||||
elif finish_reason == "MAX_TOKENS":
|
||||
finish_reason = "length"
|
||||
|
||||
usage_metadata = gemini_response.get("usageMetadata", {})
|
||||
prompt_tokens = usage_metadata.get("promptTokenCount", 0)
|
||||
completion_tokens = usage_metadata.get("candidatesTokenCount", 0)
|
||||
|
||||
return {
|
||||
"id": f"chatcmpl-{hash(str(gemini_response))}"[1:16],
|
||||
"object": "chat.completion",
|
||||
"created": int(__import__("time").time()),
|
||||
"model": model_id,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
},
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
from ..payment.models import Architecture, Model, Pricing, TopProvider
|
||||
@@ -388,34 +269,15 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
|
||||
models_list = []
|
||||
for model_data in models_data:
|
||||
model_name = model_data.get("name", "")
|
||||
if not model_name:
|
||||
model_id = model_data.get("id", "")
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
model_id = model_name.replace("models/", "")
|
||||
display_name = model_data.get("display_name", model_id)
|
||||
description = model_data.get(
|
||||
"description", f"Google {display_name} model"
|
||||
)
|
||||
display_name = model_id
|
||||
description = f"Google {display_name} model"
|
||||
|
||||
supported_methods = model_data.get(
|
||||
"supported_generation_methods", ["generateContent"]
|
||||
)
|
||||
|
||||
input_token_limit = model_data.get("input_token_limit", 32768)
|
||||
output_token_limit = model_data.get("output_token_limit", 8192)
|
||||
context_length = min(input_token_limit, 128000)
|
||||
|
||||
logger.debug(
|
||||
f"Found Gemini model: {model_id}",
|
||||
extra={
|
||||
"display_name": display_name,
|
||||
"supported_methods": supported_methods,
|
||||
"input_token_limit": input_token_limit,
|
||||
"output_token_limit": output_token_limit,
|
||||
"context_length": context_length,
|
||||
},
|
||||
)
|
||||
context_length = 128000
|
||||
output_token_limit = 8192
|
||||
|
||||
pricing_config = Pricing(
|
||||
prompt=0.000003,
|
||||
|
||||
Reference in New Issue
Block a user