diff --git a/examples/misc/multillm_memory.py b/examples/misc/multillm_memory.py new file mode 100644 index 00000000..ff8f275e --- /dev/null +++ b/examples/misc/multillm_memory.py @@ -0,0 +1,193 @@ +""" +Multi-LLM Research Team with Shared Knowledge Base + +Use Case: AI Research Team where each model has different strengths: +- GPT-4: Technical analysis and code review +- Claude: Writing and documentation + +All models share a common knowledge base, building on each other's work. +Example: GPT-4 analyzes a tech stack → Claude writes documentation → +Data analyst analyzes user data → All models can reference previous research. +""" + +from dotenv import load_dotenv +import logging +from mem0 import MemoryClient +from litellm import completion + +load_dotenv() + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + handlers=[ + logging.StreamHandler(), + logging.FileHandler('research_team.log') + ] +) +logger = logging.getLogger(__name__) + + +# Initialize memory client (platform version) +memory = MemoryClient() + +# Research team models with specialized roles +RESEARCH_TEAM = { + "tech_analyst": { + "model": "gpt-4o", + "role": "Technical Analyst - Code review, architecture, and technical decisions" + }, + "writer": { + "model": "claude-3-5-sonnet-20241022", + "role": "Documentation Writer - Clear explanations and user guides" + }, + "data_analyst": { + "model": "gpt-4o-mini", + "role": "Data Analyst - Insights, trends, and data-driven recommendations" + } +} + + +def get_team_knowledge(topic: str, project_id: str) -> str: + """Get relevant research from the team's shared knowledge base""" + memories = memory.search(query=topic, user_id=project_id, limit=5) + + if memories: + knowledge = "Team Knowledge Base:\n" + for mem in memories: + if "memory" in mem: + # Get metadata to show which team member contributed + metadata = mem.get("metadata", {}) + contributor = metadata.get("contributor", "Unknown") + knowledge += f"• [{contributor}] {mem['memory']}\n" + return knowledge + return "Team Knowledge Base: Empty - starting fresh research" + + +def research_with_specialist( + task: str, + specialist: str, + project_id: str +) -> str: + """Assign research task to specialist with access to team knowledge""" + + if specialist not in RESEARCH_TEAM: + return f"Unknown specialist. Available: {list(RESEARCH_TEAM.keys())}" + + # Get team's accumulated knowledge + team_knowledge = get_team_knowledge(task, project_id) + + # Specialist role and model + spec_info = RESEARCH_TEAM[specialist] + + system_prompt = f"""You are the {spec_info['role']}. + +{team_knowledge} + +Build upon the team's existing research. Reference previous findings when relevant. +Provide actionable insights in your area of expertise.""" + + # Call the specialist's model + response = completion( + model=spec_info["model"], + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": task} + ] + ) + + result = response.choices[0].message.content + + # Store research in shared knowledge base using both user_id and agent_id + research_entry = [ + {"role": "user", "content": f"Task: {task}"}, + {"role": "assistant", "content": result} + ] + + memory.add( + research_entry, + user_id=project_id, # Project-level memory + agent_id=specialist, # Agent-specific memory + metadata={ + "contributor": specialist, + "task_type": "research", + "model_used": spec_info["model"] + }, + output_format="v1.1" + ) + + return result + + +def show_team_knowledge(project_id: str): + """Display the team's accumulated research""" + memories = memory.get_all(user_id=project_id) + + if not memories: + logger.info("No research found for this project") + return + + logger.info(f"Team Research Summary (Project: {project_id}):") + + # Group by contributor + by_contributor = {} + for mem in memories: + if "metadata" in mem and mem["metadata"]: + contributor = mem["metadata"].get("contributor", "Unknown") + if contributor not in by_contributor: + by_contributor[contributor] = [] + by_contributor[contributor].append(mem.get("memory", "")) + + for contributor, research_items in by_contributor.items(): + logger.info(f"{contributor.upper()}:") + for i, item in enumerate(research_items[:3], 1): # Show latest 3 + logger.info(f" {i}. {item[:100]}...") + + +def demo_research_team(): + """Demo: Building a SaaS product with the research team""" + + project = "saas_product_research" + + # Define research pipeline + research_pipeline = [ + { + "stage": "Technical Architecture", + "specialist": "tech_analyst", + "task": "Analyze the best tech stack for a multi-tenant SaaS platform handling 10k+ users. Consider scalability, cost, and development speed." + }, + { + "stage": "Product Documentation", + "specialist": "writer", + "task": "Based on the technical analysis, write a clear product overview and user onboarding guide for our SaaS platform." + }, + { + "stage": "Market Analysis", + "specialist": "data_analyst", + "task": "Analyze market trends and pricing strategies for our SaaS platform. What metrics should we track?" + }, + { + "stage": "Strategic Decision", + "specialist": "tech_analyst", + "task": "Given our technical architecture, documentation, and market analysis - what should be our MVP feature priority?" + } + ] + + logger.info("AI Research Team: Building a SaaS Product") + + # Execute research pipeline + for i, step in enumerate(research_pipeline, 1): + logger.info(f"\nStage {i}: {step['stage']}") + logger.info(f"Specialist: {step['specialist']}") + + result = research_with_specialist(step['task'], step['specialist'], project) + logger.info(f"Task: {step['task']}") + logger.info(f"Result: {result[:200]}...\n") + + show_team_knowledge(project) + + +if __name__ == "__main__": + logger.info("Multi-LLM Research Team") + demo_research_team()