mirror of
https://github.com/open-webui/openapi-servers
synced 2025-06-26 18:17:04 +00:00
feat: sql server
This commit is contained in:
parent
c6a791d1e1
commit
2cc73881eb
121
servers/sql/main.py
Normal file
121
servers/sql/main.py
Normal file
@ -0,0 +1,121 @@
|
||||
import os
|
||||
from fastapi import FastAPI, HTTPException, Query
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Optional
|
||||
|
||||
# --- LLM/SQL libraries ---
|
||||
from langchain_experimental.sql import SQLDatabaseChain
|
||||
from langchain_community.llms.openai import OpenAI
|
||||
from langchain_community.utilities import SQLDatabase
|
||||
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
# -- Load DB URL from environment variable --
|
||||
DATABASE_URL = os.getenv("DATABASE_URL")
|
||||
if not DATABASE_URL:
|
||||
raise RuntimeError("DATABASE_URL environment variable must be set.")
|
||||
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") # Set this in your environment
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Pydantic models
|
||||
# -------------------------------
|
||||
class SQLChatInput(BaseModel):
|
||||
query: str = Field(
|
||||
...,
|
||||
description="Your question or task in natural language (e.g. 'Show me the top 10 customers by sales.')",
|
||||
)
|
||||
|
||||
|
||||
class SQLChatOutput(BaseModel):
|
||||
sql: str = Field(..., description="SQL that was executed")
|
||||
answer: str = Field(..., description="Answer to your query, from the database")
|
||||
raw_result: Optional[list] = Field(
|
||||
None, description="Raw result rows (list of dict/tuples)"
|
||||
)
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# API Setup
|
||||
# -------------------------------
|
||||
app = FastAPI(
|
||||
title="Chat with SQL API",
|
||||
version="1.0.0",
|
||||
description=(
|
||||
"Chat in natural language with any SQL database using LLMs. "
|
||||
"Query and analyze your data conversationally!"
|
||||
),
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# LLM + SQL Chain Setup (singleton)
|
||||
# -------------------------------
|
||||
def get_chain():
|
||||
# Initiate reflected SQLAlchemy DB
|
||||
db = SQLDatabase.from_uri(DATABASE_URL)
|
||||
# LLM instance: using OpenAI GPT (or swap for your preferred)
|
||||
llm = OpenAI(
|
||||
temperature=0, openai_api_key=OPENAI_API_KEY, model_name="gpt-3.5-turbo"
|
||||
)
|
||||
return SQLDatabaseChain.from_llm(
|
||||
llm, db, verbose=True, return_sql=True, return_intermediate_steps=True
|
||||
)
|
||||
|
||||
|
||||
sql_chain = get_chain()
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Schema endpoint
|
||||
# -------------------------------
|
||||
@app.get("/schema", summary="Get database schema overview")
|
||||
def get_db_schema():
|
||||
"""
|
||||
Returns the tables and columns for the currently connected database.
|
||||
"""
|
||||
try:
|
||||
db = sql_chain.database
|
||||
return db.get_table_info()
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Failed to retrieve schema info: {e}"
|
||||
)
|
||||
|
||||
|
||||
# -------------------------------
|
||||
# Chatting endpoint
|
||||
# -------------------------------
|
||||
@app.post(
|
||||
"/chat_sql", response_model=SQLChatOutput, summary="Chat with your SQL database"
|
||||
)
|
||||
def chat_sql(data: SQLChatInput):
|
||||
"""
|
||||
Enter a natural language instruction/question, get answer from your database.
|
||||
"""
|
||||
try:
|
||||
# Run chain
|
||||
result = sql_chain({"query": data.query})
|
||||
# result example: {'result': 'Answer in plain text', 'intermediate_steps': {'sql_cmd': sql, ...}}
|
||||
answer = result["result"]
|
||||
sql = None
|
||||
raw_result = None
|
||||
if "intermediate_steps" in result and "sql_cmd" in result["intermediate_steps"]:
|
||||
sql = result["intermediate_steps"]["sql_cmd"]
|
||||
if "intermediate_steps" in result and "result" in result["intermediate_steps"]:
|
||||
raw_result = result["intermediate_steps"]["result"]
|
||||
return SQLChatOutput(sql=sql or "", answer=answer, raw_result=raw_result)
|
||||
except SQLAlchemyError as e:
|
||||
raise HTTPException(status_code=400, detail=f"Database error: {str(e)}")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Error: {e}")
|
11
servers/sql/requirements.txt
Normal file
11
servers/sql/requirements.txt
Normal file
@ -0,0 +1,11 @@
|
||||
fastapi
|
||||
uvicorn[standard]
|
||||
pydantic
|
||||
python-multipart
|
||||
|
||||
langchain
|
||||
langchain_community
|
||||
langchain-openai
|
||||
langchain-experimental
|
||||
sentence_transformers
|
||||
sqlalchemy
|
Loading…
Reference in New Issue
Block a user