Skip to content

Commit a1e6353

Browse files
committed
Enhance Knowledge Agent Configuration and Update Version
- Updated the Knowledge Agent configuration to include a new versioning system and improved vector store settings. - Changed the collection name from "custom_knowledge" to "praison" for consistency across files. - Added user_id parameter to the Knowledge Agent for better user tracking. - Enhanced the embedding database configuration to support dynamic context management. - Bumped the version of the `praisonaiagents` package from 0.0.44 to 0.0.45 in `pyproject.toml` to reflect these updates. These changes improve the functionality and usability of the Knowledge Agent within the PraisonAI framework, particularly in knowledge management and retrieval capabilities.
1 parent 4f42d3c commit a1e6353

12 files changed

+631
-21
lines changed

agents/knowledge-agents-task.py

+28-6
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,37 @@
11
from praisonaiagents import Agent, Task, PraisonAIAgents
22
import logging
33
import os
4+
import json
45

56
# Configure logging
67
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
78
logger = logging.getLogger(__name__)
89

910
# Define the configuration for the Knowledge instance
1011
config = {
12+
"version": "v1.1",
1113
"vector_store": {
1214
"provider": "chroma",
1315
"config": {
14-
"collection_name": "knowledge_test",
15-
"path": ".praison",
16+
"collection_name": "praison",
17+
"path": os.path.abspath("./.praison"),
18+
"host": None,
19+
"port": None
20+
}
21+
},
22+
"embedder": {
23+
"provider": "openai",
24+
"config": {
25+
"model": "text-embedding-3-small",
26+
"embedding_dims": 1536
27+
}
28+
},
29+
"llm": {
30+
"provider": "openai",
31+
"config": {
32+
"model": "gpt-4o-mini",
33+
"temperature": 0,
34+
"max_tokens": 1000
1635
}
1736
}
1837
}
@@ -23,17 +42,20 @@
2342
role="Information Specialist",
2443
goal="Store and retrieve knowledge efficiently",
2544
backstory="Expert in managing and utilizing stored knowledge",
26-
knowledge=["sample.pdf"],
27-
knowledge_config=config,
2845
verbose=True
2946
)
3047

3148
# Define a task for the agent
3249
knowledge_task = Task(
3350
name="knowledge_task",
34-
description="Who is Mervin Praison?",
51+
description="KAG",
3552
expected_output="Answer to the question",
36-
agent=knowledge_agent
53+
agent=knowledge_agent,
54+
context=[
55+
{
56+
"embedding_db_config": json.dumps(config["vector_store"])
57+
}
58+
]
3759
)
3860

3961
# Create and start the agents

agents/knowledge-agents.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
"vector_store": {
55
"provider": "chroma",
66
"config": {
7-
"collection_name": "custom_knowledge",
7+
"collection_name": "praison",
88
"path": ".praison",
99
}
1010
}
@@ -14,7 +14,8 @@
1414
name="Knowledge Agent",
1515
instructions="You answer questions based on the provided knowledge.",
1616
knowledge=["large.pdf"],
17-
knowledge_config=config
17+
knowledge_config=config,
18+
user_id="user1"
1819
)
1920

2021
agent.start("What is KAG in one line?")
+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import logging
2+
from praisonaiagents.knowledge import Knowledge
3+
import os
4+
# Configure logging
5+
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
6+
logger = logging.getLogger(__name__)
7+
8+
config = {
9+
"version": "v1.1",
10+
"vector_store": {
11+
"provider": "chroma",
12+
"config": {
13+
"collection_name": "praison",
14+
"path": os.path.abspath("./.praison"),
15+
"host": None,
16+
"port": None
17+
}
18+
},
19+
"embedder": {
20+
"provider": "openai",
21+
"config": {
22+
"model": "text-embedding-3-small",
23+
"embedding_dims": 1536
24+
}
25+
},
26+
"llm": {
27+
"provider": "openai",
28+
"config": {
29+
"model": "gpt-4o-mini",
30+
"temperature": 0,
31+
"max_tokens": 1000
32+
}
33+
}
34+
}
35+
36+
knowledge = Knowledge(config)
37+
38+
# Search for memories based on a query
39+
query = "KAG"
40+
logger.info(f"Searching for memories with query: {query}")
41+
search_results = knowledge.search(query, user_id="user1")
42+
logger.info(f"Search results: {search_results}")
43+
print(f"\nSearch results for '{query}':")
44+
print(search_results)

