110 lines
4.1 KiB
Python
110 lines
4.1 KiB
Python
|
"""LLM service for workflow execution."""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
from typing import List, Dict, Any, Optional, AsyncGenerator
|
|||
|
|
from langchain_openai import ChatOpenAI
|
|||
|
|
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
|
|||
|
|
|
|||
|
|
from ..models.llm_config import LLMConfig
|
|||
|
|
from loguru import logger
|
|||
|
|
|
|||
|
|
class LLMService:
|
|||
|
|
"""LLM服务，用于工作流中的大模型调用"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
async def chat_completion(
|
|||
|
|
self,
|
|||
|
|
model_config: LLMConfig,
|
|||
|
|
messages: List[Dict[str, str]],
|
|||
|
|
temperature: Optional[float] = None,
|
|||
|
|
max_tokens: Optional[int] = None
|
|||
|
|
) -> str:
|
|||
|
|
"""调用大模型进行对话完成"""
|
|||
|
|
try:
|
|||
|
|
# 创建LangChain ChatOpenAI实例
|
|||
|
|
llm = ChatOpenAI(
|
|||
|
|
model=model_config.model_name,
|
|||
|
|
api_key=model_config.api_key,
|
|||
|
|
base_url=model_config.base_url,
|
|||
|
|
temperature=temperature or model_config.temperature,
|
|||
|
|
max_tokens=max_tokens or model_config.max_tokens,
|
|||
|
|
streaming=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 转换消息格式
|
|||
|
|
langchain_messages = []
|
|||
|
|
for msg in messages:
|
|||
|
|
role = msg.get("role", "user")
|
|||
|
|
content = msg.get("content", "")
|
|||
|
|
|
|||
|
|
if role == "system":
|
|||
|
|
langchain_messages.append(SystemMessage(content=content))
|
|||
|
|
elif role == "user":
|
|||
|
|
langchain_messages.append(HumanMessage(content=content))
|
|||
|
|
elif role == "assistant":
|
|||
|
|
langchain_messages.append(AIMessage(content=content))
|
|||
|
|
|
|||
|
|
# 调用LLM
|
|||
|
|
response = await llm.ainvoke(langchain_messages)
|
|||
|
|
|
|||
|
|
# 返回响应内容
|
|||
|
|
return response.content
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"LLM调用失败: {str(e)}")
|
|||
|
|
raise Exception(f"LLM调用失败: {str(e)}")
|
|||
|
|
|
|||
|
|
async def chat_completion_stream(
|
|||
|
|
self,
|
|||
|
|
model_config: LLMConfig,
|
|||
|
|
messages: List[Dict[str, str]],
|
|||
|
|
temperature: Optional[float] = None,
|
|||
|
|
max_tokens: Optional[int] = None
|
|||
|
|
) -> AsyncGenerator[str, None]:
|
|||
|
|
"""调用大模型进行流式对话完成"""
|
|||
|
|
try:
|
|||
|
|
# 创建LangChain ChatOpenAI实例（流式）
|
|||
|
|
llm = ChatOpenAI(
|
|||
|
|
model=model_config.model_name,
|
|||
|
|
api_key=model_config.api_key,
|
|||
|
|
base_url=model_config.base_url,
|
|||
|
|
temperature=temperature or model_config.temperature,
|
|||
|
|
max_tokens=max_tokens or model_config.max_tokens,
|
|||
|
|
streaming=True
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 转换消息格式
|
|||
|
|
langchain_messages = []
|
|||
|
|
for msg in messages:
|
|||
|
|
role = msg.get("role", "user")
|
|||
|
|
content = msg.get("content", "")
|
|||
|
|
|
|||
|
|
if role == "system":
|
|||
|
|
langchain_messages.append(SystemMessage(content=content))
|
|||
|
|
elif role == "user":
|
|||
|
|
langchain_messages.append(HumanMessage(content=content))
|
|||
|
|
elif role == "assistant":
|
|||
|
|
langchain_messages.append(AIMessage(content=content))
|
|||
|
|
|
|||
|
|
# 流式调用LLM
|
|||
|
|
async for chunk in llm.astream(langchain_messages):
|
|||
|
|
if hasattr(chunk, 'content') and chunk.content:
|
|||
|
|
yield chunk.content
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"LLM流式调用失败: {str(e)}")
|
|||
|
|
raise Exception(f"LLM流式调用失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def get_model_info(self, model_config: LLMConfig) -> Dict[str, Any]:
|
|||
|
|
"""获取模型信息"""
|
|||
|
|
return {
|
|||
|
|
"id": model_config.id,
|
|||
|
|
"name": model_config.model_name,
|
|||
|
|
"provider": model_config.provider,
|
|||
|
|
"base_url": model_config.base_url,
|
|||
|
|
"temperature": model_config.temperature,
|
|||
|
|
"max_tokens": model_config.max_tokens,
|
|||
|
|
"is_active": model_config.is_active
|
|||
|
|
}
|