feat: add distributed async dataset with unix socket server and refactor agent communication bus
Implement AsyncDataSet class supporting client-server model over Unix sockets with SQLite backend, including KV store, table management, and concurrent query handling. Rename `get_messages` to `receive_messages` in AgentCommunicationBus and update all callers. Remove deprecated `get_recommended_agent` function from agent_roles, `invalidate_tool` from tool_cache, and legacy `receive_messages` wrapper. Add multiplexer command routing in handlers with `/prompt` command support. Introduce comprehensive help documentation system for workflows. Update default API URLs to production endpoints and refactor adaptive context window calculation in AdvancedContextManager.
This commit is contained in:
@@ -1,10 +1,8 @@
|
||||
from unittest.mock import mock_open, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
from pr.core.config_loader import (
|
||||
_load_config_file,
|
||||
_parse_value,
|
||||
create_default_config,
|
||||
load_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -36,36 +34,3 @@ def test_parse_value_bool_upper():
|
||||
def test_load_config_file_not_exists(mock_exists):
|
||||
config = _load_config_file("test.ini")
|
||||
assert config == {}
|
||||
|
||||
|
||||
@patch("os.path.exists", return_value=True)
|
||||
@patch("configparser.ConfigParser")
|
||||
def test_load_config_file_exists(mock_parser_class, mock_exists):
|
||||
mock_parser = mock_parser_class.return_value
|
||||
mock_parser.sections.return_value = ["api"]
|
||||
mock_parser.items.return_value = [("key", "value")]
|
||||
config = _load_config_file("test.ini")
|
||||
assert "api" in config
|
||||
assert config["api"]["key"] == "value"
|
||||
|
||||
|
||||
@patch("pr.core.config_loader._load_config_file")
|
||||
def test_load_config(mock_load):
|
||||
mock_load.side_effect = [{"api": {"key": "global"}}, {"api": {"key": "local"}}]
|
||||
config = load_config()
|
||||
assert config["api"]["key"] == "local"
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
def test_create_default_config(mock_file):
|
||||
result = create_default_config("test.ini")
|
||||
assert result == True
|
||||
mock_file.assert_called_once_with("test.ini", "w")
|
||||
handle = mock_file()
|
||||
handle.write.assert_called_once()
|
||||
|
||||
|
||||
@patch("builtins.open", side_effect=Exception("error"))
|
||||
def test_create_default_config_error(mock_file):
|
||||
result = create_default_config("test.ini")
|
||||
assert result == False
|
||||
|
||||
@@ -0,0 +1,372 @@
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import os
|
||||
import time
|
||||
|
||||
from pr.memory.conversation_memory import ConversationMemory
|
||||
|
||||
|
||||
class TestConversationMemory:
|
||||
def setup_method(self):
|
||||
"""Set up test database for each test."""
|
||||
self.db_fd, self.db_path = tempfile.mkstemp()
|
||||
self.memory = ConversationMemory(self.db_path)
|
||||
|
||||
def teardown_method(self):
|
||||
"""Clean up test database after each test."""
|
||||
self.memory = None
|
||||
os.close(self.db_fd)
|
||||
os.unlink(self.db_path)
|
||||
|
||||
def test_init(self):
|
||||
"""Test ConversationMemory initialization."""
|
||||
assert self.memory.db_path == self.db_path
|
||||
# Verify tables were created
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
tables = [row[0] for row in cursor.fetchall()]
|
||||
assert "conversation_history" in tables
|
||||
assert "conversation_messages" in tables
|
||||
conn.close()
|
||||
|
||||
def test_create_conversation(self):
|
||||
"""Test creating a new conversation."""
|
||||
conversation_id = "test_conv_123"
|
||||
session_id = "test_session_456"
|
||||
|
||||
self.memory.create_conversation(conversation_id, session_id)
|
||||
|
||||
# Verify conversation was created
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT conversation_id, session_id FROM conversation_history WHERE conversation_id = ?",
|
||||
(conversation_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] == conversation_id
|
||||
assert row[1] == session_id
|
||||
conn.close()
|
||||
|
||||
def test_create_conversation_without_session(self):
|
||||
"""Test creating a conversation without session ID."""
|
||||
conversation_id = "test_conv_no_session"
|
||||
|
||||
self.memory.create_conversation(conversation_id)
|
||||
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT conversation_id, session_id FROM conversation_history WHERE conversation_id = ?",
|
||||
(conversation_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] == conversation_id
|
||||
assert row[1] is None
|
||||
conn.close()
|
||||
|
||||
def test_create_conversation_with_metadata(self):
|
||||
"""Test creating a conversation with metadata."""
|
||||
conversation_id = "test_conv_metadata"
|
||||
metadata = {"topic": "test", "priority": "high"}
|
||||
|
||||
self.memory.create_conversation(conversation_id, metadata=metadata)
|
||||
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT conversation_id, metadata FROM conversation_history WHERE conversation_id = ?",
|
||||
(conversation_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] == conversation_id
|
||||
assert row[1] is not None
|
||||
conn.close()
|
||||
|
||||
def test_add_message(self):
|
||||
"""Test adding a message to a conversation."""
|
||||
conversation_id = "test_conv_msg"
|
||||
message_id = "test_msg_123"
|
||||
role = "user"
|
||||
content = "Hello, world!"
|
||||
|
||||
# Create conversation first
|
||||
self.memory.create_conversation(conversation_id)
|
||||
|
||||
# Add message
|
||||
self.memory.add_message(conversation_id, message_id, role, content)
|
||||
|
||||
# Verify message was added
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT message_id, conversation_id, role, content FROM conversation_messages WHERE message_id = ?",
|
||||
(message_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] == message_id
|
||||
assert row[1] == conversation_id
|
||||
assert row[2] == role
|
||||
assert row[3] == content
|
||||
|
||||
# Verify message count was updated
|
||||
cursor.execute(
|
||||
"SELECT message_count FROM conversation_history WHERE conversation_id = ?",
|
||||
(conversation_id,),
|
||||
)
|
||||
count_row = cursor.fetchone()
|
||||
assert count_row[0] == 1
|
||||
conn.close()
|
||||
|
||||
def test_add_message_with_tool_calls(self):
|
||||
"""Test adding a message with tool calls."""
|
||||
conversation_id = "test_conv_tools"
|
||||
message_id = "test_msg_tools"
|
||||
role = "assistant"
|
||||
content = "I'll help you with that."
|
||||
tool_calls = [{"function": "test_func", "args": {"param": "value"}}]
|
||||
|
||||
self.memory.create_conversation(conversation_id)
|
||||
self.memory.add_message(conversation_id, message_id, role, content, tool_calls=tool_calls)
|
||||
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT tool_calls FROM conversation_messages WHERE message_id = ?", (message_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] is not None
|
||||
conn.close()
|
||||
|
||||
def test_add_message_with_metadata(self):
|
||||
"""Test adding a message with metadata."""
|
||||
conversation_id = "test_conv_meta"
|
||||
message_id = "test_msg_meta"
|
||||
role = "user"
|
||||
content = "Test message"
|
||||
metadata = {"tokens": 5, "model": "gpt-4"}
|
||||
|
||||
self.memory.create_conversation(conversation_id)
|
||||
self.memory.add_message(conversation_id, message_id, role, content, metadata=metadata)
|
||||
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT metadata FROM conversation_messages WHERE message_id = ?", (message_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] is not None
|
||||
conn.close()
|
||||
|
||||
def test_get_conversation_messages(self):
|
||||
"""Test retrieving conversation messages."""
|
||||
conversation_id = "test_conv_get"
|
||||
self.memory.create_conversation(conversation_id)
|
||||
|
||||
# Add multiple messages
|
||||
messages = [
|
||||
("msg1", "user", "Hello"),
|
||||
("msg2", "assistant", "Hi there"),
|
||||
("msg3", "user", "How are you?"),
|
||||
]
|
||||
|
||||
for msg_id, role, content in messages:
|
||||
self.memory.add_message(conversation_id, msg_id, role, content)
|
||||
|
||||
# Retrieve all messages
|
||||
retrieved = self.memory.get_conversation_messages(conversation_id)
|
||||
assert len(retrieved) == 3
|
||||
assert retrieved[0]["message_id"] == "msg1"
|
||||
assert retrieved[1]["message_id"] == "msg2"
|
||||
assert retrieved[2]["message_id"] == "msg3"
|
||||
|
||||
def test_get_conversation_messages_limited(self):
|
||||
"""Test retrieving limited number of conversation messages."""
|
||||
conversation_id = "test_conv_limit"
|
||||
self.memory.create_conversation(conversation_id)
|
||||
|
||||
# Add multiple messages
|
||||
for i in range(5):
|
||||
self.memory.add_message(conversation_id, f"msg{i}", "user", f"Message {i}")
|
||||
|
||||
# Retrieve limited messages
|
||||
retrieved = self.memory.get_conversation_messages(conversation_id, limit=3)
|
||||
assert len(retrieved) == 3
|
||||
# Should return most recent messages first due to DESC order
|
||||
assert retrieved[0]["message_id"] == "msg4"
|
||||
assert retrieved[1]["message_id"] == "msg3"
|
||||
assert retrieved[2]["message_id"] == "msg2"
|
||||
|
||||
def test_update_conversation_summary(self):
|
||||
"""Test updating conversation summary."""
|
||||
conversation_id = "test_conv_summary"
|
||||
self.memory.create_conversation(conversation_id)
|
||||
|
||||
summary = "This is a test conversation summary"
|
||||
topics = ["testing", "memory", "conversation"]
|
||||
|
||||
self.memory.update_conversation_summary(conversation_id, summary, topics)
|
||||
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT summary, topics, ended_at FROM conversation_history WHERE conversation_id = ?",
|
||||
(conversation_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] == summary
|
||||
assert row[1] is not None # topics should be stored
|
||||
assert row[2] is not None # ended_at should be set
|
||||
conn.close()
|
||||
|
||||
def test_search_conversations(self):
|
||||
"""Test searching conversations by content."""
|
||||
# Create conversations with different content
|
||||
conv1 = "conv_search_1"
|
||||
conv2 = "conv_search_2"
|
||||
conv3 = "conv_search_3"
|
||||
|
||||
self.memory.create_conversation(conv1)
|
||||
self.memory.create_conversation(conv2)
|
||||
self.memory.create_conversation(conv3)
|
||||
|
||||
# Add messages with searchable content
|
||||
self.memory.add_message(conv1, "msg1", "user", "Python programming tutorial")
|
||||
self.memory.add_message(conv2, "msg2", "user", "JavaScript development guide")
|
||||
self.memory.add_message(conv3, "msg3", "user", "Database design principles")
|
||||
|
||||
# Search for "programming"
|
||||
results = self.memory.search_conversations("programming")
|
||||
assert len(results) == 1
|
||||
assert results[0]["conversation_id"] == conv1
|
||||
|
||||
# Search for "development"
|
||||
results = self.memory.search_conversations("development")
|
||||
assert len(results) == 1
|
||||
assert results[0]["conversation_id"] == conv2
|
||||
|
||||
def test_get_recent_conversations(self):
|
||||
"""Test getting recent conversations."""
|
||||
# Create conversations at different times
|
||||
conv1 = "conv_recent_1"
|
||||
conv2 = "conv_recent_2"
|
||||
conv3 = "conv_recent_3"
|
||||
|
||||
self.memory.create_conversation(conv1)
|
||||
time.sleep(0.01) # Small delay to ensure different timestamps
|
||||
self.memory.create_conversation(conv2)
|
||||
time.sleep(0.01)
|
||||
self.memory.create_conversation(conv3)
|
||||
|
||||
# Get recent conversations
|
||||
recent = self.memory.get_recent_conversations(limit=2)
|
||||
assert len(recent) == 2
|
||||
# Should be ordered by started_at DESC
|
||||
assert recent[0]["conversation_id"] == conv3
|
||||
assert recent[1]["conversation_id"] == conv2
|
||||
|
||||
def test_get_recent_conversations_by_session(self):
|
||||
"""Test getting recent conversations for a specific session."""
|
||||
session1 = "session_1"
|
||||
session2 = "session_2"
|
||||
|
||||
conv1 = "conv_session_1"
|
||||
conv2 = "conv_session_2"
|
||||
conv3 = "conv_session_3"
|
||||
|
||||
self.memory.create_conversation(conv1, session1)
|
||||
self.memory.create_conversation(conv2, session2)
|
||||
self.memory.create_conversation(conv3, session1)
|
||||
|
||||
# Get conversations for session1
|
||||
session_convs = self.memory.get_recent_conversations(session_id=session1)
|
||||
assert len(session_convs) == 2
|
||||
conversation_ids = [c["conversation_id"] for c in session_convs]
|
||||
assert conv1 in conversation_ids
|
||||
assert conv3 in conversation_ids
|
||||
assert conv2 not in conversation_ids
|
||||
|
||||
def test_delete_conversation(self):
|
||||
"""Test deleting a conversation."""
|
||||
conversation_id = "conv_delete"
|
||||
self.memory.create_conversation(conversation_id)
|
||||
|
||||
# Add some messages
|
||||
self.memory.add_message(conversation_id, "msg1", "user", "Test message")
|
||||
self.memory.add_message(conversation_id, "msg2", "assistant", "Response")
|
||||
|
||||
# Delete conversation
|
||||
result = self.memory.delete_conversation(conversation_id)
|
||||
assert result is True
|
||||
|
||||
# Verify conversation and messages are gone
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM conversation_history WHERE conversation_id = ?",
|
||||
(conversation_id,),
|
||||
)
|
||||
assert cursor.fetchone()[0] == 0
|
||||
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM conversation_messages WHERE conversation_id = ?",
|
||||
(conversation_id,),
|
||||
)
|
||||
assert cursor.fetchone()[0] == 0
|
||||
conn.close()
|
||||
|
||||
def test_delete_nonexistent_conversation(self):
|
||||
"""Test deleting a non-existent conversation."""
|
||||
result = self.memory.delete_conversation("nonexistent")
|
||||
assert result is False
|
||||
|
||||
def test_get_statistics(self):
|
||||
"""Test getting memory statistics."""
|
||||
# Create some conversations and messages
|
||||
for i in range(3):
|
||||
conv_id = f"conv_stats_{i}"
|
||||
self.memory.create_conversation(conv_id)
|
||||
for j in range(2):
|
||||
self.memory.add_message(conv_id, f"msg_{i}_{j}", "user", f"Message {j}")
|
||||
|
||||
stats = self.memory.get_statistics()
|
||||
assert stats["total_conversations"] == 3
|
||||
assert stats["total_messages"] == 6
|
||||
assert stats["average_messages_per_conversation"] == 2.0
|
||||
|
||||
def test_thread_safety(self):
|
||||
"""Test that the memory can handle concurrent access."""
|
||||
import threading
|
||||
import queue
|
||||
|
||||
results = queue.Queue()
|
||||
|
||||
def worker(worker_id):
|
||||
try:
|
||||
conv_id = f"conv_thread_{worker_id}"
|
||||
self.memory.create_conversation(conv_id)
|
||||
self.memory.add_message(conv_id, f"msg_{worker_id}", "user", f"Worker {worker_id}")
|
||||
results.put(True)
|
||||
except Exception as e:
|
||||
results.put(e)
|
||||
|
||||
# Start multiple threads
|
||||
threads = []
|
||||
for i in range(5):
|
||||
t = threading.Thread(target=worker, args=(i,))
|
||||
threads.append(t)
|
||||
t.start()
|
||||
|
||||
# Wait for all threads
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Check results
|
||||
for _ in range(5):
|
||||
result = results.get()
|
||||
assert result is True
|
||||
|
||||
# Verify all conversations were created
|
||||
recent = self.memory.get_recent_conversations(limit=10)
|
||||
assert len(recent) >= 5
|
||||
@@ -0,0 +1,242 @@
|
||||
from pr.memory.fact_extractor import FactExtractor
|
||||
|
||||
|
||||
class TestFactExtractor:
|
||||
def setup_method(self):
|
||||
"""Set up test fixture."""
|
||||
self.extractor = FactExtractor()
|
||||
|
||||
def test_init(self):
|
||||
"""Test FactExtractor initialization."""
|
||||
assert self.extractor.fact_patterns is not None
|
||||
assert len(self.extractor.fact_patterns) > 0
|
||||
|
||||
def test_extract_facts_definition(self):
|
||||
"""Test extracting definition facts."""
|
||||
text = "John Smith is a software engineer. Python is a programming language."
|
||||
facts = self.extractor.extract_facts(text)
|
||||
|
||||
assert len(facts) >= 2
|
||||
# Check for definition pattern matches
|
||||
definition_facts = [f for f in facts if f["type"] == "definition"]
|
||||
assert len(definition_facts) >= 1
|
||||
|
||||
def test_extract_facts_temporal(self):
|
||||
"""Test extracting temporal facts."""
|
||||
text = "John was born in 1990. The company was founded in 2010."
|
||||
facts = self.extractor.extract_facts(text)
|
||||
|
||||
temporal_facts = [f for f in facts if f["type"] == "temporal"]
|
||||
assert len(temporal_facts) >= 1
|
||||
|
||||
def test_extract_facts_attribution(self):
|
||||
"""Test extracting attribution facts."""
|
||||
text = "John invented the widget. Mary developed the software."
|
||||
facts = self.extractor.extract_facts(text)
|
||||
|
||||
attribution_facts = [f for f in facts if f["type"] == "attribution"]
|
||||
assert len(attribution_facts) >= 1
|
||||
|
||||
def test_extract_facts_numeric(self):
|
||||
"""Test extracting numeric facts."""
|
||||
text = "The car costs $25,000. The house is worth $500,000."
|
||||
facts = self.extractor.extract_facts(text)
|
||||
|
||||
numeric_facts = [f for f in facts if f["type"] == "numeric"]
|
||||
assert len(numeric_facts) >= 1
|
||||
|
||||
def test_extract_facts_location(self):
|
||||
"""Test extracting location facts."""
|
||||
text = "John lives in San Francisco. The office is located in New York."
|
||||
facts = self.extractor.extract_facts(text)
|
||||
|
||||
location_facts = [f for f in facts if f["type"] == "location"]
|
||||
assert len(location_facts) >= 1
|
||||
|
||||
def test_extract_facts_entity(self):
|
||||
"""Test extracting entity facts from noun phrases."""
|
||||
text = "John Smith works at Google Inc. He uses Python programming."
|
||||
facts = self.extractor.extract_facts(text)
|
||||
|
||||
entity_facts = [f for f in facts if f["type"] == "entity"]
|
||||
assert len(entity_facts) >= 1
|
||||
|
||||
def test_extract_noun_phrases(self):
|
||||
"""Test noun phrase extraction."""
|
||||
text = "John Smith is a software engineer at Google. He works on Python Projects."
|
||||
phrases = self.extractor._extract_noun_phrases(text)
|
||||
|
||||
assert "John Smith" in phrases
|
||||
assert "Google" in phrases
|
||||
assert "Python Projects" in phrases
|
||||
|
||||
def test_extract_noun_phrases_capitalized(self):
|
||||
"""Test that only capitalized noun phrases are extracted."""
|
||||
text = "the quick brown fox jumps over the lazy dog"
|
||||
phrases = self.extractor._extract_noun_phrases(text)
|
||||
|
||||
# Should be empty since no capitalized words
|
||||
assert len(phrases) == 0
|
||||
|
||||
def test_extract_key_terms(self):
|
||||
"""Test key term extraction."""
|
||||
text = "Python is a programming language used for software development and data analysis."
|
||||
terms = self.extractor.extract_key_terms(text, top_k=5)
|
||||
|
||||
assert len(terms) <= 5
|
||||
# Should contain programming, language, software, development, data, analysis
|
||||
term_words = [term[0] for term in terms]
|
||||
assert "programming" in term_words
|
||||
assert "language" in term_words
|
||||
assert "software" in term_words
|
||||
|
||||
def test_extract_key_terms_stopwords_filtered(self):
|
||||
"""Test that stopwords are filtered from key terms."""
|
||||
text = "This is a test of the system that should work properly."
|
||||
terms = self.extractor.extract_key_terms(text)
|
||||
|
||||
term_words = [term[0] for term in terms]
|
||||
# Stopwords should not appear
|
||||
assert "this" not in term_words
|
||||
assert "is" not in term_words
|
||||
assert "a" not in term_words
|
||||
assert "of" not in term_words
|
||||
assert "the" not in term_words
|
||||
|
||||
def test_extract_relationships_employment(self):
|
||||
"""Test extracting employment relationships."""
|
||||
text = "John works for Google. Mary is employed by Microsoft."
|
||||
relationships = self.extractor.extract_relationships(text)
|
||||
|
||||
employment_rels = [r for r in relationships if r["type"] == "employment"]
|
||||
assert len(employment_rels) >= 1
|
||||
|
||||
def test_extract_relationships_ownership(self):
|
||||
"""Test extracting ownership relationships."""
|
||||
text = "John owns a car. Mary has a house."
|
||||
relationships = self.extractor.extract_relationships(text)
|
||||
|
||||
ownership_rels = [r for r in relationships if r["type"] == "ownership"]
|
||||
assert len(ownership_rels) >= 1
|
||||
|
||||
def test_extract_relationships_location(self):
|
||||
"""Test extracting location relationships."""
|
||||
text = "John located in New York. The factory belongs to Google."
|
||||
relationships = self.extractor.extract_relationships(text)
|
||||
|
||||
location_rels = [r for r in relationships if r["type"] == "location"]
|
||||
assert len(location_rels) >= 1
|
||||
|
||||
def test_extract_relationships_usage(self):
|
||||
"""Test extracting usage relationships."""
|
||||
text = "John uses Python. The company implements agile methodology."
|
||||
relationships = self.extractor.extract_relationships(text)
|
||||
|
||||
usage_rels = [r for r in relationships if r["type"] == "usage"]
|
||||
assert len(usage_rels) >= 1
|
||||
|
||||
def test_extract_metadata(self):
|
||||
"""Test metadata extraction."""
|
||||
text = "This is a test document. It contains some information about Python programming. You can visit https://python.org for more details. Contact john@example.com for questions. The project started in 2020 and costs $10,000."
|
||||
metadata = self.extractor.extract_metadata(text)
|
||||
|
||||
assert metadata["word_count"] > 0
|
||||
assert metadata["sentence_count"] > 0
|
||||
assert metadata["avg_words_per_sentence"] > 0
|
||||
assert len(metadata["urls"]) > 0
|
||||
assert len(metadata["email_addresses"]) > 0
|
||||
assert len(metadata["dates"]) > 0
|
||||
assert len(metadata["numeric_values"]) > 0
|
||||
assert metadata["has_code"] is False # No code in this text
|
||||
|
||||
def test_extract_metadata_with_code(self):
|
||||
"""Test metadata extraction with code content."""
|
||||
text = "Here is a function: def hello(): print('Hello, world!')"
|
||||
metadata = self.extractor.extract_metadata(text)
|
||||
|
||||
assert metadata["has_code"] is True
|
||||
|
||||
def test_extract_metadata_with_questions(self):
|
||||
"""Test metadata extraction with questions."""
|
||||
text = "What is Python? How does it work? Why use it?"
|
||||
metadata = self.extractor.extract_metadata(text)
|
||||
|
||||
assert metadata["has_questions"] is True
|
||||
|
||||
def test_categorize_content_programming(self):
|
||||
"""Test content categorization for programming."""
|
||||
text = "Python is a programming language used for code development and debugging."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "programming" in categories
|
||||
|
||||
def test_categorize_content_data(self):
|
||||
"""Test content categorization for data."""
|
||||
text = "The database contains records and tables with statistical analysis."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "data" in categories
|
||||
|
||||
def test_categorize_content_documentation(self):
|
||||
"""Test content categorization for documentation."""
|
||||
text = "This guide explains how to use the tutorial and manual."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "documentation" in categories
|
||||
|
||||
def test_categorize_content_configuration(self):
|
||||
"""Test content categorization for configuration."""
|
||||
text = "Configure the settings and setup the deployment environment."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "configuration" in categories
|
||||
|
||||
def test_categorize_content_testing(self):
|
||||
"""Test content categorization for testing."""
|
||||
text = "Run the tests to validate the functionality and verify quality."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "testing" in categories
|
||||
|
||||
def test_categorize_content_research(self):
|
||||
"""Test content categorization for research."""
|
||||
text = "The study investigates findings and results from the analysis."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "research" in categories
|
||||
|
||||
def test_categorize_content_planning(self):
|
||||
"""Test content categorization for planning."""
|
||||
text = "Plan the project schedule with milestones and timeline."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "planning" in categories
|
||||
|
||||
def test_categorize_content_general(self):
|
||||
"""Test content categorization defaults to general."""
|
||||
text = "This is some random text without specific keywords."
|
||||
categories = self.extractor.categorize_content(text)
|
||||
|
||||
assert "general" in categories
|
||||
|
||||
def test_extract_facts_empty_text(self):
|
||||
"""Test fact extraction with empty text."""
|
||||
facts = self.extractor.extract_facts("")
|
||||
assert len(facts) == 0
|
||||
|
||||
def test_extract_key_terms_empty_text(self):
|
||||
"""Test key term extraction with empty text."""
|
||||
terms = self.extractor.extract_key_terms("")
|
||||
assert len(terms) == 0
|
||||
|
||||
def test_extract_relationships_empty_text(self):
|
||||
"""Test relationship extraction with empty text."""
|
||||
relationships = self.extractor.extract_relationships("")
|
||||
assert len(relationships) == 0
|
||||
|
||||
def test_extract_metadata_empty_text(self):
|
||||
"""Test metadata extraction with empty text."""
|
||||
metadata = self.extractor.extract_metadata("")
|
||||
assert metadata["word_count"] == 0
|
||||
assert metadata["sentence_count"] == 0
|
||||
assert metadata["avg_words_per_sentence"] == 0.0
|
||||
@@ -0,0 +1,344 @@
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import os
|
||||
import time
|
||||
|
||||
from pr.memory.knowledge_store import KnowledgeStore, KnowledgeEntry
|
||||
|
||||
|
||||
class TestKnowledgeStore:
|
||||
def setup_method(self):
|
||||
"""Set up test database for each test."""
|
||||
self.db_fd, self.db_path = tempfile.mkstemp()
|
||||
self.store = KnowledgeStore(self.db_path)
|
||||
|
||||
def teardown_method(self):
|
||||
"""Clean up test database after each test."""
|
||||
self.store = None
|
||||
os.close(self.db_fd)
|
||||
os.unlink(self.db_path)
|
||||
|
||||
def test_init(self):
|
||||
"""Test KnowledgeStore initialization."""
|
||||
assert self.store.db_path == self.db_path
|
||||
# Verify tables were created
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
tables = [row[0] for row in cursor.fetchall()]
|
||||
assert "knowledge_entries" in tables
|
||||
conn.close()
|
||||
|
||||
def test_add_entry(self):
|
||||
"""Test adding a knowledge entry."""
|
||||
entry = KnowledgeEntry(
|
||||
entry_id="test_1",
|
||||
category="test",
|
||||
content="This is a test entry",
|
||||
metadata={"source": "test"},
|
||||
created_at=time.time(),
|
||||
updated_at=time.time(),
|
||||
)
|
||||
|
||||
self.store.add_entry(entry)
|
||||
|
||||
# Verify entry was added
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT entry_id, category, content FROM knowledge_entries WHERE entry_id = ?",
|
||||
("test_1",),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row[0] == "test_1"
|
||||
assert row[1] == "test"
|
||||
assert row[2] == "This is a test entry"
|
||||
conn.close()
|
||||
|
||||
def test_get_entry(self):
|
||||
"""Test retrieving a knowledge entry."""
|
||||
entry = KnowledgeEntry(
|
||||
entry_id="test_get",
|
||||
category="test",
|
||||
content="Content to retrieve",
|
||||
metadata={},
|
||||
created_at=time.time(),
|
||||
updated_at=time.time(),
|
||||
)
|
||||
|
||||
self.store.add_entry(entry)
|
||||
retrieved = self.store.get_entry("test_get")
|
||||
|
||||
assert retrieved is not None
|
||||
assert retrieved.entry_id == "test_get"
|
||||
assert retrieved.content == "Content to retrieve"
|
||||
assert retrieved.access_count == 1 # Should be incremented
|
||||
|
||||
def test_get_entry_not_found(self):
|
||||
"""Test retrieving a non-existent entry."""
|
||||
retrieved = self.store.get_entry("nonexistent")
|
||||
assert retrieved is None
|
||||
|
||||
def test_search_entries_semantic(self):
|
||||
"""Test semantic search."""
|
||||
entries = [
|
||||
KnowledgeEntry(
|
||||
"entry1", "personal", "John is a software engineer", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"entry2", "personal", "Mary works as a designer", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"entry3", "tech", "Python is a programming language", {}, time.time(), time.time()
|
||||
),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
results = self.store.search_entries("software engineer", top_k=2)
|
||||
assert len(results) >= 1
|
||||
# Should find the most relevant entry
|
||||
found_ids = [r.entry_id for r in results]
|
||||
assert "entry1" in found_ids
|
||||
|
||||
def test_search_entries_fts_exact(self):
|
||||
"""Test full-text search with exact matches."""
|
||||
entries = [
|
||||
KnowledgeEntry(
|
||||
"exact1", "test", "Python programming language", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"exact2", "test", "Java programming language", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"exact3", "test", "Database design principles", {}, time.time(), time.time()
|
||||
),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
results = self.store.search_entries("programming language", top_k=3)
|
||||
assert len(results) >= 2
|
||||
found_ids = [r.entry_id for r in results]
|
||||
assert "exact1" in found_ids
|
||||
assert "exact2" in found_ids
|
||||
|
||||
def test_search_entries_by_category(self):
|
||||
"""Test searching entries by category."""
|
||||
entries = [
|
||||
KnowledgeEntry("cat1", "personal", "John's info", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("cat2", "tech", "Python info", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("cat3", "personal", "Jane's info", {}, time.time(), time.time()),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
results = self.store.search_entries("info", category="personal", top_k=5)
|
||||
assert len(results) == 2
|
||||
found_ids = [r.entry_id for r in results]
|
||||
assert "cat1" in found_ids
|
||||
assert "cat3" in found_ids
|
||||
assert "cat2" not in found_ids
|
||||
|
||||
def test_get_by_category(self):
|
||||
"""Test getting entries by category."""
|
||||
entries = [
|
||||
KnowledgeEntry("get1", "personal", "Entry 1", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("get2", "tech", "Entry 2", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("get3", "personal", "Entry 3", {}, time.time() + 1, time.time() + 1),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
personal_entries = self.store.get_by_category("personal")
|
||||
assert len(personal_entries) == 2
|
||||
# Should be ordered by importance_score DESC, created_at DESC
|
||||
assert personal_entries[0].entry_id == "get3" # More recent
|
||||
assert personal_entries[1].entry_id == "get1"
|
||||
|
||||
def test_update_importance(self):
|
||||
"""Test updating entry importance."""
|
||||
entry = KnowledgeEntry(
|
||||
"importance_test", "test", "Test content", {}, time.time(), time.time()
|
||||
)
|
||||
self.store.add_entry(entry)
|
||||
|
||||
self.store.update_importance("importance_test", 0.8)
|
||||
|
||||
retrieved = self.store.get_entry("importance_test")
|
||||
assert retrieved.importance_score == 0.8
|
||||
|
||||
def test_delete_entry(self):
|
||||
"""Test deleting an entry."""
|
||||
entry = KnowledgeEntry("delete_test", "test", "To be deleted", {}, time.time(), time.time())
|
||||
self.store.add_entry(entry)
|
||||
|
||||
result = self.store.delete_entry("delete_test")
|
||||
assert result is True
|
||||
|
||||
# Verify it's gone
|
||||
retrieved = self.store.get_entry("delete_test")
|
||||
assert retrieved is None
|
||||
|
||||
def test_delete_entry_not_found(self):
|
||||
"""Test deleting a non-existent entry."""
|
||||
result = self.store.delete_entry("nonexistent")
|
||||
assert result is False
|
||||
|
||||
def test_get_statistics(self):
|
||||
"""Test getting store statistics."""
|
||||
entries = [
|
||||
KnowledgeEntry("stat1", "personal", "Personal info", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("stat2", "tech", "Tech info", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("stat3", "personal", "More personal info", {}, time.time(), time.time()),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
stats = self.store.get_statistics()
|
||||
assert stats["total_entries"] == 3
|
||||
assert stats["total_categories"] == 2
|
||||
assert stats["category_distribution"]["personal"] == 2
|
||||
assert stats["category_distribution"]["tech"] == 1
|
||||
assert stats["vocabulary_size"] > 0
|
||||
|
||||
def test_fts_search_exact_phrase(self):
|
||||
"""Test FTS exact phrase matching."""
|
||||
entries = [
|
||||
KnowledgeEntry(
|
||||
"fts1", "test", "Python is great for programming", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"fts2", "test", "Java is also good for programming", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"fts3", "test", "Database management is important", {}, time.time(), time.time()
|
||||
),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
fts_results = self.store._fts_search("programming")
|
||||
assert len(fts_results) == 2
|
||||
entry_ids = [entry_id for entry_id, score in fts_results]
|
||||
assert "fts1" in entry_ids
|
||||
assert "fts2" in entry_ids
|
||||
|
||||
def test_fts_search_partial_match(self):
|
||||
"""Test FTS partial word matching."""
|
||||
entries = [
|
||||
KnowledgeEntry("partial1", "test", "Python programming", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("partial2", "test", "Java programming", {}, time.time(), time.time()),
|
||||
KnowledgeEntry("partial3", "test", "Database design", {}, time.time(), time.time()),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
fts_results = self.store._fts_search("program")
|
||||
assert len(fts_results) >= 2
|
||||
|
||||
def test_combined_search_scoring(self):
|
||||
"""Test that combined semantic + FTS search produces proper scoring."""
|
||||
entries = [
|
||||
KnowledgeEntry(
|
||||
"combined1", "test", "Python programming language", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"combined2", "test", "Java programming language", {}, time.time(), time.time()
|
||||
),
|
||||
KnowledgeEntry(
|
||||
"combined3", "test", "Database management system", {}, time.time(), time.time()
|
||||
),
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.store.add_entry(entry)
|
||||
|
||||
results = self.store.search_entries("programming language")
|
||||
assert len(results) >= 2
|
||||
|
||||
# Check that results have search scores (at least one should have a positive score)
|
||||
has_positive_score = False
|
||||
for result in results:
|
||||
assert "search_score" in result.metadata
|
||||
if result.metadata["search_score"] > 0:
|
||||
has_positive_score = True
|
||||
assert has_positive_score
|
||||
|
||||
def test_empty_search(self):
|
||||
"""Test searching with no matches."""
|
||||
results = self.store.search_entries("nonexistent topic")
|
||||
assert len(results) == 0
|
||||
|
||||
def test_search_empty_query(self):
|
||||
"""Test searching with empty query."""
|
||||
results = self.store.search_entries("")
|
||||
assert len(results) == 0
|
||||
|
||||
def test_thread_safety(self):
|
||||
"""Test that the store can handle concurrent access."""
|
||||
import threading
|
||||
import queue
|
||||
|
||||
results = queue.Queue()
|
||||
|
||||
def worker(worker_id):
|
||||
try:
|
||||
entry = KnowledgeEntry(
|
||||
f"thread_{worker_id}",
|
||||
"test",
|
||||
f"Content from worker {worker_id}",
|
||||
{},
|
||||
time.time(),
|
||||
time.time(),
|
||||
)
|
||||
self.store.add_entry(entry)
|
||||
retrieved = self.store.get_entry(f"thread_{worker_id}")
|
||||
assert retrieved is not None
|
||||
results.put(True)
|
||||
except Exception as e:
|
||||
results.put(e)
|
||||
|
||||
# Start multiple threads
|
||||
threads = []
|
||||
for i in range(5):
|
||||
t = threading.Thread(target=worker, args=(i,))
|
||||
threads.append(t)
|
||||
t.start()
|
||||
|
||||
# Wait for all threads
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Check results
|
||||
for _ in range(5):
|
||||
result = results.get()
|
||||
assert result is True
|
||||
|
||||
def test_entry_to_dict(self):
|
||||
"""Test KnowledgeEntry to_dict method."""
|
||||
entry = KnowledgeEntry(
|
||||
entry_id="dict_test",
|
||||
category="test",
|
||||
content="Test content",
|
||||
metadata={"key": "value"},
|
||||
created_at=1234567890.0,
|
||||
updated_at=1234567891.0,
|
||||
access_count=5,
|
||||
importance_score=0.8,
|
||||
)
|
||||
|
||||
entry_dict = entry.to_dict()
|
||||
assert entry_dict["entry_id"] == "dict_test"
|
||||
assert entry_dict["category"] == "test"
|
||||
assert entry_dict["content"] == "Test content"
|
||||
assert entry_dict["metadata"]["key"] == "value"
|
||||
assert entry_dict["access_count"] == 5
|
||||
assert entry_dict["importance_score"] == 0.8
|
||||
@@ -0,0 +1,280 @@
|
||||
import math
|
||||
import pytest
|
||||
from pr.memory.semantic_index import SemanticIndex
|
||||
|
||||
|
||||
class TestSemanticIndex:
|
||||
def test_init(self):
|
||||
"""Test SemanticIndex initialization."""
|
||||
index = SemanticIndex()
|
||||
assert index.documents == {}
|
||||
assert index.vocabulary == set()
|
||||
assert index.idf_scores == {}
|
||||
assert index.doc_tf_scores == {}
|
||||
|
||||
def test_tokenize_basic(self):
|
||||
"""Test basic tokenization functionality."""
|
||||
index = SemanticIndex()
|
||||
tokens = index._tokenize("Hello, world! This is a test.")
|
||||
expected = ["hello", "world", "this", "is", "a", "test"]
|
||||
assert tokens == expected
|
||||
|
||||
def test_tokenize_special_characters(self):
|
||||
"""Test tokenization with special characters."""
|
||||
index = SemanticIndex()
|
||||
tokens = index._tokenize("Hello@world.com test-case_123")
|
||||
expected = ["hello", "world", "com", "test", "case", "123"]
|
||||
assert tokens == expected
|
||||
|
||||
def test_tokenize_empty_string(self):
|
||||
"""Test tokenization of empty string."""
|
||||
index = SemanticIndex()
|
||||
tokens = index._tokenize("")
|
||||
assert tokens == []
|
||||
|
||||
def test_tokenize_only_special_chars(self):
|
||||
"""Test tokenization with only special characters."""
|
||||
index = SemanticIndex()
|
||||
tokens = index._tokenize("!@#$%^&*()")
|
||||
assert tokens == []
|
||||
|
||||
def test_compute_tf_basic(self):
|
||||
"""Test TF computation for basic case."""
|
||||
index = SemanticIndex()
|
||||
tokens = ["hello", "world", "hello", "test"]
|
||||
tf_scores = index._compute_tf(tokens)
|
||||
expected = {"hello": 2 / 4, "world": 1 / 4, "test": 1 / 4} # 0.5 # 0.25 # 0.25
|
||||
assert tf_scores == expected
|
||||
|
||||
def test_compute_tf_empty(self):
|
||||
"""Test TF computation for empty tokens."""
|
||||
index = SemanticIndex()
|
||||
tf_scores = index._compute_tf([])
|
||||
assert tf_scores == {}
|
||||
|
||||
def test_compute_tf_single_token(self):
|
||||
"""Test TF computation for single token."""
|
||||
index = SemanticIndex()
|
||||
tokens = ["hello"]
|
||||
tf_scores = index._compute_tf(tokens)
|
||||
assert tf_scores == {"hello": 1.0}
|
||||
|
||||
def test_compute_idf_single_document(self):
|
||||
"""Test IDF computation with single document."""
|
||||
index = SemanticIndex()
|
||||
index.documents = {"doc1": "hello world"}
|
||||
index._compute_idf()
|
||||
assert index.idf_scores == {"hello": 1.0, "world": 1.0}
|
||||
|
||||
def test_compute_idf_multiple_documents(self):
|
||||
"""Test IDF computation with multiple documents."""
|
||||
index = SemanticIndex()
|
||||
index.documents = {"doc1": "hello world", "doc2": "hello test", "doc3": "world test"}
|
||||
index._compute_idf()
|
||||
expected = {
|
||||
"hello": math.log(3 / 2), # appears in 2/3 docs
|
||||
"world": math.log(3 / 2), # appears in 2/3 docs
|
||||
"test": math.log(3 / 2), # appears in 2/3 docs
|
||||
}
|
||||
assert index.idf_scores == expected
|
||||
|
||||
def test_compute_idf_empty_documents(self):
|
||||
"""Test IDF computation with no documents."""
|
||||
index = SemanticIndex()
|
||||
index._compute_idf()
|
||||
assert index.idf_scores == {}
|
||||
|
||||
def test_add_document_basic(self):
|
||||
"""Test adding a basic document."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
|
||||
assert "doc1" in index.documents
|
||||
assert index.documents["doc1"] == "hello world"
|
||||
assert "hello" in index.vocabulary
|
||||
assert "world" in index.vocabulary
|
||||
assert "doc1" in index.doc_tf_scores
|
||||
|
||||
def test_add_document_updates_vocabulary(self):
|
||||
"""Test that adding documents updates vocabulary."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
assert index.vocabulary == {"hello", "world"}
|
||||
|
||||
index.add_document("doc2", "hello test")
|
||||
assert index.vocabulary == {"hello", "world", "test"}
|
||||
|
||||
def test_add_document_updates_idf(self):
|
||||
"""Test that adding documents updates IDF scores."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
assert index.idf_scores == {"hello": 1.0, "world": 1.0}
|
||||
|
||||
index.add_document("doc2", "hello test")
|
||||
expected_idf = {
|
||||
"hello": math.log(2 / 2), # appears in both docs
|
||||
"world": math.log(2 / 1), # appears in 1/2 docs
|
||||
"test": math.log(2 / 1), # appears in 1/2 docs
|
||||
}
|
||||
assert index.idf_scores == expected_idf
|
||||
|
||||
def test_add_document_tf_computation(self):
|
||||
"""Test TF score computation when adding document."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world hello")
|
||||
|
||||
# TF: hello=2/3, world=1/3
|
||||
expected_tf = {"hello": 2 / 3, "world": 1 / 3}
|
||||
assert index.doc_tf_scores["doc1"] == expected_tf
|
||||
|
||||
def test_remove_document_existing(self):
|
||||
"""Test removing an existing document."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
index.add_document("doc2", "hello test")
|
||||
|
||||
initial_vocab = index.vocabulary.copy()
|
||||
initial_idf = index.idf_scores.copy()
|
||||
|
||||
index.remove_document("doc1")
|
||||
|
||||
assert "doc1" not in index.documents
|
||||
assert "doc1" not in index.doc_tf_scores
|
||||
# Vocabulary should still contain all words
|
||||
assert index.vocabulary == initial_vocab
|
||||
# IDF should be recomputed
|
||||
assert index.idf_scores != initial_idf
|
||||
assert index.idf_scores == {"hello": 1.0, "test": 1.0}
|
||||
|
||||
def test_remove_document_nonexistent(self):
|
||||
"""Test removing a non-existent document."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
|
||||
initial_state = {
|
||||
"documents": index.documents.copy(),
|
||||
"vocabulary": index.vocabulary.copy(),
|
||||
"idf_scores": index.idf_scores.copy(),
|
||||
"doc_tf_scores": index.doc_tf_scores.copy(),
|
||||
}
|
||||
|
||||
index.remove_document("nonexistent")
|
||||
|
||||
assert index.documents == initial_state["documents"]
|
||||
assert index.vocabulary == initial_state["vocabulary"]
|
||||
assert index.idf_scores == initial_state["idf_scores"]
|
||||
assert index.doc_tf_scores == initial_state["doc_tf_scores"]
|
||||
|
||||
def test_search_basic(self):
|
||||
"""Test basic search functionality."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
index.add_document("doc2", "hello test")
|
||||
index.add_document("doc3", "world test")
|
||||
|
||||
results = index.search("hello", top_k=5)
|
||||
assert len(results) == 3 # All documents are returned with similarity scores
|
||||
|
||||
# Results should be sorted by similarity (descending)
|
||||
scores = {doc_id: score for doc_id, score in results}
|
||||
assert scores["doc1"] > 0 # doc1 contains "hello"
|
||||
assert scores["doc2"] > 0 # doc2 contains "hello"
|
||||
assert scores["doc3"] == 0 # doc3 does not contain "hello"
|
||||
|
||||
def test_search_empty_query(self):
|
||||
"""Test search with empty query."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
|
||||
results = index.search("", top_k=5)
|
||||
assert results == []
|
||||
|
||||
def test_search_no_documents(self):
|
||||
"""Test search when no documents exist."""
|
||||
index = SemanticIndex()
|
||||
results = index.search("hello", top_k=5)
|
||||
assert results == []
|
||||
|
||||
def test_search_top_k_limit(self):
|
||||
"""Test search respects top_k parameter."""
|
||||
index = SemanticIndex()
|
||||
for i in range(10):
|
||||
index.add_document(f"doc{i}", f"hello world content")
|
||||
|
||||
results = index.search("content", top_k=3)
|
||||
assert len(results) == 3
|
||||
|
||||
def test_cosine_similarity_identical_vectors(self):
|
||||
"""Test cosine similarity with identical vectors."""
|
||||
index = SemanticIndex()
|
||||
vec1 = {"hello": 1.0, "world": 0.5}
|
||||
vec2 = {"hello": 1.0, "world": 0.5}
|
||||
similarity = index._cosine_similarity(vec1, vec2)
|
||||
assert similarity == pytest.approx(1.0)
|
||||
|
||||
def test_cosine_similarity_orthogonal_vectors(self):
|
||||
"""Test cosine similarity with orthogonal vectors."""
|
||||
index = SemanticIndex()
|
||||
vec1 = {"hello": 1.0}
|
||||
vec2 = {"world": 1.0}
|
||||
similarity = index._cosine_similarity(vec1, vec2)
|
||||
assert similarity == 0.0
|
||||
|
||||
def test_cosine_similarity_zero_vector(self):
|
||||
"""Test cosine similarity with zero vector."""
|
||||
index = SemanticIndex()
|
||||
vec1 = {"hello": 1.0}
|
||||
vec2 = {}
|
||||
similarity = index._cosine_similarity(vec1, vec2)
|
||||
assert similarity == 0.0
|
||||
|
||||
def test_cosine_similarity_empty_vectors(self):
|
||||
"""Test cosine similarity with empty vectors."""
|
||||
index = SemanticIndex()
|
||||
similarity = index._cosine_similarity({}, {})
|
||||
assert similarity == 0.0
|
||||
|
||||
def test_search_relevance_ordering(self):
|
||||
"""Test that search results are ordered by relevance."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello hello hello") # High TF for "hello"
|
||||
index.add_document("doc2", "hello world") # Medium relevance
|
||||
index.add_document("doc3", "world test") # No "hello"
|
||||
|
||||
results = index.search("hello", top_k=5)
|
||||
assert len(results) == 3 # All documents are returned
|
||||
|
||||
# doc1 should have higher score than doc2, and doc3 should have 0
|
||||
scores = {doc_id: score for doc_id, score in results}
|
||||
assert scores["doc1"] > scores["doc2"] > scores["doc3"]
|
||||
assert scores["doc3"] == 0
|
||||
|
||||
def test_vocabulary_persistence(self):
|
||||
"""Test that vocabulary persists even after document removal."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
index.add_document("doc2", "test case")
|
||||
|
||||
assert index.vocabulary == {"hello", "world", "test", "case"}
|
||||
|
||||
index.remove_document("doc1")
|
||||
# Vocabulary should still contain all words
|
||||
assert index.vocabulary == {"hello", "world", "test", "case"}
|
||||
|
||||
def test_idf_recomputation_after_removal(self):
|
||||
"""Test IDF recomputation after document removal."""
|
||||
index = SemanticIndex()
|
||||
index.add_document("doc1", "hello world")
|
||||
index.add_document("doc2", "hello test")
|
||||
index.add_document("doc3", "world test")
|
||||
|
||||
# Remove doc3, leaving doc1 and doc2
|
||||
index.remove_document("doc3")
|
||||
|
||||
# "hello" appears in both remaining docs, "world" and "test" in one each
|
||||
expected_idf = {
|
||||
"hello": math.log(2 / 2), # 2 docs, appears in 2
|
||||
"world": math.log(2 / 1), # 2 docs, appears in 1
|
||||
"test": math.log(2 / 1), # 2 docs, appears in 1
|
||||
}
|
||||
assert index.idf_scores == expected_idf
|
||||
@@ -168,7 +168,6 @@ class TestToolDefinitions:
|
||||
tool_names = [t["function"]["name"] for t in tools]
|
||||
|
||||
assert "run_command" in tool_names
|
||||
assert "start_interactive_session" in tool_names
|
||||
|
||||
def test_python_exec_present(self):
|
||||
tools = get_tools_definition()
|
||||
|
||||
Reference in New Issue
Block a user