xiamuceer
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
"""OpenAI Provider"""
|
|
from typing import Any, AsyncGenerator, Dict, List, Optional
|
|
|
|
from app.services.ai_clients.openai_client import OpenAIClient
|
|
from .base_provider import BaseAIProvider
|
|
|
|
|
|
class OpenAIProvider(BaseAIProvider):
|
|
"""OpenAI 提供商"""
|
|
|
|
def __init__(self, client: OpenAIClient):
|
|
self.client = client
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
model: str,
|
|
temperature: float,
|
|
max_tokens: int,
|
|
system_prompt: Optional[str] = None,
|
|
tools: Optional[List[Dict]] = None,
|
|
tool_choice: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
return await self.client.chat_completion(
|
|
messages=messages,
|
|
model=model,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
)
|
|
|
|
async def generate_stream(
|
|
self,
|
|
prompt: str,
|
|
model: str,
|
|
temperature: float,
|
|
max_tokens: int,
|
|
system_prompt: Optional[str] = None,
|
|
) -> AsyncGenerator[str, None]:
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
async for chunk in self.client.chat_completion_stream(
|
|
messages=messages,
|
|
model=model,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
):
|
|
yield chunk |