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:
2025-11-06 14:15:06 +00:00
parent ab9c29467c
commit 2c3749fd58
48 changed files with 6395 additions and 1062 deletions
+1 -36
View File
@@ -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
+372
View File
@@ -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
+242
View File
@@ -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
+344
View File
@@ -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
+280
View File
@@ -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
-1
View File
@@ -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()