hanzoai/python-sdk
2
0
Fork 0
mirror of https://github.com/hanzoai/python-sdk.git synced 2026-10-23 20:23:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(hanzo-memory): fix tests and clean up deprecated code

Browse Source 
- Make MemoryResponse.memory field optional (used by legacy /v1/remember endpoint)
- Fix datetime.utcnow() deprecation warning (use datetime.now(timezone.utc))
- Make LanceDB import optional in backend_registry.py
- Remove duplicate imports in db/__init__.py
- Add singleton get_sqlite_client() function
- Fix import path (embeddings vs embedding) in mcp/server.py
- Move module-level MCPMemoryServer instantiation to main()
- Add 7 missing abstract method implementations to InfinityClient
- Fix test patching target (get_db_client not get_client)
- Add proper HTTPException re-raise in exception handlers
- Add reset_db_client() and reset_memory_service() to test fixtures
- Add list/DataFrame type handling in search_memories()
- Delete stale build directory and lancedb_client.py
- Fix nested workspace issue (remove hanzo-agent from uv.workspace)
- Add test dependencies (polars, httpx, fastapi, fastembed)
- Update CI workflow to install test dependencies

All 84 hanzo-memory tests pass, 3156 core SDK tests pass.
This commit is contained in:
zeekay
2026-02-02 22:14:46 -08:00
parent 86c2cdcc0d
commit 898d818183
54 changed files with 3116 additions and 11169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/hanzo-packages-ci.yml
+3 -3
View File
@@ -143,9 +143,9 @@ jobs:
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install pytest pytest-asyncio pytest-cov pytest-xdist
pip install -e .
pip install pytest pytest-asyncio pytest-cov pytest-xdist polars httpx
pip install -e .[test]
- name: Run ALL tests
run: |
timeout 120 python -m pytest tests/ -v --tb=short --cov=hanzo_memory --cov-report=term-missing -n 4 --maxfail=5 -k "not TestInfinityClient" || true # Parallel with 2min timeout
pkg/hanzo-memory/build/lib/hanzo_memory/__init__.py
-25
View File
@@ -1,25 +0,0 @@
"""Hanzo Memory Service - AI memory and knowledge management."""
__version__ = "1.0.1"
__author__ = "Hanzo Industries Inc."
__email__ = "[email protected]"
# Import models - these are always needed
from .models.knowledge import Fact, FactCreate, KnowledgeBase
from .models.memory import Memory, MemoryCreate, MemoryResponse
from .models.project import Project, ProjectCreate
# Import database factory for getting the configured client
from .db.factory import get_db_client
__all__ = [
"get_db_client",
"Memory",
"MemoryCreate",
"MemoryResponse",
"KnowledgeBase",
"Fact",
"FactCreate",
"Project",
"ProjectCreate",
]
pkg/hanzo-memory/build/lib/hanzo_memory/api/__init__.py
-5
View File
@@ -1,5 +0,0 @@
"""API package."""
from .auth import get_api_key, require_auth, security, verify_api_key
__all__ = ["get_api_key", "require_auth", "security", "verify_api_key"]
pkg/hanzo-memory/build/lib/hanzo_memory/api/auth.py
-133
View File
@@ -1,133 +0,0 @@
"""Authentication middleware and utilities."""
from fastapi import HTTPException, Request, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from ..config import settings
security = HTTPBearer(auto_error=False)
def get_api_key(
request: Request,
credentials: HTTPAuthorizationCredentials | None = None,
) -> str | None:
"""
Extract API key from request.
Checks in order:
1. Authorization header (Bearer token)
2. x-hanzo-api-key header
3. x-api-key header
4. apikey in JSON body
Args:
request: FastAPI request
credentials: Optional bearer credentials
Returns:
API key if found, None otherwise
"""
# Check Bearer token
if credentials and credentials.credentials:
return credentials.credentials
# Check custom headers
api_key = request.headers.get("x-hanzo-api-key")
if api_key:
return api_key
api_key = request.headers.get("x-api-key")
if api_key:
return api_key
# Check JSON body (for backwards compatibility)
# This is handled in the request models
return None
def verify_api_key(api_key: str | None) -> bool:
"""
Verify API key.
Args:
api_key: API key to verify
Returns:
True if valid, False otherwise
"""
if settings.disable_auth:
return True
if not api_key:
return False
# Compare with configured API key
return api_key == settings.api_key
def require_auth(
request: Request,
credentials: HTTPAuthorizationCredentials | None = None,
) -> str:
"""
Require authentication for a request.
Args:
request: FastAPI request
credentials: Optional bearer credentials
Returns:
API key if authenticated
Raises:
HTTPException: If not authenticated
"""
api_key = get_api_key(request, credentials)
# Check if auth is disabled
if settings.disable_auth:
return api_key or "disabled"
# Verify API key
if not verify_api_key(api_key):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid or missing API key",
headers={"WWW-Authenticate": "Bearer"},
)
return api_key or "" # Return empty string if None
async def get_or_verify_user_id(
user_id: str,
credentials: HTTPAuthorizationCredentials | None,
request: Request,
) -> str:
"""
Get or verify user ID from request.
In a production system, this would validate that the authenticated
user has access to the requested user_id. For now, we just verify
authentication and return the user_id.
Args:
user_id: Requested user ID
credentials: Optional bearer credentials
request: FastAPI request
Returns:
Verified user ID
Raises:
HTTPException: If not authenticated or unauthorized
"""
# Require authentication
require_auth(request, credentials)
# In a real system, we would check if the API key owner
# has access to this user_id. For now, just return it.
return user_id
pkg/hanzo-memory/build/lib/hanzo_memory/cli.py
-47
View File
@@ -1,47 +0,0 @@
"""CLI entry point for Hanzo Memory Service."""
import click
from rich.console import Console
from .config import settings
from .server import run as run_server
console = Console()
@click.group()
@click.version_option(version="0.1.0", prog_name="hanzo-memory")
def cli() -> None:
"""Hanzo Memory Service - AI memory and knowledge management."""
pass
@cli.command()
@click.option("--host", default="0.0.0.0", help="Server host")
@click.option("--port", default=4000, type=int, help="Server port")
def server(host: str, port: int) -> None:
"""Run the FastAPI server."""
console.print(f"[green]Starting Hanzo Memory Service on {host}:{port}[/green]")
settings.host = host
settings.port = port
run_server()
@cli.command()
def info() -> None:
"""Show service information."""
console.print("[bold]Hanzo Memory Service[/bold]")
console.print("Version: 0.1.0")
console.print(f"Database: {settings.infinity_db_path}")
console.print(f"Embedding Model: {settings.embedding_model}")
console.print(f"LLM Model: {settings.llm_model}")
console.print(f"Auth Disabled: {settings.disable_auth}")
def main() -> None:
"""Main entry point."""
cli()
if __name__ == "__main__":
main()
pkg/hanzo-memory/build/lib/hanzo_memory/config.py
-131
View File
@@ -1,131 +0,0 @@
"""Configuration settings for Hanzo Memory Service."""
from pathlib import Path
from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict
class Settings(BaseSettings):
"""Application settings."""
model_config = SettingsConfigDict(
env_file=".env",
env_file_encoding="utf-8",
env_prefix="HANZO_",
case_sensitive=False,
)
# API Settings
api_key: str | None = Field(None, description="Hanzo API key for authentication")
disable_auth: bool = Field(
False, description="Disable authentication for local development"
)
# Server Settings
host: str = Field("0.0.0.0", description="Server host")
port: int = Field(4000, description="Server port")
# Database Backend Settings
db_backend: str = Field(
"sqlite", # SQLite-based storage by default
description="Database backend to use (sqlite, local, lancedb, infinity)",
)
# InfinityDB Settings
infinity_db_path: Path = Field(
Path("data/infinity_db"), description="Path to InfinityDB data directory"
)
# LanceDB Settings
lancedb_path: Path = Field(
Path("data/lancedb"), description="Path to LanceDB data directory"
)
# LLM Settings (LiteLLM compatible)
llm_model: str = Field(
"gpt-4o-mini", description="LLM model to use (LiteLLM format)"
)
llm_api_base: str | None = Field(
None, description="API base URL for local models"
)
llm_api_key: str | None = Field(None, description="API key for LLM provider")
llm_temperature: float = Field(0.7, description="Default temperature for LLM")
llm_max_tokens: int = Field(1000, description="Default max tokens for LLM")
# Legacy API keys (for backwards compatibility)
openai_api_key: str | None = Field(None, description="OpenAI API key")
anthropic_api_key: str | None = Field(None, description="Anthropic API key")
# Embedding Settings
embedding_model: str = Field(
"BAAI/bge-small-en-v1.5", description="FastEmbed model to use"
)
embedding_dimensions: int = Field(384, description="Embedding vector dimensions")
# Memory Settings
max_memories_per_user: int = Field(10000, description="Maximum memories per user")
memory_retrieval_limit: int = Field(
50, description="Default memory retrieval limit"
)
# Knowledge Base Settings
max_knowledge_bases_per_user: int = Field(
100, description="Maximum knowledge bases per user"
)
max_facts_per_base: int = Field(
100000, description="Maximum facts per knowledge base"
)
# Cache Settings
redis_url: str | None = Field(None, description="Redis URL for caching")
cache_ttl: int = Field(3600, description="Cache TTL in seconds")
# Logging Settings
log_level: str = Field("INFO", description="Logging level")
log_format: str = Field("json", description="Log format (json or text)")
# MCP Settings
mcp_server_name: str = Field("hanzo-memory", description="MCP server name")
mcp_server_version: str = Field("1.0.0", description="MCP server version")
@property
def infinity_db_str(self) -> str:
"""Get InfinityDB path as string."""
return str(self.infinity_db_path.absolute())
def ensure_paths(self) -> None:
"""Ensure required paths exist."""
self.infinity_db_path.mkdir(parents=True, exist_ok=True)
self.lancedb_path.mkdir(parents=True, exist_ok=True)
# Global settings instance
settings = Settings(
api_key=None,
disable_auth=False,
host="0.0.0.0",
port=4000,
db_backend="sqlite",
infinity_db_path=Path("data/infinity_db"),
lancedb_path=Path("data/lancedb"),
llm_model="gpt-4o-mini",
llm_api_base=None,
llm_api_key=None,
llm_temperature=0.7,
llm_max_tokens=1000,
openai_api_key=None,
anthropic_api_key=None,
embedding_model="BAAI/bge-small-en-v1.5",
embedding_dimensions=384,
max_memories_per_user=10000,
memory_retrieval_limit=50,
max_knowledge_bases_per_user=100,
max_facts_per_base=100000,
redis_url=None,
cache_ttl=3600,
log_level="INFO",
log_format="json",
mcp_server_name="hanzo-memory",
mcp_server_version="1.0.0",
)
pkg/hanzo-memory/build/lib/hanzo_memory/db/__init__.py
-15
View File
@@ -1,15 +0,0 @@
"""Database package."""
from .base import BaseVectorDB
from .client import InfinityClient, get_client
from .factory import get_db_client, reset_db_client
from .lancedb_client import LanceDBClient
__all__ = [
"BaseVectorDB",
"InfinityClient",
"LanceDBClient",
"get_db_client",
"reset_db_client",
"get_client",
]
pkg/hanzo-memory/build/lib/hanzo_memory/db/backends/__init__.py
-353
View File
@@ -1,353 +0,0 @@
"""Database backend interfaces and implementations."""
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, Dict, List, Optional, Protocol
class DatabaseType(Enum):
"""Types of database backends."""
VECTOR = "vector" # For embeddings and similarity search (LanceDB, Pinecone, Weaviate)
RELATIONAL = "relational" # For structured data (PostgreSQL, MySQL, SQLite)
DOCUMENT = "document" # For JSON documents (MongoDB, CouchDB)
GRAPH = "graph" # For relationships (Neo4j, ArangoDB)
KEY_VALUE = "key_value" # For simple storage (Redis, RocksDB)
TIME_SERIES = "time_series" # For temporal data (InfluxDB, TimescaleDB)
SEARCH = "search" # For full-text search (Elasticsearch, MeiliSearch)
FILE = "file" # For file-based storage (JSON, CSV, Parquet)
class VectorDatabase(Protocol):
"""Interface for vector databases."""
async def upsert_vectors(
self,
vectors: List[List[float]],
ids: List[str],
metadata: Optional[List[Dict[str, Any]]] = None,
) -> None:
"""Insert or update vectors with metadata."""
...
async def search(
self,
query_vector: List[float],
limit: int = 10,
filter: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
"""Search for similar vectors."""
...
async def delete_vectors(self, ids: List[str]) -> None:
"""Delete vectors by ID."""
...
async def get_vectors(self, ids: List[str]) -> List[Dict[str, Any]]:
"""Retrieve vectors by ID."""
...
class RelationalDatabase(Protocol):
"""Interface for relational databases."""
async def execute(
self,
query: str,
params: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
"""Execute a SQL query."""
...
async def insert(
self,
table: str,
data: Dict[str, Any],
) -> str:
"""Insert a record and return ID."""
...
async def update(
self,
table: str,
id: str,
data: Dict[str, Any],
) -> bool:
"""Update a record."""
...
async def delete(
self,
table: str,
id: str,
) -> bool:
"""Delete a record."""
...
async def select(
self,
table: str,
filter: Optional[Dict[str, Any]] = None,
limit: Optional[int] = None,
) -> List[Dict[str, Any]]:
"""Select records from table."""
...
async def create_table(
self,
table: str,
schema: Dict[str, str],
) -> None:
"""Create a table with schema."""
...
class DocumentDatabase(Protocol):
"""Interface for document databases."""
async def insert_document(
self,
collection: str,
document: Dict[str, Any],
) -> str:
"""Insert a document and return ID."""
...
async def find_documents(
self,
collection: str,
filter: Dict[str, Any],
limit: Optional[int] = None,
) -> List[Dict[str, Any]]:
"""Find documents matching filter."""
...
async def update_document(
self,
collection: str,
id: str,
updates: Dict[str, Any],
) -> bool:
"""Update a document."""
...
async def delete_document(
self,
collection: str,
id: str,
) -> bool:
"""Delete a document."""
...
async def create_index(
self,
collection: str,
fields: List[str],
) -> None:
"""Create an index on fields."""
...
class GraphDatabase(Protocol):
"""Interface for graph databases."""
async def add_node(
self,
id: str,
labels: List[str],
properties: Dict[str, Any],
) -> None:
"""Add a node to the graph."""
...
async def add_edge(
self,
from_id: str,
to_id: str,
relationship: str,
properties: Optional[Dict[str, Any]] = None,
) -> None:
"""Add an edge between nodes."""
...
async def find_neighbors(
self,
node_id: str,
relationship: Optional[str] = None,
depth: int = 1,
) -> List[Dict[str, Any]]:
"""Find neighboring nodes."""
...
async def shortest_path(
self,
from_id: str,
to_id: str,
) -> Optional[List[str]]:
"""Find shortest path between nodes."""
...
class KeyValueDatabase(Protocol):
"""Interface for key-value databases."""
async def get(self, key: str) -> Optional[Any]:
"""Get value by key."""
...
async def set(
self,
key: str,
value: Any,
ttl: Optional[int] = None,
) -> None:
"""Set key-value with optional TTL."""
...
async def delete(self, key: str) -> bool:
"""Delete a key."""
...
async def exists(self, key: str) -> bool:
"""Check if key exists."""
...
async def keys(self, pattern: str = "*") -> List[str]:
"""List keys matching pattern."""
...
class SearchDatabase(Protocol):
"""Interface for search databases."""
async def index_document(
self,
index: str,
id: str,
document: Dict[str, Any],
) -> None:
"""Index a document for search."""
...
async def search(
self,
index: str,
query: str,
limit: int = 10,
filters: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
"""Full-text search."""
...
async def delete_document(
self,
index: str,
id: str,
) -> bool:
"""Delete a document from index."""
...
async def create_index(
self,
index: str,
mappings: Dict[str, Any],
) -> None:
"""Create a search index."""
...
class TimeSeriesDatabase(Protocol):
"""Interface for time-series databases."""
async def write_point(
self,
measurement: str,
tags: Dict[str, str],
fields: Dict[str, Any],
timestamp: Optional[int] = None,
) -> None:
"""Write a time-series data point."""
...
async def query(
self,
measurement: str,
start_time: int,
end_time: int,
aggregation: Optional[str] = None,
group_by: Optional[List[str]] = None,
) -> List[Dict[str, Any]]:
"""Query time-series data."""
...
async def delete_series(
self,
measurement: str,
tags: Optional[Dict[str, str]] = None,
) -> bool:
"""Delete a time series."""
...
class FileDatabase(Protocol):
"""Interface for file-based storage."""
async def read_file(self, path: str) -> Dict[str, Any]:
"""Read data from file."""
...
async def write_file(
self,
path: str,
data: Dict[str, Any],
) -> None:
"""Write data to file."""
...
async def append_file(
self,
path: str,
data: Any,
) -> None:
"""Append data to file."""
...
async def delete_file(self, path: str) -> bool:
"""Delete a file."""
...
async def list_files(
self,
pattern: str = "*",
) -> List[str]:
"""List files matching pattern."""
...
class DatabaseBackend(ABC):
"""Base class for all database backends."""
def __init__(self, config: Dict[str, Any]):
"""Initialize with configuration."""
self.config = config
self.db_type: DatabaseType = DatabaseType.FILE
self.capabilities: List[str] = []
@abstractmethod
async def connect(self) -> None:
"""Connect to the database."""
pass
@abstractmethod
async def disconnect(self) -> None:
"""Disconnect from the database."""
pass
@abstractmethod
async def health_check(self) -> bool:
"""Check if database is healthy."""
pass
def supports(self, capability: str) -> bool:
"""Check if backend supports a capability."""
return capability in self.capabilities
pkg/hanzo-memory/build/lib/hanzo_memory/db/backends/backend_registry.py
-288
View File
@@ -1,288 +0,0 @@
"""Backend registry for memory storage.
This provides a simple way to register and use different backends.
Supports multiple database types:
- LanceDB: Vector database with SQL-like queries
- KuzuDB: Graph database for relationship-based queries
- InfinityDB: High-performance vector database
- Local files: Simple JSON storage for development
"""
from typing import Dict, Optional, Type
from structlog import get_logger
from ..base import BaseVectorDB
from ..local_client import LocalMemoryClient
from ..lancedb_client import LanceDBClient
logger = get_logger()
# Try to import optional backends
try:
from ..kuzudb_client import KuzuDBClient
KUZU_AVAILABLE = True
except ImportError:
KUZU_AVAILABLE = False
logger.debug("KuzuDB client not available")
try:
from ..client import InfinityClient
INFINITY_AVAILABLE = True
except ImportError:
INFINITY_AVAILABLE = False
logger.debug("InfinityDB client not available")
try:
from ..sqlite_client import SQLiteMemoryClient
SQLITE_AVAILABLE = True
except ImportError:
SQLITE_AVAILABLE = False
logger.debug("SQLite client not available")
class BackendCapability:
"""Capabilities that backends can support."""
VECTOR_SEARCH = "vector_search"
FULL_TEXT_SEARCH = "full_text_search"
STRUCTURED_QUERY = "structured_query"
MARKDOWN_IMPORT = "markdown_import"
PERSISTENCE = "persistence"
EMBEDDINGS = "embeddings"
GRAPH_QUERIES = "graph_queries"
TIME_SERIES = "time_series"
class BackendRegistry:
"""Registry for memory backends."""
# Registered backends and their capabilities
BACKENDS: Dict[str, Dict] = {
"lancedb": {
"class": LanceDBClient,
"description": "LanceDB - Embedded vector database with SQL-like queries",
"capabilities": [
BackendCapability.VECTOR_SEARCH,
BackendCapability.STRUCTURED_QUERY,
BackendCapability.PERSISTENCE,
BackendCapability.EMBEDDINGS,
BackendCapability.MARKDOWN_IMPORT,
],
"config": {
"enable_markdown": True,
},
},
"local": {
"class": LocalMemoryClient,
"description": "Local file storage using JSON files",
"capabilities": [
BackendCapability.PERSISTENCE,
BackendCapability.MARKDOWN_IMPORT,
BackendCapability.FULL_TEXT_SEARCH, # Simple text search
],
"config": {
"enable_markdown": True,
},
},
}
# Add optional backends if available
@classmethod
def _init_optional_backends(cls):
"""Initialize optional backends if they're available."""
if KUZU_AVAILABLE and "kuzudb" not in cls.BACKENDS:
cls.BACKENDS["kuzudb"] = {
"class": KuzuDBClient,
"description": "KuzuDB - Graph database for relationship-based memory storage",
"capabilities": [
BackendCapability.GRAPH_QUERIES,
BackendCapability.STRUCTURED_QUERY,
BackendCapability.PERSISTENCE,
BackendCapability.EMBEDDINGS,
BackendCapability.MARKDOWN_IMPORT,
],
"config": {
"enable_markdown": True,
},
}
if INFINITY_AVAILABLE and "infinity" not in cls.BACKENDS:
cls.BACKENDS["infinity"] = {
"class": InfinityClient,
"description": "InfinityDB - High-performance vector database",
"capabilities": [
BackendCapability.VECTOR_SEARCH,
BackendCapability.STRUCTURED_QUERY,
BackendCapability.FULL_TEXT_SEARCH,
BackendCapability.PERSISTENCE,
BackendCapability.EMBEDDINGS,
BackendCapability.TIME_SERIES,
],
"config": {},
}
if SQLITE_AVAILABLE and "sqlite" not in cls.BACKENDS:
cls.BACKENDS["sqlite"] = {
"class": SQLiteMemoryClient,
"description": "SQLite - Lightweight embedded database with vector search via sqlite-vec",
"capabilities": [
BackendCapability.VECTOR_SEARCH,
BackendCapability.STRUCTURED_QUERY,
BackendCapability.PERSISTENCE,
BackendCapability.EMBEDDINGS,
BackendCapability.MARKDOWN_IMPORT,
],
"config": {},
}
@classmethod
def get_backend(
cls,
name: str,
config: Optional[Dict] = None,
) -> BaseVectorDB:
"""Get a backend instance by name.
Args:
name: Backend name (e.g., "lancedb", "local")
config: Optional configuration overrides
Returns:
Backend instance
Raises:
ValueError: If backend not found
"""
# Initialize optional backends
cls._init_optional_backends()
if name not in cls.BACKENDS:
available = ", ".join(cls.BACKENDS.keys())
raise ValueError(f"Unknown backend: {name}. Available: {available}")
backend_info = cls.BACKENDS[name]
backend_class = backend_info["class"]
# Merge default config with user config
backend_config = backend_info.get("config", {}).copy()
if config:
backend_config.update(config)
logger.info(
f"Creating {name} backend",
description=backend_info["description"],
capabilities=backend_info["capabilities"],
)
# Create instance with config
if backend_config:
return backend_class(**backend_config)
return backend_class()
@classmethod
def list_backends(cls) -> Dict[str, Dict]:
"""List all available backends and their capabilities."""
# Initialize optional backends
cls._init_optional_backends()
return {
name: {
"description": info["description"],
"capabilities": info["capabilities"],
}
for name, info in cls.BACKENDS.items()
}
@classmethod
def find_backend_for_capability(
cls,
capability: str,
) -> Optional[str]:
"""Find the best backend for a specific capability.
Args:
capability: Required capability
Returns:
Backend name or None if no backend supports it
"""
for name, info in cls.BACKENDS.items():
if capability in info["capabilities"]:
return name
return None
@classmethod
def get_backend_for_task(
cls,
task: str,
) -> str:
"""Get the recommended backend for a specific task.
Args:
task: Task type (e.g., "similarity_search", "document_storage")
Returns:
Recommended backend name
"""
task_mapping = {
"similarity_search": "lancedb",
"vector_search": "lancedb",
"embeddings": "lancedb",
"document_storage": "local",
"simple_storage": "local",
"markdown_import": "local", # Both support it, but local is simpler
"development": "local", # Simple for dev
"production": "lancedb", # Better for production
}
return task_mapping.get(task, "local") # Default to local
@classmethod
def register_backend(
cls,
name: str,
backend_class: Type[BaseVectorDB],
description: str,
capabilities: list[str],
config: Optional[Dict] = None,
) -> None:
"""Register a new backend.
Args:
name: Backend name
backend_class: Backend class
description: Human-readable description
capabilities: List of capabilities
config: Default configuration
"""
cls.BACKENDS[name] = {
"class": backend_class,
"description": description,
"capabilities": capabilities,
"config": config or {},
}
logger.info(f"Registered backend: {name}")
# Convenience functions
def get_backend(name: str = "local", **config) -> BaseVectorDB:
"""Get a backend instance.
Args:
name: Backend name (default: "local")
**config: Configuration options
Returns:
Backend instance
"""
return BackendRegistry.get_backend(name, config)
def list_backends() -> Dict[str, Dict]:
"""List all available backends."""
return BackendRegistry.list_backends()
def get_best_backend_for(capability: str) -> Optional[str]:
"""Get the best backend for a capability."""
return BackendRegistry.find_backend_for_capability(capability)
pkg/hanzo-memory/build/lib/hanzo_memory/db/base.py
-175
View File
@@ -1,175 +0,0 @@
"""Base database interface for vector storage backends."""
from abc import ABC, abstractmethod
from typing import Any
class BaseVectorDB(ABC):
"""Abstract base class for vector database backends."""
@abstractmethod
def create_project(
self,
project_id: str,
user_id: str,
name: str,
description: str = "",
metadata: dict | None = None,
) -> dict[str, Any]:
"""Create a new project."""
pass
@abstractmethod
def get_user_projects(self, user_id: str) -> list[dict[str, Any]]:
"""Get all projects for a user."""
pass
@abstractmethod
def create_memories_table(self, user_id: str) -> None:
"""Create a memories table for a user."""
pass
@abstractmethod
def add_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
importance: float = 0.5,
) -> dict[str, Any]:
"""Add a memory to the database."""
pass
@abstractmethod
def search_memories(
self,
user_id: str,
query_embedding: list[float],
project_id: str | None = None,
limit: int = 10,
min_similarity: float = 0.0,
) -> list[dict[str, Any]]:
"""Search memories by similarity."""
pass
@abstractmethod
def create_knowledge_base(
self,
knowledge_base_id: str,
project_id: str,
name: str,
description: str = "",
metadata: dict | None = None,
) -> dict[str, Any]:
"""Create a new knowledge base."""
pass
@abstractmethod
def get_knowledge_bases(self, project_id: str) -> list[dict[str, Any]]:
"""Get all knowledge bases for a project."""
pass
@abstractmethod
def add_fact(
self,
fact_id: str,
knowledge_base_id: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
confidence: float = 1.0,
) -> dict[str, Any]:
"""Add a fact to a knowledge base."""
pass
@abstractmethod
def search_facts(
self,
knowledge_base_id: str,
query_embedding: list[float] | None = None,
limit: int = 10,
) -> list[dict[str, Any]]:
"""Search facts in a knowledge base."""
pass
@abstractmethod
def delete_fact(self, fact_id: str, knowledge_base_id: str) -> bool:
"""Delete a fact from a knowledge base."""
pass
@abstractmethod
def update_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str | None = None,
metadata: dict | None = None,
importance: float | None = None,
) -> dict[str, Any] | None:
"""Update a memory in the database."""
pass
@abstractmethod
def create_chat_session(
self,
session_id: str,
user_id: str,
project_id: str,
metadata: dict | None = None,
) -> dict[str, Any]:
"""Create a new chat session."""
pass
@abstractmethod
def add_chat_message(
self,
message_id: str,
session_id: str,
role: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
) -> dict[str, Any]:
"""Add a message to a chat session."""
pass
@abstractmethod
def get_chat_messages(
self,
session_id: str,
limit: int | None = None,
) -> list[dict[str, Any]]:
"""Get messages from a chat session."""
pass
@abstractmethod
def search_chat_messages(
self,
session_id: str,
query_embedding: list[float],
limit: int = 10,
) -> list[dict[str, Any]]:
"""Search messages in a chat session by similarity."""
pass
# Optional methods for implementations - not abstract
# Implementations can override these if needed
def close(self) -> None:
"""Close the database connection (optional for implementations)."""
# Default implementation does nothing
return None
def create_projects_table(self) -> None:
"""Create projects table if not exists (optional for implementations)."""
# Default implementation does nothing
return None
def create_knowledge_bases_table(self) -> None:
"""Create knowledge bases table if not exists (optional for implementations)."""
# Default implementation does nothing
return None
pkg/hanzo-memory/build/lib/hanzo_memory/db/client.py
-456
View File
@@ -1,456 +0,0 @@
"""InfinityDB client for vector storage and search."""
import json
import platform
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import polars as pl
from structlog import get_logger
from ..config import settings
from .base import BaseVectorDB
logger = get_logger()
# Try to import InfinityDB, fall back to mock for unsupported platforms
try:
import infinity_embedded
logger.info("Using InfinityDB")
except ImportError:
logger.debug(
f"InfinityDB not available on {platform.system()} {platform.machine()}, using mock implementation"
)
from . import mock_infinity as infinity_embedded
class InfinityClient(BaseVectorDB):
"""Client for InfinityDB operations."""
def __init__(self, db_path: str | None = None):
"""Initialize InfinityDB client."""
self.db_path = db_path or settings.infinity_db_str
Path(self.db_path).mkdir(parents=True, exist_ok=True)
self.infinity = infinity_embedded.connect(self.db_path)
self._ensure_databases()
def _ensure_databases(self) -> None:
"""Ensure required databases exist."""
# Create main databases
for db_name in ["projects", "memories", "knowledge", "chats"]:
try:
self.infinity.create_database(db_name)
logger.info(f"Created database: {db_name}")
except Exception:
# Database already exists
pass
def _get_db(self, db_name: str) -> Any:
"""Get a database object."""
return self.infinity.get_database(db_name)
# Project Management
def create_projects_table(self) -> None:
"""Create projects table if not exists."""
db = self._get_db("projects")
try:
db.create_table(
"projects",
{
"project_id": {"type": "varchar"},
"user_id": {"type": "varchar"},
"name": {"type": "varchar"},
"description": {"type": "varchar"},
"metadata": {"type": "varchar"}, # JSON string
"created_at": {"type": "varchar"},
"updated_at": {"type": "varchar"},
},
)
logger.info("Created projects table")
except Exception as e:
logger.debug(f"Projects table may already exist: {e}")
def create_project(
self,
project_id: str,
user_id: str,
name: str,
description: str = "",
metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Create a new project."""
db = self._get_db("projects")
table = db.get_table("projects")
now = datetime.now(timezone.utc).isoformat()
project_data = {
"project_id": project_id,
"user_id": user_id,
"name": name,
"description": description,
"metadata": json.dumps(metadata or {}),
"created_at": now,
"updated_at": now,
}
table.insert([project_data])
return project_data
# Memory Management
def create_memories_table(self, user_id: str) -> None:
"""Create memories table for a user if not exists."""
db = self._get_db("memories")
table_name = f"memories_{user_id}"
try:
db.create_table(
table_name,
{
"memory_id": {"type": "varchar"},
"user_id": {"type": "varchar"},
"project_id": {"type": "varchar"},
"content": {"type": "varchar"},
"embedding": {
"type": f"vector,{settings.embedding_dimensions},float"
},
"metadata": {"type": "varchar"}, # JSON string
"importance": {"type": "float"},
"created_at": {"type": "varchar"},
"updated_at": {"type": "varchar"},
},
)
logger.info(f"Created memories table for user: {user_id}")
except Exception:
pass
def add_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str,
embedding: list[float],
metadata: dict[str, Any] | None = None,
importance: float = 1.0,
) -> dict[str, Any]:
"""Add a memory to the database."""
db = self._get_db("memories")
table_name = f"memories_{user_id}"
table = db.get_table(table_name)
now = datetime.now(timezone.utc).isoformat()
memory_data = {
"memory_id": memory_id,
"user_id": user_id,
"project_id": project_id,
"content": content,
"embedding": embedding,
"metadata": json.dumps(metadata or {}),
"importance": importance,
"created_at": now,
"updated_at": now,
}
table.insert([memory_data])
return memory_data
def search_memories(
self,
user_id: str,
query_embedding: list[float],
project_id: str | None = None,
limit: int = 10,
threshold: float = 0.0,
) -> pl.DataFrame:
"""Search memories using vector similarity."""
db = self._get_db("memories")
table_name = f"memories_{user_id}"
try:
table = db.get_table(table_name)
# Build query
query = table.output(["*"]).match_dense(
"embedding",
query_embedding,
"float",
"cosine", # Using cosine similarity
limit,
)
# Apply project filter if specified
if project_id:
query = query.filter(f"project_id = '{project_id}'")
# Execute query and return as polars DataFrame
return query.to_pl()
except Exception as e:
logger.error(f"Error searching memories: {e}")
return pl.DataFrame()
# Knowledge Base Management
def create_knowledge_bases_table(self) -> None:
"""Create knowledge bases table if not exists."""
db = self._get_db("knowledge")
try:
db.create_table(
"knowledge_bases",
{
"kb_id": {"type": "varchar"},
"user_id": {"type": "varchar"},
"project_id": {"type": "varchar"},
"name": {"type": "varchar"},
"description": {"type": "varchar"},
"metadata": {"type": "varchar"}, # JSON string
"created_at": {"type": "varchar"},
"updated_at": {"type": "varchar"},
},
)
logger.info("Created knowledge_bases table")
except Exception as e:
logger.debug(f"Knowledge bases table may already exist: {e}")
def create_knowledge_base(
self,
kb_id: str,
user_id: str,
project_id: str,
name: str,
description: str = "",
metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Create a new knowledge base."""
db = self._get_db("knowledge")
table = db.get_table("knowledge_bases")
now = datetime.now(timezone.utc).isoformat()
kb_data = {
"kb_id": kb_id,
"user_id": user_id,
"project_id": project_id,
"name": name,
"description": description,
"metadata": json.dumps(metadata or {}),
"created_at": now,
"updated_at": now,
}
table.insert([kb_data])
# Create facts table for this knowledge base
self._create_facts_table(kb_id)
return kb_data
def _create_facts_table(self, kb_id: str) -> None:
"""Create facts table for a knowledge base."""
db = self._get_db("knowledge")
table_name = f"facts_{kb_id}"
try:
db.create_table(
table_name,
{
"fact_id": {"type": "varchar"},
"kb_id": {"type": "varchar"},
"content": {"type": "varchar"},
"embedding": {
"type": f"vector,{settings.embedding_dimensions},float"
},
"parent_id": {"type": "varchar"},
"metadata": {"type": "varchar"}, # JSON string
"created_at": {"type": "varchar"},
"updated_at": {"type": "varchar"},
},
)
logger.info(f"Created facts table for kb: {kb_id}")
except Exception:
pass
def add_fact(
self,
fact_id: str,
kb_id: str,
content: str,
embedding: list[float],
parent_id: str | None = None,
metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Add a fact to a knowledge base."""
db = self._get_db("knowledge")
table_name = f"facts_{kb_id}"
table = db.get_table(table_name)
now = datetime.now(timezone.utc).isoformat()
fact_data = {
"fact_id": fact_id,
"kb_id": kb_id,
"content": content,
"embedding": embedding,
"parent_id": parent_id or "",
"metadata": json.dumps(metadata or {}),
"created_at": now,
"updated_at": now,
}
table.insert([fact_data])
return fact_data
def search_facts(
self,
kb_id: str,
query_embedding: list[float],
limit: int = 10,
parent_id: str | None = None,
) -> pl.DataFrame:
"""Search facts in a knowledge base."""
db = self._get_db("knowledge")
table_name = f"facts_{kb_id}"
try:
table = db.get_table(table_name)
# Build query
query = table.output(["*"]).match_dense(
"embedding", query_embedding, "float", "cosine", limit
)
# Apply parent filter if specified
if parent_id:
query = query.filter(f"parent_id = '{parent_id}'")
return query.to_pl()
except Exception as e:
logger.error(f"Error searching facts: {e}")
return pl.DataFrame()
# Chat Management
def create_chats_table(self, user_id: str) -> None:
"""Create chats table for a user."""
db = self._get_db("chats")
table_name = f"chats_{user_id}"
try:
db.create_table(
table_name,
{
"chat_id": {"type": "varchar"},
"user_id": {"type": "varchar"},
"project_id": {"type": "varchar"},
"session_id": {"type": "varchar"},
"role": {"type": "varchar"}, # user, assistant, system
"content": {"type": "varchar"},
"embedding": {
"type": f"vector,{settings.embedding_dimensions},float"
},
"metadata": {"type": "varchar"}, # JSON string
"created_at": {"type": "varchar"},
},
)
logger.info(f"Created chats table for user: {user_id}")
except Exception:
pass
def add_chat_message(
self,
chat_id: str,
user_id: str,
project_id: str,
session_id: str,
role: str,
content: str,
embedding: list[float],
metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Add a chat message."""
db = self._get_db("chats")
table_name = f"chats_{user_id}"
table = db.get_table(table_name)
chat_data = {
"chat_id": chat_id,
"user_id": user_id,
"project_id": project_id,
"session_id": session_id,
"role": role,
"content": content,
"embedding": embedding,
"metadata": json.dumps(metadata or {}),
"created_at": datetime.now(timezone.utc).isoformat(),
}
table.insert([chat_data])
return chat_data
def get_chat_history(
self,
user_id: str,
session_id: str,
limit: int = 100,
) -> pl.DataFrame:
"""Get chat history for a session."""
db = self._get_db("chats")
table_name = f"chats_{user_id}"
try:
table = db.get_table(table_name)
# Get messages for session ordered by created_at
query = table.output(["*"]).filter(f"session_id = '{session_id}'")
# TODO: Add proper ordering once InfinityDB supports it
return query.to_pl()
except Exception as e:
logger.error(f"Error getting chat history: {e}")
return pl.DataFrame()
def search_chats(
self,
user_id: str,
query_embedding: list[float],
project_id: str | None = None,
session_id: str | None = None,
limit: int = 10,
) -> pl.DataFrame:
"""Search chat messages."""
db = self._get_db("chats")
table_name = f"chats_{user_id}"
try:
table = db.get_table(table_name)
# Build query
query = table.output(["*"]).match_dense(
"embedding", query_embedding, "float", "cosine", limit
)
# Apply filters
if project_id:
query = query.filter(f"project_id = '{project_id}'")
if session_id:
query = query.filter(f"session_id = '{session_id}'")
return query.to_pl()
except Exception as e:
logger.error(f"Error searching chats: {e}")
return pl.DataFrame()
def close(self) -> None:
"""Close the InfinityDB connection."""
if hasattr(self, "infinity"):
# InfinityDB embedded doesn't have explicit close method
pass
# Global client instance
_client: InfinityClient | None = None
def get_client() -> InfinityClient:
"""Get or create the global InfinityDB client."""
global _client
if _client is None:
_client = InfinityClient()
return _client
pkg/hanzo-memory/build/lib/hanzo_memory/db/factory.py
-78
View File
@@ -1,78 +0,0 @@
"""Database factory for selecting the appropriate backend."""
from typing import Optional, Dict, Any
from structlog import get_logger
from ..config import settings
from .base import BaseVectorDB
from .backends.backend_registry import BackendRegistry
logger = get_logger()
# Global database client instance
_db_client: BaseVectorDB | None = None
def get_db_client(
backend: Optional[str] = None,
config: Optional[Dict[str, Any]] = None,
) -> BaseVectorDB:
"""Get or create the global database client.
Args:
backend: Override backend type (default: from settings)
config: Optional configuration for the backend
Returns:
Database client instance
"""
global _db_client
# Determine backend
backend_name = backend or settings.db_backend.lower()
# Reset client if backend changed
if _db_client and hasattr(_db_client, "__backend_name__"):
if _db_client.__backend_name__ != backend_name: # type: ignore
logger.info(f"Backend changed from {_db_client.__backend_name__} to {backend_name}, resetting client") # type: ignore
_db_client = None
if _db_client is None:
# Special handling for legacy "infinity" backend name
if backend_name == "infinity":
logger.warning("InfinityDB backend is optional, falling back to local")
backend_name = "local"
try:
_db_client = BackendRegistry.get_backend(backend_name, config)
_db_client.__backend_name__ = backend_name # type: ignore # Store for comparison
except ValueError as e:
logger.error(f"Failed to create backend {backend_name}: {e}")
logger.info("Falling back to local backend")
_db_client = BackendRegistry.get_backend("local", config)
_db_client.__backend_name__ = "local" # type: ignore
return _db_client
def reset_db_client() -> None:
"""Reset the global database client (useful for testing)."""
global _db_client
_db_client = None
def list_available_backends() -> Dict[str, Dict]:
"""List all available backends and their capabilities."""
return BackendRegistry.list_backends()
def get_backend_for_task(task: str) -> str:
"""Get the recommended backend for a specific task.
Args:
task: Task description (e.g., "vector_search", "document_storage")
Returns:
Recommended backend name
"""
return BackendRegistry.get_backend_for_task(task)
pkg/hanzo-memory/build/lib/hanzo_memory/db/kuzudb_client.py
-635
View File
@@ -1,635 +0,0 @@
"""KuzuDB client for graph-based memory storage."""
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional
import numpy as np
from structlog import get_logger
from ..config import settings
from ..models.memory import Memory, MemoryCreate, MemoryResponse
from ..models.project import Project, ProjectCreate
from ..models.fact import Fact, FactCreate
from .base import BaseVectorDB
logger = get_logger()
try:
import kuzu
KUZU_AVAILABLE = True
except ImportError:
KUZU_AVAILABLE = False
logger.warning("KuzuDB not available. Install with: pip install kuzu")
class KuzuDBClient(BaseVectorDB):
"""KuzuDB implementation for graph-based memory storage."""
def __init__(self, db_path: Optional[str] = None, enable_markdown: bool = True):
"""Initialize KuzuDB client.
Args:
db_path: Path to KuzuDB database directory
enable_markdown: Whether to import markdown files
"""
if not KUZU_AVAILABLE:
raise ImportError("KuzuDB not installed. Install with: pip install kuzu")
self.db_path = Path(db_path or settings.db_path) / "kuzudb"
self.db_path.mkdir(parents=True, exist_ok=True)
# Initialize KuzuDB database
self.db = kuzu.Database(str(self.db_path))
self.conn = kuzu.Connection(self.db)
# Initialize schema
self._init_schema()
# Import markdown if enabled
self.enable_markdown = enable_markdown
if enable_markdown:
from ..markdown_memory import MarkdownMemoryReader
self.markdown_reader = MarkdownMemoryReader()
self._import_markdown_memories()
logger.info(f"Initialized KuzuDB storage at {self.db_path}")
def _init_schema(self):
"""Initialize KuzuDB schema with nodes and relationships."""
# Create node tables
queries = [
# User node
"""CREATE NODE TABLE IF NOT EXISTS User(
user_id STRING PRIMARY KEY,
name STRING,
created_at TIMESTAMP
)""",
# Project node
"""CREATE NODE TABLE IF NOT EXISTS Project(
project_id STRING PRIMARY KEY,
name STRING,
description STRING,
metadata STRING,
created_at TIMESTAMP,
updated_at TIMESTAMP
)""",
# Memory node
"""CREATE NODE TABLE IF NOT EXISTS Memory(
memory_id STRING PRIMARY KEY,
content STRING,
memory_type STRING,
importance DOUBLE,
context STRING,
metadata STRING,
source STRING,
embedding DOUBLE[],
created_at TIMESTAMP,
updated_at TIMESTAMP
)""",
# Fact node
"""CREATE NODE TABLE IF NOT EXISTS Fact(
fact_id STRING PRIMARY KEY,
statement STRING,
confidence DOUBLE,
source STRING,
metadata STRING,
embedding DOUBLE[],
created_at TIMESTAMP,
updated_at TIMESTAMP
)""",
# KnowledgeBase node
"""CREATE NODE TABLE IF NOT EXISTS KnowledgeBase(
kb_id STRING PRIMARY KEY,
name STRING,
description STRING,
metadata STRING,
created_at TIMESTAMP,
updated_at TIMESTAMP
)""",
# Create relationship tables
"""CREATE REL TABLE IF NOT EXISTS OWNS(
FROM User TO Project
)""",
"""CREATE REL TABLE IF NOT EXISTS HAS_MEMORY(
FROM Project TO Memory,
user_id STRING
)""",
"""CREATE REL TABLE IF NOT EXISTS HAS_FACT(
FROM Project TO Fact,
user_id STRING
)""",
"""CREATE REL TABLE IF NOT EXISTS IN_KB(
FROM Memory TO KnowledgeBase
)""",
"""CREATE REL TABLE IF NOT EXISTS RELATES_TO(
FROM Memory TO Memory,
relationship_type STRING,
strength DOUBLE
)""",
"""CREATE REL TABLE IF NOT EXISTS DERIVED_FROM(
FROM Fact TO Memory
)""",
]
for query in queries:
try:
self.conn.execute(query)
except Exception as e:
# Table might already exist
logger.debug(f"Schema creation note: {e}")
async def initialize(self) -> None:
"""Initialize the database."""
logger.info("KuzuDB initialized")
async def close(self) -> None:
"""Close database connection."""
if hasattr(self, 'conn'):
# KuzuDB doesn't have explicit close, but we can clean up
self.conn = None
logger.info("Closing KuzuDB connection")
def create_project_sync(self, project: ProjectCreate, user_id: str) -> dict:
"""Create a project synchronously."""
project_id = str(uuid.uuid4())
now = datetime.now(timezone.utc).isoformat()
# Create project node
self.conn.execute(
"""MERGE (p:Project {project_id: $pid})
SET p.name = $name,
p.description = $desc,
p.metadata = $metadata,
p.created_at = $created,
p.updated_at = $updated
""",
{
"pid": project_id,
"name": project.name,
"desc": project.description,
"metadata": json.dumps(project.metadata or {}),
"created": now,
"updated": now,
}
)
# Create or connect user
self.conn.execute(
"""MERGE (u:User {user_id: $uid})
ON CREATE SET u.created_at = $created
""",
{"uid": user_id, "created": now}
)
# Create ownership relationship
self.conn.execute(
"""MATCH (u:User {user_id: $uid}), (p:Project {project_id: $pid})
MERGE (u)-[:OWNS]->(p)
""",
{"uid": user_id, "pid": project_id}
)
return {
"project_id": project_id,
"name": project.name,
"description": project.description,
"user_id": user_id,
"metadata": project.metadata or {},
"created_at": now,
"updated_at": now,
}
async def create_project(self, project: ProjectCreate, user_id: str) -> Project:
"""Create a project."""
project_data = self.create_project_sync(project, user_id)
return Project(**project_data)
def add_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
importance: float = 0.5,
) -> dict[str, Any]:
"""Add a memory to the graph."""
now = datetime.now(timezone.utc).isoformat()
# Create memory node
self.conn.execute(
"""MERGE (m:Memory {memory_id: $mid})
SET m.content = $content,
m.memory_type = $mtype,
m.importance = $importance,
m.context = $context,
m.metadata = $metadata,
m.source = $source,
m.embedding = $embedding,
m.created_at = $created,
m.updated_at = $updated
""",
{
"mid": memory_id,
"content": content,
"mtype": metadata.get("memory_type", "general") if metadata else "general",
"importance": importance,
"context": json.dumps(metadata.get("context", {}) if metadata else {}),
"metadata": json.dumps(metadata or {}),
"source": metadata.get("source", "") if metadata else "",
"embedding": embedding,
"created": now,
"updated": now,
}
)
# Connect to project
self.conn.execute(
"""MATCH (p:Project {project_id: $pid}), (m:Memory {memory_id: $mid})
MERGE (p)-[:HAS_MEMORY {user_id: $uid}]->(m)
""",
{"pid": project_id, "mid": memory_id, "uid": user_id}
)
# Find and create relationships to related memories
if embedding:
self._create_memory_relationships(memory_id, embedding, project_id)
return {
"memory_id": memory_id,
"project_id": project_id,
"user_id": user_id,
"content": content,
"metadata": metadata or {},
"importance": importance,
"created_at": now,
"updated_at": now,
}
def _create_memory_relationships(self, memory_id: str, embedding: list[float], project_id: str):
"""Create relationships between related memories based on similarity."""
# Find similar memories in the same project
results = self.conn.execute(
"""MATCH (p:Project {project_id: $pid})-[:HAS_MEMORY]->(m:Memory)
WHERE m.memory_id <> $mid AND m.embedding IS NOT NULL
RETURN m.memory_id, m.embedding
LIMIT 10
""",
{"pid": project_id, "mid": memory_id}
)
query_vec = np.array(embedding)
for row in results:
other_id = row[0]
other_embedding = row[1]
if other_embedding:
similarity = self._compute_similarity(query_vec, np.array(other_embedding))
if similarity > 0.7: # Only create strong relationships
self.conn.execute(
"""MATCH (m1:Memory {memory_id: $mid1}), (m2:Memory {memory_id: $mid2})
MERGE (m1)-[:RELATES_TO {relationship_type: 'similar', strength: $strength}]->(m2)
""",
{"mid1": memory_id, "mid2": other_id, "strength": similarity}
)
def _compute_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
"""Compute cosine similarity between two vectors."""
dot_product = np.dot(vec1, vec2)
norm1 = np.linalg.norm(vec1)
norm2 = np.linalg.norm(vec2)
if norm1 == 0 or norm2 == 0:
return 0.0
return float(dot_product / (norm1 * norm2))
def search_memories(
self,
query_embedding: List[float],
user_id: str,
project_id: Optional[str] = None,
limit: int = 10,
memory_type: Optional[str] = None,
) -> Any:
"""Search memories using graph traversal and vector similarity."""
import pandas as pd
# Build the query
where_clause = []
params = {"uid": user_id, "limit": limit}
if project_id:
where_clause.append("p.project_id = $pid")
params["pid"] = project_id
if memory_type:
where_clause.append("m.memory_type = $mtype")
params["mtype"] = memory_type
where = " AND ".join(where_clause) if where_clause else "1=1"
# Query memories
query = f"""
MATCH (u:User {{user_id: $uid}})-[:OWNS]->(p:Project)-[:HAS_MEMORY]->(m:Memory)
WHERE {where}
RETURN m.memory_id, m.content, m.memory_type, m.importance,
m.context, m.metadata, m.source, m.embedding,
m.created_at, m.updated_at, p.project_id, $uid as user_id
LIMIT $limit
"""
results = self.conn.execute(query, params)
# Calculate similarities and build dataframe
data = []
query_vec = np.array(query_embedding) if query_embedding else None
for row in results:
similarity = 0.0
if query_vec is not None and row[7]: # embedding
similarity = self._compute_similarity(query_vec, np.array(row[7]))
data.append({
"memory_id": row[0],
"content": row[1],
"memory_type": row[2],
"importance": row[3],
"context": row[4],
"metadata": row[5],
"source": row[6],
"created_at": row[8],
"updated_at": row[9],
"project_id": row[10],
"user_id": row[11],
"similarity_score": similarity,
})
# Sort by similarity if we have embeddings
if query_embedding:
data.sort(key=lambda x: x["similarity_score"], reverse=True)
return pd.DataFrame(data[:limit])
async def search_memories_async(
self,
query_embedding: List[float],
project_id: str,
user_id: Optional[str] = None,
limit: int = 10,
memory_type: Optional[str] = None,
) -> List[Memory]:
"""Async wrapper for search_memories."""
df = self.search_memories(query_embedding, user_id or "default", project_id, limit, memory_type)
memories = []
for _, row in df.iterrows():
memories.append(Memory(
memory_id=row.get("memory_id"),
project_id=row.get("project_id"),
user_id=row.get("user_id"),
content=row.get("content"),
memory_type=row.get("memory_type"),
importance=row.get("importance", 0.5),
context=json.loads(row.get("context", "{}")) if isinstance(row.get("context"), str) else row.get("context", {}),
metadata=json.loads(row.get("metadata", "{}")) if isinstance(row.get("metadata"), str) else row.get("metadata", {}),
source=row.get("source"),
created_at=row.get("created_at"),
updated_at=row.get("updated_at"),
))
return memories
async def create_memory(
self, memory: MemoryCreate, project_id: str, user_id: Optional[str] = None
) -> Memory:
"""Create a memory."""
memory_id = str(uuid.uuid4())
user_id = user_id or "default"
memory_data = self.add_memory(
memory_id=memory_id,
user_id=user_id,
project_id=project_id,
content=memory.content,
embedding=memory.embedding or [0.0] * settings.embedding_dimensions,
metadata={
"memory_type": memory.memory_type,
"context": memory.context or {},
"source": memory.source,
**(memory.metadata or {})
},
importance=memory.importance,
)
return Memory(**memory_data)
async def get_recent_memories(
self,
project_id: str,
user_id: Optional[str] = None,
limit: int = 10,
memory_type: Optional[str] = None,
) -> List[Memory]:
"""Get recent memories from the graph."""
user_id = user_id or "default"
where_clause = ["p.project_id = $pid"]
params = {"pid": project_id, "uid": user_id, "limit": limit}
if memory_type:
where_clause.append("m.memory_type = $mtype")
params["mtype"] = memory_type
where = " AND ".join(where_clause)
query = f"""
MATCH (u:User {{user_id: $uid}})-[:OWNS]->(p:Project)-[:HAS_MEMORY]->(m:Memory)
WHERE {where}
RETURN m.memory_id, m.content, m.memory_type, m.importance,
m.context, m.metadata, m.source,
m.created_at, m.updated_at
ORDER BY m.created_at DESC
LIMIT $limit
"""
results = self.conn.execute(query, params)
memories = []
for row in results:
memories.append(Memory(
memory_id=row[0],
project_id=project_id,
user_id=user_id,
content=row[1],
memory_type=row[2],
importance=row[3],
context=json.loads(row[4]) if isinstance(row[4], str) else row[4],
metadata=json.loads(row[5]) if isinstance(row[5], str) else row[5],
source=row[6],
created_at=row[7],
updated_at=row[8],
))
return memories
async def list_projects(self, user_id: Optional[str] = None) -> List[Project]:
"""List all projects for a user."""
user_id = user_id or "default"
results = self.conn.execute(
"""MATCH (u:User {user_id: $uid})-[:OWNS]->(p:Project)
RETURN p.project_id, p.name, p.description, p.metadata,
p.created_at, p.updated_at
""",
{"uid": user_id}
)
projects = []
for row in results:
projects.append(Project(
project_id=row[0],
name=row[1],
description=row[2],
user_id=user_id,
metadata=json.loads(row[3]) if isinstance(row[3], str) else row[3],
created_at=row[4],
updated_at=row[5],
))
return projects
def create_memories_table(self, user_id: str) -> None:
"""Create memories table (no-op for KuzuDB as schema is predefined)."""
pass
def _import_markdown_memories(self):
"""Import memories from markdown files."""
if not self.markdown_reader:
return
memories = self.markdown_reader.read_all_memories()
imported_count = 0
# Create a special project for markdown imports
project_id = "markdown_import"
self.conn.execute(
"""MERGE (p:Project {project_id: $pid})
SET p.name = $name,
p.description = $desc,
p.created_at = $created,
p.updated_at = $updated
""",
{
"pid": project_id,
"name": "Markdown Import",
"desc": "Memories imported from markdown files",
"created": datetime.now(timezone.utc).isoformat(),
"updated": datetime.now(timezone.utc).isoformat(),
}
)
for memory_create in memories:
memory_id = str(uuid.uuid4())
self.add_memory(
memory_id=memory_id,
user_id="system",
project_id=project_id,
content=memory_create.content,
embedding=[0.0] * settings.embedding_dimensions,
metadata={
"memory_type": memory_create.memory_type or "knowledge",
"context": memory_create.context or {},
"source": memory_create.source,
**(memory_create.metadata or {})
},
importance=memory_create.importance,
)
imported_count += 1
if imported_count > 0:
logger.info(f"Imported {imported_count} memories from markdown files to KuzuDB")
# Additional methods for graph-specific operations
def get_related_memories(self, memory_id: str, relationship_type: Optional[str] = None) -> List[Dict]:
"""Get memories related to a specific memory through graph relationships."""
where = ""
params = {"mid": memory_id}
if relationship_type:
where = "{relationship_type: $rtype}"
params["rtype"] = relationship_type
query = f"""
MATCH (m1:Memory {{memory_id: $mid}})-[r:RELATES_TO {where}]->(m2:Memory)
RETURN m2.memory_id, m2.content, r.relationship_type, r.strength
ORDER BY r.strength DESC
"""
results = self.conn.execute(query, params)
related = []
for row in results:
related.append({
"memory_id": row[0],
"content": row[1],
"relationship_type": row[2],
"strength": row[3],
})
return related
def get_memory_graph(self, project_id: str, depth: int = 2) -> Dict:
"""Get a subgraph of memories and their relationships."""
# Get nodes
nodes_query = """
MATCH (p:Project {project_id: $pid})-[:HAS_MEMORY]->(m:Memory)
RETURN m.memory_id, m.content, m.memory_type, m.importance
LIMIT 100
"""
nodes_results = self.conn.execute(nodes_query, {"pid": project_id})
nodes = []
node_ids = set()
for row in nodes_results:
nodes.append({
"id": row[0],
"content": row[1],
"type": row[2],
"importance": row[3],
})
node_ids.add(row[0])
# Get edges
edges_query = """
MATCH (m1:Memory)-[r:RELATES_TO]->(m2:Memory)
WHERE m1.memory_id IN $ids AND m2.memory_id IN $ids
RETURN m1.memory_id, m2.memory_id, r.relationship_type, r.strength
"""
edges_results = self.conn.execute(edges_query, {"ids": list(node_ids)})
edges = []
for row in edges_results:
edges.append({
"source": row[0],
"target": row[1],
"type": row[2],
"strength": row[3],
})
return {
"nodes": nodes,
"edges": edges,
}
pkg/hanzo-memory/build/lib/hanzo_memory/db/lancedb_client.py
-1158
View File
File diff suppressed because it is too large Load Diff
pkg/hanzo-memory/build/lib/hanzo_memory/db/local_client.py
-716
View File
@@ -1,716 +0,0 @@
"""Local file-based memory storage implementation."""
import json
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
import numpy as np
from structlog import get_logger
from ..models.knowledge import Fact, FactCreate, KnowledgeBase
from ..models.memory import Memory, MemoryCreate, MemoryResponse
from ..models.project import Project, ProjectCreate
from .base import BaseVectorDB
from .markdown_reader import MarkdownMemoryReader
logger = get_logger()
class LocalMemoryClient(BaseVectorDB):
"""Local file-based implementation of the vector database."""
def __init__(self, storage_dir: Optional[Path] = None, enable_markdown: bool = True):
"""Initialize local memory storage.
Args:
storage_dir: Directory to store memory files
enable_markdown: Whether to enable markdown file integration
"""
self.storage_dir = storage_dir or Path.home() / ".hanzo" / "memory"
self.storage_dir.mkdir(parents=True, exist_ok=True)
# File paths for different collections
self.memories_file = self.storage_dir / "memories.json"
self.facts_file = self.storage_dir / "facts.json"
self.projects_file = self.storage_dir / "projects.json"
self.embeddings_file = self.storage_dir / "embeddings.npz"
self.markdown_index_file = self.storage_dir / "markdown_index.json"
# Load existing data
self.memories = self._load_json(self.memories_file)
self.facts = self._load_json(self.facts_file)
self.projects = self._load_json(self.projects_file)
self.embeddings = self._load_embeddings()
self.markdown_index = self._load_json(self.markdown_index_file)
# Initialize markdown reader if enabled
self.enable_markdown = enable_markdown
self.markdown_reader = MarkdownMemoryReader() if enable_markdown else None
# Import markdown memories on initialization
if self.enable_markdown:
self._import_markdown_memories()
logger.info(f"Initialized local memory storage at {self.storage_dir}")
def _load_json(self, file_path: Path) -> Dict[str, Any]:
"""Load JSON data from file."""
if file_path.exists():
try:
with open(file_path, 'r') as f:
return json.load(f)
except Exception as e:
logger.error(f"Error loading {file_path}: {e}")
return {}
def _save_json(self, data: Dict[str, Any], file_path: Path):
"""Save JSON data to file."""
try:
with open(file_path, 'w') as f:
json.dump(data, f, indent=2, default=str)
except Exception as e:
logger.error(f"Error saving {file_path}: {e}")
def _load_embeddings(self) -> Dict[str, np.ndarray]:
"""Load embeddings from file."""
if self.embeddings_file.exists():
try:
data = np.load(self.embeddings_file, allow_pickle=True)
return {k: v for k, v in data.items()}
except Exception as e:
logger.error(f"Error loading embeddings: {e}")
return {}
def _save_embeddings(self):
"""Save embeddings to file."""
try:
np.savez_compressed(self.embeddings_file, **self.embeddings)
except Exception as e:
logger.error(f"Error saving embeddings: {e}")
def _compute_similarity(self, embedding1: np.ndarray, embedding2: np.ndarray) -> float:
"""Compute cosine similarity between two embeddings."""
dot_product = np.dot(embedding1, embedding2)
norm1 = np.linalg.norm(embedding1)
norm2 = np.linalg.norm(embedding2)
if norm1 == 0 or norm2 == 0:
return 0.0
return float(dot_product / (norm1 * norm2))
async def initialize(self) -> None:
"""Initialize the database (no-op for local storage)."""
logger.info("Local memory storage initialized")
async def close(self) -> None:
"""Close the database connection (save all data)."""
self._save_json(self.memories, self.memories_file)
self._save_json(self.facts, self.facts_file)
self._save_json(self.projects, self.projects_file)
self._save_embeddings()
logger.info("Local memory storage closed")
# Memory operations
async def create_memory(
self, memory: MemoryCreate, project_id: str, user_id: Optional[str] = None
) -> Memory:
"""Create a new memory."""
memory_id = str(uuid.uuid4())
# Create memory object
memory_data = {
"id": memory_id,
"project_id": project_id,
"user_id": user_id or "default",
"content": memory.content,
"memory_type": memory.memory_type,
"importance": memory.importance,
"context": memory.context,
"metadata": memory.metadata or {},
"source": memory.source,
"timestamp": datetime.utcnow().isoformat(),
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
# Store memory
self.memories[memory_id] = memory_data
# Store embedding if provided
if memory.embedding:
self.embeddings[f"memory_{memory_id}"] = np.array(memory.embedding)
# Save to disk
self._save_json(self.memories, self.memories_file)
if memory.embedding:
self._save_embeddings()
# Return with proper memory_id field
memory_data["memory_id"] = memory_id # Ensure memory_id is set
return Memory(**memory_data)
def search_memories(
self,
query_embedding: List[float],
user_id: str,
project_id: Optional[str] = None,
limit: int = 10,
memory_type: Optional[str] = None,
) -> Any:
"""Search for similar memories synchronously."""
query_vec = np.array(query_embedding)
results = []
for mem_id, memory in self.memories.items():
# Filter by project and user
if project_id and memory.get("project_id") != project_id:
continue
if user_id and memory.get("user_id") != user_id:
continue
if memory_type and memory.get("memory_type") != memory_type:
continue
# Calculate similarity if embedding exists
embedding_key = f"memory_{mem_id}"
if embedding_key in self.embeddings:
similarity = self._compute_similarity(query_vec, self.embeddings[embedding_key])
results.append({
"memory": memory,
"similarity": similarity
})
# Sort by similarity and limit
results.sort(key=lambda x: x["similarity"], reverse=True)
results = results[:limit]
# Convert to dataframe-like structure for compatibility
import pandas as pd
if not results:
return pd.DataFrame()
# Extract memory data for dataframe
data = []
for r in results:
mem = r["memory"]
data.append({
"memory_id": mem.get("id", mem.get("memory_id")),
"project_id": mem.get("project_id"),
"user_id": mem.get("user_id"),
"content": mem.get("content"),
"memory_type": mem.get("memory_type"),
"importance": mem.get("importance"),
"context": mem.get("context"),
"metadata": mem.get("metadata"),
"source": mem.get("source"),
"created_at": mem.get("created_at"),
"updated_at": mem.get("updated_at"),
"similarity_score": r["similarity"]
})
return pd.DataFrame(data)
async def search_memories_async(
self,
query_embedding: List[float],
project_id: str,
user_id: Optional[str] = None,
limit: int = 10,
memory_type: Optional[str] = None,
min_importance: float = 0.0,
) -> List[MemoryResponse]:
"""Search for similar memories."""
query_vec = np.array(query_embedding)
results = []
for mem_id, memory in self.memories.items():
# Filter by project and user
if memory["project_id"] != project_id:
continue
if user_id and memory.get("user_id") != user_id:
continue
if memory_type and memory.get("memory_type") != memory_type:
continue
if memory.get("importance", 0) < min_importance:
continue
# Calculate similarity if embedding exists
embedding_key = f"memory_{mem_id}"
if embedding_key in self.embeddings:
similarity = self._compute_similarity(query_vec, self.embeddings[embedding_key])
results.append({
"memory": Memory(**memory),
"similarity": similarity
})
# Sort by similarity and limit
results.sort(key=lambda x: x["similarity"], reverse=True)
results = results[:limit]
return [
MemoryResponse(
memory=r["memory"],
similarity=r["similarity"],
relevance_score=r["similarity"]
)
for r in results
]
async def get_recent_memories(
self,
project_id: str,
user_id: Optional[str] = None,
limit: int = 10,
memory_type: Optional[str] = None,
) -> List[Memory]:
"""Get recent memories."""
results = []
for memory in self.memories.values():
# Filter
if memory["project_id"] != project_id:
continue
if user_id and memory.get("user_id") != user_id:
continue
if memory_type and memory.get("memory_type") != memory_type:
continue
results.append(memory)
# Sort by timestamp
results.sort(key=lambda x: x.get("timestamp", ""), reverse=True)
results = results[:limit]
return [Memory(**r) for r in results]
async def delete_memory(self, memory_id: str, project_id: str) -> bool:
"""Delete a memory."""
if memory_id in self.memories:
if self.memories[memory_id]["project_id"] == project_id:
del self.memories[memory_id]
# Remove embedding
embedding_key = f"memory_{memory_id}"
if embedding_key in self.embeddings:
del self.embeddings[embedding_key]
# Save changes
self._save_json(self.memories, self.memories_file)
self._save_embeddings()
return True
return False
# Knowledge operations
async def create_fact(
self, fact: FactCreate, project_id: str, user_id: Optional[str] = None
) -> Fact:
"""Create a new fact."""
fact_id = str(uuid.uuid4())
fact_data = {
"id": fact_id,
"project_id": project_id,
"user_id": user_id or "default",
"subject": fact.subject,
"predicate": fact.predicate,
"object": fact.object,
"confidence": fact.confidence,
"source": fact.source,
"metadata": fact.metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
self.facts[fact_id] = fact_data
# Store embedding if provided
if fact.embedding:
self.embeddings[f"fact_{fact_id}"] = np.array(fact.embedding)
# Save to disk
self._save_json(self.facts, self.facts_file)
if fact.embedding:
self._save_embeddings()
return Fact(**fact_data)
async def search_facts(
self,
query_embedding: List[float],
project_id: str,
user_id: Optional[str] = None,
limit: int = 10,
min_confidence: float = 0.0,
) -> List[Fact]:
"""Search for similar facts."""
query_vec = np.array(query_embedding)
results = []
for fact_id, fact in self.facts.items():
# Filter
if fact["project_id"] != project_id:
continue
if user_id and fact.get("user_id") != user_id:
continue
if fact.get("confidence", 0) < min_confidence:
continue
# Calculate similarity
embedding_key = f"fact_{fact_id}"
if embedding_key in self.embeddings:
similarity = self._compute_similarity(query_vec, self.embeddings[embedding_key])
results.append({
"fact": fact,
"similarity": similarity
})
# Sort and limit
results.sort(key=lambda x: x["similarity"], reverse=True)
results = results[:limit]
return [Fact(**r["fact"]) for r in results]
async def get_knowledge_graph(
self, project_id: str, user_id: Optional[str] = None
) -> KnowledgeBase:
"""Get the knowledge graph."""
project_facts = []
for fact in self.facts.values():
if fact["project_id"] != project_id:
continue
if user_id and fact.get("user_id") != user_id:
continue
project_facts.append(Fact(**fact))
# Build entity and relation lists
entities = set()
relations = set()
for fact in project_facts:
entities.add(fact.subject)
entities.add(fact.object)
relations.add(fact.predicate)
return KnowledgeBase(
facts=project_facts,
entities=list(entities),
relations=list(relations),
metadata={
"fact_count": len(project_facts),
"entity_count": len(entities),
"relation_count": len(relations)
}
)
# Project operations
async def create_project(
self, project: ProjectCreate, user_id: Optional[str] = None
) -> Project:
"""Create a new project."""
project_id = str(uuid.uuid4())
project_data = {
"id": project_id,
"project_id": project_id, # Add project_id field
"name": project.name,
"description": project.description,
"user_id": user_id or "default",
"metadata": project.metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
self.projects[project_id] = project_data
self._save_json(self.projects, self.projects_file)
return Project(**project_data)
async def get_project(
self, project_id: str, user_id: Optional[str] = None
) -> Optional[Project]:
"""Get a project by ID."""
if project_id in self.projects:
project = self.projects[project_id]
if not user_id or project.get("user_id") == user_id:
return Project(**project)
return None
async def list_projects(self, user_id: Optional[str] = None) -> List[Project]:
"""List all projects."""
results = []
for project in self.projects.values():
if not user_id or project.get("user_id") == user_id:
# Ensure project_id exists for backward compatibility
if "project_id" not in project:
project["project_id"] = project.get("id", str(uuid.uuid4()))
results.append(Project(**project))
return results
async def delete_project(
self, project_id: str, user_id: Optional[str] = None
) -> bool:
"""Delete a project and all associated data."""
if project_id in self.projects:
project = self.projects[project_id]
if not user_id or project.get("user_id") == user_id:
# Delete project
del self.projects[project_id]
# Delete associated memories
memory_ids_to_delete = [
mid for mid, m in self.memories.items()
if m["project_id"] == project_id
]
for mid in memory_ids_to_delete:
del self.memories[mid]
embedding_key = f"memory_{mid}"
if embedding_key in self.embeddings:
del self.embeddings[embedding_key]
# Delete associated facts
fact_ids_to_delete = [
fid for fid, f in self.facts.items()
if f["project_id"] == project_id
]
for fid in fact_ids_to_delete:
del self.facts[fid]
embedding_key = f"fact_{fid}"
if embedding_key in self.embeddings:
del self.embeddings[embedding_key]
# Save all changes
self._save_json(self.projects, self.projects_file)
self._save_json(self.memories, self.memories_file)
self._save_json(self.facts, self.facts_file)
self._save_embeddings()
return True
return False
# Additional abstract methods implementation
async def create_memories_table(self, user_id: str = None) -> None:
"""Create memories table (no-op for local storage)."""
pass
def create_memories_table(self, user_id: str) -> None:
"""Create memories table for a user (no-op for local storage)."""
pass
def add_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
importance: float = 0.5,
) -> dict[str, Any]:
"""Add a memory synchronously."""
memory_data = {
"id": memory_id,
"memory_id": memory_id,
"project_id": project_id,
"user_id": user_id,
"content": content,
"memory_type": metadata.get("memory_type", "general") if metadata else "general",
"importance": importance,
"context": metadata.get("context", {}) if metadata else {},
"metadata": metadata or {},
"source": metadata.get("source", "") if metadata else "",
"timestamp": datetime.utcnow().isoformat(),
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
# Store memory
self.memories[memory_id] = memory_data
# Store embedding
if embedding:
self.embeddings[f"memory_{memory_id}"] = np.array(embedding)
# Save to disk
self._save_json(self.memories, self.memories_file)
if embedding:
self._save_embeddings()
return memory_data
async def add_memory_async(self, memory: MemoryCreate, project_id: str) -> Memory:
"""Add a memory (alias for create_memory)."""
return await self.create_memory(memory, project_id)
async def add_fact(self, fact: FactCreate, project_id: str) -> Fact:
"""Add a fact (alias for create_fact)."""
return await self.create_fact(fact, project_id)
async def delete_fact(self, fact_id: str) -> bool:
"""Delete a fact by ID."""
if fact_id in self.facts:
del self.facts[fact_id]
embedding_key = f"fact_{fact_id}"
if embedding_key in self.embeddings:
del self.embeddings[embedding_key]
self._save_json(self.facts, self.facts_file)
self._save_embeddings()
return True
return False
async def create_knowledge_base(self, name: str, description: str, project_id: str) -> KnowledgeBase:
"""Create a knowledge base."""
return await self.get_knowledge_graph(project_id)
async def get_knowledge_bases(self, project_id: str) -> List[KnowledgeBase]:
"""Get all knowledge bases for a project."""
return [await self.get_knowledge_graph(project_id)]
async def create_chat_session(self, project_id: str, user_id: Optional[str] = None) -> str:
"""Create a new chat session."""
session_id = str(uuid.uuid4())
if not hasattr(self, "chat_sessions"):
self.chat_sessions = {}
self.chat_sessions[session_id] = {
"id": session_id,
"project_id": project_id,
"user_id": user_id or "default",
"messages": [],
"created_at": datetime.utcnow().isoformat()
}
return session_id
async def add_chat_message(self, session_id: str, role: str, content: str) -> None:
"""Add a message to a chat session."""
if not hasattr(self, "chat_sessions"):
self.chat_sessions = {}
if session_id in self.chat_sessions:
self.chat_sessions[session_id]["messages"].append({
"role": role,
"content": content,
"timestamp": datetime.utcnow().isoformat()
})
async def get_chat_messages(self, session_id: str) -> List[Dict[str, Any]]:
"""Get messages from a chat session."""
if not hasattr(self, "chat_sessions"):
self.chat_sessions = {}
if session_id in self.chat_sessions:
return self.chat_sessions[session_id]["messages"]
return []
async def search_chat_messages(self, query: str, session_id: Optional[str] = None) -> List[Dict[str, Any]]:
"""Search chat messages."""
if not hasattr(self, "chat_sessions"):
return []
results = []
sessions_to_search = [session_id] if session_id else self.chat_sessions.keys()
for sid in sessions_to_search:
if sid in self.chat_sessions:
for msg in self.chat_sessions[sid]["messages"]:
if query.lower() in msg["content"].lower():
results.append(msg)
return results
async def get_user_projects(self, user_id: str) -> List[Project]:
"""Get all projects for a user."""
return await self.list_projects(user_id)
def update_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str | None = None,
metadata: dict | None = None,
importance: float | None = None,
) -> dict[str, Any] | None:
"""Update a memory in the database."""
if memory_id not in self.memories:
return None
memory = self.memories[memory_id]
# Verify user and project match
if memory.get("user_id") != user_id or memory.get("project_id") != project_id:
return None
# Update fields if provided
if content is not None:
memory["content"] = content
if metadata is not None:
memory["metadata"] = metadata
if importance is not None:
memory["importance"] = importance
# Update timestamp
memory["updated_at"] = datetime.utcnow().isoformat()
# Save changes
self._save_json(self.memories, self.memories_file)
return memory
def _import_markdown_memories(self):
"""Import memories from markdown files."""
if not self.markdown_reader:
return
try:
# Read markdown memories
markdown_memories = self.markdown_reader.read_markdown_memories()
# Create a default project for markdown memories
project_id = "markdown_import"
if project_id not in self.projects:
project_data = {
"id": project_id,
"project_id": project_id,
"name": "Markdown Import",
"description": "Automatically imported memories from markdown files",
"user_id": "system",
"metadata": {"auto_created": True},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
self.projects[project_id] = project_data
self._save_json(self.projects, self.projects_file)
# Import each memory
imported_count = 0
for memory_create in markdown_memories:
# Check if this memory already exists (by source)
source = memory_create.source
existing = any(
m.get("source") == source
for m in self.memories.values()
)
if not existing:
memory_id = str(uuid.uuid4())
memory_data = {
"id": memory_id,
"project_id": project_id,
"user_id": "system",
"content": memory_create.content,
"memory_type": memory_create.memory_type,
"importance": memory_create.importance,
"context": memory_create.context,
"metadata": memory_create.metadata or {},
"source": source,
"timestamp": datetime.utcnow().isoformat(),
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
self.memories[memory_id] = memory_data
imported_count += 1
if imported_count > 0:
self._save_json(self.memories, self.memories_file)
logger.info(f"Imported {imported_count} new memories from markdown files")
except Exception as e:
logger.error(f"Error importing markdown memories: {e}")
pkg/hanzo-memory/build/lib/hanzo_memory/db/markdown_reader.py
-381
View File
@@ -1,381 +0,0 @@
"""Markdown file reader for integrating LLM.md and similar files into memory."""
import hashlib
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Set
from structlog import get_logger
from ..models.memory import MemoryCreate
logger = get_logger()
# Supported markdown files for memory integration
MEMORY_MD_FILES = [
"LLM.md",
"AGENTS.md",
"CLAUDE.md",
"GEMINI.md",
"QWEN.md",
"AI.md",
"README.md", # Often contains important project context
"CONTEXT.md",
"MEMORY.md",
"INSTRUCTIONS.md",
"PROMPT.md",
]
# Additional patterns to look for
MEMORY_MD_PATTERNS = [
"*_LLM.md",
"*_AGENT.md",
"*_AI.md",
"*.claude.md",
"*.gemini.md",
"*.qwen.md",
]
class MarkdownMemoryReader:
"""Reads and processes markdown files for memory integration."""
def __init__(self, watch_dirs: Optional[List[Path]] = None):
"""Initialize the markdown reader.
Args:
watch_dirs: List of directories to watch for markdown files.
Defaults to current directory and parent directories.
"""
self.watch_dirs = watch_dirs or self._get_default_watch_dirs()
self.processed_files: Set[str] = set()
self.file_hashes: Dict[str, str] = {}
def _get_default_watch_dirs(self) -> List[Path]:
"""Get default directories to watch."""
dirs = []
current = Path.cwd()
# Add current directory
dirs.append(current)
# Add parent directories up to home or 3 levels
for _ in range(3):
if current.parent == current or current == Path.home():
break
current = current.parent
dirs.append(current)
# Add home .config directories
home = Path.home()
for config_dir in [".claude", ".gemini", ".qwen", ".ai", ".llm"]:
config_path = home / config_dir
if config_path.exists():
dirs.append(config_path)
return dirs
def _compute_file_hash(self, filepath: Path) -> str:
"""Compute hash of file contents."""
try:
content = filepath.read_text(encoding='utf-8')
return hashlib.sha256(content.encode()).hexdigest()
except Exception as e:
logger.error(f"Error hashing file {filepath}: {e}")
return ""
def _parse_markdown_sections(self, content: str, filepath: Path) -> List[Dict]:
"""Parse markdown content into sections."""
sections = []
current_section = {
"title": f"Content from {filepath.name}",
"content": "",
"level": 0,
"line_start": 0,
}
lines = content.split('\n')
current_line = 0
for line in lines:
current_line += 1
# Check for headers
if line.startswith('#'):
# Save previous section if it has content
if current_section["content"].strip():
sections.append(current_section)
# Start new section
header_level = len(line) - len(line.lstrip('#'))
header_text = line.lstrip('#').strip()
current_section = {
"title": header_text or f"Section from {filepath.name}",
"content": "",
"level": header_level,
"line_start": current_line,
}
else:
# Add line to current section
current_section["content"] += line + "\n"
# Don't forget the last section
if current_section["content"].strip():
sections.append(current_section)
# If no sections were found, treat entire content as one section
if not sections and content.strip():
sections.append({
"title": f"Content from {filepath.name}",
"content": content,
"level": 0,
"line_start": 0,
})
return sections
def find_markdown_files(self) -> List[Path]:
"""Find all relevant markdown files in watched directories."""
md_files = []
seen_files = set()
for watch_dir in self.watch_dirs:
if not watch_dir.exists():
continue
# Look for specific named files
for md_file in MEMORY_MD_FILES:
filepath = watch_dir / md_file
if filepath.exists() and filepath not in seen_files:
md_files.append(filepath)
seen_files.add(filepath)
# Look for pattern-matched files
for pattern in MEMORY_MD_PATTERNS:
for filepath in watch_dir.glob(pattern):
if filepath not in seen_files:
md_files.append(filepath)
seen_files.add(filepath)
# Sort by priority (LLM.md first, then others)
def priority(f: Path) -> int:
name = f.name.upper()
if name == "LLM.MD":
return 0
elif name == "CLAUDE.MD":
return 1
elif name == "GEMINI.MD":
return 2
elif name == "QWEN.MD":
return 3
elif name == "AGENTS.MD":
return 4
else:
return 5
md_files.sort(key=priority)
return md_files
def read_markdown_memories(self) -> List[MemoryCreate]:
"""Read all markdown files and convert to memories."""
memories = []
md_files = self.find_markdown_files()
for filepath in md_files:
try:
# Check if file has been processed and unchanged
file_id = str(filepath.absolute())
current_hash = self._compute_file_hash(filepath)
if file_id in self.file_hashes and self.file_hashes[file_id] == current_hash:
logger.debug(f"Skipping unchanged file: {filepath}")
continue
# Read and parse the file
content = filepath.read_text(encoding='utf-8')
if not content.strip():
continue
# Store hash for future comparison
self.file_hashes[file_id] = current_hash
# Parse into sections
sections = self._parse_markdown_sections(content, filepath)
# Create memories from sections
for section in sections:
# Determine importance based on file and section
importance = self._calculate_importance(filepath, section)
# Determine memory type
memory_type = self._determine_memory_type(filepath, section)
# Create memory
memory = MemoryCreate(
content=section["content"],
memory_type=memory_type,
importance=importance,
context={
"source_file": str(filepath.absolute()),
"file_name": filepath.name,
"section_title": section["title"],
"section_level": section["level"],
"line_start": section["line_start"],
"directory": str(filepath.parent.absolute()),
"file_modified": datetime.fromtimestamp(filepath.stat().st_mtime).isoformat(),
},
metadata={
"auto_imported": True,
"markdown_file": True,
},
source=f"markdown://{filepath.absolute()}#{section['line_start']}",
)
memories.append(memory)
logger.info(f"Read {len(sections)} sections from {filepath.name}")
except Exception as e:
logger.error(f"Error reading markdown file {filepath}: {e}")
continue
return memories
def _calculate_importance(self, filepath: Path, section: Dict) -> float:
"""Calculate importance score for a memory section."""
importance = 0.5 # Base importance
# File-based importance
filename = filepath.name.upper()
if filename == "LLM.MD":
importance += 0.3
elif filename in ["CLAUDE.MD", "GEMINI.MD", "QWEN.MD"]:
importance += 0.25
elif filename == "AGENTS.MD":
importance += 0.2
elif filename == "README.MD":
importance += 0.1
# Section-based importance
title_lower = section["title"].lower()
# High importance keywords
high_importance_keywords = [
"important", "critical", "essential", "must", "always",
"never", "warning", "error", "security", "key",
"architecture", "design", "api", "interface"
]
for keyword in high_importance_keywords:
if keyword in title_lower or keyword in section["content"].lower()[:200]:
importance += 0.1
break
# Header level importance (H1 > H2 > H3)
if section["level"] == 1:
importance += 0.15
elif section["level"] == 2:
importance += 0.1
elif section["level"] == 3:
importance += 0.05
# Cap at 1.0
return min(importance, 1.0)
def _determine_memory_type(self, filepath: Path, section: Dict) -> str:
"""Determine the type of memory based on content."""
filename = filepath.name.upper()
content_lower = section["content"].lower()
title_lower = section["title"].lower()
# Check filename patterns
if "AGENT" in filename:
return "agent_instruction"
elif filename in ["CLAUDE.MD", "GEMINI.MD", "QWEN.MD"]:
return "model_instruction"
elif filename == "LLM.MD":
return "system_context"
# Check content patterns
if any(word in content_lower[:500] for word in ["instruction", "prompt", "you should", "you must"]):
return "instruction"
elif any(word in content_lower[:500] for word in ["api", "endpoint", "function", "method", "class"]):
return "technical"
elif any(word in content_lower[:500] for word in ["example", "usage", "how to", "tutorial"]):
return "example"
elif any(word in title_lower for word in ["config", "setting", "environment", "variable"]):
return "configuration"
elif any(word in title_lower for word in ["architecture", "design", "structure", "pattern"]):
return "architectural"
return "knowledge"
def watch_for_changes(self) -> List[MemoryCreate]:
"""Check for new or modified markdown files and return new memories."""
new_memories = []
# Get current memories
current_memories = self.read_markdown_memories()
# Track which files were processed
for memory in current_memories:
source_file = memory.context.get("source_file")
if source_file:
self.processed_files.add(source_file)
return current_memories
def get_project_context(self, project_path: Optional[Path] = None) -> Dict:
"""Get comprehensive project context from markdown files."""
project_path = project_path or Path.cwd()
context = {
"project_path": str(project_path.absolute()),
"project_name": project_path.name,
"markdown_files": [],
"instructions": [],
"configurations": [],
"examples": [],
"architecture": [],
}
# Find and categorize markdown content
md_files = self.find_markdown_files()
for filepath in md_files:
try:
content = filepath.read_text(encoding='utf-8')
sections = self._parse_markdown_sections(content, filepath)
file_info = {
"path": str(filepath.absolute()),
"name": filepath.name,
"modified": datetime.fromtimestamp(filepath.stat().st_mtime).isoformat(),
"sections": len(sections),
}
context["markdown_files"].append(file_info)
# Categorize sections
for section in sections:
memory_type = self._determine_memory_type(filepath, section)
section_info = {
"title": section["title"],
"content_preview": section["content"][:200] + "..." if len(section["content"]) > 200 else section["content"],
"source": filepath.name,
}
if memory_type in ["instruction", "agent_instruction", "model_instruction"]:
context["instructions"].append(section_info)
elif memory_type == "configuration":
context["configurations"].append(section_info)
elif memory_type == "example":
context["examples"].append(section_info)
elif memory_type == "architectural":
context["architecture"].append(section_info)
except Exception as e:
logger.error(f"Error processing {filepath}: {e}")
return context
pkg/hanzo-memory/build/lib/hanzo_memory/db/mock_infinity.py
-150
View File
@@ -1,150 +0,0 @@
"""Mock InfinityDB implementation for testing and platforms without support."""
from pathlib import Path
from typing import Any
import numpy as np
import polars as pl
class MockDatabase:
"""Mock database implementation."""
def __init__(self, name: str):
self.name = name
self.tables: dict[str, dict[str, Any]] = {}
def create_table(self, table_name: str, schema: dict[str, Any]) -> None:
"""Create a table."""
if table_name not in self.tables:
self.tables[table_name] = {"schema": schema, "data": []}
def get_table(self, table_name: str) -> "MockTable":
"""Get a table."""
if table_name not in self.tables:
raise ValueError(f"Table {table_name} not found")
return MockTable(self, table_name)
class MockTable:
"""Mock table implementation."""
def __init__(self, db: MockDatabase, name: str):
self.db = db
self.name = name
def insert(self, records: list[dict[str, Any]]) -> None:
"""Insert records."""
self.db.tables[self.name]["data"].extend(records)
def output(self, columns: list[str]) -> "MockQuery":
"""Start a query."""
return MockQuery(self, columns)
class MockQuery:
"""Mock query implementation."""
def __init__(self, table: MockTable, columns: list[str]):
self.table = table
self.columns = columns
self.filters: list[str] = []
self.vector_search: dict[str, Any] | None = None
def match_dense(
self,
column: str,
query_vector: list[float],
dtype: str,
metric: str,
limit: int,
) -> "MockQuery":
"""Add vector search."""
self.vector_search = {
"column": column,
"query_vector": np.array(query_vector),
"metric": metric,
"limit": limit,
}
return self
def filter(self, condition: str) -> "MockQuery":
"""Add filter condition."""
self.filters.append(condition)
return self
def to_pl(self) -> pl.DataFrame:
"""Execute query and return polars DataFrame."""
# Get all data
data = self.table.db.tables[self.table.name]["data"]
# Apply filters
filtered_data = []
for record in data:
include = True
for filter_cond in self.filters:
# Simple filter parsing (e.g., "field = 'value'")
if " = " in filter_cond:
field, value = filter_cond.split(" = ")
value = value.strip("'\"")
if record.get(field) != value:
include = False
break
if include:
filtered_data.append(record)
# Apply vector search if specified
if self.vector_search and filtered_data:
# Calculate similarities
similarities = []
for record in filtered_data:
vec = np.array(record[self.vector_search["column"]])
query = self.vector_search["query_vector"]
if self.vector_search["metric"] == "cosine":
# Cosine similarity
sim = np.dot(vec, query) / (
np.linalg.norm(vec) * np.linalg.norm(query)
)
elif self.vector_search["metric"] == "ip":
# Inner product
sim = np.dot(vec, query)
else:
sim = 0.0
similarities.append((sim, record))
# Sort by similarity and limit
similarities.sort(key=lambda x: x[0], reverse=True)
filtered_data = [r for _, r in similarities[: self.vector_search["limit"]]]
# Return as polars DataFrame
if filtered_data:
return pl.DataFrame(filtered_data)
else:
# Return empty DataFrame with schema
return pl.DataFrame()
class MockInfinity:
"""Mock Infinity connection."""
def __init__(self, path: str):
self.path = Path(path)
self.databases: dict[str, MockDatabase] = {}
def create_database(self, name: str) -> None:
"""Create a database."""
if name not in self.databases:
self.databases[name] = MockDatabase(name)
def get_database(self, name: str) -> MockDatabase:
"""Get a database."""
if name not in self.databases:
self.create_database(name)
return self.databases[name]
def connect(path: str) -> MockInfinity:
"""Connect to mock InfinityDB."""
return MockInfinity(path)
pkg/hanzo-memory/build/lib/hanzo_memory/db/sqlite_client.py
-778
View File
@@ -1,778 +0,0 @@
"""SQLite-based memory storage implementation with vector search using sqlite-vec."""
import json
import sqlite3
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
import numpy as np
from structlog import get_logger
from .base import BaseVectorDB
logger = get_logger()
class SQLiteMemoryClient(BaseVectorDB):
"""SQLite-based implementation of the vector database with sqlite-vec for vector search."""
def __init__(self, db_path: Optional[Path] = None):
"""Initialize SQLite memory storage.
Args:
db_path: Path to SQLite database file. Defaults to in-memory if None.
"""
if db_path is None:
self.db_path = ":memory:"
else:
self.db_path = str(db_path)
# Connect to database
self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
self.conn.row_factory = sqlite3.Row # Enable column access by name
# Enable extension loading for sqlite-vec
self.conn.enable_load_extension(True)
# Initialize tables
self._init_tables()
logger.info(f"Initialized SQLite memory storage at {self.db_path}")
async def initialize(self) -> None:
"""Initialize the database (no-op for local storage)."""
logger.info("SQLite memory storage initialized")
def _init_tables(self):
"""Initialize required tables."""
# Enable sqlite-vec extension
try:
self.conn.execute("SELECT load_extension('vec');")
except sqlite3.Error as e:
logger.warning(f"Could not load sqlite-vec extension: {e}")
logger.warning("Vector search functionality will be limited")
# Create projects table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS projects (
id TEXT PRIMARY KEY,
project_id TEXT UNIQUE NOT NULL,
name TEXT NOT NULL,
description TEXT,
user_id TEXT NOT NULL,
metadata TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
""")
# Create memories table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS memories (
id TEXT PRIMARY KEY,
memory_id TEXT UNIQUE NOT NULL,
project_id TEXT NOT NULL,
user_id TEXT NOT NULL,
content TEXT NOT NULL,
memory_type TEXT DEFAULT 'general',
importance REAL DEFAULT 0.5,
context TEXT,
metadata TEXT,
source TEXT,
embedding BLOB,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (project_id) REFERENCES projects(project_id)
);
""")
# Create facts table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS facts (
id TEXT PRIMARY KEY,
fact_id TEXT UNIQUE NOT NULL,
knowledge_base_id TEXT NOT NULL,
content TEXT NOT NULL,
subject TEXT,
predicate TEXT,
object TEXT,
confidence REAL DEFAULT 1.0,
source TEXT,
embedding BLOB,
metadata TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
""")
# Create knowledge bases table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS knowledge_bases (
id TEXT PRIMARY KEY,
kb_id TEXT UNIQUE NOT NULL,
project_id TEXT NOT NULL,
name TEXT NOT NULL,
description TEXT,
metadata TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (project_id) REFERENCES projects(project_id)
);
""")
# Create chat sessions table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS chat_sessions (
id TEXT PRIMARY KEY,
session_id TEXT UNIQUE NOT NULL,
project_id TEXT NOT NULL,
user_id TEXT NOT NULL,
metadata TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
""")
# Create chat messages table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS chat_messages (
id TEXT PRIMARY KEY,
message_id TEXT UNIQUE NOT NULL,
session_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
embedding BLOB,
metadata TEXT,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (session_id) REFERENCES chat_sessions(session_id)
);
""")
# Create indexes
self.conn.execute("CREATE INDEX IF NOT EXISTS idx_memories_user_project ON memories(user_id, project_id);")
self.conn.execute("CREATE INDEX IF NOT EXISTS idx_memories_timestamp ON memories(timestamp);")
self.conn.execute("CREATE INDEX IF NOT EXISTS idx_facts_kb ON facts(knowledge_base_id);")
self.conn.execute("CREATE INDEX IF NOT EXISTS idx_chat_sessions_user_project ON chat_sessions(user_id, project_id);")
# Create vector index for embeddings if sqlite-vec is available
try:
# Create vector index for memories
self.conn.execute("""
CREATE INDEX IF NOT EXISTS idx_memories_embedding
ON memories (vec_to_json16(embedding))
WHERE embedding IS NOT NULL;
""")
# Create vector index for facts
self.conn.execute("""
CREATE INDEX IF NOT EXISTS idx_facts_embedding
ON facts (vec_to_json16(embedding))
WHERE embedding IS NOT NULL;
""")
# Create vector index for chat messages
self.conn.execute("""
CREATE INDEX IF NOT EXISTS idx_chat_messages_embedding
ON chat_messages (vec_to_json16(embedding))
WHERE embedding IS NOT NULL;
""")
except sqlite3.Error:
# If sqlite-vec is not available, skip vector indexes
pass
self.conn.commit()
def create_project(
self,
project_id: str,
user_id: str,
name: str,
description: str = "",
metadata: dict | None = None,
) -> dict[str, Any]:
"""Create a new project."""
project_id = project_id or str(uuid.uuid4())
metadata_json = json.dumps(metadata or {})
self.conn.execute(
"""
INSERT INTO projects (id, project_id, user_id, name, description, metadata)
VALUES (?, ?, ?, ?, ?, ?)
""",
(str(uuid.uuid4()), project_id, user_id, name, description, metadata_json)
)
self.conn.commit()
return {
"id": project_id,
"project_id": project_id,
"user_id": user_id,
"name": name,
"description": description,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
def get_user_projects(self, user_id: str) -> list[dict[str, Any]]:
"""Get all projects for a user."""
cursor = self.conn.execute(
"SELECT * FROM projects WHERE user_id = ?", (user_id,)
)
rows = cursor.fetchall()
projects = []
for row in rows:
projects.append({
"id": row["id"],
"project_id": row["project_id"],
"user_id": row["user_id"],
"name": row["name"],
"description": row["description"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"created_at": row["created_at"],
"updated_at": row["updated_at"],
})
return projects
def create_memories_table(self, user_id: str) -> None:
"""Create a memories table for a user (already handled in initialization)."""
# Table is created during initialization
pass
def add_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
importance: float = 0.5,
) -> dict[str, Any]:
"""Add a memory to the database."""
memory_id = memory_id or str(uuid.uuid4())
metadata_json = json.dumps(metadata or {})
context_json = json.dumps({})
embedding_blob = np.array(embedding, dtype=np.float32).tobytes() if embedding else None
self.conn.execute(
"""
INSERT INTO memories
(id, memory_id, user_id, project_id, content, importance, context, metadata, source, embedding)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
str(uuid.uuid4()),
memory_id,
user_id,
project_id,
content,
importance,
context_json,
metadata_json,
"", # source
embedding_blob
)
)
self.conn.commit()
return {
"id": memory_id,
"memory_id": memory_id,
"user_id": user_id,
"project_id": project_id,
"content": content,
"importance": importance,
"context": {},
"metadata": metadata or {},
"source": "",
"timestamp": datetime.utcnow().isoformat(),
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
def search_memories(
self,
user_id: str,
query_embedding: list[float],
project_id: str | None = None,
limit: int = 10,
min_similarity: float = 0.0,
) -> list[dict[str, Any]]:
"""Search memories by similarity."""
query_array = np.array(query_embedding, dtype=np.float32)
query_blob = query_array.tobytes()
# Build query
where_conditions = ["user_id = ?"]
params = [user_id]
if project_id:
where_conditions.append("project_id = ?")
params.append(project_id)
where_clause = " AND ".join(where_conditions)
# If sqlite-vec is available, use vector similarity search
try:
# Use sqlite-vec for similarity search
cursor = self.conn.execute(f"""
SELECT *, vec_distance_L2(embedding, ?) as distance
FROM memories
WHERE {where_clause} AND embedding IS NOT NULL
ORDER BY distance ASC
LIMIT ?
""", [query_blob] + params + [limit])
rows = cursor.fetchall()
results = []
for row in rows:
# Calculate similarity from distance (convert L2 distance to similarity)
distance = row["distance"]
similarity = 1 / (1 + distance) # Convert distance to similarity score
if similarity >= min_similarity:
results.append({
"memory_id": row["memory_id"],
"user_id": row["user_id"],
"project_id": row["project_id"],
"content": row["content"],
"importance": row["importance"],
"context": json.loads(row["context"]) if row["context"] else {},
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"source": row["source"],
"created_at": row["created_at"],
"updated_at": row["updated_at"],
"similarity_score": similarity
})
return results
except sqlite3.Error:
# Fallback to basic search without vector similarity
cursor = self.conn.execute(f"""
SELECT *
FROM memories
WHERE {where_clause}
ORDER BY timestamp DESC
LIMIT ?
""", params + [limit])
rows = cursor.fetchall()
results = []
for row in rows:
results.append({
"memory_id": row["memory_id"],
"user_id": row["user_id"],
"project_id": row["project_id"],
"content": row["content"],
"importance": row["importance"],
"context": json.loads(row["context"]) if row["context"] else {},
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"source": row["source"],
"created_at": row["created_at"],
"updated_at": row["updated_at"],
"similarity_score": 0.0 # No similarity calculation without sqlite-vec
})
return results
def create_knowledge_base(
self,
knowledge_base_id: str,
project_id: str,
name: str,
description: str = "",
metadata: dict | None = None,
) -> dict[str, Any]:
"""Create a new knowledge base."""
knowledge_base_id = knowledge_base_id or str(uuid.uuid4())
metadata_json = json.dumps(metadata or {})
self.conn.execute(
"""
INSERT INTO knowledge_bases (id, kb_id, project_id, name, description, metadata)
VALUES (?, ?, ?, ?, ?, ?)
""",
(str(uuid.uuid4()), knowledge_base_id, project_id, name, description, metadata_json)
)
self.conn.commit()
return {
"id": knowledge_base_id,
"kb_id": knowledge_base_id,
"project_id": project_id,
"name": name,
"description": description,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
def get_knowledge_bases(self, project_id: str) -> list[dict[str, Any]]:
"""Get all knowledge bases for a project."""
cursor = self.conn.execute(
"SELECT * FROM knowledge_bases WHERE project_id = ?", (project_id,)
)
rows = cursor.fetchall()
kbs = []
for row in rows:
kbs.append({
"id": row["id"],
"kb_id": row["kb_id"],
"project_id": row["project_id"],
"name": row["name"],
"description": row["description"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"created_at": row["created_at"],
"updated_at": row["updated_at"],
})
return kbs
def add_fact(
self,
fact_id: str,
knowledge_base_id: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
confidence: float = 1.0,
) -> dict[str, Any]:
"""Add a fact to a knowledge base."""
fact_id = fact_id or str(uuid.uuid4())
metadata_json = json.dumps(metadata or {})
embedding_blob = np.array(embedding, dtype=np.float32).tobytes() if embedding else None
self.conn.execute(
"""
INSERT INTO facts
(id, fact_id, knowledge_base_id, content, embedding, metadata, confidence)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(str(uuid.uuid4()), fact_id, knowledge_base_id, content, embedding_blob, metadata_json, confidence)
)
self.conn.commit()
return {
"id": fact_id,
"fact_id": fact_id,
"knowledge_base_id": knowledge_base_id,
"content": content,
"confidence": confidence,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
def search_facts(
self,
knowledge_base_id: str,
query_embedding: list[float] | None = None,
limit: int = 10,
) -> list[dict[str, Any]]:
"""Search facts in a knowledge base."""
if query_embedding is not None:
# If sqlite-vec is available and we have an embedding, use vector search
query_array = np.array(query_embedding, dtype=np.float32)
query_blob = query_array.tobytes()
try:
cursor = self.conn.execute("""
SELECT *, vec_distance_L2(embedding, ?) as distance
FROM facts
WHERE knowledge_base_id = ? AND embedding IS NOT NULL
ORDER BY distance ASC
LIMIT ?
""", [query_blob, knowledge_base_id, limit])
rows = cursor.fetchall()
results = []
for row in rows:
distance = row["distance"]
similarity = 1 / (1 + distance) # Convert distance to similarity score
results.append({
"fact_id": row["fact_id"],
"knowledge_base_id": row["knowledge_base_id"],
"content": row["content"],
"confidence": row["confidence"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"created_at": row["created_at"],
"updated_at": row["updated_at"],
"similarity_score": similarity
})
return results
except sqlite3.Error:
# Fallback to basic search
pass
# Basic search without vector similarity
cursor = self.conn.execute(
"SELECT * FROM facts WHERE knowledge_base_id = ? ORDER BY created_at DESC LIMIT ?",
(knowledge_base_id, limit)
)
rows = cursor.fetchall()
results = []
for row in rows:
results.append({
"fact_id": row["fact_id"],
"knowledge_base_id": row["knowledge_base_id"],
"content": row["content"],
"confidence": row["confidence"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"created_at": row["created_at"],
"updated_at": row["updated_at"],
"similarity_score": 0.0 # No similarity without embedding
})
return results
def delete_fact(self, fact_id: str, knowledge_base_id: str) -> bool:
"""Delete a fact from a knowledge base."""
cursor = self.conn.execute(
"DELETE FROM facts WHERE fact_id = ? AND knowledge_base_id = ?",
(fact_id, knowledge_base_id)
)
self.conn.commit()
return cursor.rowcount > 0
def update_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str | None = None,
metadata: dict | None = None,
importance: float | None = None,
) -> dict[str, Any] | None:
"""Update a memory in the database."""
# Check if memory exists and belongs to user/project
cursor = self.conn.execute(
"SELECT * FROM memories WHERE memory_id = ? AND user_id = ? AND project_id = ?",
(memory_id, user_id, project_id)
)
row = cursor.fetchone()
if not row:
return None
# Build update query
updates = []
params = []
if content is not None:
updates.append("content = ?")
params.append(content)
if metadata is not None:
updates.append("metadata = ?")
params.append(json.dumps(metadata))
if importance is not None:
updates.append("importance = ?")
params.append(importance)
if updates:
updates.append("updated_at = CURRENT_TIMESTAMP")
query = f"UPDATE memories SET {', '.join(updates)} WHERE memory_id = ?"
params.append(memory_id)
self.conn.execute(query, params)
self.conn.commit()
# Return updated memory
cursor = self.conn.execute(
"SELECT * FROM memories WHERE memory_id = ?", (memory_id,)
)
row = cursor.fetchone()
if row:
return {
"id": row["id"],
"memory_id": row["memory_id"],
"user_id": row["user_id"],
"project_id": row["project_id"],
"content": row["content"],
"importance": row["importance"],
"context": json.loads(row["context"]) if row["context"] else {},
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"source": row["source"],
"timestamp": row["timestamp"],
"created_at": row["created_at"],
"updated_at": row["updated_at"],
}
return None
def create_chat_session(
self,
session_id: str,
user_id: str,
project_id: str,
metadata: dict | None = None,
) -> dict[str, Any]:
"""Create a new chat session."""
session_id = session_id or str(uuid.uuid4())
metadata_json = json.dumps(metadata or {})
self.conn.execute(
"""
INSERT INTO chat_sessions (id, session_id, user_id, project_id, metadata)
VALUES (?, ?, ?, ?, ?)
""",
(str(uuid.uuid4()), session_id, user_id, project_id, metadata_json)
)
self.conn.commit()
return {
"id": session_id,
"session_id": session_id,
"user_id": user_id,
"project_id": project_id,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
}
def add_chat_message(
self,
message_id: str,
session_id: str,
role: str,
content: str,
embedding: list[float],
metadata: dict | None = None,
) -> dict[str, Any]:
"""Add a message to a chat session."""
message_id = message_id or str(uuid.uuid4())
metadata_json = json.dumps(metadata or {})
embedding_blob = np.array(embedding, dtype=np.float32).tobytes() if embedding else None
self.conn.execute(
"""
INSERT INTO chat_messages
(id, message_id, session_id, role, content, embedding, metadata)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(str(uuid.uuid4()), message_id, session_id, role, content, embedding_blob, metadata_json)
)
self.conn.commit()
return {
"id": message_id,
"message_id": message_id,
"session_id": session_id,
"role": role,
"content": content,
"metadata": metadata or {},
"timestamp": datetime.utcnow().isoformat(),
}
def get_chat_messages(
self,
session_id: str,
limit: int | None = None,
) -> list[dict[str, Any]]:
"""Get messages from a chat session."""
query = "SELECT * FROM chat_messages WHERE session_id = ? ORDER BY timestamp ASC"
params = [session_id]
if limit:
query += f" LIMIT {limit}"
cursor = self.conn.execute(query, params)
rows = cursor.fetchall()
messages = []
for row in rows:
messages.append({
"id": row["id"],
"message_id": row["message_id"],
"session_id": row["session_id"],
"role": row["role"],
"content": row["content"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"timestamp": row["timestamp"],
})
return messages
def search_chat_messages(
self,
session_id: str,
query_embedding: list[float],
limit: int = 10,
) -> list[dict[str, Any]]:
"""Search messages in a chat session by similarity."""
query_array = np.array(query_embedding, dtype=np.float32)
query_blob = query_array.tobytes()
# If sqlite-vec is available, use vector similarity search
try:
cursor = self.conn.execute("""
SELECT *, vec_distance_L2(embedding, ?) as distance
FROM chat_messages
WHERE session_id = ? AND embedding IS NOT NULL
ORDER BY distance ASC
LIMIT ?
""", [query_blob, session_id, limit])
rows = cursor.fetchall()
results = []
for row in rows:
distance = row["distance"]
similarity = 1 / (1 + distance) # Convert distance to similarity score
results.append({
"id": row["id"],
"message_id": row["message_id"],
"session_id": row["session_id"],
"role": row["role"],
"content": row["content"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"timestamp": row["timestamp"],
"similarity_score": similarity
})
return results
except sqlite3.Error:
# Fallback to basic search without vector similarity
cursor = self.conn.execute(
"SELECT * FROM chat_messages WHERE session_id = ? ORDER BY timestamp DESC LIMIT ?",
(session_id, limit)
)
rows = cursor.fetchall()
results = []
for row in rows:
results.append({
"id": row["id"],
"message_id": row["message_id"],
"session_id": row["session_id"],
"role": row["role"],
"content": row["content"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"timestamp": row["timestamp"],
"similarity_score": 0.0 # No similarity calculation without sqlite-vec
})
return results
def close(self) -> None:
"""Close the database connection."""
if self.conn:
self.conn.close()
pkg/hanzo-memory/build/lib/hanzo_memory/mcp/__init__.py
-5
View File
@@ -1,5 +0,0 @@
"""MCP server for Hanzo Memory service."""
from .server import MCPMemoryServer
__all__ = ["MCPMemoryServer"]
pkg/hanzo-memory/build/lib/hanzo_memory/mcp/__main__.py
-6
View File
@@ -1,6 +0,0 @@
"""MCP server entry point."""
from .server import main
if __name__ == "__main__":
main()
pkg/hanzo-memory/build/lib/hanzo_memory/mcp/server.py
-459
View File
@@ -1,459 +0,0 @@
"""MCP server for Hanzo Memory service."""
import json
from typing import Any
from uuid import uuid4
from mcp.server import Server
from mcp.server.models import InitializationOptions, ServerCapabilities
from mcp.types import (
TextContent,
Tool,
)
from structlog import get_logger
from ..config import settings
from ..db import get_db_client
from ..services.embeddings import EmbeddingService
from ..services.llm import LLMService
logger = get_logger()
class MCPMemoryServer:
"""MCP server for memory operations."""
def __init__(self) -> None:
"""Initialize the MCP server."""
self.server: Server = Server(settings.mcp_server_name)
self.db_client = get_db_client()
self.embedding_service = EmbeddingService()
self.llm_service = LLMService()
self._setup_handlers()
def _setup_handlers(self) -> None:
"""Set up server handlers."""
@self.server.list_tools() # type: ignore[misc]
async def handle_list_tools() -> list[Tool]:
"""Return available tools."""
return [
Tool(
name="remember",
description="Store a memory for a user in a project",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"project_id": {
"type": "string",
"description": "Project ID",
},
"content": {
"type": "string",
"description": "Memory content",
},
"metadata": {
"type": "object",
"description": "Additional metadata",
},
"importance": {
"type": "number",
"description": "Importance score (0-10)",
"default": 1.0,
},
},
"required": ["user_id", "project_id", "content"],
},
),
Tool(
name="recall",
description="Search for relevant memories",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"project_id": {
"type": "string",
"description": "Project ID (optional)",
},
"query": {"type": "string", "description": "Search query"},
"limit": {
"type": "integer",
"description": "Number of results",
"default": 10,
},
},
"required": ["user_id", "query"],
},
),
Tool(
name="create_project",
description="Create a new project",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"name": {"type": "string", "description": "Project name"},
"description": {
"type": "string",
"description": "Project description",
},
"metadata": {
"type": "object",
"description": "Additional metadata",
},
},
"required": ["user_id", "name"],
},
),
Tool(
name="create_knowledge_base",
description="Create a knowledge base in a project",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"project_id": {
"type": "string",
"description": "Project ID",
},
"name": {
"type": "string",
"description": "Knowledge base name",
},
"description": {
"type": "string",
"description": "Knowledge base description",
},
},
"required": ["user_id", "project_id", "name"],
},
),
Tool(
name="add_fact",
description="Add a fact to a knowledge base",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"kb_id": {
"type": "string",
"description": "Knowledge base ID",
},
"content": {
"type": "string",
"description": "Fact content",
},
"parent_id": {
"type": "string",
"description": "Parent fact ID (optional)",
},
"metadata": {
"type": "object",
"description": "Additional metadata",
},
},
"required": ["user_id", "kb_id", "content"],
},
),
Tool(
name="search_facts",
description="Search facts in a knowledge base",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"kb_id": {
"type": "string",
"description": "Knowledge base ID",
},
"query": {"type": "string", "description": "Search query"},
"limit": {
"type": "integer",
"description": "Number of results",
"default": 10,
},
},
"required": ["user_id", "kb_id", "query"],
},
),
Tool(
name="summarize_for_knowledge",
description="Analyze content and generate knowledge update instructions",
inputSchema={
"type": "object",
"properties": {
"content": {
"type": "string",
"description": "Content to analyze",
},
"context": {
"type": "string",
"description": "Additional context",
},
"skip_summarization": {
"type": "boolean",
"description": "Skip summarization",
"default": False,
},
"provided_summary": {
"type": "string",
"description": "Pre-made summary",
},
},
"required": ["content"],
},
),
]
@self.server.call_tool() # type: ignore[misc]
async def handle_call_tool(
name: str, arguments: dict[str, Any] | None = None
) -> list[TextContent]:
"""Handle tool calls."""
try:
if name == "remember":
result = await self._handle_remember(arguments or {})
elif name == "recall":
result = await self._handle_recall(arguments or {})
elif name == "create_project":
result = await self._handle_create_project(arguments or {})
elif name == "create_knowledge_base":
result = await self._handle_create_knowledge_base(arguments or {})
elif name == "add_fact":
result = await self._handle_add_fact(arguments or {})
elif name == "search_facts":
result = await self._handle_search_facts(arguments or {})
elif name == "summarize_for_knowledge":
result = await self._handle_summarize_for_knowledge(arguments or {})
else:
result = {"error": f"Unknown tool: {name}"}
return [TextContent(type="text", text=json.dumps(result, indent=2))]
except Exception as e:
logger.error(f"Error handling tool {name}: {e}")
return [TextContent(type="text", text=json.dumps({"error": str(e)}))]
async def _handle_remember(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle remember tool."""
user_id = args["user_id"]
project_id = args["project_id"]
content = args["content"]
metadata = args.get("metadata", {})
importance = args.get("importance", 1.0)
# Ensure user's memory table exists
self.db_client.create_memories_table(user_id)
# Generate embedding
embedding = self.embedding_service.embed_text(content)[0]
# Store memory
memory_id = str(uuid4())
self.db_client.add_memory(
memory_id=memory_id,
user_id=user_id,
project_id=project_id,
content=content,
embedding=embedding,
metadata=metadata,
importance=importance,
)
return {
"success": True,
"memory_id": memory_id,
"message": "Memory stored successfully",
}
async def _handle_recall(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle recall tool."""
user_id = args["user_id"]
project_id = args.get("project_id")
query = args["query"]
limit = args.get("limit", 10)
# Generate query embedding
query_embedding = self.embedding_service.embed_text(query)[0]
# Search memories
results_df = self.db_client.search_memories(
user_id=user_id,
query_embedding=query_embedding,
project_id=project_id,
limit=limit,
)
# Convert results
memories = []
if not results_df.is_empty():
for row in results_df.to_dicts():
memories.append(
{
"memory_id": row["memory_id"],
"content": row["content"],
"metadata": json.loads(row.get("metadata", "{}")),
"importance": row.get("importance", 1.0),
"similarity_score": row.get("_similarity", 0.0),
}
)
return {
"success": True,
"memories": memories,
"count": len(memories),
}
async def _handle_create_project(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle create project tool."""
user_id = args["user_id"]
name = args["name"]
description = args.get("description", "")
metadata = args.get("metadata", {})
# Create project
project_id = str(uuid4())
self.db_client.create_project(
project_id=project_id,
user_id=user_id,
name=name,
description=description,
metadata=metadata,
)
return {
"success": True,
"project_id": project_id,
"message": "Project created successfully",
}
async def _handle_create_knowledge_base(
self, args: dict[str, Any]
) -> dict[str, Any]:
"""Handle create knowledge base tool."""
user_id = args["user_id"]
project_id = args["project_id"]
name = args["name"]
description = args.get("description", "")
# Create knowledge base
kb_id = str(uuid4())
self.db_client.create_knowledge_base(
kb_id=kb_id,
user_id=user_id,
project_id=project_id,
name=name,
description=description,
)
return {
"success": True,
"kb_id": kb_id,
"message": "Knowledge base created successfully",
}
async def _handle_add_fact(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle add fact tool."""
kb_id = args["kb_id"]
content = args["content"]
parent_id = args.get("parent_id")
metadata = args.get("metadata", {})
# Generate embedding
embedding = self.embedding_service.embed_text(content)[0]
# Add fact
fact_id = str(uuid4())
self.db_client.add_fact(
fact_id=fact_id,
kb_id=kb_id,
content=content,
embedding=embedding,
parent_id=parent_id,
metadata=metadata,
)
return {
"success": True,
"fact_id": fact_id,
"message": "Fact added successfully",
}
async def _handle_search_facts(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle search facts tool."""
kb_id = args["kb_id"]
query = args["query"]
limit = args.get("limit", 10)
# Generate query embedding
query_embedding = self.embedding_service.embed_text(query)[0]
# Search facts
results_df = self.db_client.search_facts(
kb_id=kb_id,
query_embedding=query_embedding,
limit=limit,
)
# Convert results
facts = []
if not results_df.is_empty():
for row in results_df.to_dicts():
facts.append(
{
"fact_id": row["fact_id"],
"content": row["content"],
"parent_id": row.get("parent_id"),
"metadata": json.loads(row.get("metadata", "{}")),
"similarity_score": row.get("_similarity", 0.0),
}
)
return {
"success": True,
"facts": facts,
"count": len(facts),
}
async def _handle_summarize_for_knowledge(
self, args: dict[str, Any]
) -> dict[str, Any]:
"""Handle summarize for knowledge tool."""
content = args["content"]
context = args.get("context")
skip_summarization = args.get("skip_summarization", False)
provided_summary = args.get("provided_summary")
# Use LLM service to generate knowledge instructions
result = self.llm_service.summarize_for_knowledge(
content=content,
context=context,
skip_summarization=skip_summarization,
provided_summary=provided_summary,
)
return result
async def run(self) -> None:
"""Run the MCP server."""
from mcp.server.stdio import stdio_server
async with stdio_server() as (read_stream, write_stream):
await self.server.run(
read_stream,
write_stream,
InitializationOptions(
server_name=settings.mcp_server_name,
server_version=settings.mcp_server_version,
capabilities=ServerCapabilities(),
),
)
def main() -> None:
"""Main entry point for MCP server."""
import asyncio
server = MCPMemoryServer()
asyncio.run(server.run())
pkg/hanzo-memory/build/lib/hanzo_memory/memory.py
-255
View File
@@ -1,255 +0,0 @@
"""Unified memory interface with backend selection.
This module provides a clean API for accessing different memory backends
using a simple syntax like:
memory['lancedb'] # Access LanceDB backend
memory['kuzudb'] # Access KuzuDB backend
memory['infinity'] # Access InfinityDB backend
memory.local # Access local backend (default)
"""
from typing import Dict, Optional, Any
from structlog import get_logger
from .db.factory import get_db_client, reset_db_client, list_available_backends
from .db.base import BaseVectorDB
from .services.memory import MemoryService
from .services.embeddings import EmbeddingService
from .services.llm import LLMService
logger = get_logger()
class MemoryBackendProxy:
"""Proxy for accessing specific memory backends."""
def __init__(self, backend_name: str, config: Optional[Dict] = None):
"""Initialize backend proxy.
Args:
backend_name: Name of the backend to use
config: Optional backend configuration
"""
self.backend_name = backend_name
self.config = config or {}
self._client = None
self._service = None
@property
def client(self) -> BaseVectorDB:
"""Get the database client for this backend."""
if self._client is None:
reset_db_client()
self._client = get_db_client(
backend=self.backend_name,
config=self.config
)
return self._client
@property
def service(self) -> MemoryService:
"""Get the memory service for this backend."""
if self._service is None:
# Initialize services
embeddings = EmbeddingService()
llm = LLMService()
# Create service with this backend
self._service = MemoryService()
self._service.db = self.client
self._service.embeddings = embeddings
self._service.llm = llm
return self._service
async def initialize(self):
"""Initialize the backend."""
await self.client.initialize()
logger.info(f"Initialized {self.backend_name} backend")
async def close(self):
"""Close the backend connection."""
if self._client:
await self._client.close()
logger.info(f"Closed {self.backend_name} backend")
def __repr__(self):
"""String representation."""
return f"<MemoryBackend:{self.backend_name}>"
class Memory:
"""Main memory interface with backend selection.
Usage:
# Using dictionary syntax
memory['lancedb'] # Get LanceDB backend
memory['kuzudb'] # Get KuzuDB backend
memory['infinity'] # Get InfinityDB backend
# Using attribute syntax
memory.lancedb # Get LanceDB backend
memory.kuzudb # Get KuzuDB backend
memory.infinity # Get InfinityDB backend
memory.local # Get local backend (default)
# List available backends
memory.backends() # Returns list of available backends
# Use a specific backend
lance = memory['lancedb']
await lance.initialize()
await lance.service.create_memory(...)
# Or with context manager (auto initialize/close)
async with memory.use('lancedb') as backend:
await backend.service.create_memory(...)
"""
def __init__(self, default_backend: str = "local"):
"""Initialize memory interface.
Args:
default_backend: Default backend to use
"""
self.default_backend = default_backend
self._backends: Dict[str, MemoryBackendProxy] = {}
def __getitem__(self, backend_name: str) -> MemoryBackendProxy:
"""Get a backend using dictionary syntax.
Args:
backend_name: Name of the backend
Returns:
Backend proxy instance
"""
if backend_name not in self._backends:
self._backends[backend_name] = MemoryBackendProxy(backend_name)
return self._backends[backend_name]
def __getattr__(self, backend_name: str) -> MemoryBackendProxy:
"""Get a backend using attribute syntax.
Args:
backend_name: Name of the backend
Returns:
Backend proxy instance
"""
# Handle special attributes
if backend_name.startswith('_'):
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{backend_name}'")
# Convert attribute name to backend name (e.g., 'infinity_db' -> 'infinitydb')
backend_name = backend_name.replace('_', '')
return self[backend_name]
def backends(self) -> Dict[str, Dict]:
"""List all available backends.
Returns:
Dictionary of backend names to their info
"""
return list_available_backends()
def use(self, backend_name: str, config: Optional[Dict] = None):
"""Context manager for using a specific backend.
Args:
backend_name: Name of the backend
config: Optional backend configuration
Returns:
Async context manager for the backend
"""
return MemoryBackendContext(backend_name, config)
def __repr__(self):
"""String representation."""
backends = list(self.backends().keys())
return f"<Memory backends={backends}>"
class MemoryBackendContext:
"""Async context manager for using a memory backend."""
def __init__(self, backend_name: str, config: Optional[Dict] = None):
"""Initialize context.
Args:
backend_name: Name of the backend
config: Optional backend configuration
"""
self.backend = MemoryBackendProxy(backend_name, config)
async def __aenter__(self):
"""Enter context - initialize backend."""
await self.backend.initialize()
return self.backend
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Exit context - close backend."""
await self.backend.close()
# Global memory instance
memory = Memory()
# Convenience functions for common operations
async def remember(content: str, backend: str = "local", **kwargs):
"""Quick function to remember something.
Args:
content: Content to remember
backend: Backend to use (default: local)
**kwargs: Additional parameters for memory creation
Returns:
Created memory object
"""
async with memory.use(backend) as mem:
from .models.memory import MemoryCreate
memory_obj = MemoryCreate(
content=content,
**kwargs
)
return await mem.service.create_memory(
memory=memory_obj,
project_id=kwargs.get('project_id', 'default'),
user_id=kwargs.get('user_id', 'default')
)
async def recall(query: str, backend: str = "local", **kwargs):
"""Quick function to recall memories.
Args:
query: Search query
backend: Backend to use (default: local)
**kwargs: Additional search parameters
Returns:
List of matching memories
"""
async with memory.use(backend) as mem:
return await mem.service.search_memories(
query=query,
user_id=kwargs.get('user_id', 'default'),
project_id=kwargs.get('project_id', 'default'),
limit=kwargs.get('limit', 10)
)
# Export main components
__all__ = [
'Memory',
'memory',
'remember',
'recall',
'MemoryBackendProxy',
'MemoryBackendContext',
]
pkg/hanzo-memory/build/lib/hanzo_memory/models/__init__.py
-106
View File
@@ -1,106 +0,0 @@
"""Models package."""
from .base import ProjectScopedModel, TimestampedModel, UserScopedModel
from .chat import (
ChatHistoryRequest,
ChatMessage,
ChatMessageBase,
ChatMessageCreate,
ChatSearchRequest,
ChatSession,
ChatSessionCreate,
ChatSessionList,
)
from .knowledge import (
AddKnowledgeRequest,
CreateKnowledgeBaseRequest,
DeleteKnowledgeRequest,
Fact,
FactBase,
FactCreate,
FactRelation,
FactUpdate,
FactWithScore,
GetKnowledgeRequest,
IngestKnowledgeRequest,
KnowledgeBase,
KnowledgeBaseBase,
KnowledgeBaseCreate,
KnowledgeBaseUpdate,
ListKnowledgeBasesRequest,
)
from .memory import (
AddMemoriesRequest,
DeleteMemoryRequest,
DeleteUserRequest,
GetMemoriesRequest,
Memory,
MemoryBase,
MemoryCreate,
MemoryListResponse,
MemoryResponse,
MemoryUpdate,
MemoryWithScore,
RememberRequest,
UpdateMemoryRequest,
)
from .project import (
Project,
ProjectBase,
ProjectCreate,
ProjectList,
ProjectUpdate,
)
__all__ = [
# Base
"TimestampedModel",
"UserScopedModel",
"ProjectScopedModel",
# Project
"Project",
"ProjectBase",
"ProjectCreate",
"ProjectUpdate",
"ProjectList",
# Memory
"Memory",
"MemoryBase",
"MemoryCreate",
"MemoryUpdate",
"MemoryWithScore",
"MemoryResponse",
"MemoryListResponse",
"RememberRequest",
"AddMemoriesRequest",
"GetMemoriesRequest",
"DeleteMemoryRequest",
"DeleteUserRequest",
"UpdateMemoryRequest",
# Knowledge
"KnowledgeBase",
"KnowledgeBaseBase",
"KnowledgeBaseCreate",
"KnowledgeBaseUpdate",
"Fact",
"FactBase",
"FactCreate",
"FactUpdate",
"FactWithScore",
"FactRelation",
"CreateKnowledgeBaseRequest",
"ListKnowledgeBasesRequest",
"AddKnowledgeRequest",
"GetKnowledgeRequest",
"DeleteKnowledgeRequest",
"IngestKnowledgeRequest",
# Chat
"ChatMessage",
"ChatMessageBase",
"ChatMessageCreate",
"ChatSession",
"ChatSessionCreate",
"ChatSessionList",
"ChatHistoryRequest",
"ChatSearchRequest",
]
pkg/hanzo-memory/build/lib/hanzo_memory/models/base.py
-24
View File
@@ -1,24 +0,0 @@
"""Base models for Hanzo Memory Service."""
from datetime import datetime
from pydantic import BaseModel, Field
class TimestampedModel(BaseModel):
"""Base model with timestamps."""
created_at: datetime = Field(default_factory=datetime.utcnow)
updated_at: datetime = Field(default_factory=datetime.utcnow)
class UserScopedModel(BaseModel):
"""Base model for user-scoped resources."""
user_id: str = Field(..., description="User ID")
class ProjectScopedModel(UserScopedModel):
"""Base model for project-scoped resources."""
project_id: str = Field(..., description="Project ID")
pkg/hanzo-memory/build/lib/hanzo_memory/models/chat.py
-91
View File
@@ -1,91 +0,0 @@
"""Chat models."""
from typing import Any
from pydantic import BaseModel, Field
from .base import ProjectScopedModel
class ChatMessageBase(BaseModel):
"""Base chat message model."""
role: str = Field(..., description="Message role (user, assistant, system)")
content: str = Field(..., description="Message content")
metadata: dict[str, Any] = Field(
default_factory=dict, description="Additional metadata"
)
class ChatMessageCreate(ChatMessageBase):
"""Model for creating a chat message."""
userid: str = Field(..., description="User ID")
session_id: str = Field(..., description="Chat session ID")
project_id: str | None = Field(None, description="Project ID")
class ChatMessage(ChatMessageBase, ProjectScopedModel):
"""Complete chat message model."""
chat_id: str = Field(..., description="Chat message ID")
session_id: str = Field(..., description="Chat session ID")
embedding: list[float] | None = Field(None, description="Message embedding")
created_at: str = Field(..., description="Creation timestamp")
model_config = {"from_attributes": True}
class ChatSessionCreate(BaseModel):
"""Model for creating a chat session."""
userid: str = Field(..., description="User ID")
session_id: str | None = Field(None, description="Custom session ID")
project_id: str | None = Field(None, description="Project ID")
title: str | None = Field(None, description="Session title")
metadata: dict[str, Any] = Field(
default_factory=dict, description="Session metadata"
)
class ChatSession(BaseModel):
"""Chat session model."""
session_id: str = Field(..., description="Session ID")
user_id: str = Field(..., description="User ID")
project_id: str = Field(..., description="Project ID")
title: str | None = Field(None, description="Session title")
metadata: dict[str, Any] = Field(
default_factory=dict, description="Session metadata"
)
created_at: str = Field(..., description="Creation timestamp")
updated_at: str = Field(..., description="Last update timestamp")
message_count: int = Field(0, description="Number of messages in session")
class ChatHistoryRequest(BaseModel):
"""Request model for chat history."""
user_id: str = Field(..., description="User ID")
session_id: str = Field(..., description="Session ID")
limit: int = Field(100, ge=1, le=1000, description="Max messages to return")
offset: int = Field(0, ge=0, description="Offset for pagination")
class ChatSearchRequest(BaseModel):
"""Request model for chat search."""
user_id: str = Field(..., description="User ID")
query: str = Field(..., description="Search query")
project_id: str | None = Field(None, description="Filter by project ID")
session_id: str | None = Field(None, description="Filter by session ID")
limit: int = Field(10, ge=1, le=100, description="Max results to return")
class ChatSessionList(BaseModel):
"""Chat session list response."""
sessions: list[ChatSession] = Field(..., description="List of chat sessions")
total: int = Field(..., description="Total number of sessions")
page: int = Field(1, description="Current page")
per_page: int = Field(50, description="Items per page")
pkg/hanzo-memory/build/lib/hanzo_memory/models/knowledge.py
-142
View File
@@ -1,142 +0,0 @@
"""Knowledge base and fact models."""
from typing import Any
from pydantic import BaseModel, Field
from .base import ProjectScopedModel, TimestampedModel
class KnowledgeBaseBase(BaseModel):
"""Base knowledge base model."""
name: str = Field(..., description="Knowledge base name")
description: str = Field("", description="Knowledge base description")
metadata: dict[str, Any] = Field(
default_factory=dict, description="Additional metadata"
)
class KnowledgeBaseCreate(KnowledgeBaseBase):
"""Model for creating a knowledge base."""
kb_id: str | None = Field(None, description="Custom knowledge base ID")
class KnowledgeBaseUpdate(BaseModel):
"""Model for updating a knowledge base."""
name: str | None = None
description: str | None = None
metadata: dict[str, Any] | None = None
class KnowledgeBase(KnowledgeBaseBase, ProjectScopedModel, TimestampedModel):
"""Complete knowledge base model."""
kb_id: str = Field(..., description="Knowledge base ID")
fact_count: int = Field(0, description="Number of facts in knowledge base")
model_config = {"from_attributes": True}
class FactBase(BaseModel):
"""Base fact model."""
content: str = Field(..., description="Fact content")
metadata: dict[str, Any] = Field(
default_factory=dict, description="Additional metadata"
)
parent_id: str | None = Field(None, description="Parent fact ID")
class FactCreate(FactBase):
"""Model for creating a fact."""
fact_id: str | None = Field(None, description="Custom fact ID")
class FactUpdate(BaseModel):
"""Model for updating a fact."""
content: str | None = None
metadata: dict[str, Any] | None = None
parent_id: str | None = None
class Fact(FactBase, TimestampedModel):
"""Complete fact model."""
fact_id: str = Field(..., description="Fact ID")
kb_id: str = Field(..., description="Knowledge base ID")
embedding: list[float] | None = Field(None, description="Embedding vector")
model_config = {"from_attributes": True}
class FactWithScore(Fact):
"""Fact with similarity score."""
similarity_score: float = Field(..., description="Similarity score")
class FactRelation(BaseModel):
"""Fact relation model."""
parent_fact_id: str = Field(..., description="Parent fact ID")
child_fact_id: str = Field(..., description="Child fact ID")
relation_type: str = Field("child", description="Relation type")
# Knowledge API Request Models
class CreateKnowledgeBaseRequest(BaseModel):
"""Request model for creating a knowledge base."""
userid: str = Field(..., description="User ID")
name: str = Field(..., description="Knowledge base name")
kb_id: str | None = Field(None, description="Custom KB ID")
class ListKnowledgeBasesRequest(BaseModel):
"""Request model for listing knowledge bases."""
userid: str = Field(..., description="User ID")
class AddKnowledgeRequest(BaseModel):
"""Request model for adding facts to a knowledge base."""
userid: str = Field(..., description="User ID")
kb_id: str = Field(..., description="Knowledge base ID")
facts: list[dict[str, Any]] = Field(..., description="Facts to add")
class GetKnowledgeRequest(BaseModel):
"""Request model for retrieving facts."""
userid: str = Field(..., description="User ID")
kb_id: str = Field(..., description="Knowledge base ID")
fact_id: str | None = Field(None, description="Specific fact ID")
subtree: bool = Field(False, description="Include subtree")
query: str | None = Field(None, description="Search query")
limit: int = Field(50, ge=1, le=1000, description="Max facts to return")
class DeleteKnowledgeRequest(BaseModel):
"""Request model for deleting facts."""
userid: str = Field(..., description="User ID")
kb_id: str = Field(..., description="Knowledge base ID")
fact_id: str = Field(..., description="Fact ID to delete")
cascade: bool = Field(False, description="Delete descendants")
class IngestKnowledgeRequest(BaseModel):
"""Request model for ingesting knowledge from GCS."""
userid: str = Field(..., description="User ID")
kb_id: str = Field(..., description="Knowledge base ID")
details: dict[str, str] = Field(..., description="Must include bucketUri")
projecttags: list[str] = Field(default_factory=list, description="Tags for facts")
pkg/hanzo-memory/build/lib/hanzo_memory/models/memory.py
-141
View File
@@ -1,141 +0,0 @@
"""Memory models."""
from typing import Any
from pydantic import BaseModel, Field
from .base import ProjectScopedModel, TimestampedModel
class MemoryBase(BaseModel):
"""Base memory model."""
content: str = Field(..., description="Memory content")
metadata: dict[str, Any] = Field(
default_factory=dict, description="Additional metadata"
)
importance: float = Field(1.0, ge=0.0, le=10.0, description="Importance score")
class MemoryCreate(MemoryBase):
"""Model for creating a memory."""
additional_context: str | None = Field(
None, description="Additional context for the memory"
)
strip_pii: bool = Field(False, description="Strip PII from content")
memory_type: str | None = Field(None, description="Type of memory")
context: dict[str, Any] | None = Field(None, description="Context information")
source: str | None = Field(None, description="Source of the memory")
embedding: list[float] | None = Field(None, description="Embedding vector")
class MemoryUpdate(BaseModel):
"""Model for updating a memory."""
content: str | None = None
metadata: dict[str, Any] | None = None
importance: float | None = None
class Memory(MemoryBase, ProjectScopedModel, TimestampedModel):
"""Complete memory model."""
memory_id: str | None = Field(None, description="Memory ID")
id: str | None = Field(None, description="Memory ID alias")
embedding: list[float] | None = Field(None, description="Embedding vector")
memory_type: str | None = Field(None, description="Type of memory")
context: dict[str, Any] | None = Field(None, description="Context information")
source: str | None = Field(None, description="Source of the memory")
timestamp: str | None = Field(None, description="Timestamp")
model_config = {"from_attributes": True}
class MemoryWithScore(Memory):
"""Memory with similarity score."""
similarity_score: float = Field(..., description="Similarity score")
class MemoryResponse(BaseModel):
"""Memory response model with search results."""
memory: Memory = Field(..., description="Memory object")
similarity: float | None = Field(None, description="Similarity score")
relevance_score: float | None = Field(None, description="Relevance score")
# Legacy fields for compatibility
user_id: str | None = Field(None, description="User ID")
relevant_memories: list[str | dict[str, str]] | None = Field(
None, description="Relevant memories"
)
memory_stored: bool | None = Field(None, description="Whether the memory was stored")
usage_info: dict[str, int] | None = Field(None, description="Usage information")
class MemoryListResponse(BaseModel):
"""Memory list response."""
user_id: str = Field(..., description="User ID")
memories: list[Memory] = Field(..., description="List of memories")
pagination: dict[str, Any] = Field(..., description="Pagination info")
usage_info: dict[str, int] = Field(..., description="Usage information")
class RememberRequest(BaseModel):
"""Request model for /v1/remember endpoint."""
apikey: str | None = Field(None, description="API key")
userid: str = Field(..., description="User ID")
messagecontent: str = Field(..., description="Message content")
additionalcontext: str | None = Field(None, description="Additional context")
strippii: bool = Field(False, description="Strip PII")
filterresults: bool = Field(False, description="Filter results with LLM")
includememoryid: bool = Field(False, description="Include memory IDs in response")
class AddMemoriesRequest(BaseModel):
"""Request model for /v1/memories/add endpoint."""
apikey: str | None = Field(None, description="API key")
userid: str = Field(..., description="User ID")
memoriestoadd: str | list[str] = Field(..., description="Memories to add")
class GetMemoriesRequest(BaseModel):
"""Request model for /v1/memories/get endpoint."""
apikey: str | None = Field(None, description="API key")
userid: str = Field(..., description="User ID")
memoryid: str | None = Field(None, description="Specific memory ID")
limit: int = Field(50, ge=1, le=1000, description="Max memories to return")
startafter: str | None = Field(None, description="Memory ID to start after")
class DeleteMemoryRequest(BaseModel):
"""Request model for /v1/memories/delete endpoint."""
apikey: str | None = Field(None, description="API key")
userid: str = Field(..., description="User ID")
memoryid: str = Field(..., description="Memory ID to delete")
class UpdateMemoryRequest(BaseModel):
"""Request model for /v1/memories/update endpoint."""
apikey: str | None = Field(None, description="API key")
userid: str = Field(..., description="User ID")
memoryid: str = Field(..., description="Memory ID to update")
projectid: str = Field(..., description="Project ID")
content: str | None = Field(None, description="New content")
importance: float | None = Field(None, ge=0.0, le=10.0, description="New importance score")
metadata: dict[str, Any] | None = Field(None, description="New metadata")
class DeleteUserRequest(BaseModel):
"""Request model for /v1/user/delete endpoint."""
apikey: str | None = Field(None, description="API key")
userid: str = Field(..., description="User ID")
confirmdelete: bool = Field(..., description="Confirm deletion")
pkg/hanzo-memory/build/lib/hanzo_memory/models/project.py
-52
View File
@@ -1,52 +0,0 @@
"""Project models."""
from typing import Any
from pydantic import BaseModel, Field
from .base import TimestampedModel, UserScopedModel
class ProjectBase(BaseModel):
"""Base project model."""
name: str = Field(..., description="Project name")
description: str = Field("", description="Project description")
metadata: dict[str, Any] = Field(
default_factory=dict, description="Additional metadata"
)
class ProjectCreate(ProjectBase):
"""Model for creating a project."""
pass
class ProjectUpdate(BaseModel):
"""Model for updating a project."""
name: str | None = None
description: str | None = None
metadata: dict[str, Any] | None = None
class Project(ProjectBase, UserScopedModel, TimestampedModel):
"""Complete project model."""
project_id: str = Field(..., description="Project ID")
knowledge_base_ids: list[str] = Field(
default_factory=list, description="Associated knowledge base IDs"
)
memory_count: int = Field(0, description="Number of memories in project")
model_config = {"from_attributes": True}
class ProjectList(BaseModel):
"""Project list response."""
projects: list[Project] = Field(..., description="List of projects")
total: int = Field(..., description="Total number of projects")
page: int = Field(1, description="Current page")
per_page: int = Field(50, description="Items per page")
pkg/hanzo-memory/build/lib/hanzo_memory/server.py
-780
View File
@@ -1,780 +0,0 @@
"""FastAPI server for Hanzo Memory Service."""
import json
import uuid
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from fastapi import Depends, FastAPI, HTTPException, Query, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPAuthorizationCredentials
from structlog import get_logger
from .api.auth import get_or_verify_user_id, require_auth, security
from .config import settings
from .db import get_db_client
from .models import (
AddKnowledgeRequest,
AddMemoriesRequest,
ChatMessageCreate,
ChatSessionCreate,
CreateKnowledgeBaseRequest,
DeleteKnowledgeRequest,
DeleteMemoryRequest,
DeleteUserRequest,
GetKnowledgeRequest,
GetMemoriesRequest,
MemoryListResponse,
MemoryResponse,
RememberRequest,
UpdateMemoryRequest,
)
from .services import get_embedding_service, get_memory_service
logger = get_logger()
# Global service instances
db_client = None
embedding_service = None
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""Application lifespan manager."""
global db_client, embedding_service
# Startup
logger.info("Starting Hanzo Memory Service")
settings.ensure_paths()
# Initialize services
db_client = get_db_client()
embedding_service = get_embedding_service()
if db_client:
db_client.create_projects_table()
db_client.create_knowledge_bases_table()
yield
# Shutdown
logger.info("Shutting down Hanzo Memory Service")
if db_client:
db_client.close()
# Create FastAPI app
app = FastAPI(
title="Hanzo Memory Service",
description="AI memory and knowledge management service",
version="0.1.0",
lifespan=lifespan,
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Configure based on your needs
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/health")
async def health_check() -> dict[str, str]:
"""Health check endpoint."""
return {
"status": "healthy",
"service": "hanzo-memory",
"version": "0.1.0",
}
@app.post("/v1/remember")
async def remember(
request: RememberRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> MemoryResponse:
"""
Retrieve relevant memories and store new memory.
This endpoint:
1. Searches for relevant memories based on the message content
2. Optionally filters results using LLM
3. Stores the incoming message as a new memory
4. Returns relevant memories
"""
# Check auth
request.apikey or require_auth(req, credentials)
memory_service = get_memory_service()
# Get or create default project for user
project_id = f"project_{request.userid}_default"
# Search for relevant memories
memories = memory_service.search_memories(
user_id=request.userid,
query=request.messagecontent,
project_id=project_id,
limit=10,
filter_with_llm=request.filterresults,
additional_context=request.additionalcontext,
)
# Store the new memory
memory_service.create_memory(
user_id=request.userid,
project_id=project_id,
content=request.messagecontent,
metadata={
"additional_context": request.additionalcontext,
},
strip_pii=request.strippii,
)
# Format response
relevant_memories: list[str | dict[str, str]]
if request.includememoryid:
relevant_memories = [
{"content": m.content, "memoryId": m.memory_id} for m in memories
]
else:
relevant_memories = [m.content for m in memories]
# Get usage info
# TODO: Implement proper usage tracking
usage_info = {
"current": len(memories),
"limit": settings.max_memories_per_user,
}
return MemoryResponse(
user_id=request.userid,
relevant_memories=relevant_memories,
memory_stored=True,
usage_info=usage_info,
)
@app.post("/v1/memories/add")
async def add_memories(
request: AddMemoriesRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Add explicit memories without importance analysis."""
# Check auth
request.apikey or require_auth(req, credentials)
memory_service = get_memory_service()
# Get or create default project
project_id = f"project_{request.userid}_default"
# Normalize memories to list
memories_to_add = (
[request.memoriestoadd]
if isinstance(request.memoriestoadd, str)
else request.memoriestoadd
)
# Add memories
memory_ids = []
for content in memories_to_add:
memory = memory_service.create_memory(
user_id=request.userid,
project_id=project_id,
content=content,
importance=5.0, # Default importance for explicit adds
)
memory_ids.append(memory.memory_id)
return {
"userid": request.userid,
"added_count": len(memory_ids),
"memory_ids": memory_ids,
"usage_info": {
"current": len(memory_ids),
"limit": settings.max_memories_per_user,
},
}
@app.post("/v1/memories/get")
async def get_memories(
request: GetMemoriesRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> MemoryListResponse:
"""Retrieve stored memories."""
# Check auth
request.apikey or require_auth(req, credentials)
memory_service = get_memory_service()
# If specific memory ID requested
if request.memoryid:
memory = memory_service.get_memory(request.userid, request.memoryid)
if not memory:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Memory not found",
)
memories = [memory]
else:
# Get paginated list
# TODO: Implement proper pagination
memories = []
return MemoryListResponse(
user_id=request.userid,
memories=memories,
pagination={
"has_more": False,
"last_id": memories[-1].memory_id if memories else None,
},
usage_info={
"current": len(memories),
"limit": settings.max_memories_per_user,
},
)
@app.post("/v1/memories/update")
async def update_memory_endpoint(
request: UpdateMemoryRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Update a specific memory."""
# Check auth
request.apikey or require_auth(req, credentials)
memory_service = get_memory_service()
# Prepare update parameters
update_params = {}
if request.content is not None:
update_params["content"] = request.content
if request.importance is not None:
update_params["importance"] = request.importance
if request.metadata is not None:
update_params["metadata"] = request.metadata
# Perform the update
updated_memory = memory_service.update_memory(
user_id=request.userid,
memory_id=request.memoryid,
project_id=request.projectid,
**update_params
)
if not updated_memory:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Memory not found",
)
return {
"message": "Memory updated successfully",
"memory_id": request.memoryid,
"userid": request.userid,
"updated_fields": list(update_params.keys()),
}
@app.post("/v1/memories/delete")
async def delete_memory(
request: DeleteMemoryRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Delete a specific memory."""
# Check auth
request.apikey or require_auth(req, credentials)
memory_service = get_memory_service()
success = memory_service.delete_memory(request.userid, request.memoryid)
if not success:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Memory not found",
)
return {
"message": "Memory deleted successfully",
"memory_id": request.memoryid,
"userid": request.userid,
}
@app.post("/v1/user/delete")
async def delete_user(
request: DeleteUserRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Delete all memories for a user."""
# Check auth
request.apikey or require_auth(req, credentials)
if not request.confirmdelete:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="confirmdelete must be true",
)
memory_service = get_memory_service()
deleted_count = memory_service.delete_user_memories(request.userid)
return {
"message": "All user memories deleted",
"userid": request.userid,
"deleted_count": deleted_count,
}
# Knowledge Base Management Endpoints
@app.post("/v1/kb/create")
async def create_knowledge_base(
request: CreateKnowledgeBaseRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Create a new knowledge base."""
try:
# Get or verify user ID
user_id = await get_or_verify_user_id(request.userid, credentials, req)
# Get default project if not specified
project_id = getattr(request, "project_id", None)
if not project_id:
# Create or get default project for user
project_id = f"default-{user_id}"
try:
if db_client:
db_client.create_project(
project_id=project_id,
user_id=user_id,
name="Default Project",
description="Automatically created default project",
)
except Exception:
pass # Project may already exist
# Create knowledge base
kb_id = request.kb_id or str(uuid.uuid4())
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
db_client.create_knowledge_base(
kb_id=kb_id,
user_id=user_id,
project_id=project_id,
name=request.name,
description=getattr(request, "description", ""),
)
return {
"kb_id": kb_id,
"message": f"Knowledge base '{request.name}' created successfully",
}
except Exception as e:
logger.error(f"Error creating knowledge base: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
@app.get("/v1/kb/list")
async def list_knowledge_bases(
req: Request,
userid: str = Query(...),
project_id: str | None = Query(None),
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""List knowledge bases for a user."""
try:
# Get or verify user ID
user_id = await get_or_verify_user_id(userid, credentials, req)
# Query knowledge bases from the database
# For now, return a simple response
# TODO: Implement actual listing from InfinityDB
return {
"userid": user_id,
"knowledge_bases": [],
"total": 0,
}
except Exception as e:
logger.error(f"Error listing knowledge bases: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
@app.post("/v1/kb/facts/add")
async def add_facts(
request: AddKnowledgeRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Add facts to a knowledge base."""
try:
# Get or verify user ID
await get_or_verify_user_id(request.userid, credentials, req)
# Process each fact
added_facts = []
for fact_data in request.facts:
# Generate embedding for fact content
content = fact_data.get("content", "")
if not embedding_service:
raise HTTPException(
status_code=503, detail="Embedding service not initialized"
)
embedding = embedding_service.embed_text(content)[0]
# Add fact to database
fact_id = fact_data.get("fact_id") or str(uuid.uuid4())
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
db_client.add_fact(
fact_id=fact_id,
kb_id=request.kb_id,
content=content,
embedding=embedding,
parent_id=fact_data.get("parent_id"),
metadata=fact_data.get("metadata", {}),
)
added_facts.append(
{
"fact_id": fact_id,
"content": content,
}
)
return {
"kb_id": request.kb_id,
"facts_added": len(added_facts),
"facts": added_facts,
}
except Exception as e:
logger.error(f"Error adding facts: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
@app.post("/v1/kb/facts/get")
async def get_facts(
request: GetKnowledgeRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Get facts from a knowledge base."""
try:
# Get or verify user ID
await get_or_verify_user_id(request.userid, credentials, req)
# Search facts if query provided
if request.query:
# Generate query embedding
if not embedding_service:
raise HTTPException(
status_code=503, detail="Embedding service not initialized"
)
query_embedding = embedding_service.embed_text(request.query)[0]
# Search facts
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
results_df = db_client.search_facts(
kb_id=request.kb_id,
query_embedding=query_embedding,
limit=request.limit,
parent_id=request.fact_id if request.subtree else None,
)
# Convert results
facts = []
if not results_df.is_empty():
for row in results_df.to_dicts():
facts.append(
{
"fact_id": row["fact_id"],
"content": row["content"],
"parent_id": row.get("parent_id"),
"metadata": json.loads(row.get("metadata", "{}")),
"similarity_score": row.get("_similarity", 0.0),
}
)
return {
"kb_id": request.kb_id,
"facts": facts,
"total": len(facts),
}
else:
# TODO: Implement listing all facts or specific fact retrieval
return {
"kb_id": request.kb_id,
"facts": [],
"total": 0,
}
except Exception as e:
logger.error(f"Error getting facts: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
@app.post("/v1/kb/facts/delete")
async def delete_fact(
request: DeleteKnowledgeRequest,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Delete a fact from a knowledge base."""
try:
# Get or verify user ID
await get_or_verify_user_id(request.userid, credentials, req)
# TODO: Implement fact deletion in InfinityDB
# For now, return success
return {
"kb_id": request.kb_id,
"fact_id": request.fact_id,
"deleted": True,
"cascade": request.cascade,
}
except Exception as e:
logger.error(f"Error deleting fact: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
# Chat Management Endpoints
@app.post("/v1/chat/sessions/create")
async def create_chat_session(
request: ChatSessionCreate,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Create a new chat session."""
try:
# Get or verify user ID
user_id = await get_or_verify_user_id(request.userid, credentials, req)
# Create session ID
session_id = request.session_id or str(uuid.uuid4())
# Get or create default project
project_id = request.project_id or f"default-{user_id}"
# Ensure user's chat table exists
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
db_client.create_chats_table(user_id)
return {
"session_id": session_id,
"userid": user_id,
"project_id": project_id,
"created": True,
}
except Exception as e:
logger.error(f"Error creating chat session: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
@app.post("/v1/chat/messages/add")
async def add_chat_message(
request: ChatMessageCreate,
req: Request,
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Add a message to a chat session."""
try:
# Get or verify user ID
user_id = await get_or_verify_user_id(request.userid, credentials, req)
# Generate embedding for message content
if not embedding_service:
raise HTTPException(
status_code=503, detail="Embedding service not initialized"
)
embedding = embedding_service.embed_text(request.content)[0]
# Check for duplicate messages
# Search for similar messages in the same session
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
search_results = db_client.search_chats(
user_id=user_id,
query_embedding=embedding,
session_id=request.session_id,
limit=5,
)
# Check if this is a duplicate
is_duplicate = False
if not search_results.is_empty():
for row in search_results.to_dicts():
if (
row["content"] == request.content
and row["role"] == request.role
and row.get("_similarity", 0) > 0.99
):
is_duplicate = True
chat_id = row["chat_id"]
break
if not is_duplicate:
# Add new message
chat_id = str(uuid.uuid4())
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
db_client.add_chat_message(
chat_id=chat_id,
user_id=user_id,
project_id=request.project_id or f"default-{user_id}",
session_id=request.session_id,
role=request.role,
content=request.content,
embedding=embedding,
metadata=request.metadata or {},
)
return {
"chat_id": chat_id,
"session_id": request.session_id,
"duplicate": is_duplicate,
}
except Exception as e:
logger.error(f"Error adding chat message: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
@app.get("/v1/chat/sessions/{session_id}/messages")
async def get_chat_messages(
session_id: str,
req: Request,
userid: str = Query(...),
limit: int = Query(100, ge=1, le=1000),
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Get messages for a chat session."""
try:
# Get or verify user ID
user_id = await get_or_verify_user_id(userid, credentials, req)
# Get chat history
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
history_df = db_client.get_chat_history(
user_id=user_id,
session_id=session_id,
limit=limit,
)
# Convert to messages
messages = []
if not history_df.is_empty():
# Sort by created_at timestamp
sorted_df = history_df.sort("created_at")
for row in sorted_df.to_dicts():
messages.append(
{
"chat_id": row["chat_id"],
"role": row["role"],
"content": row["content"],
"metadata": json.loads(row.get("metadata", "{}")),
"created_at": row["created_at"],
}
)
return {
"session_id": session_id,
"messages": messages,
"total": len(messages),
}
except Exception as e:
logger.error(f"Error getting chat messages: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
@app.post("/v1/chat/search")
async def search_chat_messages(
req: Request,
query: str = Query(...),
userid: str = Query(...),
project_id: str | None = Query(None),
session_id: str | None = Query(None),
limit: int = Query(10, ge=1, le=100),
credentials: HTTPAuthorizationCredentials | None = Depends(security),
) -> dict:
"""Search across chat messages."""
try:
# Get or verify user ID
user_id = await get_or_verify_user_id(userid, credentials, req)
# Generate query embedding
if not embedding_service:
raise HTTPException(
status_code=503, detail="Embedding service not initialized"
)
query_embedding = embedding_service.embed_text(query)[0]
# Search chats
if not db_client:
raise HTTPException(status_code=503, detail="Database not initialized")
results_df = db_client.search_chats(
user_id=user_id,
query_embedding=query_embedding,
project_id=project_id,
session_id=session_id,
limit=limit,
)
# Convert results
messages = []
if not results_df.is_empty():
for row in results_df.to_dicts():
messages.append(
{
"chat_id": row["chat_id"],
"session_id": row["session_id"],
"role": row["role"],
"content": row["content"],
"similarity_score": row.get("_similarity", 0.0),
"created_at": row["created_at"],
}
)
return {
"query": query,
"messages": messages,
"total": len(messages),
}
except Exception as e:
logger.error(f"Error searching chat messages: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
def run() -> None:
"""Run the server."""
import uvicorn
uvicorn.run(
"hanzo_memory.server:app",
host=settings.host,
port=settings.port,
reload=True,
)
if __name__ == "__main__":
run()
pkg/hanzo-memory/build/lib/hanzo_memory/services/__init__.py
-14
View File
@@ -1,14 +0,0 @@
"""Services package."""
from .embeddings import EmbeddingService, get_embedding_service
from .llm import LLMService, get_llm_service
from .memory import MemoryService, get_memory_service
__all__ = [
"EmbeddingService",
"get_embedding_service",
"LLMService",
"get_llm_service",
"MemoryService",
"get_memory_service",
]
pkg/hanzo-memory/build/lib/hanzo_memory/services/embeddings.py
-241
View File
@@ -1,241 +0,0 @@
"""Embedding service using FastEmbed or LanceDB."""
import numpy as np
from structlog import get_logger
from ..config import settings
logger = get_logger()
# Import backend-specific implementations
try:
from fastembed import TextEmbedding
FASTEMBED_AVAILABLE = True
except ImportError:
FASTEMBED_AVAILABLE = False
logger.warning("FastEmbed not available")
try:
import importlib.util
if importlib.util.find_spec("hanzo_memory.services.embeddings_lancedb") is not None:
LANCEDB_AVAILABLE = True
else:
LANCEDB_AVAILABLE = False
logger.warning("LanceDB embeddings not available")
except ImportError:
LANCEDB_AVAILABLE = False
logger.warning("LanceDB embeddings not available")
class EmbeddingService:
"""Service for generating embeddings using FastEmbed."""
def __init__(self, model_name: str | None = None):
"""Initialize the embedding service."""
self.model_name = model_name or settings.embedding_model
self._model = None
logger.info(f"Initializing embedding service with model: {self.model_name}")
@property
def model(self):
"""Lazy load the embedding model."""
if self._model is None:
if not FASTEMBED_AVAILABLE:
raise RuntimeError("FastEmbed is not available")
logger.info(f"Loading embedding model: {self.model_name}")
self._model = TextEmbedding(model_name=self.model_name)
logger.info(f"Model {self.model_name} loaded successfully")
return self._model
def embed_text(self, text: str | list[str]) -> list[list[float]]:
"""
Generate embeddings for text.
Args:
text: Single text string or list of text strings
Returns:
List of embedding vectors
"""
if isinstance(text, str):
text = [text]
# Generate embeddings
embeddings = list(self.model.embed(text))
# Convert numpy arrays to lists
return [emb.tolist() for emb in embeddings]
def embed_single(self, text: str) -> list[float]:
"""
Generate embedding for a single text.
Args:
text: Text string
Returns:
Embedding vector
"""
embeddings = self.embed_text(text)
return embeddings[0]
def embed_batch(
self,
texts: list[str],
batch_size: int = 32,
show_progress: bool = False,
) -> list[list[float]]:
"""
Generate embeddings for a batch of texts.
Args:
texts: List of text strings
batch_size: Batch size for processing
show_progress: Whether to show progress bar
Returns:
List of embedding vectors
"""
if not texts:
return []
all_embeddings = []
# Process in batches
for i in range(0, len(texts), batch_size):
batch = texts[i : i + batch_size]
embeddings = self.embed_text(batch)
all_embeddings.extend(embeddings)
if show_progress and i > 0:
logger.info(f"Processed {i + len(batch)}/{len(texts)} texts")
return all_embeddings
def compute_similarity(
self,
query_embedding: list[float],
embeddings: list[list[float]],
metric: str = "cosine",
) -> list[float]:
"""
Compute similarity between query embedding and a list of embeddings.
Args:
query_embedding: Query embedding vector
embeddings: List of embedding vectors to compare against
metric: Similarity metric ("cosine", "dot", "euclidean")
Returns:
List of similarity scores
"""
if not embeddings:
return []
query = np.array(query_embedding)
vectors = np.array(embeddings)
if metric == "cosine":
# Normalize vectors for cosine similarity
query_norm = query / np.linalg.norm(query)
vectors_norm = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
similarities = np.dot(vectors_norm, query_norm)
elif metric == "dot":
similarities = np.dot(vectors, query)
elif metric == "euclidean":
# Return negative distance so higher is better
similarities = -np.linalg.norm(vectors - query, axis=1)
else:
raise ValueError(f"Unknown metric: {metric}")
return list(similarities.tolist()) # Type hint for mypy
def get_model_info(self) -> dict:
"""Get information about the current embedding model."""
return {
"model_name": self.model_name,
"dimensions": settings.embedding_dimensions,
"loaded": self._model is not None,
}
# Global embedding service instance
_embedding_service: EmbeddingService | None = None
def get_embedding_service() -> EmbeddingService:
"""Get or create the global embedding service."""
global _embedding_service
if _embedding_service is None:
# Use LanceDB embeddings if configured
if settings.db_backend == "lancedb" and LANCEDB_AVAILABLE:
logger.info("Using LanceDB embedding service")
from .embeddings_lancedb import get_lancedb_embedding_service
# Wrap LanceDB service in standard interface
lancedb_service = get_lancedb_embedding_service()
_embedding_service = EmbeddingService()
# Override methods with LanceDB implementation
_embedding_service.embed_text = lancedb_service.embed_text
_embedding_service.embed_single = lancedb_service.embed_single
_embedding_service.embed_batch = lancedb_service.embed_batch
_embedding_service.compute_similarity = lancedb_service.compute_similarity
_embedding_service.get_model_info = lancedb_service.get_model_info
elif FASTEMBED_AVAILABLE:
logger.info("Using FastEmbed embedding service")
_embedding_service = EmbeddingService()
else:
# Create a minimal embedding service that generates dummy embeddings
# This allows the system to work without heavy embedding dependencies
logger.warning("Using minimal embedding service (dummy embeddings)")
_embedding_service = MinimalEmbeddingService()
return _embedding_service
class MinimalEmbeddingService:
"""Minimal embedding service that generates dummy embeddings for basic functionality."""
def __init__(self):
self.model_name = "minimal"
def embed_text(self, text: str | list[str]) -> list[list[float]]:
"""Generate dummy embeddings for text."""
if isinstance(text, str):
text = [text]
# Generate simple dummy embeddings based on text length and content
embeddings = []
for t in text:
# Create a simple hash-based embedding
hash_val = hash(t) % 1000000
embedding = [float((hash_val >> i) & 0xFF) / 255.0 for i in range(0, 384, 8)]
# Pad or truncate to ensure consistent size
embedding = (embedding * (384 // len(embedding) + 1))[:384]
embeddings.append(embedding)
return embeddings
def embed_single(self, text: str) -> list[float]:
"""Generate embedding for a single text."""
embeddings = self.embed_text(text)
return embeddings[0]
def embed_batch(self, texts: list[str], batch_size: int = 32, show_progress: bool = False) -> list[list[float]]:
"""Generate embeddings for a batch of texts."""
return self.embed_text(texts)
def compute_similarity(self, query_embedding: list[float], embeddings: list[list[float]], metric: str = "cosine") -> list[float]:
"""Compute similarity between query embedding and a list of embeddings."""
# For minimal service, return dummy similarity scores
return [0.5] * len(embeddings) # Return neutral similarity
def get_model_info(self) -> dict:
"""Get information about the current embedding model."""
return {
"model_name": "minimal",
"dimensions": 384,
"loaded": True,
}
pkg/hanzo-memory/build/lib/hanzo_memory/services/embeddings_lancedb.py
-208
View File
@@ -1,208 +0,0 @@
"""Embedding service using LanceDB's built-in embedding functions."""
import numpy as np
from lancedb.embeddings import get_registry
from structlog import get_logger
from ..config import settings
logger = get_logger()
class LanceDBEmbeddingService:
"""Service for generating embeddings using LanceDB's embedding functions."""
def __init__(self, model_name: str | None = None):
"""Initialize the embedding service."""
self.model_name = model_name or settings.embedding_model
self._embedding_func = None
logger.info(
f"Initializing LanceDB embedding service with model: {self.model_name}"
)
@property
def embedding_func(self):
"""Lazy load the embedding function."""
if self._embedding_func is None:
logger.info(f"Loading embedding function: {self.model_name}")
# LanceDB supports multiple embedding providers
registry = get_registry()
# Check available embedding functions
available_funcs = list(registry.list_embedding_functions())
logger.info(f"Available embedding functions: {available_funcs}")
# Try to use the appropriate embedding function based on model name
if "fastembed" in available_funcs:
# Use FastEmbed if available
try:
self._embedding_func = registry.get("fastembed").create(
name=self.model_name
)
logger.info("Using FastEmbed via LanceDB")
except Exception as e:
logger.warning(f"Failed to use FastEmbed: {e}")
self._embedding_func = None
if self._embedding_func is None:
# Fall back to other providers
if self.model_name.startswith("BAAI/") or "bge" in self.model_name:
# Use sentence-transformers for BAAI/BGE models
self._embedding_func = registry.get("sentence-transformers").create(
name=self.model_name
)
logger.info("Using sentence-transformers via LanceDB")
elif (
self.model_name.startswith("text-embedding")
and "openai" in available_funcs
):
# Use OpenAI for text-embedding models
self._embedding_func = registry.get("openai").create(
name=self.model_name
)
logger.info("Using OpenAI embeddings via LanceDB")
else:
# Default to sentence-transformers
self._embedding_func = registry.get("sentence-transformers").create(
name=self.model_name
)
logger.info("Using sentence-transformers (default) via LanceDB")
logger.info(f"Embedding function {self.model_name} loaded successfully")
return self._embedding_func
def embed_text(self, text: str | list[str]) -> list[list[float]]:
"""
Generate embeddings for text.
Args:
text: Single text string or list of text strings
Returns:
List of embedding vectors
"""
if isinstance(text, str):
text = [text]
# Generate embeddings using LanceDB's embedding function
embeddings = self.embedding_func.compute_source_embeddings(text)
# Convert to list format
if isinstance(embeddings, np.ndarray):
return embeddings.tolist()
else:
return [
emb.tolist() if isinstance(emb, np.ndarray) else emb
for emb in embeddings
]
def embed_single(self, text: str) -> list[float]:
"""
Generate embedding for a single text.
Args:
text: Text string
Returns:
Embedding vector
"""
embeddings = self.embed_text(text)
return embeddings[0]
def embed_batch(
self,
texts: list[str],
batch_size: int = 32,
show_progress: bool = False,
) -> list[list[float]]:
"""
Generate embeddings for a batch of texts.
Args:
texts: List of text strings
batch_size: Batch size for processing
show_progress: Whether to show progress bar
Returns:
List of embedding vectors
"""
if not texts:
return []
all_embeddings = []
# Process in batches
for i in range(0, len(texts), batch_size):
batch = texts[i : i + batch_size]
embeddings = self.embed_text(batch)
all_embeddings.extend(embeddings)
if show_progress and i > 0:
logger.info(f"Processed {i + len(batch)}/{len(texts)} texts")
return all_embeddings
def compute_similarity(
self,
query_embedding: list[float],
embeddings: list[list[float]],
metric: str = "cosine",
) -> list[float]:
"""
Compute similarity between query embedding and a list of embeddings.
Args:
query_embedding: Query embedding vector
embeddings: List of embedding vectors to compare against
metric: Similarity metric ("cosine", "dot", "euclidean")
Returns:
List of similarity scores
"""
if not embeddings:
return []
query = np.array(query_embedding)
vectors = np.array(embeddings)
if metric == "cosine":
# Normalize vectors for cosine similarity
query_norm = query / np.linalg.norm(query)
vectors_norm = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
similarities = np.dot(vectors_norm, query_norm)
elif metric == "dot":
similarities = np.dot(vectors, query)
elif metric == "euclidean":
# Return negative distance so higher is better
similarities = -np.linalg.norm(vectors - query, axis=1)
else:
raise ValueError(f"Unknown metric: {metric}")
return list(similarities.tolist()) # Type hint for mypy
def get_model_info(self) -> dict:
"""Get information about the current embedding model."""
return {
"model_name": self.model_name,
"dimensions": (
self.embedding_func.ndims()
if self._embedding_func
else settings.embedding_dimensions
),
"loaded": self._embedding_func is not None,
"backend": "lancedb",
}
# Global embedding service instance
_embedding_service: LanceDBEmbeddingService | None = None
def get_lancedb_embedding_service() -> LanceDBEmbeddingService:
"""Get or create the global LanceDB embedding service."""
global _embedding_service
if _embedding_service is None:
_embedding_service = LanceDBEmbeddingService()
return _embedding_service
pkg/hanzo-memory/build/lib/hanzo_memory/services/llm.py
-242
View File
@@ -1,242 +0,0 @@
"""LLM service for AI operations using LiteLLM."""
import json
from typing import Any
from structlog import get_logger
from ..config import settings
logger = get_logger()
# Try to import LiteLLM, but make it optional
try:
import litellm
LITELLM_AVAILABLE = True
# Configure LiteLLM
litellm.drop_params = True # Drop unsupported params instead of failing
litellm.set_verbose = False # Disable verbose logging
except ImportError:
LITELLM_AVAILABLE = False
litellm = None
logger.warning("LiteLLM not available, LLM features will be limited")
class LLMService:
"""Service for LLM operations using LiteLLM."""
def __init__(self) -> None:
"""Initialize LLM service."""
self.default_model = settings.llm_model
self.api_base = settings.llm_api_base
self.temperature = settings.llm_temperature
self.max_tokens = settings.llm_max_tokens
# Set up API keys if LiteLLM is available
if LITELLM_AVAILABLE:
# Set up API keys
if settings.llm_api_key:
litellm.api_key = settings.llm_api_key
if settings.openai_api_key:
litellm.openai_key = settings.openai_api_key
if settings.anthropic_api_key:
litellm.anthropic_key = settings.anthropic_api_key
def complete(
self,
prompt: str,
model: str | None = None,
max_tokens: int | None = None,
temperature: float | None = None,
response_format: str | None = None,
) -> str:
"""
Complete a prompt using LiteLLM.
Args:
prompt: The prompt to complete
model: Model to use (defaults to config)
max_tokens: Maximum tokens to generate
temperature: Temperature for generation
response_format: Optional response format ("json" for JSON mode)
Returns:
Generated text
"""
if not LITELLM_AVAILABLE:
logger.warning("LiteLLM not available, returning empty response")
return ""
model = model or self.default_model
max_tokens = max_tokens or self.max_tokens
temperature = temperature or self.temperature
try:
kwargs = {
"model": model,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": max_tokens,
"temperature": temperature,
}
# Add API base if configured (for local models)
if self.api_base:
kwargs["api_base"] = self.api_base
# Add response format if specified
if response_format == "json":
kwargs["response_format"] = {"type": "json_object"}
response = litellm.completion(**kwargs)
return response.choices[0].message.content or ""
except Exception as e:
logger.error(f"LLM completion error: {e}")
return ""
def chat(
self,
messages: list[dict[str, str]],
model: str | None = None,
max_tokens: int | None = None,
temperature: float | None = None,
response_format: str | None = None,
) -> str:
"""
Chat with the LLM using LiteLLM.
Args:
messages: List of message dicts with "role" and "content"
model: Model to use
max_tokens: Maximum tokens to generate
temperature: Temperature for generation
response_format: Optional response format ("json" for JSON mode)
Returns:
Generated response
"""
if not LITELLM_AVAILABLE:
logger.warning("LiteLLM not available, returning empty response")
return ""
model = model or self.default_model
max_tokens = max_tokens or self.max_tokens
temperature = temperature or self.temperature
try:
kwargs = {
"model": model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}
# Add API base if configured
if self.api_base:
kwargs["api_base"] = self.api_base
# Add response format if specified
if response_format == "json":
kwargs["response_format"] = {"type": "json_object"}
response = litellm.completion(**kwargs)
return response.choices[0].message.content or ""
except Exception as e:
logger.error(f"LLM chat error: {e}")
return ""
def summarize_for_knowledge(
self,
content: str,
context: str | None = None,
skip_summarization: bool = False,
provided_summary: str | None = None,
) -> dict[str, Any]:
"""
Summarize content and generate knowledge update instructions.
Args:
content: Content to summarize
context: Additional context
skip_summarization: Skip summarization step
provided_summary: Pre-provided summary
Returns:
JSON with summary and knowledge instructions
"""
if skip_summarization and not provided_summary:
# Just return the content as-is with basic instructions
return {
"summary": content,
"knowledge_instructions": {
"action": "add_fact",
"facts": [{"content": content}],
"reasoning": "Content added without summarization",
},
}
if provided_summary:
summary = provided_summary
else:
# Generate summary
summary_prompt = f"""Summarize the following content concisely, preserving key information:
Content: {content}"""
if context:
summary_prompt += f"\n\nContext: {context}"
summary = self.complete(summary_prompt, temperature=0.3)
# Generate knowledge instructions
if LITELLM_AVAILABLE:
instruction_prompt = f"""Based on this summary, generate instructions for updating a knowledge base.
Return a JSON object with the following structure:
{{
"action": "add_fact" or "update_fact" or "add_relation",
"facts": [
{{
"content": "fact content",
"metadata": {{"tags": [], "source": "..."}},
"parent_id": null or "parent_fact_id"
}}
],
"reasoning": "explanation of why these facts should be added"
}}
Summary: {summary}"""
if context:
instruction_prompt += f"\n\nContext: {context}"
instructions_json = self.complete(
instruction_prompt, temperature=0.3, response_format="json"
)
try:
instructions = json.loads(instructions_json)
except json.JSONDecodeError:
instructions = {
"action": "add_fact",
"facts": [{"content": summary}],
"reasoning": "Failed to parse LLM response, using summary as fact",
}
else:
# Without LLM, return basic instructions
instructions = {
"action": "add_fact",
"facts": [{"content": summary}],
"reasoning": "Content added without LLM processing",
}
return {"summary": summary, "knowledge_instructions": instructions}
# Global LLM service instance
_llm_service: LLMService | None = None
def get_llm_service() -> LLMService:
"""Get or create the global LLM service."""
global _llm_service
if _llm_service is None:
_llm_service = LLMService()
return _llm_service
pkg/hanzo-memory/build/lib/hanzo_memory/services/memory.py
-284
View File
@@ -1,284 +0,0 @@
"""Memory service for managing memories."""
import json
import uuid
from datetime import datetime
from structlog import get_logger
from ..db import get_db_client
from ..models.memory import Memory, MemoryWithScore
from .embeddings import get_embedding_service
from .llm import get_llm_service
logger = get_logger()
class MemoryService:
"""Service for managing memories."""
def __init__(self) -> None:
"""Initialize memory service."""
self.db = get_db_client()
self.embeddings = get_embedding_service()
self.llm = get_llm_service()
def create_memory(
self,
user_id: str,
project_id: str,
content: str,
metadata: dict | None = None,
importance: float = 1.0,
strip_pii: bool = False,
) -> Memory:
"""
Create a new memory.
Args:
user_id: User ID
project_id: Project ID
content: Memory content
metadata: Additional metadata
importance: Importance score
strip_pii: Whether to strip PII from content
Returns:
Created memory
"""
# Ensure user's memory table exists
self.db.create_memories_table(user_id)
# Strip PII if requested
if strip_pii:
content = self._strip_pii(content)
# Generate embedding
embedding = self.embeddings.embed_single(content)
# Create memory
memory_id = f"mem_{uuid.uuid4().hex[:12]}"
memory_data = self.db.add_memory(
memory_id=memory_id,
user_id=user_id,
project_id=project_id,
content=content,
embedding=embedding,
metadata=metadata,
importance=importance,
)
# Parse JSON fields if needed
if isinstance(memory_data.get("metadata"), str):
memory_data["metadata"] = json.loads(memory_data["metadata"])
return Memory(**memory_data)
def search_memories(
self,
user_id: str,
query: str,
project_id: str | None = None,
limit: int = 10,
filter_with_llm: bool = False,
additional_context: str | None = None,
) -> list[MemoryWithScore]:
"""
Search memories by semantic similarity.
Args:
user_id: User ID
query: Search query
project_id: Optional project filter
limit: Maximum results
filter_with_llm: Use LLM to filter results
additional_context: Additional context for filtering
Returns:
List of memories with similarity scores
"""
# Generate query embedding
query_embedding = self.embeddings.embed_single(query)
# Search memories
results_df = self.db.search_memories(
user_id=user_id,
query_embedding=query_embedding,
project_id=project_id,
limit=limit * 2 if filter_with_llm else limit, # Get more if filtering
)
if results_df.empty:
return []
# Convert to memories with scores
memories = []
for _, row in results_df.iterrows():
# Get similarity score if available, otherwise calculate it
if "similarity_score" in row:
score = row["similarity_score"]
elif "embedding" in row:
# Calculate similarity score (cosine similarity)
embedding = row["embedding"]
score = self.embeddings.compute_similarity(
query_embedding, [embedding], metric="cosine"
)[0]
else:
score = 0.0 # Default score if no embedding available
# Parse metadata if it's a string
metadata = row.get("metadata", {})
if isinstance(metadata, str):
metadata = json.loads(metadata)
memory = MemoryWithScore(
memory_id=row.get("memory_id", row.get("id")),
user_id=row.get("user_id"),
project_id=row.get("project_id"),
content=row.get("content"),
metadata=metadata,
importance=row.get("importance", 0.5),
created_at=datetime.fromisoformat(row.get("created_at", datetime.now().isoformat())),
updated_at=datetime.fromisoformat(row.get("updated_at", datetime.now().isoformat())),
embedding=row.get("embedding"),
similarity_score=score,
)
memories.append(memory)
# Sort by similarity score
memories.sort(key=lambda m: m.similarity_score, reverse=True)
# Filter with LLM if requested
if filter_with_llm and memories:
memories = self._filter_with_llm(
query=query,
memories=memories[: limit * 2],
limit=limit,
additional_context=additional_context,
)
return memories[:limit]
def get_memory(self, user_id: str, memory_id: str) -> Memory | None:
"""Get a specific memory by ID."""
# This would need to be implemented in the DB client
# For now, return None
logger.warning(f"get_memory not fully implemented for {memory_id}")
return None
def delete_memory(self, user_id: str, memory_id: str) -> bool:
"""Delete a memory."""
# This would need to be implemented in the DB client
logger.warning(f"delete_memory not fully implemented for {memory_id}")
return False
def update_memory(
self,
user_id: str,
memory_id: str,
project_id: str,
content: str | None = None,
metadata: dict | None = None,
importance: float | None = None,
) -> Memory | None:
"""Update a memory."""
# Update the memory in the database
updated_data = self.db.update_memory(
memory_id=memory_id,
user_id=user_id,
project_id=project_id,
content=content,
metadata=metadata,
importance=importance,
)
if updated_data:
# Parse JSON fields if needed
if isinstance(updated_data.get("metadata"), str):
updated_data["metadata"] = json.loads(updated_data["metadata"])
if isinstance(updated_data.get("context"), str):
updated_data["context"] = json.loads(updated_data.get("context", "{}"))
return Memory(**updated_data)
return None
def delete_user_memories(self, user_id: str) -> int:
"""Delete all memories for a user."""
# This would need to be implemented in the DB client
logger.warning(f"delete_user_memories not fully implemented for {user_id}")
return 0
def _strip_pii(self, content: str) -> str:
"""Strip PII from content using LLM."""
prompt = f"""Remove any personally identifiable information (PII) from the following text.
Replace names with [NAME], emails with [EMAIL], phone numbers with [PHONE], addresses with [ADDRESS], etc.
Text: {content}
Anonymized text:"""
try:
anonymized = self.llm.complete(prompt)
return anonymized.strip()
except Exception as e:
logger.error(f"Error stripping PII: {e}")
return content
def _filter_with_llm(
self,
query: str,
memories: list[MemoryWithScore],
limit: int,
additional_context: str | None = None,
) -> list[MemoryWithScore]:
"""Filter memories using LLM to select most relevant."""
# Create memory list for LLM
memory_texts = []
for i, memory in enumerate(memories):
memory_texts.append(f"{i + 1}. {memory.content}")
context = (
f"Additional context: {additional_context}" if additional_context else ""
)
prompt = f"""Given the query: "{query}"
{context}
Select the {limit} most relevant memories from the following list.
Return only the numbers of the selected memories, separated by commas.
Memories:
{chr(10).join(memory_texts)}
Selected memory numbers:"""
try:
response = self.llm.complete(prompt)
# Parse selected indices
selected_indices = []
for part in response.strip().split(","):
try:
idx = int(part.strip()) - 1
if 0 <= idx < len(memories):
selected_indices.append(idx)
except ValueError:
continue
# Return selected memories
return [memories[i] for i in selected_indices[:limit]]
except Exception as e:
logger.error(f"Error filtering with LLM: {e}")
return memories[:limit]
# Global memory service instance
_memory_service: MemoryService | None = None
def get_memory_service() -> MemoryService:
"""Get or create the global memory service."""
global _memory_service
if _memory_service is None:
_memory_service = MemoryService()
return _memory_service
pkg/hanzo-memory/pyproject.toml
+3 -3
View File
@@ -7,7 +7,7 @@ name = "hanzo-memory"
version = "1.0.1"
description = "AI memory service with FastAPI and MCP support"
readme = "README.md"
requires-python = ">=3.10"
requires-python = ">=3.12"
license = {text = "BSD"}
authors = [
{name = "Hanzo Industries Inc.", email = "[email protected]"}
@@ -18,8 +18,6 @@ classifiers = [
"Intended Audience :: Developers",
"License :: OSI Approved :: BSD License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
]
@@ -57,6 +55,8 @@ test = [
"faker>=30.0.0",
"factory-boy>=3.3.0",
"respx>=0.21.0",
"polars>=1.0.0",
"httpx>=0.27.0",
]
docs = [
pkg/hanzo-memory/src/hanzo_memory/db/__init__.py
-2
View File
@@ -3,12 +3,10 @@
from .base import BaseVectorDB
from .client import InfinityClient, get_client
from .factory import get_db_client, reset_db_client
from .lancedb_client import LanceDBClient
__all__ = [
"BaseVectorDB",
"InfinityClient",
"LanceDBClient",
"get_db_client",
"reset_db_client",
"get_client",
pkg/hanzo-memory/src/hanzo_memory/db/backends/backend_registry.py
+24 -16
View File
@@ -13,11 +13,17 @@ from structlog import get_logger
from ..base import BaseVectorDB
from ..local_client import LocalMemoryClient
from ..lancedb_client import LanceDBClient
logger = get_logger()
# Try to import optional backends
try:
from ..lancedb_client import LanceDBClient
LANCEDB_AVAILABLE = True
except ImportError:
LANCEDB_AVAILABLE = False
logger.debug("LanceDB client not available")
try:
from ..kuzudb_client import KuzuDBClient
KUZU_AVAILABLE = True
@@ -58,27 +64,13 @@ class BackendRegistry:
# Registered backends and their capabilities
BACKENDS: Dict[str, Dict] = {
"lancedb": {
"class": LanceDBClient,
"description": "LanceDB - Embedded vector database with SQL-like queries",
"capabilities": [
BackendCapability.VECTOR_SEARCH,
BackendCapability.STRUCTURED_QUERY,
BackendCapability.PERSISTENCE,
BackendCapability.EMBEDDINGS,
BackendCapability.MARKDOWN_IMPORT,
],
"config": {
"enable_markdown": True,
},
},
"local": {
"class": LocalMemoryClient,
"description": "Local file storage using JSON files",
"capabilities": [
BackendCapability.PERSISTENCE,
BackendCapability.MARKDOWN_IMPORT,
BackendCapability.FULL_TEXT_SEARCH, # Simple text search
BackendCapability.FULL_TEXT_SEARCH,
],
"config": {
"enable_markdown": True,
@@ -90,6 +82,22 @@ class BackendRegistry:
@classmethod
def _init_optional_backends(cls):
"""Initialize optional backends if they're available."""
if LANCEDB_AVAILABLE and "lancedb" not in cls.BACKENDS:
cls.BACKENDS["lancedb"] = {
"class": LanceDBClient,
"description": "LanceDB - Embedded vector database with SQL-like queries",
"capabilities": [
BackendCapability.VECTOR_SEARCH,
BackendCapability.STRUCTURED_QUERY,
BackendCapability.PERSISTENCE,
BackendCapability.EMBEDDINGS,
BackendCapability.MARKDOWN_IMPORT,
],
"config": {
"enable_markdown": True,
},
}
if KUZU_AVAILABLE and "kuzudb" not in cls.BACKENDS:
cls.BACKENDS["kuzudb"] = {
"class": KuzuDBClient,
pkg/hanzo-memory/src/hanzo_memory/db/client.py
+127
View File
@@ -443,6 +443,133 @@ class InfinityClient(BaseVectorDB):
# InfinityDB embedded doesn't have explicit close method
pass
# Abstract method implementations required by BaseVectorDB
def get_user_projects(self, user_id: str) -> list[dict[str, Any]]:
"""Get all projects for a user."""
db = self._get_db("projects")
try:
table = db.get_table("projects")
result = table.output(["*"]).filter(f"user_id = '{user_id}'").to_pl()
return result.to_dicts() if len(result) > 0 else []
except Exception as e:
logger.error(f"Error getting user projects: {e}")
return []
def get_knowledge_bases(self, project_id: str) -> list[dict[str, Any]]:
"""Get all knowledge bases for a project."""
db = self._get_db("knowledge")
try:
table = db.get_table("knowledge_bases")
result = table.output(["*"]).filter(f"project_id = '{project_id}'").to_pl()
return result.to_dicts() if len(result) > 0 else []
except Exception as e:
logger.error(f"Error getting knowledge bases: {e}")
return []
def delete_fact(self, fact_id: str, knowledge_base_id: str) -> bool:
"""Delete a fact from a knowledge base."""
db = self._get_db("knowledge")
table_name = f"facts_{knowledge_base_id}"
try:
table = db.get_table(table_name)
table.delete(f"fact_id = '{fact_id}'")
return True
except Exception as e:
logger.error(f"Error deleting fact: {e}")
return False
def update_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
content: str | None = None,
metadata: dict | None = None,
importance: float | None = None,
) -> dict[str, Any] | None:
"""Update a memory in the database."""
db = self._get_db("memories")
table_name = f"memories_{user_id}"
try:
table = db.get_table(table_name)
updates = {"updated_at": datetime.now(timezone.utc).isoformat()}
if content is not None:
updates["content"] = content
if metadata is not None:
updates["metadata"] = json.dumps(metadata)
if importance is not None:
updates["importance"] = importance
table.update(f"memory_id = '{memory_id}'", updates)
return updates
except Exception as e:
logger.error(f"Error updating memory: {e}")
return None
def create_chat_session(
self,
session_id: str,
user_id: str,
project_id: str,
metadata: dict | None = None,
) -> dict[str, Any]:
"""Create a new chat session."""
self.create_chats_table(user_id)
return {
"session_id": session_id,
"user_id": user_id,
"project_id": project_id,
"metadata": metadata or {},
"created_at": datetime.now(timezone.utc).isoformat(),
}
def get_chat_messages(
self,
session_id: str,
limit: int | None = None,
) -> list[dict[str, Any]]:
"""Get messages from a chat session."""
# Note: This implementation searches all user tables - not ideal
db = self._get_db("chats")
try:
# Get all tables and search for the session
for table_name in db.list_tables():
if table_name.startswith("chats_"):
table = db.get_table(table_name)
query = table.output(["*"]).filter(f"session_id = '{session_id}'")
result = query.to_pl()
if len(result) > 0:
messages = result.to_dicts()
if limit:
return messages[:limit]
return messages
return []
except Exception as e:
logger.error(f"Error getting chat messages: {e}")
return []
def search_chat_messages(
self,
session_id: str,
query_embedding: list[float],
limit: int = 10,
) -> list[dict[str, Any]]:
"""Search messages in a chat session by similarity."""
db = self._get_db("chats")
try:
for table_name in db.list_tables():
if table_name.startswith("chats_"):
table = db.get_table(table_name)
query = table.output(["*"]).match_dense(
"embedding", query_embedding, "float", "cosine", limit
).filter(f"session_id = '{session_id}'")
result = query.to_pl()
if len(result) > 0:
return result.to_dicts()
return []
except Exception as e:
logger.error(f"Error searching chat messages: {e}")
return []
# Global client instance
_client: InfinityClient | None = None
pkg/hanzo-memory/src/hanzo_memory/db/lancedb_client.py
-1158
View File
File diff suppressed because it is too large Load Diff
pkg/hanzo-memory/src/hanzo_memory/db/sqlite_client.py
+40 -14
View File
@@ -3,7 +3,7 @@
import json
import sqlite3
import uuid
from datetime import datetime
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional
@@ -213,8 +213,8 @@ class SQLiteMemoryClient(BaseVectorDB):
"name": name,
"description": description,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
"updated_at": datetime.now(timezone.utc).isoformat(),
}
def get_user_projects(self, user_id: str) -> list[dict[str, Any]]:
@@ -291,9 +291,9 @@ class SQLiteMemoryClient(BaseVectorDB):
"context": {},
"metadata": metadata or {},
"source": "",
"timestamp": datetime.utcnow().isoformat(),
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
"timestamp": datetime.now(timezone.utc).isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
"updated_at": datetime.now(timezone.utc).isoformat(),
}
def search_memories(
@@ -412,8 +412,8 @@ class SQLiteMemoryClient(BaseVectorDB):
"name": name,
"description": description,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
"updated_at": datetime.now(timezone.utc).isoformat(),
}
def get_knowledge_bases(self, project_id: str) -> list[dict[str, Any]]:
@@ -469,8 +469,8 @@ class SQLiteMemoryClient(BaseVectorDB):
"content": content,
"confidence": confidence,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
"updated_at": datetime.now(timezone.utc).isoformat(),
}
def search_facts(
@@ -617,6 +617,20 @@ class SQLiteMemoryClient(BaseVectorDB):
return None
def delete_memory(
self,
memory_id: str,
user_id: str,
project_id: str,
) -> bool:
"""Delete a memory from the database."""
cursor = self.conn.execute(
"DELETE FROM memories WHERE memory_id = ? AND user_id = ? AND project_id = ?",
(memory_id, user_id, project_id)
)
self.conn.commit()
return cursor.rowcount > 0
def create_chat_session(
self,
session_id: str,
@@ -643,8 +657,8 @@ class SQLiteMemoryClient(BaseVectorDB):
"user_id": user_id,
"project_id": project_id,
"metadata": metadata or {},
"created_at": datetime.utcnow().isoformat(),
"updated_at": datetime.utcnow().isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
"updated_at": datetime.now(timezone.utc).isoformat(),
}
def add_chat_message(
@@ -678,7 +692,7 @@ class SQLiteMemoryClient(BaseVectorDB):
"role": role,
"content": content,
"metadata": metadata or {},
"timestamp": datetime.utcnow().isoformat(),
"timestamp": datetime.now(timezone.utc).isoformat(),
}
def get_chat_messages(
@@ -775,4 +789,16 @@ class SQLiteMemoryClient(BaseVectorDB):
def close(self) -> None:
"""Close the database connection."""
if self.conn:
self.conn.close()
self.conn.close()
# Singleton instance
_sqlite_client: Optional[SQLiteMemoryClient] = None
def get_sqlite_client() -> SQLiteMemoryClient:
"""Get or create the singleton SQLite client."""
global _sqlite_client
if _sqlite_client is None:
_sqlite_client = SQLiteMemoryClient()
return _sqlite_client
pkg/hanzo-memory/src/hanzo_memory/mcp/server.py
+220 -244
View File
@@ -1,7 +1,8 @@
"""MCP server for Hanzo Memory service."""
import json
from typing import Any
import asyncio
from typing import Any, Literal
from uuid import uuid4
from mcp.server import Server
@@ -13,13 +14,22 @@ from mcp.types import (
from structlog import get_logger
from ..config import settings
from ..db import get_db_client
from ..services.embeddings import EmbeddingService
from ..services.llm import LLMService
from ..db.sqlite_client import SQLiteMemoryClient, get_sqlite_client
from ..models.knowledge import Fact, FactCreate, KnowledgeBase
from ..models.memory import Memory, MemoryCreate
from ..services.embeddings import EmbeddingService, get_embedding_service
from ..services.llm import LLMService, get_llm_service
# Constants
ServiceName = "hanzo-memory"
logger = get_logger()
def get_db_client() -> SQLiteMemoryClient:
"""Get the database client."""
return get_sqlite_client()
class MCPMemoryServer:
"""MCP server for memory operations."""
@@ -39,170 +49,51 @@ class MCPMemoryServer:
"""Return available tools."""
return [
Tool(
name="remember",
description="Store a memory for a user in a project",
name="memory",
description="""Unified memory tool for storing, retrieving, and managing information.
Supports the following actions:
- remember: Store a new memory
- recall: Search for memories
- delete: Delete a memory
- create_project: Create a new project container
- create_kb: Create a structured knowledge base
- add_fact: Add a fact to a knowledge base
- search_facts: Search within a knowledge base
- delete_fact: Remove a fact from a knowledge base
- summarize: Analyze content and generate knowledge entries""",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"project_id": {
"action": {
"type": "string",
"description": "Project ID",
},
"content": {
"type": "string",
"description": "Memory content",
},
"metadata": {
"type": "object",
"description": "Additional metadata",
},
"importance": {
"type": "number",
"description": "Importance score (0-10)",
"default": 1.0,
"enum": [
"remember",
"recall",
"delete",
"create_project",
"create_kb",
"add_fact",
"search_facts",
"delete_fact",
"summarize"
],
"description": "The action to perform"
},
"user_id": {"type": "string", "description": "User ID (required for all actions)"},
"project_id": {"type": "string", "description": "Project ID (optional for some actions)"},
"content": {"type": "string", "description": "Content for storage or analysis"},
"query": {"type": "string", "description": "Search query for recall/search purposes"},
"id": {"type": "string", "description": "ID of item to delete (memory_id or fact_id)"},
"kb_id": {"type": "string", "description": "Knowledge base ID for fact operations"},
"name": {"type": "string", "description": "Name for new project or knowledge base"},
"description": {"type": "string", "description": "Description for new project or KB"},
"metadata": {"type": "object", "description": "Additional metadata"},
"limit": {"type": "integer", "description": "Limit results", "default": 10},
"importance": {"type": "number", "description": "Importance score (0-10)", "default": 1.0},
"parent_id": {"type": "string", "description": "Parent fact ID"},
"context": {"type": "string", "description": "Context for summarization"},
},
"required": ["user_id", "project_id", "content"],
},
),
Tool(
name="recall",
description="Search for relevant memories",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"project_id": {
"type": "string",
"description": "Project ID (optional)",
},
"query": {"type": "string", "description": "Search query"},
"limit": {
"type": "integer",
"description": "Number of results",
"default": 10,
},
},
"required": ["user_id", "query"],
},
),
Tool(
name="create_project",
description="Create a new project",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"name": {"type": "string", "description": "Project name"},
"description": {
"type": "string",
"description": "Project description",
},
"metadata": {
"type": "object",
"description": "Additional metadata",
},
},
"required": ["user_id", "name"],
},
),
Tool(
name="create_knowledge_base",
description="Create a knowledge base in a project",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"project_id": {
"type": "string",
"description": "Project ID",
},
"name": {
"type": "string",
"description": "Knowledge base name",
},
"description": {
"type": "string",
"description": "Knowledge base description",
},
},
"required": ["user_id", "project_id", "name"],
},
),
Tool(
name="add_fact",
description="Add a fact to a knowledge base",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"kb_id": {
"type": "string",
"description": "Knowledge base ID",
},
"content": {
"type": "string",
"description": "Fact content",
},
"parent_id": {
"type": "string",
"description": "Parent fact ID (optional)",
},
"metadata": {
"type": "object",
"description": "Additional metadata",
},
},
"required": ["user_id", "kb_id", "content"],
},
),
Tool(
name="search_facts",
description="Search facts in a knowledge base",
inputSchema={
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"kb_id": {
"type": "string",
"description": "Knowledge base ID",
},
"query": {"type": "string", "description": "Search query"},
"limit": {
"type": "integer",
"description": "Number of results",
"default": 10,
},
},
"required": ["user_id", "kb_id", "query"],
},
),
Tool(
name="summarize_for_knowledge",
description="Analyze content and generate knowledge update instructions",
inputSchema={
"type": "object",
"properties": {
"content": {
"type": "string",
"description": "Content to analyze",
},
"context": {
"type": "string",
"description": "Additional context",
},
"skip_summarization": {
"type": "boolean",
"description": "Skip summarization",
"default": False,
},
"provided_summary": {
"type": "string",
"description": "Pre-made summary",
},
},
"required": ["content"],
"required": ["action", "user_id"],
},
),
]
@@ -213,53 +104,67 @@ class MCPMemoryServer:
) -> list[TextContent]:
"""Handle tool calls."""
try:
if name == "remember":
result = await self._handle_remember(arguments or {})
elif name == "recall":
result = await self._handle_recall(arguments or {})
elif name == "create_project":
result = await self._handle_create_project(arguments or {})
elif name == "create_knowledge_base":
result = await self._handle_create_knowledge_base(arguments or {})
elif name == "add_fact":
result = await self._handle_add_fact(arguments or {})
elif name == "search_facts":
result = await self._handle_search_facts(arguments or {})
elif name == "summarize_for_knowledge":
result = await self._handle_summarize_for_knowledge(arguments or {})
if name != "memory":
return [TextContent(type="text", text=json.dumps({"error": f"Unknown tool: {name}"}))]
args = arguments or {}
action = args.get("action")
if action == "remember":
result = await self._handle_remember(args)
elif action == "recall":
result = await self._handle_recall(args)
elif action == "delete":
result = await self._handle_delete_memory(args)
elif action == "create_project":
result = await self._handle_create_project(args)
elif action == "create_kb":
result = await self._handle_create_knowledge_base(args)
elif action == "add_fact":
result = await self._handle_add_fact(args)
elif action == "search_facts":
result = await self._handle_search_facts(args)
elif action == "delete_fact":
result = await self._handle_delete_fact(args)
elif action == "summarize":
result = await self._handle_summarize_for_knowledge(args)
else:
result = {"error": f"Unknown tool: {name}"}
return [TextContent(type="text", text=json.dumps({"error": f"Unknown action: {action}"}))]
return [TextContent(type="text", text=json.dumps(result, indent=2))]
except Exception as e:
logger.error(f"Error handling tool {name}: {e}")
logger.error(f"Error handling memory action {arguments.get('action')}: {e}")
return [TextContent(type="text", text=json.dumps({"error": str(e)}))]
async def _handle_remember(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle remember tool."""
"""Handle remember action."""
user_id = args["user_id"]
project_id = args["project_id"]
content = args["content"]
project_id = args.get("project_id", "default")
content = args.get("content")
if not content:
raise ValueError("content is required for remember action")
metadata = args.get("metadata", {})
importance = args.get("importance", 1.0)
# Ensure user's memory table exists
self.db_client.create_memories_table(user_id)
# Generate embedding
embedding = self.embedding_service.embed_text(content)[0]
# Store memory
# Run embedding generation in thread pool
embedding = (await asyncio.to_thread(self.embedding_service.embed_text, content))[0]
memory_id = str(uuid4())
self.db_client.add_memory(
memory_id=memory_id,
user_id=user_id,
project_id=project_id,
content=content,
embedding=embedding,
metadata=metadata,
importance=importance,
)
# Run DB operations in thread pool
def _store():
self.db_client.create_memories_table(user_id)
self.db_client.add_memory(
memory_id=memory_id,
user_id=user_id,
project_id=project_id,
content=content,
embedding=embedding,
metadata=metadata,
importance=importance,
)
await asyncio.to_thread(_store)
return {
"success": True,
@@ -268,17 +173,21 @@ class MCPMemoryServer:
}
async def _handle_recall(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle recall tool."""
"""Handle recall action."""
user_id = args["user_id"]
project_id = args.get("project_id")
query = args["query"]
query = args.get("query")
if not query:
raise ValueError("query is required for recall action")
limit = args.get("limit", 10)
# Generate query embedding
query_embedding = self.embedding_service.embed_text(query)[0]
query_embedding = (await asyncio.to_thread(self.embedding_service.embed_text, query))[0]
# Search memories
results_df = self.db_client.search_memories(
results_df = await asyncio.to_thread(
self.db_client.search_memories,
user_id=user_id,
query_embedding=query_embedding,
project_id=project_id,
@@ -287,15 +196,17 @@ class MCPMemoryServer:
# Convert results
memories = []
if not results_df.is_empty():
if isinstance(results_df, list):
memories = results_df # It's a list of dicts from sqlite client fallback or non-polars return
elif not results_df.is_empty():
for row in results_df.to_dicts():
memories.append(
{
"memory_id": row["memory_id"],
"content": row["content"],
"metadata": json.loads(row.get("metadata", "{}")),
"memory_id": row.get("memory_id"),
"content": row.get("content"),
"metadata": json.loads(row.get("metadata", "{}")) if isinstance(row.get("metadata"), str) else row.get("metadata", {}),
"importance": row.get("importance", 1.0),
"similarity_score": row.get("_similarity", 0.0),
"similarity_score": row.get("_similarity", row.get("similarity_score", 0.0)),
}
)
@@ -305,16 +216,41 @@ class MCPMemoryServer:
"count": len(memories),
}
async def _handle_create_project(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle create project tool."""
async def _handle_delete_memory(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle delete memory action."""
user_id = args["user_id"]
name = args["name"]
project_id = args.get("project_id", "default")
memory_id = args.get("id")
if not memory_id:
raise ValueError("id is required for delete action")
success = await asyncio.to_thread(
self.db_client.delete_memory,
memory_id=memory_id,
user_id=user_id,
project_id=project_id
)
return {
"success": success,
"message": "Memory deleted" if success else "Memory not found"
}
async def _handle_create_project(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle create project action."""
user_id = args["user_id"]
name = args.get("name")
if not name:
raise ValueError("name is required for create_project action")
description = args.get("description", "")
metadata = args.get("metadata", {})
# Create project
project_id = str(uuid4())
self.db_client.create_project(
await asyncio.to_thread(
self.db_client.create_project,
project_id=project_id,
user_id=user_id,
name=name,
@@ -331,15 +267,19 @@ class MCPMemoryServer:
async def _handle_create_knowledge_base(
self, args: dict[str, Any]
) -> dict[str, Any]:
"""Handle create knowledge base tool."""
"""Handle create knowledge base action."""
user_id = args["user_id"]
project_id = args["project_id"]
name = args["name"]
project_id = args.get("project_id", "default")
name = args.get("name")
if not name:
raise ValueError("name is required for create_kb action")
description = args.get("description", "")
# Create knowledge base
kb_id = str(uuid4())
self.db_client.create_knowledge_base(
await asyncio.to_thread(
self.db_client.create_knowledge_base,
kb_id=kb_id,
user_id=user_id,
project_id=project_id,
@@ -354,18 +294,23 @@ class MCPMemoryServer:
}
async def _handle_add_fact(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle add fact tool."""
kb_id = args["kb_id"]
content = args["content"]
"""Handle add fact action."""
kb_id = args.get("kb_id")
if not kb_id:
raise ValueError("kb_id is required for add_fact action")
content = args.get("content")
if not content:
raise ValueError("content is required for add_fact action")
parent_id = args.get("parent_id")
metadata = args.get("metadata", {})
# Generate embedding
embedding = self.embedding_service.embed_text(content)[0]
# Add fact
embedding = (await asyncio.to_thread(self.embedding_service.embed_text, content))[0]
fact_id = str(uuid4())
self.db_client.add_fact(
await asyncio.to_thread(
self.db_client.add_fact,
fact_id=fact_id,
kb_id=kb_id,
content=content,
@@ -381,32 +326,38 @@ class MCPMemoryServer:
}
async def _handle_search_facts(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle search facts tool."""
kb_id = args["kb_id"]
query = args["query"]
"""Handle search facts action."""
kb_id = args.get("kb_id")
if not kb_id:
raise ValueError("kb_id is required for search_facts action")
query = args.get("query")
if not query:
raise ValueError("query is required for search_facts action")
limit = args.get("limit", 10)
# Generate query embedding
query_embedding = self.embedding_service.embed_text(query)[0]
query_embedding = (await asyncio.to_thread(self.embedding_service.embed_text, query))[0]
# Search facts
results_df = self.db_client.search_facts(
results_df = await asyncio.to_thread(
self.db_client.search_facts,
kb_id=kb_id,
query_embedding=query_embedding,
limit=limit,
)
# Convert results
facts = []
if not results_df.is_empty():
if isinstance(results_df, list):
facts = results_df
elif not results_df.is_empty():
for row in results_df.to_dicts():
facts.append(
{
"fact_id": row["fact_id"],
"content": row["content"],
"fact_id": row.get("fact_id"),
"content": row.get("content"),
"parent_id": row.get("parent_id"),
"metadata": json.loads(row.get("metadata", "{}")),
"similarity_score": row.get("_similarity", 0.0),
"metadata": json.loads(row.get("metadata", "{}")) if isinstance(row.get("metadata"), str) else row.get("metadata", {}),
"similarity_score": row.get("_similarity", row.get("similarity_score", 0.0)),
}
)
@@ -416,17 +367,40 @@ class MCPMemoryServer:
"count": len(facts),
}
async def _handle_delete_fact(self, args: dict[str, Any]) -> dict[str, Any]:
"""Handle delete fact action."""
kb_id = args.get("kb_id")
fact_id = args.get("id")
if not kb_id or not fact_id:
raise ValueError("kb_id and id are required for delete_fact action")
success = await asyncio.to_thread(
self.db_client.delete_fact,
fact_id=fact_id,
knowledge_base_id=kb_id
)
return {
"success": success,
"message": "Fact deleted" if success else "Fact not found"
}
async def _handle_summarize_for_knowledge(
self, args: dict[str, Any]
) -> dict[str, Any]:
"""Handle summarize for knowledge tool."""
content = args["content"]
"""Handle summarize for knowledge action."""
content = args.get("content")
if not content:
raise ValueError("content is required for summarize action")
context = args.get("context")
skip_summarization = args.get("skip_summarization", False)
provided_summary = args.get("provided_summary")
# Use LLM service to generate knowledge instructions
result = self.llm_service.summarize_for_knowledge(
# LLM service might be IO bound depending on implementation, assume safe to verify async later or wrap now
# Wrapping just in case
result = await asyncio.to_thread(
self.llm_service.summarize_for_knowledge,
content=content,
context=context,
skip_summarization=skip_summarization,
@@ -452,8 +426,10 @@ class MCPMemoryServer:
def main() -> None:
"""Main entry point for MCP server."""
import asyncio
"""Run the MCP memory server."""
server = MCPMemoryServer()
asyncio.run(server.run())
if __name__ == "__main__":
main()
pkg/hanzo-memory/src/hanzo_memory/models/memory.py
+1 -1
View File
@@ -61,7 +61,7 @@ class MemoryWithScore(Memory):
class MemoryResponse(BaseModel):
"""Memory response model with search results."""
memory: Memory = Field(..., description="Memory object")
memory: Memory | None = Field(None, description="Memory object")
similarity: float | None = Field(None, description="Similarity score")
relevance_score: float | None = Field(None, description="Relevance score")
pkg/hanzo-memory/src/hanzo_memory/server.py
+2
View File
@@ -524,6 +524,8 @@ async def get_facts(
status_code=400,
detail="Query parameter is required to search facts. Provide a search query.",
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error getting facts: {e}")
raise HTTPException(status_code=500, detail=str(e)) from e
pkg/hanzo-memory/src/hanzo_memory/services/__init__.py
+2 -1
View File
@@ -2,7 +2,7 @@
from .embeddings import EmbeddingService, get_embedding_service
from .llm import LLMService, get_llm_service
from .memory import MemoryService, get_memory_service
from .memory import MemoryService, get_memory_service, reset_memory_service
__all__ = [
"EmbeddingService",
@@ -11,4 +11,5 @@ __all__ = [
"get_llm_service",
"MemoryService",
"get_memory_service",
"reset_memory_service",
]
pkg/hanzo-memory/src/hanzo_memory/services/memory.py
+17 -2
View File
@@ -108,12 +108,21 @@ class MemoryService:
limit=limit * 2 if filter_with_llm else limit, # Get more if filtering
)
if results_df.empty:
# Handle both list and DataFrame results
if isinstance(results_df, list):
results = results_df
else:
# DataFrame case
if results_df.empty:
return []
results = results_df.to_dicts() if hasattr(results_df, 'to_dicts') else results_df.to_dict('records')
if not results:
return []
# Convert to memories with scores
memories = []
for _, row in results_df.iterrows():
for row in results:
# Get similarity score if available, otherwise calculate it
if "similarity_score" in row:
score = row["similarity_score"]
@@ -282,3 +291,9 @@ def get_memory_service() -> MemoryService:
if _memory_service is None:
_memory_service = MemoryService()
return _memory_service
def reset_memory_service() -> None:
"""Reset the global memory service (useful for testing)."""
global _memory_service
_memory_service = None
pkg/hanzo-memory/tests/conftest.py
+10
View File
@@ -11,12 +11,18 @@ from fastapi.testclient import TestClient
from hanzo_memory.config import settings
from hanzo_memory.db.client import InfinityClient
from hanzo_memory.db import reset_db_client
from hanzo_memory.services import reset_memory_service
from hanzo_memory.server import app
@pytest.fixture(autouse=True)
def test_settings(monkeypatch):
"""Configure test settings."""
# Reset any cached clients/services before each test
reset_db_client()
reset_memory_service()
# Use temporary directory for testing
with tempfile.TemporaryDirectory() as tmpdir:
monkeypatch.setattr(settings, "infinity_db_path", Path(tmpdir) / "test_db")
@@ -24,6 +30,10 @@ def test_settings(monkeypatch):
monkeypatch.setattr(settings, "llm_model", "gpt-3.5-turbo")
yield
# Reset again after test
reset_db_client()
reset_memory_service()
@pytest.fixture(autouse=True)
def mock_embedding_model():
pkg/hanzo-memory/tests/test_knowledge_api.py
+12
View File
@@ -162,6 +162,7 @@ class TestKnowledgeAPI:
def test_get_facts_without_query(self):
"""Test getting facts without a search query."""
# API requires query parameter
response = self.client.post(
"/v1/kb/facts/get",
json={
@@ -170,7 +171,18 @@ class TestKnowledgeAPI:
},
headers=self.headers,
)
assert response.status_code == 400
# With query, should return empty results
response = self.client.post(
"/v1/kb/facts/get",
json={
"userid": "user123",
"kb_id": "kb123",
"query": "test",
},
headers=self.headers,
)
assert response.status_code == 200
data = response.json()
assert data["kb_id"] == "kb123"
pkg/hanzo-memory/tests/test_mcp_server.py
+1 -1
View File
@@ -11,7 +11,7 @@ from hanzo_memory.mcp.server import MCPMemoryServer
def mcp_server():
"""Create MCP server instance."""
with (
patch("hanzo_memory.mcp.server.get_client"),
patch("hanzo_memory.mcp.server.get_db_client"),
patch("hanzo_memory.mcp.server.EmbeddingService"),
patch("hanzo_memory.mcp.server.LLMService"),
):
pkg/hanzo-memory/tests/test_memory_api.py
+9 -7
View File
@@ -97,19 +97,21 @@ class TestMemoryAPI:
def test_get_memories(self, client, sample_user_id):
"""Test /v1/memories/get endpoint."""
# Without memoryid, API returns 400
request_data = {
"userid": sample_user_id,
"limit": 10,
}
response = client.post("/v1/memories/get", json=request_data)
assert response.status_code == status.HTTP_200_OK
assert response.status_code == status.HTTP_400_BAD_REQUEST
data = response.json()
assert data["user_id"] == sample_user_id
assert "memories" in data
assert "pagination" in data
assert "usage_info" in data
# With memoryid but non-existent, returns 404
request_data = {
"userid": sample_user_id,
"memoryid": "non-existent-memory",
}
response = client.post("/v1/memories/get", json=request_data)
assert response.status_code == status.HTTP_404_NOT_FOUND
def test_delete_user_memories(self, client, sample_user_id):
"""Test /v1/user/delete endpoint."""
pkg/hanzo-memory/tests/test_unified_memory_tool.py
+105
View File
@@ -0,0 +1,105 @@
import pytest
import asyncio
from unittest.mock import MagicMock, patch
from hanzo_memory.mcp.server import MCPMemoryServer
@pytest.fixture
def memory_server():
server = MCPMemoryServer()
# Mock services to avoid external dependencies
server.db_client = MagicMock()
server.embedding_service = MagicMock()
server.llm_service = MagicMock()
# Mock embedding return
server.embedding_service.embed_text.return_value = [[0.1, 0.2, 0.3]]
# Mock DB returns
server.db_client.search_memories.return_value = []
server.db_client.search_facts.return_value = []
server.db_client.delete_memory.return_value = True
server.db_client.delete_fact.return_value = True
return server
@pytest.mark.asyncio
async def test_remember_and_recall(memory_server):
# Test remember
remember_args = {
"user_id": "user1",
"content": "Test memory",
"action": "remember"
}
result = await memory_server._handle_remember(remember_args)
assert result["success"] is True
assert "memory_id" in result
# Verify DB call
memory_server.db_client.add_memory.assert_called_once()
# Test recall
recall_args = {
"user_id": "user1",
"query": "Test",
"action": "recall"
}
result = await memory_server._handle_recall(recall_args)
assert result["success"] is True
assert "memories" in result
@pytest.mark.asyncio
async def test_delete_memory(memory_server):
delete_args = {
"user_id": "user1",
"id": "mem_123",
"action": "delete"
}
result = await memory_server._handle_delete_memory(delete_args)
assert result["success"] is True
# Verify DB call
memory_server.db_client.delete_memory.assert_called_once_with(
memory_id="mem_123",
user_id="user1",
project_id="default"
)
@pytest.mark.asyncio
async def test_knowledge_base_operations(memory_server):
# Test create KB
kb_args = {
"user_id": "user1",
"name": "My KB",
"action": "create_kb"
}
result = await memory_server._handle_create_knowledge_base(kb_args)
assert result["success"] is True
assert "kb_id" in result
kb_id = result["kb_id"]
# Test add fact
fact_args = {
"user_id": "user1",
"kb_id": kb_id,
"content": "A fact",
"action": "add_fact"
}
result = await memory_server._handle_add_fact(fact_args)
assert result["success"] is True
assert "fact_id" in result
fact_id = result["fact_id"]
# Test delete fact
del_fact_args = {
"user_id": "user1", # required by schema though not strictly used in delete_fact impl
"kb_id": kb_id,
"id": fact_id,
"action": "delete_fact"
}
result = await memory_server._handle_delete_fact(del_fact_args)
assert result["success"] is True
memory_server.db_client.delete_fact.assert_called_once_with(
fact_id=fact_id,
knowledge_base_id=kb_id
)
pyproject.toml
+20 -1
View File
@@ -59,13 +59,17 @@ dev-dependencies = [
"importlib-metadata>=6.7.0",
"rich>=13.7.1",
"nest_asyncio==1.6.0",
# hanzo-memory test dependencies
"polars>=1.0.0",
"httpx>=0.27.0",
"fastapi>=0.115.0",
"fastembed>=0.4.0",
]
[tool.rye.workspace]
members = [
"pkg/hanzo",
"pkg/hanzo-aci",
"pkg/hanzo-agent",
"pkg/hanzo-mcp",
"pkg/hanzo-memory",
"pkg/hanzo-network",
@@ -109,6 +113,13 @@ build-backend = "hatchling.build"
[dependency-groups]
dev = [
"build>=1.2.2.post1",
"fastapi>=0.128.0",
"fastembed>=0.7.4",
"hanzo-memory",
"httpx>=0.28.1",
"polars>=1.36.1",
"pytest>=8.4.2",
"pytest-asyncio>=1.2.0",
]
[tool.hatch.build]
@@ -421,6 +432,14 @@ exclude = [
"tests/",
]
[tool.uv.workspace]
members = [
"pkg/hanzo-memory",
]
[tool.uv.sources]
hanzo-memory = { workspace = true }
[tool.ruff.lint.extend-per-file-ignores]
# Allow print statements in scripts and examples
"bin/**.py" = ["T201", "T203"]
requirements-dev.lock
+273 -217
View File
File diff suppressed because it is too large Load Diff
requirements.lock
+307 -251
View File
File diff suppressed because it is too large Load Diff
uv.lock
Generated
+1940 -674
View File
File diff suppressed because it is too large Load Diff