Build a Long-Term Memory Agent with Python and pgvector

# config.py
import os
import logging
from dotenv import load_dotenv

load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Database Configuration
POSTGRES_HOST = os.environ.get("POSTGRES_HOST", "localhost")
POSTGRES_PORT = os.environ.get("POSTGRES_PORT", "5432")
POSTGRES_DB = os.environ.get("POSTGRES_DB", "agent_memory_db")
POSTGRES_USER = os.environ.get("POSTGRES_USER", "user")
POSTGRES_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "password")
POSTGRES_URI = f"postgresql+psycopg2://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}"

# LLM Configuration
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
OPENAI_MODEL_NAME = os.environ.get("OPENAI_MODEL_NAME", "gpt-4o")

# LangSmith Configuration (Optional but recommended for observability)
LANGCHAIN_TRACING_V2 = os.environ.get("LANGCHAIN_TRACING_V2", "false").lower() == "true"
LANGCHAIN_API_KEY = os.environ.get("LANGCHAIN_API_KEY")
LANGCHAIN_PROJECT = os.environ.get("LANGCHAIN_PROJECT", "long-term-memory-agent")

if not OPENAI_API_KEY:
    logger.error("OPENAI_API_KEY not set. Please set it in your environment or .env file.")
    raise ValueError("OPENAI_API_KEY is required.")

logger.info("Configuration loaded successfully.")

# db/init_db.sql
-- Enable pgvector extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Table for Episodic Memory
CREATE TABLE IF NOT EXISTS episodic_memory (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    timestamp TIMESTAMPTZ DEFAULT NOW(),
    event_type VARCHAR(255) NOT NULL,
    content TEXT NOT NULL,
    embedding VECTOR(1536) NOT NULL, -- text-embedding-3-small / text-embedding-ada-002 produce 1536-dim vectors
    metadata JSONB DEFAULT '{}'
);

-- Table for Semantic Memory (Knowledge Graph triples)
CREATE TABLE IF NOT EXISTS semantic_memory (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    subject TEXT NOT NULL,
    predicate TEXT NOT NULL,
    object TEXT NOT NULL,
    source_episodic_id UUID, -- Link to original episodic memory if consolidated
    timestamp TIMESTAMPTZ DEFAULT NOW(),
    CONSTRAINT fk_source_episodic
        FOREIGN KEY(source_episodic_id)
        REFERENCES episodic_memory(id)
        ON DELETE SET NULL
);

-- Index for efficient semantic memory retrieval
CREATE INDEX IF NOT EXISTS idx_semantic_subject ON semantic_memory (subject);
CREATE INDEX IF NOT EXISTS idx_semantic_predicate ON semantic_memory (predicate);
CREATE INDEX IF NOT EXISTS idx_semantic_object ON semantic_memory (object);

-- Index for efficient episodic memory search
CREATE INDEX IF NOT EXISTS idx_episodic_timestamp ON episodic_memory (timestamp DESC);
CREATE INDEX IF NOT EXISTS idx_episodic_event_type ON episodic_memory (event_type);

# components/memory.py
import os
import logging
import uuid
from typing import List, Dict, Any, Optional
from datetime import datetime

import psycopg2
from psycopg2.extras import DictCursor, register_uuid
from sqlalchemy import create_engine, text
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_community.vectorstores import PGVector

from config import POSTGRES_URI, OPENAI_API_KEY

logger = logging.getLogger(__name__)
register_uuid() # Enable UUID support for psycopg2

