Skip to content

Commit b8f4fc0

Browse files
fix: resolve MongoDB integration issues
- Fix memory record citation duplication (lines 861, 880, 933 in memory.py) - Make vector dimensions configurable based on embedding model - Add _get_embedding_dimensions method for proper dimension detection - Support multiple embedding models (text-embedding-3-small, 3-large, ada-002) - Improve citation format to use record IDs instead of duplicating text 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com>
1 parent c1a99fd commit b8f4fc0

File tree

5 files changed

+180
-5
lines changed

5 files changed

+180
-5
lines changed

src/praisonai-agents/praisonaiagents/knowledge/knowledge.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,25 @@ def _init_embedding_model(self):
107107
except Exception as e:
108108
raise Exception(f"Failed to initialize embedding model: {e}")
109109

110+
def _get_embedding_dimensions(self, model_name: str) -> int:
111+
"""Get embedding dimensions based on model name."""
112+
# Common embedding model dimensions
113+
model_dimensions = {
114+
"text-embedding-3-small": 1536,
115+
"text-embedding-3-large": 3072,
116+
"text-embedding-ada-002": 1536,
117+
"text-embedding-002": 1536,
118+
# Add more models as needed
119+
}
120+
121+
# Check if model name contains known model identifiers
122+
for model_key, dimensions in model_dimensions.items():
123+
if model_key in model_name.lower():
124+
return dimensions
125+
126+
# Default to 1536 for unknown models (OpenAI standard)
127+
return 1536
128+
110129
def _create_indexes(self):
111130
"""Create necessary indexes for MongoDB."""
112131
try:
@@ -133,7 +152,7 @@ def _create_vector_index(self):
133152
"fields": {
134153
"embedding": {
135154
"type": "knnVector",
136-
"dimensions": 1536,
155+
"dimensions": self._get_embedding_dimensions(self.embedding_model_name),
137156
"similarity": "cosine"
138157
}
139158
}

src/praisonai-agents/praisonaiagents/memory/memory.py

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,10 @@ def __init__(self, config: Dict[str, Any], verbose: int = 0):
144144

145145
self._log_verbose(f"Using embedding model: {self.embedding_model}")
146146

147+
# Determine embedding dimensions based on model
148+
self.embedding_dimensions = self._get_embedding_dimensions(self.embedding_model)
149+
self._log_verbose(f"Using embedding dimensions: {self.embedding_dimensions}")
150+
147151
# Create .praison directory if it doesn't exist
148152
os.makedirs(".praison", exist_ok=True)
149153

@@ -351,7 +355,7 @@ def _create_vector_search_indexes(self):
351355
"fields": {
352356
"embedding": {
353357
"type": "knnVector",
354-
"dimensions": 1536, # OpenAI embedding dimensions
358+
"dimensions": self.embedding_dimensions,
355359
"similarity": "cosine"
356360
}
357361
}
@@ -398,6 +402,25 @@ def _get_embedding(self, text: str) -> List[float]:
398402
self._log_verbose(f"Error getting embedding: {e}", logging.ERROR)
399403
return None
400404

