%pip install --upgrade pip
%pip install mem0ai pyautogen flaml

# Set up ENV Vars
import os

os.environ["OPENAI_API_KEY"] = "sk-xxx"

# AutoGen GPTAssistantAgent Capabilities:
# - Generates code based on user requirements and preferences.
# - Analyzes, refactors, and debugs existing code efficiently.
# - Maintains consistent coding standards across multiple sessions.
# - Remembers project-specific conventions and architectural decisions.
# - Learns from past interactions to improve future code suggestions.
# - Reduces repetitive explanations of coding preferences, enhancing productivity.
# - Adapts to team-specific practices for a more cohesive development process.

import logging
import os

from autogen import AssistantAgent, UserProxyAgent
from autogen.agentchat.contrib.gpt_assistant_agent import GPTAssistantAgent

logger = logging.getLogger(__name__)
logger.setLevel(logging.WARNING)

assistant_id = os.environ.get("ASSISTANT_ID", None)

# LLM Configuration
CACHE_SEED = 42  # choose your poison
llm_config = {
    "config_list": [{"model": "gpt-4o", "api_key": os.environ["OPENAI_API_KEY"]}],
    "cache_seed": CACHE_SEED,
    "timeout": 120,
    "temperature": 0.0,
}

assistant_config = {"assistant_id": assistant_id}

gpt_assistant = GPTAssistantAgent(
    name="assistant",
    instructions=AssistantAgent.DEFAULT_SYSTEM_MESSAGE,
    llm_config=llm_config,
    assistant_config=assistant_config,
)

user_proxy = UserProxyAgent(
    name="user_proxy",
    code_execution_config={
        "work_dir": "coding",
        "use_docker": False,
    },  # Please set use_docker=True if docker is available to run the generated code. Using docker is safer than running the generated code directly.
    is_termination_msg=lambda msg: "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
    max_consecutive_auto_reply=1,
    llm_config=llm_config,
)

user_query = "Write a Python function that reverses a string."
# Initiate Chat w/o Memory
user_proxy.initiate_chat(gpt_assistant, message=user_query)

# Benefits of Preference Memory in AutoGen Agents:
# - Personalization: Tailors responses to individual user or team preferences.
# - Consistency: Maintains uniform coding style and standards across sessions.
# - Efficiency: Reduces need to restate preferences, saving time in each interaction.
# - Adaptability: Evolves understanding of user needs over multiple conversations.
# - Context Retention: Keeps project-specific details accessible without repetition.
# - Improved Recommendations: Suggests solutions aligned with past preferences.
# - Long-term Learning: Accumulates knowledge to enhance future interactions.
# - Reduced Cognitive Load: Users don't need to remember and restate all preferences.


# Setting memory (preference) for the user
from mem0 import Memory

# Initialize Mem0
MEM0_MEMORY_CLIENT = Memory()

USER_ID = "chicory.ai.user"
MEMORY_DATA = """
* Preference for readability: The user prefers code to be explicitly written with clear variable names.
* Preference for comments: The user prefers comments explaining each step.
* Naming convention: The user prefers camelCase for variable names.
* Docstrings: The user prefers functions to have a descriptive docstring.
"""
AGENT_ID = "chicory.ai"

# Add preference data to memory
MEM0_MEMORY_CLIENT.add(MEMORY_DATA, user_id=USER_ID)
MEM0_MEMORY_CLIENT.add(MEMORY_DATA, agent_id=AGENT_ID)

# Retrieve the memory
relevant_memories = MEM0_MEMORY_CLIENT.search(user_query, user_id=USER_ID, limit=3)
relevant_memories_text = "\n".join(mem["memory"] for mem in relevant_memories)
print("Relevant memories:")
print(relevant_memories_text)

prompt = f"{user_query}\n Coding Preferences: \n{relevant_memories_text}"
browse_result = user_proxy.initiate_chat(gpt_assistant, message=prompt)

# UserProxyAgent in AutoGen:
# - Acts as intermediary between humans and AI agents in the AutoGen framework.
# - Simulates user behavior and interactions within multi-agent conversations.
# - Can be configured to execute code blocks received in messages.
# - Supports flexible human input modes (e.g., ALWAYS, TERMINATE, NEVER).
# - Customizable for specific interaction patterns and behaviors.
# - Can be integrated with memory systems like mem0 for enhanced functionality.
# - Capable of fetching relevant memories before processing a query.
# - Enables more context-aware and personalized agent responses.
# - Bridges the gap between human input and AI processing in complex workflows.


class Mem0ProxyCoderAgent(UserProxyAgent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.memory = MEM0_MEMORY_CLIENT
        self.agent_id = kwargs.get("name")

    def initiate_chat(self, assistant, message):
        # Retrieve memory for the agent
        agent_memories = self.memory.search(message, agent_id=self.agent_id, limit=3)
        agent_memories_txt = "\n".join(mem["memory"] for mem in agent_memories)
        prompt = f"{message}\n Coding Preferences: \n{str(agent_memories_txt)}"
        response = super().initiate_chat(assistant, message=prompt)
        # Add new memory after processing the message
        response_dist = response.__dict__ if not isinstance(response, dict) else response
        MEMORY_DATA = [{"role": "user", "content": message}, {"role": "assistant", "content": response_dist}]
        self.memory.add(MEMORY_DATA, agent_id=self.agent_id)
        return response

mem0_user_proxy = Mem0ProxyCoderAgent(
    name=AGENT_ID,
    code_execution_config={
        "work_dir": "coding",
        "use_docker": False,
    },  # Please set use_docker=True if docker is available to run the generated code. Using docker is safer than running the generated code directly.
    is_termination_msg=lambda msg: "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
    max_consecutive_auto_reply=1,
)
code_result = mem0_user_proxy.initiate_chat(gpt_assistant, message=user_query)

# building on top of existing Teachability package from autogen
# from autogen.agentchat.contrib.capabilities.teachability import Teachability

# AutoGen Teachability Feature:
# - Enables agents to learn and remember across multiple chat sessions.
# - Addresses the limitation of traditional LLMs forgetting after conversations end.
# - Uses vector database to store "memos" of taught information.
# - Can remember facts, preferences, and even complex skills.
# - Allows for cumulative learning and knowledge retention over time.
# - Enhances personalization and adaptability of AI assistants.
# - Can be integrated with mem0 for improved memory management.
# - Potential for more efficient and context-aware information retrieval.
# - Enables creation of AI agents with long-term memory and learning abilities.
# - Improves consistency and reduces repetition in user-agent interactions.

from cookbooks.helper.mem0_teachability import Mem0Teachability

teachability = Mem0Teachability(
    verbosity=2,  # for visibility of what's happening
    recall_threshold=0.5,
    reset_db=False,  # Use True to force-reset the memo DB, and False to use an existing DB.
    agent_id=AGENT_ID,
    memory_client=MEM0_MEMORY_CLIENT,
)
teachability.add_to_agent(user_proxy)

# Initiate Chat w/ Teachability + Memory
user_proxy.initiate_chat(gpt_assistant, message=user_query)