ChromaDB+OpenAI Embeddings处理DataFrame整数列及条件查询问题求助

问题根源分析 #

1. 整数列无法获取有效响应：原代码仅将name列设为page_content_column，乘客的其他属性（如sibsp这类整数列）未被纳入向量嵌入的内容，向量库无法感知这些数值信息，自然无法针对这类属性给出有效响应。
2. 条件查询仅返回单个结果：similarity_search默认返回语义最相似的少量结果（默认4个），且你只取了第一个结果；同时纯向量检索是语义匹配，不是精确的数值条件筛选，无法直接完成“超过3个兄弟姐妹”这类精确条件的查询。

具体解决步骤 #

1. 调整DataFrame加载逻辑，整合所有关键字段到文档内容 #

将乘客的多列信息（姓名、兄弟姐妹数量、存活状态等）拼接成一段完整文本作为每个文档的内容，确保向量嵌入时包含所有关键属性。

2. 优化条件查询的实现方式 #

提供两种可行方案：

• 方案一：向量检索后手动过滤：先获取足够多的相似结果，再通过代码筛选符合数值条件的条目；
• 方案二：结合大模型处理查询：用LangChain的PandasAgent让大模型直接解析自然语言查询，从DataFrame中提取符合要求的结果，同时返回数量和全部姓名。

修改后的代码示例 #

方案一：向量检索后过滤结果 #

import pandas as pd
from langchain.document_loaders import DataFrameLoader
from langchain.vectorstores import Chroma
from langchain.embeddings.openai import OpenAIEmbeddings

titanic_data = pd.read_csv("/content/titanic.csv")
MAX_PASSENGERS = 1000
subset_titanic = titanic_data.head(MAX_PASSENGERS)

# 为每个乘客生成包含所有关键属性的文本
subset_titanic["combined_info"] = subset_titanic.apply(
    lambda row: f"姓名: {row['name']}, 兄弟姐妹/配偶数量: {row['sibsp']}, 是否存活: {row['survived']}, 舱位等级: {row['pclass']}",
    axis=1
)

# 用整合后的列加载文档
df_loader = DataFrameLoader(subset_titanic, page_content_column="combined_info")
df_documents = df_loader.load()

# 创建向量库
db = Chroma.from_documents(df_documents, OpenAIEmbeddings(openai_api_key='TOKEN'))

# 查询并筛选结果
query = "灾难中有超过3个兄弟姐妹的人员"
# 获取足够多的相似结果，确保覆盖所有符合条件的条目
docs = db.similarity_search(query, k=50)

# 筛选符合条件的乘客
qualified_passengers = []
for doc in docs:
    # 提取兄弟姐妹数量
    sibsp_part = [p for p in doc.page_content.split(", ") if "兄弟姐妹/配偶数量" in p][0]
    sibsp = int(sibsp_part.split(": ")[1])
    if sibsp > 3:
        name = [p for p in doc.page_content.split(", ") if "姓名" in p][0].split(": ")[1]
        qualified_passengers.append(name)

# 去重
qualified_passengers = list(set(qualified_passengers))

print(f"符合条件的人员数量: {len(qualified_passengers)}")
print("符合条件的人员姓名:")
for name in qualified_passengers:
    print(f"- {name}")

方案二：结合大模型处理自然语言查询 #

import pandas as pd
from langchain.agents import create_pandas_dataframe_agent
from langchain.llms import OpenAI

titanic_data = pd.read_csv("/content/titanic.csv")
MAX_PASSENGERS = 1000
subset_titanic = titanic_data.head(MAX_PASSENGERS)

# 创建Pandas Agent，让大模型直接操作DataFrame
agent = create_pandas_dataframe_agent(
    OpenAI(temperature=0, openai_api_key='TOKEN'),
    subset_titanic,
    verbose=True
)

# 执行查询
query = "找出灾难中有超过3个兄弟姐妹的人员，告诉我数量和全部姓名"
agent.run(query)

方案说明 #

• 方案一适合需要精确控制过滤逻辑的场景，无需依赖大模型的推理能力；
• 方案二适合复杂的自然语言查询，大模型会自动解析条件并提取结果，灵活性更强。

内容的提问来源于stack exchange，提问作者Aravind Arjun