utils.py

from py2neo import Graph
from config import *
from neo4j import GraphDatabase
import os
from dotenv import load_dotenv
load_dotenv()
from langchain_openai import OpenAIEmbeddings
from langchain_openai import ChatOpenAI

def get_embeddings_model():
    model_map = {
        'openai': OpenAIEmbeddings(
            model = os.getenv('OPENAI_EMBEDDINGS_MODEL')
        )
    }
    return model_map.get(os.getenv('EMBEDDINGS_MODEL'))

def get_llm_model():
    model_map = {
        'openai': ChatOpenAI(
            model = os.getenv('OPENAI_LLM_MODEL'),
            temperature = os.getenv('TEMPERATURE'),
            max_tokens = os.getenv('MAX_TOKENS'),
            openai_api_base=os.getenv('OPENAI_API_BASE'),
        )
    }
    return model_map.get(os.getenv('LLM_MODEL'))


def structured_output_parser(response_schemas):
    text = '''
    请从以下文本中，抽取出实体信息，并按json格式输出，json包含首尾的 "```json" 和 "```"。
    以下是字段含义和类型，要求输出json中，必须包含下列所有字段：\n
    '''
    for schema in response_schemas:
        text += schema.name + ' 字段，表示：' + schema.description + '，类型为：' + schema.type + '\n'
    return text


def replace_token_in_string(string, slots):
    #print(type(slots))
    for key, value in slots:
        # print(key, value)
        string = string.replace('%'+key+'%', value)
    return string

def replace_graph_token(string, slots):
    for key, value in slots.items():
        # print(key, value)
        string = string.replace('%'+key+'%', value)
    return string
   

# def get_neo4j_conn():
#     return Graph(
#         os.getenv('NEO4J_URI'), 
#         auth = (os.getenv('NEO4J_USERNAME'), os.getenv('NEO4J_PASSWORD'))
#     )

# def get_neo4j_conn():
#     with GraphDatabase.driver(os.getenv('NEO4J_URI'), auth=(os.getenv('NEO4J_USERNAME'), os.getenv('NEO4J_PASSWORD'))) as driver:
#         driver.verify_connectivity()
#         return driver

if __name__ == '__main__':
    llm_model = get_llm_model()
    #print(llm_model.predict('陈华编程是什么？'))