from langchain_community.embeddings import HuggingFaceEmbeddings
from sparkai.llm.llm import ChatSparkLLM
from sparkai.core.messages import ChatMessage

# 星火 API 配置信息
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v1.1/chat'
SPARKAI_APP_ID = 'ebf052b9'
SPARKAI_API_SECRET = 'NzIwNzYyOWM0ZmM3OTk3YWVmMDlmZTVj'
SPARKAI_API_KEY = '930eb4722c8d5bc3254779893f7a34e4'
SPARKAI_DOMAIN = 'general'

def get_chat_model(model="spark-gpt", temperature=0):
    """
    获取星火聊天模型
    """
    return ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False
    )

# 使用模型名称初始化 HuggingFaceEmbeddings
embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

def matter(txt):
    """
    调用星火 API 生成结果
    """
    spark = get_chat_model()
    messages = [ChatMessage(
        role="user",
        content=txt
    )]
    a = spark.generate([messages])
    return a.generations[0][0].text