用PydanticAI寫AI Agent太簡(jiǎn)單:全流程CRUD
AI代理簡(jiǎn)介
你可能聽說過“生成式AI”這個(gè)術(shù)語——但什么是AI代理?本質(zhì)上,它們是利用高級(jí)語言模型來更自主地處理任務(wù)的下一代工具。它們能夠解析用戶查詢,挑選相關(guān)信息,并以結(jié)構(gòu)化的方式與外部服務(wù)交互。
上圖展示了AI如何從基本的生成模型發(fā)展到能夠與各種工具協(xié)作的復(fù)雜AI代理。
PydanticAI簡(jiǎn)介
PydanticAI是一個(gè)Python代理框架,旨在簡(jiǎn)化利用生成式AI開發(fā)生產(chǎn)級(jí)應(yīng)用的流程。由Pydantic背后的團(tuán)隊(duì)開發(fā)——Pydantic是許多Python庫和框架中不可或缺的驗(yàn)證層——PydanticAI強(qiáng)調(diào)類型安全,并與mypy和pyright等靜態(tài)類型檢查器無縫集成。
詳細(xì)內(nèi)容可以閱讀上一篇文章:??全新代理構(gòu)建器與框架 - PydanticAI??
PydanticAI的關(guān)鍵特性包括:
?結(jié)構(gòu)化響應(yīng):利用Pydantic驗(yàn)證和結(jié)構(gòu)化模型輸出,確保跨運(yùn)行的一致性。?依賴注入系統(tǒng):提供可選系統(tǒng),為代理的系統(tǒng)提示、工具和結(jié)果驗(yàn)證器提供數(shù)據(jù)和服務(wù),增強(qiáng)測(cè)試和迭代開發(fā)。?流式響應(yīng):支持LLM輸出的持續(xù)流式傳輸,并進(jìn)行即時(shí)驗(yàn)證,以獲得快速且準(zhǔn)確的結(jié)果。
目前PydanticAI處于早期測(cè)試階段,其API可能會(huì)發(fā)生變化,開發(fā)團(tuán)隊(duì)歡迎反饋以優(yōu)化和增強(qiáng)其功能。
環(huán)境設(shè)置
在安裝pydantic-ai之前,請(qǐng)確認(rèn)你的Python版本為3.9或更高:
python --version
然后,創(chuàng)建并激活虛擬環(huán)境,接著安裝pydantic-ai:
virtualenv skoloenv
source skoloenv/bin/activate
pip install pydantic-ai
理解PydanticAI的核心
在PydanticAI中,核心工作組件是??Agent?
?類。通過它,你可以在多種模型上運(yùn)行查詢。你可以在官方文檔中查看支持的模型完整列表。以下是一個(gè)使用OpenAI模型初始化Agent的簡(jiǎn)單代碼片段:
from pydantic_ai importAgent
from pydantic_ai.models.openai importOpenAIModel
# 引用你所需的模型
model =OpenAIModel('gpt-4o', api_key='add-your-api-key-here')
agent =Agent(model)
result = agent.run_sync("什么是比特幣?")
print(result.data)
我們建議將API密鑰存儲(chǔ)為環(huán)境變量:
export OPENAI_API_KEY='your-api-key'
以下是PydanticAI支持的所有可用模型列表:
KnownModelName=Literal[
"openai:gpt-4o",
"openai:gpt-4o-mini",
"openai:gpt-4-turbo",
"openai:gpt-4",
"openai:o1-preview",
"openai:o1-mini",
"openai:o1",
"openai:gpt-3.5-turbo",
"groq:llama-3.3-70b-versatile",
"groq:llama-3.1-70b-versatile",
"groq:llama3-groq-70b-8192-tool-use-preview",
"groq:llama3-groq-8b-8192-tool-use-preview",
"groq:llama-3.1-70b-specdec",
"groq:llama-3.1-8b-instant",
"groq:llama-3.2-1b-preview",
"groq:llama-3.2-3b-preview",
"groq:llama-3.2-11b-vision-preview",
"groq:llama-3.2-90b-vision-preview",
"groq:llama3-70b-8192",
"groq:llama3-8b-8192",
"groq:mixtral-8x7b-32768",
"groq:gemma2-9b-it",
"groq:gemma-7b-it",
"gemini-1.5-flash",
"gemini-1.5-pro",
"gemini-2.0-flash-exp",
"vertexai:gemini-1.5-flash",
"vertexai:gemini-1.5-pro",
"mistral:mistral-small-latest",
"mistral:mistral-large-latest",
"mistral:codestral-latest",
"mistral:mistral-moderation-latest",
"ollama:codellama",
"ollama:gemma",
"ollama:gemma2",
"ollama:llama3",
"ollama:llama3.1",
"ollama:llama3.2",
"ollama:llama3.2-vision",
"ollama:llama3.3",
"ollama:mistral",
"ollama:mistral-nemo",
"ollama:mixtral",
"ollama:phi3",
"ollama:qwq",
"ollama:qwen",
"ollama:qwen2",
"ollama:qwen2.5",
"ollama:starcoder2",
"claude-3-5-haiku-latest",
"claude-3-5-sonnet-latest",
"claude-3-opus-latest",
"test",
]
使用外部工具擴(kuò)展PydanticAI
我們將使用PostgreSQL數(shù)據(jù)庫并構(gòu)建一個(gè)數(shù)據(jù)庫連接類,作為PydanticAI代理的依賴項(xiàng),使代理能夠執(zhí)行數(shù)據(jù)庫功能。
PostgreSQL設(shè)置
首先,你需要一個(gè)正常運(yùn)行且干凈的PostgreSQL數(shù)據(jù)庫。在上文鏈接的視頻中,我展示了如何在DigitalOcean上設(shè)置一個(gè)PostgreSQL數(shù)據(jù)庫(注意:你可以使用任何PostgreSQL數(shù)據(jù)庫)。
你可以使用以下DigitalOcean的推薦鏈接開始:https://m.do.co/c/7d9a2c75356d
一旦你有了數(shù)據(jù)庫,獲取連接字符串——你將在下一步中使用它。
然后安裝以下內(nèi)容:
pip install psycopg2
pip install asyncpg
我們將創(chuàng)建幾個(gè)Python函數(shù)來設(shè)置一個(gè)“notes”表并檢查其是否存在:
import psycopg2
DB_DSN ="database-connection-string"
def create_notes_table():
"""
如果“notes”表不存在,則創(chuàng)建它,包含“id”、“title”和“text”字段。
"""
create_table_query ="""
CREATE TABLE IF NOT EXISTS notes (
id SERIAL PRIMARY KEY,
title VARCHAR(200) UNIQUE NOT NULL,
text TEXT NOT NULL
);
"""
try:
connection = psycopg2.connect(DB_DSN)
cursor = connection.cursor()
cursor.execute(create_table_query)
connection.commit()
print("成功創(chuàng)建或驗(yàn)證了‘notes’表。")
except psycopg2.Erroras e:
print(f"創(chuàng)建表時(shí)出錯(cuò):{e}")
finally:
if connection:
cursor.close()
connection.close()
def check_table_exists(table_name: str)->bool:
"""
檢查指定表是否在數(shù)據(jù)庫中存在。
"""
query ="""
SELECT EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema ='public'
AND table_name =%s
);
"""
try:
connection = psycopg2.connect(DB_DSN)
cursor = connection.cursor()
cursor.execute(query,(table_name,))
exists = cursor.fetchone()[0]
return exists
except psycopg2.Erroras e:
print(f"檢查表時(shí)出錯(cuò):{e}")
returnFalse
finally:
if connection:
cursor.close()
connection.close()
確保你可以運(yùn)行??check_table_exists("notes")?
?函數(shù)并得到“True”響應(yīng)。這表明你的數(shù)據(jù)庫連接正常工作,并且“notes”表已成功創(chuàng)建。
接下來,我們將引入一個(gè)異步類來管理筆記操作,例如添加筆記、檢索筆記和列出標(biāo)題。這個(gè)類將由“Agent”使用。
import asyncpg
from typing importOptional,List
classDatabaseConn:
def __init__(self):
"""
存儲(chǔ)用于連接的DSN(數(shù)據(jù)源名稱)。
"""
self.dsn = DB_DSN
async def _connect(self):
"""
打開與PostgreSQL的異步連接。
"""
return await asyncpg.connect(self.dsn)
async def add_note(self, title: str, text: str)->bool:
"""
插入具有給定標(biāo)題和文本的筆記。
如果存在同名筆記,則不會(huì)覆蓋。
"""
query ="""
INSERT INTO notes (title, text)
VALUES ($1, $2)
ON CONFLICT (title) DO NOTHING;
"""
conn = await self._connect()
try:
result = await conn.execute(query, title, text)
return result =="INSERT 0 1"
finally:
await conn.close()
async def get_note_by_title(self, title: str)->Optional[dict]:
"""
檢索與指定標(biāo)題匹配的筆記。返回字典或None。
"""
query ="SELECT title, text FROM notes WHERE title = $1;"
conn = await self._connect()
try:
record = await conn.fetchrow(query, title)
if record:
return{"title": record["title"],"text": record["text"]}
returnNone
finally:
await conn.close()
async def list_all_titles(self)->List[str]:
"""
獲取并返回所有筆記標(biāo)題。
"""
query ="SELECT title FROM notes ORDER BY title;"
conn = await self._connect()
try:
results = await conn.fetch(query)
return[row["title"]for row in results]
finally:
await conn.close()
將筆記與PydanticAI集成
為了將這些組件整合在一起,我們將創(chuàng)建兩個(gè)不同的代理:
?意圖提取代理:判斷用戶是想創(chuàng)建、列出還是檢索筆記。?動(dòng)作處理代理:使用我們的數(shù)據(jù)庫代碼實(shí)際處理數(shù)據(jù)。
以下是??main.py?
?的示例結(jié)構(gòu):
from dataclasses import dataclass
from pydantic importBaseModel
from pydantic_ai importAgent,RunContext
from typing importOptional,List
from database importDatabaseConn
from pydantic_ai.models.openai importOpenAIModel
OPENAI_API_KEY ="enter-your-openai-api-key-here"
@dataclass
classNoteIntent:
action: str
title:Optional[str]=None
text:Optional[str]=None
@dataclass
classNoteDependencies:
db:DatabaseConn
classNoteResponse(BaseModel):
message: str
note:Optional[dict]=None
titles:Optional[List[str]]=None
# 1. 用于解析用戶意圖的代理
intent_model =OpenAIModel('gpt-4o-mini', api_key=OPENAI_API_KEY)
intent_agent =Agent(
intent_model,
result_type=NoteIntent,
system_prompt=(
"你是一個(gè)意圖提取助手。理解用戶想要做什么(例如創(chuàng)建、檢索、列出),并提取相關(guān)數(shù)據(jù),如標(biāo)題和文本。"
"你的輸出格式必須是JSON-like結(jié)構(gòu),包含鍵:action、title、text。"
)
)
# 2. 用于執(zhí)行識(shí)別動(dòng)作的代理
action_model =OpenAIModel('gpt-4o-mini', api_key=OPENAI_API_KEY)
action_agent =Agent(
action_model,
deps_type=NoteDependencies,
result_type=NoteResponse,
system_prompt=(
"根據(jù)識(shí)別的用戶意圖,在筆記存儲(chǔ)上執(zhí)行請(qǐng)求的操作。"
"操作可以包括:‘create’(添加筆記)、‘retrieve’(獲取筆記)或‘list’(列出所有筆記)。"
)
)
# action_agent的工具
@action_agent.tool
async def create_note_tool(ctx:RunContext[NoteDependencies], title: str, text: str)->NoteResponse:
db = ctx.deps.db
success = await db.add_note(title, text)
returnNoteResponse(message="CREATED:SUCCESS"if success else"CREATED:FAILED")
@action_agent.tool
async def retrieve_note_tool(ctx:RunContext[NoteDependencies], title: str)->NoteResponse:
db = ctx.deps.db
note = await db.get_note_by_title(title)
returnNoteResponse(message="GET:SUCCESS", note=note)if note elseNoteResponse(message="GET:FAILED")
@action_agent.tool
async def list_notes_tool(ctx:RunContext[NoteDependencies])->NoteResponse:
db = ctx.deps.db
all_titles = await db.list_all_titles()
returnNoteResponse(message="LIST:SUCCESS", titles=all_titles)
async def handle_user_query(user_input: str, deps:NoteDependencies)->NoteResponse:
# 確定用戶意圖
intent = await intent_agent.run(user_input)
print(intent.data)
if intent.data.action =="create":
query = f"創(chuàng)建名為‘{intent.data.title}’的筆記,文本為‘{intent.data.text}’。"
response = await action_agent.run(query, deps=deps)
return response.data
elif intent.data.action =="retrieve":
query = f"檢索標(biāo)題為‘{intent.data.title}’的筆記。"
response = await action_agent.run(query, deps=deps)
return response.data
elif intent.data.action =="list":
query ="列出所有筆記的標(biāo)題。"
response = await action_agent.run(query, deps=deps)
return response.data
else:
returnNoteResponse(message="無法識(shí)別的操作。")
async def ask(query: str):
db_conn =DatabaseConn()
note_deps =NoteDependencies(db=db_conn)
return await handle_user_query(query, note_deps)
構(gòu)建Streamlit前端
最后一步是通過一個(gè)簡(jiǎn)單的Web界面使所有功能可訪問。安裝Streamlit非常簡(jiǎn)單:
pip install streamlit
然后創(chuàng)建一個(gè)??app.py?
?文件:
import asyncio
import streamlit as st
from main import ask # 從main.py導(dǎo)入ask函數(shù)
st.set_page_config(page_title="筆記管理器", layout="centered")
st.title("我的筆記儀表板")
st.write("在下方輸入指令以創(chuàng)建、檢索或列出筆記。")
user_input = st.text_area("你想做什么?", placeholder="例如,‘創(chuàng)建一篇關(guān)于我周一會(huì)議的筆記。’")
if st.button("提交"):
ifnot user_input.strip():
st.error("請(qǐng)輸入內(nèi)容。")
else:
with st.spinner("正在處理..."):
try:
response = asyncio.run(ask(user_input))
if response.note isnotNone:
st.success(response.message)
st.subheader(f"筆記標(biāo)題:{response.note.get('title', '')}")
st.write(response.note.get('text','未找到內(nèi)容。'))
elif response.titles isnotNone:
st.success(response.message)
if response.titles:
st.subheader("當(dāng)前標(biāo)題:")
for t in response.titles:
st.write(f"- {t}")
else:
st.info("尚無可用筆記。")
else:
st.info(response.message)
exceptExceptionas e:
st.error(f"錯(cuò)誤:{e}")
然后可以通過以下命令啟動(dòng):
streamlit run app.py
總結(jié)
通過一點(diǎn)努力,我們構(gòu)建了一個(gè)強(qiáng)大的筆記管理工具,使用了以下技術(shù):
?PydanticAI:用于解析用戶請(qǐng)求和結(jié)構(gòu)化數(shù)據(jù)?PostgreSQL:用于存儲(chǔ)筆記?Streamlit:提供一個(gè)流暢、交互式的Web界面
本文轉(zhuǎn)載自??????PyTorch研習(xí)社???,作者:AI研究生
