ML Vector

import os
from typing import TypedDict, Annotated, List
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_ollama import ChatOllama
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.sqlite import SqliteSaver

# Load environment variables from a .env file
from dotenv import load_dotenv
load_dotenv()

# --- Configuration ---
LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower() # 'openai' or 'ollama'
OLLAMA_MODEL_CHAT = os.getenv("OLLAMA_MODEL_CHAT", "llama2").lower() # e.g., 'llama2', 'mistral'

# --- LLM Initialization ---
def initialize_llm(provider: str, model_name: str = None, temp: float = 0.7):
    """Initializes and returns the ChatLargeLanguageModel based on provider."""
    if provider == "openai":
        if not os.getenv("OPENAI_API_KEY"):
            raise ValueError("OPENAI_API_KEY not set for OpenAI provider.")
        return ChatOpenAI(model=model_name or "gpt-3.5-turbo", temperature=temp)
    elif provider == "ollama":
        try:
            llm_instance = ChatOllama(model=model_name or OLLAMA_MODEL_CHAT, temperature=temp)
            llm_instance.invoke("Hello!", config={"stream": False}) # Test connection
            return llm_instance
        except Exception as e:
            print(f"Error connecting to Ollama LLM or model '{model_name or OLLAMA_MODEL_CHAT}' not found: {e}")
            print("Please ensure Ollama is running and the specified model is pulled (e.g., 'ollama pull llama2').")
            exit()
    else:
        raise ValueError(f"Invalid LLM provider: {provider}. Must be 'openai' or 'ollama'.")

# Initialize the chosen LLM
llm = initialize_llm(LLM_PROVIDER)
print(f"Using LLM: {LLM_PROVIDER} ({llm.model_name if hasattr(llm, 'model_name') else OLLAMA_MODEL_CHAT})\n")


# --- 1. Define Agent State ---
class AgentState(TypedDict):
    """
    Represents the state of our conversational agent.
    'messages' will automatically manage conversation history thanks to add_messages.
    'facts_gathered' is a custom field to demonstrate persistence of internal agent data.
    """
    messages: Annotated[List[BaseMessage], add_messages]
    facts_gathered: List[str] # To store internal facts gathered during conversation


# --- 2. Define Agent Node ---
def research_assistant_node(state: AgentState) -> AgentState:
    """
    A simple research assistant node that responds to user questions and
    simulates gathering specific facts.
    """
    print("\n--- Node: Research Assistant ---")
    messages = state['messages']
    # The last message in the `messages` list should be the current user's input
    last_user_message = messages[-1].content 

    # Define the LLM prompt for the research assistant
    # We include the full message history in the prompt to maintain conversation context
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful and knowledgeable research assistant. Answer the user's question clearly and concisely. If the question implies gathering information, add a short, relevant 'fact' to your internal knowledge base about the topic. Keep your responses brief and directly answer the question."),
        ("human", "{question}") # {question} is implicitly the last message in `messages` here
    ])
    
    llm_with_history = prompt | llm # Chain the prompt with the LLM

    # Invoke the LLM with the current conversation history.
    # LangGraph will automatically pass `state['messages']` as `messages` to the prompt.
    response = llm_with_history.invoke({"question": last_user_message}) 
    ai_response_content = response.content

    # Simulate gathering a 'fact' based on certain keywords in the user's question
    new_fact = ""
    if "Python" in last_user_message or "programming" in last_user_message:
        new_fact = "Python is a versatile, high-level programming language known for its readability."
    elif "history" in last_user_message or "event" in last_user_message:
        new_fact = "The fall of the Berlin Wall in 1989 was a pivotal event in modern history."
    elif "space" in last_user_message or "planet" in last_user_message:
        new_fact = "Mars is often called the 'Red Planet' due to its iron oxide-rich surface."

    # Retrieve existing facts and append the new one if unique
    facts_gathered = state.get('facts_gathered', [])
    if new_fact and new_fact not in facts_gathered:
        facts_gathered.append(new_fact)
        print(f"  (Simulated) Gathered fact: '{new_fact}'")

    print(f"  AI Response: {ai_response_content}")
    # Return the updated state. LangGraph will automatically save this.
    return {
        "messages": [AIMessage(content=ai_response_content)], # Add AI's response to messages
        "facts_gathered": facts_gathered # Update the list of facts
    }

# --- 3. Build the LangGraph Workflow with Persistence ---
print("--- Building the Stateful Conversational Agent Workflow ---")

