from langchain_community.utilities import SQLDatabase
from langchain_core.callbacks import BaseCallbackHandler
from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
from uuid import UUID
from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_openai import OpenAIEmbeddings
from langchain.agents.agent_toolkits import create_retriever_tool
from langchain_core.output_parsers import JsonOutputParser
import os
from langchain_core.prompts import (
    ChatPromptTemplate,
    FewShotPromptTemplate,
    MessagesPlaceholder,
    PromptTemplate,
    SystemMessagePromptTemplate,
)
import ast
import re
from utils import query_as_list, get_answer
import gradio as gr
from fewshot import examples

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=os.environ['API_KEY'])
example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    OpenAIEmbeddings(model="text-embedding-3-small", api_key=os.environ['API_KEY']),
    Chroma(persist_directory="data"),
    # Chroma,
    k=5,
    input_keys=["input"],
)

db = SQLDatabase.from_uri("sqlite:///attendance_system.db")

employee = query_as_list(db, "SELECT FullName FROM Employee")

vector_db = Chroma.from_texts(employee, OpenAIEmbeddings(model="text-embedding-3-small", api_key=os.environ['API_KEY']))
retriever = vector_db.as_retriever(search_kwargs={"k": 15})
description = """Use to look up values to filter on. Input is an approximate spelling of the proper noun, output is \
valid proper nouns. Use the noun most similar to the search."""
retriever_tool = create_retriever_tool(
    retriever,
    name="search_proper_nouns",
    description=description,
)


if __name__ == "__main__":
    demo = gr.Interface(fn=get_answer, inputs="text", outputs="text")
    demo.launch()