memoryfn 0.0.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.
- memoryfn-0.0.1/.gitignore +48 -0
- memoryfn-0.0.1/PKG-INFO +59 -0
- memoryfn-0.0.1/README.md +40 -0
- memoryfn-0.0.1/memoryfn/__init__.py +5 -0
- memoryfn-0.0.1/memoryfn/core/config.py +9 -0
- memoryfn-0.0.1/memoryfn/core/pipeline.py +90 -0
- memoryfn-0.0.1/memoryfn/core/types.py +60 -0
- memoryfn-0.0.1/memoryfn/extraction/facts.py +25 -0
- memoryfn-0.0.1/memoryfn/providers/openai.py +34 -0
- memoryfn-0.0.1/memoryfn/storage/postgres.py +138 -0
- memoryfn-0.0.1/pyproject.toml +32 -0
- memoryfn-0.0.1/tests/test_pipeline.py +58 -0
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
# Dependencies
|
|
2
|
+
node_modules/
|
|
3
|
+
.pnp
|
|
4
|
+
.pnp.js
|
|
5
|
+
|
|
6
|
+
# Build outputs
|
|
7
|
+
dist/
|
|
8
|
+
build/
|
|
9
|
+
*.tsbuildinfo
|
|
10
|
+
|
|
11
|
+
# Logs
|
|
12
|
+
logs
|
|
13
|
+
*.log
|
|
14
|
+
npm-debug.log*
|
|
15
|
+
yarn-debug.log*
|
|
16
|
+
yarn-error.log*
|
|
17
|
+
lerna-debug.log*
|
|
18
|
+
.pnpm-debug.log*
|
|
19
|
+
|
|
20
|
+
# Environment variables
|
|
21
|
+
.env
|
|
22
|
+
.env.*
|
|
23
|
+
!.env.example
|
|
24
|
+
|
|
25
|
+
# Testing
|
|
26
|
+
coverage/
|
|
27
|
+
.nyc_output
|
|
28
|
+
|
|
29
|
+
# Editors
|
|
30
|
+
.vscode/
|
|
31
|
+
.idea/
|
|
32
|
+
*.swp
|
|
33
|
+
*.swo
|
|
34
|
+
*~
|
|
35
|
+
.DS_Store
|
|
36
|
+
|
|
37
|
+
# Turbo
|
|
38
|
+
.turbo/
|
|
39
|
+
|
|
40
|
+
# Misc
|
|
41
|
+
*.tgz
|
|
42
|
+
.cache/
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
_conduct/.meta
|
|
46
|
+
|
|
47
|
+
.agent
|
|
48
|
+
.venv
|
memoryfn-0.0.1/PKG-INFO
ADDED
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: memoryfn
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: Self-hostable AI memory system for agents and applications
|
|
5
|
+
Author: 21n
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Requires-Python: >=3.10
|
|
8
|
+
Requires-Dist: httpx>=0.27
|
|
9
|
+
Requires-Dist: openai>=1.0.0
|
|
10
|
+
Requires-Dist: pgvector>=0.2.0
|
|
11
|
+
Requires-Dist: psycopg[binary]>=3.1
|
|
12
|
+
Requires-Dist: pydantic>=2.6
|
|
13
|
+
Requires-Dist: tenacity>=8.0.0
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
|
|
16
|
+
Requires-Dist: pytest>=8.0.0; extra == 'dev'
|
|
17
|
+
Requires-Dist: ruff>=0.3.0; extra == 'dev'
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
|
|
20
|
+
# MemoryFn Python SDK
|
|
21
|
+
|
|
22
|
+
> Self-hostable AI memory system for agents and applications
|
|
23
|
+
|
|
24
|
+
## Installation
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
pip install memoryfn
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
## Usage
|
|
31
|
+
|
|
32
|
+
```python
|
|
33
|
+
import asyncio
|
|
34
|
+
from memoryfn import MemoryFn, Config
|
|
35
|
+
|
|
36
|
+
async def main():
|
|
37
|
+
memory = MemoryFn(Config(
|
|
38
|
+
storage_url="postgresql://user:pass@localhost:5432/memoryfn",
|
|
39
|
+
openai_api_key="sk-..."
|
|
40
|
+
))
|
|
41
|
+
|
|
42
|
+
# Add memory
|
|
43
|
+
await memory.add(
|
|
44
|
+
content="User prefers dark mode",
|
|
45
|
+
container_tags=["user:alice"],
|
|
46
|
+
type="profile_static"
|
|
47
|
+
)
|
|
48
|
+
|
|
49
|
+
# Search
|
|
50
|
+
results = await memory.search(
|
|
51
|
+
q="preferences",
|
|
52
|
+
container_tags=["user:alice"]
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
print(results)
|
|
56
|
+
|
|
57
|
+
if __name__ == "__main__":
|
|
58
|
+
asyncio.run(main())
|
|
59
|
+
```
|
memoryfn-0.0.1/README.md
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
# MemoryFn Python SDK
|
|
2
|
+
|
|
3
|
+
> Self-hostable AI memory system for agents and applications
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install memoryfn
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Usage
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
import asyncio
|
|
15
|
+
from memoryfn import MemoryFn, Config
|
|
16
|
+
|
|
17
|
+
async def main():
|
|
18
|
+
memory = MemoryFn(Config(
|
|
19
|
+
storage_url="postgresql://user:pass@localhost:5432/memoryfn",
|
|
20
|
+
openai_api_key="sk-..."
|
|
21
|
+
))
|
|
22
|
+
|
|
23
|
+
# Add memory
|
|
24
|
+
await memory.add(
|
|
25
|
+
content="User prefers dark mode",
|
|
26
|
+
container_tags=["user:alice"],
|
|
27
|
+
type="profile_static"
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
# Search
|
|
31
|
+
results = await memory.search(
|
|
32
|
+
q="preferences",
|
|
33
|
+
container_tags=["user:alice"]
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
print(results)
|
|
37
|
+
|
|
38
|
+
if __name__ == "__main__":
|
|
39
|
+
asyncio.run(main())
|
|
40
|
+
```
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
from pydantic import BaseModel, Field
|
|
2
|
+
from typing import Optional, Literal, Dict
|
|
3
|
+
|
|
4
|
+
class Config(BaseModel):
|
|
5
|
+
storage_url: str = Field(alias="storageUrl")
|
|
6
|
+
openai_api_key: Optional[str] = Field(default=None, alias="openaiApiKey")
|
|
7
|
+
openai_model: str = Field(default="gpt-4o-mini", alias="openaiModel")
|
|
8
|
+
embedding_model: str = Field(default="text-embedding-3-small", alias="embeddingModel")
|
|
9
|
+
embedding_dims: int = Field(default=1536, alias="embeddingDims")
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
from memoryfn.core.config import Config
|
|
2
|
+
from memoryfn.core.types import AddMemoryInput, SearchMemoryInput, SearchMemoryResult, Memory
|
|
3
|
+
from memoryfn.storage.postgres import PostgresAdapter
|
|
4
|
+
from memoryfn.providers.openai import OpenAIProvider
|
|
5
|
+
from memoryfn.extraction.facts import FactExtractor
|
|
6
|
+
import time
|
|
7
|
+
|
|
8
|
+
class MemoryFn:
|
|
9
|
+
def __init__(self, config: Config):
|
|
10
|
+
self.config = config
|
|
11
|
+
self.storage = PostgresAdapter(config.storage_url)
|
|
12
|
+
self.provider = None
|
|
13
|
+
self.extractor = None
|
|
14
|
+
|
|
15
|
+
if config.openai_api_key:
|
|
16
|
+
self.provider = OpenAIProvider(
|
|
17
|
+
api_key=config.openai_api_key,
|
|
18
|
+
embedding_model=config.embedding_model,
|
|
19
|
+
llm_model=config.openai_model
|
|
20
|
+
)
|
|
21
|
+
self.extractor = FactExtractor(self.provider)
|
|
22
|
+
|
|
23
|
+
async def add(self, input: AddMemoryInput) -> List[Memory]:
|
|
24
|
+
# 1. Normalize
|
|
25
|
+
content = input.content
|
|
26
|
+
memories_to_persist = []
|
|
27
|
+
|
|
28
|
+
# 2. Extract
|
|
29
|
+
if self.extractor:
|
|
30
|
+
facts = await self.extractor.extract(content)
|
|
31
|
+
if facts:
|
|
32
|
+
for fact in facts:
|
|
33
|
+
memories_to_persist.append(Memory(
|
|
34
|
+
tenantId=input.tenant_id,
|
|
35
|
+
containerTags=input.container_tags + fact.get("tags", []),
|
|
36
|
+
type=fact.get("type", "conversational"),
|
|
37
|
+
content=fact.get("content"),
|
|
38
|
+
metadata={**input.metadata, "confidence": fact.get("confidence", 1.0)}
|
|
39
|
+
))
|
|
40
|
+
else:
|
|
41
|
+
# Fallback
|
|
42
|
+
memories_to_persist.append(Memory(
|
|
43
|
+
tenantId=input.tenant_id,
|
|
44
|
+
containerTags=input.container_tags,
|
|
45
|
+
type=input.type,
|
|
46
|
+
content=content,
|
|
47
|
+
metadata=input.metadata
|
|
48
|
+
))
|
|
49
|
+
else:
|
|
50
|
+
memories_to_persist.append(Memory(
|
|
51
|
+
tenantId=input.tenant_id,
|
|
52
|
+
containerTags=input.container_tags,
|
|
53
|
+
type=input.type,
|
|
54
|
+
content=content,
|
|
55
|
+
metadata=input.metadata
|
|
56
|
+
))
|
|
57
|
+
|
|
58
|
+
# 3. Embed
|
|
59
|
+
if self.provider:
|
|
60
|
+
texts = [m.content for m in memories_to_persist]
|
|
61
|
+
embeddings = await self.provider.embed_batch(texts)
|
|
62
|
+
for i, mem in enumerate(memories_to_persist):
|
|
63
|
+
mem.embedding = embeddings[i]
|
|
64
|
+
|
|
65
|
+
# 4. Persist
|
|
66
|
+
saved = await self.storage.insert_memories(memories_to_persist)
|
|
67
|
+
return saved
|
|
68
|
+
|
|
69
|
+
async def search(self, input: SearchMemoryInput) -> SearchMemoryResult:
|
|
70
|
+
embedding = []
|
|
71
|
+
if self.provider:
|
|
72
|
+
embedding = await self.provider.embed(input.q)
|
|
73
|
+
else:
|
|
74
|
+
# Mock or error
|
|
75
|
+
embedding = [0.0] * self.config.embedding_dims
|
|
76
|
+
|
|
77
|
+
results = await self.storage.search_vectors(
|
|
78
|
+
embedding=embedding,
|
|
79
|
+
container_tags=input.container_tags,
|
|
80
|
+
top_k=input.limit,
|
|
81
|
+
threshold=input.threshold
|
|
82
|
+
)
|
|
83
|
+
|
|
84
|
+
return SearchMemoryResult(
|
|
85
|
+
results=results,
|
|
86
|
+
totalFound=len(results)
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
async def close(self):
|
|
90
|
+
await self.storage.close()
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
from typing import List, Optional, Dict, Any, Literal
|
|
2
|
+
from pydantic import BaseModel, Field
|
|
3
|
+
from datetime import datetime
|
|
4
|
+
import time
|
|
5
|
+
|
|
6
|
+
MemoryType = Literal[
|
|
7
|
+
"profile_static",
|
|
8
|
+
"profile_dynamic",
|
|
9
|
+
"conversational",
|
|
10
|
+
"procedural",
|
|
11
|
+
"document",
|
|
12
|
+
"derived"
|
|
13
|
+
]
|
|
14
|
+
|
|
15
|
+
RelationType = Literal["updates", "extends", "contradicts", "derives", "related"]
|
|
16
|
+
|
|
17
|
+
class Memory(BaseModel):
|
|
18
|
+
id: Optional[str] = None
|
|
19
|
+
tenant_id: str = Field(default="default", alias="tenantId")
|
|
20
|
+
container_tags: List[str] = Field(default_factory=list, alias="containerTags")
|
|
21
|
+
type: MemoryType = "conversational"
|
|
22
|
+
content: str
|
|
23
|
+
embedding: Optional[List[float]] = None
|
|
24
|
+
metadata: Dict[str, Any] = Field(default_factory=dict)
|
|
25
|
+
is_latest: bool = Field(default=True, alias="isLatest")
|
|
26
|
+
created_at: float = Field(default_factory=lambda: time.time() * 1000, alias="createdAt")
|
|
27
|
+
updated_at: float = Field(default_factory=lambda: time.time() * 1000, alias="updatedAt")
|
|
28
|
+
|
|
29
|
+
class Config:
|
|
30
|
+
populate_by_name = True
|
|
31
|
+
|
|
32
|
+
class MemoryRelationship(BaseModel):
|
|
33
|
+
id: Optional[str] = None
|
|
34
|
+
from_id: str = Field(alias="fromId")
|
|
35
|
+
to_id: str = Field(alias="toId")
|
|
36
|
+
type: RelationType
|
|
37
|
+
confidence: float = 1.0
|
|
38
|
+
reasoning: Optional[str] = None
|
|
39
|
+
created_at: float = Field(default_factory=lambda: time.time() * 1000, alias="createdAt")
|
|
40
|
+
|
|
41
|
+
class Config:
|
|
42
|
+
populate_by_name = True
|
|
43
|
+
|
|
44
|
+
class AddMemoryInput(BaseModel):
|
|
45
|
+
content: str
|
|
46
|
+
container_tags: List[str] = Field(default_factory=list, alias="containerTags")
|
|
47
|
+
type: MemoryType = "conversational"
|
|
48
|
+
metadata: Dict[str, Any] = Field(default_factory=dict)
|
|
49
|
+
tenant_id: str = Field(default="default", alias="tenantId")
|
|
50
|
+
|
|
51
|
+
class SearchMemoryInput(BaseModel):
|
|
52
|
+
q: str
|
|
53
|
+
container_tags: List[str] = Field(default_factory=list, alias="containerTags")
|
|
54
|
+
filters: Optional[Dict[str, Any]] = None
|
|
55
|
+
limit: int = 10
|
|
56
|
+
threshold: float = 0.7
|
|
57
|
+
|
|
58
|
+
class SearchMemoryResult(BaseModel):
|
|
59
|
+
results: List[Memory]
|
|
60
|
+
total_found: int = Field(alias="totalFound")
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
from memoryfn.providers.openai import OpenAIProvider
|
|
2
|
+
from typing import List, Dict
|
|
3
|
+
|
|
4
|
+
class FactExtractor:
|
|
5
|
+
def __init__(self, provider: OpenAIProvider):
|
|
6
|
+
self.provider = provider
|
|
7
|
+
|
|
8
|
+
async def extract(self, text: str) -> List[Dict]:
|
|
9
|
+
prompt = f"""
|
|
10
|
+
Analyze the following text and extract discrete, atomic facts.
|
|
11
|
+
Return a JSON object with a key "facts" containing an array of objects.
|
|
12
|
+
Each object should have:
|
|
13
|
+
- "content": The fact as a standalone sentence.
|
|
14
|
+
- "type": One of "profile_static", "profile_dynamic", "conversational".
|
|
15
|
+
- "confidence": A number between 0 and 1.
|
|
16
|
+
- "tags": Array of relevant keywords.
|
|
17
|
+
|
|
18
|
+
Text: "{text}"
|
|
19
|
+
"""
|
|
20
|
+
try:
|
|
21
|
+
result = await self.provider.generate_json(prompt)
|
|
22
|
+
return result.get("facts", [])
|
|
23
|
+
except Exception as e:
|
|
24
|
+
print(f"Extraction failed: {e}")
|
|
25
|
+
return []
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
from typing import List
|
|
2
|
+
from openai import AsyncOpenAI
|
|
3
|
+
|
|
4
|
+
class OpenAIProvider:
|
|
5
|
+
def __init__(self, api_key: str, embedding_model: str, llm_model: str):
|
|
6
|
+
self.client = AsyncOpenAI(api_key=api_key)
|
|
7
|
+
self.embedding_model = embedding_model
|
|
8
|
+
self.llm_model = llm_model
|
|
9
|
+
|
|
10
|
+
async def embed(self, text: str) -> List[float]:
|
|
11
|
+
resp = await self.client.embeddings.create(
|
|
12
|
+
model=self.embedding_model,
|
|
13
|
+
input=text
|
|
14
|
+
)
|
|
15
|
+
return resp.data[0].embedding
|
|
16
|
+
|
|
17
|
+
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
|
18
|
+
resp = await self.client.embeddings.create(
|
|
19
|
+
model=self.embedding_model,
|
|
20
|
+
input=texts
|
|
21
|
+
)
|
|
22
|
+
return [d.embedding for d in resp.data]
|
|
23
|
+
|
|
24
|
+
async def generate_json(self, prompt: str) -> dict:
|
|
25
|
+
resp = await self.client.chat.completions.create(
|
|
26
|
+
model=self.llm_model,
|
|
27
|
+
messages=[
|
|
28
|
+
{"role": "system", "content": "You are a helpful assistant that outputs JSON."},
|
|
29
|
+
{"role": "user", "content": prompt}
|
|
30
|
+
],
|
|
31
|
+
response_format={"type": "json_object"}
|
|
32
|
+
)
|
|
33
|
+
import json
|
|
34
|
+
return json.loads(resp.choices[0].message.content)
|
|
@@ -0,0 +1,138 @@
|
|
|
1
|
+
from typing import List, Optional
|
|
2
|
+
import psycopg
|
|
3
|
+
from pgvector.psycopg import register_vector
|
|
4
|
+
from memoryfn.core.types import Memory, MemoryRelationship
|
|
5
|
+
import json
|
|
6
|
+
from datetime import datetime
|
|
7
|
+
|
|
8
|
+
class PostgresAdapter:
|
|
9
|
+
def __init__(self, dsn: str):
|
|
10
|
+
self.dsn = dsn
|
|
11
|
+
self._conn = None
|
|
12
|
+
|
|
13
|
+
async def connect(self):
|
|
14
|
+
if not self._conn:
|
|
15
|
+
# Note: In a real async app we'd use psycopg_pool or AsyncConnection
|
|
16
|
+
# Here we use standard sync connection for simplicity or psycopg.AsyncConnection
|
|
17
|
+
self._conn = await psycopg.AsyncConnection.connect(self.dsn, autocommit=True)
|
|
18
|
+
# Enable vector extension support
|
|
19
|
+
await register_vector(self._conn)
|
|
20
|
+
|
|
21
|
+
async def insert_memories(self, memories: List[Memory]) -> List[Memory]:
|
|
22
|
+
await self.connect()
|
|
23
|
+
async with self._conn.cursor() as cur:
|
|
24
|
+
results = []
|
|
25
|
+
for mem in memories:
|
|
26
|
+
# Using timestamp in ms for domain, but DB uses timestamptz
|
|
27
|
+
created_at = datetime.fromtimestamp(mem.created_at / 1000.0)
|
|
28
|
+
updated_at = datetime.fromtimestamp(mem.updated_at / 1000.0)
|
|
29
|
+
|
|
30
|
+
await cur.execute(
|
|
31
|
+
"""
|
|
32
|
+
INSERT INTO memories (
|
|
33
|
+
tenant_id, container_tags, type, content, embedding,
|
|
34
|
+
metadata, is_latest, created_at, updated_at
|
|
35
|
+
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
36
|
+
RETURNING id
|
|
37
|
+
""",
|
|
38
|
+
(
|
|
39
|
+
mem.tenant_id,
|
|
40
|
+
mem.container_tags,
|
|
41
|
+
mem.type,
|
|
42
|
+
mem.content,
|
|
43
|
+
mem.embedding,
|
|
44
|
+
json.dumps(mem.metadata),
|
|
45
|
+
mem.is_latest,
|
|
46
|
+
created_at,
|
|
47
|
+
updated_at
|
|
48
|
+
)
|
|
49
|
+
)
|
|
50
|
+
row = await cur.fetchone()
|
|
51
|
+
mem.id = str(row[0])
|
|
52
|
+
results.append(mem)
|
|
53
|
+
return results
|
|
54
|
+
|
|
55
|
+
async def insert_relationships(self, relationships: List[MemoryRelationship]) -> List[MemoryRelationship]:
|
|
56
|
+
await self.connect()
|
|
57
|
+
async with self._conn.cursor() as cur:
|
|
58
|
+
results = []
|
|
59
|
+
for rel in relationships:
|
|
60
|
+
created_at = datetime.fromtimestamp(rel.created_at / 1000.0)
|
|
61
|
+
await cur.execute(
|
|
62
|
+
"""
|
|
63
|
+
INSERT INTO memory_relationships (
|
|
64
|
+
from_id, to_id, type, confidence, reasoning, created_at
|
|
65
|
+
) VALUES (%s, %s, %s, %s, %s, %s)
|
|
66
|
+
RETURNING id
|
|
67
|
+
""",
|
|
68
|
+
(
|
|
69
|
+
rel.from_id,
|
|
70
|
+
rel.to_id,
|
|
71
|
+
rel.type,
|
|
72
|
+
rel.confidence,
|
|
73
|
+
rel.reasoning,
|
|
74
|
+
created_at
|
|
75
|
+
)
|
|
76
|
+
)
|
|
77
|
+
row = await cur.fetchone()
|
|
78
|
+
rel.id = str(row[0])
|
|
79
|
+
results.append(rel)
|
|
80
|
+
return results
|
|
81
|
+
|
|
82
|
+
async def search_vectors(
|
|
83
|
+
self,
|
|
84
|
+
embedding: List[float],
|
|
85
|
+
container_tags: List[str],
|
|
86
|
+
top_k: int = 10,
|
|
87
|
+
threshold: float = 0.7
|
|
88
|
+
) -> List[Memory]:
|
|
89
|
+
await self.connect()
|
|
90
|
+
|
|
91
|
+
# Build query
|
|
92
|
+
# Using <=> for cosine distance. Distance = 1 - similarity.
|
|
93
|
+
# Similarity >= threshold => Distance <= 1 - threshold
|
|
94
|
+
max_dist = 1.0 - threshold
|
|
95
|
+
|
|
96
|
+
query = """
|
|
97
|
+
SELECT id, tenant_id, container_tags, type, content, embedding,
|
|
98
|
+
metadata, is_latest, created_at, updated_at,
|
|
99
|
+
1 - (embedding <=> %s::vector) as similarity
|
|
100
|
+
FROM memories
|
|
101
|
+
WHERE 1=1
|
|
102
|
+
"""
|
|
103
|
+
params = [embedding]
|
|
104
|
+
|
|
105
|
+
if container_tags:
|
|
106
|
+
query += " AND container_tags && %s"
|
|
107
|
+
params.append(container_tags)
|
|
108
|
+
|
|
109
|
+
if threshold:
|
|
110
|
+
query += f" AND (embedding <=> %s::vector) <= {max_dist}"
|
|
111
|
+
params.append(embedding)
|
|
112
|
+
|
|
113
|
+
query += " ORDER BY embedding <=> %s::vector LIMIT %s"
|
|
114
|
+
params.extend([embedding, top_k])
|
|
115
|
+
|
|
116
|
+
async with self._conn.cursor() as cur:
|
|
117
|
+
await cur.execute(query, params)
|
|
118
|
+
rows = await cur.fetchall()
|
|
119
|
+
|
|
120
|
+
results = []
|
|
121
|
+
for row in rows:
|
|
122
|
+
results.append(Memory(
|
|
123
|
+
id=str(row[0]),
|
|
124
|
+
tenantId=row[1],
|
|
125
|
+
containerTags=row[2],
|
|
126
|
+
type=row[3],
|
|
127
|
+
content=row[4],
|
|
128
|
+
embedding=row[5].tolist() if hasattr(row[5], 'tolist') else row[5],
|
|
129
|
+
metadata=row[6],
|
|
130
|
+
isLatest=row[7],
|
|
131
|
+
createdAt=row[8].timestamp() * 1000,
|
|
132
|
+
updatedAt=row[9].timestamp() * 1000
|
|
133
|
+
))
|
|
134
|
+
return results
|
|
135
|
+
|
|
136
|
+
async def close(self):
|
|
137
|
+
if self._conn:
|
|
138
|
+
await self._conn.close()
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "memoryfn"
|
|
7
|
+
version = "0.0.1"
|
|
8
|
+
description = "Self-hostable AI memory system for agents and applications"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "21n" },
|
|
14
|
+
]
|
|
15
|
+
dependencies = [
|
|
16
|
+
"pydantic>=2.6",
|
|
17
|
+
"httpx>=0.27",
|
|
18
|
+
"pgvector>=0.2.0",
|
|
19
|
+
"psycopg[binary]>=3.1",
|
|
20
|
+
"openai>=1.0.0",
|
|
21
|
+
"tenacity>=8.0.0"
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.optional-dependencies]
|
|
25
|
+
dev = [
|
|
26
|
+
"pytest>=8.0.0",
|
|
27
|
+
"pytest-asyncio>=0.23.0",
|
|
28
|
+
"ruff>=0.3.0"
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
[tool.hatch.build.targets.wheel]
|
|
32
|
+
packages = ["memoryfn"]
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
import pytest
|
|
2
|
+
from unittest.mock import AsyncMock, MagicMock
|
|
3
|
+
from memoryfn.core.pipeline import MemoryFn
|
|
4
|
+
from memoryfn.core.config import Config
|
|
5
|
+
from memoryfn.core.types import AddMemoryInput, SearchMemoryInput, Memory
|
|
6
|
+
|
|
7
|
+
@pytest.mark.asyncio
|
|
8
|
+
async def test_memoryfn_pipeline():
|
|
9
|
+
# Mock Config
|
|
10
|
+
config = Config(
|
|
11
|
+
storageUrl="postgresql://mock:mock@localhost:5432/mock",
|
|
12
|
+
openaiApiKey="mock-key"
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
# Mock Storage
|
|
16
|
+
mock_storage = AsyncMock()
|
|
17
|
+
mock_storage.insert_memories.return_value = [
|
|
18
|
+
Memory(id="1", content="Test content", type="conversational")
|
|
19
|
+
]
|
|
20
|
+
mock_storage.search_vectors.return_value = [
|
|
21
|
+
Memory(id="1", content="Test content", type="conversational")
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
# Mock Provider
|
|
25
|
+
mock_provider = AsyncMock()
|
|
26
|
+
mock_provider.embed.return_value = [0.1] * 1536
|
|
27
|
+
mock_provider.embed_batch.return_value = [[0.1] * 1536]
|
|
28
|
+
mock_provider.generate_json.return_value = {"facts": [{"content": "Test content", "type": "conversational", "confidence": 0.9}]}
|
|
29
|
+
|
|
30
|
+
# Initialize MemoryFn and inject mocks
|
|
31
|
+
memory = MemoryFn(config)
|
|
32
|
+
memory.storage = mock_storage
|
|
33
|
+
memory.provider = mock_provider
|
|
34
|
+
memory.extractor.provider = mock_provider # type: ignore
|
|
35
|
+
|
|
36
|
+
# Test Add
|
|
37
|
+
input_data = AddMemoryInput(
|
|
38
|
+
content="Test content",
|
|
39
|
+
containerTags=["user:test"],
|
|
40
|
+
tenantId="test-tenant"
|
|
41
|
+
)
|
|
42
|
+
result = await memory.add(input_data)
|
|
43
|
+
|
|
44
|
+
assert len(result) == 1
|
|
45
|
+
assert result[0].content == "Test content"
|
|
46
|
+
mock_storage.insert_memories.assert_called_once()
|
|
47
|
+
mock_provider.embed_batch.assert_called_once()
|
|
48
|
+
|
|
49
|
+
# Test Search
|
|
50
|
+
search_input = SearchMemoryInput(
|
|
51
|
+
q="query",
|
|
52
|
+
containerTags=["user:test"]
|
|
53
|
+
)
|
|
54
|
+
search_result = await memory.search(search_input)
|
|
55
|
+
|
|
56
|
+
assert len(search_result.results) == 1
|
|
57
|
+
assert search_result.total_found == 1
|
|
58
|
+
mock_storage.search_vectors.assert_called_once()
|