# Configure the memory (checkpointer).
# Using ':memory:' creates an in-memory database, which is good for quick tests.
# For persistent file-based storage across runs, use a file path:
# memory = SqliteSaver.from_conn_string("sqlite:///langgraph_checkpoints.sqlite")
memory = SqliteSaver.from_conn_string(":memory:") 

workflow = StateGraph(AgentState)

# Add the research assistant node to the workflow
workflow.add_node("research_assistant", research_assistant_node)

# Set the entry point of the graph. The conversation starts with the assistant.
workflow.set_entry_point("research_assistant")

# Define the edge: After the research assistant responds, the graph reaches an END state.
# For multi-turn conversations, the persistence mechanism handles resuming.
workflow.add_edge("research_assistant", END) 

# Compile the graph and enable checkpointing by passing the 'memory' object.
conversational_app = workflow.compile(checkpointer=memory)
print("Stateful Conversational Agent workflow compiled successfully.\n")


# --- 4. Demonstrate State Persistence ---
print("--- Demonstrating State Persistence Across Turns and Sessions ---")

# --- Scenario 1: Start a new conversation (Thread ID: 'user_session_abc') ---
# This is a unique identifier for this conversation/session.
thread_id_1 = "user_session_abc"
print(f"\n--- Scenario 1: Starting new conversation (Thread ID: '{thread_id_1}') ---")

print("User: What is the capital of France?")
app_state_1_turn_1 = conversational_app.invoke(
    {"messages": [HumanMessage(content="What is the capital of France?")]},
    config={"configurable": {"thread_id": thread_id_1}}
)
print(f"Current Facts Gathered: {app_state_1_turn_1['facts_gathered']}")
print(f"Messages count after Turn 1: {len(app_state_1_turn_1['messages'])}")

print("\nUser: Tell me about Python programming.")
app_state_1_turn_2 = conversational_app.invoke(
    {"messages": [HumanMessage(content="Tell me about Python programming.")]},
    config={"configurable": {"thread_id": thread_id_1}}
)
print(f"Current Facts Gathered: {app_state_1_turn_2['facts_gathered']}")
print(f"Messages count after Turn 2: {len(app_state_1_turn_2['messages'])}")
# At this point, the entire state (messages, facts_gathered) for thread_id_1 is checkpointed.

# --- Scenario 2: Resume the conversation (Thread ID: 'user_session_abc') ---
# Simulates a user coming back later or a new turn in the same session.
# LangGraph will automatically load the last saved state for this thread_id.
print(f"\n--- Scenario 2: Resuming conversation (Thread ID: '{thread_id_1}') ---")

print("User: What historical event happened in 1989?")
app_state_1_turn_3_resumed = conversational_app.invoke(
    {"messages": [HumanMessage(content="What historical event happened in 1989?")]},
    config={"configurable": {"thread_id": thread_id_1}}
)
print(f"Current Facts Gathered: {app_state_1_turn_3_resumed['facts_gathered']}")
print(f"Messages count after Resumed Turn 3: {len(app_state_1_turn_3_resumed['messages'])}")

# Verify previous history and facts are present:
print(f"  (Verification) First user message in history: '{app_state_1_turn_3_resumed['messages'][0].content}'")
if "Python is a versatile" in str(app_state_1_turn_3_resumed['facts_gathered']):
    print("  (Verification) 'Python' fact from previous turn is present in facts_gathered.")


# --- Scenario 3: Start a new, separate conversation (Thread ID: 'new_user_session_xyz') ---
# This demonstrates that different thread IDs maintain independent states.
thread_id_2 = "new_user_session_xyz"
print(f"\n--- Scenario 3: Starting new, separate conversation (Thread ID: '{thread_id_2}') ---")
print("User: Tell me about the planet Mars.")
app_state_2_turn_1 = conversational_app.invoke(
    {"messages": [HumanMessage(content="Tell me about the planet Mars.")]},
    config={"configurable": {"thread_id": thread_id_2}}
)
print(f"Current Facts Gathered: {app_state_2_turn_1['facts_gathered']}")
print(f"Messages count after Thread 2, Turn 1: {len(app_state_2_turn_1['messages'])}")


print("\n" + "="*60)
print("--- Persistence Demonstration Complete ---")
print("Notice how 'facts_gathered' and message history correctly persisted and were loaded")
print("across invocations within the same thread_id, but remained separate for different thread_ids.")
print("To see disk persistence (data saved to a file), change ':memory:' to a file path")
print("like 'sqlite:///langgraph_checkpoints.sqlite' in the code.")
print("="*60 + "\n")