langchain structured sql reply

llm-structured-response.py

import streamlit as st
import os


#os.environ["OPENAI_API_KEY"] = st.secrets['OPENAI_API_KEY'] # for gpt
os.environ["OPENAI_API_KEY"] = st.secrets['LITELLM_KEY']     # for litellm

from pydantic import BaseModel, Field
class SQLResponse(BaseModel):
    """Defines the structure for SQL response."""
    question: str = Field(description="The question asked by the user")
    sql_query: str = Field(description="The SQL query generated by the assistant.")
    explanation: str = Field(description="A detailed explanation of how the SQL query answers the input question.")

from langchain_openai import ChatOpenAI

#llm = ChatOpenAI(model="gpt-4", temperature=0)

llm = ChatOpenAI(model = "gorilla", temperature=0, base_url="https://llm.nrp-nautilus.io/")
#llm = ChatOpenAI(model = "groq-tools", temperature=0, base_url="https://llm.nrp-nautilus.io/")
#llm = ChatOpenAI(model = "llama3", temperature=0, base_url="https://llm.nrp-nautilus.io/")

structured_llm = llm.with_structured_output(SQLResponse)

structured_llm.invoke("Show the first few rows of a table using a SQL query.")