app.py

# Standard library imports
import os.path

# Third-party imports
import gradio as gr
from dotenv import load_dotenv
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.prompts import ChatPromptTemplate
from langchain.schema import HumanMessage, AIMessage
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain_openai import ChatOpenAI

# Local imports
from init_db import Project, db_description, json_structure

# Load environment variables
load_dotenv()

# Check if database exists
if not os.path.exists('estate.db'):
    raise FileNotFoundError(
        "Database file 'estate.db' not found. Please run 'uv run init_db.py' first."
    )

# Initialize model and database
model = ChatOpenAI(model="o3-mini-2025-01-31", streaming=True)
db = SQLDatabase.from_uri("sqlite:///estate.db")

# Set up SQL toolkit and tools
toolkit = SQLDatabaseToolkit(db=db, llm=model)
tools = toolkit.get_tools()

sql_distance_query = """
SELECT 
    id, 
    url,
    structure->>'$.lat' AS lat,
    structure->>'$.lng' AS lng,
    (6371000 * acos(
        cos(radians(50.08804)) * cos(radians(CAST(structure->>'$.lat' AS FLOAT))) * 
        cos(radians(CAST(structure->>'$.lng' AS FLOAT)) - radians(14.42076)) + 
        sin(radians(50.08804)) * sin(radians(CAST(structure->>'$.lat' AS FLOAT)))
    )) AS distance
FROM project
WHERE structure->>'$.lat' IS NOT NULL 
AND structure->>'$.lng' IS NOT NULL
ORDER BY distance ASC
LIMIT 5;
"""

# Get the OpenAI tools agent prompt
prompt = ChatPromptTemplate.from_messages([
  ("system", f"You are a helpful assistant that presents information about real estate projects in database for investors. Before executing queries make surey youhave schema of dabatase and first rows. This JSON structure info can help you: {json_structure}."),
  ("placeholder", "{chat_history}"),
  ("human", "{input}"),
  ("placeholder", "{agent_scratchpad}"),
])

# Create the agent with OpenAI tools format
agent = create_openai_tools_agent(
    llm=model,
    tools=tools,
    prompt=prompt
)

# Create agent executor
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True
)

def chat_with_sql(message, history):
    try:
      history_langchain_format = []
      for msg in history:
          if msg['role'] == "user":
              history_langchain_format.append(HumanMessage(content=msg['content']))
          elif msg['role'] == "assistant":
              history_langchain_format.append(AIMessage(content=msg['content']))
      history_langchain_format.append(HumanMessage(content=message))
      response = agent_executor.invoke({"input": message, "history": history_langchain_format})
      for i in range(len(response["output"])):
          yield response["output"][:i+1]
    except Exception as e:
        yield f"Error: {str(e)}"

# Create the Gradio interface
count_projects = Project.select().where(Project.structure.is_null(False)).count()
demo = gr.ChatInterface(
    fn=chat_with_sql,
    title="Estate Chat",
    description=f"Planning to buy a new apartment? Ask me about {count_projects} new buildings (novostavby) in Prague 🇨🇿",
    examples=[
        "Project with the lowest deposit before completion",
        "Cheapest apartment in Prague",
        "Best price per sqm in Prague"
    ],
    type="messages"
)

if __name__ == "__main__":
    demo.launch(share=False, debug=True)