405+
def _get_embedding_dimensions(self, model_name: str) -> int:
406+
"""Get embedding dimensions based on model name."""
407+
# Common embedding model dimensions
408+
model_dimensions = {
409+
"text-embedding-3-small": 1536,
410+
"text-embedding-3-large": 3072,
411+
"text-embedding-ada-002": 1536,
412+
"text-embedding-002": 1536,
413+
# Add more models as needed
414+
}
415+
416+
# Check if model name contains known model identifiers
417+
for model_key, dimensions in model_dimensions.items():
418+
if model_key in model_name.lower():
419+
return dimensions
420+
421+
# Default to 1536 for unknown models (OpenAI standard)
422+
return 1536
423+
401424
# -------------------------------------------------------------------------
402425
# Basic Quality Score Computation
403426
# -------------------------------------------------------------------------
@@ -858,7 +881,7 @@ def search_long_term(
858881
text = doc["content"]
859882
# Add memory record citation
860883
if "(Memory record:" not in text:
861-
text = f"{text} (Memory record: {text})"
884+
text = f"{text} (Memory record: {str(doc['_id'])})"
862885
results.append({
863886
"id": str(doc["_id"]),
864887
"text": text,
@@ -877,7 +900,7 @@ def search_long_term(
877900
text = doc["content"]
878901
# Add memory record citation
879902
if "(Memory record:" not in text:
880-
text = f"{text} (Memory record: {text})"
903+
text = f"{text} (Memory record: {str(doc['_id'])})"
881904
results.append({
882905
"id": str(doc["_id"]),
883906
"text": text,
@@ -930,7 +953,7 @@ def search_long_term(
930953
metadata = resp["metadatas"][0][i] if "metadatas" in resp else {}
931954
text = resp["documents"][0][i]
932955
# Add memory record citation
933-
text = f"{text} (Memory record: {text})"
956+
text = f"{text} (Memory record: {resp['ids'][0][i]})"
934957
found.append({
935958
"id": resp["ids"][0][i],
936959
"text": text,

test_mongodb_direct.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#!/usr/bin/env python3
2+
"""Direct test of MongoDB tools"""
3+
import sys
4+
sys.path.insert(0, 'src/praisonai-agents')
5+
6+
try:
7+
# Test direct import of MongoDB tools
8+
from praisonaiagents.tools.mongodb_tools import MongoDBTools
9+
print('MongoDB tools import successful')
10+
11+
# Test basic initialization
12+
tools = MongoDBTools()
13+
print('MongoDB tools initialization successful')
14+
15+
# Test the _get_pymongo method which should handle missing dependency gracefully
16+
pymongo = tools._get_pymongo()
17+
if pymongo is None:
18+
print('pymongo not available - graceful fallback working')
19+
else:
20+
print('pymongo available')
21+
22+
# Test client connection (should fail gracefully)
23+
client = tools._get_client()
24+
if client is None:
25+
print('MongoDB client connection failed gracefully (expected without pymongo)')
26+
else:
27+
print('MongoDB client connection successful')
28+
29+
print('All MongoDB direct tests passed')
30+
except Exception as e:
31+
print(f'MongoDB direct test failed: {e}')
32+
import traceback
33+
traceback.print_exc()

test_mongodb_integration.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#!/usr/bin/env python3
2+
"""Test MongoDB integration"""
3+
import sys
4+
sys.path.insert(0, 'src/praisonai-agents')
5+
6+
try:
7+
from praisonaiagents.tools.mongodb_tools import MongoDBTools
8+
print('MongoDB tools import successful')
9+
10+
# Test basic initialization (should work even without MongoDB running)
11+
tools = MongoDBTools()
12+
print('MongoDB tools initialization successful')
13+
14+
# Test function imports
15+
from praisonaiagents.tools import mongodb_tools
16+
print('MongoDB tools module import successful')
17+
18+
# Test individual function imports
19+
from praisonaiagents.tools import insert_document, find_documents, vector_search
20+
print('Individual MongoDB function imports successful')
21+
22+
print('All MongoDB integration tests passed')
23+
except Exception as e:
24+
print(f'MongoDB integration test failed: {e}')
25+
import traceback
26+
traceback.print_exc()

test_mongodb_isolated.py

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
#!/usr/bin/env python3
2+
"""Isolated test of MongoDB tools"""
3+
import sys
4+
import os
5+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src/praisonai-agents'))
6+
7+
# Test if the MongoDB tools can be directly imported and used
8+
try:
9+
# Read the mongodb_tools.py file content and check for issues
10+
with open('src/praisonai-agents/praisonaiagents/tools/mongodb_tools.py', 'r') as f:
11+
content = f.read()
12+
13+
# Check for potential syntax or import issues
14+
print("Checking MongoDB tools file...")
15+
16+
# Check imports
17+
if 'from importlib import util' in content:
18+
print('✓ importlib import present')
19+
if 'from typing import' in content:
20+
print('✓ typing imports present')
21+
if 'TYPE_CHECKING' in content:
22+
print('✓ TYPE_CHECKING conditional imports present')
23+
24+
# Check class structure
25+
if 'class MongoDBTools:' in content:
26+
print('✓ MongoDBTools class present')
27+
if 'def _get_pymongo' in content:
28+
print('✓ _get_pymongo method present')
29+
if 'def _get_client' in content:
30+
print('✓ _get_client method present')
31+
32+
# Check MongoDB operations
33+
if 'def insert_document' in content:
34+
print('✓ insert_document method present')
35+
if 'def vector_search' in content:
36+
print('✓ vector_search method present')
37+
if 'def create_vector_index' in content:
38+
print('✓ create_vector_index method present')
39+
40+
# Check error handling
41+
if 'try:' in content and 'except' in content:
42+
print('✓ Error handling present')
43+
44+
# Check graceful fallback
45+
if 'if util.find_spec' in content:
46+
print('✓ Graceful dependency checking present')
47+
48+
print('MongoDB tools file structure looks good!')
49+
50+
except Exception as e:
51+
print(f'Error checking MongoDB tools: {e}')
52+
53+
# Check if imports work correctly
54+
try:
55+
# Test individual components
56+
import logging
57+
from typing import List, Dict, Any, Optional, Union
58+
from importlib import util
59+
import json
60+
import time
61+
from datetime import datetime
62+
63+
print('✓ All required imports for MongoDB tools available')
64+
65+
# Test util.find_spec functionality
66+
if util.find_spec('pymongo') is None:
67+
print('✓ pymongo not available - graceful fallback should work')
68+
else:
69+
print('✓ pymongo available')
70+
71+
print('MongoDB tools dependencies check passed')
72+
73+
except Exception as e:
74+
print(f'Error with MongoDB tools dependencies: {e}')

0 commit comments

Comments
 (0)