Database Tools

from langchain_nebius import ChatNebius
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import os

def get_llm():
    return ChatNebius(
        model="zai-org/GLM-4.5-Air",
        temperature=0.1,
        top_p=0.95,
        api_key=os.getenv("NEBIUS_API_KEY")
    )

def translate_to_sql(natural_question: str) -> str

from ai_services import translate_to_sql

question = "What are the top 5 most expensive products?"
sql_query = translate_to_sql(question)
print(sql_query)
# Output: SELECT `id`, `name`, `price` FROM `product` ORDER BY `price` DESC LIMIT 5;

DB_SCHEMA = """
Database: Ecommerce (MySQL)
Tables:
- category (id, uuid, name, description, date_created, date_updated)
- product (id, uuid, name, description, price, stock_quantity, category_id, date_created, date_updated)
- product_category (id, product_id, category_id)
- user (id, uuid, name, email, address, date_created, date_updated)
- order (id, uuid, user_id, total_amount, status, order_date, date_created, date_updated)
- order_item (id, order_id, product_id, quantity, unit_price, date_created, date_updated)

Relationships:
- product.category_id -> category.id
- product_category.product_id -> product.id
- product_category.category_id -> category.id
- order.user_id -> user.id
- order_item.order_id -> order.id
- order_item.product_id -> product.id

Note: All tables use MySQL syntax with backticks for identifiers.
"""

import pymysql
from urllib.parse import urlparse

def parse_connection_string(connection_string):
    """Parse MySQL connection string and return database config"""
    parsed = urlparse(connection_string)
    
    return {
        "host": parsed.hostname,
        "user": parsed.username,
        "password": parsed.password,
        "database": parsed.path.lstrip("/"),
        "port": parsed.port or 3306,
        "ssl": {"ssl": {}}
    }

def get_database_connection():
    """Create and return a database connection"""
    db_config = get_database_config()
    connection = pymysql.connect(**db_config)
    return connection

def execute_query(sql_query: str) -> tuple[list[dict] | None, str | None]

from database import execute_query

sql = "SELECT * FROM product LIMIT 5"
results, error = execute_query(sql)

if error:
    print(f"Error: {error}")
else:
    for row in results:
        print(row)

import asyncio
from agno.agent import Agent
from agno.models.nebius import Nebius
from agno.tools.mcp import MCPTools
import os

async def run_database_agent(message: str):
    mcp_tools = MCPTools(
        "uvx --from gibson-cli@latest gibson mcp run",
        timeout_seconds=300
    )
    
    await mcp_tools.connect()
    
    agent = Agent(
        name="GibsonAIAgent",
        model=Nebius(
            id="meta-llama/Meta-Llama-3.1-70B-Instruct",
            api_key=os.getenv("NEBIUS_API_KEY")
        ),
        tools=[mcp_tools],
        description="Agent for managing database projects and schemas",
        instructions="""You are a GibsonAI database assistant. Help users manage their database projects and schemas.

Your capabilities include:
- Creating new GibsonAI projects
- Managing database schemas (tables, columns, relationships)
- Deploying schema changes to hosted databases
- Querying database schemas and data
- Providing insights about database structure and best practices
""",
        markdown=True,
        show_tool_calls=True
    )
    
    await agent.aprint_response(message, stream=True)
    await mcp_tools.close()

# Example usage
asyncio.run(
    run_database_agent(
        "Create a new GibsonAI project for my Blog Application."
    )
)

def explain_results(results: list[dict], original_question: str) -> str

from ai_services import translate_to_sql, explain_results
from database import execute_query

question = "How many orders do we have?"
sql = translate_to_sql(question)
results, error = execute_query(sql)

if not error:
    explanation = explain_results(results, question)
    print(explanation)
    # Output: "You currently have 150 orders in your database."

import streamlit as st
from ai_services import translate_to_sql, explain_results
from database import execute_query, parse_connection_string

# Streamlit interface
st.title("Talk to Your Database")

# Database connection
connection_string = st.text_input(
    "Database Connection String",
    type="password",
    placeholder="mysql://user:password@host:port/database"
)

if connection_string:
    st.session_state.db_config = parse_connection_string(connection_string)
    
    # Natural language query
    user_question = st.text_input(
        "Ask a question about your database",
        placeholder="What are the top selling products?"
    )
    
    if user_question:
        # Convert to SQL
        with st.spinner("Generating SQL..."):
            sql_query = translate_to_sql(user_question)
            st.code(sql_query, language="sql")
        
        # Execute query
        with st.spinner("Executing query..."):
            results, error = execute_query(sql_query)
        
        if error:
            st.error(f"Error: {error}")
        else:
            # Show results
            st.dataframe(results)
            
            # Explain results
            explanation = explain_results(results, user_question)
            st.info(explanation)

prompt = ChatPromptTemplate.from_messages([
    (
        "system",
        """You are a MySQL SQL expert. Convert natural language questions to MySQL queries.

Database Schema:
{db_schema}

Rules:
1. For SELECT queries: Use appropriate JOINs when needed
2. Use proper MySQL syntax with backticks for identifiers
3. Return ONLY the SQL query, no explanations
4. Do NOT include any thinking or reasoning
5. Include LIMIT 100 for large result sets

Example queries:
Question: "What are the product categories?"
SELECT `id`, `name`, `description` FROM `category` ORDER BY `name`;

Question: "Show products with their categories"
SELECT p.`name`, p.`price`, c.`name` as category 
FROM `product` p 
LEFT JOIN `category` c ON p.`category_id` = c.`id`;
"""
    ),
    ("human", "Question: {question}\n\nGenerate the SQL query:")
])

chain = prompt | llm | StrOutputParser()

# Nebius API for LLM
NEBIUS_API_KEY=your_nebius_api_key

# Database connection (example)
DATABASE_URL=mysql://user:password@host:port/database