test: 补充知识处理、检索与 Agent 回归用例
覆盖切块、向量检索、工具调用和历史预算；模型相关用例使用桩。
This commit is contained in:
WangShen007 committed
1 parent
2dd062787d
commit
d9fe727c3f
123 files changed
+29005
No files matched your search
@@ -0,0 +1,262 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from app.repositories.knowledge import KnowledgeRepository
|
||||
from app.models.knowledge import Knowledge, SourceType
|
||||
from app.models.knowledge_to_agent import KnowledgeToAgent
|
||||
from app.models.agent import Agent, AgentType
|
||||
from uuid import uuid4
|
||||
|
||||
@pytest.fixture
|
||||
def knowledge_repo(db):
|
||||
return KnowledgeRepository(db)
|
||||
|
||||
@pytest.fixture
|
||||
def test_knowledge(db, test_organization_id):
|
||||
"""Create a test knowledge source"""
|
||||
knowledge = Knowledge(
|
||||
organization_id=test_organization_id,
|
||||
source="/test/path/document.pdf",
|
||||
source_type=SourceType.FILE,
|
||||
schema="test_schema",
|
||||
table_name="test_table"
|
||||
)
|
||||
db.add(knowledge)
|
||||
db.commit()
|
||||
db.refresh(knowledge)
|
||||
return knowledge
|
||||
|
||||
@pytest.fixture
|
||||
def test_agent(db, test_organization_id):
|
||||
"""Create a test agent"""
|
||||
agent = Agent(
|
||||
name="Test Agent",
|
||||
display_name="Test Agent",
|
||||
description="A test agent",
|
||||
agent_type=AgentType.GENERAL,
|
||||
instructions=["Test instruction"],
|
||||
organization_id=test_organization_id,
|
||||
is_active=True
|
||||
)
|
||||
db.add(agent)
|
||||
db.commit()
|
||||
db.refresh(agent)
|
||||
return agent
|
||||
|
||||
def test_create_knowledge(knowledge_repo, test_organization_id):
|
||||
"""Test creating a new knowledge source"""
|
||||
knowledge = Knowledge(
|
||||
organization_id=test_organization_id,
|
||||
source="/test/path/new.pdf",
|
||||
source_type=SourceType.FILE,
|
||||
schema="test_schema",
|
||||
table_name="test_table"
|
||||
)
|
||||
|
||||
created = knowledge_repo.create(knowledge)
|
||||
assert created is not None
|
||||
assert created.source == "/test/path/new.pdf"
|
||||
assert created.source_type == SourceType.FILE
|
||||
|
||||
def test_create_duplicate_knowledge(knowledge_repo, test_knowledge):
|
||||
"""Test creating a duplicate knowledge source returns existing one"""
|
||||
duplicate = Knowledge(
|
||||
organization_id=test_knowledge.organization_id,
|
||||
source=test_knowledge.source,
|
||||
source_type=test_knowledge.source_type,
|
||||
schema=test_knowledge.schema,
|
||||
table_name=test_knowledge.table_name
|
||||
)
|
||||
|
||||
created = knowledge_repo.create(duplicate)
|
||||
assert created.id == test_knowledge.id
|
||||
|
||||
def test_get_by_id(knowledge_repo, test_knowledge):
|
||||
"""Test retrieving a knowledge source by ID"""
|
||||
knowledge = knowledge_repo.get_by_id(test_knowledge.id)
|
||||
assert knowledge is not None
|
||||
assert knowledge.id == test_knowledge.id
|
||||
assert knowledge.source == test_knowledge.source
|
||||
|
||||
def test_get_by_id_nonexistent(knowledge_repo):
|
||||
"""Test retrieving a nonexistent knowledge source"""
|
||||
knowledge = knowledge_repo.get_by_id(999)
|
||||
assert knowledge is None
|
||||
|
||||
def test_get_by_agent(knowledge_repo, test_knowledge, test_agent, db):
|
||||
"""Test retrieving knowledge sources for an agent"""
|
||||
# Create a link to the test knowledge
|
||||
link = KnowledgeToAgent(
|
||||
knowledge_id=test_knowledge.id,
|
||||
agent_id=test_agent.id
|
||||
)
|
||||
db.add(link)
|
||||
|
||||
# Create another knowledge source and link
|
||||
another_knowledge = Knowledge(
|
||||
organization_id=test_knowledge.organization_id,
|
||||
source="/test/path/another.pdf",
|
||||
source_type=SourceType.FILE,
|
||||
schema="test_schema",
|
||||
table_name="test_table"
|
||||
)
|
||||
db.add(another_knowledge)
|
||||
db.flush()
|
||||
|
||||
another_link = KnowledgeToAgent(
|
||||
knowledge_id=another_knowledge.id,
|
||||
agent_id=test_agent.id
|
||||
)
|
||||
db.add(another_link)
|
||||
db.commit()
|
||||
|
||||
# Test pagination
|
||||
knowledge_sources = knowledge_repo.get_by_agent(test_agent.id, skip=0, limit=1)
|
||||
assert len(knowledge_sources) == 1
|
||||
|
||||
knowledge_sources = knowledge_repo.get_by_agent(test_agent.id, skip=0, limit=10)
|
||||
assert len(knowledge_sources) == 2
|
||||
assert all(k.id in [test_knowledge.id, another_knowledge.id] for k in knowledge_sources)
|
||||
|
||||
def test_count_by_agent(knowledge_repo, test_knowledge, test_agent, db):
|
||||
"""Test counting knowledge sources for an agent"""
|
||||
# Create a link to the test knowledge
|
||||
link = KnowledgeToAgent(
|
||||
knowledge_id=test_knowledge.id,
|
||||
agent_id=test_agent.id
|
||||
)
|
||||
db.add(link)
|
||||
db.commit()
|
||||
|
||||
count = knowledge_repo.count_by_agent(test_agent.id)
|
||||
assert count == 1
|
||||
|
||||
def test_get_by_org(knowledge_repo, test_knowledge, db):
|
||||
"""Test retrieving knowledge sources for an organization"""
|
||||
# Create another knowledge source for the same org
|
||||
another_knowledge = Knowledge(
|
||||
organization_id=test_knowledge.organization_id,
|
||||
source="/test/path/another.pdf",
|
||||
source_type=SourceType.FILE,
|
||||
schema="test_schema",
|
||||
table_name="test_table"
|
||||
)
|
||||
db.add(another_knowledge)
|
||||
db.commit()
|
||||
|
||||
knowledge_sources = knowledge_repo.get_by_org(test_knowledge.organization_id)
|
||||
assert len(knowledge_sources) == 2
|
||||
assert all(k.organization_id == test_knowledge.organization_id for k in knowledge_sources)
|
||||
|
||||
def test_get_by_organization_paginated(knowledge_repo, test_knowledge, db):
|
||||
"""Test retrieving paginated knowledge sources for an organization"""
|
||||
# Create another knowledge source for the same org
|
||||
another_knowledge = Knowledge(
|
||||
organization_id=test_knowledge.organization_id,
|
||||
source="/test/path/another.pdf",
|
||||
source_type=SourceType.FILE,
|
||||
schema="test_schema",
|
||||
table_name="test_table"
|
||||
)
|
||||
db.add(another_knowledge)
|
||||
db.commit()
|
||||
|
||||
# Test pagination
|
||||
knowledge_sources = knowledge_repo.get_by_organization(
|
||||
test_knowledge.organization_id,
|
||||
skip=0,
|
||||
limit=1
|
||||
)
|
||||
assert len(knowledge_sources) == 1
|
||||
|
||||
knowledge_sources = knowledge_repo.get_by_organization(
|
||||
test_knowledge.organization_id,
|
||||
skip=0,
|
||||
limit=10
|
||||
)
|
||||
assert len(knowledge_sources) == 2
|
||||
|
||||
def test_count_by_organization(knowledge_repo, test_knowledge, db):
|
||||
"""Test counting knowledge sources for an organization"""
|
||||
# Create another knowledge source for the same org
|
||||
another_knowledge = Knowledge(
|
||||
organization_id=test_knowledge.organization_id,
|
||||
source="/test/path/another.pdf",
|
||||
source_type=SourceType.FILE,
|
||||
schema="test_schema",
|
||||
table_name="test_table"
|
||||
)
|
||||
db.add(another_knowledge)
|
||||
db.commit()
|
||||
|
||||
count = knowledge_repo.count_by_organization(test_knowledge.organization_id)
|
||||
assert count == 2
|
||||
|
||||
def test_get_by_sources(knowledge_repo, test_knowledge, db):
|
||||
"""Test retrieving knowledge sources by source URLs"""
|
||||
# Create another knowledge source
|
||||
another_knowledge = Knowledge(
|
||||
organization_id=test_knowledge.organization_id,
|
||||
source="/test/path/another.pdf",
|
||||
source_type=SourceType.FILE,
|
||||
schema="test_schema",
|
||||
table_name="test_table"
|
||||
)
|
||||
db.add(another_knowledge)
|
||||
db.commit()
|
||||
|
||||
sources = [test_knowledge.source, another_knowledge.source]
|
||||
knowledge_sources = knowledge_repo.get_by_sources(test_knowledge.organization_id, sources)
|
||||
assert len(knowledge_sources) == 2
|
||||
assert all(k.source in sources for k in knowledge_sources)
|
||||
|
||||
def test_delete(knowledge_repo, test_knowledge):
|
||||
"""Test deleting a knowledge source"""
|
||||
success = knowledge_repo.delete(test_knowledge.id)
|
||||
assert success is True
|
||||
|
||||
# Verify deletion
|
||||
knowledge = knowledge_repo.get_by_id(test_knowledge.id)
|
||||
assert knowledge is None
|
||||
|
||||
def test_delete_nonexistent(knowledge_repo):
|
||||
"""Test deleting a nonexistent knowledge source"""
|
||||
success = knowledge_repo.delete(999)
|
||||
assert success is False
|
||||
|
||||
def test_delete_with_data(knowledge_repo, test_knowledge, test_agent, db):
|
||||
"""Test deleting a knowledge source with its data"""
|
||||
# Create a link to test cascade deletion
|
||||
link = KnowledgeToAgent(
|
||||
knowledge_id=test_knowledge.id,
|
||||
agent_id=test_agent.id
|
||||
)
|
||||
db.add(link)
|
||||
db.commit()
|
||||
|
||||
success = knowledge_repo.delete_with_data(test_knowledge.id)
|
||||
assert success is True
|
||||
|
||||
# Verify knowledge source deletion
|
||||
knowledge = knowledge_repo.get_by_id(test_knowledge.id)
|
||||
assert knowledge is None
|
||||
|
||||
# Verify link deletion
|
||||
link = db.query(KnowledgeToAgent).filter(
|
||||
KnowledgeToAgent.knowledge_id == test_knowledge.id
|
||||
).first()
|
||||
assert link is None
|
||||
Reference in new issue
Block a user