1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
|
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_core.tools import tool
import psycopg2
from datetime import datetime, timedelta
import os
from typing import Dict, Any, List
class LibraryAgent:
def __init__(self):
self.db = SQLDatabase.from_uri(
f"postgresql://{os.getenv('DB_USER')}:{os.getenv('DB_PASSWORD')}@{os.getenv('DB_HOST')}/{os.getenv('DB_NAME')}"
)
self.llm = ChatOpenAI(
model="gpt-4-turbo",
temperature=0.3,
api_key=os.getenv("OPENAI_API_KEY")
)
self.toolkit = SQLDatabaseToolkit(db=self.db, llm=self.llm)
self.tools = [
self.borrow_book_tool,
self.return_book_tool,
self.check_availability_tool,
self.get_user_info_tool,
self.search_books_tool
] + self.toolkit.get_tools()
self.agent = self.create_agent()
self.agent_executor = AgentExecutor(
agent=self.agent,
tools=self.tools,
verbose=True,
max_iterations=10,
early_stopping_method="force"
)
def create_agent(self):
"""创建图书馆管理Agent"""
prompt = ChatPromptTemplate.from_messages([
("system", """你是一个专业的图书馆管理助手，负责处理图书借阅、归还、查询等业务。
请根据用户请求选择合适的工具执行操作，并提供清晰、友好的回复。
可用工具：
- borrow_book_tool: 处理图书借阅
- return_book_tool: 处理图书归还
- check_availability_tool: 检查图书可用性
- get_user_info_tool: 获取用户信息
- search_books_tool: 搜索图书
重要规则：
1. 借书前必须检查用户资格和图书可用性
2. 归还图书时需要验证借阅记录
3. 所有操作都需要记录日志
4. 遇到错误时提供友好的错误提示"""),
("human", "{input}"),
("assistant", "{agent_scratchpad}")
])
return create_tool_calling_agent(self.llm, self.tools, prompt)
@tool
def borrow_book_tool(self, user_id: str, book_id: str) -> Dict[str, Any]:
"""借阅图书工具"""
try:
user_info = self.get_user_info_tool(user_id)
if not user_info["eligible"]:
return {"success": False, "message": "用户借阅资格不足"}
availability = self.check_availability_tool(book_id)
if availability["available_copies"] <= 0:
return {"success": False, "message": "图书暂无库存"}
conn = self.get_db_connection()
cursor = conn.cursor()
borrow_date = datetime.now().date()
due_date = borrow_date + timedelta(days=14)
cursor.execute("""
INSERT INTO borrow_records (user_id, book_id, borrow_date, due_date, status)
VALUES (%s, %s, %s, %s, 'borrowed')
RETURNING id, due_date
""", (user_id, book_id, borrow_date, due_date))
result = cursor.fetchone()
transaction_id, due_date = result
conn.commit()
cursor.close()
conn.close()
return {
"success": True,
"transaction_id": transaction_id,
"due_date": due_date.strftime("%Y-%m-%d"),
"message": f"借阅成功！请在{due_date.strftime('%Y-%m-%d')}前归还。"
}
except Exception as e:
return {"success": False, "message": f"借阅失败: {str(e)}"}
@tool
def return_book_tool(self, user_id: str, book_id: str) -> Dict[str, Any]:
"""归还图书工具"""
try:
conn = self.get_db_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT id FROM borrow_records
WHERE user_id = %s AND book_id = %s AND status = 'borrowed'
ORDER BY borrow_date DESC LIMIT 1
""", (user_id, book_id))
record = cursor.fetchone()
if not record:
return {"success": False, "message": "未找到相关借阅记录"}
record_id = record[0]
return_date = datetime.now().date()
cursor.execute("""
UPDATE borrow_records
SET status = 'returned', return_date = %s
WHERE id = %s
""", (return_date, record_id))
conn.commit()
cursor.close()
conn.close()
return {
"success": True,
"message": "图书归还成功！感谢您的使用。"
}
except Exception as e:
return {"success": False, "message": f"归还失败: {str(e)}"}
@tool
def check_availability_tool(self, book_id: str) -> Dict[str, Any]:
"""检查图书可用性"""
try:
conn = self.get_db_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT title, author, total_copies, available_copies
FROM books WHERE id = %s
""", (book_id,))
book = cursor.fetchone()
cursor.close()
conn.close()
if not book:
return {"available": False, "message": "图书不存在"}
title, author, total, available = book
return {
"available": available > 0,
"book_title": title,
"author": author,
"total_copies": total,
"available_copies": available,
"message": f"《{title}》当前有{available}本可借"
}
except Exception as e:
return {"error": str(e)}
@tool
def get_user_info_tool(self, user_id: str) -> Dict[str, Any]:
"""获取用户信息"""
try:
conn = self.get_db_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT name, email, max_borrow_limit FROM users WHERE id = %s
""", (user_id,))
user = cursor.fetchone()
if not user:
return {"eligible": False, "message": "用户不存在"}
name, email, max_limit = user
cursor.execute("""
SELECT COUNT(*) FROM borrow_records
WHERE user_id = %s AND status = 'borrowed'
""", (user_id,))
current_borrows = cursor.fetchone()[0]
cursor.close()
conn.close()
eligible = current_borrows < max_limit
return {
"eligible": eligible,
"name": name,
"email": email,
"max_borrow_limit": max_limit,
"current_borrows": current_borrows,
"remaining_limit": max_limit - current_borrows,
"message": eligible and "用户借阅资格正常" or "已达到最大借阅数量"
}
except Exception as e:
return {"error": str(e)}
@tool
def search_books_tool(self, query: str, category: str = None) -> List[Dict[str, Any]]:
"""搜索图书"""
try:
conn = self.get_db_connection()
cursor = conn.cursor()
conditions = []
params = [f"%{query}%"]
sql = """
SELECT id, title, author, isbn, category, available_copies
FROM books
WHERE title ILIKE %s OR author ILIKE %s OR isbn ILIKE %s
"""
conditions.extend([f"%{query}%", f"%{query}%", f"%{query}%"])
if category:
sql += " AND category = %s"
conditions.append(category)
sql += " ORDER BY title LIMIT 10"
cursor.execute(sql, tuple(conditions))
books = cursor.fetchall()
cursor.close()
conn.close()
return [
{
"id": book[0],
"title": book[1],
"author": book[2],
"isbn": book[3],
"category": book[4],
"available_copies": book[5]
}
for book in books
]
except Exception as e:
return [{"error": str(e)}]
def get_db_connection(self):
"""获取数据库连接"""
return psycopg2.connect(
host=os.getenv('DB_HOST', 'localhost'),
database=os.getenv('DB_NAME', 'library_db'),
user=os.getenv('DB_USER', 'library_user'),
password=os.getenv('DB_PASSWORD', ''),
port=os.getenv('DB_PORT', '5432')
)
def run(self, user_input: str) -> str:
"""运行Agent处理用户输入"""
result = self.agent_executor.invoke({"input": user_input})
return result["output"]
if __name__ == "__main__":
os.environ["OPENAI_API_KEY"] = "your-api-key"
os.environ["DB_HOST"] = "localhost"
os.environ["DB_NAME"] = "library_db"
os.environ["DB_USER"] = "library_user"
os.environ["DB_PASSWORD"] = "your-db-password"
agent = LibraryAgent()
test_inputs = [
"我想借一本《人工智能导论》",
"用户user123可以借多少本书？",
"帮我查找所有编程类的图书",
"我要归还《机器学习实战》"
]
for user_input in test_inputs:
print(f"\n用户: {user_input}")
response = agent.run(user_input)
print(f"助手: {response}")
|
