DB-GPT/dbgpt/model/llm/base.py
FangYin Cheng cd725db1fb
refactor: The first refactored version for sdk release (#907)
Co-authored-by: chengfangyin2 <chengfangyin3@jd.com>
2023-12-08 14:45:59 +08:00

38 lines
691 B
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from dataclasses import dataclass
from typing import TypedDict
class Message(TypedDict):
"""Vicuna Message object containing a role and the message content"""
role: str
content: str
@dataclass
class ModelInfo:
"""Struct for model information.
Would be lovely to eventually get this directly from APIs
"""
name: str
max_tokens: int
@dataclass
class LLMResponse:
"""Standard response struct for a response from a LLM model."""
model_info = ModelInfo
@dataclass
class ChatModelResponse(LLMResponse):
"""Standard response struct for a response from an LLM model."""
content: str = None