memory_engine.py

import time
import uuid
from qdrant_client import QdrantClient, models
from fastembed import TextEmbedding
from deepface import DeepFace

class MemoryEngine:
    def __init__(self):
        # Initialize Qdrant (In-memory for Hackathon demo)
        self.client = QdrantClient(location=":memory:")
        self.text_embedder = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
        self._init_collections()

    def _init_collections(self):
        # 1. Visual Collection (Faces)
        if not self.client.collection_exists("semantic_identity"):
            self.client.recreate_collection(
                collection_name="semantic_identity",
                vectors_config=models.VectorParams(size=4096, distance=models.Distance.COSINE)
            )

        # 2. Episodic Collection (Context/Events)
        if not self.client.collection_exists("episodic_memory"):
            self.client.recreate_collection(
                collection_name="episodic_memory",
                vectors_config=models.VectorParams(size=384, distance=models.Distance.COSINE)
            )

    def get_face_vector(self, img_path):
        try:
            embedding_objs = DeepFace.represent(
                img_path=img_path, 
                model_name="VGG-Face", 
                enforce_detection=False
            )
            return embedding_objs[0]["embedding"]
        except Exception as e:
            return None

    def memorize_identity(self, img_path, name, relationship):
        """Stores a new identity in the visual memory."""
        vector = self.get_face_vector(img_path)
        if not vector: return None

        point_id = str(uuid.uuid4())
        self.client.upsert(
            collection_name="semantic_identity",
            points=[models.PointStruct(
                id=point_id,
                vector=vector,
                payload={
                    "name": name, 
                    "relationship": relationship,
                    "last_seen": int(time.time()),
                    "type": "identity"
                }
            )]
        )
        return point_id

    def update_identity(self, name, new_relationship):
        """Updates the relationship status for an existing identity."""
        hits = self.client.scroll(
            collection_name="semantic_identity",
            scroll_filter=models.Filter(
                must=[models.FieldCondition(key="name", match=models.MatchValue(value=name))]
            )
        )
        
        if hits[0]:
            point_id = hits[0][0].id
            self.client.set_payload(
                collection_name="semantic_identity",
                payload={"relationship": new_relationship},
                points=[point_id]
            )
            return True
        return False

    def memorize_event(self, text, related_person_name):
        """Links a text memory to a specific person."""
        vector = list(self.text_embedder.embed([text]))[0]
        
        self.client.upsert(
            collection_name="episodic_memory",
            points=[models.PointStruct(
                id=str(uuid.uuid4()),
                vector=vector,
                payload={
                    "content": text,
                    "related_name": related_person_name,
                    "timestamp": int(time.time())
                }
            )]
        )

    def recall(self, img_path):
        query_vector = self.get_face_vector(img_path)
        if not query_vector:
            return {"status": "error", "msg": "No face detected in the image."}

        try:
            face_hits = self.client.search(
                collection_name="semantic_identity",
                query_vector=query_vector,
                limit=1
            )
        except AttributeError:
            print("Switching to low-level query_points API...")
            result = self.client.query_points(
                collection_name="semantic_identity",
                query=query_vector,
                limit=1
            )
            face_hits = result.points

        if not face_hits or face_hits[0].score < 0.60:
            return {
                "status": "unknown", 
                "confidence": face_hits[0].score if face_hits else 0.0,
                "msg": "Safety Alert: Identity confidence below threshold."
            }

        person = face_hits[0].payload
        
        context_vector = list(self.text_embedder.embed(["recent updates"]))[0]
        context_filter = models.Filter(
            must=[models.FieldCondition(
                key="related_name", 
                match=models.MatchValue(value=person['name'])
            )]
        )

        try:
            context_hits = self.client.search(
                collection_name="episodic_memory",
                query_vector=context_vector,
                query_filter=context_filter,
                limit=1
            )
        except AttributeError:
             result = self.client.query_points(
                collection_name="episodic_memory",
                query=context_vector,
                query_filter=context_filter,
                limit=1
            )
             context_hits = result.points
        
        last_memory = context_hits[0].payload['content'] if context_hits else "No specific recent records found."
        
        return {
            "status": "known",
            "name": person['name'],
            "relationship": person['relationship'],
            "confidence": face_hits[0].score,
            "context": last_memory,
            "id": face_hits[0].id
        }