class MemoryManager:
    """Manages interactions with episodic and semantic memory."""

    def __init__(self):
        if not OPENAI_API_KEY:
            raise ValueError("OPENAI_API_KEY is not set. Cannot initialize embeddings.")
        
        self.embedding_model = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY, model="text-embedding-3-small")
        self.pg_engine = create_engine(POSTGRES_URI)
        
        # Initialize PGVector for episodic memory
        self.episodic_vector_store = PGVector(
            connection_string=POSTGRES_URI,
            embedding_function=self.embedding_model,
            collection_name="episodic_memory",
            table_name="episodic_memory", # Ensure this matches our table name
            pre_delete_collection=False # Do not delete existing table on init
        )
        logger.info("MemoryManager initialized with PostgreSQL and OpenAIEmbeddings.")

    def _get_db_connection(self):
        """Establishes a new database connection."""
        try:
            conn = psycopg2.connect(POSTGRES_URI.replace("postgresql+psycopg2", "postgresql"))
            conn.autocommit = True
            return conn
        except Exception as e:
            logger.error(f"Failed to connect to database: {e}")
            raise

    def add_episodic_memory(self, event_type: str, content: str, metadata: Optional[Dict[str, Any]] = None) -> str:
        """Adds an event to episodic memory, generates an embedding, and stores it."""
        try:
            embedding = self.embedding_model.embed_query(content)
            
            with self._get_db_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        """
                        INSERT INTO episodic_memory (event_type, content, embedding, metadata)
                        VALUES (%s, %s, %s, %s) RETURNING id;
                        """,
                        (event_type, content, embedding, metadata if metadata else {})
                    )
                    memory_id = cur.fetchone()[0]
            logger.info(f"Added episodic memory (ID: {memory_id}, Type: {event_type})")
            return str(memory_id)
        except Exception as e:
            logger.error(f"Error adding episodic memory: {e}")
            raise

    def retrieve_episodic_memory(self, query: str, k: int = 5) -> List[Document]:
        """Retrieves relevant episodic memories based on a semantic query."""
        try:
            results = self.episodic_vector_store.similarity_search(query, k=k)
            logger.info(f"Retrieved {len(results)} episodic memories for query: '{query}'")
            return results
        except Exception as e:
            logger.error(f"Error retrieving episodic memory for query '{query}': {e}")
            return []

    def add_semantic_triple(self, subject: str, predicate: str, object: str, source_episodic_id: Optional[str] = None) -> str:
        """Adds a semantic triple (subject, predicate, object) to the knowledge graph."""
        try:
            with self._get_db_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        """
                        INSERT INTO semantic_memory (subject, predicate, object, source_episodic_id)
                        VALUES (%s, %s, %s, %s) RETURNING id;
                        """,
                        (subject, predicate, object, uuid.UUID(source_episodic_id) if source_episodic_id else None)
                    )
                    triple_id = cur.fetchone()[0]
            logger.info(f"Added semantic triple (ID: {triple_id}): {subject} - {predicate} - {object}")
            return str(triple_id)
        except Exception as e:
            logger.error(f"Error adding semantic triple '{subject} {predicate} {object}': {e}")
            raise

    def retrieve_semantic_memory(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Retrieves semantic memory triples based on a keyword query. 
           Can extend to full knowledge graph traversal if needed."""
        try:
            with self.pg_engine.connect() as conn:
                # Simple keyword search for demonstration; can be enhanced with graph queries
                search_query = f"%{query.lower()}%"
                result = conn.execute(
                    text("""
                        SELECT id, subject, predicate, object, source_episodic_id, timestamp
                        FROM semantic_memory
                        WHERE LOWER(subject) LIKE :search_query
                           OR LOWER(predicate) LIKE :search_query
                           OR LOWER(object) LIKE :search_query
                        LIMIT :k;
                    """),
                    {"search_query": search_query, "k": k}
                ).fetchall()
                
                triples = []
                for row in result:
                    triples.append({
                        "id": str(row[0]),
                        "subject": row[1],
                        "predicate": row[2],
                        "object": row[3],
                        "source_episodic_id": str(row[4]) if row[4] else None,
                        "timestamp": row[5].isoformat()
                    })
            logger.info(f"Retrieved {len(triples)} semantic memories for query: '{query}'")
            return triples
        except Exception as e:
            logger.error(f"Error retrieving semantic memory for query '{query}': {e}")
            return []

    def get_all_episodic_memories(self, limit: int = 1000) -> List[Dict[str, Any]]:
        """Retrieves all episodic memories up to a limit. Useful for consolidation."""
        try:
            with self._get_db_connection() as conn:
                with conn.cursor(cursor_factory=DictCursor) as cur:
                    cur.execute("SELECT id, content, timestamp, event_type, metadata FROM episodic_memory ORDER BY timestamp DESC LIMIT %s;", (limit,))
                    memories = cur.fetchall()
            logger.info(f"Retrieved {len(memories)} episodic memories for consolidation.")
            return [dict(m) for m in memories]
        except Exception as e:
            logger.error(f"Error retrieving all episodic memories: {e}")
            return []

# components/tools.py
import logging
from typing import List, Dict, Any, Optional

from langchain.tools import tool
from langchain_core.documents import Document

from components.memory import MemoryManager

logger = logging.getLogger(__name__)

class AgentTools:
    """A collection of tools for the agent to interact with its memory systems."""

    def __init__(self, memory_manager: MemoryManager):
        self.memory_manager = memory_manager

    @tool
    def add_episodic_memory_tool(self, event_type: str, content: str, metadata: Dict[str, Any] = None) -> str:
        """
        Adds a new event or observation to the agent's episodic memory.
        Use this tool to record important interactions, observations, or facts learned during a conversation.
        The `event_type` should categorize the memory (e.g., "conversation", "observation", "action").
        The `content` is the detailed description of the event.
        Optional `metadata` can include context like user_id, topic, etc.
        Returns the ID of the added memory.
        """
        try:
            memory_id = self.memory_manager.add_episodic_memory(event_type, content, metadata)
            logger.info(f"Tool: Successfully added episodic memory with ID {memory_id}")
            return f"Episodic memory added with ID: {memory_id}"
        except Exception as e:
            logger.error(f"Tool: Failed to add episodic memory: {e}")
            return f"Error: Failed to add episodic memory - {e}"

    @tool
    def retrieve_episodic_memory_tool(self, query: str, k: int = 3) -> List[str]:
        """
        Retrieves relevant past events or observations from episodic memory based on a semantic query.
        Use this tool when you need to recall specific past interactions or details.
        `query` is the natural language question or description of what to retrieve.
        `k` is the maximum number of relevant memories to retrieve (default 3).
        Returns a list of strings, where each string is the content of a retrieved memory.
        """
        try:
            results: List[Document] = self.memory_manager.retrieve_episodic_memory(query, k)
            if not results:
                return ["No relevant episodic memories found."]
            
            formatted_results = [
                f"Memory (ID: {doc.metadata.get('id')}, Timestamp: {doc.metadata.get('timestamp')}): {doc.page_content}"
                for doc in results
            ]
            logger.info(f"Tool: Successfully retrieved {len(formatted_results)} episodic memories.")
            return formatted_results
        except Exception as e:
            logger.error(f"Tool: Failed to retrieve episodic memory: {e}")
            return [f"Error: Failed to retrieve episodic memory - {e}"]

    @tool
    def add_semantic_triple_tool(self, subject: str, predicate: str, object: str, source_episodic_id: Optional[str] = None) -> str:
        """
        Adds a structured fact (subject, predicate, object) to the agent's semantic memory (knowledge graph).
        Use this tool to store important, consolidated facts that represent long-term knowledge.
        `subject`, `predicate`, `object` define the fact.
        `source_episodic_id` is an optional ID linking to the episodic memory from which this fact was derived.
        Returns the ID of the added semantic triple.
        """
        try:
            triple_id = self.memory_manager.add_semantic_triple(subject, predicate, object, source_episodic_id)
            logger.info(f"Tool: Successfully added semantic triple with ID {triple_id}")
            return f"Semantic triple added with ID: {triple_id}"
        except Exception as e:
            logger.error(f"Tool: Failed to add semantic triple: {e}")
            return f"Error: Failed to add semantic triple - {e}"

    @tool
    def retrieve_semantic_memory_tool(self, query: str, k: int = 3) -> List[str]:
        """
        Retrieves consolidated facts from the agent's semantic memory (knowledge graph) based on a keyword query.
        Use this tool to access high-level, generalized knowledge.
        `query` is a keyword or phrase to search for in subjects, predicates, or objects.
        `k` is the maximum number of relevant triples to retrieve (default 3).
        Returns a list of strings, where each string represents a retrieved semantic triple.
        """
        try:
            results: List[Dict[str, Any]] = self.memory_manager.retrieve_semantic_memory(query, k)
            if not results:
                return ["No relevant semantic memories found."]
            
            formatted_results = [
                f"Fact (ID: {triple['id']}, Source: {triple['source_episodic_id'] if triple['source_episodic_id'] else 'N/A'}): {triple['subject']} {triple['predicate']} {triple['object']}"
                for triple in results
            ]
            logger.info(f"Tool: Successfully retrieved {len(formatted_results)} semantic memories.")
            return formatted_results
        except Exception as e:
            logger.error(f"Tool: Failed to retrieve semantic memory: {e}")
            return [f"Error: Failed to retrieve semantic memory - {e}"]

# components/agent.py
import logging
import os
from typing import TypedDict, List, Annotated, Sequence, Union, Dict, Any
from operator import add

from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langgraph.graph import StateGraph, END

from components.memory import MemoryManager
from components.tools import AgentTools
from config import OPENAI_API_KEY, OPENAI_MODEL_NAME, LANGCHAIN_TRACING_V2, LANGCHAIN_API_KEY, LANGCHAIN_PROJECT

logger = logging.getLogger(__name__)

# Set up LangSmith for tracing if enabled
if LANGCHAIN_TRACING_V2:
    os.environ["LANGCHAIN_TRACING_V2"] = "true"
    os.environ["LANGCHAIN_API_KEY"] = LANGCHAIN_API_KEY
    os.environ["LANGCHAIN_PROJECT"] = LANGCHAIN_PROJECT
    logger.info(f"LangSmith tracing enabled for project: {LANGCHAIN_PROJECT}")

class AgentState(TypedDict):
    """
    Represents the state of our agent in LangGraph.
    - messages: A list of messages making up the conversation history.
    """
    messages: Annotated[Sequence[BaseMessage], add]

class LongTermMemoryAgent:
    def __init__(self):
        if not OPENAI_API_KEY:
            raise ValueError("OPENAI_API_KEY is not set. Cannot initialize LLM.")
        
        self.memory_manager = MemoryManager()
        self.agent_tools = AgentTools(self.memory_manager)
        self.tools = [
            self.agent_tools.add_episodic_memory_tool,
            self.agent_tools.retrieve_episodic_memory_tool,
            self.agent_tools.add_semantic_triple_tool,
            self.agent_tools.retrieve_semantic_memory_tool,
        ]
        
        self.llm = ChatOpenAI(model=OPENAI_MODEL_NAME, temperature=0, api_key=OPENAI_API_KEY)
        self.agent_executor = self._create_agent_executor()
        self.graph = self._build_graph()
        logger.info("LongTermMemoryAgent initialized.")

    def _create_agent_executor(self) -> AgentExecutor:
        """Creates an AgentExecutor with the LLM and tools."""
        
        # Define the prompt for the agent
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", """You are a sophisticated AI assistant capable of managing and utilizing a multi-layered memory system.
                 Your memory consists of:
                 1. Working Memory (current conversation context).
                 2. Episodic Memory (vector database for past events, interactions, and observations).
                 3. Semantic Memory (knowledge graph for consolidated, high-level facts).
                 
                 **When to use tools:**
                 - Use `add_episodic_memory_tool` to record important new information from the conversation or observations you make that might be useful later. This is for storing specific events.
                 - Use `retrieve_episodic_memory_tool` when the user asks about past interactions, specific events, or details that might have occurred previously.
                 - Use `add_semantic_triple_tool` *only* if you identify a highly generalized, factual piece of information that is worth consolidating into long-term knowledge. This tool is primarily used by the nightly consolidation job, but you can use it if a new, universally true fact is established in the conversation.
                 - Use `retrieve_semantic_memory_tool` when the user asks for general facts, definitions, or high-level knowledge that would likely be stored in a knowledge graph.
                 
                 Always strive to provide helpful and accurate information. If a question requires recalling specific past events, use episodic memory. If it requires general knowledge, use semantic memory. If you learn something new and specific, record it episodically. If you confirm a general truth, record it semantically.
                 """),
                MessagesPlaceholder(variable_name="chat_history"),
                MessagesPlaceholder(variable_name="agent_scratchpad"),
            ]
        )

        # Create the LangChain agent
        agent = create_openai_tools_agent(self.llm, self.tools, prompt)
        return AgentExecutor(agent=agent, tools=self.tools, verbose=True)

    def _agent_node(self, state: AgentState) -> Dict[str, Any]:
        """The main agent node that invokes the LLM and its tools."""
        logger.info("Agent node: Invoking agent executor.")
        result = self.agent_executor.invoke({"chat_history": state["messages"]})
        
        # AgentExecutor handles tool calls internally and always returns {output: final_answer}
        if "output" in result:
            return {"messages": [AIMessage(content=result["output"])]}
        else:
            logger.warning("Agent executor returned an unexpected format.")
            return {"messages": [AIMessage(content="An unexpected error occurred during agent execution.")]}

    def _tool_node(self, state: AgentState) -> Dict[str, Any]:
        """The tool node that executes tools."""
        logger.info("Tool node: Executing tools.")
        last_message = state["messages"][-1]
        tool_outputs = []
        
        if isinstance(last_message, AIMessage) and last_message.tool_calls:
            for tool_call in last_message.tool_calls:
                try:
                    tool_name = tool_call["name"]
                    tool_args = tool_call["args"]
                    
                    # Find and execute the tool
                    selected_tool = next((t for t in self.tools if t.name == tool_name), None)
                    if selected_tool:
                        output = selected_tool.invoke(tool_args)
                        tool_outputs.append(ToolMessage(content=str(output), tool_call_id=tool_call["id"]))
                    else:
                        tool_outputs.append(ToolMessage(content=f"Error: Tool '{tool_name}' not found.", tool_call_id=tool_call["id"]))
                        logger.error(f"Tool '{tool_name}' not found for execution.")
                except Exception as e:
                    tool_outputs.append(ToolMessage(content=f"Error executing tool '{tool_call.get('name', 'unknown')}': {e}", tool_call_id=tool_call["id"]))
                    logger.error(f"Error executing tool '{tool_call.get('name', 'unknown')}': {e}")
        else:
            logger.warning("Tool node received a message without tool calls.")

        return {"messages": tool_outputs}

    def _should_continue(self, state: AgentState) -> str:
        """Determines whether the agent should continue or end."""
        last_message = state["messages"][-1]
        if isinstance(last_message, AIMessage) and not last_message.tool_calls:
            logger.info("Should continue: Ending conversation (AI message, no tool calls).")
            return "end"
        logger.info("Should continue: Continuing with tool execution.")
        return "continue"

    def _build_graph(self):
        """Builds the LangGraph state machine."""
        workflow = StateGraph(AgentState)

        # Define nodes
        workflow.add_node("agent", self._agent_node)
        workflow.add_node("tools", self._tool_node)

        # Define edges
        workflow.set_entry_point("agent")
        workflow.add_conditional_edges(
            "agent",
            self._should_continue,
            {
                "continue": "tools",
                "end": END
            }
        )
        workflow.add_edge("tools", "agent")

        app = workflow.compile()
        logger.info("LangGraph workflow compiled successfully.")
        return app

    def run(self, input_message: str, chat_history: List[BaseMessage] = None) -> str:
        """Runs the agent with a given input message."""
        if chat_history is None:
            chat_history = []

        # Add the human message to the chat history
        current_chat_history = chat_history + [HumanMessage(content=input_message)]
        
        final_state = None
        try:
            for s in self.graph.stream({"messages": current_chat_history}):
                final_state = s
            
            if final_state and "messages" in final_state:
                ai_messages = [m for m in final_state["messages"] if isinstance(m, AIMessage)]
                if ai_messages:
                    final_response_message = next((m for m in reversed(ai_messages) if not m.tool_calls), None)
                    if final_response_message:
                        logger.info(f"Agent finished with response: {final_response_message.content}")
                        return final_response_message.content
                
            logger.warning("Agent finished without a clear final AI message.")
            return "I'm sorry, I couldn't process that request fully."

        except Exception as e:
            logger.error(f"Error running agent: {e}")
            return f"An error occurred: {e}"

# agent_app.py
import logging
import os
from typing import List

from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from components.agent import LongTermMemoryAgent
from config import LANGCHAIN_TRACING_V2, LANGCHAIN_API_KEY, LANGCHAIN_PROJECT, OPENAI_API_KEY # Import OPENAI_API_KEY for OTEL check

# OpenTelemetry imports (add to requirements.txt)
from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.langchain import LangchainInstrumentor
from opentelemetry.instrumentation.openai import OpenAIInstrumentor

logger = logging.getLogger(__name__)

# Configure LangSmith if enabled
if LANGCHAIN_TRACING_V2:
    os.environ["LANGCHAIN_TRACING_V2"] = "true"
    os.environ["LANGCHAIN_API_KEY"] = LANGCHAIN_API_KEY
    os.environ["LANGCHAIN_PROJECT"] = LANGCHAIN_PROJECT

# --- OpenTelemetry Configuration ---
def configure_opentelemetry():
    """Configures OpenTelemetry for the application."""
    logger.info("Configuring OpenTelemetry...")
    resource = Resource.create({
        "service.name": "long-term-memory-agent",
        "service.version": "1.0.0",
        "deployment.environment": os.environ.get("ENV", "development")
    })

    provider = TracerProvider(resource=resource)
    trace.set_tracer_provider(provider)

    otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318/v1/traces")
    exporter = OTLPSpanExporter(endpoint=otlp_endpoint)
    
    span_processor = BatchSpanProcessor(exporter)
    provider.add_span_processor(span_processor)

    LangchainInstrumentor().instrument()
    OpenAIInstrumentor().instrument()
    logger.info(f"OpenTelemetry configured with OTLP exporter to {otlp_endpoint}")

if OPENAI_API_KEY:
    configure_opentelemetry()
else:
    logger.warning("OPENAI_API_KEY not set. OpenTelemetry for OpenAI/LangChain will not be fully functional.")

# --- Main application logic ---
def main():
    """Main function to run the conversational agent."""
    logger.info("Starting Long Term Memory Agent application...")
    agent = LongTermMemoryAgent() # This will also set up LangSmith if enabled
    
    chat_history: List[BaseMessage] = []

    print("Hello! I am an AI agent with long-term memory. How can I help you today?")
    print("Type 'exit' to end the conversation.")

    while True:
        user_input = input("You: ")
        if user_input.lower() == 'exit':
            print("Goodbye!")
            break
        
        # Create a new trace for each user interaction
        with trace.get_tracer(__name__).start_as_current_span("user_interaction") as span:
            span.set_attribute("user.input", user_input)
            response = agent.run(user_input, chat_history)
            span.set_attribute("agent.response", response)
        
        print(f"Agent: {response}")
        
        # Update chat history with user input and agent response for the next turn
        chat_history.append(HumanMessage(content=user_input))
        chat_history.append(AIMessage(content=response))

if __name__ == "__main__":
    main()

# scripts/consolidate_memory.py
import logging
import os
import time
from typing import Dict, Any, List

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_core.output_parsers import JsonOutputParser

from components.memory import MemoryManager
from config import OPENAI_API_KEY, OPENAI_MODEL_NAME, LANGCHAIN_TRACING_V2, LANGCHAIN_API_KEY, LANGCHAIN_PROJECT

# OpenTelemetry imports
from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.langchain import LangchainInstrumentor
from opentelemetry.instrumentation.openai import OpenAIInstrumentor

logger = logging.getLogger(__name__)

# Configure LangSmith for tracing if enabled
if LANGCHAIN_TRACING_V2:
    os.environ["LANGCHAIN_TRACING_V2"] = "true"
    os.environ["LANGCHAIN_API_KEY"] = LANGCHAIN_API_KEY
    os.environ["LANGCHAIN_PROJECT"] = LANGCHAIN_PROJECT
    logger.info(f"LangSmith tracing enabled for project: {LANGCHAIN_PROJECT}")

# --- OpenTelemetry Configuration for the script ---
def configure_opentelemetry_script():
    """Configures OpenTelemetry for the consolidation script."""
    logger.info("Configuring OpenTelemetry for consolidation script...")
    resource = Resource.create({
        "service.name": "memory-consolidation-job",
        "service.version": "1.0.0",
        "deployment.environment": os.environ.get("ENV", "development")
    })
    provider = TracerProvider(resource=resource)
    trace.set_tracer_provider(provider)
    otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318/v1/traces")
    exporter = OTLPSpanExporter(endpoint=otlp_endpoint)
    span_processor = BatchSpanProcessor(exporter)
    provider.add_span_processor(span_processor)
    LangchainInstrumentor().instrument()
    OpenAIInstrumentor().instrument()
    logger.info(f"OpenTelemetry configured for script with OTLP exporter to {otlp_endpoint}")

if OPENAI_API_KEY:
    configure_opentelemetry_script()
else:
    logger.warning("OPENAI_API_KEY not set. OpenTelemetry for OpenAI/LangChain in consolidator will not be fully functional.")

class MemoryConsolidator:
    """
    A service that reviews episodic memories and promotes important facts
    to semantic memory (knowledge graph).
    """
    def __init__(self, memory_manager: MemoryManager):
        if not OPENAI_API_KEY:
            raise ValueError("OPENAI_API_KEY is not set. Cannot initialize LLM for consolidation.")
        
        self.memory_manager = memory_manager
        self.llm = ChatOpenAI(model=OPENAI_MODEL_NAME, temperature=0.3, api_key=OPENAI_API_KEY)
        self.consolidation_prompt = ChatPromptTemplate.from_messages([
            SystemMessage(content="""You are an expert knowledge engineer. Your task is to review recent episodic memories (raw events or observations) and extract general, important, and universally true facts that should be stored in a long-term knowledge graph.
            Represent these facts as (subject, predicate, object) triples.
            
            **Rules for extraction:**
            1.  Focus on generalizable knowledge, not specific temporary conversational turns.
            2.  Only extract facts that are universally true or highly significant for long-term understanding.
            3.  If a memory describes a user's personal detail (e.g., "My name is Alice"), consider if it's a fact about the user that should be remembered, or just a temporary context. For this exercise, prioritize facts about the world or persistent entities.
            4.  Each fact must be a complete (subject, predicate, object) triple.
            5.  Output only a JSON list of triples, like: `[{"subject": "X", "predicate": "Y", "object": "Z"}, ...]`. If no generalizable facts are found, return an empty JSON list: `[]`.
            """),
            HumanMessage(content="""Review the following episodic memories and extract semantic triples (subject, predicate, object) in JSON format.
            Memories:
            {memories_content}
            
            Output JSON:
            """)
        ])
        logger.info("MemoryConsolidator initialized.")

    def consolidate_recent_memories(self, limit: int = 100) -> int:
        """
        Fetches recent episodic memories, uses an LLM to extract semantic triples,
        and stores them in semantic memory.
        """
        logger.info(f"Starting memory consolidation job for the last {limit} episodic memories.")
        episodic_memories = self.memory_manager.get_all_episodic_memories(limit=limit)
        
        if not episodic_memories:
            logger.info("No new episodic memories to consolidate.")
            return 0

        memories_content = "
".join([
            f"ID: {m['id']}, Timestamp: {m['timestamp'].isoformat()}, Content: {m['content']}" 
            for m in episodic_memories
        ])

        try:
            # Wrap consolidation logic in a span for better tracing
            with trace.get_tracer(__name__).start_as_current_span("consolidate_memories_llm_extraction") as span:
                span.set_attribute("memory.count", len(episodic_memories))
                
                chain = self.consolidation_prompt | self.llm | JsonOutputParser()
                extracted_triples = chain.invoke({"memories_content": memories_content})
                
                span.set_attribute("extracted_triples.count", len(extracted_triples))
            
            if not extracted_triples:
                logger.info("LLM extracted no semantic triples from recent memories.")
                return 0

            added_count = 0
            for triple_data in extracted_triples:
                subject = triple_data.get("subject")
                predicate = triple_data.get("predicate")
                object_val = triple_data.get("object")
                
                if subject and predicate and object_val:
                    # For simplicity, we link to the first memory from the batch as source
                    source_id = str(episodic_memories[0]['id']) if episodic_memories else None
                    self.memory_manager.add_semantic_triple(subject, predicate, object_val, source_id)
                    added_count += 1
                else:
                    logger.warning(f"Skipping malformed triple: {triple_data}")
            
            logger.info(f"Consolidation complete. Added {added_count} new semantic triples.")
            return added_count

        except Exception as e:
            logger.error(f"Error during memory consolidation: {e}")
            return 0

if __name__ == "__main__":
    with trace.get_tracer(__name__).start_as_current_span("memory_consolidation_job_run"):
        memory_manager_instance = MemoryManager()
        consolidator = MemoryConsolidator(memory_manager_instance)
        
        logger.info("Running memory consolidation script...")
        consolidated_triples_count = consolidator.consolidate_recent_memories(limit=50)
        logger.info(f"Successfully consolidated {consolidated_triples_count} triples.")
        
        print("
Verifying semantic memory after consolidation:")
        retrieved_facts = memory_manager_instance.retrieve_semantic_memory("general knowledge")
        for fact in retrieved_facts:
            print(f" - {fact['subject']} {fact['predicate']} {fact['object']}")
    # Give exporter time to send spans
    time.sleep(2)