70 lines
2.7 KiB
Python
70 lines
2.7 KiB
Python
from django.shortcuts import render
|
||
|
||
# Create your views here.
|
||
|
||
from pydantic import Field
|
||
from langchain_core.language_models import LLM
|
||
from langchain_core.outputs import LLMResult, Generation
|
||
import requests
|
||
|
||
class CustomLLM(LLM):
|
||
model_url: str
|
||
|
||
def _call(self, prompt: str, stop: list = None) -> str:
|
||
data = {
|
||
"model": "glm4",
|
||
"messages": [
|
||
{
|
||
"role": "system",
|
||
"content": "你是一个 SQL 助手,严格遵循以下规则:\n"
|
||
"1. 只返回 PostgreSQL 语法 SQL 语句。\n"
|
||
"2. 严格禁止添加任何解释、注释、Markdown 代码块标记(包括 ```sql 和 ```)。\n"
|
||
"3. 输出必须是纯 SQL,且可直接执行,无需任何额外处理。"
|
||
"4. 在 SQL 中如有多个表,请始终使用表名前缀引用字段,避免字段歧义。"
|
||
},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
"stream": False
|
||
}
|
||
response = requests.post(self.model_url, json=data, timeout=600)
|
||
response.raise_for_status()
|
||
content = response.json()["choices"][0]["message"]["content"]
|
||
clean_sql = self.strip_sql_markdown(content)
|
||
print('clean content----------', clean_sql)
|
||
return clean_sql
|
||
|
||
def _generate(self, prompts: list, stop: list = None) -> LLMResult:
|
||
generations = []
|
||
for prompt in prompts:
|
||
text = self._call(prompt, stop)
|
||
generations.append([Generation(text=text)])
|
||
return LLMResult(generations=generations)
|
||
|
||
def strip_sql_markdown(self, content: str) -> str:
|
||
import re
|
||
# 去掉包裹在 ```sql 或 ``` 中的内容
|
||
match = re.search(r"```sql\s*(.*?)```", content, re.DOTALL | re.IGNORECASE)
|
||
if match:
|
||
return match.group(1).strip()
|
||
match = re.search(r"```\s*(.*?)```", content, re.DOTALL)
|
||
if match:
|
||
return match.group(1).strip()
|
||
return content.strip()
|
||
|
||
@property
|
||
def _llm_type(self) -> str:
|
||
return "custom_llm"
|
||
# 实例化
|
||
# model_url = "http://14.22.88.72:11025/v1/chat/completions"
|
||
model_url = "http://139.159.180.64:11434/v1/chat/completions"
|
||
llm = CustomLLM(model_url=model_url)
|
||
|
||
# 测试数据库查询
|
||
from langchain_experimental.sql import SQLDatabaseChain
|
||
from langchain_community.utilities import SQLDatabase
|
||
|
||
db = SQLDatabase.from_uri("postgresql+pg8000://postgres:zcDsj%402024@127.0.0.1:5432/factory")
|
||
db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)
|
||
result = db_chain.invoke("帮我生成一份数据库的表结构")
|
||
print('res-------------------------------',result)
|