Aimff/Diesel_Fault_Diagnosis_System

1#from langchain_openai import OpenAIEmbeddings2from langchain_openai import ChatOpenAI3from py2neo import Graph4from config import *5from langchain_community.embeddings import HuggingFaceEmbeddings6import os7from dotenv import load_dotenv8load_dotenv()9 10 11def get_embeddings_model():12    model_map = {13        'text2vec': HuggingFaceEmbeddings(14            model_name=os.getenv("HuggingFaceEmbedding_MODEL"),  # 或你下载到本地的模型路径15            model_kwargs={"device": os.getenv("DEVICE_TYPE", "cpu")},  # 或 "cuda" 如果你有GPU16            encode_kwargs={"normalize_embeddings": True}17        )18    }19    return model_map.get(os.getenv('EMBEDDINGS_MODEL', 'text2vec'))20 21 22def get_llm_model():23    return ChatOpenAI(24        openai_api_key=os.getenv("DEEPSEEK_API_KEY"),25        openai_api_base="https://api.deepseek.com/v1",26        model=os.getenv("DEEPSEEK_LLM_MODEL"),27        temperature=float(os.getenv("TEMPERATURE", 0)),  # 强制转为浮点28        max_tokens=int(os.getenv("MAX_TOKENS", 2000))    # 强制转为整数29    )30 31def structured_output_parser(response_schemas):32    text = '''33    请从以下文本中，抽取出实体信息，并按json格式输出，json包含首尾的 "```json" 和 "```"。34    以下是字段含义和类型，要求输出json中，必须包含下列所有字段：\n35    '''36    for schema in response_schemas:37        text += schema.name + ' 字段，表示：' + schema.description + '，类型为：' + schema.type + '\n'38    return text39 40 41def replace_token_in_string(string, slots):42    for key, value in slots:43        string = string.replace('%'+key+'%', value)44    return string45 46 47def get_neo4j_conn():48    return Graph(49        os.getenv('NEO4J_URI'), 50        auth = (os.getenv('NEO4J_USERNAME'), os.getenv('NEO4J_PASSWORD'))51    )52 53if __name__ == '__main__':54    llm_model = get_llm_model()55    print(llm_model.predict('常见柴油机故障有哪些？'))