agents/knowledge-embedding-agents.py

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
from praisonaiagents import Agent
2+
3+
config = {
4+
"vector_store": {
5+
"provider": "chroma",
6+
"config": {
7+
"collection_name": "praison",
8+
"path": ".praison",
9+
}
10+
}
11+
}
12+
13+
agent = Agent(
14+
name="Knowledge Agent",
15+
instructions="You answer questions based on the provided knowledge.",
16+
knowledge=["small.pdf"],
17+
knowledge_config=config
18+
)
19+
20+
agent.start("What is KAG in one line?")

agents/memory-add-search.py

+220
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
from mem0 import Memory
2+
import logging
3+
import json
4+
import os
5+
import time
6+
from datetime import datetime
7+
from typing import Dict, Any, Optional
8+
9+
# Configure logging
10+
logging.basicConfig(
11+
level=logging.INFO,
12+
format='%(asctime)s - %(levelname)s - %(message)s'
13+
)
14+
logger = logging.getLogger(__name__)
15+
16+
class MemoryManager:
17+
def __init__(self):
18+
"""Initialize Mem0 Memory manager with configuration"""
19+
# Define the complete configuration with only supported fields
20+
config = {
21+
"version": "v1.1",
22+
"custom_prompt": None,
23+
"vector_store": {
24+
"provider": "chroma",
25+
"config": {
26+
"collection_name": "praison",
27+
"path": os.path.abspath("./.praison/storage"),
28+
"host": None,
29+
"port": None
30+
}
31+
},
32+
"embedder": {
33+
"provider": "openai",
34+
"config": {
35+
"model": "text-embedding-3-small",
36+
"embedding_dims": 1536
37+
}
38+
},
39+
"llm": {
40+
"provider": "openai",
41+
"config": {
42+
"model": "gpt-4o-mini",
43+
"temperature": 0,
44+
"max_tokens": 1000
45+
}
46+
}
47+
}
48+
49+
# Ensure storage directory exists
50+
storage_path = os.path.abspath("./.praison/storage")
51+
if os.path.exists(storage_path):
52+
logger.info(f"Clearing existing storage at {storage_path}")
53+
import shutil
54+
shutil.rmtree(storage_path)
55+
os.makedirs(storage_path, exist_ok=True)
56+
logger.info(f"Created storage path: {storage_path}")
57+
58+
# Initialize memory
59+
self.memory = Memory.from_config(config)
60+
logger.info("Memory manager initialized")
61+
62+
def pretty_print(self, data: Any) -> None:
63+
"""Helper function to print memory data in a readable format"""
64+
if isinstance(data, (dict, list)):
65+
print(json.dumps(data, indent=2))
66+
else:
67+
print(data)
68+
69+
def extract_memory_id(self, result: Any) -> Optional[str]:
70+
"""Extract memory ID from the result, handling different formats"""
71+
try:
72+
if isinstance(result, dict):
73+
# v1.1 format
74+
memories = result.get('results', [])
75+
return memories[0]['id'] if memories else None
76+
elif isinstance(result, list):
77+
# v1.0 format
78+
return result[0]['id'] if result else None
79+
return None
80+
except (KeyError, IndexError) as e:
81+
logger.warning(f"Could not extract memory ID from result: {e}")
82+
return None
83+
84+
def prepare_metadata(self, metadata: Optional[Dict]) -> Dict:
85+
"""Prepare metadata by converting values to acceptable types"""
86+
if not metadata:
87+
return {}
88+
89+
clean_metadata = {}
90+
for key, value in metadata.items():
91+
if isinstance(value, (str, int, float, bool)):
92+
clean_metadata[key] = value
93+
elif isinstance(value, list):
94+
clean_metadata[key] = ','.join(map(str, value))
95+
elif isinstance(value, dict):
96+
clean_metadata[key] = json.dumps(value)
97+
else:
98+
clean_metadata[key] = str(value)
99+
return clean_metadata
100+
101+
def verify_memory_addition(self, user_id: str, max_retries: int = 3) -> bool:
102+
"""Verify that memory was added successfully"""
103+
for attempt in range(max_retries):
104+
try:
105+
time.sleep(0.5) # Brief pause to allow indexing
106+
memories = self.memory.get_all(user_id=user_id)
107+
if len(memories.get('results', [])) > 0:
108+
return True
109+
logger.warning(f"Memory not found on attempt {attempt + 1}")
110+
except Exception as e:
111+
logger.warning(f"Verification attempt {attempt + 1} failed: {e}")
112+
if attempt < max_retries - 1:
113+
time.sleep(1)
114+
return False
115+
116+
def add_memory(self, text: str, user_id: Optional[str] = None, metadata: Optional[Dict] = None) -> Dict:
117+
"""Add a new memory with error handling and verification"""
118+
try:
119+
logger.info(f"Adding memory for user {user_id}")
120+
clean_metadata = self.prepare_metadata(metadata)
121+
logger.debug(f"Prepared metadata: {clean_metadata}")
122+
123+
messages = [{"role": "user", "content": text}]
124+
result = self.memory.add(
125+
messages=messages,
126+
user_id=user_id,
127+
metadata=clean_metadata
128+
)
129+
130+
if self.verify_memory_addition(user_id):
131+
logger.info("Memory added and verified successfully")
132+
else:
133+
logger.warning("Memory addition could not be verified")
134+
135+
return result
136+
137+
except Exception as e:
138+
logger.error(f"Failed to add memory: {str(e)}")
139+
raise
140+
141+
def get_all_memories(self, user_id: Optional[str] = None) -> Dict:
142+
"""Retrieve all memories for a user"""
143+
try:
144+
logger.info(f"Retrieving all memories for user {user_id}")
145+
memories = self.memory.get_all(user_id=user_id)
146+
count = len(memories.get('results', []))
147+
logger.info(f"Retrieved {count} memories")
148+
return memories
149+
except Exception as e:
150+
logger.error(f"Failed to retrieve memories: {str(e)}")
151+
raise
152+
153+
def search_memories(self, query: str, user_id: Optional[str] = None, limit: int = 5) -> Dict:
154+
"""Search memories with a query"""
155+
try:
156+
logger.info(f"Searching memories with query: {query}")
157+
results = self.memory.search(
158+
query=query,
159+
user_id=user_id,
160+
limit=limit
161+
)
162+
count = len(results.get('results', []))
163+
logger.info(f"Found {count} matching memories")
164+
return results
165+
except Exception as e:
166+
logger.error(f"Search failed: {str(e)}")
167+
raise
168+
169+
def main():
170+
try:
171+
# Initialize memory manager
172+
manager = MemoryManager()
173+
174+
# Add a test memory
175+
print("\n1. Adding a new memory:")
176+
metadata = {
177+
"category": "hobbies",
178+
"timestamp": datetime.now().isoformat(),
179+
"tags": "tennis,learning",
180+
"priority": 1,
181+
"source": "user_input"
182+
}
183+
184+
result = manager.add_memory(
185+
text="I am working on improving my tennis skills. Suggest some online courses.",
186+
user_id="alice",
187+
metadata=metadata
188+
)
189+
manager.pretty_print(result)
190+
191+
# Wait briefly for indexing
192+
time.sleep(1)
193+
194+
print("\n2. Retrieving all memories:")
195+
all_memories = manager.get_all_memories(user_id="alice")
196+
manager.pretty_print(all_memories)
197+
198+
print("\n3. Searching memories:")
199+
search_results = manager.search_memories(
200+
query="What are Alice's hobbies?",
201+
user_id="alice",
202+
limit=5
203+
)
204+
manager.pretty_print(search_results)
205+
206+
# Print summary
207+
print("\nMemory Store Summary:")
208+
total_memories = len(all_memories.get('results', []))
209+
print(f"Total memories stored: {total_memories}")
210+
if total_memories > 0:
211+
print("Available memories:")
212+
for mem in all_memories['results']:
213+
print(f"- {mem.get('memory', 'No content')}")
214+
215+
except Exception as e:
216+
logger.error(f"Main process failed: {str(e)}", exc_info=True)
217+
raise
218+
219+
if __name__ == "__main__":
220+
main()

0 commit comments

Comments
 (0)