diff --git a/Dockerfile b/Dockerfile index 7f7b119..9c365e7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -3,10 +3,10 @@ FROM python:3.10-slim WORKDIR /app # 复制所需文件到容器中 -COPY ./app /app/app COPY ./requirements.txt /app RUN pip install --no-cache-dir -r requirements.txt +COPY ./app /app/app ENV API_KEYS='["your_api_key_1"]' ENV ALLOWED_TOKENS='["your_token_1"]' ENV BASE_URL=https://generativelanguage.googleapis.com/v1beta diff --git a/app/api/gemini_routes.py b/app/api/gemini_routes.py index a50889a..8905257 100644 --- a/app/api/gemini_routes.py +++ b/app/api/gemini_routes.py @@ -63,7 +63,7 @@ async def generate_content( logger.info(f"Using API key: {api_key}") try: - response = chat_service.generate_content( + response = await chat_service.generate_content( model=model_name, request=request, api_key=api_key @@ -122,7 +122,7 @@ async def verify_key(api_key: str): ) ] ) - response = chat_service.generate_content(settings.TEST_MODEL,gemini_requset, api_key) + response = await chat_service.generate_content(settings.TEST_MODEL,gemini_requset, api_key) if response: return JSONResponse({"status": "valid"}) return JSONResponse({"status": "invalid"}) diff --git a/app/services/chat/api_client.py b/app/services/chat/api_client.py index df7f0cb..0a834d6 100644 --- a/app/services/chat/api_client.py +++ b/app/services/chat/api_client.py @@ -24,13 +24,13 @@ class GeminiApiClient(ApiClient): self.base_url = base_url self.timeout = timeout - def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]: + 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) if model.endswith("-search"): model = model[:-7] - with httpx.Client(timeout=timeout) as client: + async with httpx.AsyncClient(timeout=timeout) as client: url = f"{self.base_url}/models/{model}:generateContent?key={api_key}" - response = client.post(url, json=payload) + response = await client.post(url, json=payload) if response.status_code != 200: error_content = response.text raise Exception(f"API call failed with status code {response.status_code}, {error_content}") diff --git a/app/services/gemini_chat_service.py b/app/services/gemini_chat_service.py index 48d4991..d281679 100644 --- a/app/services/gemini_chat_service.py +++ b/app/services/gemini_chat_service.py @@ -79,10 +79,10 @@ class GeminiChatService: self.key_manager = key_manager self.response_handler = GeminiResponseHandler() - def generate_content(self, model: str, request: GeminiRequest, api_key: str) -> Dict[str, Any]: + async def generate_content(self, model: str, request: GeminiRequest, api_key: str) -> Dict[str, Any]: """生成内容""" payload = _build_payload(model, request) - response = self.api_client.generate_content(payload, model, api_key) + response = await self.api_client.generate_content(payload, model, api_key) return self.response_handler.handle_response(response, model, stream=False) async def stream_generate_content(self, model: str, request: GeminiRequest, api_key: str) -> AsyncGenerator[str, None]: diff --git a/app/services/openai_chat_service.py b/app/services/openai_chat_service.py index 5852eb2..bd8c566 100644 --- a/app/services/openai_chat_service.py +++ b/app/services/openai_chat_service.py @@ -144,13 +144,13 @@ class OpenAIChatService: if request.stream: return self._handle_stream_completion(request.model, payload, api_key) - return self._handle_normal_completion(request.model, payload, api_key) + return await self._handle_normal_completion(request.model, payload, api_key) - def _handle_normal_completion( + async def _handle_normal_completion( self, model: str, payload: Dict[str, Any], api_key: str ) -> Dict[str, Any]: """处理普通聊天完成""" - response = self.api_client.generate_content(payload, model, api_key) + response = await self.api_client.generate_content(payload, model, api_key) return self.response_handler.handle_response( response, model, stream=False, finish_reason="stop" ) diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..965b3fa --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,9 @@ +version: '3' + +services: + gemini-balance: + build: . + ports: + - "8000:8000" + env_file: + - .env