You've already forked DataMate
feat: Implement data synthesis task management with database models and API endpoints (#122)
This commit is contained in:
@@ -0,0 +1,29 @@
|
||||
from typing import Optional
|
||||
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_openai import ChatOpenAI
|
||||
from pydantic import SecretStr
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.db.models.model_config import ModelConfig
|
||||
|
||||
|
||||
async def get_model_by_id(db: AsyncSession, model_id: str) -> Optional[ModelConfig]:
|
||||
"""根据模型ID获取 ModelConfig 记录。"""
|
||||
result = await db.execute(select(ModelConfig).where(ModelConfig.id == model_id))
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
def get_chat_client(model: ModelConfig) -> BaseChatModel:
|
||||
return ChatOpenAI(
|
||||
model=model.model_name,
|
||||
base_url=model.base_url,
|
||||
api_key=SecretStr(model.api_key),
|
||||
)
|
||||
|
||||
|
||||
def chat(model: BaseChatModel, prompt: str) -> str:
|
||||
"""使用指定模型进行聊天"""
|
||||
response = model.invoke(prompt)
|
||||
return response.content
|
||||
Reference in New Issue
Block a user