58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
|
import dashscope
|
||
|
|
from dashscope import Generation
|
||
|
|
from typing import AsyncGenerator, Optional, Generator
|
||
|
|
|
||
|
|
|
||
|
|
class LLMService:
|
||
|
|
def __init__(self):
|
||
|
|
from app.core.config import settings
|
||
|
|
self.model = settings.llm_model
|
||
|
|
dashscope.api_key = settings.dashscope_api_key
|
||
|
|
|
||
|
|
def generate_stream(
|
||
|
|
self,
|
||
|
|
prompt: str,
|
||
|
|
system_prompt: Optional[str] = None,
|
||
|
|
) -> Generator[str, None, None]:
|
||
|
|
"""流式生成文本"""
|
||
|
|
messages = []
|
||
|
|
if system_prompt:
|
||
|
|
messages.append({"role": "system", "content": system_prompt})
|
||
|
|
messages.append({"role": "user", "content": prompt})
|
||
|
|
|
||
|
|
responses = Generation.call(
|
||
|
|
model=self.model,
|
||
|
|
messages=messages,
|
||
|
|
result_format="message",
|
||
|
|
stream=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
for response in responses:
|
||
|
|
if response.status_code == 200:
|
||
|
|
content = response.output.choices[0].message.content
|
||
|
|
if content:
|
||
|
|
yield content
|
||
|
|
|
||
|
|
async def generate(
|
||
|
|
self,
|
||
|
|
prompt: str,
|
||
|
|
system_prompt: Optional[str] = None,
|
||
|
|
) -> str:
|
||
|
|
"""一次性生成文本"""
|
||
|
|
messages = []
|
||
|
|
if system_prompt:
|
||
|
|
messages.append({"role": "system", "content": system_prompt})
|
||
|
|
messages.append({"role": "user", "content": prompt})
|
||
|
|
|
||
|
|
response = Generation.call(
|
||
|
|
model=self.model,
|
||
|
|
messages=messages,
|
||
|
|
result_format="message",
|
||
|
|
)
|
||
|
|
|
||
|
|
if response.status_code == 200:
|
||
|
|
return response.output.choices[0].message.content
|
||
|
|
raise Exception(f"LLM generation failed: {response.code}")
|
||
|
|
|
||
|
|
|
||
|
|
llm_service = LLMService()
|