Based on the proposal, I'll create an end-to-end POC implementation of the Contextual AI Assistant using Mem0. Let me build this step by step.
I'll create a working prototype that demonstrates how to implement the proposed system using Mem0's hybrid memory architecture. This POC will focus on:
First, let's set up our development environment with the necessary dependencies:
# contextual_assistant_poc.py
import os
import json
import datetime
import uuid
from typing import Dict, List, Any, Optional
import numpy as np
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
import uvicorn
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Initialize FastAPI app
app = FastAPI(title="Contextual AI Assistant")
# Mock Mem0 client since we're building a POC
class MockMem0Client:
def __init__(self):
# Initialize our storage systems
self.vector_store = {} # Simulates vector database
self.key_value_store = {} # Simulates key-value store
self.graph_store = {} # Simulates graph database
# Track relationships between entities
self.relationships = []
def add(self, entry: Dict[str, Any], user_id: str):
"""Add data to all appropriate stores"""
# Generate unique ID for the entry
entry_id = str(uuid.uuid4())
# Store in vector store (in a real system, we'd compute embeddings)
if "raw_content" in entry:
self.vector_store[entry_id] = {
"content": entry["raw_content"],
"metadata": {k: v for k, v in entry.items() if k != "raw_content"},
"user_id": user_id,
"timestamp": entry.get("timestamp", datetime.datetime.now().isoformat())
}
# Store entities in key-value store
if "entities" in entry:
for entity in entry["entities"]:
entity_id = entity.get("id", str(uuid.uuid4()))
self.key_value_store[entity_id] = {
**entity,
"user_id": user_id,
"entry_id": entry_id,
"timestamp": entry.get("timestamp", datetime.datetime.now().isoformat())
}
# Store relationships in graph store
if "relationships" in entry:
for rel in entry["relationships"]:
rel_id = str(uuid.uuid4())
self.graph_store[rel_id] = {
**rel,
"user_id": user_id,
"entry_id": entry_id,
"timestamp": entry.get("timestamp", datetime.datetime.now().isoformat())
}
self.relationships.append(rel)
return entry_id
def search(self, query: str, user_id: str, context_type: str = "vector",
filters: Dict[str, Any] = None, limit: int = 10):
"""Search for relevant information across stores"""
results = []
filters = filters or {}
if context_type == "vector":
# Simulate vector search with simple keyword matching
for entry_id, entry in self.vector_store.items():
if user_id != entry["user_id"]:
continue
# Simple keyword matching (in a real system, we'd use embeddings)
score = self._simple_similarity(query, entry["content"])
# Apply filters
if filters and not self._matches_filters(entry["metadata"], filters):
continue
if score > 0.2: # Arbitrary threshold
results.append({
"id": entry_id,
"content": entry["content"],
"score": score,
"metadata": entry["metadata"],
"timestamp": entry["timestamp"]
})
elif context_type == "key_value":
# Direct lookup with filters
for entity_id, entity in self.key_value_store.items():
if user_id != entity["user_id"]:
continue
if filters and not self._matches_filters(entity, filters):
continue
# Simple keyword matching on entity attributes
score = 0
for k, v in entity.items():
if isinstance(v, str) and any(term.lower() in v.lower() for term in query.split()):
score += 0.3
if score > 0.2: # Arbitrary threshold
results.append({
"id": entity_id,
**entity,
"score": score
})
elif context_type == "graph":
# Search relationships
for rel_id, rel in self.graph_store.items():
if user_id != rel["user_id"]:
continue
if filters and not self._matches_filters(rel, filters):
continue
# Simple keyword matching on relationship attributes
score = 0
for k, v in rel.items():
if isinstance(v, str) and any(term.lower() in v.lower() for term in query.split()):
score += 0.3
if score > 0.2: # Arbitrary threshold
results.append({
"id": rel_id,
**rel,
"score": score
})
# Also include connected entities
source_id = rel.get("source")
target_id = rel.get("target")
if source_id in self.key_value_store:
results.append({
"id": source_id,
**self.key_value_store[source_id],
"score": score * 0.9,
"connection": "source"
})
if target_id in self.key_value_store:
results.append({
"id": target_id,
**self.key_value_store[target_id],
"score": score * 0.9,
"connection": "target"
})
# Sort by score and limit results
results.sort(key=lambda x: x.get("score", 0), reverse=True)
return results[:limit]
def get(self, entity_id: str, user_id: str, context_type: str = "key_value"):
"""Get a specific entity by ID"""
if context_type == "key_value" and entity_id in self.key_value_store:
entity = self.key_value_store[entity_id]
if entity["user_id"] == user_id:
return entity
return None
def get_many(self, entity_ids: List[str], user_id: str, context_type: str = "key_value"):
"""Get multiple entities by ID"""
results = []
for entity_id in entity_ids:
entity = self.get(entity_id, user_id, context_type)
if entity:
results.append(entity)
return results
def _simple_similarity(self, query: str, text: str) -> float:
"""Simple text similarity function (placeholder for vector similarity)"""
query_terms = set(query.lower().split())
text_terms = set(text.lower().split())
# Jaccard similarity
intersection = query_terms.intersection(text_terms)
union = query_terms.union(text_terms)
if not union:
return 0.0
return len(intersection) / len(union)
def _matches_filters(self, item: Dict[str, Any], filters: Dict[str, Any]) -> bool:
"""Check if an item matches the given filters"""
for k, v in filters.items():
# Handle nested keys with dot notation
if "." in k:
parts = k.split(".")
current = item
for part in parts[:-1]:
if part not in current:
return False
current = current[part]
if parts[-1] not in current or current[parts[-1]] != v:
return False
# Handle direct keys
elif k not in item or item[k] != v:
return False
return True
# Initialize our mock Mem0 client
memory = MockMem0Client()
Now, let's define the data models for our API endpoints:
# Data models for API endpoints
class ScreenData(BaseModel):
"""Model for screen data sent to the ingestion API"""
content: str
metadata: Optional[Dict[str, Any]] = None
class ChatQuery(BaseModel):
"""Model for queries sent to the chat completion API"""
query: str
user_id: str
class ChatResponse(BaseModel):
"""Model for responses from the chat completion API"""
response: str
sources: Optional[List[Dict[str, Any]]] = None
Next, let's implement the ingestion pipeline that processes screen data: