solana-agent 31.0.0__tar.gz → 31.1.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {solana_agent-31.0.0 → solana_agent-31.1.1}/PKG-INFO +32 -27
- {solana_agent-31.0.0 → solana_agent-31.1.1}/README.md +22 -17
- {solana_agent-31.0.0 → solana_agent-31.1.1}/pyproject.toml +13 -13
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/client/solana_agent.py +4 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/domains/agent.py +7 -1
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/factories/agent_factory.py +37 -6
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/providers/memory.py +12 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/services/query.py +2 -0
- solana_agent-31.1.1/solana_agent/repositories/memory.py +276 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/services/agent.py +33 -1
- solana_agent-31.1.1/solana_agent/services/query.py +743 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/services/routing.py +19 -13
- solana_agent-31.0.0/solana_agent/repositories/memory.py +0 -216
- solana_agent-31.0.0/solana_agent/services/query.py +0 -460
- {solana_agent-31.0.0 → solana_agent-31.1.1}/LICENSE +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/adapters/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/adapters/mongodb_adapter.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/adapters/openai_adapter.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/adapters/pinecone_adapter.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/cli.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/client/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/domains/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/domains/routing.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/factories/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/guardrails/pii.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/client/client.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/guardrails/guardrails.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/plugins/plugins.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/providers/data_storage.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/providers/llm.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/providers/vector_storage.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/services/agent.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/services/knowledge_base.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/interfaces/services/routing.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/plugins/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/plugins/manager.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/plugins/registry.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/plugins/tools/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/plugins/tools/auto_tool.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/repositories/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/services/__init__.py +0 -0
- {solana_agent-31.0.0 → solana_agent-31.1.1}/solana_agent/services/knowledge_base.py +0 -0
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.3
|
2
2
|
Name: solana-agent
|
3
|
-
Version: 31.
|
3
|
+
Version: 31.1.1
|
4
4
|
Summary: AI Agents for Solana
|
5
5
|
License: MIT
|
6
6
|
Keywords: solana,solana ai,solana agent,ai,ai agent,ai agents
|
@@ -14,20 +14,20 @@ Classifier: Programming Language :: Python :: 3
|
|
14
14
|
Classifier: Programming Language :: Python :: 3.12
|
15
15
|
Classifier: Programming Language :: Python :: 3.13
|
16
16
|
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
17
|
-
Requires-Dist: instructor (==1.
|
18
|
-
Requires-Dist: llama-index-core (==0.13.
|
17
|
+
Requires-Dist: instructor (==1.11.2)
|
18
|
+
Requires-Dist: llama-index-core (==0.13.3)
|
19
19
|
Requires-Dist: llama-index-embeddings-openai (==0.5.0)
|
20
|
-
Requires-Dist: logfire (==4.
|
21
|
-
Requires-Dist: openai (==1.
|
20
|
+
Requires-Dist: logfire (==4.3.6)
|
21
|
+
Requires-Dist: openai (==1.102.0)
|
22
22
|
Requires-Dist: pillow (==11.3.0)
|
23
|
-
Requires-Dist: pinecone (==7.3.0)
|
23
|
+
Requires-Dist: pinecone[asyncio] (==7.3.0)
|
24
24
|
Requires-Dist: pydantic (>=2)
|
25
|
-
Requires-Dist: pymongo (==4.14.
|
26
|
-
Requires-Dist: pypdf (==
|
25
|
+
Requires-Dist: pymongo (==4.14.1)
|
26
|
+
Requires-Dist: pypdf (==6.0.0)
|
27
27
|
Requires-Dist: rich (>=13,<14.0)
|
28
28
|
Requires-Dist: scrubadub (==2.0.1)
|
29
|
-
Requires-Dist: typer (==0.
|
30
|
-
Requires-Dist: zep-cloud (==3.
|
29
|
+
Requires-Dist: typer (==0.17.3)
|
30
|
+
Requires-Dist: zep-cloud (==3.4.3)
|
31
31
|
Project-URL: Documentation, https://docs.solana-agent.com
|
32
32
|
Project-URL: Homepage, https://solana-agent.com
|
33
33
|
Project-URL: Repository, https://github.com/truemagic-coder/solana-agent
|
@@ -63,7 +63,7 @@ Build your AI agents in three lines of code!
|
|
63
63
|
* Extensible Tooling
|
64
64
|
* Autonomous Operation
|
65
65
|
* Smart Workflows
|
66
|
-
*
|
66
|
+
* Agentic Forms
|
67
67
|
* Knowledge Base
|
68
68
|
* MCP Support
|
69
69
|
* Guardrails
|
@@ -112,7 +112,7 @@ Smart workflows are as easy as combining your tools and prompts.
|
|
112
112
|
* Integrated Knowledge Base with semantic search and automatic PDF chunking
|
113
113
|
* Input and output guardrails for content filtering, safety, and data sanitization
|
114
114
|
* Generate custom images based on text prompts with storage on S3 compatible services
|
115
|
-
*
|
115
|
+
* Deterministic agentic form filling in natural conversation
|
116
116
|
* Combine with event-driven systems to create autonomous agents
|
117
117
|
|
118
118
|
## Stack
|
@@ -344,7 +344,9 @@ async for response in solana_agent.process("user123", "What is in this image? De
|
|
344
344
|
print(response, end="")
|
345
345
|
```
|
346
346
|
|
347
|
-
###
|
347
|
+
### Agentic Forms
|
348
|
+
|
349
|
+
You can attach a JSON Schema to any agent in your config so it can collect structured data conversationally.
|
348
350
|
|
349
351
|
```python
|
350
352
|
from solana_agent import SolanaAgent
|
@@ -355,25 +357,28 @@ config = {
|
|
355
357
|
},
|
356
358
|
"agents": [
|
357
359
|
{
|
358
|
-
"name": "
|
359
|
-
"instructions": "You
|
360
|
-
"specialization": "
|
360
|
+
"name": "customer_support",
|
361
|
+
"instructions": "You provide friendly, helpful customer support responses.",
|
362
|
+
"specialization": "Customer inquiries",
|
363
|
+
"capture_name": "contact_info",
|
364
|
+
"capture_mode": "once",
|
365
|
+
"capture_schema": {
|
366
|
+
"type": "object",
|
367
|
+
"properties": {
|
368
|
+
"email": { "type": "string" },
|
369
|
+
"phone": { "type": "string" },
|
370
|
+
"newsletter_subscribe": { "type": "boolean" }
|
371
|
+
},
|
372
|
+
"required": ["email"]
|
373
|
+
}
|
361
374
|
}
|
362
|
-
]
|
375
|
+
]
|
363
376
|
}
|
364
377
|
|
365
378
|
solana_agent = SolanaAgent(config=config)
|
366
379
|
|
367
|
-
|
368
|
-
|
369
|
-
abstract: str
|
370
|
-
key_points: list[str]
|
371
|
-
|
372
|
-
full_response = None
|
373
|
-
async for response in solana_agent.process("user123", "Research the life of Ben Franklin - the founding Father.", output_model=ResearchProposal):
|
374
|
-
full_response = response
|
375
|
-
|
376
|
-
print(full_response.model_dump())
|
380
|
+
async for response in solana_agent.process("user123", "What are the latest AI developments?"):
|
381
|
+
print(response, end="")
|
377
382
|
```
|
378
383
|
|
379
384
|
### Command Line Interface (CLI)
|
@@ -28,7 +28,7 @@ Build your AI agents in three lines of code!
|
|
28
28
|
* Extensible Tooling
|
29
29
|
* Autonomous Operation
|
30
30
|
* Smart Workflows
|
31
|
-
*
|
31
|
+
* Agentic Forms
|
32
32
|
* Knowledge Base
|
33
33
|
* MCP Support
|
34
34
|
* Guardrails
|
@@ -77,7 +77,7 @@ Smart workflows are as easy as combining your tools and prompts.
|
|
77
77
|
* Integrated Knowledge Base with semantic search and automatic PDF chunking
|
78
78
|
* Input and output guardrails for content filtering, safety, and data sanitization
|
79
79
|
* Generate custom images based on text prompts with storage on S3 compatible services
|
80
|
-
*
|
80
|
+
* Deterministic agentic form filling in natural conversation
|
81
81
|
* Combine with event-driven systems to create autonomous agents
|
82
82
|
|
83
83
|
## Stack
|
@@ -309,7 +309,9 @@ async for response in solana_agent.process("user123", "What is in this image? De
|
|
309
309
|
print(response, end="")
|
310
310
|
```
|
311
311
|
|
312
|
-
###
|
312
|
+
### Agentic Forms
|
313
|
+
|
314
|
+
You can attach a JSON Schema to any agent in your config so it can collect structured data conversationally.
|
313
315
|
|
314
316
|
```python
|
315
317
|
from solana_agent import SolanaAgent
|
@@ -320,25 +322,28 @@ config = {
|
|
320
322
|
},
|
321
323
|
"agents": [
|
322
324
|
{
|
323
|
-
"name": "
|
324
|
-
"instructions": "You
|
325
|
-
"specialization": "
|
325
|
+
"name": "customer_support",
|
326
|
+
"instructions": "You provide friendly, helpful customer support responses.",
|
327
|
+
"specialization": "Customer inquiries",
|
328
|
+
"capture_name": "contact_info",
|
329
|
+
"capture_mode": "once",
|
330
|
+
"capture_schema": {
|
331
|
+
"type": "object",
|
332
|
+
"properties": {
|
333
|
+
"email": { "type": "string" },
|
334
|
+
"phone": { "type": "string" },
|
335
|
+
"newsletter_subscribe": { "type": "boolean" }
|
336
|
+
},
|
337
|
+
"required": ["email"]
|
338
|
+
}
|
326
339
|
}
|
327
|
-
]
|
340
|
+
]
|
328
341
|
}
|
329
342
|
|
330
343
|
solana_agent = SolanaAgent(config=config)
|
331
344
|
|
332
|
-
|
333
|
-
|
334
|
-
abstract: str
|
335
|
-
key_points: list[str]
|
336
|
-
|
337
|
-
full_response = None
|
338
|
-
async for response in solana_agent.process("user123", "Research the life of Ben Franklin - the founding Father.", output_model=ResearchProposal):
|
339
|
-
full_response = response
|
340
|
-
|
341
|
-
print(full_response.model_dump())
|
345
|
+
async for response in solana_agent.process("user123", "What are the latest AI developments?"):
|
346
|
+
print(response, end="")
|
342
347
|
```
|
343
348
|
|
344
349
|
### Command Line Interface (CLI)
|
@@ -1,6 +1,6 @@
|
|
1
1
|
[tool.poetry]
|
2
2
|
name = "solana-agent"
|
3
|
-
version = "31.
|
3
|
+
version = "31.1.1"
|
4
4
|
description = "AI Agents for Solana"
|
5
5
|
authors = ["Bevan Hunt <bevan@bevanhunt.com>"]
|
6
6
|
license = "MIT"
|
@@ -19,23 +19,23 @@ classifiers = [
|
|
19
19
|
]
|
20
20
|
packages = [{ include = "solana_agent" }]
|
21
21
|
|
22
|
-
[tool.pytest.
|
23
|
-
|
22
|
+
[tool.pytest.ini_options]
|
23
|
+
testpaths = ["tests"]
|
24
24
|
|
25
25
|
[tool.poetry.dependencies]
|
26
26
|
python = ">=3.12,<4.0"
|
27
|
-
openai = "1.
|
27
|
+
openai = "1.102.0"
|
28
28
|
pydantic = ">=2"
|
29
|
-
pymongo = "4.14.
|
30
|
-
zep-cloud = "3.
|
31
|
-
instructor = "1.
|
32
|
-
pinecone = "7.3.0"
|
33
|
-
llama-index-core = "0.13.
|
29
|
+
pymongo = "4.14.1"
|
30
|
+
zep-cloud = "3.4.3"
|
31
|
+
instructor = "1.11.2"
|
32
|
+
pinecone = { version = "7.3.0", extras = ["asyncio"] }
|
33
|
+
llama-index-core = "0.13.3"
|
34
34
|
llama-index-embeddings-openai = "0.5.0"
|
35
|
-
pypdf = "
|
35
|
+
pypdf = "6.0.0"
|
36
36
|
scrubadub = "2.0.1"
|
37
|
-
logfire = "4.
|
38
|
-
typer = "0.
|
37
|
+
logfire = "4.3.6"
|
38
|
+
typer = "0.17.3"
|
39
39
|
rich = ">=13,<14.0"
|
40
40
|
pillow = "11.3.0"
|
41
41
|
|
@@ -50,7 +50,7 @@ sphinx-rtd-theme = "^3.0.2"
|
|
50
50
|
myst-parser = "^4.0.1"
|
51
51
|
sphinx-autobuild = "^2024.10.3"
|
52
52
|
mongomock = "^4.3.0"
|
53
|
-
ruff = "^0.12.
|
53
|
+
ruff = "^0.12.10"
|
54
54
|
|
55
55
|
[tool.poetry.scripts]
|
56
56
|
solana-agent = "solana_agent.cli:app"
|
@@ -49,6 +49,8 @@ class SolanaAgent(SolanaAgentInterface):
|
|
49
49
|
user_id: str,
|
50
50
|
message: Union[str, bytes],
|
51
51
|
prompt: Optional[str] = None,
|
52
|
+
capture_schema: Optional[Dict[str, Any]] = None,
|
53
|
+
capture_name: Optional[str] = None,
|
52
54
|
output_format: Literal["text", "audio"] = "text",
|
53
55
|
audio_voice: Literal[
|
54
56
|
"alloy",
|
@@ -103,6 +105,8 @@ class SolanaAgent(SolanaAgentInterface):
|
|
103
105
|
prompt=prompt,
|
104
106
|
router=router,
|
105
107
|
output_model=output_model,
|
108
|
+
capture_schema=capture_schema,
|
109
|
+
capture_name=capture_name,
|
106
110
|
):
|
107
111
|
yield chunk
|
108
112
|
|
@@ -5,7 +5,7 @@ This module defines the core domain models for representing
|
|
5
5
|
AI agents, human agents, and business mission/values.
|
6
6
|
"""
|
7
7
|
|
8
|
-
from typing import List, Dict
|
8
|
+
from typing import List, Dict, Optional, Any
|
9
9
|
from pydantic import BaseModel, Field, field_validator
|
10
10
|
|
11
11
|
|
@@ -53,6 +53,12 @@ class AIAgent(BaseModel):
|
|
53
53
|
name: str = Field(..., description="Unique agent identifier name")
|
54
54
|
instructions: str = Field(..., description="Base instructions for the agent")
|
55
55
|
specialization: str = Field(..., description="Agent's specialized domain")
|
56
|
+
capture_name: Optional[str] = Field(
|
57
|
+
default=None, description="Optional capture name for structured data"
|
58
|
+
)
|
59
|
+
capture_schema: Optional[Dict[str, Any]] = Field(
|
60
|
+
default=None, description="Optional JSON schema for structured capture"
|
61
|
+
)
|
56
62
|
|
57
63
|
@field_validator("name", "specialization")
|
58
64
|
@classmethod
|
@@ -133,21 +133,38 @@ class SolanaAgentFactory:
|
|
133
133
|
voice=org_config.get("voice", ""),
|
134
134
|
)
|
135
135
|
|
136
|
+
# Build capture modes from agent config if provided
|
137
|
+
capture_modes: Dict[str, str] = {}
|
138
|
+
for agent in config.get("agents", []):
|
139
|
+
mode = agent.get("capture_mode")
|
140
|
+
if mode in {"once", "multiple"} and agent.get("name"):
|
141
|
+
capture_modes[agent["name"]] = mode
|
142
|
+
|
136
143
|
# Create repositories
|
137
144
|
memory_provider = None
|
138
145
|
|
139
146
|
if "zep" in config and "mongo" in config:
|
140
|
-
|
141
|
-
mongo_adapter
|
142
|
-
|
147
|
+
mem_kwargs: Dict[str, Any] = {
|
148
|
+
"mongo_adapter": db_adapter,
|
149
|
+
"zep_api_key": config["zep"].get("api_key"),
|
150
|
+
}
|
151
|
+
if capture_modes: # pragma: no cover
|
152
|
+
mem_kwargs["capture_modes"] = capture_modes
|
153
|
+
memory_provider = MemoryRepository(**mem_kwargs)
|
143
154
|
|
144
155
|
if "mongo" in config and "zep" not in config:
|
145
|
-
|
156
|
+
mem_kwargs = {"mongo_adapter": db_adapter}
|
157
|
+
if capture_modes:
|
158
|
+
mem_kwargs["capture_modes"] = capture_modes
|
159
|
+
memory_provider = MemoryRepository(**mem_kwargs)
|
146
160
|
|
147
161
|
if "zep" in config and "mongo" not in config:
|
148
162
|
if "api_key" not in config["zep"]:
|
149
163
|
raise ValueError("Zep API key is required.")
|
150
|
-
|
164
|
+
mem_kwargs = {"zep_api_key": config["zep"].get("api_key")}
|
165
|
+
if capture_modes: # pragma: no cover
|
166
|
+
mem_kwargs["capture_modes"] = capture_modes
|
167
|
+
memory_provider = MemoryRepository(**mem_kwargs)
|
151
168
|
|
152
169
|
guardrail_config = config.get("guardrails", {})
|
153
170
|
input_guardrails: List[InputGuardrail] = SolanaAgentFactory._create_guardrails(
|
@@ -169,9 +186,16 @@ class SolanaAgentFactory:
|
|
169
186
|
)
|
170
187
|
|
171
188
|
# Create routing service
|
189
|
+
# Optional routing model override (use small, cheap model by default in service)
|
190
|
+
routing_model = (
|
191
|
+
config.get("openai", {}).get("routing_model")
|
192
|
+
if isinstance(config.get("openai"), dict)
|
193
|
+
else None
|
194
|
+
)
|
172
195
|
routing_service = RoutingService(
|
173
196
|
llm_provider=llm_adapter,
|
174
197
|
agent_service=agent_service,
|
198
|
+
model=routing_model,
|
175
199
|
)
|
176
200
|
|
177
201
|
# Debug the agent service tool registry
|
@@ -191,11 +215,18 @@ class SolanaAgentFactory:
|
|
191
215
|
loaded_plugins = 0
|
192
216
|
|
193
217
|
# Register predefined agents
|
194
|
-
for agent_config in config.get("agents", []):
|
218
|
+
for agent_config in config.get("agents", []): # pragma: no cover
|
219
|
+
extra_kwargs = {}
|
220
|
+
if "capture_name" in agent_config:
|
221
|
+
extra_kwargs["capture_name"] = agent_config.get("capture_name")
|
222
|
+
if "capture_schema" in agent_config:
|
223
|
+
extra_kwargs["capture_schema"] = agent_config.get("capture_schema")
|
224
|
+
|
195
225
|
agent_service.register_ai_agent(
|
196
226
|
name=agent_config["name"],
|
197
227
|
instructions=agent_config["instructions"],
|
198
228
|
specialization=agent_config["specialization"],
|
229
|
+
**extra_kwargs,
|
199
230
|
)
|
200
231
|
|
201
232
|
# Register tools for this agent
|
@@ -36,3 +36,15 @@ class MemoryProvider(ABC):
|
|
36
36
|
def count_documents(self, collection: str, query: Dict) -> int:
|
37
37
|
"""Count documents matching query."""
|
38
38
|
pass
|
39
|
+
|
40
|
+
@abstractmethod
|
41
|
+
async def save_capture(
|
42
|
+
self,
|
43
|
+
user_id: str,
|
44
|
+
capture_name: str,
|
45
|
+
agent_name: Optional[str],
|
46
|
+
data: Dict[str, Any],
|
47
|
+
schema: Optional[Dict[str, Any]] = None,
|
48
|
+
) -> Optional[str]:
|
49
|
+
"""Persist a structured capture for a user and return its ID if available."""
|
50
|
+
pass
|
@@ -38,6 +38,8 @@ class QueryService(ABC):
|
|
38
38
|
router: Optional[RoutingInterface] = None,
|
39
39
|
images: Optional[List[Union[str, bytes]]] = None,
|
40
40
|
output_model: Optional[Type[BaseModel]] = None,
|
41
|
+
capture_schema: Optional[Dict[str, Any]] = None,
|
42
|
+
capture_name: Optional[str] = None,
|
41
43
|
) -> AsyncGenerator[Union[str, bytes, BaseModel], None]:
|
42
44
|
"""Process the user request and generate a response."""
|
43
45
|
pass
|
@@ -0,0 +1,276 @@
|
|
1
|
+
import logging
|
2
|
+
from typing import List, Dict, Optional, Tuple, Any
|
3
|
+
from datetime import datetime, timezone
|
4
|
+
from copy import deepcopy
|
5
|
+
|
6
|
+
from zep_cloud.client import AsyncZep as AsyncZepCloud
|
7
|
+
from zep_cloud.types import Message
|
8
|
+
|
9
|
+
from solana_agent.interfaces.providers.memory import MemoryProvider
|
10
|
+
from solana_agent.adapters.mongodb_adapter import MongoDBAdapter
|
11
|
+
|
12
|
+
logger = logging.getLogger(__name__)
|
13
|
+
|
14
|
+
|
15
|
+
class MemoryRepository(MemoryProvider):
|
16
|
+
"""Combined Zep and MongoDB implementation of MemoryProvider."""
|
17
|
+
|
18
|
+
def __init__(
|
19
|
+
self,
|
20
|
+
mongo_adapter: Optional[MongoDBAdapter] = None,
|
21
|
+
zep_api_key: Optional[str] = None,
|
22
|
+
capture_modes: Optional[Dict[str, str]] = None,
|
23
|
+
):
|
24
|
+
self.capture_modes: Dict[str, str] = capture_modes or {}
|
25
|
+
|
26
|
+
# Mongo setup
|
27
|
+
if not mongo_adapter:
|
28
|
+
self.mongo = None
|
29
|
+
self.collection = None
|
30
|
+
self.captures_collection = "captures"
|
31
|
+
else:
|
32
|
+
self.mongo = mongo_adapter
|
33
|
+
self.collection = "conversations"
|
34
|
+
try:
|
35
|
+
self.mongo.create_collection(self.collection)
|
36
|
+
self.mongo.create_index(self.collection, [("user_id", 1)])
|
37
|
+
self.mongo.create_index(self.collection, [("timestamp", 1)])
|
38
|
+
except Exception as e:
|
39
|
+
logger.error(f"Error initializing MongoDB: {e}")
|
40
|
+
|
41
|
+
try:
|
42
|
+
self.captures_collection = "captures"
|
43
|
+
self.mongo.create_collection(self.captures_collection)
|
44
|
+
# Basic indexes
|
45
|
+
self.mongo.create_index(self.captures_collection, [("user_id", 1)])
|
46
|
+
self.mongo.create_index(self.captures_collection, [("capture_name", 1)])
|
47
|
+
self.mongo.create_index(self.captures_collection, [("agent_name", 1)])
|
48
|
+
self.mongo.create_index(self.captures_collection, [("timestamp", 1)])
|
49
|
+
# Unique only when mode == 'once'
|
50
|
+
try:
|
51
|
+
self.mongo.create_index(
|
52
|
+
self.captures_collection,
|
53
|
+
[("user_id", 1), ("agent_name", 1), ("capture_name", 1)],
|
54
|
+
unique=True,
|
55
|
+
partialFilterExpression={"mode": "once"},
|
56
|
+
)
|
57
|
+
except Exception as e:
|
58
|
+
logger.error(
|
59
|
+
f"Error creating partial unique index for captures: {e}"
|
60
|
+
)
|
61
|
+
except Exception as e:
|
62
|
+
logger.error(f"Error initializing MongoDB captures collection: {e}")
|
63
|
+
self.captures_collection = "captures"
|
64
|
+
|
65
|
+
# Zep setup
|
66
|
+
self.zep = AsyncZepCloud(api_key=zep_api_key) if zep_api_key else None
|
67
|
+
|
68
|
+
async def store(self, user_id: str, messages: List[Dict[str, Any]]) -> None:
|
69
|
+
if not user_id or not isinstance(user_id, str):
|
70
|
+
raise ValueError("User ID cannot be None or empty")
|
71
|
+
if not messages or not isinstance(messages, list):
|
72
|
+
raise ValueError("Messages must be a non-empty list")
|
73
|
+
if not all(
|
74
|
+
isinstance(m, dict) and "role" in m and "content" in m for m in messages
|
75
|
+
):
|
76
|
+
raise ValueError(
|
77
|
+
"All messages must be dictionaries with 'role' and 'content' keys"
|
78
|
+
)
|
79
|
+
for m in messages:
|
80
|
+
if m["role"] not in ["user", "assistant"]:
|
81
|
+
raise ValueError(
|
82
|
+
"Invalid role in message. Only 'user' and 'assistant' are accepted."
|
83
|
+
)
|
84
|
+
|
85
|
+
# Persist last user/assistant pair to Mongo
|
86
|
+
if self.mongo and len(messages) >= 2:
|
87
|
+
try:
|
88
|
+
user_msg = None
|
89
|
+
assistant_msg = None
|
90
|
+
for m in reversed(messages):
|
91
|
+
if m.get("role") == "user" and not user_msg:
|
92
|
+
user_msg = m.get("content")
|
93
|
+
elif m.get("role") == "assistant" and not assistant_msg:
|
94
|
+
assistant_msg = m.get("content")
|
95
|
+
if user_msg and assistant_msg:
|
96
|
+
break
|
97
|
+
if user_msg and assistant_msg:
|
98
|
+
self.mongo.insert_one(
|
99
|
+
self.collection,
|
100
|
+
{
|
101
|
+
"user_id": user_id,
|
102
|
+
"user_message": user_msg,
|
103
|
+
"assistant_message": assistant_msg,
|
104
|
+
"timestamp": datetime.now(timezone.utc),
|
105
|
+
},
|
106
|
+
)
|
107
|
+
except Exception as e:
|
108
|
+
logger.error(f"MongoDB storage error: {e}")
|
109
|
+
|
110
|
+
# Zep
|
111
|
+
if not self.zep:
|
112
|
+
return
|
113
|
+
|
114
|
+
zep_messages: List[Message] = []
|
115
|
+
for m in messages:
|
116
|
+
content = (
|
117
|
+
self._truncate(deepcopy(m.get("content"))) if "content" in m else None
|
118
|
+
)
|
119
|
+
if content is None: # pragma: no cover
|
120
|
+
continue
|
121
|
+
role_type = "user" if m.get("role") == "user" else "assistant"
|
122
|
+
zep_messages.append(Message(content=content, role=role_type))
|
123
|
+
|
124
|
+
if zep_messages:
|
125
|
+
try:
|
126
|
+
await self.zep.thread.add_messages(
|
127
|
+
thread_id=user_id, messages=zep_messages
|
128
|
+
)
|
129
|
+
except Exception:
|
130
|
+
try:
|
131
|
+
try:
|
132
|
+
await self.zep.user.add(user_id=user_id)
|
133
|
+
except Exception as e:
|
134
|
+
logger.error(f"Zep user addition error: {e}")
|
135
|
+
try:
|
136
|
+
await self.zep.thread.create(thread_id=user_id, user_id=user_id)
|
137
|
+
except Exception as e:
|
138
|
+
logger.error(f"Zep thread creation error: {e}")
|
139
|
+
await self.zep.thread.add_messages(
|
140
|
+
thread_id=user_id, messages=zep_messages
|
141
|
+
)
|
142
|
+
except Exception as e:
|
143
|
+
logger.error(f"Zep memory addition error: {e}")
|
144
|
+
|
145
|
+
async def retrieve(self, user_id: str) -> str:
|
146
|
+
try:
|
147
|
+
memories = ""
|
148
|
+
if self.zep:
|
149
|
+
memory = await self.zep.thread.get_user_context(thread_id=user_id)
|
150
|
+
if memory and memory.context:
|
151
|
+
memories = memory.context
|
152
|
+
return memories
|
153
|
+
except Exception as e:
|
154
|
+
logger.error(f"Error retrieving memories: {e}")
|
155
|
+
return ""
|
156
|
+
|
157
|
+
async def delete(self, user_id: str) -> None:
|
158
|
+
if self.mongo:
|
159
|
+
try:
|
160
|
+
self.mongo.delete_all(self.collection, {"user_id": user_id})
|
161
|
+
except Exception as e:
|
162
|
+
logger.error(f"MongoDB deletion error: {e}")
|
163
|
+
if not self.zep:
|
164
|
+
return
|
165
|
+
try:
|
166
|
+
await self.zep.thread.delete(thread_id=user_id)
|
167
|
+
except Exception as e:
|
168
|
+
logger.error(f"Zep memory deletion error: {e}")
|
169
|
+
try:
|
170
|
+
await self.zep.user.delete(user_id=user_id)
|
171
|
+
except Exception as e:
|
172
|
+
logger.error(f"Zep user deletion error: {e}")
|
173
|
+
|
174
|
+
def find(
|
175
|
+
self,
|
176
|
+
collection: str,
|
177
|
+
query: Dict,
|
178
|
+
sort: Optional[List[Tuple]] = None,
|
179
|
+
limit: int = 0,
|
180
|
+
skip: int = 0,
|
181
|
+
) -> List[Dict]: # pragma: no cover
|
182
|
+
if not self.mongo:
|
183
|
+
return []
|
184
|
+
try:
|
185
|
+
return self.mongo.find(collection, query, sort=sort, limit=limit, skip=skip)
|
186
|
+
except Exception as e:
|
187
|
+
logger.error(f"MongoDB find error: {e}")
|
188
|
+
return []
|
189
|
+
|
190
|
+
def count_documents(self, collection: str, query: Dict) -> int:
|
191
|
+
if not self.mongo:
|
192
|
+
return 0
|
193
|
+
return self.mongo.count_documents(collection, query)
|
194
|
+
|
195
|
+
def _truncate(self, text: str, limit: int = 2500) -> str:
|
196
|
+
if text is None:
|
197
|
+
raise AttributeError("Cannot truncate None text")
|
198
|
+
if not text:
|
199
|
+
return ""
|
200
|
+
if len(text) <= limit:
|
201
|
+
return text
|
202
|
+
last_period = text.rfind(".", 0, limit)
|
203
|
+
if last_period > 0:
|
204
|
+
return text[: last_period + 1]
|
205
|
+
return text[: limit - 3] + "..."
|
206
|
+
|
207
|
+
async def save_capture(
|
208
|
+
self,
|
209
|
+
user_id: str,
|
210
|
+
capture_name: str,
|
211
|
+
agent_name: Optional[str],
|
212
|
+
data: Dict[str, Any],
|
213
|
+
schema: Optional[Dict[str, Any]] = None,
|
214
|
+
) -> Optional[str]:
|
215
|
+
if not self.mongo: # pragma: no cover
|
216
|
+
logger.warning("MongoDB not configured; cannot save capture.")
|
217
|
+
return None
|
218
|
+
if not user_id or not isinstance(user_id, str):
|
219
|
+
raise ValueError("user_id must be a non-empty string")
|
220
|
+
if not capture_name or not isinstance(capture_name, str):
|
221
|
+
raise ValueError("capture_name must be a non-empty string")
|
222
|
+
if not isinstance(data, dict):
|
223
|
+
raise ValueError("data must be a dictionary")
|
224
|
+
|
225
|
+
try:
|
226
|
+
mode = self.capture_modes.get(agent_name, "once") if agent_name else "once"
|
227
|
+
now = datetime.now(timezone.utc)
|
228
|
+
if mode == "multiple":
|
229
|
+
doc = {
|
230
|
+
"user_id": user_id,
|
231
|
+
"agent_name": agent_name,
|
232
|
+
"capture_name": capture_name,
|
233
|
+
"data": data or {},
|
234
|
+
"schema": schema or {},
|
235
|
+
"mode": "multiple",
|
236
|
+
"timestamp": now,
|
237
|
+
"created_at": now,
|
238
|
+
}
|
239
|
+
return self.mongo.insert_one(self.captures_collection, doc)
|
240
|
+
else:
|
241
|
+
key = {
|
242
|
+
"user_id": user_id,
|
243
|
+
"agent_name": agent_name,
|
244
|
+
"capture_name": capture_name,
|
245
|
+
}
|
246
|
+
existing = self.mongo.find_one(self.captures_collection, key)
|
247
|
+
merged_data: Dict[str, Any] = {}
|
248
|
+
if existing and isinstance(existing.get("data"), dict):
|
249
|
+
merged_data.update(existing.get("data", {}))
|
250
|
+
merged_data.update(data or {})
|
251
|
+
update_doc = {
|
252
|
+
"$set": {
|
253
|
+
"user_id": user_id,
|
254
|
+
"agent_name": agent_name,
|
255
|
+
"capture_name": capture_name,
|
256
|
+
"data": merged_data,
|
257
|
+
"schema": (
|
258
|
+
schema
|
259
|
+
if schema is not None
|
260
|
+
else existing.get("schema")
|
261
|
+
if existing
|
262
|
+
else {}
|
263
|
+
),
|
264
|
+
"mode": "once",
|
265
|
+
"timestamp": now,
|
266
|
+
},
|
267
|
+
"$setOnInsert": {"created_at": now},
|
268
|
+
}
|
269
|
+
self.mongo.update_one(
|
270
|
+
self.captures_collection, key, update_doc, upsert=True
|
271
|
+
)
|
272
|
+
doc = self.mongo.find_one(self.captures_collection, key)
|
273
|
+
return str(doc.get("_id")) if doc and doc.get("_id") else None
|
274
|
+
except Exception as e: # pragma: no cover
|
275
|
+
logger.error(f"MongoDB save_capture error: {e}")
|
276
|
+
return None
|