Spaces:
Runtime error
Runtime error
from langchain.tools import BaseTool | |
from langchain_openai import ChatOpenAI | |
from langchain.agents import AgentExecutor, create_openai_tools_agent | |
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder | |
from langchain_community.utilities import SQLDatabase | |
from langchain.schema import SystemMessage, HumanMessage, AIMessage | |
import os | |
import psycopg2 | |
import gradio as gr | |
# Set up database connection details | |
pg_uri = "postgresql://ridhima:0skESLQ9D6c3m7smqwG47peapk7HzVvu@dpg-cq2h613v2p9s73esp8eg-a.singapore-postgres.render.com/hr_qugd" | |
def get_db_connection(): | |
conn = psycopg2.connect(pg_uri) | |
return conn | |
class LeaveRequestInfoTool(BaseTool): | |
name = "leave_request_information" | |
description = "Provides information about the leave request process and database schema." | |
def _run(self, query: str) -> str: | |
return """ | |
Leave Request Process: | |
1. Collect employee ID | |
2. Ask for leave type | |
3. Get reason for leave | |
4. Get start date of leave | |
5. Get end date of leave | |
6. Calculate duration (in days) | |
7. Insert data into request table | |
8. Confirm submission to user | |
Database Schema: | |
Table: request | |
Columns: | |
- leave_id (auto-increment integer) | |
- employee_id (integer) | |
- leave_type (text) | |
- reason (text) | |
- start_of_leave (date) | |
- end_of_leave (date) | |
- duration (integer, calculated in days) | |
- leave_status (text, default 'Pending') | |
Instructions: | |
- Collect all necessary information from the user one by one. | |
- Calculate the duration as the number of days between start_of_leave and end_of_leave. | |
- Once all information is collected, formulate an SQL INSERT statement for the 'request' table. | |
- REMEMBER TO EXECUTE THE INSERT QUERY FOR EACH REQUEST ONLY ONCE. | |
- CRITICAL: EXECUTE THE INSERT QUERY FOR EACH REQUEST ONLY ONCE. | |
- If you receive an error message saying an insertion has already been made, DO NOT attempt to insert again. | |
- Instead, inform the user that their request has been submitted and ask if they need anything else. | |
- After insertion, confirm to the user that their request has been submitted. | |
""" | |
class SQLAgentTool(BaseTool): | |
name = "sql_agent" | |
description = "Use this tool to interact with the database and execute SQL queries." | |
def _run(self, query: str) -> str: | |
print("Executing Query: ", query) | |
conn = get_db_connection() | |
cur = conn.cursor() | |
try: | |
cur.execute(query) | |
conn.commit() | |
result = "Your leave request has been submitted successfully." | |
except Exception as e: | |
conn.rollback() | |
result = f"An error occurred: {str(e)}" | |
finally: | |
cur.close() | |
conn.close() | |
return result | |
prompt = ChatPromptTemplate.from_messages([ | |
SystemMessage(content="""You are an HR assistant. You can help with leave requests and provide information about company policies. | |
- For leave requests, ask for each piece of information one at a time. After collecting all information, use the sql_agent tool to INSERT the data into the 'request' table. | |
- For information queries about company policies, use the rag_info tool to provide accurate information from the RAG-trained model. | |
The pieces of information you need to collect for leave requests are: | |
1. Employee ID | |
2. Leave Type | |
3. Reason for Leave | |
4. Start Date of Leave (YYYY-MM-DD) | |
5. End Date of Leave (YYYY-MM-DD) | |
Calculate the duration as the number of days between start and end dates using PostgreSQL functions and convert duration to integer. | |
Use 'Pending' as the default leave_status. | |
Remember to use the correct column names as per the table structure: | |
request(leave_id, employee_id, leave_type, reason, start_of_leave, end_of_leave, duration, leave_status) | |
Where leave_id is auto-increment and should not be included in the INSERT statement. | |
For company policies, provide detailed and accurate information based on the RAG-trained model. | |
"""), | |
MessagesPlaceholder(variable_name="chat_history"), | |
("human", "{input}"), | |
MessagesPlaceholder(variable_name="agent_scratchpad") | |
]) | |