70 lines
1.6 KiB
Python
70 lines
1.6 KiB
Python
import json
|
||
from openai import OpenAI
|
||
|
||
def generate_json_via_llm(
|
||
prompt: str,
|
||
base_url: str,
|
||
api_key: str,
|
||
model_id: str
|
||
) -> str:
|
||
client = OpenAI(
|
||
api_key=api_key,
|
||
base_url=base_url
|
||
)
|
||
|
||
try:
|
||
|
||
|
||
response = client.chat.completions.create(
|
||
model=model_id,
|
||
messages=[
|
||
{
|
||
"role": "user",
|
||
"content": prompt
|
||
}
|
||
],
|
||
response_format={
|
||
'type': 'json_object'
|
||
}
|
||
)
|
||
|
||
return response.choices[0].message.content
|
||
except Exception as e:
|
||
raise RuntimeError(f"API请求失败: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
import os
|
||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||
from config.llm import load_config
|
||
# 将项目根目录添加到 sys.path 中
|
||
|
||
# 示例用法
|
||
try:
|
||
config = load_config()
|
||
print(config)
|
||
result = generate_json_via_llm(
|
||
prompt="""测试，随便生成点什么，返回json格式的字符串,格式如下
|
||
{
|
||
"dataset":[
|
||
{
|
||
"question":"",
|
||
"answer":""
|
||
},
|
||
{
|
||
"question":"",
|
||
"answer":""
|
||
}
|
||
......
|
||
]
|
||
}
|
||
""",
|
||
base_url=config["openai"]["base_url"],
|
||
api_key=config["openai"]["api_key"],
|
||
model_id=config["openai"]["model_id"],
|
||
)
|
||
print(result)
|
||
except Exception as e:
|
||
print(f"错误: {e}")
|