import os
from langgraph.graph import StateGraph, END
from langchain_core.runnables import Runnable
from langchain_core.messages import HumanMessage
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI

class LangGraphAgent:
    def __init__(self):
        print("LangGraphAgent initialized.")

        open_key = os.environ.get("OPEN_AI_KEY")

        # Define a simple LangGraph that uses OpenAI to answer questions
        self.llm = ChatOpenAI(model="gpt-4.1", temperature=0, api_key=open_key))
        
        def answer_fn(state):
            question = state["question"]
            print(f"[LangGraphAgent] Received: {question[:50]}...")
            response = self.llm.invoke([HumanMessage(content=question)])
            return {"answer": response.content}

        builder = StateGraph(state={"question": str, "answer": str})
        builder.add_node("answer_question", RunnableLambda(answer_fn))
        builder.set_entry_point("answer_question")
        builder.set_finish_point("answer_question")
        self.graph = builder.compile()

    def __call__(self, question: str) -> str:
        output = self.graph.invoke({"question": question})
        return output["answer"]