Memory System

async def extract_memories(
    self,
    messages: List[Dict[str, Any]],
    account_id: str,
    thread_id: str
) -> List[ExtractedMemory]:
    if not config.ENABLE_MEMORY:
        return []
    
    conversation_text = self._format_conversation(messages)
    
    if len(conversation_text.strip()) < 20:
        return []
    
    prompt = MEMORY_EXTRACTION_PROMPT.format(conversation=conversation_text)
    
    response = await self.client.acompletion(
        model=resolved_model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
        max_tokens=2000,
        timeout=60
    )
    
    # Parse and validate extracted memories
    result = json.loads(response.choices[0].message.content)
    
    if not result.get('worth_extracting', True):
        return []
    
    memories = []
    for mem in result.get('memories', []):
        memories.append(ExtractedMemory(
            content=mem['content'],
            memory_type=MemoryType(mem['memory_type']),
            confidence_score=float(mem.get('confidence_score', 0.8)),
            metadata=mem.get('metadata', {})
        ))
    
    return memories

async def retrieve_memories(
    self,
    account_id: str,
    query_text: str,
    tier_name: str,
    similarity_threshold: float = 0.1
) -> List[MemoryItem]:
    if not config.ENABLE_MEMORY or not is_memory_enabled(tier_name):
        return []
    
    memory_config = get_memory_config(tier_name)
    retrieval_limit = memory_config.get('retrieval_limit', 0)
    
    # Check cache first
    cache_key = f"memories:retrieved:{account_id}:{hash(query_text)}"
    cached = await Cache.get(cache_key)
    if cached:
        return [self._dict_to_memory_item(m) for m in cached]
    
    # Generate embedding for query
    query_embedding = await self.embedding_service.embed_text(query_text)
    
    # Search using vector similarity
    result = await client.rpc(
        'search_memories_by_similarity',
        {
            'p_account_id': account_id,
            'p_query_embedding': query_embedding,
            'p_limit': retrieval_limit,
            'p_similarity_threshold': similarity_threshold
        }
    ).execute()
    
    memories = []
    for row in result.data:
        memory = MemoryItem(
            memory_id=row['memory_id'],
            account_id=account_id,
            content=row['content'],
            memory_type=MemoryType(row['memory_type']),
            confidence_score=row['confidence_score'],
            metadata=row.get('metadata', {})
        )
        memories.append(memory)
    
    # Cache results
    await Cache.set(cache_key, [self._memory_item_to_dict(m) for m in memories])
    
    return memories

def format_memories_for_prompt(self, memories: List[MemoryItem]) -> str:
    if not memories:
        return ""
    
    sections = {
        MemoryType.FACT: [],
        MemoryType.PREFERENCE: [],
        MemoryType.CONTEXT: [],
        MemoryType.CONVERSATION_SUMMARY: []
    }
    
    for memory in memories:
        sections[memory.memory_type].append(memory.content)
    
    formatted_parts = []
    
    if sections[MemoryType.FACT]:
        formatted_parts.append("Personal Facts:\n- " + "\n- ".join(sections[MemoryType.FACT]))
    
    if sections[MemoryType.PREFERENCE]:
        formatted_parts.append("Preferences:\n- " + "\n- ".join(sections[MemoryType.PREFERENCE]))
    
    if sections[MemoryType.CONTEXT]:
        formatted_parts.append("Context:\n- " + "\n- ".join(sections[MemoryType.CONTEXT]))
    
    if sections[MemoryType.CONVERSATION_SUMMARY]:
        formatted_parts.append("Past Conversations:\n- " + "\n- ".join(sections[MemoryType.CONVERSATION_SUMMARY]))
    
    return "# What You Remember About This User\n\n" + "\n\n".join(formatted_parts)

async def process_memory_extraction_job(
    thread_id: str,
    agent_run_id: str,
    account_id: str
):
    # Fetch conversation messages
    messages = await fetch_messages(thread_id, agent_run_id)
    
    # Extract memories using AI
    memories = await extraction_service.extract_memories(
        messages=messages,
        account_id=account_id,
        thread_id=thread_id
    )
    
    # Store with embeddings
    for memory in memories:
        embedding = await embedding_service.embed_text(memory.content)
        await store_memory(account_id, memory, embedding)
    
    # Invalidate cache
    await invalidate_cache(account_id)