import pandas as pd import streamlit as st from operator import itemgetter from langchain.chains.openai_tools import create_extraction_chain_pydantic from langchain_core.pydantic_v1 import BaseModel, Field from langchain_openai import ChatOpenAI from typing import List class Table(BaseModel): """Table in SQL database.""" name: str = Field(description="Name of table in SQL database.") def get_tables(tables: List[Table]) -> List[str]: return [table.name for table in tables] @st.cache_data def get_table_details(): try: table_description = pd.read_excel("database_table_descriptions.xlsx") table_details = "" for index, row in table_description.iterrows(): table_details += f"Table Name:{row['Table']}\nTable Description:{row['Description']}\n\n" return table_details except Exception as e: st.error(f"Error reading table descriptions: {str(e)}") return "" def create_table_chain(api_key): table_details = get_table_details() table_details_prompt = f"""Return the names of ALL the SQL tables that MIGHT be relevant to the user question. \ The tables are: {table_details} Remember to include ALL POTENTIALLY RELEVANT tables, even if you're not sure that they're needed.""" llm = ChatOpenAI(temperature=0.7, model="gpt-3.5-turbo", api_key=api_key) return ( {"input": itemgetter("question")} | create_extraction_chain_pydantic(Table, llm, system_message=table_details_prompt) | get_tables )