Свяжите одного агента и два инструмента в индексе ламы с помощью pgvector ⇐ Python

[Anonymous]

У меня есть следующий код, который позволяет мне задавать вопросы о конкретных инцидентах в моем векторном хранилище, а также об агрегированных показателях из таблицы метаданных SQL. Мне бы хотелось объединить оба инструмента, чтобы, например, мой запрос:
Какова причина инцидента, который произошел в Париже в 2021 году и который подразумевал наводнение Сены,
Агент
1/ запросит с помощью структурированного_запроса идентификаторы инцидентов с критериями city=Paris и датой = 2021
2/ использует semantic_search для «Наводнения в Сене» для связанных инцидентовИли наоборот. На данный момент этот тип запроса не обрабатывается моделью правильно
import pandas as pd
import psycopg2
from sqlalchemy import make_url, create_engine

from llama_index.agent.openai import OpenAIAgent
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.postgres import PGVectorStore
from llama_index.core import Document, Settings, SQLDatabase, VectorStoreIndex, StorageContext,
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.query_engine import NLSQLTableQueryEngine

embed_model = OpenAIEmbedding(
model="text-embedding-3-small",
dimensions=1536 ,
embed_batch_size=20
)
llm = OpenAI(model="gpt-4o")
Settings.embed_model = embed_model
Settings.llm = llm

df = pd.read_pickle('all_incidents.pkl')

def sql_connection():
connection_string = "postgresql://postgres@localhost:5432/incidents"
conn = psycopg2.connect(connection_string)
conn.autocommit = True
url = make_url(connection_string)
return url

# Create a SQL database
engine = create_engine(sql_connection())
df.to_sql('incidents', con=engine, if_exists='replace', index='public')

# Create the SQLStructStoreIndex
sql_database = SQLDatabase(engine, include_tables=["incidents"])

# Create documents for the VectorStoreIndex
documents = [
Document(
text=row['text'],
metadata={
"title": row['title'],
"city": row['city'],
"introduction": row['introduction'],
'level': row['level'],
"type": row['type'],
"date": str(row['date']),
},
)
for _, row in df.iterrows()
]

def open_vector_store():
url = sql_connection()
db_name = "incidents"
vector_store = PGVectorStore.from_params(
database=db_name,
host=url.host,
password=url.password,
port=url.port,
user=url.username,
table_name="incidents_vectors",
embed_dim=1536,
hnsw_kwargs={
"hnsw_m": 16,
"hnsw_ef_construction": 64,
"hnsw_ef_search": 40,
"hnsw_dist_method": "vector_cosine_ops",
},
)
return vector_store

vector_store = open_vector_store()
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
documents, strage_context=storage_context, show_progress=True
)
vector_query_engine = index.as_query_engine()

sql_query_engine = NLSQLTableQueryEngine(
sql_database=sql_database, tables=["incidents"], llm=llm
)

query_engine_tools = [
QueryEngineTool(
query_engine=vector_query_engine,
metadata=ToolMetadata(
name= "semantic_search",
description= "Returns informations about specific incidents"
)
),
QueryEngineTool(
query_engine=sql_query_engine,
metadata=ToolMetadata(
name= "structured_query",
description= "Allows queries for synthesis on the number of incidents, their distribution etc"
),
),
]

agent = OpenAIAgent.from_tools(query_engine_tools, verbose=True)
agent.chat_repl()



Подробнее здесь: https://stackoverflow.com/questions/791 ... h-pgvector