laantungir/routstr-core
1
0
Fork 0
mirror of https://github.com/Routstr/routstr-core.git synced 2026-08-12 12:13:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

improve performance

Browse Source
This commit is contained in:
9qeklajc
2025-11-24 21:20:40 +01:00
parent 12515cb0e3
commit afa4a59bb6
2 changed files with 65 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
routstr/upstream/clients/gemini.py
+8 -1
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import time
from typing import Any, AsyncGenerator
from typing import Any, AsyncGenerator, Awaitable, Callable
import google.generativeai as genai
@@ -95,6 +95,9 @@ class GeminiClient(BaseAPIClient):
messages: list[dict[str, Any]],
temperature: float | None = None,
max_tokens: int | None = None,
usage_callback: Callable[[dict[str, Any]], None] | None = None,
completion_callback: Callable[[str, dict[str, Any] | None], Awaitable[None]]
| None = None,
**kwargs: Any,
) -> AsyncGenerator[dict[str, Any], None]:
"""Generate content using Gemini API (streaming)."""
@@ -171,10 +174,14 @@ class GeminiClient(BaseAPIClient):
if usage_data:
chunk_data["usage"] = usage_data
if usage_callback:
usage_callback(usage_data)
yield chunk_data
if finish_reason == "stop":
if completion_callback:
await completion_callback(model, usage_data)
break
async def list_models(self) -> list[dict[str, Any]]:
routstr/upstream/gemini.py
+57 -52
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import json
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
from fastapi import Request
from fastapi.responses import Response, StreamingResponse
@@ -188,70 +188,73 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
)
if is_streaming:
final_usage_data: dict | None = None
def usage_callback(usage_data: dict[str, Any]) -> None:
"""Callback to capture usage data during streaming"""
nonlocal final_usage_data
final_usage_data = usage_data
async def completion_callback(
model: str, usage_data: dict[str, Any] | None
) -> None:
"""Callback to handle payment when streaming completes"""
nonlocal final_usage_data
if usage_data:
final_usage_data = usage_data
payment_data = {
"model": model,
"usage": final_usage_data,
}
from ..auth import adjust_payment_for_tokens
from ..core.db import create_session
async with create_session() as new_session:
fresh_key = await new_session.get(key.__class__, key.hashed_key)
if fresh_key:
try:
cost_data = await adjust_payment_for_tokens(
fresh_key,
payment_data,
new_session,
max_cost_for_model,
)
logger.info(
"Gemini streaming payment finalized",
extra={
"cost_data": cost_data,
"usage_data": final_usage_data,
"key_hash": key.hashed_key[:8] + "...",
},
)
except Exception as cost_error:
logger.error(
"Error finalizing Gemini streaming payment",
extra={
"error": str(cost_error),
"key_hash": key.hashed_key[:8] + "...",
},
)
response_generator = self.client.generate_content_stream(
model=model_obj.id,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
top_p=top_p,
usage_callback=usage_callback,
completion_callback=completion_callback,
)
async def stream_with_cost() -> AsyncGenerator[bytes, None]:
last_model_seen: str = model_obj.id
final_usage_data: dict | None = None
try:
async for chunk in response_generator:
if chunk.get("usage"):
final_usage_data = chunk["usage"]
sse_data = f"data: {json.dumps(chunk)}\n\n"
yield sse_data.encode()
if (
chunk.get("choices", [{}])[0].get("finish_reason")
== "stop"
):
break
payment_data = {
"model": last_model_seen,
"usage": final_usage_data,
}
from ..auth import adjust_payment_for_tokens
from ..core.db import create_session
async with create_session() as new_session:
fresh_key = await new_session.get(
key.__class__, key.hashed_key
)
if fresh_key:
try:
cost_data = await adjust_payment_for_tokens(
fresh_key,
payment_data,
new_session,
max_cost_for_model,
)
logger.info(
"Gemini streaming payment finalized",
extra={
"cost_data": cost_data,
"usage_data": final_usage_data,
"key_hash": key.hashed_key[:8] + "...",
},
)
except Exception as cost_error:
logger.error(
"Error finalizing Gemini streaming payment",
extra={
"error": str(cost_error),
"key_hash": key.hashed_key[:8] + "...",
},
)
except Exception as e:
logger.error(
"Error in Gemini streaming response",
@@ -395,7 +398,9 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
"description", f"Google {display_name} model"
)
supported_methods = model_data.get("supported_generation_methods", ['generateContent'])
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)